From fbe8059cf57c46e13d75fef6929ac4cee7cedf3c Mon Sep 17 00:00:00 2001 From: Ollie Agent Date: Sun, 9 Aug 2026 21:45:56 +0200 Subject: [PATCH] maintainability: docs, contracts, structural cleanup MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - doc/boot-sequence.md: init chain from main → fs.NewRoot → session.Init - Event protocol: comprehensive table on Event struct (all roles, semantics) - RPC server moved from cmd/ollie-remote into toolsrv/rpcserver.go (binary is now a 93-line thin wrapper) - SetSessionEnv folded into NewAgent (no post-construction wiring needed) - AgentCfg → AgentParams (distinguish from AgentConfig JSON schema) - ToolResult/ToolResultContent shared types in toolsrv/rpcwire.go - toolsrv/shell.go → toolsrv/exec.go (name matches purpose) - Resume unified: single path, always through Keeper - AGENTS.md: fix fs/builder.go → fsedsl/builder.go, handlers_*.go - fsedsl/Tree: document dual-mode (EDSL lazy vs Manual Mount) - Truncation stack documented at defaultToolResultMaxBytes - Integration test rewritten to use srv.ServeRPC directly --- AGENTS.md | 10 +- agent/agent.go | 38 +- agent/loop.go | 10 + agent/runtime.go | 24 +- cmd/ollie-remote/main.go | 272 +------------- doc/boot-sequence.md | 54 +++ fsedsl/tree.go | 17 +- session/persist.go | 2 +- session/registry.go | 2 +- session/session.go | 61 ++-- toolsrv/{shell.go => exec.go} | 0 toolsrv/rpc_integration_test.go | 625 +++++--------------------------- toolsrv/rpcserver.go | 263 ++++++++++++++ toolsrv/rpcwire.go | 29 ++ toolsrv/server.go | 10 +- 15 files changed, 539 insertions(+), 878 deletions(-) create mode 100644 doc/boot-sequence.md rename toolsrv/{shell.go => exec.go} (100%) create mode 100644 toolsrv/rpcserver.go diff --git a/AGENTS.md b/AGENTS.md index d6e3354..7d22769 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,10 +15,12 @@ ollie/ ← you are here ├── agent/ (Go) Agent loop, history, hooks, commands ├── toolsrv/ (Go) Tool server, registry, sandboxed execution ├── session/ (Go) Session lifecycle, config, persistence -├── fs/ (Go) 9P filesystem tree (EDSL-declared, flat package) +├── fs/ (Go) 9P filesystem tree (EDSL-declared) │ ├── spec.go Single source of truth — entire namespace declared here │ ├── fsnode.go Core FsNodeDecl type + Dir/Leaf/TemplateDir constructors -│ └── builder.go BuildTree — walks spec, validates, wires handlers +│ └── handlers_*.go Read/write handlers grouped by scope (session, agent, root) +├── fsedsl/ (Go) Filesystem declaration EDSL (generic, reusable) +│ └── builder.go BuildTree — walks spec, validates, produces *Tree ├── cmd/ (Go) Binaries (olliesrv, ollie-9p, ollie-remote) ├── tools/ (Go) Tool implementations: │ └── lsp/ LSP bridge + cmd/ binaries (definition, hover, rename, ...) @@ -98,7 +100,7 @@ just test-remote 4. **Sandbox** (`sandbox/`): Landlock-based. Config in `sandbox/*.yaml` defines filesystem access per profile. Escape via bypass broker (`bypass/`). 5. **Backends** (`backend/`): Ollama, OpenAI-compatible, Anthropic, Copilot, Kiro, Gemini, CodeWhisperer. Selectable per-session. 6. **Prompts assembled at runtime**: Agent JSON `prompt` array specifies which prompt files to concatenate. Static prompt files can be included directly; the base system prompt is embedded in the binary and always prepended. -7. **9P namespace declared via EDSL**: The entire filesystem is a single recursive `FsNodeDecl` tree in `fs/spec.go`, built by `BuildTree()` in `fs/builder.go`. The `fs/` package is flat — handlers in `handlers.go`, ctl dispatch in `ctl.go`. See `doc/edsl.md` for the full reference. +7. **9P namespace declared via EDSL**: The entire filesystem is a single recursive `FsNodeDecl` tree in `fs/spec.go`, built by `BuildTree()` in `fsedsl/builder.go`. The `fs/` package is flat — handlers in `handlers_*.go`, ctl dispatch in `ctl.go`. See `doc/edsl.md` for the full reference. ## Key Files | What | Where | |------|-------| @@ -107,7 +109,7 @@ just test-remote | Remote execution | `toolsrv/remote.go`, `cmd/ollie-remote/` | | Sandbox enforcement | `sandbox/` | | 9P filesystem (EDSL spec) | `fs/spec.go` | -| 9P filesystem (builder) | `fs/builder.go` | +| 9P filesystem (builder) | `fsedsl/builder.go` | | Session management | `session/session.go` | | System prompt template | Embedded in binary | | Agent configs | `agents/*.json` | diff --git a/agent/agent.go b/agent/agent.go index 506f008..8e13602 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -335,6 +335,34 @@ func (ag *Agent) SwitchProfile(name string) error { // Event is a typed output event emitted during an agent turn or in response // to a command. +// Event carries a single piece of output from the agent loop to consumers +// (frontends, loggers, the turn orchestrator in turn.go). +// +// The Role field determines the event semantics: +// +// Role Name Content Emitted by +// ──────────── ───────── ────────────────────────── ────────── +// "user" — user input text turn.go (before run) +// "assistant" — streamed LLM text chunk loop.go (streamResponse) +// "reasoning" — … chunks loop.go (streamResponse) +// "call" tool name JSON args loop.go (execOne) +// "tool" tool name result text (may stream) loop.go (execOne) +// "usage" — "in out est cost cached creation" loop.go (streamResponse) +// "state" — "thinking"|"compacting"|… loop.go / turn.go +// "limitretry" — — loop.go (rate limit hit) +// "retry" — "HH:MM:SS" countdown loop.go (retryCountdown) +// "maxsteps" — step count as string loop.go (budget exhausted) +// "error" — error message turn.go / loop.go +// "info" — informational text turn.go (compaction, cost) +// +// turn.go intercepts events before forwarding to the external handler: +// - "assistant" → accumulates reply text +// - "call" → updates agent state to "calling: " +// - "tool" → logs result +// - "state" → calls ag.SetState() +// - "limitretry"→ sets state to "limitretry" +// - "usage" → parses token counts, updates history, checks spending cap +// - "error" → logs type Event struct { Role string Name string @@ -609,8 +637,9 @@ func (ag *Agent) DismissDetached(pid int) bool { return ag.runtime.ToolServer.DismissDetached(pid) } -// AgentCfg is the configuration for constructing a new Agent. -type AgentCfg struct { +// AgentParams holds the runtime dependencies for constructing a new Agent. +// Not to be confused with AgentConfig, which is the on-disk JSON schema. +type AgentParams struct { ID string // unique agent identity (uname) SessionID string // id of session the agent belongs to Profile string // config profile name (e.g. "default" → agents/default.json) @@ -630,7 +659,9 @@ type AgentCfg struct { } // NewAgent constructs an Agent from the given configuration. -func NewAgent(cfg AgentCfg) *Agent { +// Post-construction wiring (SetSessionEnv, wireToolsChanged) is handled +// internally — no additional calls are required after construction. +func NewAgent(cfg AgentParams) *Agent { ag := &Agent{ history: cfg.History, runtime: cfg.Runtime, @@ -653,5 +684,6 @@ func NewAgent(cfg AgentCfg) *Agent { } ag.signalCh = make(chan struct{}) ag.wireToolsChanged() + ag.SetSessionEnv(cfg.SessionID) return ag } diff --git a/agent/loop.go b/agent/loop.go index b55a53b..376399c 100644 --- a/agent/loop.go +++ b/agent/loop.go @@ -660,6 +660,16 @@ func contextBudget(ctx context.Context, b backend.Backend) int { } // defaultToolResultMaxBytes caps tool result content sent back to the model. +// This is the single semantic limit for what enters the LLM context window. +// +// Truncation stack (from outer to inner): +// - Server 10MB (toolsrv/shell.go limitedWriter): OOM safety. Prevents runaway +// process output from exhausting memory. Well above model limit. +// - Agent 128KB (this constant): The model-facing cap. Applied in two places: +// streaming path (caps real-time chunks sent to UI) and post-hoc (caps stored +// result). Both use this constant — it's one policy, two enforcement points. +// - Detach 64KB ring (toolsrv/detach.go): Separate system for background +// processes. Unrelated to the model context. const defaultToolResultMaxBytes = 131072 // MemoryTier classifies how long a tool result stays in the hot message list. diff --git a/agent/runtime.go b/agent/runtime.go index ad04d27..131ecef 100644 --- a/agent/runtime.go +++ b/agent/runtime.go @@ -244,29 +244,9 @@ func toolInfosToBackend(infos []toolsrv.ToolInfo) []backend.Tool { // extractToolResult parses the tool server's JSON-RPC response into text // and optional content blocks. // -// Protocol: tool scripts write JSON to stdout matching this schema: -// -// { -// "isError": false, -// "content": [ -// {"type": "text", "text": "..."}, -// {"type": "image", "media_type": "image/png", "data": ""} -// ] -// } -// -// If the response is not valid JSON (e.g. a plain string from a simple script), -// it is returned as-is with isError=false. This fallback allows tools to emit -// raw text without wrapping it in the content-block envelope. +// See toolsrv.ToolResult for the wire format. func extractToolResult(raw json.RawMessage) (text string, contentBlocks []backend.ContentBlock, isError bool) { - var result struct { - IsError bool `json:"isError"` - Content []struct { - Type string `json:"type"` - Text string `json:"text"` - MediaType string `json:"media_type"` - Data string `json:"data"` - } `json:"content"` - } + var result toolsrv.ToolResult if err := json.Unmarshal(raw, &result); err != nil { return string(raw), nil, false } diff --git a/cmd/ollie-remote/main.go b/cmd/ollie-remote/main.go index d547c42..89ebacf 100644 --- a/cmd/ollie-remote/main.go +++ b/cmd/ollie-remote/main.go @@ -14,20 +14,14 @@ package main import ( "context" - "encoding/json" "flag" "fmt" - "io" - "net" "os" "os/signal" "path/filepath" - "runtime" - "sync" "syscall" "ollie/env" - "ollie/paths" "ollie/toolsrv" ) @@ -44,19 +38,14 @@ func main() { } flag.CommandLine.Parse(os.Args[2:]) - // Load env file and set defaults (same as olliesrv). env.EnsureDefaults() - // Prepend our bin dir to PATH so landrun is found (deployed alongside us) + // Prepend our bin dir to PATH so landrun is found (deployed alongside us). home, _ := os.UserHomeDir() binDir := filepath.Join(home, ".config", "ollie", "bin") os.Setenv("PATH", binDir+":"+os.Getenv("PATH")) - // Build the tool server with tool registry. toolReg, _ := toolsrv.NewRegistry() - - // OLLIE_SESSION_ID may be set in the inherited environment (local spawn) - // or arrive later via set_env RPC (SSH bootstrap). Accept both paths. sessionID := os.Getenv("OLLIE_SESSION_ID") var opts []toolsrv.Option @@ -71,8 +60,7 @@ func main() { o(server) } - // If tool registry is available but session ID was not set yet, - // wire up OnEnvSet to attach it when OLLIE_SESSION_ID arrives. + // Wire up environment propagation for late-arriving session IDs. server.OnEnvSet = func(key, value string) { switch key { case "OLLIE_SESSION_ID": @@ -95,257 +83,11 @@ func main() { }() if *listenPath != "" { - serveSocket(ctx, server, *listenPath) + if err := server.ServeSocket(ctx, *listenPath); err != nil { + fmt.Fprintf(os.Stderr, "%v\n", err) + os.Exit(1) + } } else { - serveRPC(ctx, server, os.Stdin, os.Stdout) + server.ServeRPC(ctx, os.Stdin, os.Stdout) } } - -func serveSocket(ctx context.Context, srv *toolsrv.Server, sockPath string) { - os.Remove(sockPath) - os.MkdirAll(filepath.Dir(sockPath), 0700) - - ln, err := net.Listen("unix", sockPath) - if err != nil { - fmt.Fprintf(os.Stderr, "listen %s: %v\n", sockPath, err) - os.Exit(1) - } - defer ln.Close() - defer os.Remove(sockPath) - - fmt.Fprintln(os.Stderr, "ListenReady") - - var mu sync.Mutex - var activeConns []net.Conn - var wg sync.WaitGroup - - go func() { - <-ctx.Done() - ln.Close() - mu.Lock() - for _, c := range activeConns { - c.Close() - } - mu.Unlock() - }() - - for { - conn, err := ln.Accept() - if err != nil { - if ctx.Err() != nil { - break - } - continue - } - mu.Lock() - activeConns = append(activeConns, conn) - mu.Unlock() - - wg.Add(1) - go func() { - defer wg.Done() - defer conn.Close() - serveRPC(ctx, srv, conn, conn) - mu.Lock() - for i, c := range activeConns { - if c == conn { - activeConns = append(activeConns[:i], activeConns[i+1:]...) - break - } - } - mu.Unlock() - }() - } - wg.Wait() -} - -type rpcResponse struct { - JSONRPC string `json:"jsonrpc"` - ID json.RawMessage `json:"id"` - Result any `json:"result,omitempty"` - Error *toolsrv.RPCError `json:"error,omitempty"` - Stream bool `json:"stream,omitempty"` -} - -type lockedEncoder struct { - mu sync.Mutex - enc *json.Encoder -} - -func (e *lockedEncoder) Encode(v any) error { - e.mu.Lock() - defer e.mu.Unlock() - return e.enc.Encode(v) -} - -func serveRPC(ctx context.Context, srv *toolsrv.Server, in io.Reader, out io.Writer) { - dec := json.NewDecoder(in) - enc := &lockedEncoder{enc: json.NewEncoder(out)} - - for { - var req toolsrv.RPCRequest - if err := dec.Decode(&req); err != nil { - if err == io.EOF || ctx.Err() != nil { - return - } - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &toolsrv.RPCError{Code: -32700, Message: "parse error: " + err.Error()}, - }) - continue - } - - go handleRPC(ctx, srv, req, enc) - } -} - -func handleRPC(ctx context.Context, srv *toolsrv.Server, req toolsrv.RPCRequest, enc *lockedEncoder) { - switch req.Method { - case "list_tools": - tools, err := srv.ListTools() - if err != nil { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &toolsrv.RPCError{Code: -32000, Message: err.Error()}, - }) - } else { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: tools, - }) - } - - case "host_info": - info := map[string]any{ - "platform": runtime.GOOS, - "arch": runtime.GOARCH, - "is_git_repo": paths.IsGitRepo(*cwd), - } - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: info, - }) - - case "ping": - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: "pong", - }) - - case "set_env": - var params struct { - Key string `json:"key"` - Value string `json:"value"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - srv.SetEnv(params.Key, params.Value) - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - - case "set_cwd": - var params struct { - Dir string `json:"dir"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - srv.SetCWD(params.Dir) - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - - case "detach": - ok := srv.Detach() - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) - - case "list_detached": - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: srv.ListDetachedRaw()}) - - case "signal_detached": - var params struct { - PID int `json:"pid"` - Signal int `json:"signal"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - if err := srv.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) - } - - case "get_detached_output": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - out, err := srv.GetDetachedOutput(params.PID) - if err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32000, Message: err.Error()}}) - } else { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: out}) - } - - case "dismiss_detached": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - ok := srv.DismissDetached(params.PID) - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) - - case "tool_load": - var params struct { - Name string `json:"name"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32602, Message: err.Error()}}) - return - } - if err := srv.LoadTool(params.Name); err != nil { - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &toolsrv.RPCError{Code: -32000, Message: err.Error()}}) - } else { - res, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name}) - enc.Encode(rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(res)}) - } - - default: - streamCtx := toolsrv.WithOutputStream(ctx, func(data string) { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: toolsrv.OutputNotification{Data: data}, - Stream: true, - }) - }) - result, err := srv.CallTool(streamCtx, req.Method, req.Params) - if err != nil { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &toolsrv.RPCError{Code: -32000, Message: err.Error()}, - }) - } else { - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Result: json.RawMessage(result), - }) - } - } -} - diff --git a/doc/boot-sequence.md b/doc/boot-sequence.md new file mode 100644 index 0000000..b8329db --- /dev/null +++ b/doc/boot-sequence.md @@ -0,0 +1,54 @@ +# Boot Sequence + +How `olliesrv` starts up, from `main()` to serving 9P requests. + +## Call chain + +``` +main() + │ + ├── env.EnsureDefaults() — populate $XDG_CONFIG_HOME/ollie/ if missing + ├── olog.NewSink() — stderr/stdout log sink + ├── fs.NewModelCache() — async LLM model list (prefetched later) + │ + ├── fs.NewRoot(cfg) ← composition root + │ │ + │ ├── session.Init(...) — set package-level config (ctx, dirs, logger) + │ ├── build RootState — allocate SessionNodes map + │ ├── fsedsl.BuildTree(spec) — walk fs/spec.go, produce *Tree with lazy resolution + │ └── session.RestoreAllSessions() + │ └── for each .json in sessions/active/ + │ ├── unmarshal PersistedSession + │ ├── rebuild Session + Agents (paused) + │ └── register in session registry + │ + ├── bypass.NewBroker(...) — start bypass socket (sandbox escape approval) + ├── New(serverCfg) — create 9P server wrapping the tree + ├── net.Listen("unix", sockPath) — bind the 9P socket + └── accept loop — each conn gets srv.Start(conn) +``` + +## Why `fs.NewRoot` is the composition root + +The 9P filesystem IS the runtime. Every session, agent, and tool server is +reachable only through the namespace. So the filesystem package owns +initialization: it creates the tree, initializes the session registry, and +restores persisted state — all before a single 9P connection is accepted. + +`main.go` is deliberately thin: parse flags, set up logging, call `fs.NewRoot`, +start the listener. All domain logic lives in the packages. + +## Key ordering constraints + +1. `session.Init()` must be called before `RestoreAllSessions()` (sets pkg config). +2. `RootState` must be allocated before `BuildTree()` (closures capture the pointer). +3. `BuildTree()` must complete before `RestoreAllSessions()` (restored sessions + create nodes in the tree). +4. The bypass broker starts AFTER the tree is built but BEFORE the accept loop + (tool execution needs the broker socket). + +## Session restoration + +Restored sessions come back **paused**. They have agents with history but no +running tool server or LLM connection. The first `Resume` (via ctl write) +re-dials the tool server and reconnects the backend. diff --git a/fsedsl/tree.go b/fsedsl/tree.go index 7dbf3ac..d04898e 100644 --- a/fsedsl/tree.go +++ b/fsedsl/tree.go @@ -6,10 +6,21 @@ import ( "path/filepath" ) -// Tree is a synthetic filesystem node built by BuildTree. -// Handlers are injected via functional options; there are no backing directories. +// Tree is a synthetic filesystem node. +// +// Two usage modes: +// +// 1. EDSL (BuildTree): Returns a single root node. Children are resolved lazily +// via closures over the FsNodeDecl spec. The children map is NOT populated; +// all path resolution happens through the handler closures (openFn, statFn, etc). +// +// 2. Manual (Mount): Children are attached explicitly via Mount(). The split() +// method routes path lookups through the children map before falling through +// to handler closures. +// +// In production, only mode 1 is used. Mode 2 exists for composability and testing. type Tree struct { - Tree []string // backing dirs (legacy; nil in EDSL trees) + Tree []string // backing dirs (nil in EDSL trees) perm os.FileMode listFn func([]string) ([]os.DirEntry, error) createFn func([]string, string, os.FileMode) error diff --git a/session/persist.go b/session/persist.go index 543e19f..4b489d6 100644 --- a/session/persist.go +++ b/session/persist.go @@ -202,7 +202,7 @@ func restoreMultiAgentSession(ps *PersistedSession) (*RestoredSession, error) { // For paused sessions, create a minimal agent without tool server if ps.Paused { history := agent.RestoreHistoryFromMessages(pa.Messages) - ag := agent.NewAgent(agent.AgentCfg{ + ag := agent.NewAgent(agent.AgentParams{ History: history, Runtime: &agent.Runtime{}, Profile: pa.Profile, diff --git a/session/registry.go b/session/registry.go index 22a43b5..63a3981 100644 --- a/session/registry.go +++ b/session/registry.go @@ -419,7 +419,7 @@ func buildAgent(sess *Session, p AgentParams) (*agent.Agent, error) { rt.Backend.SetModel(cfg.Model) } - ag := agent.NewAgent(agent.AgentCfg{ + ag := agent.NewAgent(agent.AgentParams{ History: p.History, Runtime: rt, Profile: p.Profile, diff --git a/session/session.go b/session/session.go index 04da804..e95d55f 100644 --- a/session/session.go +++ b/session/session.go @@ -114,7 +114,6 @@ func (s *Session) Agents() []*agent.Agent { return s.agents } // AddAgent appends an agent to the session. func (s *Session) AddAgent(ag *agent.Agent) { s.agents = append(s.agents, ag) - ag.SetSessionEnv(s.ID) } // RemoveAgent removes an agent by ID and closes it. @@ -229,45 +228,39 @@ func (s *Session) Resume() error { s.Ctx = ctx s.Cancel = cancel - if s.Keeper != nil { + // Ensure we always have a Keeper. If the session was restored from disk + // without infra, spawn a new tool server now. + if s.Keeper == nil { + cwd := "" + if len(s.agents) > 0 { + cwd = s.agents[0].Cwd() + } + if cwd == "" { + cwd, _ = os.Getwd() + } + infra, err := SetupToolServer(ToolServerConfig{ + Ctx: ctx, + CWD: cwd, + RemoteTarget: s.Remote, + Yolo: pkgYolo, + }) + if err != nil { + return fmt.Errorf("resume: %w", err) + } + s.Proc = infra.Proc + s.Keeper = infra.Keeper + s.toolsConn = infra.ToolsConn + for _, ag := range s.agents { + ag.SetToolServer(infra.NewToolServer, infra.NewToolServer()) + } + } else { + // Keeper exists — just re-dial. s.Keeper.SetContext(ctx) conn, err := s.Keeper.Dial() if err != nil { return fmt.Errorf("resume failed: %w", err) } s.toolsConn = conn - } else { - // Session was restored paused without infra — set up per-agent. - for i, ag := range s.agents { - cwd := ag.Cwd() - if cwd == "" { - cwd, _ = os.Getwd() - } - var reuseFrom *InfraConfig - if s.Proc != nil { - reuseFrom = &InfraConfig{ - Proc: s.Proc, - Keeper: s.Keeper, - ToolsConn: s.toolsConn, - } - } - infra, err := SetupToolServer(ToolServerConfig{ - Ctx: ctx, - CWD: cwd, - RemoteTarget: s.Remote, - Yolo: pkgYolo, - ReuseFrom: reuseFrom, - }) - if err != nil { - return fmt.Errorf("resume setup for agent %s failed: %w", ag.ID(), err) - } - if i == 0 { - s.Proc = infra.Proc - s.Keeper = infra.Keeper - s.toolsConn = infra.ToolsConn - } - ag.SetToolServer(infra.NewToolServer, infra.ToolsConn) - } } s.paused = false diff --git a/toolsrv/shell.go b/toolsrv/exec.go similarity index 100% rename from toolsrv/shell.go rename to toolsrv/exec.go diff --git a/toolsrv/rpc_integration_test.go b/toolsrv/rpc_integration_test.go index bb5ef3b..d05e45f 100644 --- a/toolsrv/rpc_integration_test.go +++ b/toolsrv/rpc_integration_test.go @@ -3,7 +3,6 @@ package toolsrv_test import ( "context" "encoding/json" - "io" "net" "os" "path/filepath" @@ -11,14 +10,13 @@ import ( "sync" "syscall" "testing" - "time" "ollie/toolsrv" ) // --- Test infrastructure --- -// rpcRequest mirrors the JSON-RPC 2.0 request format. +// rpcRequest mirrors the JSON-RPC 2.0 request format (test-local). type rpcRequest struct { JSONRPC string `json:"jsonrpc"` ID int64 `json:"id"` @@ -26,7 +24,7 @@ type rpcRequest struct { Params json.RawMessage `json:"params,omitempty"` } -// rpcResponse mirrors the JSON-RPC 2.0 response format. +// rpcResponse mirrors the JSON-RPC 2.0 response format (test-local). type rpcResponse struct { JSONRPC string `json:"jsonrpc"` ID int64 `json:"id"` @@ -40,23 +38,18 @@ type rpcError struct { Message string `json:"message"` } -type outputNotif struct { - Data string `json:"data"` -} - -// testServer wraps a toolsrv.Server and provides RPC dispatch for testing. -// This mimics the cmd/ollie-remote serveRPC loop. +// testServer wraps a toolsrv.Server with a pipe-based RPC connection. type testServer struct { - srv *toolsrv.Server - registry *toolsrv.Registry - cwd string + srv *toolsrv.Server + conn net.Conn // client side + enc *json.Encoder + dec *json.Decoder + wg sync.WaitGroup } func newTestServer(t *testing.T) *testServer { t.Helper() cwd := t.TempDir() - - // Create a .git directory to make isGitRepo return true os.MkdirAll(filepath.Join(cwd, ".git"), 0755) reg, err := toolsrv.NewRegistry() @@ -68,186 +61,44 @@ func newTestServer(t *testing.T) *testServer { srv := toolsrv.New(cwd) srv.SetToolRegistry(reg, sessionID) - return &testServer{ - srv: srv, - registry: reg, - cwd: cwd, - } -} - -// handleRPC dispatches a single RPC request and returns the response. -// This mirrors cmd/ollie-remote handleRPC but returns instead of encoding. -func (ts *testServer) handleRPC(ctx context.Context, req rpcRequest) rpcResponse { - switch req.Method { - case "ping": - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`"pong"`)} - - case "host_info": - info := map[string]any{ - "platform": runtime.GOOS, - "arch": runtime.GOARCH, - "is_git_repo": isGitRepo(ts.cwd), - } - result, _ := json.Marshal(info) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "list_tools": - tools, err := ts.srv.ListTools() - if err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}} - } - result, _ := json.Marshal(tools) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "set_env": - var params struct { - Key string `json:"key"` - Value string `json:"value"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - ts.srv.SetEnv(params.Key, params.Value) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)} - - case "set_cwd": - var params struct { - Dir string `json:"dir"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - ts.srv.SetCWD(params.Dir) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)} - - case "detach": - ok := ts.srv.Detach() - result, _ := json.Marshal(ok) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "list_detached": - list := ts.srv.ListDetachedRaw() - result, _ := json.Marshal(list) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "signal_detached": - var params struct { - PID int `json:"pid"` - Signal int `json:"signal"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - if err := ts.srv.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}} - } - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(`true`)} - - case "get_detached_output": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - out, err := ts.srv.GetDetachedOutput(params.PID) - if err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}} - } - result, _ := json.Marshal(out) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "dismiss_detached": - var params struct { - PID int `json:"pid"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - ok := ts.srv.DismissDetached(params.PID) - result, _ := json.Marshal(ok) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - case "tool_load": - var params struct { - Name string `json:"name"` - } - if err := json.Unmarshal(req.Params, ¶ms); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32602, Message: err.Error()}} - } - if err := ts.srv.LoadTool(params.Name); err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}} - } - result, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name}) - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - - default: - // Unknown method = tool call - result, err := ts.srv.CallTool(ctx, req.Method, req.Params) - if err != nil { - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Error: &rpcError{Code: -32000, Message: err.Error()}} - } - return rpcResponse{JSONRPC: "2.0", ID: req.ID, Result: result} - } -} - -func isGitRepo(dir string) bool { - _, err := os.Stat(filepath.Join(dir, ".git")) - return err == nil -} - -// --- Pipe-based RPC client/server --- - -// pipeConn creates an in-memory connection for testing RPC client/server interaction. -// Returns client and server ends of a bidirectional pipe. -func pipeConn() (client, server net.Conn) { clientConn, serverConn := net.Pipe() - return clientConn, serverConn + ts := &testServer{ + srv: srv, + conn: clientConn, + enc: json.NewEncoder(clientConn), + dec: json.NewDecoder(clientConn), + } + + ctx, cancel := context.WithCancel(context.Background()) + ts.wg.Add(1) + go func() { + defer ts.wg.Done() + srv.ServeRPC(ctx, serverConn, serverConn) + }() + + t.Cleanup(func() { + cancel() + clientConn.Close() + serverConn.Close() + ts.wg.Wait() + }) + + return ts } -// runTestRPCServer runs the test server's RPC handler over the given connection. -// Blocks until the connection is closed or context is cancelled. -func runTestRPCServer(ctx context.Context, ts *testServer, conn net.Conn) { - dec := json.NewDecoder(conn) - enc := json.NewEncoder(conn) - var encMu sync.Mutex - +// call sends an RPC request and returns the response (skipping stream frames). +func (ts *testServer) call(method string, params any) rpcResponse { + var raw json.RawMessage + if params != nil { + raw, _ = json.Marshal(params) + } + ts.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: 1, Method: method, Params: raw}) for { - select { - case <-ctx.Done(): - return - default: + var resp rpcResponse + ts.dec.Decode(&resp) + if !resp.Stream { + return resp } - - // Set read deadline to allow periodic context checks - conn.SetReadDeadline(time.Now().Add(100 * time.Millisecond)) - - var req rpcRequest - if err := dec.Decode(&req); err != nil { - if netErr, ok := err.(net.Error); ok && netErr.Timeout() { - continue // Just a timeout, check context and retry - } - if err == io.EOF || ctx.Err() != nil { - return - } - encMu.Lock() - enc.Encode(rpcResponse{ - JSONRPC: "2.0", - ID: req.ID, - Error: &rpcError{Code: -32700, Message: "parse error: " + err.Error()}, - }) - encMu.Unlock() - continue - } - - // Handle request (synchronously for simplicity in tests) - resp := ts.handleRPC(ctx, req) - encMu.Lock() - if err := enc.Encode(resp); err != nil { - encMu.Unlock() - return - } - encMu.Unlock() } } @@ -255,19 +106,13 @@ func runTestRPCServer(ctx context.Context, ts *testServer, conn net.Conn) { func TestRPC_Ping(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "ping"} - resp := ts.handleRPC(ctx, req) + resp := ts.call("ping", nil) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } var result string - if err := json.Unmarshal(resp.Result, &result); err != nil { - t.Fatalf("unmarshal result: %v", err) - } + json.Unmarshal(resp.Result, &result) if result != "pong" { t.Errorf("expected 'pong', got %q", result) } @@ -275,20 +120,14 @@ func TestRPC_Ping(t *testing.T) { func TestRPC_HostInfo(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "host_info"} - resp := ts.handleRPC(ctx, req) + resp := ts.call("host_info", nil) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } var info toolsrv.HostInfo - if err := json.Unmarshal(resp.Result, &info); err != nil { - t.Fatalf("unmarshal result: %v", err) - } + json.Unmarshal(resp.Result, &info) if info.Platform != runtime.GOOS { t.Errorf("platform: expected %q, got %q", runtime.GOOS, info.Platform) } @@ -302,21 +141,14 @@ func TestRPC_HostInfo(t *testing.T) { func TestRPC_ListTools_Empty(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "list_tools"} - resp := ts.handleRPC(ctx, req) + resp := ts.call("list_tools", nil) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } var tools []toolsrv.ToolInfo - if err := json.Unmarshal(resp.Result, &tools); err != nil { - t.Fatalf("unmarshal result: %v", err) - } - // No tools loaded yet + json.Unmarshal(resp.Result, &tools) if len(tools) != 0 { t.Errorf("expected 0 tools, got %d", len(tools)) } @@ -324,30 +156,17 @@ func TestRPC_ListTools_Empty(t *testing.T) { func TestRPC_SetEnv(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - // Hook to verify env is set var envKey, envValue string ts.srv.OnEnvSet = func(k, v string) { envKey, envValue = k, v } - ctx := context.Background() - params, _ := json.Marshal(map[string]string{"key": "TEST_VAR", "value": "test_value"}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: params} - resp := ts.handleRPC(ctx, req) + resp := ts.call("set_env", map[string]string{"key": "TEST_VAR", "value": "test_value"}) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } - - var result bool - if err := json.Unmarshal(resp.Result, &result); err != nil { - t.Fatalf("unmarshal result: %v", err) - } - if !result { - t.Error("expected result=true") - } if envKey != "TEST_VAR" || envValue != "test_value" { t.Errorf("env not set: got key=%q value=%q", envKey, envValue) } @@ -355,23 +174,13 @@ func TestRPC_SetEnv(t *testing.T) { func TestRPC_SetCWD(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - newDir := "/tmp/test-cwd" - - ctx := context.Background() - params, _ := json.Marshal(map[string]string{"dir": newDir}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_cwd", Params: params} - resp := ts.handleRPC(ctx, req) + resp := ts.call("set_cwd", map[string]string{"dir": "/tmp/test-cwd"}) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } - var result bool - if err := json.Unmarshal(resp.Result, &result); err != nil { - t.Fatalf("unmarshal result: %v", err) - } + json.Unmarshal(resp.Result, &result) if !result { t.Error("expected result=true") } @@ -379,21 +188,13 @@ func TestRPC_SetCWD(t *testing.T) { func TestRPC_Detach_NoRunningProcess(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "detach"} - resp := ts.handleRPC(ctx, req) + resp := ts.call("detach", nil) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } - var result bool - if err := json.Unmarshal(resp.Result, &result); err != nil { - t.Fatalf("unmarshal result: %v", err) - } - // No running process to detach + json.Unmarshal(resp.Result, &result) if result { t.Error("expected false when no process running") } @@ -401,20 +202,13 @@ func TestRPC_Detach_NoRunningProcess(t *testing.T) { func TestRPC_ListDetached_Empty(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "list_detached"} - resp := ts.handleRPC(ctx, req) + resp := ts.call("list_detached", nil) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } - var list []any - if err := json.Unmarshal(resp.Result, &list); err != nil { - t.Fatalf("unmarshal result: %v", err) - } + json.Unmarshal(resp.Result, &list) if len(list) != 0 { t.Errorf("expected empty list, got %d items", len(list)) } @@ -422,14 +216,8 @@ func TestRPC_ListDetached_Empty(t *testing.T) { func TestRPC_SignalDetached_NotFound(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() + resp := ts.call("signal_detached", map[string]any{"pid": 99999, "signal": int(syscall.SIGTERM)}) - ctx := context.Background() - params, _ := json.Marshal(map[string]any{"pid": 99999, "signal": int(syscall.SIGTERM)}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "signal_detached", Params: params} - resp := ts.handleRPC(ctx, req) - - // Should error because no such detached process if resp.Error == nil { t.Fatal("expected error for nonexistent PID") } @@ -440,14 +228,8 @@ func TestRPC_SignalDetached_NotFound(t *testing.T) { func TestRPC_GetDetachedOutput_NotFound(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() + resp := ts.call("get_detached_output", map[string]any{"pid": 99999}) - ctx := context.Background() - params, _ := json.Marshal(map[string]any{"pid": 99999}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "get_detached_output", Params: params} - resp := ts.handleRPC(ctx, req) - - // Should error because no such detached process if resp.Error == nil { t.Fatal("expected error for nonexistent PID") } @@ -455,22 +237,13 @@ func TestRPC_GetDetachedOutput_NotFound(t *testing.T) { func TestRPC_DismissDetached_NotFound(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - params, _ := json.Marshal(map[string]any{"pid": 99999}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "dismiss_detached", Params: params} - resp := ts.handleRPC(ctx, req) + resp := ts.call("dismiss_detached", map[string]any{"pid": 99999}) if resp.Error != nil { t.Fatalf("unexpected error: %v", resp.Error.Message) } - var result bool - if err := json.Unmarshal(resp.Result, &result); err != nil { - t.Fatalf("unmarshal result: %v", err) - } - // Should return false for nonexistent PID + json.Unmarshal(resp.Result, &result) if result { t.Error("expected false for nonexistent PID") } @@ -478,14 +251,8 @@ func TestRPC_DismissDetached_NotFound(t *testing.T) { func TestRPC_CallTool_NotLoaded(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() + resp := ts.call("nonexistent_tool", map[string]any{"arg": "value"}) - ctx := context.Background() - params, _ := json.Marshal(map[string]any{"arg": "value"}) - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "nonexistent_tool", Params: params} - resp := ts.handleRPC(ctx, req) - - // Should error because tool not loaded if resp.Error == nil { t.Fatal("expected error for unloaded tool") } @@ -496,52 +263,52 @@ func TestRPC_CallTool_NotLoaded(t *testing.T) { func TestRPC_InvalidParams(t *testing.T) { ts := newTestServer(t) - defer ts.srv.Close() - - ctx := context.Background() - // Send malformed JSON as params - req := rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: json.RawMessage(`{invalid`)} - resp := ts.handleRPC(ctx, req) - - // Should error with invalid params - if resp.Error == nil { - t.Fatal("expected error for invalid params") - } - if resp.Error.Code != -32602 { - t.Errorf("expected error code -32602 (invalid params), got %d", resp.Error.Code) + // Send an array where an object is expected — json.Unmarshal will fail. + ts.enc.Encode(rpcRequest{JSONRPC: "2.0", ID: 1, Method: "set_env", Params: json.RawMessage(`[1,2,3]`)}) + for { + var resp rpcResponse + if err := ts.dec.Decode(&resp); err != nil { + t.Fatalf("decode: %v", err) + } + if resp.Stream { + continue + } + if resp.Error == nil { + t.Fatal("expected error for invalid params") + } + if resp.Error.Code != -32602 { + t.Errorf("expected error code -32602 (invalid params), got %d", resp.Error.Code) + } + break } } // --- Full pipe-based client/server tests --- func TestRPC_PipeConnection(t *testing.T) { - ts := newTestServer(t) - defer ts.srv.Close() + cwd := t.TempDir() + os.MkdirAll(filepath.Join(cwd, ".git"), 0755) + + reg, _ := toolsrv.NewRegistry() + srv := toolsrv.New(cwd) + srv.SetToolRegistry(reg, "pipe-test") ctx, cancel := context.WithCancel(context.Background()) defer cancel() - clientConn, serverConn := pipeConn() + clientConn, serverConn := net.Pipe() defer clientConn.Close() defer serverConn.Close() - // Start server in background - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - runTestRPCServer(ctx, ts, serverConn) - }() + go srv.ServeRPC(ctx, serverConn, serverConn) - // Create client connection using toolsrv.Conn + // Use toolsrv.Conn as the client conn := toolsrv.NewConn(clientConn, nil) - // Test ping if err := conn.Ping(); err != nil { t.Errorf("Ping failed: %v", err) } - // Test list tools (should be empty) tools, err := conn.ListTools() if err != nil { t.Errorf("ListTools failed: %v", err) @@ -550,228 +317,6 @@ func TestRPC_PipeConnection(t *testing.T) { t.Errorf("expected 0 tools, got %d", len(tools)) } - // Test set env conn.SetEnv("TEST_KEY", "TEST_VALUE") - - // Test set cwd conn.SetCWD("/tmp/test") - - // Clean up - cancel() - conn.Close() - wg.Wait() -} - -func TestRPC_ConcurrentRequests(t *testing.T) { - ts := newTestServer(t) - defer ts.srv.Close() - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - clientConn, serverConn := pipeConn() - defer clientConn.Close() - defer serverConn.Close() - - // Start server - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - runTestRPCServer(ctx, ts, serverConn) - }() - - conn := toolsrv.NewConn(clientConn, nil) - - // Fire off concurrent ping requests - const numRequests = 10 - var reqWg sync.WaitGroup - errors := make(chan error, numRequests) - - for i := 0; i < numRequests; i++ { - reqWg.Add(1) - go func() { - defer reqWg.Done() - if err := conn.Ping(); err != nil { - errors <- err - } - }() - } - - reqWg.Wait() - close(errors) - - for err := range errors { - t.Errorf("concurrent ping failed: %v", err) - } - - cancel() - conn.Close() - wg.Wait() -} - -func TestRPC_ContextCancellation(t *testing.T) { - ts := newTestServer(t) - defer ts.srv.Close() - - ctx, cancel := context.WithCancel(context.Background()) - - clientConn, serverConn := pipeConn() - defer clientConn.Close() - defer serverConn.Close() - - // Start server - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - runTestRPCServer(ctx, ts, serverConn) - }() - - conn := toolsrv.NewConn(clientConn, nil) - - // Cancel context - cancel() - - // Allow some time for cancellation to propagate - time.Sleep(50 * time.Millisecond) - - // Clean up - conn.Close() - wg.Wait() -} - -func TestRPC_ConnectionClose(t *testing.T) { - ts := newTestServer(t) - defer ts.srv.Close() - - ctx, cancel := context.WithCancel(context.Background()) - defer cancel() - - clientConn, serverConn := pipeConn() - - // Start server - var wg sync.WaitGroup - wg.Add(1) - go func() { - defer wg.Done() - runTestRPCServer(ctx, ts, serverConn) - }() - - conn := toolsrv.NewConn(clientConn, nil) - - // Verify ping works - if err := conn.Ping(); err != nil { - t.Fatalf("initial Ping failed: %v", err) - } - - // Close the client connection - this should cause the server to detect closure - conn.Close() - clientConn.Close() - serverConn.Close() - - // Cancel context to ensure server exits - cancel() - - // Wait with timeout - done := make(chan struct{}) - go func() { - wg.Wait() - close(done) - }() - - select { - case <-done: - // Good, server exited - case <-time.After(2 * time.Second): - t.Error("server did not exit after connection close") - } -} - -// TestRPC_JSONRPCProtocol tests raw JSON-RPC protocol compliance -func TestRPC_JSONRPCProtocol(t *testing.T) { - ts := newTestServer(t) - defer ts.srv.Close() - - tests := []struct { - name string - request rpcRequest - wantErr bool - errCode int - checkResult func(t *testing.T, result json.RawMessage) - }{ - { - name: "ping returns pong", - request: rpcRequest{JSONRPC: "2.0", ID: 1, Method: "ping"}, - checkResult: func(t *testing.T, result json.RawMessage) { - var s string - json.Unmarshal(result, &s) - if s != "pong" { - t.Errorf("expected 'pong', got %q", s) - } - }, - }, - { - name: "set_env with valid params", - request: rpcRequest{ - JSONRPC: "2.0", - ID: 2, - Method: "set_env", - Params: json.RawMessage(`{"key":"K","value":"V"}`), - }, - checkResult: func(t *testing.T, result json.RawMessage) { - var b bool - json.Unmarshal(result, &b) - if !b { - t.Error("expected true") - } - }, - }, - { - name: "set_env with invalid params", - request: rpcRequest{ - JSONRPC: "2.0", - ID: 3, - Method: "set_env", - Params: json.RawMessage(`not valid json`), - }, - wantErr: true, - errCode: -32602, - }, - { - name: "unknown tool returns error", - request: rpcRequest{ - JSONRPC: "2.0", - ID: 4, - Method: "fake_tool", - Params: json.RawMessage(`{}`), - }, - wantErr: true, - errCode: -32000, - }, - } - - ctx := context.Background() - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - resp := ts.handleRPC(ctx, tt.request) - - if tt.wantErr { - if resp.Error == nil { - t.Fatal("expected error, got none") - } - if resp.Error.Code != tt.errCode { - t.Errorf("error code: expected %d, got %d", tt.errCode, resp.Error.Code) - } - return - } - - if resp.Error != nil { - t.Fatalf("unexpected error: %v", resp.Error.Message) - } - if tt.checkResult != nil { - tt.checkResult(t, resp.Result) - } - }) - } } diff --git a/toolsrv/rpcserver.go b/toolsrv/rpcserver.go new file mode 100644 index 0000000..cb599a5 --- /dev/null +++ b/toolsrv/rpcserver.go @@ -0,0 +1,263 @@ +package toolsrv + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "os" + "path/filepath" + "runtime" + "sync" + "syscall" + + "ollie/paths" +) + +// RPCResponse is a JSON-RPC 2.0 response. +type RPCResponse struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Result any `json:"result,omitempty"` + Error *RPCError `json:"error,omitempty"` + Stream bool `json:"stream,omitempty"` +} + +// lockedEncoder serializes JSON writes to a single writer. +type lockedEncoder struct { + mu sync.Mutex + enc *json.Encoder +} + +func (e *lockedEncoder) Encode(v any) error { + e.mu.Lock() + defer e.mu.Unlock() + return e.enc.Encode(v) +} + +// ServeSocket accepts JSON-RPC connections on a Unix socket. Blocks until +// ctx is cancelled. +func (s *Server) ServeSocket(ctx context.Context, sockPath string) error { + os.Remove(sockPath) + os.MkdirAll(filepath.Dir(sockPath), 0700) + + ln, err := net.Listen("unix", sockPath) + if err != nil { + return fmt.Errorf("listen %s: %w", sockPath, err) + } + defer ln.Close() + defer os.Remove(sockPath) + + fmt.Fprintln(os.Stderr, "ListenReady") + + var mu sync.Mutex + var activeConns []net.Conn + var wg sync.WaitGroup + + go func() { + <-ctx.Done() + ln.Close() + mu.Lock() + for _, c := range activeConns { + c.Close() + } + mu.Unlock() + }() + + for { + conn, err := ln.Accept() + if err != nil { + if ctx.Err() != nil { + break + } + continue + } + mu.Lock() + activeConns = append(activeConns, conn) + mu.Unlock() + + wg.Add(1) + go func() { + defer wg.Done() + defer conn.Close() + s.ServeRPC(ctx, conn, conn) + mu.Lock() + for i, c := range activeConns { + if c == conn { + activeConns = append(activeConns[:i], activeConns[i+1:]...) + break + } + } + mu.Unlock() + }() + } + wg.Wait() + return nil +} + +// ServeRPC reads JSON-RPC 2.0 requests from in and writes responses to out. +// Blocks until in is closed or ctx is cancelled. +func (s *Server) ServeRPC(ctx context.Context, in io.Reader, out io.Writer) { + dec := json.NewDecoder(in) + enc := &lockedEncoder{enc: json.NewEncoder(out)} + + for { + var req RPCRequest + if err := dec.Decode(&req); err != nil { + if err == io.EOF || ctx.Err() != nil { + return + } + enc.Encode(RPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Error: &RPCError{Code: -32700, Message: "parse error: " + err.Error()}, + }) + continue + } + + go s.handleRPC(ctx, req, enc) + } +} + +// CWD returns the server's current working directory. +func (s *Server) CWD() string { + s.wdMu.RLock() + defer s.wdMu.RUnlock() + return s.cwd +} + +// handleRPC dispatches a single JSON-RPC request. +// +// Reserved method names (control plane): +// - list_tools, host_info, ping, set_env, set_cwd +// - detach, list_detached, signal_detached, get_detached_output, dismiss_detached +// - tool_load +// +// Any other method name is treated as a tool invocation (data plane). +// A tool with the same name as a reserved method is unreachable. +func (s *Server) handleRPC(ctx context.Context, req RPCRequest, enc *lockedEncoder) { + switch req.Method { + case "list_tools": + tools, err := s.ListTools() + if err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) + } else { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: tools}) + } + + case "host_info": + info := HostInfo{ + Platform: runtime.GOOS, + Arch: runtime.GOARCH, + IsGitRepo: paths.IsGitRepo(s.CWD()), + } + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: info}) + + case "ping": + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: "pong"}) + + case "set_env": + var params struct { + Key string `json:"key"` + Value string `json:"value"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + s.SetEnv(params.Key, params.Value) + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) + + case "set_cwd": + var params struct { + Dir string `json:"dir"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + s.SetCWD(params.Dir) + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) + + case "detach": + ok := s.Detach() + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) + + case "list_detached": + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: s.ListDetachedRaw()}) + + case "signal_detached": + var params struct { + PID int `json:"pid"` + Signal int `json:"signal"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + if err := s.SignalDetached(params.PID, syscall.Signal(params.Signal)); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) + } else { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: true}) + } + + case "tool_load": + var params struct { + Name string `json:"name"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + if err := s.LoadTool(params.Name); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) + } else { + res, _ := json.Marshal(map[string]string{"result": "loaded: " + params.Name}) + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(res)}) + } + + case "get_detached_output": + var params struct { + PID int `json:"pid"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + out, err := s.GetDetachedOutput(params.PID) + if err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) + } else { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: out}) + } + + case "dismiss_detached": + var params struct { + PID int `json:"pid"` + } + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32602, Message: err.Error()}}) + return + } + ok := s.DismissDetached(params.PID) + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: ok}) + + default: + // Data plane: method name is the tool name. + streamCtx := WithOutputStream(ctx, func(data string) { + enc.Encode(RPCResponse{ + JSONRPC: "2.0", + ID: req.ID, + Result: OutputNotification{Data: data}, + Stream: true, + }) + }) + result, err := s.CallTool(streamCtx, req.Method, req.Params) + if err != nil { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Error: &RPCError{Code: -32000, Message: err.Error()}}) + } else { + enc.Encode(RPCResponse{JSONRPC: "2.0", ID: req.ID, Result: json.RawMessage(result)}) + } + } +} diff --git a/toolsrv/rpcwire.go b/toolsrv/rpcwire.go index 03e4d3d..ea07b46 100644 --- a/toolsrv/rpcwire.go +++ b/toolsrv/rpcwire.go @@ -20,3 +20,32 @@ type RPCError struct { type OutputNotification struct { Data string `json:"data"` } + +// ToolResult is the JSON envelope returned by tool execution. +// Tool scripts write this to stdout; the agent loop parses it via +// extractToolResult in agent/runtime.go. +// +// Schema: +// +// { +// "isError": false, +// "content": [ +// {"type": "text", "text": "..."}, +// {"type": "image", "media_type": "image/png", "data": ""} +// ] +// } +// +// If a tool emits raw text (not JSON), the agent treats it as a single +// text content block with isError=false. +type ToolResult struct { + IsError bool `json:"isError,omitempty"` + Content []ToolResultContent `json:"content"` +} + +// ToolResultContent is a single content block within a ToolResult. +type ToolResultContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + MediaType string `json:"media_type,omitempty"` + Data string `json:"data,omitempty"` +} diff --git a/toolsrv/server.go b/toolsrv/server.go index 153c329..2d2eeb3 100644 --- a/toolsrv/server.go +++ b/toolsrv/server.go @@ -244,13 +244,13 @@ func (s *Server) callPromotedTool(ctx context.Context, tool string, args json.Ra result, err = s.executeWithStdin(ctx, code, "bash", timeout, sandboxName, false, stdinData) } if err != nil { - return json.Marshal(map[string]interface{}{ - "isError": true, - "content": []map[string]string{{"type": "text", "text": result + ": " + err.Error()}}, + return json.Marshal(ToolResult{ + IsError: true, + Content: []ToolResultContent{{Type: "text", Text: result + ": " + err.Error()}}, }) } - return json.Marshal(map[string]interface{}{ - "content": []map[string]string{{"type": "text", "text": result}}, + return json.Marshal(ToolResult{ + Content: []ToolResultContent{{Type: "text", Text: result}}, }) }