lkn
/
ollie
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

maintainability: docs, contracts, structural cleanup 

- 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
This commit is contained in:
Ollie Agent 2026-08-09 21:45:56 +02:00
parent 2d74ffb95e
commit fbe8059cf5
15 changed files with 539 additions and 878 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
AGENTS.md
View File

@ -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` |

38
agent/agent.go
View File

@ -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" — <think>…</think> 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: <name>"
// - "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
}

10
agent/loop.go
View File

@ -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.

24
agent/runtime.go
View File

@ -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": "<base64>"}
// ]
// }
//
// 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
}

272
cmd/ollie-remote/main.go
View File

@ -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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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),
})
}
}
}

54
doc/boot-sequence.md Normal file
View File

@ -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.

17
fsedsl/tree.go
View File

@ -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

2
session/persist.go
View File

@ -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,

2
session/registry.go
View File

@ -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,

61
session/session.go
View File

@ -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

0
toolsrv/shell.go → toolsrv/exec.go
View File

625
toolsrv/rpc_integration_test.go
View File

@ -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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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)
}
})
}
}

263
toolsrv/rpcserver.go Normal file
View File

@ -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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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, &params); 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)})
}
}
}

29
toolsrv/rpcwire.go
View File

@ -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": "<base64>"}
// ]
// }
//
// 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"`
}

10
toolsrv/server.go
View File

@ -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}},
})
}