From 85b0d7f6fa0eb314d969552e34bb52d009446f10 Mon Sep 17 00:00:00 2001 From: Levi Neely Date: Tue, 21 Apr 2026 13:25:32 +0200 Subject: [PATCH] remove MCP support, add execute_code tests - Remove pkg/mcp/ (client, transport, JSON-RPC handshake) - Remove mcpServer adapter and NewServer() from tools.go - Remove MCPServers/ServerConfig from config - Remove MCP server connection loop from BuildAgentEnv - Remove ListServers() from Core interface and agentCore - Remove /mcp command and help line - Rename extractMCPResult -> extractToolResult - Rename mcpToolsToBackend -> toolInfosToBackend - Add execute_test.go: 28 tests covering detectLanguage, injectArgs, ansiCEscape, ValidateCode, allowed, ReadTool, Execute, Dispatch (pipeline, parallel, tool steps), limitedWriter, rate limiting, SetEnv, SetCWD (68.2% coverage from 0%) - Add testdata/sandbox/default.yaml for real sandboxed test execution --- README.md | 20 +- pkg/agent/core.go | 78 +--- pkg/agent/core_test.go | 42 +- pkg/agent/types.go | 4 - pkg/config/config.go | 8 - pkg/mcp/client.go | 123 ------ pkg/mcp/transport.go | 59 --- pkg/tools/execute/execute_test.go | 387 ++++++++++++++++++ .../execute/testdata/sandbox/default.yaml | 19 + pkg/tools/tools.go | 49 --- 10 files changed, 431 insertions(+), 358 deletions(-) delete mode 100644 pkg/mcp/client.go delete mode 100644 pkg/mcp/transport.go create mode 100644 pkg/tools/execute/execute_test.go create mode 100644 pkg/tools/execute/testdata/sandbox/default.yaml diff --git a/README.md b/README.md index d18c4f4..d6362dd 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # ollie -A Go library for building agentic systems. Provides a sandboxed `execute_code` tool, a common LLM backend interface, MCP server support, and a skill system for domain-specific capabilities. +A Go library for building agentic systems. Provides a sandboxed `execute_code` tool, a common LLM backend interface, and a skill system for domain-specific capabilities. Works well with [anvillm](https://github.com/lneely/anvillm), which provides a skill system, tool scripts, and multi-agent infrastructure. @@ -8,7 +8,7 @@ The reference frontend is [ollie-tui](https://github.com/lneely/ollie-tui), a te ## Primitives -**`agent.Core`** — the central interface for a running agent session. Exposes `Submit` (send a prompt, stream events back), `Interrupt`, `Inject`, `Queue`/`PopQueue` (buffered prompt delivery), `State`, `Reply`, `SystemPrompt`, `Usage`, `CtxSz`, `ListModels`, `ListServers`, `CWD`/`SetCWD`, `SetSessionID`, `IsRunning`, and `Close`. +**`agent.Core`** — the central interface for a running agent session. Exposes `Submit` (send a prompt, stream events back), `Interrupt`, `Inject`, `Queue`/`PopQueue` (buffered prompt delivery), `State`, `Reply`, `SystemPrompt`, `Usage`, `CtxSz`, `ListModels`, `CWD`/`SetCWD`, `SetSessionID`, `IsRunning`, and `Close`. **`agent.Session`** — the conversation turn accumulator. Tracks message history, token usage, context compaction, and session persistence. Supports `compact` (summarize-and-truncate) and `PreCompactionSnapshot`. @@ -18,7 +18,7 @@ The reference frontend is [ollie-tui](https://github.com/lneely/ollie-tui), a te **`backend.Backend`** — the LLM interface: `ChatStream`, `Models`, `ContextLength`, `Name`, `Model`/`SetModel`, `DefaultModel`. Implementations: Ollama, OpenAI-compatible, Anthropic, Copilot, Kiro. -**`tools.Server`** — interface for a tool provider: `ListTools`, `CallTool`, `Close`. The only built-in implementation is `execute.Server`. MCP servers are wrapped via `tools.NewServer(client)`. Custom servers implement this interface directly. +**`tools.Server`** — interface for a tool provider: `ListTools`, `CallTool`, `Close`. The only built-in implementation is `execute.Server`. Custom servers implement this interface directly. **`tools.Dispatcher`** — routes tool calls to the correct server by name. Built via `NewDispatcher` or `NewDispatcherFunc` (from a map of `Decl` factories). Supports `AddServer`, `GetServer`, `ListTools`, `Dispatch`. @@ -28,7 +28,6 @@ The reference frontend is [ollie-tui](https://github.com/lneely/ollie-tui), a te pkg/agent/ — Core interface, agent loop, session management pkg/backend/ — Backend interface + implementations (Ollama, OpenAI, Anthropic, Copilot, Kiro) pkg/config/ — Config struct and loader -pkg/mcp/ — MCP client pkg/tools/ — Server and Dispatcher interfaces; tool definitions pkg/tools/execute/ — execute.Server: execute_code ``` @@ -47,15 +46,6 @@ No build step — ollie-core is a library. ```json { - "mcpServers": { - "my-server": { - "command": "my-mcp-server", - "args": [], - "env": { - "API_TOKEN": "${API_TOKEN}" - } - } - }, "hooks": { "agentSpawn": [ "$OLLIE/x/prime sys-base", @@ -68,8 +58,6 @@ No build step — ollie-core is a library. } ``` -MCP server `env` values support `${VAR}` expansion from the parent environment. - Hook values accept a string or an array of strings. Commands run in order; each command's stdout is appended to the system prompt context. Exit code 2 blocks the triggering action; any other non-zero exit is a non-blocking warning. ### System prompt @@ -92,8 +80,6 @@ One built-in tool via `execute.Server`: Additional capabilities (file I/O, memory, reasoning, task planning, sub-agents) are implemented as tool scripts in `OLLIE_TOOLS_PATH`, invoked via `execute_code` `{tool}` steps. -MCP server tools are discovered at startup and available alongside the built-ins. - Each tool server exports a `Decl` function that returns a `func() tools.Server` factory. `execute.Decl(cwd)` accepts a working directory used as `cmd.Dir` for sandboxed commands and for `{CWD}` expansion in the sandbox config; pass `""` to fall back to `os.Getwd()`. Frontends register servers by passing Decl results to `tools.NewDispatcherFunc`. Adding a new tool means implementing `tools.Server`, exporting a `Decl` function, and registering it — no frontend changes required. ## Skills diff --git a/pkg/agent/core.go b/pkg/agent/core.go index e6505f4..22cbb00 100644 --- a/pkg/agent/core.go +++ b/pkg/agent/core.go @@ -20,7 +20,6 @@ import ( "ollie/pkg/backend" "ollie/pkg/config" olog "ollie/pkg/log" - "ollie/pkg/mcp" "ollie/pkg/paths" "ollie/pkg/tools" ) @@ -48,22 +47,6 @@ type AgentEnv struct { func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string) AgentEnv { var messages []string - if cfg != nil { - for name, serverCfg := range cfg.MCPServers { - if serverCfg.Disabled || serverCfg.Command == "" { - continue - } - transport := mcp.NewSTDIOTransport(serverCfg.Command, serverCfg.Args, serverCfg.Env) - client, err := transport.Connect() - if err != nil { - messages = append(messages, fmt.Sprintf("MCP %s: failed to connect: %v", name, err)) - continue - } - d.AddServer(name, tools.NewServer(client)) - messages = append(messages, fmt.Sprintf("MCP %s: connected", name)) - } - } - allToolInfos, listErr := d.ListTools() if listErr != nil { messages = append(messages, fmt.Sprintf("list tools: %v", listErr)) @@ -74,7 +57,7 @@ func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string) AgentEnv serverOf[t.Name] = t.Server } - allTools := mcpToolsToBackend(allToolInfos) + allTools := toolInfosToBackend(allToolInfos) hooks := Hooks{} agentPrompt := "" @@ -118,7 +101,7 @@ func BuildAgentEnv(cfg *config.Config, d tools.Dispatcher, cwd string) AgentEnv if err != nil { return "", err } - text, isErr := extractMCPResult(raw) + text, isErr := extractToolResult(raw) if isErr { return "", fmt.Errorf("%s", text) } @@ -646,50 +629,6 @@ func (s *agentCore) ListModels() string { return strings.Join(models, "\n") } -func (s *agentCore) ListServers() string { - clog.Debug("ListServers()") - if s.dispatcher == nil { - return "no dispatcher" - } - allTools, err := s.dispatcher.ListTools() - if err != nil { - return "error: " + err.Error() - } - if len(allTools) == 0 { - return "no servers registered" - } - - // Group tools by server, preserving first-seen order. - type serverEntry struct { - name string - tools []tools.ToolInfo - } - index := map[string]int{} - var servers []serverEntry - for _, t := range allTools { - i, ok := index[t.Server] - if !ok { - i = len(servers) - index[t.Server] = i - servers = append(servers, serverEntry{name: t.Server}) - } - servers[i].tools = append(servers[i].tools, t) - } - - var sb strings.Builder - for si, srv := range servers { - if si > 0 { - sb.WriteByte('\n') - } - fmt.Fprintf(&sb, "%s\n", srv.name) - for _, t := range srv.tools { - desc := firstSentence(t.Description) - fmt.Fprintf(&sb, " %-22s %s\n", t.Name, desc) - } - } - return strings.TrimRight(sb.String(), "\n") -} - // firstSentence returns the first sentence of s (up to the first period or // newline), trimmed. Falls back to s truncated at 80 chars if no sentence end // is found. @@ -1226,10 +1165,6 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve "/skills": func(args []string) { listMountDir("sk") }, "/tools": func(args []string) { listMountDir("t") }, - "/mcp": func(args []string) { - handler(infoEvent(s.ListServers())) - }, - "/sp": func(args []string) { handler(infoEvent(s.loopcfg.systemPrompt)) }, @@ -1245,7 +1180,6 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve " /models - list available models", " /skills - list available skills", " /tools - list available tools", - " /mcp - list registered tool servers and their tools", " /cwd [path] - show or change working directory", " /i - inject prompt into the running turn", " /irw - rewrite the pending inject", @@ -1276,9 +1210,9 @@ func (s *agentCore) handleCommand(ctx context.Context, input string, handler Eve return true } -func mcpToolsToBackend(mcpTools []tools.ToolInfo) []backend.Tool { - out := make([]backend.Tool, len(mcpTools)) - for i, t := range mcpTools { +func toolInfosToBackend(infos []tools.ToolInfo) []backend.Tool { + out := make([]backend.Tool, len(infos)) + for i, t := range infos { out[i] = backend.Tool{ Name: t.Name, Description: t.Description, @@ -1288,7 +1222,7 @@ func mcpToolsToBackend(mcpTools []tools.ToolInfo) []backend.Tool { return out } -func extractMCPResult(raw json.RawMessage) (text string, isError bool) { +func extractToolResult(raw json.RawMessage) (text string, isError bool) { var result struct { IsError bool `json:"isError"` Content []struct { diff --git a/pkg/agent/core_test.go b/pkg/agent/core_test.go index 42a700a..73c1a89 100644 --- a/pkg/agent/core_test.go +++ b/pkg/agent/core_test.go @@ -1389,16 +1389,6 @@ func TestCommand_History_WithSession(t *testing.T) { } } -// --- /mcp --- - -func TestCommand_MCP(t *testing.T) { - c := newCore(t, nil, nil) - evs := collectEvents(context.Background(), c, "/mcp") - if len(byRole(evs, "info")) == 0 { - t.Error("/mcp: no info events emitted") - } -} - // --- /sp --- func TestCommand_SP(t *testing.T) { @@ -1940,11 +1930,11 @@ func TestCommand_Backend_Error(t *testing.T) { } } -// --- extractMCPResult --- +// --- extractToolResult --- -func TestExtractMCPResult_Success(t *testing.T) { +func TestExtractToolResult_Success(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"text","text":"hello"}]}`) - text, isErr := extractMCPResult(raw) + text, isErr := extractToolResult(raw) if text != "hello" { t.Errorf("text = %q; want %q", text, "hello") } @@ -1953,9 +1943,9 @@ func TestExtractMCPResult_Success(t *testing.T) { } } -func TestExtractMCPResult_IsError(t *testing.T) { +func TestExtractToolResult_IsError(t *testing.T) { raw := json.RawMessage(`{"isError":true,"content":[{"type":"text","text":"something failed"}]}`) - text, isErr := extractMCPResult(raw) + text, isErr := extractToolResult(raw) if text != "something failed" { t.Errorf("text = %q; want %q", text, "something failed") } @@ -1964,25 +1954,25 @@ func TestExtractMCPResult_IsError(t *testing.T) { } } -func TestExtractMCPResult_MultipleContentItems(t *testing.T) { +func TestExtractToolResult_MultipleContentItems(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]}`) - text, _ := extractMCPResult(raw) + text, _ := extractToolResult(raw) if text != "a\nb" { t.Errorf("text = %q; want %q", text, "a\nb") } } -func TestExtractMCPResult_NonTextItemsSkipped(t *testing.T) { +func TestExtractToolResult_NonTextItemsSkipped(t *testing.T) { raw := json.RawMessage(`{"isError":false,"content":[{"type":"image","text":"ignored"},{"type":"text","text":"kept"}]}`) - text, _ := extractMCPResult(raw) + text, _ := extractToolResult(raw) if text != "kept" { t.Errorf("text = %q; want %q", text, "kept") } } -func TestExtractMCPResult_InvalidJSON(t *testing.T) { +func TestExtractToolResult_InvalidJSON(t *testing.T) { raw := json.RawMessage(`not json`) - text, isErr := extractMCPResult(raw) + text, isErr := extractToolResult(raw) if text != "not json" { t.Errorf("text = %q; want raw input on parse failure", text) } @@ -1991,15 +1981,15 @@ func TestExtractMCPResult_InvalidJSON(t *testing.T) { } } -// --- mcpToolsToBackend --- +// --- toolInfosToBackend --- -func TestMcpToolsToBackend(t *testing.T) { +func TestToolInfosToBackend(t *testing.T) { schema := json.RawMessage(`{"type":"object"}`) infos := []tools.ToolInfo{ {Name: "tool_a", Description: "does A.", InputSchema: schema}, {Name: "tool_b", Description: "does B.", InputSchema: schema}, } - got := mcpToolsToBackend(infos) + got := toolInfosToBackend(infos) if len(got) != 2 { t.Fatalf("len = %d; want 2", len(got)) } @@ -2011,8 +2001,8 @@ func TestMcpToolsToBackend(t *testing.T) { } } -func TestMcpToolsToBackend_Empty(t *testing.T) { - got := mcpToolsToBackend(nil) +func TestToolInfosToBackend_Empty(t *testing.T) { + got := toolInfosToBackend(nil) if len(got) != 0 { t.Errorf("expected empty slice, got %v", got) } diff --git a/pkg/agent/types.go b/pkg/agent/types.go index b21c66f..61203ed 100644 --- a/pkg/agent/types.go +++ b/pkg/agent/types.go @@ -83,10 +83,6 @@ type Core interface { // ListModels returns available model names, one per line. ListModels() string - // ListServers returns all registered tool servers and their tools, - // grouped by server name. - ListServers() string - // CWD returns the current working directory used for tool execution. CWD() string diff --git a/pkg/config/config.go b/pkg/config/config.go index f162099..8dd7693 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -24,7 +24,6 @@ func (h *HookCmds) UnmarshalJSON(data []byte) error { } type Config struct { - MCPServers map[string]ServerConfig `json:"mcpServers,omitempty"` Hooks map[string]HookCmds `json:"hooks,omitempty"` Prompt string `json:"prompt,omitempty"` TrustedTools []string `json:"trustedTools,omitempty"` @@ -34,13 +33,6 @@ type Config struct { PresencePenalty *float64 `json:"presencePenalty,omitempty"` } -type ServerConfig struct { - Command string `json:"command,omitempty"` - Args []string `json:"args,omitempty"` - Env map[string]string `json:"env,omitempty"` - Disabled bool `json:"disabled,omitempty"` -} - func Load(path string) (*Config, error) { data, err := os.ReadFile(path) if err != nil { diff --git a/pkg/mcp/client.go b/pkg/mcp/client.go deleted file mode 100644 index 897056b..0000000 --- a/pkg/mcp/client.go +++ /dev/null @@ -1,123 +0,0 @@ -package mcp - -import ( - "bufio" - "encoding/json" - "fmt" - "io" - "os/exec" - "sync" -) - -// Message is a JSON-RPC 2.0 message. -type Message struct { - JSONRPC string `json:"jsonrpc"` - ID int `json:"id,omitempty"` - Method string `json:"method,omitempty"` - Params interface{} `json:"params,omitempty"` - Result json.RawMessage `json:"result,omitempty"` - Error *RPCError `json:"error,omitempty"` -} - -// RPCError is a JSON-RPC error object. -type RPCError struct { - Code int `json:"code"` - Message string `json:"message"` -} - -func (e *RPCError) Error() string { - return fmt.Sprintf("rpc error %d: %s", e.Code, e.Message) -} - -// Client is a JSON-RPC client over a pair of streams (typically stdin/stdout -// of a subprocess). -type Client struct { - mu sync.Mutex - wc io.WriteCloser - w io.Writer - sc *bufio.Scanner - nextID int - cmd *exec.Cmd -} - -// NewClient creates a Client from a reader and writer. -func NewClient(r io.Reader, w io.Writer) *Client { - return &Client{w: w, sc: bufio.NewScanner(r)} -} - -// newClientWithProcess creates a Client that owns the given subprocess. -// Close() will shut it down. -func newClientWithProcess(r io.Reader, wc io.WriteCloser, cmd *exec.Cmd) *Client { - return &Client{wc: wc, w: wc, sc: bufio.NewScanner(r), cmd: cmd} -} - -// Close shuts down the underlying subprocess: closes its stdin so it receives -// EOF, then kills the process if it is still running. -func (c *Client) Close() error { - var err error - if c.wc != nil { - err = c.wc.Close() - } - if c.cmd != nil && c.cmd.Process != nil { - c.cmd.Process.Kill() - c.cmd.Wait() // reap the process; ignore errors - } - return err -} - -// Notify sends a JSON-RPC notification (no ID, no response expected). -func (c *Client) Notify(method string, params interface{}) error { - c.mu.Lock() - defer c.mu.Unlock() - - msg := Message{ - JSONRPC: "2.0", - Method: method, - Params: params, - } - data, err := json.Marshal(msg) - if err != nil { - return err - } - _, err = fmt.Fprintf(c.w, "%s\n", data) - return err -} - -// Call sends a JSON-RPC request and returns the result. -func (c *Client) Call(method string, params interface{}) (json.RawMessage, error) { - c.mu.Lock() - defer c.mu.Unlock() - - c.nextID++ - id := c.nextID - - msg := Message{ - JSONRPC: "2.0", - ID: id, - Method: method, - Params: params, - } - data, err := json.Marshal(msg) - if err != nil { - return nil, err - } - if _, err := fmt.Fprintf(c.w, "%s\n", data); err != nil { - return nil, err - } - - if !c.sc.Scan() { - if err := c.sc.Err(); err != nil { - return nil, err - } - return nil, io.EOF - } - - var resp Message - if err := json.Unmarshal(c.sc.Bytes(), &resp); err != nil { - return nil, err - } - if resp.Error != nil { - return nil, resp.Error - } - return resp.Result, nil -} diff --git a/pkg/mcp/transport.go b/pkg/mcp/transport.go deleted file mode 100644 index 52006f2..0000000 --- a/pkg/mcp/transport.go +++ /dev/null @@ -1,59 +0,0 @@ -package mcp - -import ( - "fmt" - "os" - "os/exec" -) - -// STDIOTransport launches a subprocess and connects a Client to its stdin/stdout. -type STDIOTransport struct { - command string - args []string - env map[string]string -} - -// NewSTDIOTransport creates a transport that will launch the given command. -// env values may contain ${VAR} references expanded against the process environment. -func NewSTDIOTransport(command string, args []string, env map[string]string) *STDIOTransport { - return &STDIOTransport{command: command, args: args, env: env} -} - -// Connect launches the subprocess and returns a connected Client. -func (t *STDIOTransport) Connect() (*Client, error) { - cmd := exec.Command(t.command, t.args...) - if len(t.env) > 0 { - cmd.Env = os.Environ() - for k, v := range t.env { - cmd.Env = append(cmd.Env, fmt.Sprintf("%s=%s", k, os.ExpandEnv(v))) - } - } - - stdin, err := cmd.StdinPipe() - if err != nil { - return nil, err - } - stdout, err := cmd.StdoutPipe() - if err != nil { - return nil, err - } - if err := cmd.Start(); err != nil { - return nil, err - } - - client := newClientWithProcess(stdout, stdin, cmd) - - // MCP initialization handshake. - if _, err := client.Call("initialize", map[string]interface{}{ - "protocolVersion": "2024-11-05", - "capabilities": map[string]interface{}{}, - "clientInfo": map[string]interface{}{"name": "ollie", "version": "0.1"}, - }); err != nil { - return nil, fmt.Errorf("initialize: %w", err) - } - if err := client.Notify("notifications/initialized", nil); err != nil { - return nil, fmt.Errorf("notifications/initialized: %w", err) - } - - return client, nil -} diff --git a/pkg/tools/execute/execute_test.go b/pkg/tools/execute/execute_test.go new file mode 100644 index 0000000..3f8b7a5 --- /dev/null +++ b/pkg/tools/execute/execute_test.go @@ -0,0 +1,387 @@ +package execute + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +// TestMain points OLLIE_CFG_PATH at testdata so loadSandboxConfig uses the +// minimal test sandbox config rather than the user's ~/.config/ollie/sandbox/. +func TestMain(m *testing.M) { + wd, err := os.Getwd() + if err != nil { + panic(err) + } + os.Setenv("OLLIE_CFG_PATH", filepath.Join(wd, "testdata")) + os.Exit(m.Run()) +} + +// newServer returns a Server with Confirm always returning true. +func newServer(t *testing.T) *Server { + t.Helper() + s := New(t.TempDir()) + s.Confirm = func(string) bool { return true } + return s +} + +// callCode is a helper that invokes execute_code via Dispatch. +func callCode(t *testing.T, s *Server, steps []map[string]any, extra ...map[string]any) (string, error) { + t.Helper() + payload := map[string]any{"steps": steps} + if len(extra) > 0 { + for k, v := range extra[0] { + payload[k] = v + } + } + raw, _ := json.Marshal(payload) + return s.Dispatch(context.Background(), "execute_code", raw) +} + +// ---- detectLanguage ---- + +func TestDetectLanguage(t *testing.T) { + cases := []struct{ code, want string }{ + {"echo hi", "bash"}, + {"#!/bin/bash\necho hi", "bash"}, + {"#!/usr/bin/env python3\nprint(1)", "python3"}, + {"#!/usr/bin/python\nprint(1)", "python3"}, + {"#!/usr/bin/perl\nprint 1", "perl"}, + {"#!/usr/bin/awk -f\n{print}", "awk"}, + {"#!/usr/bin/env gawk\n{print}", "awk"}, + {"#!/usr/bin/sed -f\ns/a/b/", "sed"}, + {"#!/usr/bin/ed\n,p", "ed"}, + {"#!/usr/bin/env jq\n.", "jq"}, + {"#!/usr/bin/env lua\nprint(1)", "lua"}, + } + for _, c := range cases { + if got := detectLanguage(c.code); got != c.want { + t.Errorf("detectLanguage(%q) = %q, want %q", c.code[:min(20, len(c.code))], got, c.want) + } + } +} + +func min(a, b int) int { + if a < b { + return a + } + return b +} + +// ---- ansiCEscape ---- + +func TestAnsiCEscape(t *testing.T) { + got := ansiCEscape("a\\b'c\nd\te") + want := `a\\b\'c\nd\te` + if got != want { + t.Errorf("got %q, want %q", got, want) + } +} + +// ---- injectArgs ---- + +func TestInjectArgsBash(t *testing.T) { + out := injectArgs("bash", "myscript", []string{"hello", "world"}, "echo $1 $2") + if !strings.HasPrefix(out, "set -- ") { + t.Errorf("bash inject should start with 'set --', got: %q", out) + } + if !strings.Contains(out, "echo $1 $2") { + t.Errorf("bash inject should contain original code") + } +} + +func TestInjectArgsPython(t *testing.T) { + out := injectArgs("python3", "s", []string{"a"}, "print(sys.argv)") + if !strings.HasPrefix(out, "import sys\n") { + t.Errorf("python inject should start with 'import sys', got: %q", out) + } +} + +// ---- ToolsPath / PluginsPath ---- + +func TestToolsPathEnv(t *testing.T) { + t.Setenv("OLLIE_TOOLS_PATH", "/custom/tools:/other") + if got := ToolsPath(); got != "/custom/tools" { + t.Errorf("got %q, want /custom/tools", got) + } +} + +func TestPluginsPathEnv(t *testing.T) { + t.Setenv("OLLIE_PLUGINS_PATH", "/custom/plugins:/other") + if got := PluginsPath(); got != "/custom/plugins" { + t.Errorf("got %q, want /custom/plugins", got) + } +} + +// ---- ReadTool ---- + +func TestReadTool(t *testing.T) { + dir := t.TempDir() + t.Setenv("OLLIE_TOOLS_PATH", dir) + if err := os.WriteFile(filepath.Join(dir, "mytool"), []byte("#!/bin/bash\necho ok"), 0644); err != nil { + t.Fatal(err) + } + code, err := ReadTool("mytool") + if err != nil { + t.Fatal(err) + } + if !strings.Contains(code, "echo ok") { + t.Errorf("unexpected content: %q", code) + } +} + +func TestReadToolNotFound(t *testing.T) { + t.Setenv("OLLIE_TOOLS_PATH", t.TempDir()) + _, err := ReadTool("nonexistent") + if err == nil || !strings.Contains(err.Error(), "tool not found") { + t.Errorf("expected 'tool not found' error, got %v", err) + } +} + +func TestReadToolInvalidName(t *testing.T) { + for _, name := range []string{"../etc/passwd", "foo/bar"} { + _, err := ReadTool(name) + if err == nil { + t.Errorf("expected error for name %q", name) + } + } +} + +// ---- ValidateCode ---- + +func TestValidateCodeDangerous(t *testing.T) { + s := newServer(t) + cases := []struct{ code, lang string }{ + {"sudo rm -rf /", "bash"}, + {"rm -rf /home", "bash"}, + {"mkfs /dev/sda", "bash"}, + {"dd if=/dev/zero of=/dev/sda", "bash"}, + {"shutil.rmtree('/')", "python3"}, + } + for _, c := range cases { + if err := s.ValidateCode(c.code, c.lang); err == nil { + t.Errorf("expected dangerous pattern error for %q (%s)", c.code, c.lang) + } + } +} + +func TestValidateCodeSafe(t *testing.T) { + s := newServer(t) + if err := s.ValidateCode("echo hello", "bash"); err != nil { + t.Errorf("unexpected error: %v", err) + } +} + +// ---- allowed ---- + +func TestAllowedWithConfirm(t *testing.T) { + s := New(t.TempDir()) + s.Confirm = func(string) bool { return true } + if !s.allowed("execute_code", "echo hi") { + t.Error("expected allowed=true when Confirm returns true") + } +} + +func TestAllowedDeniedNoConfirm(t *testing.T) { + s := New(t.TempDir()) + if s.allowed("execute_code", "echo hi") { + t.Error("expected allowed=false when Confirm is nil") + } +} + +func TestAllowedTrustedTool(t *testing.T) { + s := New(t.TempDir()) + s.SetTrustedTools([]string{"execute_code"}) + if !s.allowed("execute_code", "anything") { + t.Error("expected trusted tool to be allowed without Confirm") + } +} + +// ---- Execute (integration, requires bash) ---- + +func TestExecuteSimpleBash(t *testing.T) { + s := newServer(t) + out, err := s.Execute(context.Background(), "echo hello", "bash", 10, "default", true) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.TrimSpace(out) != "hello" { + t.Errorf("got %q, want %q", out, "hello") + } +} + +func TestExecuteStdin(t *testing.T) { + s := newServer(t) + out, err := s.executeWithStdin(context.Background(), "cat", "bash", 10, "default", true, "piped input") + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.TrimSpace(out) != "piped input" { + t.Errorf("got %q, want %q", out, "piped input") + } +} + +func TestExecuteTimeout(t *testing.T) { + s := newServer(t) + _, err := s.Execute(context.Background(), "sleep 10", "bash", 1, "default", true) + if err == nil || !strings.Contains(err.Error(), "timeout") { + t.Errorf("expected timeout error, got %v", err) + } +} + +func TestExecuteNonZeroExit(t *testing.T) { + s := newServer(t) + _, err := s.Execute(context.Background(), "exit 1", "bash", 10, "default", true) + if err == nil { + t.Error("expected error for non-zero exit") + } +} + +func TestExecuteUnsupportedLanguage(t *testing.T) { + s := newServer(t) + _, err := s.Execute(context.Background(), "code", "cobol", 10, "default", true) + if err == nil || !strings.Contains(err.Error(), "unsupported language") { + t.Errorf("expected unsupported language error, got %v", err) + } +} + +// ---- Dispatch / execute_code ---- + +func TestDispatchUnknownTool(t *testing.T) { + s := newServer(t) + _, err := s.Dispatch(context.Background(), "unknown_tool", json.RawMessage(`{}`)) + if err == nil { + t.Error("expected error for unknown tool") + } +} + +func TestDispatchNoSteps(t *testing.T) { + s := newServer(t) + _, err := callCode(t, s, []map[string]any{}) + if err == nil || !strings.Contains(err.Error(), "at least one step") { + t.Errorf("expected 'at least one step' error, got %v", err) + } +} + +func TestDispatchDenied(t *testing.T) { + s := New(t.TempDir()) // no Confirm, not trusted → denied + _, err := callCode(t, s, []map[string]any{{"code": "echo hi"}}) + if err == nil || !strings.Contains(err.Error(), "denied") { + t.Errorf("expected denied error, got %v", err) + } +} + +func TestDispatchSingleStep(t *testing.T) { + s := newServer(t) + out, err := callCode(t, s, []map[string]any{{"code": "echo single"}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(out, "single") { + t.Errorf("got %q, want output containing 'single'", out) + } +} + +func TestDispatchPipeline(t *testing.T) { + s := newServer(t) + out, err := callCode(t, s, []map[string]any{ + {"code": "printf 'a\\nb\\nc'"}, + {"code": "grep b"}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.TrimSpace(out) != "b" { + t.Errorf("got %q, want %q", out, "b") + } +} + +func TestDispatchParallel(t *testing.T) { + s := newServer(t) + out, err := callCode(t, s, []map[string]any{ + {"parallel": []map[string]any{ + {"code": "echo A"}, + {"code": "echo B"}, + }}, + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(out, "A") || !strings.Contains(out, "B") { + t.Errorf("parallel output missing A or B: %q", out) + } +} + +func TestDispatchToolStep(t *testing.T) { + dir := t.TempDir() + t.Setenv("OLLIE_TOOLS_PATH", dir) + if err := os.WriteFile(filepath.Join(dir, "greet"), []byte("#!/bin/bash\necho hello-from-tool"), 0755); err != nil { + t.Fatal(err) + } + s := newServer(t) + out, err := callCode(t, s, []map[string]any{{"tool": "greet"}}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !strings.Contains(out, "hello-from-tool") { + t.Errorf("got %q, want output containing 'hello-from-tool'", out) + } +} + +// ---- limitedWriter ---- + +func TestLimitedWriter(t *testing.T) { + var buf strings.Builder + lw := &limitedWriter{w: &buf, limit: 5} + lw.Write([]byte("hello world")) + if buf.String() != "hello" { + t.Errorf("got %q, want %q", buf.String(), "hello") + } + if !lw.truncated { + t.Error("expected truncated=true") + } +} + +// ---- rate limiting ---- + +func TestRateLimitBlocks(t *testing.T) { + s := newServer(t) + // Trigger maxFailures validation failures. + for i := 0; i < maxFailures; i++ { + s.recordValidationFailure() + } + if err := s.checkRateLimit(); err == nil { + t.Error("expected rate limit error after max failures") + } +} + +// ---- SetEnv / SetCWD ---- + +func TestSetEnvInjected(t *testing.T) { + s := newServer(t) + s.SetEnv("MY_TEST_VAR", "injected_value") + out, err := s.Execute(context.Background(), "echo $MY_TEST_VAR", "bash", 10, "default", true) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if strings.TrimSpace(out) != "injected_value" { + t.Errorf("got %q, want %q", out, "injected_value") + } +} + +func TestSetCWD(t *testing.T) { + dir := t.TempDir() + s := newServer(t) + s.SetCWD(dir) + out, err := s.Execute(context.Background(), "pwd", "bash", 10, "default", true) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + // TempDir may use symlinks; compare base names. + if !strings.Contains(out, filepath.Base(dir)) { + t.Errorf("pwd output %q doesn't contain expected dir %q", out, dir) + } +} diff --git a/pkg/tools/execute/testdata/sandbox/default.yaml b/pkg/tools/execute/testdata/sandbox/default.yaml new file mode 100644 index 0000000..7c9cc84 --- /dev/null +++ b/pkg/tools/execute/testdata/sandbox/default.yaml @@ -0,0 +1,19 @@ +filesystem: + rwx: + - "{CWD}" + - "/tmp" + - "/usr" + - "/bin" + - "/lib" + - "/lib64" + - "/etc" + ro: + - "/proc" + - "/dev" +env: + - HOME + - PATH + - TMPDIR + - MY_TEST_VAR +network: + unrestricted: true diff --git a/pkg/tools/tools.go b/pkg/tools/tools.go index ae1d0a2..7a9b6ff 100644 --- a/pkg/tools/tools.go +++ b/pkg/tools/tools.go @@ -6,8 +6,6 @@ import ( "context" "encoding/json" "fmt" - - "ollie/pkg/mcp" ) // ToolInfo describes a tool provided by a server. @@ -117,50 +115,3 @@ type EnvSetter interface { type TrustedToolsSetter interface { SetTrustedTools(tools []string) } - -// NewServer wraps an mcp.Client as a Server. -func NewServer(client *mcp.Client) Server { - return &mcpServer{client: client} -} - -// mcpServer wraps an mcp.Client as a Server. -type mcpServer struct { - client *mcp.Client -} - -func (m *mcpServer) ListTools() ([]ToolInfo, error) { - result, err := m.client.Call("tools/list", nil) - if err != nil { - return nil, err - } - var resp struct { - Tools []struct { - Name string `json:"name"` - Description string `json:"description"` - InputSchema json.RawMessage `json:"inputSchema"` - } `json:"tools"` - } - if err := json.Unmarshal(result, &resp); err != nil { - return nil, err - } - var tools []ToolInfo - for _, t := range resp.Tools { - tools = append(tools, ToolInfo{ - Name: t.Name, - Description: t.Description, - InputSchema: t.InputSchema, - }) - } - return tools, nil -} - -func (m *mcpServer) CallTool(_ context.Context, tool string, args json.RawMessage) (json.RawMessage, error) { - return m.client.Call("tools/call", map[string]any{ - "name": tool, - "arguments": args, - }) -} - -func (m *mcpServer) Close() { - m.client.Close() -}