codemode.go

// path: codemode/codemode_mcp.go
package codemode

import (
	"bytes"
	"context"
	"fmt"
	"reflect"
	"sync"
	"time"

	"github.com/mark3labs/mcp-go/mcp"

	"github.com/traefik/yaegi/interp"
	"github.com/traefik/yaegi/stdlib"
)

const CodeModeToolName = "codemode.run_code"

var (
	minimalStdlibOnce  sync.Once
	minimalStdlibCache map[string]map[string]reflect.Value
)

func getMinimalStdlib() map[string]map[string]reflect.Value {
	minimalStdlibOnce.Do(func() {
		minimalStdlibCache = map[string]map[string]reflect.Value{}

		neededPackages := []string{
			"context/context",
			"fmt/fmt",
			"reflect/reflect",
		}

		for _, pkg := range neededPackages {
			if symbols, ok := stdlib.Symbols[pkg]; ok {
				minimalStdlibCache[pkg] = symbols
			}
		}
	})

	return minimalStdlibCache
}

type CodeModeArgs struct {
	Code    string `json:"code"`
	Timeout int    `json:"timeout"`
}

type CodeModeResult struct {
	Value  any    `json:"value"`
	Stdout string `json:"stdout"`
	Stderr string `json:"stderr"`
}

type MCPClient interface {
	ListTools(
		ctx context.Context,
		request mcp.ListToolsRequest,
	) (*mcp.ListToolsResult, error)

	CallTool(
		ctx context.Context,
		request mcp.CallToolRequest,
	) (*mcp.CallToolResult, error)
}

type CodeModeMCP struct {
	client MCPClient
}

func NewCodeModeMCP(client MCPClient) *CodeModeMCP {
	return &CodeModeMCP{
		client: client,
	}
}

func newInterpreter() (*interp.Interpreter, *bytes.Buffer, *bytes.Buffer) {
	var stdout, stderr bytes.Buffer

	i := interp.New(interp.Options{
		Stdout: &stdout,
		Stderr: &stderr,
	})

	return i, &stdout, &stderr
}

func injectHelpers(
	i *interp.Interpreter,
	mcpClient MCPClient,
) error {

	if err := i.Use(getMinimalStdlib()); err != nil {
		return err
	}

	exports := interp.Exports{
		"codemode_helpers/codemode_helpers": map[string]reflect.Value{

			"CallTool": reflect.ValueOf(
				func(
					name string,
					args map[string]any,
				) (any, error) {

					res, err := mcpClient.CallTool(
						context.Background(),
						mcp.CallToolRequest{
							Params: mcp.CallToolParams{
								Name:      name,
								Arguments: args,
							},
						},
					)
					if err != nil {
						return nil, err
					}

					if res.IsError {
						var errMsg string
						for _, c := range res.Content {
							if tc, ok := c.(mcp.TextContent); ok {
								errMsg += tc.Text
							} else if tc, ok := c.(*mcp.TextContent); ok {
								errMsg += tc.Text
							}
						}
						if errMsg == "" {
							errMsg = "tool returned an error result without message"
						}
						return nil, fmt.Errorf("tool error: %s", errMsg)
					}

					if res.StructuredContent != nil {
						return res.StructuredContent, nil
					}

					var out string
					for _, c := range res.Content {
						if tc, ok := c.(mcp.TextContent); ok {
							out += tc.Text
						} else if tc, ok := c.(*mcp.TextContent); ok {
							out += tc.Text
						}
					}

					return out, nil
				},
			),

			"ListTools": reflect.ValueOf(
				func() ([]mcp.Tool, error) {

					res, err := mcpClient.ListTools(
						context.Background(),
						mcp.ListToolsRequest{},
					)
					if err != nil {
						return nil, err
					}

					return res.Tools, nil
				},
			),

			"Errorf":  reflect.ValueOf(fmt.Errorf),
			"Sprintf": reflect.ValueOf(fmt.Sprintf),
		},
	}

	return i.Use(exports)
}

func (c *CodeModeMCP) Execute(
	ctx context.Context,
	args CodeModeArgs,
) (CodeModeResult, error) {

	timeoutMs := args.Timeout
	if timeoutMs <= 0 {
		timeoutMs = 30000
	}

	ctx, cancel := context.WithTimeout(
		ctx,
		time.Duration(timeoutMs)*time.Millisecond,
	)
	defer cancel()

	i, stdout, stderr := newInterpreter()

	if err := injectHelpers(i, c.client); err != nil {
		return CodeModeResult{}, err
	}

	wrapped := wrapIntoProgram(
		preprocessUserCode(args.Code),
	)

	type evalResult struct {
		val reflect.Value
		err error
	}

	done := make(chan evalResult, 1)

	go func() {
		defer func() {
			if r := recover(); r != nil {
				done <- evalResult{
					err: fmt.Errorf(
						"interpreter panic: %v",
						r,
					),
				}
			}
		}()

		if _, err := i.Eval(wrapped); err != nil {
			done <- evalResult{
				err: fmt.Errorf(
					"compilation failed: %w",
					err,
				),
			}
			return
		}

		v, err := i.Eval(`main.run()`)

		done <- evalResult{
			val: v,
			err: err,
		}
	}()

	select {

	case <-ctx.Done():
		return CodeModeResult{
				Stdout: stdout.String(),
				Stderr: stderr.String(),
			}, fmt.Errorf(
				"execution timed out after %dms",
				timeoutMs,
			)

	case res := <-done:

		if res.err != nil {
			return CodeModeResult{
				Stdout: stdout.String(),
				Stderr: stderr.String(),
			}, res.err
		}

		return CodeModeResult{
			Value:  res.val.Interface(),
			Stdout: stdout.String(),
			Stderr: stderr.String(),
		}, nil
	}
}