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
This commit is contained in:
Levi Neely
parent
46272f22a6
commit
85b0d7f6fa
20
README.md
20
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
|
||||
|
|
|
|||
|
|
@ -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 <prompt> - inject prompt into the running turn",
|
||||
" /irw <prompt> - 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 {
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
Reference in New Issue