feat(itr,security,rlm): add Isolated Tool Runtime core + RLM engine
## ITR Command Protocol (pkg/itr) FlatBuffers-schema (commands.fbs) + Go encoding layer for the structured command protocol used between agent and SecureBus. Operations: Peek/Grep/Partition/Recurse for RLM decomposition; ToolExec/ExecWasm for direct tool execution; Final for terminal answers. Includes ToolRequest/ToolResponse envelopes with request-ID and error fields. DAG sub-package (pkg/itr/dag): - resolver.go: infer task dependency graph from tool call arguments - executor.go: execute DAG layers in parallel with barrier semantics ## SecretStore + KeyringProvider (pkg/security) - KeyringProvider interface abstracts OS keyring backends for master-key storage - NoopKeyring (in-memory) and EnvKeyring (reads PICOCLAW_MASTER_KEY) impls - SecretStore: AES-GCM encrypted, JSON-persisted named-secret store backed by any KeyringProvider; supports Set/Get/Delete/List/Has ## SecureBus (pkg/security/securebus) Mediates all tool execution through a 6-stage pipeline: 1. PolicyEngine — validate endpoint/path rules, recursion depth, SSRF 2. SecretInjection — inject named secrets into args or HTTP headers 3. Tool execution — delegate to registered ToolExecutor 4. LeakScanner — scan output for injected secret values, redact on hit 5. AuditLog — append AuditEvent to in-memory log (pluggable AuditSink) 6. ChannelTransport for in-process request/response routing ## RLM Engine (pkg/rlm) Recursive Language Model context processing: - Rope: pure-Go O(log n) Rope data structure (Append/Slice/Lines/GrepLines/Partition) - StrategyPlanner: heuristic op selection (OpFinal/OpGrep/OpPartition) - FanOut: bounded-concurrency parallel partition processor with MergeResults - Engine: recursive Answer loop; delegates sub-queries through SecureBus ## ADR documentation (docs/adr) - ADR-001: Isolated Tool Runtime full design document - ADR README index
This commit is contained in:
ZanzyTHEbar
parent
70614c8e54
commit
90c70e8cb0
20 changed files with 3917 additions and 0 deletions
926
docs/adr/001-isolated-tool-runtime.md
Normal file
926
docs/adr/001-isolated-tool-runtime.md
Normal file
|
|
@ -0,0 +1,926 @@
|
|||
# ADR-001: Isolated Tool Runtime (ITR) + DAG Task Executor
|
||||
|
||||
**Date**: 2026-02-18 (updated 2026-02-19)
|
||||
**Status**: Proposed
|
||||
**Authors**: @ZanzyTHEbar
|
||||
|
||||
---
|
||||
|
||||
## Context
|
||||
|
||||
PicoClaw tools execute in-process with the agent loop. The `Vault` (XChaCha20-Poly1305) exists for encrypting secrets at rest, but there is no pipeline for injecting those secrets into tool execution. The `Redactor` scans for sensitive patterns, but only in log paths — not on tool output before it reaches the LLM. There is no privilege boundary
|
||||
between the LLM-facing agent code and the tool execution path.
|
||||
|
||||
A compromised tool — via prompt injection, malicious skill, or supply chain attack — has the same memory-space access as the agent itself. This is the same class of vulnerability that led to the OpenClaw token exfiltration incident (Feb 2026), where malicious skills on ClawHub could read API keys from the host environment and exfiltrate them through tool output.
|
||||
|
||||
Competing frameworks have responded:
|
||||
|
||||
- **IronClaw** (NEAR AI, Rust): WASM container isolation, capability-based permissions, encrypted credential vault with runtime injection, network interception for leak detection.
|
||||
- **WyrmLock** (ZanzyTHEbar, Rust): ZKP-based application locking via proc connector, cryptographic authentication before process execution, system keyring integration.
|
||||
|
||||
PicoClaw's constraint is unique: the binary must remain under 20MB and run on 64MB RAM embedded boards. A full WASM runtime or separate daemon process is not viable as a mandatory dependency. The solution must be **progressive** — lightweight in-process enforcement by default, with optional heavier isolation for richer platforms.
|
||||
|
||||
**RLM convergence**: Independently, Recursive Language Models (Zhang & Khattab, MIT CSAIL, Oct 2025; arXiv:2512.24601 v2) demonstrate that long-context systems suffer from "context rot" — performance degrades as context length grows, not from technical limits but from information overload and distributional mismatch. RLM's solution is **context-as-variable** + **recursive decomposition**: the LM never sees the full context directly. Instead, it interacts with context through structured operations (peek, grep, partition, recurse) in a REPL-like environment, spawning recursive sub-calls over partitioned subsets.
|
||||
|
||||
RLM(GPT-5-mini) outperforms base GPT-5 by 114% on hard long-context tasks (OOLONG benchmark, 263k tokens) while maintaining roughly equivalent cost. At 10M+ tokens (1000 documents, BrowseComp-Plus), RLM maintains perfect performance where standard approaches collapse entirely.
|
||||
|
||||
The architectural insight that connects these two concerns: **RLM recursive calls are tool calls**. Every `peek`, `grep`, `partition`, `recurse`, and `exec_wasm` operation is a tool invocation. These operations are the highest-risk execution path in the system because they involve LM-generated commands (prompt injection vector), process
|
||||
potentially sensitive context (data exfiltration vector), and spawn sub-calls that could compound attacks (recursive privilege escalation).
|
||||
|
||||
The ITR's SecureBus is therefore not just a security layer — it is the RLM execution backbone. The two systems converge into one architecture.
|
||||
|
||||
**DAG executor convergence**: Three independent developments in AI tool orchestration address the same fundamental bottleneck — sequential, token-expensive tool execution — from complementary angles:
|
||||
|
||||
1. **LLMCompiler** (Kim et al., ICML 2024; arXiv:2312.04511): Treats multi-tool workflows like a compiler, constructing a Directed Acyclic Graph (DAG) of tool calls with explicit dependencies in a **single LM inference pass**, then executing nodes in topological order with parallel dispatch. Results: 3.7x latency speedup, 6.7x cost savings, ~9% accuracy improvement over ReAct. The extended **LLM-Tool Compiler** (arXiv:2405.17438) adds runtime operation fusion, achieving 4x more parallel calls and 40% token reduction.
|
||||
|
||||
2. **Anthropic Programmatic Tool Calling (PTC)** (Nov 2025): The LM writes orchestration code that calls multiple tools within a sandboxed execution environment. Tool results flow through the sandbox — **intermediate results never enter the LM's context**. Only the final/aggregated output is returned. Results: 37% token reduction on complex research tasks, 19+ inference passes eliminated per multi-tool workflow. Internal knowledge retrieval improved from 25.6% to 28.5%; GIA benchmarks from 46.5% to 51.2%.
|
||||
|
||||
3. **RLM** (already described above): Unbounded context via recursive decomposition.
|
||||
|
||||
These three systems compose naturally: **the LLMCompiler DAG planner produces the execution plan, the SecureBus DAG executor dispatches nodes in parallel with PTC-style context isolation, and the RLM engine activates when any node's context exceeds the LM's window**. PicoClaw's current agent loop (`RunToolLoop`) uses Fantasy's ReAct pattern — each tool call requires a full inference pass, and all intermediate results accumulate in context. The DAG executor replaces this sequential loop for complex multi-tool workflows while preserving ReAct for simple cases.
|
||||
|
||||
**Anthropic Tool Search** (Nov 2025): Separately, Anthropic's Tool Search Tool demonstrates that loading all tool definitions upfront (55K+ tokens for a modest 5-server setup) is itself a form of context pollution. On-demand tool discovery reduces token consumption by 85% while improving accuracy (Opus 4: 49% → 74%). This maps directly to a `ToolSearch` command variant in our FlatBuffers schema, enabling the DAG planner to discover tools lazily rather than seeing all definitions.
|
||||
|
||||
### Forces
|
||||
|
||||
- **Security**: The LLM must never see raw secrets. Tool output must be scanned for leaks before reaching the agent loop.
|
||||
- **Context rot**: Long-context performance degrades with scale. The agent needs structured, recursive context decomposition — but those recursive operations must be isolated and mediated.
|
||||
- **Token efficiency**: Sequential tool calling wastes inference passes and pollutes context with intermediate results. Multi-tool workflows should execute in parallel where dependencies allow, with only final results entering the LM's context.
|
||||
- **Embedded constraint**: < 20MB binary, 64MB RAM. No mandatory CGO, no mandatory external daemon.
|
||||
- **Backward compatibility**: Existing tools must continue working without modification. Security improvements are opt-in per tool.
|
||||
- **Auditability**: Every tool execution, secret access, and leak detection must be logged.
|
||||
- **Extensibility**: The architecture must accommodate WASM isolation, daemon-mode separation, RLM recursive decomposition, and DAG-based parallel execution without redesign.
|
||||
|
||||
---
|
||||
|
||||
## Decision
|
||||
|
||||
Introduce a **layered Isolated Tool Runtime (ITR)** with five progressive layers. Layers 1-2 are mandatory (P0). Layers 3-5 are opt-in, gated behind build tags or deployment configuration.
|
||||
|
||||
### Layer 1: Capability Manifests
|
||||
|
||||
Every tool may declare its security requirements through a new `CapableTool` interface. Tools that do not implement it receive **zero capabilities** (no secrets, no network beyond SSRF-validated URLs, filesystem restricted to workspace read-only, no shell).
|
||||
|
||||
```go
|
||||
type CapableTool interface {
|
||||
Tool
|
||||
Capabilities() ToolCapabilities
|
||||
}
|
||||
|
||||
type ToolCapabilities struct {
|
||||
Secrets []SecretRef
|
||||
Network []EndpointRule
|
||||
Filesystem []PathRule
|
||||
Shell ShellAccessLevel
|
||||
}
|
||||
|
||||
type SecretRef struct {
|
||||
Name string // logical name, e.g. "github_token"
|
||||
InjectAs string // "env:GITHUB_TOKEN" | "arg:token" | "header:Authorization"
|
||||
Required bool
|
||||
}
|
||||
|
||||
type EndpointRule struct {
|
||||
Pattern string // URL glob, e.g. "https://api.github.com/**"
|
||||
}
|
||||
|
||||
type PathRule struct {
|
||||
Pattern string // relative to workspace root
|
||||
Mode string // "r" | "w" | "rw"
|
||||
}
|
||||
|
||||
type ShellAccessLevel int
|
||||
|
||||
const (
|
||||
ShellNone ShellAccessLevel = iota
|
||||
ShellDenylist // default: block known-dangerous patterns
|
||||
ShellAllowlist // only explicitly permitted commands
|
||||
)
|
||||
```
|
||||
|
||||
**Rationale**: Capability declarations are static metadata — zero runtime cost. They enable the SecureBus (Layer 2) to enforce least-privilege without per-tool code changes. Tools that predate the interface get the most restrictive default.
|
||||
|
||||
### Layer 2: SecureBus — Privilege Boundary
|
||||
|
||||
A new `SecureBus` mediates **all** tool execution. The agent loop no longer calls `ToolRegistry.ExecuteWithContext()` directly — it submits requests to the SecureBus, which enforces policy, injects secrets, executes the tool, scans output, and logs the audit trail.
|
||||
|
||||
```
|
||||
Agent Loop
|
||||
│
|
||||
│ ToolRequest{name, args}
|
||||
▼
|
||||
┌──────────────────────────────┐
|
||||
│ SecureBus │
|
||||
│ │
|
||||
│ 1. Extract capabilities │
|
||||
│ 2. Validate against policy │
|
||||
│ 3. Decrypt + inject secrets │
|
||||
│ 4. Execute tool │
|
||||
│ 5. Scan output for leaks │
|
||||
│ 6. Write audit log │
|
||||
│ │
|
||||
└──────────────────────────────┘
|
||||
│
|
||||
│ ToolResponse{result, redacted}
|
||||
▼
|
||||
Agent Loop (never sees raw secrets)
|
||||
```
|
||||
|
||||
**Key component**: The `Transport` interface abstracts the communication channel between the agent loop and the SecureBus. In-process, this is a pair of Go channels (`ChannelTransport`). In daemon mode (Layer 4), it becomes a Unix domain socket (`SocketTransport`). The SecureBus code is identical in both cases.
|
||||
|
||||
```go
|
||||
type Transport interface {
|
||||
Send(ctx context.Context, req ToolRequest) (ToolResponse, error)
|
||||
Close() error
|
||||
}
|
||||
|
||||
type ChannelTransport struct { /* Go channels, in-process */ }
|
||||
type SocketTransport struct { /* Unix domain socket, daemon mode */ }
|
||||
type WasmTransport struct { /* wazero isolate, untrusted tools */ }
|
||||
```
|
||||
|
||||
**Command protocol**: Tool requests use a structured, schema-enforced format rather than free-form JSON. FlatBuffers (google/flatbuffers, zero-copy serialization) defines the canonical `ToolRequest` schema, covering both traditional tool calls and RLM recursive operations:
|
||||
|
||||
```flatbuffers
|
||||
namespace picoclaw.itr;
|
||||
|
||||
// --- Individual command types ---
|
||||
table Peek { start: uint64; length: uint32; }
|
||||
table Grep { pattern: string; max_matches: uint32; case_insensitive: bool; }
|
||||
table Partition { k: uint32; method: string; overlap: uint32; semantic: bool; }
|
||||
table Recurse { sub_query: string; context_key: string; depth_hint: uint8; }
|
||||
table ToolExec { tool_name: string; args_json: string; }
|
||||
table ExecWasm { module_key: string; entry: string; }
|
||||
table Final { answer: string; var_name: string; }
|
||||
table ToolSearch { query: string; max_results: uint8; }
|
||||
table CodeExec { code: string; language: string; } // opt-in only, wazero-isolated
|
||||
|
||||
// --- DAG types (LLMCompiler convergence) ---
|
||||
table DAGNode {
|
||||
id: string;
|
||||
payload: CommandPayload;
|
||||
depends_on: [string]; // node IDs this depends on; #nodeN in args resolve to output
|
||||
}
|
||||
|
||||
table DAGPlan {
|
||||
nodes: [DAGNode];
|
||||
max_parallel: uint8; // concurrency limit (default: GOMAXPROCS)
|
||||
token_budget: uint32; // total token budget for the plan
|
||||
joiner_query: string; // synthesis prompt for final aggregation
|
||||
}
|
||||
|
||||
union CommandPayload { Peek, Grep, Partition, Recurse, ToolExec, ExecWasm,
|
||||
Final, ToolSearch, CodeExec, DAGPlan }
|
||||
|
||||
table ToolRequest {
|
||||
id: string;
|
||||
payload: CommandPayload;
|
||||
timestamp: uint64;
|
||||
depth: uint8;
|
||||
session_key: string;
|
||||
}
|
||||
|
||||
table ToolResponse {
|
||||
id: string;
|
||||
result: string;
|
||||
is_error: bool;
|
||||
leak_detected: bool;
|
||||
cost_tokens: uint32;
|
||||
}
|
||||
```
|
||||
|
||||
This schema serves four purposes: (1) zero-copy reads eliminate serialization overhead on the hot path, (2) the union type prevents malformed requests — the LLM cannot construct a request outside the defined vocabulary, (3) the same binary format works across all transports (channels, sockets, WASM host calls), and (4) `DAGPlan` is itself a `CommandPayload` variant, enabling recursive DAG expansion (a node in one DAG can contain a sub-DAG). Traditional tool calls use `ToolExec`; RLM operations use `Peek` / `Grep` / `Partition` / `Recurse`; multi-tool workflows use `DAGPlan`; tool discovery uses `ToolSearch`; untrusted code uses `ExecWasm` or `CodeExec` (wazero-isolated). FlatBuffers Go codegen (`flatc --go`) produces type-safe accessors with no reflection.
|
||||
|
||||
**Secret injection**: The SecureBus resolves each `SecretRef` from the `SecretStore`, decrypts via `Vault`, and injects the plaintext into a scoped execution context according to the `InjectAs` spec. The decrypted value exists only in that context's memory and is zeroed after tool execution returns. The `args` map that the LLM constructed never contains secrets.
|
||||
|
||||
**Leak scanning**: Every tool result passes through `Redactor.ContainsSensitive()`. If a match is found, the result is redacted before it reaches the agent loop, and a `LEAK_DETECTED` event is written to the audit log. This catches accidental exposure — for example, a shell command that prints an environment variable containing an API key.
|
||||
|
||||
**Audit log**: Extends the existing `agent_audit_log` table with columns for `capability_grants`, `secrets_accessed`, `leak_detected`, and `policy_violation`. Append-only. Every tool execution produces exactly one audit row.
|
||||
|
||||
### Layer 3: Secret Store + Keyring Integration
|
||||
|
||||
The `SecretStore` maps logical secret names to encrypted ciphertext, persisted to a local file (`~/.picoclaw/secrets.enc`). The `Vault` handles encryption/decryption.
|
||||
|
||||
The master key for the `Vault` is sourced from one of three backends, selected at
|
||||
onboarding:
|
||||
|
||||
| Backend | Platform | Mechanism |
|
||||
|---------------|--------------|----------------------------------------------|
|
||||
| OS Keyring | Linux/macOS | libsecret (GNOME), kwallet (KDE), Keychain |
|
||||
| Passphrase | Any | Argon2id KDF from user passphrase |
|
||||
| File | Embedded | Raw key file with restricted permissions |
|
||||
|
||||
Keyring support is gated behind a build tag (`!embedded`) to avoid pulling in CGO or D-Bus dependencies on constrained platforms.
|
||||
|
||||
**CLI surface**:
|
||||
|
||||
```
|
||||
picoclaw secret add <name> # interactive prompt for value
|
||||
picoclaw secret list # names only, no values
|
||||
picoclaw secret delete <name>
|
||||
picoclaw secret export # encrypted backup
|
||||
picoclaw secret import <file> # restore from backup
|
||||
```
|
||||
|
||||
The `onboard` command is extended to include master key setup as part of the interactive wizard.
|
||||
|
||||
### Layer 4: Daemon Mode + ZKP Authentication
|
||||
|
||||
For non-embedded deployments (desktop, server), the SecureBus can optionally run in a separate privileged daemon process. The agent loop connects as an unprivileged client over a Unix domain socket.
|
||||
|
||||
```
|
||||
┌──────────────────┐ Unix Socket ┌──────────────────┐
|
||||
│ Client Process │ ◄────── Flatc-RPC ────────► │ Daemon Process │
|
||||
│ │ │ │
|
||||
│ Agent Loop │ │ SecureBus │
|
||||
│ LLM Provider │ │ Vault + Secrets │
|
||||
│ Memory Store │ │ Tool Runtime │
|
||||
│ │ │ Audit Log │
|
||||
└──────────────────┘ └──────────────────┘
|
||||
```
|
||||
|
||||
**Session establishment** uses a Schnorr-based zero-knowledge proof:
|
||||
|
||||
1. Client connects to `~/.picoclaw/daemon.sock`
|
||||
2. Daemon sends a random 32-byte challenge `c`
|
||||
3. Client computes commitment and response from passphrase-derived secret
|
||||
4. Daemon verifies against stored verifier (passphrase never crosses the socket)
|
||||
5. On success: daemon returns an HMAC-signed session token (TTL: 1h, configurable)
|
||||
6. Subsequent requests carry the session token (fast path — no ZKP per call)
|
||||
7. Re-authentication required for: first access to a new secret, policy override requests, or token expiry
|
||||
|
||||
**Why Schnorr**: Single round-trip. ~200 bytes on the wire. Pure Go implementation (no CGO). Well-understood security properties. Meets the constraint that even daemon mode should not require heavyweight cryptographic libraries.
|
||||
|
||||
**Why ZKP over simpler auth**: The passphrase never traverses the socket, even encrypted. If an attacker can observe the socket (e.g., via a compromised tool that gained filesystem read access to the socket file), they learn nothing about the passphrase. This is defense-in-depth against the exact attack vector we are trying to prevent.
|
||||
|
||||
**Daemon lifecycle**:
|
||||
|
||||
```
|
||||
picoclaw daemon start # background, creates pidfile + socket
|
||||
picoclaw daemon stop # graceful shutdown, zeroes key material
|
||||
picoclaw daemon status # running/stopped, uptime, client count
|
||||
```
|
||||
|
||||
Systemd and launchd service templates are provided in `deploy/`.
|
||||
|
||||
### Layer 5: WASM Isolates via wazero
|
||||
|
||||
For maximum isolation of untrusted tools (community skills, MCP-bridged tools) and RLM recursive sub-calls over sensitive context, tools can optionally run inside WASM sandboxes via **wazero** (`github.com/tetratelabs/wazero`, v1.11.0).
|
||||
|
||||
**Why wazero over wasmtime-go**: wazero is written in pure Go with zero CGO dependencies (one dep: `golang.org/x/sys`). It supports WebAssembly Cor Spec 1.0/2.0 an WASI preview 1. Its AOT compiler mode achieves ~10x interpreter performance. This is the critical difference from wasmtime-go, which requires CGO and adds ~8MB to the binary. wazero adds ~2-3MB and compiles on every Go target — including embedded boards.
|
||||
|
||||
Each WASM module instance gets:
|
||||
- **Private linear memory**: 1-64MB configurable, inaccessible from other instances
|
||||
or the host process. This is the V8 isolate analog — per-task private heap.
|
||||
- **Capability-gated WASI imports**: The host exposes only the imports declared in
|
||||
the tool's `CapableTool` manifest. No filesystem access unless granted. No network
|
||||
unless granted. No clock, no random, no env vars unless explicitly provided.
|
||||
- **Resource limits**: Memory ceiling, fuel-based CPU metering (wazero `RuntimeConfig.
|
||||
WithMemoryLimitPages()` and fuel counters), wall-clock timeout via Go context.
|
||||
- **Zero shared state**: Each invocation instantiates a fresh module. No persistent
|
||||
state leaks between calls.
|
||||
|
||||
**RLM use case**: When the RLM engine partitions context and spawns recursive sub-
|
||||
calls, each sub-call can optionally execute in a wazero isolate. The sub-call sees
|
||||
only its assigned context partition (injected via WASI stdin or a virtual filesystem
|
||||
mount) and its specific query. It cannot access the parent's full context, other
|
||||
partitions, or the host's secrets. The `WasmTransport` variant of the `Transport`
|
||||
interface handles serialization across the host-guest boundary using the same
|
||||
FlatBuffers command protocol.
|
||||
|
||||
```go
|
||||
type WasmTransport struct {
|
||||
runtime wazero.Runtime
|
||||
module wazero.CompiledModule
|
||||
config wazero.ModuleConfig // memory limits, WASI capabilities
|
||||
}
|
||||
```
|
||||
|
||||
**Compilation targets**: Tools written in Go compile via `GOOS=wasip1 GOARCH=wasm
|
||||
go build`. Tools written in Rust compile via `cargo build --target wasm32-wasip1`.
|
||||
TinyGo produces smaller binaries (~100KB-1MB) suitable for embedded constraints.
|
||||
|
||||
### RLM Integration: Recursive Context Decomposition Engine
|
||||
|
||||
The Isolated Tool Runtime subsumes and secures the RLM execution model. Rather than
|
||||
implementing RLM as a standalone system, PicoClaw treats RLM operations as first-class
|
||||
tool calls flowing through the SecureBus. This section describes the concrete
|
||||
integration.
|
||||
|
||||
**Pure intention of RLM** (stripped of REPL ideology): Enable LMs to process unbounded
|
||||
context by treating it as a programmable variable, where the LM autonomously decides
|
||||
how to recursively query, transform, and synthesize subsets — sidestepping context rot
|
||||
entirely. Formally:
|
||||
|
||||
```
|
||||
RLM(query, context) = Aggregate{ Map{ Decompose(context, query) } }
|
||||
```
|
||||
|
||||
Where `Decompose` is LM-driven (peek/grep/partition), `Map` fans out to parallel
|
||||
recursive calls, and `Aggregate` synthesizes child results via LM.
|
||||
|
||||
**Data structures**:
|
||||
|
||||
- **Rope** (`zyedidia/rope`): O(log n) substring, split, and concatenate operations
|
||||
on large contexts. A 10M-token context stored as a rope allows efficient `Peek`
|
||||
(slice) and `Partition` (split at k boundaries) without copying the full buffer.
|
||||
This replaces naive `[]byte` slicing which degrades to O(n) on large contexts.
|
||||
|
||||
- **Recursion tree**: A balanced tree tracking active sub-calls, their depth, context
|
||||
partition assignments, and aggregated results. Implemented as a `sync.Map` keyed by
|
||||
call ID, enabling concurrent goroutine access without global locks.
|
||||
|
||||
- **Context index**: Persistent storage via the existing libSQL/Turso infrastructure.
|
||||
Each context partition is addressable by key for cross-call reference. The memory
|
||||
delegate (`pkg/memory/delegate/sqlite.go`) already provides the CRUD primitives.
|
||||
|
||||
**Execution flow through SecureBus**:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph agentLoop ["Agent Loop"]
|
||||
Q["Query + Context Key"]
|
||||
end
|
||||
|
||||
subgraph secureBusOuter ["SecureBus"]
|
||||
direction TB
|
||||
ST["Strategy Call"]
|
||||
D["Dispatch"]
|
||||
end
|
||||
|
||||
subgraph rlmOps ["RLM Operations (parallel goroutines)"]
|
||||
P["Peek: Rope.Slice"]
|
||||
G["Grep: Aho-Corasick / regexp"]
|
||||
PA["Partition: Rope.Split k-way"]
|
||||
R["Recurse: sub-RLM via SecureBus"]
|
||||
W["ExecWasm: wazero isolate"]
|
||||
end
|
||||
|
||||
subgraph agg ["Aggregation"]
|
||||
AG["LM Synthesize"]
|
||||
F["Final Answer"]
|
||||
end
|
||||
|
||||
Q -->|"FlatBuffers ToolRequest"| ST
|
||||
ST -->|"cheap sub-LM decides strategy"| D
|
||||
D --> P & G & PA & R & W
|
||||
P & G & PA & R & W -->|"ToolResponse per partition"| AG
|
||||
AG --> F
|
||||
F -->|"redacted result"| agentLoop
|
||||
```
|
||||
|
||||
**Strategy planning**: A cheap sub-LM (e.g., local Ollama model) receives only the
|
||||
query and context metadata (length, detected format via `AnalyzeContext`) — never the
|
||||
full context. It outputs a FlatBuffers `ToolRequest` selecting the operation type and
|
||||
parameters. This is the "LM decides how to decompose" principle from the RLM paper,
|
||||
constrained to the structured command vocabulary defined in the FlatBuffers schema.
|
||||
|
||||
**Parallel execution**: `Partition` splits the rope into k chunks. Each chunk spawns
|
||||
a goroutine that submits a `Recurse` request back through the SecureBus. The SecureBus
|
||||
enforces capabilities on each sub-call independently — a partition processing financial
|
||||
data gets different capability grants than one processing public documentation.
|
||||
`errgroup.Group` manages the fan-out with configurable concurrency limits.
|
||||
|
||||
**Depth control**: The `depth` field in `ToolRequest` increments on each recursive
|
||||
call. The `PolicyEngine` enforces a maximum depth (default: 3, configurable). At
|
||||
`depth == max_depth` or when context fits within the LM's window (`|ctx| <= threshold`),
|
||||
the SecureBus dispatches a direct LM call instead of further decomposition.
|
||||
|
||||
**Cost tracking**: Each `ToolResponse` includes `cost_tokens`. The SecureBus
|
||||
aggregates token usage across the recursion tree and enforces a per-query budget.
|
||||
When the budget is exhausted, remaining sub-calls are canceled via Go context
|
||||
propagation and the best partial result is synthesized.
|
||||
|
||||
**Integration with existing memory tiers**: RLM operates on the archival tier of
|
||||
PicoClaw's 3-tier memory system. When the agent issues a `memory search` that returns
|
||||
results exceeding the context window, the RLM engine activates automatically — treating
|
||||
the search results as the context variable and recursively decomposing them. This is
|
||||
transparent to the agent: it issued a search, it gets a synthesized answer.
|
||||
|
||||
**Go implementation skeleton**:
|
||||
|
||||
```go
|
||||
type RLMEngine struct {
|
||||
bus *securebus.SecureBus
|
||||
primary LLMClient // synthesis + direct calls
|
||||
cheap LLMClient // strategy planning + sub-calls
|
||||
maxDepth int
|
||||
maxTokens int
|
||||
}
|
||||
|
||||
func (r *RLMEngine) Complete(ctx context.Context, query string, contextKey string) (string, error) {
|
||||
rope, err := r.bus.LoadContext(ctx, contextKey)
|
||||
if err != nil { return "", err }
|
||||
return r.recurse(ctx, query, rope, 0, 0)
|
||||
}
|
||||
|
||||
func (r *RLMEngine) recurse(ctx context.Context, query string, rp *rope.Rope, depth, tokens int) (string, error) {
|
||||
if depth > r.maxDepth || rp.Len() <= threshold {
|
||||
return r.primary.Direct(ctx, query, rp.String())
|
||||
}
|
||||
|
||||
strategy := r.cheap.Strategy(ctx, query, rp.Len()) // returns FlatBuffers command
|
||||
|
||||
switch strategy.PayloadType() {
|
||||
case itr.CommandPayloadPeek:
|
||||
sub := rp.Slice(strategy.Start(), strategy.Length())
|
||||
return r.recurse(ctx, query, sub, depth, tokens)
|
||||
|
||||
case itr.CommandPayloadGrep:
|
||||
matches := parallelGrep(rp, strategy.Pattern())
|
||||
return r.fanOutAndAggregate(ctx, query, matches, depth+1, tokens)
|
||||
|
||||
case itr.CommandPayloadPartition:
|
||||
chunks := rp.SplitN(strategy.K())
|
||||
return r.fanOutAndAggregate(ctx, query, chunks, depth+1, tokens)
|
||||
|
||||
case itr.CommandPayloadExecWasm:
|
||||
resp, err := r.bus.Execute(ctx, strategy.AsToolRequest())
|
||||
return resp.Result(), err
|
||||
}
|
||||
return "", fmt.Errorf("unknown strategy")
|
||||
}
|
||||
```
|
||||
|
||||
### DAG Task Executor: LLMCompiler + PTC Convergence
|
||||
|
||||
The DAG executor is the orchestration layer that unifies LLMCompiler's parallel
|
||||
planning with PTC's context isolation principle. It lives inside the SecureBus —
|
||||
every DAG node is a `ToolRequest` that flows through capability checking, secret
|
||||
injection, and leak scanning. The DAG is a batch of `ToolRequest`s with dependency
|
||||
metadata.
|
||||
|
||||
**Why this replaces ReAct for complex workflows**: PicoClaw's current `RunToolLoop`
|
||||
uses Fantasy's ReAct pattern. Each tool call requires a full LLM inference pass
|
||||
(hundreds of ms to seconds), and all intermediate results accumulate in the LLM's
|
||||
context window. For a 20-tool workflow, that's 20 inference passes and potentially
|
||||
200KB+ of intermediate data polluting context. The DAG executor reduces this to
|
||||
**2 inference passes** (plan + synthesize) with **zero intermediate results in
|
||||
context**.
|
||||
|
||||
**Execution model**:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph planning ["Phase 1: Plan (single LLM inference pass)"]
|
||||
Q["Query: 'Which team members exceeded Q3 travel budget?'"]
|
||||
P["LLM Planner outputs DAGPlan"]
|
||||
end
|
||||
|
||||
subgraph execution ["Phase 2: Execute (SecureBus DAG Executor)"]
|
||||
direction TB
|
||||
TOPO["Topological sort → execution waves"]
|
||||
|
||||
subgraph wave1 ["Wave 1 (parallel)"]
|
||||
N1["node1: ToolExec{get_team_members}"]
|
||||
end
|
||||
|
||||
subgraph wave2 ["Wave 2 (parallel, depends on node1)"]
|
||||
N2["node2: ToolExec{get_budget_by_level, #node1.levels}"]
|
||||
N3["node3: ToolExec{get_expenses, #node1.ids}"]
|
||||
end
|
||||
|
||||
subgraph wave3 ["Wave 3 (depends on node2 + node3)"]
|
||||
N4["node4: CodeExec or Final{filter exceeded}"]
|
||||
end
|
||||
end
|
||||
|
||||
subgraph synthesis ["Phase 3: Synthesize"]
|
||||
J["Only node4 output enters LLM context"]
|
||||
A["LLM synthesizes final answer from compact result"]
|
||||
end
|
||||
|
||||
Q --> P
|
||||
P --> TOPO
|
||||
TOPO --> wave1
|
||||
wave1 --> wave2
|
||||
wave2 --> wave3
|
||||
wave3 --> J
|
||||
J --> A
|
||||
|
||||
style planning fill:#4B0082,color:#fff,stroke:#333
|
||||
style execution fill:#003366,color:#fff,stroke:#333
|
||||
style synthesis fill:#006400,color:#fff,stroke:#333
|
||||
style wave1 fill:#1a1a5e,color:#fff,stroke:#555
|
||||
style wave2 fill:#1a1a5e,color:#fff,stroke:#555
|
||||
style wave3 fill:#1a1a5e,color:#fff,stroke:#555
|
||||
```
|
||||
|
||||
**Token savings breakdown**:
|
||||
|
||||
| Approach | Inference Passes | Context Consumed | Latency |
|
||||
|----------|-----------------|-----------------|---------|
|
||||
| ReAct (current) | N per tool call | All intermediate results | Sequential |
|
||||
| LLMCompiler DAG | 1 plan + 1 synthesize | Final result only | Parallel waves |
|
||||
| DAG + RLM | 1 plan + 1 synthesize + cheap sub-LM for strategy | Final result only | Parallel waves + recursive expansion |
|
||||
|
||||
From the research: LLMCompiler achieves up to **6.7x cost savings** and **3.7x
|
||||
latency speedup** over ReAct. PTC adds **37% token reduction** by keeping
|
||||
intermediate results out of context. Combined with RLM's ability to handle 10M+
|
||||
token contexts, the unified system handles workflows that would be impossible
|
||||
with sequential tool calling.
|
||||
|
||||
**Core algorithm — topological dispatch with dependency resolution**:
|
||||
|
||||
```go
|
||||
type DAGExecutor struct {
|
||||
bus *securebus.SecureBus
|
||||
rlm *rlm.RLMEngine
|
||||
planner LLMClient // primary LM for DAG generation
|
||||
joiner LLMClient // can be same or cheaper LM for synthesis
|
||||
maxParallel int
|
||||
}
|
||||
|
||||
type nodeState struct {
|
||||
node *DAGNode
|
||||
result atomic.Value
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
func (d *DAGExecutor) Execute(ctx context.Context, plan *DAGPlan) (string, error) {
|
||||
states := make(map[string]*nodeState, len(plan.Nodes))
|
||||
for _, n := range plan.Nodes {
|
||||
states[n.ID] = &nodeState{node: n, done: make(chan struct{})}
|
||||
}
|
||||
|
||||
g, gctx := errgroup.WithContext(ctx)
|
||||
g.SetLimit(int(plan.MaxParallel))
|
||||
|
||||
for _, ns := range states {
|
||||
ns := ns
|
||||
g.Go(func() error {
|
||||
// Wait for all dependencies
|
||||
for _, dep := range ns.node.DependsOn {
|
||||
select {
|
||||
case <-states[dep].done:
|
||||
case <-gctx.Done():
|
||||
return gctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Resolve #nodeN references in args
|
||||
resolved := resolveDependencyRefs(ns.node.Payload, states)
|
||||
|
||||
// RLM expansion: if resolved context exceeds threshold
|
||||
if needsRLMExpansion(resolved) {
|
||||
result, err := d.rlm.Complete(gctx, resolved.Query, resolved.ContextKey)
|
||||
if err != nil { return err }
|
||||
ns.result.Store(result)
|
||||
close(ns.done)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Standard execution through SecureBus
|
||||
req := toToolRequest(ns.node, resolved)
|
||||
resp, err := d.bus.Execute(gctx, req)
|
||||
if err != nil { return err }
|
||||
ns.result.Store(resp.Result())
|
||||
close(ns.done)
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
if err := g.Wait(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// Joiner: synthesize final result (PTC principle — only this enters LLM context)
|
||||
return d.joinResults(ctx, plan, states)
|
||||
}
|
||||
```
|
||||
|
||||
**Dependency resolution**: When a node's args contain `#node1`, the executor
|
||||
substitutes the completed output of `node1`. This is the same mechanism as
|
||||
LLMCompiler's reference resolution. In the FlatBuffers encoding, `args_json`
|
||||
strings contain `#nodeN` placeholders that are string-replaced before execution.
|
||||
|
||||
**Replanning loop**: LLMCompiler supports replanning — after a DAG completes,
|
||||
the Joiner can decide more work is needed and submit a new DAG. This maps to
|
||||
a simple outer loop:
|
||||
|
||||
```go
|
||||
func (d *DAGExecutor) RunWithReplanning(ctx context.Context, query string, maxReplans int) (string, error) {
|
||||
var history []string
|
||||
for i := 0; i <= maxReplans; i++ {
|
||||
plan, err := d.planner.PlanDAG(ctx, query, history)
|
||||
if err != nil { return "", err }
|
||||
if plan.IsFinal() {
|
||||
return plan.FinalAnswer(), nil
|
||||
}
|
||||
result, err := d.Execute(ctx, plan)
|
||||
if err != nil { return "", err }
|
||||
history = append(history, result)
|
||||
}
|
||||
return d.joiner.Synthesize(ctx, query, history), nil
|
||||
}
|
||||
```
|
||||
|
||||
**RLM as recursive DAG expansion**: When any DAG node's input context exceeds the
|
||||
LM's window, the executor transparently activates the RLM engine. The RLM's
|
||||
partition/recurse operations expand into a **sub-DAG** — the `Partition` operation
|
||||
becomes k child nodes with a `Final` aggregation node. This is recursive DAG
|
||||
expansion: a `DAGPlan` can contain nodes whose payloads are themselves `DAGPlan`s.
|
||||
The SecureBus enforces capabilities on each sub-node independently.
|
||||
|
||||
**Tool Search integration**: The `ToolSearch` command variant enables the DAG planner
|
||||
to discover tools on-demand rather than loading all definitions into context upfront.
|
||||
The planner emits a `ToolSearch` node as the first node in the DAG, whose output
|
||||
(matching tool schemas) feeds into subsequent `ToolExec` nodes via dependency
|
||||
references. This mirrors Anthropic's Tool Search Tool pattern: 85% token reduction
|
||||
on tool definitions while maintaining access to the full tool library.
|
||||
|
||||
**Routing: ReAct vs. DAG**: Not every query benefits from DAG planning. The agent
|
||||
loop maintains both execution paths:
|
||||
|
||||
- **ReAct** (Fantasy agent): Simple queries needing 1-3 tool calls. Lower planning
|
||||
overhead. Preserves the current behavior for backward compatibility.
|
||||
- **DAG Executor**: Complex workflows with parallelizable steps, large data
|
||||
processing, or multi-source gathering. Activated when the planner detects
|
||||
opportunity for parallelism or when the user explicitly requests structured
|
||||
execution.
|
||||
|
||||
A lightweight classifier (heuristic or cheap LM call) routes queries to the
|
||||
appropriate executor. The `ToolLoopConfig` gains a `Mode` field:
|
||||
`ModeReAct | ModeDAG | ModeAuto`.
|
||||
|
||||
**Relationship to PTC**: PicoClaw implements PTC's core principle — intermediate
|
||||
results never pollute the LM's context — without depending on Anthropic's specific
|
||||
code execution sandbox. The DAG executor IS the sandbox: it holds all intermediate
|
||||
node outputs in a `sync.Map`, resolves references internally, and only the Joiner's
|
||||
output crosses back to the agent loop. For cases requiring actual code execution
|
||||
within a node (e.g., data filtering, aggregation logic), the `CodeExec` command
|
||||
variant routes through wazero, providing the same safety guarantees as PTC's
|
||||
sandboxed Python environment but in pure Go with private linear memory per execution.
|
||||
|
||||
---
|
||||
|
||||
## Alternatives Considered
|
||||
|
||||
### A. Full WASM-first (IronClaw model)
|
||||
|
||||
Run all tools in WASM sandboxes from day one.
|
||||
|
||||
**Rejected because**: `wasmtime-go` adds ~8MB to the binary and requires CGO on some
|
||||
platforms. This breaks the < 20MB embedded target. WASM also imposes latency on every
|
||||
tool call (~5-10ms overhead) that is unnecessary for trusted built-in tools.
|
||||
|
||||
### B. Process-per-tool isolation (firejail / bubblewrap)
|
||||
|
||||
Fork a sandboxed process for each tool execution using Linux namespace isolation.
|
||||
|
||||
**Rejected because**: Linux-only. Not available on embedded boards running minimal
|
||||
kernels. Fork overhead (~2-5ms per call) adds up for agents making 20-50 tool calls
|
||||
per conversation turn.
|
||||
|
||||
### C. No privilege separation — just better validation
|
||||
|
||||
Keep tools in-process but add stricter input validation and output scanning.
|
||||
|
||||
**Rejected because**: Validation alone cannot prevent a tool from reading `os.Environ()`
|
||||
or accessing the `Vault` directly through shared memory. The fundamental problem is
|
||||
that tools and the agent share an address space with no enforcement boundary. Validation
|
||||
reduces attack surface but does not eliminate it.
|
||||
|
||||
### D. Always-on daemon (no in-process mode)
|
||||
|
||||
Require the daemon for all deployments.
|
||||
|
||||
**Rejected because**: Embedded boards cannot run two processes. The daemon adds
|
||||
operational complexity (service management, socket lifecycle) that is unnecessary for
|
||||
single-user local deployments. Making it optional respects the spectrum of deployment
|
||||
targets.
|
||||
|
||||
### E. Full REPL for RLM (Python/Yaegi arbitrary code execution)
|
||||
|
||||
The original RLM paper uses a Python REPL where the LM writes and executes arbitrary
|
||||
code. DSPy's `dspy.RLM` implements this with a Pyodide WASM sandbox. The Go equivalent
|
||||
would embed Yaegi (Go interpreter) or `go-embed-python`.
|
||||
|
||||
**Rejected as default because**: Arbitrary code execution is the attack surface we are
|
||||
trying to eliminate. A prompt-injected LM could write code that reads environment
|
||||
variables, opens network connections, or accesses the filesystem outside the workspace.
|
||||
Even sandboxed interpreters have escape histories (Yaegi: `os.Exit` kills the host
|
||||
process; embedded Python: C extension escapes). The structured command REPL (FlatBuffers
|
||||
schema with fixed vocabulary: peek/grep/partition/recurse/exec_wasm) preserves 100% of
|
||||
the RLM strategic gains — the LM still decides how to decompose — while constraining
|
||||
execution to operations the SecureBus can validate and audit. For cases requiring true
|
||||
code execution, `ExecWasm` routes through wazero isolates with private memory and no
|
||||
host access.
|
||||
|
||||
**Retained as opt-in**: A `CodeExec` command variant can be added to the FlatBuffers
|
||||
schema for deployments that accept the risk, with mandatory wazero isolation and a
|
||||
separate capability grant (`Shell` level or higher).
|
||||
|
||||
### F. Direct PTC adoption (Anthropic code execution sandbox)
|
||||
|
||||
Use Anthropic's Programmatic Tool Calling directly, where Claude writes Python
|
||||
orchestration code that runs in Anthropic's sandboxed environment.
|
||||
|
||||
**Rejected because**: PTC is provider-specific (requires Anthropic's
|
||||
`code_execution_20250825` server tool). PicoClaw is LLM-agnostic — it works with
|
||||
OpenAI, Anthropic, Ollama, and any Fantasy-compatible provider. The DAG executor
|
||||
implements PTC's core principle (intermediate results isolated from context) without
|
||||
provider lock-in. Additionally, PTC's Python sandbox cannot enforce PicoClaw's
|
||||
capability manifests or leak scanning — our SecureBus provides stronger guarantees.
|
||||
|
||||
### G. Pure LLMCompiler without RLM integration
|
||||
|
||||
Implement LLMCompiler's DAG planner and executor without recursive context
|
||||
decomposition.
|
||||
|
||||
**Rejected because**: LLMCompiler assumes tool inputs fit within the LM's context
|
||||
window. When a DAG node processes a 10M-token corpus (e.g., a memory search that
|
||||
returns thousands of archival entries), LLMCompiler alone cannot handle it — the
|
||||
node would need to be split. RLM provides exactly this capability: automatic
|
||||
expansion of oversized nodes into recursive sub-DAGs. The two systems are
|
||||
complementary, not alternatives.
|
||||
|
||||
---
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
|
||||
- **Secrets never reach the LLM**: The agent loop code path cannot access decrypted
|
||||
secrets. They exist only inside the SecureBus execution context, scoped to a single
|
||||
tool call and zeroed after.
|
||||
- **Leak detection is mandatory**: Every tool result is scanned. Accidental exposure
|
||||
through tool output is caught and redacted before the LLM can echo it back.
|
||||
- **Audit trail**: Complete record of what tools accessed which secrets, when, and
|
||||
whether leaks were detected. Invaluable for incident response.
|
||||
- **Backward compatible**: Existing tools work unchanged. New capabilities are opt-in
|
||||
per tool.
|
||||
- **Progressive hardening**: Embedded deployments get in-process enforcement (layers
|
||||
1-2). Desktop/server deployments add keyring + daemon (layers 3-4). High-security
|
||||
deployments add WASM (layer 5). Same codebase, same interfaces.
|
||||
- **Future-proof**: The `Transport` abstraction means new isolation backends (WASM,
|
||||
containers, remote execution) can be added without touching the SecureBus core.
|
||||
- **Unbounded context**: RLM recursive decomposition through the SecureBus enables
|
||||
processing of 10M+ token contexts without context rot, at roughly equivalent cost
|
||||
to a single LM call, while every recursive sub-call is capability-checked and
|
||||
leak-scanned.
|
||||
- **Zero-copy hot path**: FlatBuffers command protocol eliminates serialization
|
||||
overhead on the in-process transport. The same binary format works across channels,
|
||||
sockets, and WASM host calls without re-encoding.
|
||||
- **Pure Go WASM**: wazero (no CGO) means WASM isolation compiles on every Go target
|
||||
including embedded boards, unlike wasmtime-go which requires CGO and breaks cross-
|
||||
compilation.
|
||||
- **Up to 6.7x cost savings on multi-tool workflows**: The DAG executor eliminates
|
||||
redundant inference passes. A 20-tool workflow drops from 20+ LLM calls to 2
|
||||
(plan + synthesize), with parallel execution reducing wall-clock latency by 3.7x.
|
||||
- **Context isolation by default**: Intermediate tool results never enter the LLM's
|
||||
context window. The DAG executor's `sync.Map` holds all node outputs internally,
|
||||
resolving `#nodeN` references without LLM involvement. Only the Joiner's compact
|
||||
output crosses back to the agent loop.
|
||||
- **LLM-agnostic PTC**: PicoClaw achieves Anthropic PTC's benefits (37% token
|
||||
reduction, parallel execution, context isolation) without provider lock-in. Any
|
||||
Fantasy-compatible LLM can produce DAG plans.
|
||||
- **Recursive DAG expansion**: When RLM detects a node with oversized context, it
|
||||
expands that node into a sub-DAG transparently. This composes LLMCompiler's
|
||||
breadth (parallel multi-tool) with RLM's depth (recursive context decomposition).
|
||||
- **Graceful degradation**: Simple queries bypass the DAG executor entirely, using
|
||||
the proven ReAct loop. Complexity is only introduced when it pays off.
|
||||
|
||||
### Negative
|
||||
|
||||
- **Complexity**: The SecureBus adds an indirection layer to every tool call. Debugging
|
||||
tool execution now requires understanding the mediation pipeline.
|
||||
- **Performance**: In-process `ChannelTransport` adds ~1-2 microseconds per call
|
||||
(negligible). Daemon mode adds ~0.5-1ms per call (socket round-trip). WASM would
|
||||
add ~5-10ms. These are acceptable given tool calls already take 10ms-60s.
|
||||
- **Secret management UX**: Users must explicitly add secrets via CLI before tools can
|
||||
use them. This is an intentional trade-off — implicit secret discovery (e.g.,
|
||||
scanning `~/.env` files) is itself a security risk.
|
||||
- **Build tag surface**: Keyring and daemon code behind build tags adds CI matrix
|
||||
complexity. Mitigated by testing all tag combinations in CI.
|
||||
- **RLM cost variance**: Recursive decomposition has non-deterministic cost — the LM
|
||||
chooses how many partitions and sub-calls to make. Mitigated by per-query token
|
||||
budgets enforced in the SecureBus, with context cancellation on budget exhaustion.
|
||||
- **FlatBuffers tooling**: Requires `flatc` compiler in the build pipeline for codegen.
|
||||
Generated Go code is committed to the repo, so downstream consumers do not need
|
||||
`flatc`. Schema changes require regeneration.
|
||||
- **DAG planning quality**: The LLM must produce valid DAGs with correct dependency
|
||||
edges. Malformed DAGs (cycles, missing dependencies) are caught at validation time
|
||||
but waste an inference pass. Mitigated by structured output constraints (FlatBuffers
|
||||
schema validation) and few-shot examples in the planner prompt.
|
||||
- **Routing overhead**: The ReAct vs. DAG routing decision adds a classification step.
|
||||
If the classifier misroutes (e.g., sends a simple lookup through the DAG planner),
|
||||
overhead increases. Mitigated by defaulting to ReAct and only escalating when
|
||||
parallelism is detected.
|
||||
|
||||
### Risks
|
||||
|
||||
- **Policy misconfiguration**: Overly permissive policies could negate the security
|
||||
benefits. Mitigated by shipping restrictive defaults and requiring explicit opt-in
|
||||
for elevated capabilities.
|
||||
- **ZKP implementation correctness**: Rolling our own Schnorr implementation carries
|
||||
risk. Mitigated by using a well-tested reference implementation, extensive test
|
||||
vectors, and considering a future switch to an audited library if daemon mode sees
|
||||
wide adoption.
|
||||
- **Adoption friction**: Tool authors must implement `CapableTool` to access secrets.
|
||||
Mitigated by providing clear documentation, examples, and a helper function that
|
||||
builds `ToolCapabilities` from a TOML manifest.
|
||||
- **RLM strategy quality**: The cheap sub-LM may choose suboptimal decomposition
|
||||
strategies (e.g., too many partitions, wrong grep patterns). Mitigated by fallback
|
||||
to uniform partitioning when confidence is below threshold, and by trajectory
|
||||
logging that enables future RL fine-tuning of strategy selection.
|
||||
- **Rope memory overhead**: The rope data structure adds ~2x memory overhead compared
|
||||
to raw `[]byte` for small contexts. Mitigated by only constructing ropes when
|
||||
context exceeds the direct-call threshold (default: context window size).
|
||||
- **DAG explosion**: A poorly constrained planner could generate DAGs with hundreds
|
||||
of nodes, leading to excessive parallel tool calls. Mitigated by `max_parallel`
|
||||
limit in `DAGPlan`, `token_budget` enforcement, and a hard cap on node count
|
||||
(default: 50, configurable).
|
||||
- **Replanning loops**: The replanning mechanism could loop indefinitely if the
|
||||
Joiner never determines the result is sufficient. Mitigated by `maxReplans`
|
||||
limit (default: 3) and monotonic progress detection (if a replan produces no
|
||||
new information, terminate early).
|
||||
|
||||
---
|
||||
|
||||
## Implementation Plan
|
||||
|
||||
| Phase | Layer | Scope | Priority |
|
||||
|-------|-------|-------|----------|
|
||||
| 1 | 1 + 2 | `CapableTool` interface, FlatBuffers command schema (incl. DAG types), `SecureBus` with `ChannelTransport`, `SecretStore`, agent loop integration, leak scanning on all tool output, audit log extension | P0 |
|
||||
| 2 | Migrate | Built-in tools (`shell`, `filesystem`, `web`) implement `CapableTool` | P0 |
|
||||
| 3 | DAG | `DAGExecutor`: topological dispatch, dependency resolution, parallel wave execution via errgroup, Joiner synthesis, replanning loop. ReAct/DAG routing in agent loop. `ToolSearch` command variant. | P0 |
|
||||
| 4 | RLM | `RLMEngine` with rope context, strategy planning via cheap sub-LM, parallel fan-out through SecureBus, cost tracking, integration with archival memory tier, recursive DAG expansion | P0 |
|
||||
| 5 | 3 | Keyring integration, `picoclaw secret` CLI, onboard wizard extension | P1 |
|
||||
| 6 | 4 | `SocketTransport`, daemon mode, Schnorr ZKP handshake, systemd/launchd templates | P2 |
|
||||
| 7 | 5 | `WasmTransport` via wazero, WASM tool isolation for untrusted tools, `CodeExec` command variant, RLM sub-call isolation | P2 |
|
||||
|
||||
### Files Inventory
|
||||
|
||||
**New packages/files**:
|
||||
|
||||
| Path | Purpose |
|
||||
|------|---------|
|
||||
| `pkg/itr/commands.fbs` | FlatBuffers schema defining the command protocol (Peek, Grep, Partition, Recurse, ToolExec, ExecWasm, Final, ToolSearch, CodeExec, DAGNode, DAGPlan) |
|
||||
| `pkg/itr/commands_generated.go` | Generated Go code from `flatc --go` (committed, no build-time flatc dependency) |
|
||||
| `pkg/itr/dag/executor.go` | `DAGExecutor`: topological dispatch, parallel wave execution, dependency resolution, Joiner synthesis |
|
||||
| `pkg/itr/dag/planner.go` | `DAGPlanner`: LLM-driven DAG generation with structured output, few-shot examples, validation |
|
||||
| `pkg/itr/dag/resolver.go` | Dependency reference resolver: `#nodeN` substitution, type coercion, error propagation |
|
||||
| `pkg/itr/dag/replan.go` | Replanning loop: iterative DAG execution with progress detection and budget enforcement |
|
||||
| `pkg/itr/dag/router.go` | Query router: classifies queries as ReAct-suitable or DAG-suitable (`ModeReAct | ModeDAG | ModeAuto`) |
|
||||
| `pkg/security/securebus/bus.go` | SecureBus core: orchestrates capability check, secret injection, execution, leak scan, audit |
|
||||
| `pkg/security/securebus/policy.go` | PolicyEngine: validates tool capabilities against configured rules, depth limits, token budgets |
|
||||
| `pkg/security/securebus/transport.go` | `Transport` interface, `ChannelTransport` (in-process), `SocketTransport` (daemon), `WasmTransport` (wazero isolate) |
|
||||
| `pkg/security/securebus/audit.go` | Audit log writer (extends `agent_audit_log` table) |
|
||||
| `pkg/security/secretstore.go` | `SecretStore`: named secrets, encrypted persistence, CRUD operations |
|
||||
| `pkg/security/keyring.go` | `KeyringProvider` interface + OS implementations (build-tag gated) |
|
||||
| `pkg/security/keyring_noop.go` | No-op keyring for embedded builds |
|
||||
| `pkg/security/zkp.go` | Schnorr ZKP prover/verifier (daemon mode only, P2) |
|
||||
| `pkg/rlm/engine.go` | `RLMEngine`: recursive context decomposition through SecureBus, recursive DAG expansion |
|
||||
| `pkg/rlm/strategy.go` | Strategy planning: cheap sub-LM selects decomposition op, context analysis heuristics |
|
||||
| `pkg/rlm/rope.go` | Rope context wrapper: O(log n) slice/split/concat over `zyedidia/rope` |
|
||||
| `pkg/rlm/fanout.go` | Parallel fan-out: errgroup-managed goroutines with concurrency limits and cost tracking |
|
||||
| `docs/adr/001-isolated-tool-runtime.md` | This document |
|
||||
|
||||
**Modified files**:
|
||||
|
||||
| Path | Change |
|
||||
|------|--------|
|
||||
| `pkg/tools/base.go` | Add `CapableTool` interface, `ToolCapabilities` types |
|
||||
| `pkg/tools/shell.go` | Implement `CapableTool` |
|
||||
| `pkg/tools/filesystem.go` | Implement `CapableTool` |
|
||||
| `pkg/tools/web.go` | Implement `CapableTool` |
|
||||
| `pkg/tools/registry.go` | Add `ExtractCapabilities(Tool) ToolCapabilities` helper; add `Search(query) []ToolInfo` for ToolSearch |
|
||||
| `pkg/tools/toolloop.go` | Add `ToolLoopMode` enum, DAG executor integration, routing logic |
|
||||
| `pkg/agent/loop.go` | Route tool execution through SecureBus; integrate DAGExecutor + RLMEngine; ReAct/DAG mode selection |
|
||||
| `pkg/memory/store/memory_store.go` | Add RLM activation when search results exceed context window |
|
||||
| `cmd/picoclaw/main.go` | Wire SecureBus + DAGExecutor + RLMEngine, add `secret` and `daemon` subcommands |
|
||||
| `go.mod` | Add `tetratelabs/wazero`, `zyedidia/rope`, `google/flatbuffers` |
|
||||
| `ROADMAP.md` | Update to reference DAG executor convergence and this ADR |
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
### Security
|
||||
|
||||
- [IronClaw](https://github.com/nicholasgriffintn/ironclaw) — WASM isolation, capability-based permissions, encrypted vault with runtime injection
|
||||
- [WyrmLock](https://github.com/ZanzyTHEbar/WyrmLock) — ZKP-based application locking, proc connector interception, system keyring integration
|
||||
- [OpenClaw CVE-2026-XXXX](https://github.com/openclaw/openclaw/security/advisories) — Token exfiltration via malicious skills reading host environment
|
||||
- Schnorr Identification Protocol — C.P. Schnorr, "Efficient Signature Generation by Smart Cards", Journal of Cryptology, 1991
|
||||
- [XChaCha20-Poly1305](https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-xchacha) — Extended-nonce AEAD used by the existing `Vault`
|
||||
|
||||
### RLM and Context Management
|
||||
|
||||
- [Recursive Language Models](https://arxiv.org/abs/2512.24601) — Zhang & Khattab, MIT CSAIL, arXiv:2512.24601 v2 (Jan 2026). 114% uplift on OOLONG, perfect performance at 10M+ tokens.
|
||||
- [DSPy RLM](https://github.com/stanfordnlp/dspy) — `dspy.RLM` canonical production implementation (DSPy >= 3.1.2, experimental)
|
||||
- [dspy-go](https://github.com/XiaoConstantine/dspy-go) — Go port of DSPy including `pkg/modules/rlm` with Yaegi REPL, adaptive iteration, SubRLM
|
||||
|
||||
### DAG Execution and Programmatic Tool Calling
|
||||
|
||||
- [LLMCompiler: An LLM Compiler for Parallel Function Calling](https://arxiv.org/abs/2312.04511) — Kim et al., ICML 2024. 3.7x latency speedup, 6.7x cost savings over ReAct.
|
||||
- [LLM-Tool Compiler: Fused Parallel Function Calling](https://arxiv.org/abs/2405.17438) — Runtime operation fusion, 4x more parallel calls, 40% token reduction.
|
||||
- [Introducing Advanced Tool Use](https://www.anthropic.com/engineering/advanced-tool-use) — Anthropic, Nov 2025. Programmatic Tool Calling, Tool Search Tool, Tool Use Examples.
|
||||
- [Programmatic Tool Calling docs](https://platform.claude.com/docs/en/agents-and-tools/tool-use/programmatic-tool-calling) — Anthropic API reference for PTC.
|
||||
- [LLMCompiler reference implementation](https://github.com/SqueezeAILab/LLMCompiler) — Python/LangGraph, SqueezeAILab.
|
||||
|
||||
### Dependencies
|
||||
|
||||
- [wazero](https://github.com/tetratelabs/wazero) v1.11.0 — Pure Go WebAssembly runtime, zero CGO, WASI preview 1, AOT compiler
|
||||
- [FlatBuffers](https://github.com/google/flatbuffers) v25.12.19 — Zero-copy serialization with Go codegen
|
||||
- [zyedidia/rope](https://github.com/zyedidia/rope) — Rope data structure for O(log n) string operations on large contexts
|
||||
20
docs/adr/README.md
Normal file
20
docs/adr/README.md
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
# Architecture Decision Records
|
||||
|
||||
This directory contains Architecture Decision Records (ADRs) for PicoClaw.
|
||||
|
||||
ADRs document significant technical decisions — the context, the options considered, the chosen approach, and the trade-offs accepted. They create a historical record of how the system evolved and why.
|
||||
|
||||
## Format
|
||||
|
||||
Each ADR follows [Michael Nygard's template](https://cognitect.com/blog/2011/11/15/documenting-architecture-decisions):
|
||||
|
||||
- **Status**: Proposed | Accepted | Superseded | Deprecated
|
||||
- **Context**: What forces are at play
|
||||
- **Decision**: What we chose
|
||||
- **Consequences**: What changes as a result
|
||||
|
||||
## Index
|
||||
|
||||
| ADR | Title | Status |
|
||||
|-----|-------|--------|
|
||||
| [001](001-isolated-tool-runtime.md) | Isolated Tool Runtime (ITR) + DAG Task Executor + RLM Engine | Proposed |
|
||||
99
pkg/itr/commands.fbs
Normal file
99
pkg/itr/commands.fbs
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
// PicoClaw Isolated Tool Runtime — Command Protocol
|
||||
// Namespace: picoclaw.itr
|
||||
//
|
||||
// This schema defines the binary command protocol between the agent loop
|
||||
// and the SecureBus. All tool invocations, including RLM recursive
|
||||
// decomposition operations, are serialized using this schema.
|
||||
//
|
||||
// Generated Go code lives in commands_generated.go (committed).
|
||||
// To regenerate: flatc --go --gen-mutable -o pkg/itr pkg/itr/commands.fbs
|
||||
//
|
||||
// FlatBuffers provides:
|
||||
// - Zero-copy reads (no deserialization on the hot path)
|
||||
// - Schema-enforced vocabulary (LLM cannot construct out-of-vocab requests)
|
||||
// - Single binary format across all transports (channel, socket, WASM)
|
||||
|
||||
namespace picoclaw.itr;
|
||||
|
||||
// ── RLM Decomposition Operations ────────────────────────────────────────────
|
||||
|
||||
// Peek reads a byte range from the context rope.
|
||||
table Peek {
|
||||
start: uint64; // byte offset
|
||||
length: uint32; // bytes to read (0 = to end)
|
||||
}
|
||||
|
||||
// Grep searches the context for a pattern.
|
||||
table Grep {
|
||||
pattern: string (required);
|
||||
max_matches: uint32 = 50;
|
||||
case_insensitive: bool = false;
|
||||
}
|
||||
|
||||
// Partition splits the context into k roughly-equal chunks.
|
||||
table Partition {
|
||||
k: uint32 = 4; // number of partitions
|
||||
method: string = "uniform"; // "uniform" | "semantic"
|
||||
overlap: uint32 = 0; // token overlap between partitions
|
||||
semantic: bool = false; // use semantic boundaries (sentence/paragraph)
|
||||
}
|
||||
|
||||
// Recurse spawns a sub-RLM call over a context partition.
|
||||
table Recurse {
|
||||
sub_query: string (required); // the question to answer over the partition
|
||||
context_key: string (required); // identifies which partition to use
|
||||
depth_hint: uint8 = 3; // max recursion depth
|
||||
}
|
||||
|
||||
// ── Traditional Tool Execution ───────────────────────────────────────────────
|
||||
|
||||
// ToolExec invokes a named tool with JSON-serialized arguments.
|
||||
table ToolExec {
|
||||
tool_name: string (required);
|
||||
args_json: string (required); // JSON object
|
||||
}
|
||||
|
||||
// ExecWasm runs a WASM module (Layer 5 — optional, wazero isolate).
|
||||
table ExecWasm {
|
||||
module_key: string (required); // key into the WASM module registry
|
||||
entry: string (required); // exported function name
|
||||
input_json: string; // JSON input for the module
|
||||
}
|
||||
|
||||
// Final records the terminal answer from an RLM recursion.
|
||||
table Final {
|
||||
answer: string (required);
|
||||
var_name: string; // optional variable name for multi-part answers
|
||||
}
|
||||
|
||||
// ── Envelope ─────────────────────────────────────────────────────────────────
|
||||
|
||||
union CommandPayload {
|
||||
Peek,
|
||||
Grep,
|
||||
Partition,
|
||||
Recurse,
|
||||
ToolExec,
|
||||
ExecWasm,
|
||||
Final
|
||||
}
|
||||
|
||||
table ToolRequest {
|
||||
id: string (required); // UUIDv7 for correlation
|
||||
payload: CommandPayload;
|
||||
timestamp: uint64; // Unix nanoseconds
|
||||
depth: uint8 = 0; // recursion depth (0 = top-level)
|
||||
session_key: string; // opaque session/conversation identifier
|
||||
tool_call_id: string; // links back to LLM tool_call.id
|
||||
}
|
||||
|
||||
table ToolResponse {
|
||||
id: string (required); // mirrors request id
|
||||
result: string; // JSON-serialized result
|
||||
is_error: bool = false;
|
||||
leak_detected: bool = false; // set when Redactor found a match
|
||||
cost_tokens: uint32 = 0; // tokens consumed (RLM operations)
|
||||
redacted_keys: [string]; // keys that were redacted, for audit
|
||||
}
|
||||
|
||||
root_type ToolRequest;
|
||||
361
pkg/itr/commands.go
Normal file
361
pkg/itr/commands.go
Normal file
|
|
@ -0,0 +1,361 @@
|
|||
// Package itr defines the binary command protocol for the PicoClaw Isolated
|
||||
// Tool Runtime (ITR). All tool invocations — traditional tools and RLM
|
||||
// recursive decomposition operations — are serialized as ToolRequest /
|
||||
// ToolResponse pairs.
|
||||
//
|
||||
// The canonical schema lives in commands.fbs. This file provides the Go
|
||||
// types and encoding helpers that replace the raw flatc-generated code while
|
||||
// keeping the same wire format semantics. The encoding uses encoding/json
|
||||
// on the initial implementation path; the FlatBuffers binary encoding is
|
||||
// available via the flatbuffers package when performance demands it.
|
||||
//
|
||||
// Migration path: when the codebase moves to the full FlatBuffers binary
|
||||
// encoding, replace the JSON marshal/unmarshal calls with flatbuffers builder
|
||||
// calls without changing any call sites — the public types remain stable.
|
||||
package itr
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CommandType identifies which command variant is carried in a ToolRequest.
|
||||
type CommandType string
|
||||
|
||||
const (
|
||||
// RLM decomposition operations
|
||||
CmdPeek CommandType = "peek"
|
||||
CmdGrep CommandType = "grep"
|
||||
CmdPartition CommandType = "partition"
|
||||
CmdRecurse CommandType = "recurse"
|
||||
|
||||
// Traditional tool execution
|
||||
CmdToolExec CommandType = "tool_exec"
|
||||
CmdExecWasm CommandType = "exec_wasm"
|
||||
|
||||
// Terminal answer from an RLM recursion
|
||||
CmdFinal CommandType = "final"
|
||||
|
||||
// On-demand tool discovery (mirrors Anthropic Tool Search)
|
||||
CmdToolSearch CommandType = "tool_search"
|
||||
|
||||
// Sandboxed code execution (wazero-isolated)
|
||||
CmdCodeExec CommandType = "code_exec"
|
||||
|
||||
// DAG plan: a composite payload containing multiple nodes with dependencies
|
||||
CmdDAGPlan CommandType = "dag_plan"
|
||||
)
|
||||
|
||||
// ── Payload types ────────────────────────────────────────────────────────────
|
||||
|
||||
// Peek reads a byte range from the context rope.
|
||||
type Peek struct {
|
||||
Start uint64 `json:"start"` // byte offset into context
|
||||
Length uint32 `json:"length,omitempty"` // 0 = read to end
|
||||
}
|
||||
|
||||
// Grep searches the context for a pattern.
|
||||
type Grep struct {
|
||||
Pattern string `json:"pattern"`
|
||||
MaxMatches uint32 `json:"max_matches,omitempty"` // 0 → 50
|
||||
CaseInsensitive bool `json:"case_insensitive,omitempty"` // default false
|
||||
}
|
||||
|
||||
// Partition splits the context into k chunks.
|
||||
type Partition struct {
|
||||
K uint32 `json:"k,omitempty"` // 0 → 4
|
||||
Method string `json:"method,omitempty"` // "uniform" | "semantic"
|
||||
Overlap uint32 `json:"overlap,omitempty"` // token overlap between chunks
|
||||
Semantic bool `json:"semantic,omitempty"` // use sentence/paragraph boundaries
|
||||
}
|
||||
|
||||
// Recurse spawns a sub-RLM call over a context partition.
|
||||
type Recurse struct {
|
||||
SubQuery string `json:"sub_query"`
|
||||
ContextKey string `json:"context_key"`
|
||||
DepthHint uint8 `json:"depth_hint,omitempty"` // 0 → 3
|
||||
}
|
||||
|
||||
// ToolExec invokes a named tool with JSON-serialized arguments.
|
||||
type ToolExec struct {
|
||||
ToolName string `json:"tool_name"`
|
||||
ArgsJSON string `json:"args_json"` // JSON object
|
||||
}
|
||||
|
||||
// ExecWasm runs a function in a wazero WASM isolate (Layer 5).
|
||||
type ExecWasm struct {
|
||||
ModuleKey string `json:"module_key"`
|
||||
Entry string `json:"entry"`
|
||||
InputJSON string `json:"input_json,omitempty"`
|
||||
}
|
||||
|
||||
// Final records the terminal answer from an RLM recursion.
|
||||
type Final struct {
|
||||
Answer string `json:"answer"`
|
||||
VarName string `json:"var_name,omitempty"`
|
||||
}
|
||||
|
||||
// ToolSearch discovers tools matching a query. The executor returns a list
|
||||
// of ToolInfo entries the DAG planner can reference in subsequent ToolExec nodes.
|
||||
type ToolSearch struct {
|
||||
Query string `json:"query"`
|
||||
MaxResults uint8 `json:"max_results,omitempty"` // 0 → 10
|
||||
}
|
||||
|
||||
// CodeExec runs code in a wazero-isolated sandbox (opt-in only, Layer 5).
|
||||
type CodeExec struct {
|
||||
Code string `json:"code"`
|
||||
Language string `json:"language"` // e.g. "javascript", "python-wasm"
|
||||
}
|
||||
|
||||
// DAGNode is a single unit of work within a DAGPlan. Each node carries a
|
||||
// command payload and declares dependencies on other nodes by ID.
|
||||
// Args may contain "#nodeN" references that are resolved to the output of
|
||||
// the dependency node at execution time.
|
||||
type DAGNode struct {
|
||||
ID string `json:"id"`
|
||||
Type CommandType `json:"type"`
|
||||
Payload interface{} `json:"payload"`
|
||||
DependsOn []string `json:"depends_on,omitempty"`
|
||||
}
|
||||
|
||||
// DAGPlan is a composite command that contains multiple DAGNodes forming a
|
||||
// Directed Acyclic Graph. The executor dispatches nodes in topological order,
|
||||
// running independent nodes in parallel.
|
||||
type DAGPlan struct {
|
||||
Nodes []DAGNode `json:"nodes"`
|
||||
MaxParallel uint8 `json:"max_parallel,omitempty"` // 0 → GOMAXPROCS
|
||||
TokenBudget uint32 `json:"token_budget,omitempty"`
|
||||
JoinerQuery string `json:"joiner_query,omitempty"` // synthesis prompt
|
||||
}
|
||||
|
||||
// ── Envelope ─────────────────────────────────────────────────────────────────
|
||||
|
||||
// ToolRequest is the envelope transmitted from the agent loop to the SecureBus.
|
||||
// Exactly one of the command payload fields should be non-nil.
|
||||
type ToolRequest struct {
|
||||
ID string `json:"id"` // UUIDv7
|
||||
Type CommandType `json:"type"` // discriminator
|
||||
Payload interface{} `json:"payload"` // one of the Cmd* types
|
||||
Timestamp int64 `json:"timestamp"` // Unix nanoseconds
|
||||
Depth uint8 `json:"depth,omitempty"` // recursion depth
|
||||
SessionKey string `json:"session_key,omitempty"` // conversation scope
|
||||
ToolCallID string `json:"tool_call_id,omitempty"` // LLM tool_call.id
|
||||
}
|
||||
|
||||
// ToolResponse is returned from the SecureBus to the agent loop.
|
||||
type ToolResponse struct {
|
||||
ID string `json:"id"` // mirrors request ID
|
||||
Result string `json:"result,omitempty"` // JSON-serialised result
|
||||
IsError bool `json:"is_error,omitempty"`
|
||||
LeakDetected bool `json:"leak_detected,omitempty"` // Redactor triggered
|
||||
CostTokens uint32 `json:"cost_tokens,omitempty"` // tokens consumed
|
||||
RedactedKeys []string `json:"redacted_keys,omitempty"` // keys that were redacted
|
||||
}
|
||||
|
||||
// ── Constructors ─────────────────────────────────────────────────────────────
|
||||
|
||||
func NewToolExecRequest(id, sessionKey, toolCallID, toolName, argsJSON string) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdToolExec,
|
||||
Payload: ToolExec{ToolName: toolName, ArgsJSON: argsJSON},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
SessionKey: sessionKey,
|
||||
ToolCallID: toolCallID,
|
||||
}
|
||||
}
|
||||
|
||||
func NewPeekRequest(id, sessionKey string, depth uint8, start uint64, length uint32) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdPeek,
|
||||
Payload: Peek{Start: start, Length: length},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
Depth: depth,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewGrepRequest(id, sessionKey string, depth uint8, pattern string, maxMatches uint32, caseInsensitive bool) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdGrep,
|
||||
Payload: Grep{Pattern: pattern, MaxMatches: maxMatches, CaseInsensitive: caseInsensitive},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
Depth: depth,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewPartitionRequest(id, sessionKey string, depth uint8, k uint32, method string, overlap uint32, semantic bool) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdPartition,
|
||||
Payload: Partition{K: k, Method: method, Overlap: overlap, Semantic: semantic},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
Depth: depth,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewRecurseRequest(id, sessionKey string, depth uint8, subQuery, contextKey string, depthHint uint8) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdRecurse,
|
||||
Payload: Recurse{SubQuery: subQuery, ContextKey: contextKey, DepthHint: depthHint},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
Depth: depth,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewFinalRequest(id, sessionKey string, depth uint8, answer, varName string) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdFinal,
|
||||
Payload: Final{Answer: answer, VarName: varName},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
Depth: depth,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewToolSearchRequest(id, sessionKey string, query string, maxResults uint8) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdToolSearch,
|
||||
Payload: ToolSearch{Query: query, MaxResults: maxResults},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewCodeExecRequest(id, sessionKey string, code, language string) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdCodeExec,
|
||||
Payload: CodeExec{Code: code, Language: language},
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewDAGPlanRequest(id, sessionKey string, plan DAGPlan) ToolRequest {
|
||||
return ToolRequest{
|
||||
ID: id,
|
||||
Type: CmdDAGPlan,
|
||||
Payload: plan,
|
||||
Timestamp: time.Now().UnixNano(),
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
|
||||
func NewSuccessResponse(id, result string, costTokens uint32) ToolResponse {
|
||||
return ToolResponse{ID: id, Result: result, CostTokens: costTokens}
|
||||
}
|
||||
|
||||
func NewErrorResponse(id, errMsg string) ToolResponse {
|
||||
return ToolResponse{ID: id, Result: errMsg, IsError: true}
|
||||
}
|
||||
|
||||
func NewLeakResponse(id, redactedResult string, redactedKeys []string) ToolResponse {
|
||||
return ToolResponse{
|
||||
ID: id,
|
||||
Result: redactedResult,
|
||||
LeakDetected: true,
|
||||
RedactedKeys: redactedKeys,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Serialisation helpers ────────────────────────────────────────────────────
|
||||
|
||||
// Marshal encodes a ToolRequest to JSON bytes for transmission.
|
||||
func (r ToolRequest) Marshal() ([]byte, error) {
|
||||
return json.Marshal(r)
|
||||
}
|
||||
|
||||
// Marshal encodes a ToolResponse to JSON bytes for transmission.
|
||||
func (r ToolResponse) Marshal() ([]byte, error) {
|
||||
return json.Marshal(r)
|
||||
}
|
||||
|
||||
// UnmarshalRequest decodes a ToolRequest from JSON bytes.
|
||||
func UnmarshalRequest(data []byte) (ToolRequest, error) {
|
||||
var raw struct {
|
||||
ID string `json:"id"`
|
||||
Type CommandType `json:"type"`
|
||||
Payload json.RawMessage `json:"payload"`
|
||||
Timestamp int64 `json:"timestamp"`
|
||||
Depth uint8 `json:"depth"`
|
||||
SessionKey string `json:"session_key"`
|
||||
ToolCallID string `json:"tool_call_id"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &raw); err != nil {
|
||||
return ToolRequest{}, err
|
||||
}
|
||||
|
||||
req := ToolRequest{
|
||||
ID: raw.ID,
|
||||
Type: raw.Type,
|
||||
Timestamp: raw.Timestamp,
|
||||
Depth: raw.Depth,
|
||||
SessionKey: raw.SessionKey,
|
||||
ToolCallID: raw.ToolCallID,
|
||||
}
|
||||
|
||||
var err error
|
||||
switch raw.Type {
|
||||
case CmdPeek:
|
||||
var p Peek
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdGrep:
|
||||
var p Grep
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdPartition:
|
||||
var p Partition
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdRecurse:
|
||||
var p Recurse
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdToolExec:
|
||||
var p ToolExec
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdExecWasm:
|
||||
var p ExecWasm
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdFinal:
|
||||
var p Final
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdToolSearch:
|
||||
var p ToolSearch
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdCodeExec:
|
||||
var p CodeExec
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
case CmdDAGPlan:
|
||||
var p DAGPlan
|
||||
err = json.Unmarshal(raw.Payload, &p)
|
||||
req.Payload = p
|
||||
default:
|
||||
return ToolRequest{}, fmt.Errorf("unknown command type: %q", raw.Type)
|
||||
}
|
||||
|
||||
return req, err
|
||||
}
|
||||
|
||||
// UnmarshalResponse decodes a ToolResponse from JSON bytes.
|
||||
func UnmarshalResponse(data []byte) (ToolResponse, error) {
|
||||
var r ToolResponse
|
||||
return r, json.Unmarshal(data, &r)
|
||||
}
|
||||
253
pkg/itr/dag/executor.go
Normal file
253
pkg/itr/dag/executor.go
Normal file
|
|
@ -0,0 +1,253 @@
|
|||
// Package dag implements the DAG Task Executor for PicoClaw. It receives a
|
||||
// DAGPlan (a set of nodes with dependency edges), dispatches them in
|
||||
// topological wave order via the SecureBus, and synthesises a final result
|
||||
// using the Joiner pattern.
|
||||
//
|
||||
// Design principles (from ADR-001):
|
||||
// - LLMCompiler: single-pass DAG planning, parallel topological execution
|
||||
// - Anthropic PTC: intermediate results never enter the LLM's context;
|
||||
// only the Joiner's output crosses back to the agent loop
|
||||
// - RLM integration: when a node's context exceeds a threshold, the
|
||||
// RLMEngine expands it into a sub-DAG transparently
|
||||
package dag
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
"github.com/sipeed/picoclaw/pkg/security/securebus"
|
||||
)
|
||||
|
||||
// JoinerFunc synthesises a final answer from all node results.
|
||||
// systemPrompt provides instruction context; resultSummary contains the
|
||||
// concatenated node outputs for synthesis.
|
||||
type JoinerFunc func(ctx context.Context, systemPrompt, resultSummary string) (string, uint32, error)
|
||||
|
||||
// Executor runs a DAGPlan through the SecureBus with topological dispatch.
|
||||
type Executor struct {
|
||||
bus *securebus.Bus
|
||||
joiner JoinerFunc
|
||||
maxParallel int
|
||||
}
|
||||
|
||||
// NewExecutor creates a DAG executor.
|
||||
// joiner is called after all nodes complete to synthesise the final answer.
|
||||
func NewExecutor(bus *securebus.Bus, joiner JoinerFunc) *Executor {
|
||||
return &Executor{
|
||||
bus: bus,
|
||||
joiner: joiner,
|
||||
maxParallel: runtime.GOMAXPROCS(0),
|
||||
}
|
||||
}
|
||||
|
||||
// ExecuteResult holds the output of a DAG execution.
|
||||
type ExecuteResult struct {
|
||||
FinalAnswer string
|
||||
NodeResults map[string]string
|
||||
TotalTokens uint32
|
||||
}
|
||||
|
||||
// Execute runs the full DAGPlan: topological dispatch, parallel execution,
|
||||
// dependency resolution, and Joiner synthesis.
|
||||
func (e *Executor) Execute(ctx context.Context, sessionKey string, plan *itr.DAGPlan) (*ExecuteResult, error) {
|
||||
if len(plan.Nodes) == 0 {
|
||||
return &ExecuteResult{}, nil
|
||||
}
|
||||
|
||||
states := make(map[string]*nodeState, len(plan.Nodes))
|
||||
for i := range plan.Nodes {
|
||||
n := &plan.Nodes[i]
|
||||
states[n.ID] = newNodeState(n.ID, n.DependsOn)
|
||||
}
|
||||
|
||||
waves, err := topologicalOrder(states)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
maxPar := int(plan.MaxParallel)
|
||||
if maxPar <= 0 {
|
||||
maxPar = e.maxParallel
|
||||
}
|
||||
|
||||
var totalTokens uint32
|
||||
var tokensMu sync.Mutex
|
||||
|
||||
for _, wave := range waves {
|
||||
sem := make(chan struct{}, maxPar)
|
||||
var wg sync.WaitGroup
|
||||
var waveErr error
|
||||
var errOnce sync.Once
|
||||
|
||||
for _, nodeID := range wave {
|
||||
nodeID := nodeID
|
||||
ns := states[nodeID]
|
||||
|
||||
wg.Add(1)
|
||||
sem <- struct{}{}
|
||||
go func() {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
|
||||
// Wait for all dependencies (should already be done since
|
||||
// we execute wave-by-wave, but handles cross-wave edges).
|
||||
for _, dep := range ns.DependsOn {
|
||||
depState, ok := states[dep]
|
||||
if !ok {
|
||||
ns.setResult("", fmt.Errorf("unknown dependency: %s", dep))
|
||||
errOnce.Do(func() { waveErr = fmt.Errorf("unknown dependency: %s", dep) })
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-depState.done:
|
||||
if _, depErr := depState.getResult(); depErr != nil {
|
||||
ns.setResult("", fmt.Errorf("dependency %s failed: %w", dep, depErr))
|
||||
errOnce.Do(func() { waveErr = depErr })
|
||||
return
|
||||
}
|
||||
case <-ctx.Done():
|
||||
ns.setResult("", ctx.Err())
|
||||
errOnce.Do(func() { waveErr = ctx.Err() })
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Find the original node to get its payload
|
||||
var node *itr.DAGNode
|
||||
for i := range plan.Nodes {
|
||||
if plan.Nodes[i].ID == nodeID {
|
||||
node = &plan.Nodes[i]
|
||||
break
|
||||
}
|
||||
}
|
||||
if node == nil {
|
||||
ns.setResult("", fmt.Errorf("node %s not found in plan", nodeID))
|
||||
return
|
||||
}
|
||||
|
||||
resp := e.executeNode(ctx, sessionKey, node, states)
|
||||
|
||||
tokensMu.Lock()
|
||||
totalTokens += resp.CostTokens
|
||||
tokensMu.Unlock()
|
||||
|
||||
if resp.IsError {
|
||||
ns.setResult(resp.Result, fmt.Errorf("node %s: %s", nodeID, resp.Result))
|
||||
} else {
|
||||
ns.setResult(resp.Result, nil)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
|
||||
if waveErr != nil {
|
||||
return nil, fmt.Errorf("dag execution failed: %w", waveErr)
|
||||
}
|
||||
}
|
||||
|
||||
// Collect all node results.
|
||||
nodeResults := make(map[string]string, len(states))
|
||||
for id, ns := range states {
|
||||
result, _ := ns.getResult()
|
||||
nodeResults[id] = result
|
||||
}
|
||||
|
||||
// Joiner synthesis: combine all node results into a final answer.
|
||||
finalAnswer, joinerTokens, err := e.join(ctx, plan, nodeResults)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("joiner synthesis failed: %w", err)
|
||||
}
|
||||
totalTokens += joinerTokens
|
||||
|
||||
return &ExecuteResult{
|
||||
FinalAnswer: finalAnswer,
|
||||
NodeResults: nodeResults,
|
||||
TotalTokens: totalTokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// executeNode dispatches a single node through the SecureBus after resolving
|
||||
// dependency references in its payload.
|
||||
func (e *Executor) executeNode(ctx context.Context, sessionKey string, node *itr.DAGNode, states map[string]*nodeState) itr.ToolResponse {
|
||||
req := nodeToRequest(sessionKey, node, states)
|
||||
return e.bus.Execute(ctx, req)
|
||||
}
|
||||
|
||||
// nodeToRequest converts a DAGNode into a ToolRequest, resolving #nodeN
|
||||
// references in tool arguments.
|
||||
func nodeToRequest(sessionKey string, node *itr.DAGNode, states map[string]*nodeState) itr.ToolRequest {
|
||||
switch node.Type {
|
||||
case itr.CmdToolExec:
|
||||
te, ok := node.Payload.(itr.ToolExec)
|
||||
if !ok {
|
||||
if m, ok := node.Payload.(map[string]interface{}); ok {
|
||||
b, _ := json.Marshal(m)
|
||||
_ = json.Unmarshal(b, &te)
|
||||
}
|
||||
}
|
||||
te.ArgsJSON = resolveToolExecArgs(te.ArgsJSON, states)
|
||||
return itr.NewToolExecRequest(node.ID, sessionKey, node.ID, te.ToolName, te.ArgsJSON)
|
||||
|
||||
case itr.CmdToolSearch:
|
||||
ts, ok := node.Payload.(itr.ToolSearch)
|
||||
if !ok {
|
||||
if m, ok := node.Payload.(map[string]interface{}); ok {
|
||||
b, _ := json.Marshal(m)
|
||||
_ = json.Unmarshal(b, &ts)
|
||||
}
|
||||
}
|
||||
return itr.NewToolSearchRequest(node.ID, sessionKey, ts.Query, ts.MaxResults)
|
||||
|
||||
default:
|
||||
return itr.ToolRequest{
|
||||
ID: node.ID,
|
||||
Type: node.Type,
|
||||
Payload: node.Payload,
|
||||
SessionKey: sessionKey,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// join calls the JoinerFunc to synthesise a final answer from node results.
|
||||
// If no JoinerFunc is configured, concatenates results.
|
||||
func (e *Executor) join(ctx context.Context, plan *itr.DAGPlan, nodeResults map[string]string) (string, uint32, error) {
|
||||
if e.joiner == nil || plan.JoinerQuery == "" {
|
||||
return concatenateResults(nodeResults), 0, nil
|
||||
}
|
||||
|
||||
var sb strings.Builder
|
||||
sb.WriteString("Node results:\n")
|
||||
for id, result := range nodeResults {
|
||||
sb.WriteString(fmt.Sprintf("- %s: %s\n", id, truncate(result, 2000)))
|
||||
}
|
||||
|
||||
systemPrompt := "You are synthesizing results from parallel tool executions into a coherent final answer. Be concise and precise."
|
||||
userQuery := plan.JoinerQuery + "\n\n" + sb.String()
|
||||
|
||||
return e.joiner(ctx, systemPrompt, userQuery)
|
||||
}
|
||||
|
||||
func concatenateResults(results map[string]string) string {
|
||||
var parts []string
|
||||
for _, v := range results {
|
||||
if v != "" {
|
||||
parts = append(parts, v)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
func truncate(s string, max int) string {
|
||||
if len(s) <= max {
|
||||
return s
|
||||
}
|
||||
return s[:max] + "..."
|
||||
}
|
||||
121
pkg/itr/dag/resolver.go
Normal file
121
pkg/itr/dag/resolver.go
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
package dag
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
var nodeRefPattern = regexp.MustCompile(`#node([A-Za-z0-9_-]+)`)
|
||||
|
||||
// nodeState tracks execution state for a single DAG node.
|
||||
type nodeState struct {
|
||||
ID string
|
||||
DependsOn []string
|
||||
result string
|
||||
err error
|
||||
done chan struct{}
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func newNodeState(id string, deps []string) *nodeState {
|
||||
return &nodeState{
|
||||
ID: id,
|
||||
DependsOn: deps,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
func (ns *nodeState) setResult(result string, err error) {
|
||||
ns.mu.Lock()
|
||||
ns.result = result
|
||||
ns.err = err
|
||||
ns.mu.Unlock()
|
||||
close(ns.done)
|
||||
}
|
||||
|
||||
func (ns *nodeState) getResult() (string, error) {
|
||||
ns.mu.Lock()
|
||||
defer ns.mu.Unlock()
|
||||
return ns.result, ns.err
|
||||
}
|
||||
|
||||
// resolveRefs replaces all "#nodeXYZ" references in argsJSON with the
|
||||
// actual output from the referenced node. Returns the resolved JSON string.
|
||||
func resolveRefs(argsJSON string, states map[string]*nodeState) string {
|
||||
return nodeRefPattern.ReplaceAllStringFunc(argsJSON, func(match string) string {
|
||||
nodeID := match[5:] // strip "#node"
|
||||
ns, ok := states[nodeID]
|
||||
if !ok {
|
||||
return match
|
||||
}
|
||||
result, _ := ns.getResult()
|
||||
return escapeForJSON(result)
|
||||
})
|
||||
}
|
||||
|
||||
// resolveToolExecArgs resolves #nodeN references within a ToolExec payload's
|
||||
// ArgsJSON field, returning the resolved args as a JSON string.
|
||||
func resolveToolExecArgs(argsJSON string, states map[string]*nodeState) string {
|
||||
if !strings.Contains(argsJSON, "#node") {
|
||||
return argsJSON
|
||||
}
|
||||
return resolveRefs(argsJSON, states)
|
||||
}
|
||||
|
||||
// escapeForJSON makes a string safe for embedding into a JSON value.
|
||||
func escapeForJSON(s string) string {
|
||||
b, err := json.Marshal(s)
|
||||
if err != nil {
|
||||
return s
|
||||
}
|
||||
// Strip surrounding quotes since we're replacing within an existing string.
|
||||
return string(b[1 : len(b)-1])
|
||||
}
|
||||
|
||||
// topologicalOrder returns node IDs in topological execution order.
|
||||
// Returns an error if the graph contains a cycle.
|
||||
func topologicalOrder(nodes map[string]*nodeState) ([][]string, error) {
|
||||
inDegree := make(map[string]int, len(nodes))
|
||||
dependents := make(map[string][]string, len(nodes))
|
||||
|
||||
for id, ns := range nodes {
|
||||
inDegree[id] = len(ns.DependsOn)
|
||||
for _, dep := range ns.DependsOn {
|
||||
dependents[dep] = append(dependents[dep], id)
|
||||
}
|
||||
}
|
||||
|
||||
var waves [][]string
|
||||
for {
|
||||
var wave []string
|
||||
for id, deg := range inDegree {
|
||||
if deg == 0 {
|
||||
wave = append(wave, id)
|
||||
}
|
||||
}
|
||||
if len(wave) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
waves = append(waves, wave)
|
||||
for _, id := range wave {
|
||||
delete(inDegree, id)
|
||||
for _, dep := range dependents[id] {
|
||||
inDegree[dep]--
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(inDegree) > 0 {
|
||||
var remaining []string
|
||||
for id := range inDegree {
|
||||
remaining = append(remaining, id)
|
||||
}
|
||||
return nil, fmt.Errorf("cycle detected among nodes: %v", remaining)
|
||||
}
|
||||
|
||||
return waves, nil
|
||||
}
|
||||
155
pkg/rlm/engine.go
Normal file
155
pkg/rlm/engine.go
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
package rlm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
"github.com/sipeed/picoclaw/pkg/security/securebus"
|
||||
)
|
||||
|
||||
// CallModelFunc is the signature for making a direct LLM call.
|
||||
// The RLMEngine uses this for the leaf-level answer when the context fits
|
||||
// within the direct threshold, and for merge synthesis.
|
||||
type CallModelFunc func(ctx context.Context, systemPrompt, userQuery string) (string, uint32, error)
|
||||
|
||||
// EngineConfig configures the RLMEngine.
|
||||
type EngineConfig struct {
|
||||
Strategy StrategyConfig
|
||||
MaxConcurrency int // goroutines per fan-out level; 0 = unbounded
|
||||
}
|
||||
|
||||
// DefaultEngineConfig returns sensible defaults.
|
||||
func DefaultEngineConfig() EngineConfig {
|
||||
return EngineConfig{
|
||||
Strategy: DefaultStrategyConfig(),
|
||||
MaxConcurrency: 4,
|
||||
}
|
||||
}
|
||||
|
||||
// Engine is the RLM recursive decomposition engine. It processes arbitrarily
|
||||
// long contexts by recursively partitioning them and answering sub-queries,
|
||||
// all mediated through the SecureBus for capability enforcement and leak scanning.
|
||||
//
|
||||
// Architecture (from ADR-001):
|
||||
//
|
||||
// Query + Context
|
||||
// │
|
||||
// ▼
|
||||
// StrategyPlanner.PlanNext()
|
||||
// │
|
||||
// ├── OpFinal → callModel(context, query) → answer
|
||||
// ├── OpGrep → grep context, recurse over matches
|
||||
// └── OpPartition → split into k partitions
|
||||
// │
|
||||
// └── FanOut → k × Engine.Recurse (depth+1)
|
||||
// │
|
||||
// └── MergeResults → final answer
|
||||
type Engine struct {
|
||||
cfg EngineConfig
|
||||
planner *StrategyPlanner
|
||||
bus *securebus.Bus
|
||||
callModel CallModelFunc
|
||||
}
|
||||
|
||||
// NewEngine creates an RLMEngine. bus may be nil in tests (only FanOut logic
|
||||
// is exercised). callModel is required.
|
||||
func NewEngine(cfg EngineConfig, bus *securebus.Bus, callModel CallModelFunc) *Engine {
|
||||
return &Engine{
|
||||
cfg: cfg,
|
||||
planner: NewStrategyPlanner(cfg.Strategy),
|
||||
bus: bus,
|
||||
callModel: callModel,
|
||||
}
|
||||
}
|
||||
|
||||
// Answer is the entry point for an RLM query. It recursively decomposes the
|
||||
// context until each partition fits within the direct threshold, then synthesises
|
||||
// the sub-answers bottom-up.
|
||||
//
|
||||
// sessionKey is passed through to the SecureBus for audit tracing.
|
||||
func (e *Engine) Answer(ctx context.Context, sessionKey, query, context_ string) (string, uint32, error) {
|
||||
return e.recurse(ctx, sessionKey, query, context_, 0)
|
||||
}
|
||||
|
||||
// recurse is the recursive heart of the engine.
|
||||
func (e *Engine) recurse(ctx context.Context, sessionKey, query, ctxContent string, depth uint8) (string, uint32, error) {
|
||||
rope := NewRope(ctxContent)
|
||||
op := e.planner.PlanNext(rope.Len(), query, depth)
|
||||
|
||||
switch op.Type {
|
||||
case OpFinal:
|
||||
return e.callModel(ctx, ctxContent, query)
|
||||
|
||||
case OpGrep:
|
||||
matches := rope.GrepLines(op.GrepQuery, 50, true)
|
||||
if len(matches) == 0 {
|
||||
// No matches — fall back to direct answer.
|
||||
return e.callModel(ctx, ctxContent, query)
|
||||
}
|
||||
var sb strings.Builder
|
||||
for _, m := range matches {
|
||||
sb.WriteString(m.Line)
|
||||
sb.WriteByte('\n')
|
||||
}
|
||||
return e.recurse(ctx, sessionKey, query, sb.String(), depth+1)
|
||||
|
||||
case OpPartition:
|
||||
k := op.PartitionK
|
||||
if k <= 0 {
|
||||
k = e.cfg.Strategy.DefaultPartitionK
|
||||
}
|
||||
partitions := rope.Partition(k)
|
||||
|
||||
// Fan out: process each partition concurrently.
|
||||
results := FanOut(ctx, partitions, e.cfg.MaxConcurrency,
|
||||
func(ctx context.Context, idx int, contextKey, partition string) PartitionResult {
|
||||
if e.bus != nil {
|
||||
// Route through SecureBus for audit logging.
|
||||
req := itr.NewRecurseRequest(
|
||||
fmt.Sprintf("%s-d%d-p%d", sessionKey, depth, idx),
|
||||
sessionKey,
|
||||
depth+1,
|
||||
query,
|
||||
contextKey,
|
||||
depth+1,
|
||||
)
|
||||
resp := e.bus.Execute(ctx, req)
|
||||
if resp.IsError {
|
||||
return PartitionResult{
|
||||
PartitionIdx: idx,
|
||||
ContextKey: contextKey,
|
||||
Err: fmt.Errorf("securebus: %s", resp.Result),
|
||||
}
|
||||
}
|
||||
// The bus returned a stub response for RLM commands —
|
||||
// recurse directly to get the actual answer.
|
||||
}
|
||||
answer, tokens, err := e.recurse(ctx, sessionKey, query, partition, depth+1)
|
||||
return PartitionResult{
|
||||
PartitionIdx: idx,
|
||||
ContextKey: contextKey,
|
||||
Answer: answer,
|
||||
Tokens: tokens,
|
||||
Err: err,
|
||||
}
|
||||
})
|
||||
|
||||
merged := MergeResults(results)
|
||||
totalTok := TotalTokens(results)
|
||||
|
||||
// If we got something useful, do a final synthesis pass to produce
|
||||
// a coherent answer from the merged sub-answers.
|
||||
if merged != "" && int64(depth) < int64(e.cfg.Strategy.MaxDepth)-1 {
|
||||
synth, syntTok, err := e.callModel(ctx,
|
||||
"You are synthesising answers from multiple text partitions. Be concise.",
|
||||
"Original question: "+query+"\n\nPartition answers:\n"+merged)
|
||||
return synth, totalTok + syntTok, err
|
||||
}
|
||||
return merged, totalTok, nil
|
||||
|
||||
default:
|
||||
return "", 0, fmt.Errorf("rlm: unknown op %q", op.Type)
|
||||
}
|
||||
}
|
||||
184
pkg/rlm/engine_test.go
Normal file
184
pkg/rlm/engine_test.go
Normal file
|
|
@ -0,0 +1,184 @@
|
|||
package rlm_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/rlm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// echoModel is a deterministic callModel that returns the first 50 chars
|
||||
// of the context plus the query, for tracing purposes.
|
||||
func echoModel(_ context.Context, ctxContent, query string) (string, uint32, error) {
|
||||
preview := ctxContent
|
||||
if len(preview) > 50 {
|
||||
preview = preview[:50]
|
||||
}
|
||||
return fmt.Sprintf("Q=%s CTX=%s", query, preview), 10, nil
|
||||
}
|
||||
|
||||
// errorModel always returns an error.
|
||||
func errorModel(_ context.Context, _, _ string) (string, uint32, error) {
|
||||
return "", 0, fmt.Errorf("model error")
|
||||
}
|
||||
|
||||
func TestEngine_SmallContext_AnswerDirectly(t *testing.T) {
|
||||
cfg := rlm.DefaultEngineConfig()
|
||||
cfg.Strategy.DirectThreshold = 10000 // larger than our test context
|
||||
engine := rlm.NewEngine(cfg, nil, echoModel)
|
||||
|
||||
answer, tokens, err := engine.Answer(context.Background(), "sess", "what?", "short context")
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, answer)
|
||||
assert.Greater(t, tokens, uint32(0))
|
||||
}
|
||||
|
||||
func TestEngine_LargeContext_Partitions(t *testing.T) {
|
||||
// Force partitioning by setting DirectThreshold very low.
|
||||
cfg := rlm.DefaultEngineConfig()
|
||||
cfg.Strategy.DirectThreshold = 10
|
||||
cfg.Strategy.DefaultPartitionK = 2
|
||||
cfg.Strategy.MaxDepth = 2
|
||||
cfg.MaxConcurrency = 2
|
||||
|
||||
largeCtx := strings.Repeat("hello world ", 100) // ~1200 bytes
|
||||
engine := rlm.NewEngine(cfg, nil, echoModel)
|
||||
|
||||
answer, tokens, err := engine.Answer(context.Background(), "sess", "summarise", largeCtx)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, answer)
|
||||
assert.Greater(t, tokens, uint32(0))
|
||||
}
|
||||
|
||||
func TestEngine_MaxDepthTerminates(t *testing.T) {
|
||||
cfg := rlm.DefaultEngineConfig()
|
||||
cfg.Strategy.DirectThreshold = 0 // always partition
|
||||
cfg.Strategy.MaxDepth = 3
|
||||
cfg.MaxConcurrency = 1
|
||||
|
||||
content := strings.Repeat("x", 500)
|
||||
engine := rlm.NewEngine(cfg, nil, echoModel)
|
||||
|
||||
// Should terminate without stack overflow.
|
||||
_, _, err := engine.Answer(context.Background(), "sess", "any", content)
|
||||
assert.NoError(t, err)
|
||||
}
|
||||
|
||||
func TestEngine_ModelError_Propagates(t *testing.T) {
|
||||
cfg := rlm.DefaultEngineConfig()
|
||||
cfg.Strategy.DirectThreshold = 10000
|
||||
engine := rlm.NewEngine(cfg, nil, errorModel)
|
||||
|
||||
_, _, err := engine.Answer(context.Background(), "sess", "q", "ctx")
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestEngine_GrepQuery_NarrowsContext(t *testing.T) {
|
||||
cfg := rlm.DefaultEngineConfig()
|
||||
cfg.Strategy.DirectThreshold = 10
|
||||
cfg.Strategy.DefaultPartitionK = 2
|
||||
cfg.Strategy.MaxDepth = 3
|
||||
cfg.MaxConcurrency = 1
|
||||
|
||||
// Content has one line matching a grep pattern.
|
||||
content := "foo\nfunc myFunction() {}\nbar\nbaz"
|
||||
engine := rlm.NewEngine(cfg, nil, echoModel)
|
||||
|
||||
answer, _, err := engine.Answer(context.Background(), "sess", "find func definition", content)
|
||||
require.NoError(t, err)
|
||||
assert.NotEmpty(t, answer)
|
||||
}
|
||||
|
||||
func TestFanOut_AllPartitions_Processed(t *testing.T) {
|
||||
partitions := []string{"A", "B", "C", "D"}
|
||||
results := rlm.FanOut(context.Background(), partitions, 2,
|
||||
func(_ context.Context, idx int, key, part string) rlm.PartitionResult {
|
||||
return rlm.PartitionResult{PartitionIdx: idx, ContextKey: key, Answer: "ans-" + part, Tokens: 1}
|
||||
})
|
||||
|
||||
assert.Len(t, results, 4)
|
||||
for i, r := range results {
|
||||
assert.Equal(t, i, r.PartitionIdx)
|
||||
assert.Equal(t, "ans-"+partitions[i], r.Answer)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFanOut_UnboundedConcurrency(t *testing.T) {
|
||||
parts := make([]string, 20)
|
||||
for i := range parts {
|
||||
parts[i] = fmt.Sprintf("part-%d", i)
|
||||
}
|
||||
results := rlm.FanOut(context.Background(), parts, 0,
|
||||
func(_ context.Context, idx int, key, part string) rlm.PartitionResult {
|
||||
return rlm.PartitionResult{PartitionIdx: idx, Answer: part, Tokens: 2}
|
||||
})
|
||||
assert.Len(t, results, 20)
|
||||
assert.Equal(t, uint32(40), rlm.TotalTokens(results))
|
||||
}
|
||||
|
||||
func TestFanOut_Empty(t *testing.T) {
|
||||
results := rlm.FanOut(context.Background(), nil, 4,
|
||||
func(_ context.Context, _ int, _, _ string) rlm.PartitionResult {
|
||||
return rlm.PartitionResult{}
|
||||
})
|
||||
assert.Nil(t, results)
|
||||
}
|
||||
|
||||
func TestMergeResults_Deduplication(t *testing.T) {
|
||||
results := []rlm.PartitionResult{
|
||||
{Answer: "alpha"},
|
||||
{Answer: "beta"},
|
||||
{Answer: "alpha"}, // duplicate
|
||||
{Answer: ""}, // empty — should be skipped
|
||||
}
|
||||
merged := rlm.MergeResults(results)
|
||||
assert.Equal(t, "alpha\nbeta", merged)
|
||||
}
|
||||
|
||||
func TestMergeResults_WithErrors(t *testing.T) {
|
||||
results := []rlm.PartitionResult{
|
||||
{Answer: "good"},
|
||||
{Err: fmt.Errorf("failed"), Answer: "should be skipped"},
|
||||
}
|
||||
merged := rlm.MergeResults(results)
|
||||
assert.Equal(t, "good", merged)
|
||||
}
|
||||
|
||||
func TestStrategyPlanner_SmallContext_OpFinal(t *testing.T) {
|
||||
cfg := rlm.DefaultStrategyConfig()
|
||||
cfg.DirectThreshold = 1000
|
||||
planner := rlm.NewStrategyPlanner(cfg)
|
||||
|
||||
op := planner.PlanNext(500, "any query", 0)
|
||||
assert.Equal(t, rlm.OpFinal, op.Type)
|
||||
}
|
||||
|
||||
func TestStrategyPlanner_MaxDepth_OpFinal(t *testing.T) {
|
||||
cfg := rlm.DefaultStrategyConfig()
|
||||
cfg.MaxDepth = 3
|
||||
planner := rlm.NewStrategyPlanner(cfg)
|
||||
|
||||
op := planner.PlanNext(100000, "any query", 3) // depth == MaxDepth
|
||||
assert.Equal(t, rlm.OpFinal, op.Type)
|
||||
}
|
||||
|
||||
func TestStrategyPlanner_KeywordQuery_OpGrep(t *testing.T) {
|
||||
planner := rlm.NewStrategyPlanner(rlm.DefaultStrategyConfig())
|
||||
|
||||
// "find" prefix should trigger grep.
|
||||
op := planner.PlanNext(100000, "find myFunction in code", 0)
|
||||
assert.Equal(t, rlm.OpGrep, op.Type)
|
||||
assert.NotEmpty(t, op.GrepQuery)
|
||||
}
|
||||
|
||||
func TestStrategyPlanner_LargeContext_OpPartition(t *testing.T) {
|
||||
planner := rlm.NewStrategyPlanner(rlm.DefaultStrategyConfig())
|
||||
|
||||
op := planner.PlanNext(100000, "summarise everything", 0)
|
||||
assert.Equal(t, rlm.OpPartition, op.Type)
|
||||
assert.Greater(t, op.PartitionK, 0)
|
||||
}
|
||||
103
pkg/rlm/fanout.go
Normal file
103
pkg/rlm/fanout.go
Normal file
|
|
@ -0,0 +1,103 @@
|
|||
package rlm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// PartitionResult holds the answer produced for one context partition.
|
||||
type PartitionResult struct {
|
||||
PartitionIdx int
|
||||
ContextKey string
|
||||
Answer string
|
||||
Err error
|
||||
Tokens uint32
|
||||
}
|
||||
|
||||
// FanOut executes fn over each partition concurrently and collects results.
|
||||
// maxConcurrency limits how many goroutines run simultaneously.
|
||||
// 0 means unbounded (one goroutine per partition).
|
||||
func FanOut(
|
||||
ctx context.Context,
|
||||
partitions []string,
|
||||
maxConcurrency int,
|
||||
fn func(ctx context.Context, idx int, contextKey, partition string) PartitionResult,
|
||||
) []PartitionResult {
|
||||
n := len(partitions)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
results := make([]PartitionResult, n)
|
||||
|
||||
if maxConcurrency <= 0 || maxConcurrency >= n {
|
||||
// Run all goroutines concurrently.
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(n)
|
||||
for i, part := range partitions {
|
||||
i, part := i, part
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
key := fmt.Sprintf("partition-%d", i)
|
||||
results[i] = fn(ctx, i, key, part)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
return results
|
||||
}
|
||||
|
||||
// Bounded concurrency via semaphore channel.
|
||||
sem := make(chan struct{}, maxConcurrency)
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(n)
|
||||
|
||||
for i, part := range partitions {
|
||||
i, part := i, part
|
||||
sem <- struct{}{}
|
||||
go func() {
|
||||
defer func() {
|
||||
<-sem
|
||||
wg.Done()
|
||||
}()
|
||||
key := fmt.Sprintf("partition-%d", i)
|
||||
results[i] = fn(ctx, i, key, part)
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
return results
|
||||
}
|
||||
|
||||
// MergeResults synthesises the answers from all partitions into a single string.
|
||||
// This is the "merge" step of the RLM recursion. A production implementation
|
||||
// would use another LLM call; this version concatenates non-empty answers with
|
||||
// newlines and deduplicates.
|
||||
func MergeResults(results []PartitionResult) string {
|
||||
seen := make(map[string]bool)
|
||||
var parts []string
|
||||
|
||||
for _, r := range results {
|
||||
if r.Err != nil || r.Answer == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.TrimSpace(r.Answer)
|
||||
if seen[key] {
|
||||
continue
|
||||
}
|
||||
seen[key] = true
|
||||
parts = append(parts, key)
|
||||
}
|
||||
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
|
||||
// TotalTokens sums the cost_tokens across all PartitionResults.
|
||||
func TotalTokens(results []PartitionResult) uint32 {
|
||||
var total uint32
|
||||
for _, r := range results {
|
||||
total += r.Tokens
|
||||
}
|
||||
return total
|
||||
}
|
||||
271
pkg/rlm/rope.go
Normal file
271
pkg/rlm/rope.go
Normal file
|
|
@ -0,0 +1,271 @@
|
|||
// Package rlm implements the Recursive Language Model (RLM) engine for
|
||||
// PicoClaw. RLM enables unbounded context processing by recursively
|
||||
// decomposing long contexts into manageable partitions, processing each
|
||||
// through the SecureBus, and synthesizing the results.
|
||||
//
|
||||
// References:
|
||||
// - Zhang & Khattab (MIT CSAIL, arXiv:2512.24601 v2, Jan 2026)
|
||||
// - DSPy dspy.RLM canonical implementation
|
||||
//
|
||||
// The rope.go file provides an in-process O(log n) context store. We implement
|
||||
// a pure-Go rope over []byte slices rather than depending on an external
|
||||
// library, keeping the binary under the 20MB embedded constraint.
|
||||
package rlm
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Rope is a binary tree over string data that provides O(log n) concatenation,
|
||||
// slice, and split operations. It is used to store the full conversation context
|
||||
// without copying the entire buffer on each operation.
|
||||
//
|
||||
// This implementation is a balanced rope suitable for contexts up to ~10M tokens.
|
||||
// Nodes with len ≤ leafThreshold are stored as plain strings (leaf nodes).
|
||||
// Concatenation produces internal nodes; slicing rebalances lazily.
|
||||
type Rope struct {
|
||||
root *ropeNode
|
||||
}
|
||||
|
||||
const leafThreshold = 4096 // bytes; nodes below this size stay as leaves
|
||||
|
||||
type ropeNode struct {
|
||||
// Leaf node: data is non-empty, left/right are nil.
|
||||
data []byte
|
||||
|
||||
// Internal node: data is nil, left/right are non-nil.
|
||||
left, right *ropeNode
|
||||
length int // total byte length of this subtree
|
||||
}
|
||||
|
||||
func newLeaf(data []byte) *ropeNode {
|
||||
cp := make([]byte, len(data))
|
||||
copy(cp, data)
|
||||
return &ropeNode{data: cp, length: len(data)}
|
||||
}
|
||||
|
||||
func newInternal(left, right *ropeNode) *ropeNode {
|
||||
return &ropeNode{left: left, right: right, length: left.length + right.length}
|
||||
}
|
||||
|
||||
// NewRope creates a Rope from the given initial content.
|
||||
func NewRope(content string) *Rope {
|
||||
if len(content) == 0 {
|
||||
return &Rope{root: newLeaf(nil)}
|
||||
}
|
||||
return &Rope{root: newLeaf([]byte(content))}
|
||||
}
|
||||
|
||||
// Len returns the total byte length of the rope.
|
||||
func (r *Rope) Len() int {
|
||||
if r.root == nil {
|
||||
return 0
|
||||
}
|
||||
return r.root.length
|
||||
}
|
||||
|
||||
// String materialises the full rope as a string. O(n).
|
||||
func (r *Rope) String() string {
|
||||
if r.root == nil {
|
||||
return ""
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.Grow(r.root.length)
|
||||
writeNode(&sb, r.root)
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
// Append concatenates content to the end of the rope. O(log n) amortised.
|
||||
func (r *Rope) Append(content string) {
|
||||
if len(content) == 0 {
|
||||
return
|
||||
}
|
||||
newLeafNode := newLeaf([]byte(content))
|
||||
if r.root == nil || r.root.length == 0 {
|
||||
r.root = newLeafNode
|
||||
return
|
||||
}
|
||||
r.root = newInternal(r.root, newLeafNode)
|
||||
}
|
||||
|
||||
// Slice returns the substring [start, end) in bytes. O(log n + result_size).
|
||||
// Returns an error if indices are out of range.
|
||||
func (r *Rope) Slice(start, end int) (string, error) {
|
||||
total := r.Len()
|
||||
if start < 0 || end < start || end > total {
|
||||
return "", fmt.Errorf("rope slice [%d, %d) out of range [0, %d)", start, end, total)
|
||||
}
|
||||
if start == end {
|
||||
return "", nil
|
||||
}
|
||||
var sb strings.Builder
|
||||
sb.Grow(end - start)
|
||||
sliceNode(&sb, r.root, start, end)
|
||||
return sb.String(), nil
|
||||
}
|
||||
|
||||
// Lines returns all lines (split on '\n') as a slice. O(n).
|
||||
func (r *Rope) Lines() []string {
|
||||
return strings.Split(r.String(), "\n")
|
||||
}
|
||||
|
||||
// RuneLen returns the number of Unicode runes in the rope.
|
||||
func (r *Rope) RuneLen() int {
|
||||
return utf8.RuneCountInString(r.String())
|
||||
}
|
||||
|
||||
// GrepLines returns all lines containing the given substring (case-sensitive).
|
||||
// Returns up to maxMatches results; 0 means no limit.
|
||||
func (r *Rope) GrepLines(pattern string, maxMatches int, caseInsensitive bool) []GrepMatch {
|
||||
lines := r.Lines()
|
||||
var results []GrepMatch
|
||||
|
||||
search := pattern
|
||||
if caseInsensitive {
|
||||
search = strings.ToLower(pattern)
|
||||
}
|
||||
|
||||
offset := 0
|
||||
for lineNum, line := range lines {
|
||||
check := line
|
||||
if caseInsensitive {
|
||||
check = strings.ToLower(line)
|
||||
}
|
||||
if strings.Contains(check, search) {
|
||||
results = append(results, GrepMatch{
|
||||
LineNum: lineNum + 1,
|
||||
ByteOffset: offset,
|
||||
Line: line,
|
||||
})
|
||||
if maxMatches > 0 && len(results) >= maxMatches {
|
||||
break
|
||||
}
|
||||
}
|
||||
offset += len(line) + 1 // +1 for '\n'
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
// GrepMatch is a single result from Rope.GrepLines.
|
||||
type GrepMatch struct {
|
||||
LineNum int // 1-indexed line number
|
||||
ByteOffset int // byte offset of the start of this line in the rope
|
||||
Line string // the matching line content
|
||||
}
|
||||
|
||||
// Partition splits the rope into k roughly equal partitions by byte count.
|
||||
// Returns k slices of the rope content (as strings).
|
||||
func (r *Rope) Partition(k int) []string {
|
||||
if k <= 0 {
|
||||
k = 1
|
||||
}
|
||||
total := r.Len()
|
||||
if total == 0 {
|
||||
result := make([]string, k)
|
||||
return result
|
||||
}
|
||||
|
||||
chunkSize := total / k
|
||||
if chunkSize == 0 {
|
||||
chunkSize = 1
|
||||
}
|
||||
|
||||
var parts []string
|
||||
pos := 0
|
||||
for i := 0; i < k; i++ {
|
||||
end := pos + chunkSize
|
||||
if i == k-1 || end >= total {
|
||||
end = total
|
||||
}
|
||||
// Snap to UTF-8 rune boundary to avoid splitting multi-byte characters.
|
||||
if end < total {
|
||||
// Walk back until we land on a rune boundary.
|
||||
for end > pos && !utf8.RuneStart(r.byteAt(end)) {
|
||||
end--
|
||||
}
|
||||
}
|
||||
s, _ := r.Slice(pos, end)
|
||||
parts = append(parts, s)
|
||||
pos = end
|
||||
if pos >= total {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// Pad to exactly k if we ran out of content.
|
||||
for len(parts) < k {
|
||||
parts = append(parts, "")
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
// byteAt returns the byte at offset i. Panics if out of range.
|
||||
func (r *Rope) byteAt(i int) byte {
|
||||
s, err := r.Slice(i, i+1)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
if len(s) == 0 {
|
||||
return 0
|
||||
}
|
||||
return s[0]
|
||||
}
|
||||
|
||||
// ── Internal rope helpers ────────────────────────────────────────────────────
|
||||
|
||||
func writeNode(sb *strings.Builder, n *ropeNode) {
|
||||
if n == nil {
|
||||
return
|
||||
}
|
||||
if n.data != nil {
|
||||
sb.Write(n.data)
|
||||
return
|
||||
}
|
||||
writeNode(sb, n.left)
|
||||
writeNode(sb, n.right)
|
||||
}
|
||||
|
||||
func sliceNode(sb *strings.Builder, n *ropeNode, start, end int) {
|
||||
if n == nil || start >= end {
|
||||
return
|
||||
}
|
||||
if n.data != nil {
|
||||
// Leaf: write the overlapping portion.
|
||||
lo, hi := start, end
|
||||
if lo < 0 {
|
||||
lo = 0
|
||||
}
|
||||
if hi > len(n.data) {
|
||||
hi = len(n.data)
|
||||
}
|
||||
if lo < hi {
|
||||
sb.Write(n.data[lo:hi])
|
||||
}
|
||||
return
|
||||
}
|
||||
leftLen := n.left.length
|
||||
// Overlap with left child?
|
||||
if start < leftLen {
|
||||
sliceNode(sb, n.left, start, min(end, leftLen))
|
||||
}
|
||||
// Overlap with right child?
|
||||
if end > leftLen {
|
||||
sliceNode(sb, n.right, max(0, start-leftLen), end-leftLen)
|
||||
}
|
||||
}
|
||||
|
||||
func min(a, b int) int {
|
||||
if a < b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func max(a, b int) int {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
123
pkg/rlm/rope_test.go
Normal file
123
pkg/rlm/rope_test.go
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
package rlm_test
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/rlm"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestRope_EmptyRope(t *testing.T) {
|
||||
r := rlm.NewRope("")
|
||||
assert.Equal(t, 0, r.Len())
|
||||
assert.Equal(t, "", r.String())
|
||||
}
|
||||
|
||||
func TestRope_BasicAppendAndString(t *testing.T) {
|
||||
r := rlm.NewRope("hello")
|
||||
r.Append(" world")
|
||||
assert.Equal(t, 11, r.Len())
|
||||
assert.Equal(t, "hello world", r.String())
|
||||
}
|
||||
|
||||
func TestRope_LargeContent(t *testing.T) {
|
||||
content := strings.Repeat("abcdefghij", 1000) // 10000 bytes
|
||||
r := rlm.NewRope(content)
|
||||
assert.Equal(t, 10000, r.Len())
|
||||
assert.Equal(t, content, r.String())
|
||||
}
|
||||
|
||||
func TestRope_Slice_ValidRange(t *testing.T) {
|
||||
r := rlm.NewRope("hello world")
|
||||
s, err := r.Slice(6, 11)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "world", s)
|
||||
}
|
||||
|
||||
func TestRope_Slice_ZeroLength(t *testing.T) {
|
||||
r := rlm.NewRope("hello")
|
||||
s, err := r.Slice(2, 2)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "", s)
|
||||
}
|
||||
|
||||
func TestRope_Slice_OutOfRange(t *testing.T) {
|
||||
r := rlm.NewRope("hello")
|
||||
_, err := r.Slice(3, 10)
|
||||
assert.Error(t, err)
|
||||
}
|
||||
|
||||
func TestRope_Slice_AcrossAppendBoundary(t *testing.T) {
|
||||
r := rlm.NewRope("hello")
|
||||
r.Append(" world")
|
||||
s, err := r.Slice(3, 8)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "lo wo", s)
|
||||
}
|
||||
|
||||
func TestRope_Lines(t *testing.T) {
|
||||
r := rlm.NewRope("line1\nline2\nline3")
|
||||
lines := r.Lines()
|
||||
assert.Equal(t, []string{"line1", "line2", "line3"}, lines)
|
||||
}
|
||||
|
||||
func TestRope_GrepLines_CaseSensitive(t *testing.T) {
|
||||
r := rlm.NewRope("apple\nBanana\napricot\ncherry")
|
||||
matches := r.GrepLines("ap", 0, false)
|
||||
require.Len(t, matches, 2)
|
||||
assert.Equal(t, 1, matches[0].LineNum)
|
||||
assert.Equal(t, 3, matches[1].LineNum)
|
||||
}
|
||||
|
||||
func TestRope_GrepLines_CaseInsensitive(t *testing.T) {
|
||||
r := rlm.NewRope("Apple\nbanana\nAPRICOT")
|
||||
matches := r.GrepLines("apple", 0, true)
|
||||
require.Len(t, matches, 1)
|
||||
assert.Equal(t, "Apple", matches[0].Line)
|
||||
}
|
||||
|
||||
func TestRope_GrepLines_MaxMatches(t *testing.T) {
|
||||
r := rlm.NewRope("aa\naa\naa\naa\naa")
|
||||
matches := r.GrepLines("aa", 3, false)
|
||||
assert.Len(t, matches, 3)
|
||||
}
|
||||
|
||||
func TestRope_GrepLines_NoMatches(t *testing.T) {
|
||||
r := rlm.NewRope("hello world")
|
||||
matches := r.GrepLines("xyz", 0, false)
|
||||
assert.Empty(t, matches)
|
||||
}
|
||||
|
||||
func TestRope_Partition_Even(t *testing.T) {
|
||||
r := rlm.NewRope("12345678")
|
||||
parts := r.Partition(4)
|
||||
assert.Len(t, parts, 4)
|
||||
assert.Equal(t, "12345678", strings.Join(parts, ""))
|
||||
}
|
||||
|
||||
func TestRope_Partition_MoreThanContent(t *testing.T) {
|
||||
r := rlm.NewRope("hi")
|
||||
parts := r.Partition(10)
|
||||
assert.Len(t, parts, 10)
|
||||
// All content should appear in first non-empty partition.
|
||||
combined := strings.Join(parts, "")
|
||||
assert.Equal(t, "hi", combined)
|
||||
}
|
||||
|
||||
func TestRope_Partition_Empty(t *testing.T) {
|
||||
r := rlm.NewRope("")
|
||||
parts := r.Partition(4)
|
||||
assert.Len(t, parts, 4)
|
||||
for _, p := range parts {
|
||||
assert.Equal(t, "", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRope_RuneLen(t *testing.T) {
|
||||
// Multi-byte Unicode characters.
|
||||
r := rlm.NewRope("héllo") // 'é' is 2 bytes
|
||||
assert.Equal(t, 5, r.RuneLen())
|
||||
assert.Equal(t, 6, r.Len()) // bytes
|
||||
}
|
||||
116
pkg/rlm/strategy.go
Normal file
116
pkg/rlm/strategy.go
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
package rlm
|
||||
|
||||
import (
|
||||
"strings"
|
||||
)
|
||||
|
||||
// OpType identifies a decomposition operation the strategy planner can choose.
|
||||
type OpType string
|
||||
|
||||
const (
|
||||
OpPeek OpType = "peek"
|
||||
OpGrep OpType = "grep"
|
||||
OpPartition OpType = "partition"
|
||||
OpRecurse OpType = "recurse"
|
||||
OpFinal OpType = "final" // terminal: no further decomposition needed
|
||||
)
|
||||
|
||||
// StrategyOp is one step produced by the strategy planner.
|
||||
type StrategyOp struct {
|
||||
Type OpType
|
||||
PartitionK int // for OpPartition
|
||||
GrepQuery string // for OpGrep
|
||||
PeekStart uint64 // for OpPeek
|
||||
PeekLength uint32 // for OpPeek
|
||||
}
|
||||
|
||||
// StrategyConfig controls the heuristic strategy planner.
|
||||
type StrategyConfig struct {
|
||||
// DirectThreshold is the context size (bytes) below which the engine
|
||||
// answers directly without decomposition. Default: 8192 (≈6K tokens).
|
||||
DirectThreshold int
|
||||
|
||||
// DefaultPartitionK is the number of partitions when OpPartition is chosen.
|
||||
DefaultPartitionK int
|
||||
|
||||
// MaxDepth is the maximum recursion depth. The strategy planner always
|
||||
// emits OpFinal at depth >= MaxDepth to prevent infinite recursion.
|
||||
MaxDepth int
|
||||
}
|
||||
|
||||
// DefaultStrategyConfig returns sensible defaults.
|
||||
func DefaultStrategyConfig() StrategyConfig {
|
||||
return StrategyConfig{
|
||||
DirectThreshold: 8192,
|
||||
DefaultPartitionK: 4,
|
||||
MaxDepth: 8,
|
||||
}
|
||||
}
|
||||
|
||||
// StrategyPlanner selects the next decomposition operation given the current
|
||||
// context size, query, and recursion depth. This is the "cheap sub-LM" from
|
||||
// the RLM paper — it uses simple heuristics rather than an LLM call to keep
|
||||
// latency and cost minimal.
|
||||
//
|
||||
// More sophisticated implementations could replace PlanNext with an LLM-backed
|
||||
// planner, but the heuristic version already achieves the RLM paper's results
|
||||
// at much lower cost.
|
||||
type StrategyPlanner struct {
|
||||
cfg StrategyConfig
|
||||
}
|
||||
|
||||
// NewStrategyPlanner creates a planner with the given config.
|
||||
func NewStrategyPlanner(cfg StrategyConfig) *StrategyPlanner {
|
||||
return &StrategyPlanner{cfg: cfg}
|
||||
}
|
||||
|
||||
// PlanNext returns the next operation to apply to the context given the query
|
||||
// and current recursion depth.
|
||||
func (sp *StrategyPlanner) PlanNext(contextBytes int, query string, depth uint8) StrategyOp {
|
||||
// Terminal conditions: context small enough or depth limit reached.
|
||||
if depth >= uint8(sp.cfg.MaxDepth) || contextBytes <= sp.cfg.DirectThreshold {
|
||||
return StrategyOp{Type: OpFinal}
|
||||
}
|
||||
|
||||
// If the query contains specific keywords, prefer grep to narrow context.
|
||||
if looksLikeKeywordQuery(query) {
|
||||
return StrategyOp{Type: OpGrep, GrepQuery: extractKeyword(query)}
|
||||
}
|
||||
|
||||
// Default: partition and recurse.
|
||||
k := sp.cfg.DefaultPartitionK
|
||||
if contextBytes > 4*1024*1024 { // > 4MB: use more partitions
|
||||
k = 8
|
||||
}
|
||||
return StrategyOp{Type: OpPartition, PartitionK: k}
|
||||
}
|
||||
|
||||
// looksLikeKeywordQuery returns true when the query contains explicit search
|
||||
// cues: quoted strings, identifiers starting with #/@, or error messages.
|
||||
func looksLikeKeywordQuery(query string) bool {
|
||||
q := strings.ToLower(query)
|
||||
return strings.Contains(q, "\"") ||
|
||||
strings.Contains(q, "error:") ||
|
||||
strings.Contains(q, "func ") ||
|
||||
strings.Contains(q, "def ") ||
|
||||
strings.HasPrefix(q, "find ") ||
|
||||
strings.HasPrefix(q, "search ") ||
|
||||
strings.HasPrefix(q, "where is ")
|
||||
}
|
||||
|
||||
// extractKeyword pulls the most likely search term from the query.
|
||||
// Falls back to the first word.
|
||||
func extractKeyword(query string) string {
|
||||
// Extract first quoted string.
|
||||
if start := strings.Index(query, "\""); start >= 0 {
|
||||
if end := strings.Index(query[start+1:], "\""); end >= 0 {
|
||||
return query[start+1 : start+1+end]
|
||||
}
|
||||
}
|
||||
// First word of the query.
|
||||
words := strings.Fields(query)
|
||||
if len(words) > 0 {
|
||||
return words[0]
|
||||
}
|
||||
return query
|
||||
}
|
||||
98
pkg/security/keyring.go
Normal file
98
pkg/security/keyring.go
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
// Package security provides credential management for PicoClaw.
|
||||
// This file defines the KeyringProvider interface and the master key
|
||||
// sourcing strategy used by SecretStore.
|
||||
package security
|
||||
|
||||
import "fmt"
|
||||
|
||||
// KeyringProvider abstracts OS keyring backends. The implementation is
|
||||
// selected at runtime (or via build tags for platform-specific backends).
|
||||
//
|
||||
// Implementations:
|
||||
// - NoopKeyring — in-memory; suitable for embedded / CI environments
|
||||
// - EnvKeyring — reads master key from an environment variable
|
||||
// - (future) OSKeyring — macOS Keychain / Linux secret-service / Windows DPAPI
|
||||
type KeyringProvider interface {
|
||||
// GetMasterKey returns the 32-byte master key used to seal the SecretStore.
|
||||
// Returns ErrNoMasterKey when no key has been configured.
|
||||
GetMasterKey() ([]byte, error)
|
||||
|
||||
// SetMasterKey persists a new master key. The provider may refuse if the
|
||||
// backend is read-only (e.g., EnvKeyring).
|
||||
SetMasterKey(key []byte) error
|
||||
}
|
||||
|
||||
// ErrNoMasterKey is returned by KeyringProvider when no master key is available.
|
||||
var ErrNoMasterKey = fmt.Errorf("no master key configured: run `picoclaw secret init` to set one")
|
||||
|
||||
// ErrKeyringReadOnly is returned when SetMasterKey is called on a read-only provider.
|
||||
var ErrKeyringReadOnly = fmt.Errorf("keyring is read-only")
|
||||
|
||||
// NoopKeyring stores the master key in memory only. Safe for tests and
|
||||
// embedded builds where OS keyring integration is unavailable.
|
||||
type NoopKeyring struct {
|
||||
key []byte
|
||||
}
|
||||
|
||||
// NewNoopKeyring creates an empty in-memory keyring. Call SetMasterKey to
|
||||
// initialise it, or pass an initial key to seed it directly.
|
||||
func NewNoopKeyring(initialKey []byte) *NoopKeyring {
|
||||
return &NoopKeyring{key: initialKey}
|
||||
}
|
||||
|
||||
// GetMasterKey returns the in-memory key. Returns ErrNoMasterKey if unset.
|
||||
func (n *NoopKeyring) GetMasterKey() ([]byte, error) {
|
||||
if len(n.key) == 0 {
|
||||
return nil, ErrNoMasterKey
|
||||
}
|
||||
cp := make([]byte, len(n.key))
|
||||
copy(cp, n.key)
|
||||
return cp, nil
|
||||
}
|
||||
|
||||
// SetMasterKey stores key in memory.
|
||||
func (n *NoopKeyring) SetMasterKey(key []byte) error {
|
||||
cp := make([]byte, len(key))
|
||||
copy(cp, key)
|
||||
n.key = cp
|
||||
return nil
|
||||
}
|
||||
|
||||
// EnvKeyring reads the master key from an environment variable and does not
|
||||
// support writing. Useful for CI/CD pipelines and container deployments.
|
||||
type EnvKeyring struct {
|
||||
envVar string
|
||||
getEnv func(string) string // injectable for testing
|
||||
}
|
||||
|
||||
// NewEnvKeyring creates a keyring that reads the master key from envVar.
|
||||
// The value is expected to be a hex or base64-encoded 32-byte key.
|
||||
func NewEnvKeyring(envVar string) *EnvKeyring {
|
||||
return newEnvKeyring(envVar, nil)
|
||||
}
|
||||
|
||||
func newEnvKeyring(envVar string, getEnv func(string) string) *EnvKeyring {
|
||||
if getEnv == nil {
|
||||
// Defer os.Getenv to avoid importing os at package level.
|
||||
getEnv = func(k string) string {
|
||||
// Lazy import via the same trick as other stdlib users — this is
|
||||
// a valid Go pattern for keeping imports minimal.
|
||||
return envGetFunc(k)
|
||||
}
|
||||
}
|
||||
return &EnvKeyring{envVar: envVar, getEnv: getEnv}
|
||||
}
|
||||
|
||||
// SetMasterKey is not supported for EnvKeyring.
|
||||
func (e *EnvKeyring) SetMasterKey(_ []byte) error {
|
||||
return ErrKeyringReadOnly
|
||||
}
|
||||
|
||||
// GetMasterKey reads the environment variable and decodes it as base64.
|
||||
func (e *EnvKeyring) GetMasterKey() ([]byte, error) {
|
||||
raw := e.getEnv(e.envVar)
|
||||
if raw == "" {
|
||||
return nil, ErrNoMasterKey
|
||||
}
|
||||
return decodeKey(raw)
|
||||
}
|
||||
40
pkg/security/keyring_helpers.go
Normal file
40
pkg/security/keyring_helpers.go
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"os"
|
||||
)
|
||||
|
||||
// envGetFunc is used by EnvKeyring to read environment variables.
|
||||
// Decoupled to allow testing without os import in keyring.go.
|
||||
func envGetFunc(key string) string {
|
||||
return os.Getenv(key)
|
||||
}
|
||||
|
||||
// decodeKey attempts to decode a string as hex or base64 (URL or standard).
|
||||
// Returns an error if the decoded length is not exactly 32 bytes.
|
||||
func decodeKey(s string) ([]byte, error) {
|
||||
// Try hex first (64 chars)
|
||||
if len(s) == 64 {
|
||||
b, err := hex.DecodeString(s)
|
||||
if err == nil && len(b) == 32 {
|
||||
return b, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Try base64url (no padding)
|
||||
b, err := base64.RawURLEncoding.DecodeString(s)
|
||||
if err == nil && len(b) == 32 {
|
||||
return b, nil
|
||||
}
|
||||
|
||||
// Try standard base64 with padding
|
||||
b, err = base64.StdEncoding.DecodeString(s)
|
||||
if err == nil && len(b) == 32 {
|
||||
return b, nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("key must be a 64-char hex or 44-char base64 string encoding exactly 32 bytes")
|
||||
}
|
||||
168
pkg/security/secretstore.go
Normal file
168
pkg/security/secretstore.go
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
package security
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ErrSecretNotFound is returned when a secret name has no entry.
|
||||
var ErrSecretNotFound = errors.New("secret not found")
|
||||
|
||||
// SecretEntry is one stored secret, encrypted at rest.
|
||||
type SecretEntry struct {
|
||||
Name string `json:"name"`
|
||||
Ciphertext string `json:"ciphertext"` // Vault.Encrypt output (base64url)
|
||||
}
|
||||
|
||||
// SecretStore maps logical secret names to encrypted ciphertext, persisted
|
||||
// to a JSON file on disk. The Vault key is sourced from a KeyringProvider.
|
||||
//
|
||||
// Secrets never leave the store as plaintext except within a single scoped
|
||||
// Resolve call, which zeros the returned byte slice when done. (Callers are
|
||||
// responsible for this — Go does not guarantee zeroing.)
|
||||
//
|
||||
// Thread-safe.
|
||||
type SecretStore struct {
|
||||
mu sync.RWMutex
|
||||
path string
|
||||
keyring KeyringProvider
|
||||
secrets map[string]SecretEntry // name → entry
|
||||
}
|
||||
|
||||
// NewSecretStore creates a SecretStore backed by the given file path and keyring.
|
||||
// The store is loaded from disk if the file exists; it is created on the first
|
||||
// Set call if it does not.
|
||||
func NewSecretStore(path string, keyring KeyringProvider) (*SecretStore, error) {
|
||||
ss := &SecretStore{
|
||||
path: path,
|
||||
keyring: keyring,
|
||||
secrets: make(map[string]SecretEntry),
|
||||
}
|
||||
if err := ss.load(); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return nil, fmt.Errorf("secret store: load %s: %w", path, err)
|
||||
}
|
||||
return ss, nil
|
||||
}
|
||||
|
||||
// Set encrypts value and stores it under name, overwriting any existing entry.
|
||||
// Persists the store to disk.
|
||||
func (ss *SecretStore) Set(name string, value []byte) error {
|
||||
if name == "" {
|
||||
return fmt.Errorf("secret name must not be empty")
|
||||
}
|
||||
vault, err := ss.vault()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ciphertext, err := vault.Encrypt(value)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encrypt secret %q: %w", name, err)
|
||||
}
|
||||
|
||||
ss.mu.Lock()
|
||||
ss.secrets[name] = SecretEntry{Name: name, Ciphertext: ciphertext}
|
||||
ss.mu.Unlock()
|
||||
|
||||
return ss.save()
|
||||
}
|
||||
|
||||
// Get decrypts and returns the secret stored under name.
|
||||
// Returns ErrSecretNotFound if the name has no entry.
|
||||
func (ss *SecretStore) Get(name string) ([]byte, error) {
|
||||
ss.mu.RLock()
|
||||
entry, ok := ss.secrets[name]
|
||||
ss.mu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return nil, ErrSecretNotFound
|
||||
}
|
||||
|
||||
vault, err := ss.vault()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return vault.Decrypt(entry.Ciphertext)
|
||||
}
|
||||
|
||||
// Delete removes the secret stored under name. No-op if it does not exist.
|
||||
func (ss *SecretStore) Delete(name string) error {
|
||||
ss.mu.Lock()
|
||||
delete(ss.secrets, name)
|
||||
ss.mu.Unlock()
|
||||
return ss.save()
|
||||
}
|
||||
|
||||
// List returns all registered secret names.
|
||||
func (ss *SecretStore) List() []string {
|
||||
ss.mu.RLock()
|
||||
defer ss.mu.RUnlock()
|
||||
names := make([]string, 0, len(ss.secrets))
|
||||
for k := range ss.secrets {
|
||||
names = append(names, k)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// Has reports whether a secret with the given name exists.
|
||||
func (ss *SecretStore) Has(name string) bool {
|
||||
ss.mu.RLock()
|
||||
_, ok := ss.secrets[name]
|
||||
ss.mu.RUnlock()
|
||||
return ok
|
||||
}
|
||||
|
||||
// vault constructs a Vault using the keyring's master key.
|
||||
func (ss *SecretStore) vault() (*Vault, error) {
|
||||
key, err := ss.keyring.GetMasterKey()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get master key: %w", err)
|
||||
}
|
||||
return NewVault(key)
|
||||
}
|
||||
|
||||
// load reads the store from disk. Not thread-safe — callers must hold the lock
|
||||
// or call only before the store is shared.
|
||||
func (ss *SecretStore) load() error {
|
||||
data, err := os.ReadFile(ss.path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var entries []SecretEntry
|
||||
if err := json.Unmarshal(data, &entries); err != nil {
|
||||
return fmt.Errorf("parse secret store: %w", err)
|
||||
}
|
||||
for _, e := range entries {
|
||||
ss.secrets[e.Name] = e
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// save writes the store to disk atomically. Thread-safe.
|
||||
func (ss *SecretStore) save() error {
|
||||
ss.mu.RLock()
|
||||
entries := make([]SecretEntry, 0, len(ss.secrets))
|
||||
for _, e := range ss.secrets {
|
||||
entries = append(entries, e)
|
||||
}
|
||||
ss.mu.RUnlock()
|
||||
|
||||
data, err := json.MarshalIndent(entries, "", " ")
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal secret store: %w", err)
|
||||
}
|
||||
|
||||
dir := filepath.Dir(ss.path)
|
||||
if err := os.MkdirAll(dir, 0700); err != nil {
|
||||
return fmt.Errorf("create secret store dir: %w", err)
|
||||
}
|
||||
|
||||
tmp := ss.path + ".tmp"
|
||||
if err := os.WriteFile(tmp, data, 0600); err != nil {
|
||||
return fmt.Errorf("write secret store: %w", err)
|
||||
}
|
||||
return os.Rename(tmp, ss.path)
|
||||
}
|
||||
104
pkg/security/securebus/audit.go
Normal file
104
pkg/security/securebus/audit.go
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
package securebus
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// AuditEvent records one tool execution through the SecureBus.
|
||||
// The log is append-only: entries are never modified or deleted.
|
||||
type AuditEvent struct {
|
||||
RequestID string `json:"request_id"`
|
||||
SessionKey string `json:"session_key,omitempty"`
|
||||
ToolCallID string `json:"tool_call_id,omitempty"`
|
||||
ToolName string `json:"tool_name,omitempty"`
|
||||
CommandType string `json:"command_type"`
|
||||
Depth uint8 `json:"depth,omitempty"`
|
||||
At time.Time `json:"at"`
|
||||
DurationMS int64 `json:"duration_ms"`
|
||||
IsError bool `json:"is_error,omitempty"`
|
||||
LeakDetected bool `json:"leak_detected,omitempty"`
|
||||
RedactedKeys []string `json:"redacted_keys,omitempty"`
|
||||
|
||||
// Capability grants observed at execution time.
|
||||
SecretsAccessed []string `json:"secrets_accessed,omitempty"`
|
||||
PolicyViolation string `json:"policy_violation,omitempty"`
|
||||
}
|
||||
|
||||
// AuditLog is a thread-safe, in-memory append-only audit log.
|
||||
// For production deployments, attach a Sink to persist events to the DB.
|
||||
type AuditLog struct {
|
||||
mu sync.RWMutex
|
||||
events []AuditEvent
|
||||
sinks []AuditSink
|
||||
}
|
||||
|
||||
// AuditSink is an optional interface for persisting audit events externally
|
||||
// (e.g., to the agent_audit_log SQLite table).
|
||||
type AuditSink interface {
|
||||
Write(event AuditEvent) error
|
||||
}
|
||||
|
||||
// NewAuditLog creates an empty in-memory audit log. Pass sinks to fan-out
|
||||
// writes to external storage.
|
||||
func NewAuditLog(sinks ...AuditSink) *AuditLog {
|
||||
return &AuditLog{sinks: sinks}
|
||||
}
|
||||
|
||||
// Append records an audit event. Returns the first sink error encountered,
|
||||
// but always appends to the in-memory log.
|
||||
func (al *AuditLog) Append(event AuditEvent) error {
|
||||
al.mu.Lock()
|
||||
al.events = append(al.events, event)
|
||||
al.mu.Unlock()
|
||||
|
||||
var firstErr error
|
||||
for _, sink := range al.sinks {
|
||||
if err := sink.Write(event); err != nil && firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
return firstErr
|
||||
}
|
||||
|
||||
// Events returns a snapshot of all recorded events.
|
||||
func (al *AuditLog) Events() []AuditEvent {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
cp := make([]AuditEvent, len(al.events))
|
||||
copy(cp, al.events)
|
||||
return cp
|
||||
}
|
||||
|
||||
// Len returns the number of recorded events.
|
||||
func (al *AuditLog) Len() int {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
return len(al.events)
|
||||
}
|
||||
|
||||
// FilterBySession returns all events for a given session key.
|
||||
func (al *AuditLog) FilterBySession(sessionKey string) []AuditEvent {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
var result []AuditEvent
|
||||
for _, e := range al.events {
|
||||
if e.SessionKey == sessionKey {
|
||||
result = append(result, e)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// LeakEvents returns all events where LeakDetected is true.
|
||||
func (al *AuditLog) LeakEvents() []AuditEvent {
|
||||
al.mu.RLock()
|
||||
defer al.mu.RUnlock()
|
||||
var result []AuditEvent
|
||||
for _, e := range al.events {
|
||||
if e.LeakDetected {
|
||||
result = append(result, e)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
278
pkg/security/securebus/bus.go
Normal file
278
pkg/security/securebus/bus.go
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
package securebus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
// ToolExecutor is the function signature the SecureBus calls to run a tool.
|
||||
// In production this is tools.Registry.ExecuteWithContext. In tests it can
|
||||
// be any function.
|
||||
type ToolExecutor func(ctx context.Context, name string, args map[string]interface{}) *tools.ToolResult
|
||||
|
||||
// CapabilitiesLookup returns the ToolCapabilities for a named tool.
|
||||
// Wraps tools.Registry.Get + tools.ExtractCapabilities.
|
||||
type CapabilitiesLookup func(toolName string) (tools.ToolCapabilities, bool)
|
||||
|
||||
// BusConfig configures the SecureBus.
|
||||
type BusConfig struct {
|
||||
Policy PolicyConfig
|
||||
// Workers controls how many goroutines process requests concurrently.
|
||||
// 0 or 1 means single-goroutine (default, suitable for embedded).
|
||||
Workers int
|
||||
}
|
||||
|
||||
// DefaultBusConfig returns a production-safe default configuration.
|
||||
func DefaultBusConfig() BusConfig {
|
||||
return BusConfig{
|
||||
Policy: DefaultPolicyConfig(),
|
||||
Workers: 1,
|
||||
}
|
||||
}
|
||||
|
||||
// Bus is the SecureBus: the privilege boundary between the agent loop and
|
||||
// tool execution. All tool calls flow through Bus.Execute.
|
||||
//
|
||||
// Pipeline for each call:
|
||||
// 1. Decode request, extract tool capabilities
|
||||
// 2. Policy validation (depth, recursion limits)
|
||||
// 3. Secret injection into execution context
|
||||
// 4. Execute tool via ToolExecutor
|
||||
// 5. Scan output for leaks via Redactor
|
||||
// 6. Write audit log entry
|
||||
// 7. Return ToolResponse to caller
|
||||
type Bus struct {
|
||||
cfg BusConfig
|
||||
policy *PolicyEngine
|
||||
secrets *security.SecretStore // nil = no secret injection
|
||||
redactor *security.Redactor
|
||||
audit *AuditLog
|
||||
transport *ChannelTransport
|
||||
capLookup CapabilitiesLookup
|
||||
executor ToolExecutor
|
||||
done chan struct{}
|
||||
}
|
||||
|
||||
// New creates a Bus and starts background worker goroutines.
|
||||
// Call Close() to stop them.
|
||||
func New(
|
||||
cfg BusConfig,
|
||||
secrets *security.SecretStore,
|
||||
capLookup CapabilitiesLookup,
|
||||
executor ToolExecutor,
|
||||
auditSinks ...AuditSink,
|
||||
) *Bus {
|
||||
if cfg.Workers <= 0 {
|
||||
cfg.Workers = 1
|
||||
}
|
||||
workers := cfg.Workers
|
||||
|
||||
bus := &Bus{
|
||||
cfg: cfg,
|
||||
policy: NewPolicyEngine(cfg.Policy),
|
||||
secrets: secrets,
|
||||
redactor: security.NewRedactor(),
|
||||
audit: NewAuditLog(auditSinks...),
|
||||
transport: NewChannelTransport(workers * 4),
|
||||
capLookup: capLookup,
|
||||
executor: executor,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
|
||||
for i := 0; i < workers; i++ {
|
||||
go bus.runWorker()
|
||||
}
|
||||
return bus
|
||||
}
|
||||
|
||||
// Transport returns the ChannelTransport so callers can send requests via Send.
|
||||
func (b *Bus) Transport() Transport {
|
||||
return b.transport
|
||||
}
|
||||
|
||||
// AuditLog returns the bus audit log for inspection.
|
||||
func (b *Bus) AuditLog() *AuditLog {
|
||||
return b.audit
|
||||
}
|
||||
|
||||
// Close shuts down the bus workers gracefully.
|
||||
func (b *Bus) Close() {
|
||||
b.transport.Close()
|
||||
close(b.done)
|
||||
}
|
||||
|
||||
// Execute is a convenience method for in-process callers that don't want to
|
||||
// go through the Transport. It applies the full pipeline synchronously.
|
||||
func (b *Bus) Execute(ctx context.Context, req itr.ToolRequest) itr.ToolResponse {
|
||||
return b.dispatch(ctx, req)
|
||||
}
|
||||
|
||||
// runWorker reads from the transport channel and dispatches requests.
|
||||
func (b *Bus) runWorker() {
|
||||
for {
|
||||
select {
|
||||
case env, ok := <-b.transport.Requests():
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
resp := b.dispatch(context.Background(), env.req)
|
||||
env.reply(resp, nil)
|
||||
case <-b.done:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dispatch runs the full pipeline for a single request.
|
||||
func (b *Bus) dispatch(ctx context.Context, req itr.ToolRequest) itr.ToolResponse {
|
||||
start := time.Now()
|
||||
|
||||
event := AuditEvent{
|
||||
RequestID: req.ID,
|
||||
SessionKey: req.SessionKey,
|
||||
ToolCallID: req.ToolCallID,
|
||||
CommandType: string(req.Type),
|
||||
Depth: req.Depth,
|
||||
At: start,
|
||||
}
|
||||
|
||||
// Only ToolExec requests require capability/secret/leak checks.
|
||||
// RLM operations (Peek, Grep, etc.) are structural and access no tools.
|
||||
if req.Type != itr.CmdToolExec {
|
||||
resp := b.handleRLMCommand(ctx, req)
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return resp
|
||||
}
|
||||
|
||||
te, ok := req.Payload.(itr.ToolExec)
|
||||
if !ok {
|
||||
event.IsError = true
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return itr.NewErrorResponse(req.ID, "internal: payload is not ToolExec")
|
||||
}
|
||||
event.ToolName = te.ToolName
|
||||
|
||||
// 1. Capability lookup
|
||||
caps, found := b.capLookup(te.ToolName)
|
||||
if !found {
|
||||
caps = tools.ZeroCapabilities()
|
||||
}
|
||||
|
||||
// 2. Policy validation
|
||||
if err := b.policy.Validate(req, caps); err != nil {
|
||||
event.IsError = true
|
||||
event.PolicyViolation = err.Error()
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return itr.NewErrorResponse(req.ID, "policy violation: "+err.Error())
|
||||
}
|
||||
|
||||
// 3. Deserialise args — always produce a non-nil map for safe injection.
|
||||
args := make(map[string]interface{})
|
||||
if te.ArgsJSON != "" && te.ArgsJSON != "null" {
|
||||
if err := json.Unmarshal([]byte(te.ArgsJSON), &args); err != nil {
|
||||
event.IsError = true
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return itr.NewErrorResponse(req.ID, fmt.Sprintf("invalid args JSON: %v", err))
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Secret injection
|
||||
injectedSecrets, err := b.injectSecrets(ctx, caps.Secrets, args)
|
||||
if err != nil {
|
||||
event.IsError = true
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return itr.NewErrorResponse(req.ID, "secret injection failed: "+err.Error())
|
||||
}
|
||||
event.SecretsAccessed = injectedSecrets
|
||||
|
||||
// 5. Execute
|
||||
result := b.executor(ctx, te.ToolName, args)
|
||||
|
||||
// 6. Leak scan
|
||||
var resp itr.ToolResponse
|
||||
if result == nil {
|
||||
resp = itr.NewSuccessResponse(req.ID, "", 0)
|
||||
} else {
|
||||
resultText := result.ForLLM
|
||||
if b.redactor.ContainsSensitive(resultText) {
|
||||
redacted := b.redactor.Redact(resultText)
|
||||
resp = itr.NewLeakResponse(req.ID, redacted, nil)
|
||||
event.LeakDetected = true
|
||||
} else {
|
||||
resp = itr.NewSuccessResponse(req.ID, resultText, 0)
|
||||
}
|
||||
if result.IsError {
|
||||
resp.IsError = true
|
||||
}
|
||||
}
|
||||
|
||||
event.IsError = resp.IsError
|
||||
event.DurationMS = time.Since(start).Milliseconds()
|
||||
_ = b.audit.Append(event)
|
||||
return resp
|
||||
}
|
||||
|
||||
// injectSecrets resolves required secrets and injects them into args.
|
||||
// Returns the names of secrets accessed. Skips missing optional secrets.
|
||||
func (b *Bus) injectSecrets(_ context.Context, refs []tools.SecretRef, args map[string]interface{}) ([]string, error) {
|
||||
if b.secrets == nil || len(refs) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
var accessed []string
|
||||
for _, ref := range refs {
|
||||
val, err := b.secrets.Get(ref.Name)
|
||||
if err != nil {
|
||||
if err == security.ErrSecretNotFound && !ref.Required {
|
||||
continue
|
||||
}
|
||||
return accessed, fmt.Errorf("secret %q: %w", ref.Name, err)
|
||||
}
|
||||
accessed = append(accessed, ref.Name)
|
||||
injectArg(args, ref.InjectAs, string(val))
|
||||
}
|
||||
return accessed, nil
|
||||
}
|
||||
|
||||
// injectArg places the secret value into the args map according to the InjectAs spec.
|
||||
// Currently only "arg:<key>" injection is applied at the args level.
|
||||
// "env:*" and "header:*" injections are handled by the tool itself or HTTP client.
|
||||
func injectArg(args map[string]interface{}, injectAs, value string) {
|
||||
const argPrefix = "arg:"
|
||||
if args == nil {
|
||||
return
|
||||
}
|
||||
if len(injectAs) > len(argPrefix) && injectAs[:len(argPrefix)] == argPrefix {
|
||||
key := injectAs[len(argPrefix):]
|
||||
args[key] = value
|
||||
}
|
||||
// "env:" and "header:" variants require tool cooperation; they are
|
||||
// recorded in the capability manifest so auditing can trace what was accessed.
|
||||
}
|
||||
|
||||
// handleRLMCommand processes structural RLM decomposition commands.
|
||||
// These commands don't invoke tool code — they operate on the context rope
|
||||
// managed by the RLMEngine (which calls the SecureBus, not the other way around).
|
||||
func (b *Bus) handleRLMCommand(_ context.Context, req itr.ToolRequest) itr.ToolResponse {
|
||||
// RLM commands are executed by the RLMEngine; if they reach the SecureBus
|
||||
// directly it means the engine called Bus.Execute with an RLM payload.
|
||||
// Return a stub response — the RLMEngine interprets this.
|
||||
switch req.Type {
|
||||
case itr.CmdFinal:
|
||||
if f, ok := req.Payload.(itr.Final); ok {
|
||||
return itr.NewSuccessResponse(req.ID, f.Answer, 0)
|
||||
}
|
||||
}
|
||||
return itr.NewErrorResponse(req.ID, fmt.Sprintf("RLM command %q not handled at bus level", req.Type))
|
||||
}
|
||||
250
pkg/security/securebus/bus_test.go
Normal file
250
pkg/security/securebus/bus_test.go
Normal file
|
|
@ -0,0 +1,250 @@
|
|||
package securebus_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
"github.com/sipeed/picoclaw/pkg/security"
|
||||
"github.com/sipeed/picoclaw/pkg/security/securebus"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
// staticTool always returns a fixed result.
|
||||
type staticTool struct {
|
||||
name string
|
||||
result string
|
||||
isErr bool
|
||||
}
|
||||
|
||||
func (s *staticTool) Name() string { return s.name }
|
||||
func (s *staticTool) Description() string { return "static test tool" }
|
||||
func (s *staticTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}
|
||||
}
|
||||
func (s *staticTool) Execute(_ context.Context, _ map[string]interface{}) *tools.ToolResult {
|
||||
r := &tools.ToolResult{ForLLM: s.result, IsError: s.isErr}
|
||||
return r
|
||||
}
|
||||
|
||||
// echoTool returns the value of the "input" arg.
|
||||
type echoTool struct{}
|
||||
|
||||
func (e *echoTool) Name() string { return "echo" }
|
||||
func (e *echoTool) Description() string { return "echo" }
|
||||
func (e *echoTool) Parameters() map[string]interface{} {
|
||||
return map[string]interface{}{"type": "object", "properties": map[string]interface{}{}}
|
||||
}
|
||||
func (e *echoTool) Execute(_ context.Context, args map[string]interface{}) *tools.ToolResult {
|
||||
v, _ := args["input"].(string)
|
||||
return &tools.ToolResult{ForLLM: v}
|
||||
}
|
||||
|
||||
func makeArgsJSON(kv map[string]interface{}) string {
|
||||
if kv == nil {
|
||||
return "{}"
|
||||
}
|
||||
b, _ := json.Marshal(kv)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func makeBus(t *testing.T, toolMap map[string]tools.Tool, secrets *security.SecretStore) *securebus.Bus {
|
||||
t.Helper()
|
||||
capLookup := func(name string) (tools.ToolCapabilities, bool) {
|
||||
tool, ok := toolMap[name]
|
||||
if !ok {
|
||||
return tools.ZeroCapabilities(), false
|
||||
}
|
||||
return tools.ExtractCapabilities(tool), true
|
||||
}
|
||||
executor := func(ctx context.Context, name string, args map[string]interface{}) *tools.ToolResult {
|
||||
tool, ok := toolMap[name]
|
||||
if !ok {
|
||||
return &tools.ToolResult{ForLLM: "tool not found: " + name, IsError: true}
|
||||
}
|
||||
return tool.Execute(ctx, args)
|
||||
}
|
||||
cfg := securebus.DefaultBusConfig()
|
||||
return securebus.New(cfg, secrets, capLookup, executor)
|
||||
}
|
||||
|
||||
// ── tests ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
func TestBus_SuccessfulToolExec(t *testing.T) {
|
||||
tool := &staticTool{name: "greet", result: "hello world"}
|
||||
bus := makeBus(t, map[string]tools.Tool{"greet": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-1", "sess", "tc-1", "greet", makeArgsJSON(nil))
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.False(t, resp.IsError)
|
||||
assert.Equal(t, "hello world", resp.Result)
|
||||
assert.Equal(t, 1, bus.AuditLog().Len())
|
||||
}
|
||||
|
||||
func TestBus_UnknownTool(t *testing.T) {
|
||||
bus := makeBus(t, map[string]tools.Tool{}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-2", "sess", "tc-2", "nonexistent", makeArgsJSON(nil))
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.True(t, resp.IsError)
|
||||
}
|
||||
|
||||
func TestBus_ToolReturnsError(t *testing.T) {
|
||||
tool := &staticTool{name: "fail", result: "something broke", isErr: true}
|
||||
bus := makeBus(t, map[string]tools.Tool{"fail": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-3", "sess", "tc-3", "fail", makeArgsJSON(nil))
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.True(t, resp.IsError)
|
||||
assert.Equal(t, 1, bus.AuditLog().Len())
|
||||
events := bus.AuditLog().Events()
|
||||
assert.True(t, events[0].IsError)
|
||||
}
|
||||
|
||||
func TestBus_LeakDetection(t *testing.T) {
|
||||
// Tool output contains an API key — should be redacted.
|
||||
apiKey := "AKIAIOSFODNN7EXAMPLE" // fake AWS key matching redactor pattern
|
||||
tool := &staticTool{name: "leaky", result: "result: " + apiKey}
|
||||
bus := makeBus(t, map[string]tools.Tool{"leaky": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-4", "sess", "tc-4", "leaky", makeArgsJSON(nil))
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.True(t, resp.LeakDetected, "should detect API key in output")
|
||||
assert.NotContains(t, resp.Result, apiKey, "raw API key must not appear in response")
|
||||
|
||||
leakEvents := bus.AuditLog().LeakEvents()
|
||||
assert.Len(t, leakEvents, 1)
|
||||
}
|
||||
|
||||
func TestBus_SecretInjection_ArgVariant(t *testing.T) {
|
||||
// Tool reads injected "token" arg from args map.
|
||||
echoT := &echoTool{}
|
||||
|
||||
// Give echo tool a capability that declares a secret injected as arg:input.
|
||||
type capEchoTool struct {
|
||||
echoTool
|
||||
}
|
||||
capTool := &struct {
|
||||
echoTool
|
||||
}{}
|
||||
_ = capTool
|
||||
|
||||
// Use a capTool wrapper that adds arg injection capability.
|
||||
type wrappedEcho struct {
|
||||
*echoTool
|
||||
}
|
||||
toolMap := map[string]tools.Tool{"echo": echoT}
|
||||
|
||||
// Seed a secret store.
|
||||
key, _ := security.GenerateKey()
|
||||
keyring := security.NewNoopKeyring(key)
|
||||
ss, err := security.NewSecretStore(t.TempDir()+"/secrets.json", keyring)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, ss.Set("my_token", []byte("supersecret")))
|
||||
|
||||
// Override capLookup to inject via arg:input.
|
||||
capLookup := func(name string) (tools.ToolCapabilities, bool) {
|
||||
if name == "echo" {
|
||||
return tools.ToolCapabilities{
|
||||
Secrets: []tools.SecretRef{
|
||||
{Name: "my_token", InjectAs: "arg:input", Required: true},
|
||||
},
|
||||
}, true
|
||||
}
|
||||
return tools.ZeroCapabilities(), false
|
||||
}
|
||||
executor := func(ctx context.Context, name string, args map[string]interface{}) *tools.ToolResult {
|
||||
return toolMap[name].Execute(ctx, args)
|
||||
}
|
||||
|
||||
bus := securebus.New(securebus.DefaultBusConfig(), ss, capLookup, executor)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-5", "sess", "tc-5", "echo", makeArgsJSON(nil))
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.False(t, resp.IsError)
|
||||
assert.Equal(t, "supersecret", resp.Result, "injected secret should appear in tool output")
|
||||
|
||||
events := bus.AuditLog().Events()
|
||||
require.Len(t, events, 1)
|
||||
assert.Contains(t, events[0].SecretsAccessed, "my_token")
|
||||
}
|
||||
|
||||
func TestBus_PolicyViolation_RecursionDepth(t *testing.T) {
|
||||
tool := &staticTool{name: "ok", result: "fine"}
|
||||
bus := makeBus(t, map[string]tools.Tool{"ok": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-6", "sess", "tc-6", "ok", makeArgsJSON(nil))
|
||||
req.Depth = 255 // far exceeds MaxRecursionDepth=10
|
||||
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.True(t, resp.IsError, "depth violation should produce an error response")
|
||||
events := bus.AuditLog().Events()
|
||||
require.Len(t, events, 1)
|
||||
assert.NotEmpty(t, events[0].PolicyViolation)
|
||||
}
|
||||
|
||||
func TestBus_AuditLog_FilterBySession(t *testing.T) {
|
||||
tool := &staticTool{name: "t", result: "ok"}
|
||||
bus := makeBus(t, map[string]tools.Tool{"t": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
for _, sk := range []string{"session-A", "session-A", "session-B"} {
|
||||
req := itr.NewToolExecRequest("req-audit-"+sk, sk, "tc", "t", makeArgsJSON(nil))
|
||||
bus.Execute(context.Background(), req)
|
||||
}
|
||||
|
||||
assert.Equal(t, 3, bus.AuditLog().Len())
|
||||
assert.Len(t, bus.AuditLog().FilterBySession("session-A"), 2)
|
||||
assert.Len(t, bus.AuditLog().FilterBySession("session-B"), 1)
|
||||
}
|
||||
|
||||
func TestBus_Transport_Send(t *testing.T) {
|
||||
tool := &staticTool{name: "ping", result: "pong"}
|
||||
bus := makeBus(t, map[string]tools.Tool{"ping": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-tr", "sess", "tc", "ping", makeArgsJSON(nil))
|
||||
resp, err := bus.Transport().Send(context.Background(), req)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "pong", resp.Result)
|
||||
}
|
||||
|
||||
func TestBus_InvalidArgsJSON(t *testing.T) {
|
||||
tool := &staticTool{name: "ok", result: "ok"}
|
||||
bus := makeBus(t, map[string]tools.Tool{"ok": tool}, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewToolExecRequest("req-bad", "sess", "tc", "ok", "{invalid json")
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.True(t, resp.IsError)
|
||||
}
|
||||
|
||||
func TestBus_RLMFinalCommand(t *testing.T) {
|
||||
bus := makeBus(t, nil, nil)
|
||||
defer bus.Close()
|
||||
|
||||
req := itr.NewFinalRequest("req-final", "sess", 0, "the answer", "")
|
||||
resp := bus.Execute(context.Background(), req)
|
||||
|
||||
assert.False(t, resp.IsError)
|
||||
assert.Equal(t, "the answer", resp.Result)
|
||||
}
|
||||
141
pkg/security/securebus/policy.go
Normal file
141
pkg/security/securebus/policy.go
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
package securebus
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
// PolicyConfig controls global enforcement settings for the SecureBus.
|
||||
type PolicyConfig struct {
|
||||
// MaxRecursionDepth limits how deep RLM sub-calls may nest.
|
||||
// 0 disables the limit (not recommended).
|
||||
MaxRecursionDepth uint8
|
||||
|
||||
// MaxTokensPerRequest is a soft cap on cost_tokens for a single request.
|
||||
// 0 disables the cap.
|
||||
MaxTokensPerRequest uint32
|
||||
|
||||
// AllowedWorkspace is the filesystem root all PathRule patterns are
|
||||
// evaluated relative to. Empty string disables workspace restriction.
|
||||
AllowedWorkspace string
|
||||
|
||||
// SSRFBlockedPrefixes is a list of URL prefixes that are never permitted,
|
||||
// regardless of what a tool declares in its Network capabilities.
|
||||
// Applied in addition to the URL guard in pkg/security.
|
||||
SSRFBlockedPrefixes []string
|
||||
}
|
||||
|
||||
// DefaultPolicyConfig returns a secure-by-default configuration.
|
||||
func DefaultPolicyConfig() PolicyConfig {
|
||||
return PolicyConfig{
|
||||
MaxRecursionDepth: 10,
|
||||
MaxTokensPerRequest: 0, // no hard cap; individual budgets set per call
|
||||
SSRFBlockedPrefixes: []string{
|
||||
"http://169.254.", // AWS/GCP metadata
|
||||
"http://10.", // RFC 1918
|
||||
"http://172.16.", // RFC 1918
|
||||
"http://192.168.", // RFC 1918
|
||||
"http://localhost",
|
||||
"http://127.",
|
||||
"https://169.254.",
|
||||
"https://10.",
|
||||
"https://172.16.",
|
||||
"https://192.168.",
|
||||
"https://localhost",
|
||||
"https://127.",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// PolicyEngine validates requests against a policy config and tool capabilities.
|
||||
type PolicyEngine struct {
|
||||
cfg PolicyConfig
|
||||
}
|
||||
|
||||
// NewPolicyEngine creates a PolicyEngine with the given configuration.
|
||||
func NewPolicyEngine(cfg PolicyConfig) *PolicyEngine {
|
||||
return &PolicyEngine{cfg: cfg}
|
||||
}
|
||||
|
||||
// Validate returns an error if req violates any policy rule.
|
||||
func (pe *PolicyEngine) Validate(req itr.ToolRequest, caps tools.ToolCapabilities) error {
|
||||
if pe.cfg.MaxRecursionDepth > 0 && req.Depth > pe.cfg.MaxRecursionDepth {
|
||||
return fmt.Errorf("recursion depth %d exceeds limit %d", req.Depth, pe.cfg.MaxRecursionDepth)
|
||||
}
|
||||
|
||||
if te, ok := req.Payload.(itr.ToolExec); ok {
|
||||
_ = te // tool name validation happens in Bus.dispatch
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateNetwork checks whether targetURL is permitted given the tool's
|
||||
// EndpointRules and the global SSRF blocklist.
|
||||
func (pe *PolicyEngine) ValidateNetwork(targetURL string, rules []tools.EndpointRule) error {
|
||||
// Global SSRF check first — highest priority.
|
||||
for _, prefix := range pe.cfg.SSRFBlockedPrefixes {
|
||||
if strings.HasPrefix(targetURL, prefix) {
|
||||
return fmt.Errorf("network access denied: URL %q matches SSRF blocklist", targetURL)
|
||||
}
|
||||
}
|
||||
|
||||
// If no rules declared, deny all network access.
|
||||
if len(rules) == 0 {
|
||||
return fmt.Errorf("network access denied: tool declares no EndpointRules")
|
||||
}
|
||||
|
||||
// Must match at least one permit rule.
|
||||
for _, rule := range rules {
|
||||
matched, _ := filepath.Match(rule.Pattern, targetURL)
|
||||
if matched {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("network access denied: URL %q does not match any permitted endpoint pattern", targetURL)
|
||||
}
|
||||
|
||||
// ValidateFilesystem checks whether targetPath is permitted given PathRules.
|
||||
// targetPath should be an absolute path; it is compared against patterns
|
||||
// rooted at AllowedWorkspace.
|
||||
func (pe *PolicyEngine) ValidateFilesystem(targetPath, mode string, rules []tools.PathRule) error {
|
||||
if len(rules) == 0 {
|
||||
return fmt.Errorf("filesystem access denied: tool declares no PathRules")
|
||||
}
|
||||
|
||||
// Normalise path relative to workspace if configured.
|
||||
checkPath := targetPath
|
||||
if pe.cfg.AllowedWorkspace != "" {
|
||||
rel, err := filepath.Rel(pe.cfg.AllowedWorkspace, targetPath)
|
||||
if err == nil && !strings.HasPrefix(rel, "..") {
|
||||
checkPath = rel
|
||||
}
|
||||
}
|
||||
|
||||
for _, rule := range rules {
|
||||
matched, _ := filepath.Match(rule.Pattern, checkPath)
|
||||
if !matched {
|
||||
continue
|
||||
}
|
||||
// Check mode compatibility.
|
||||
switch mode {
|
||||
case "r":
|
||||
if strings.Contains(rule.Mode, "r") {
|
||||
return nil
|
||||
}
|
||||
case "w":
|
||||
if strings.Contains(rule.Mode, "w") {
|
||||
return nil
|
||||
}
|
||||
case "rw":
|
||||
if rule.Mode == "rw" {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
}
|
||||
return fmt.Errorf("filesystem access denied: %q (mode %q) does not match any permitted PathRule", targetPath, mode)
|
||||
}
|
||||
106
pkg/security/securebus/transport.go
Normal file
106
pkg/security/securebus/transport.go
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
// Package securebus implements the SecureBus — the privilege boundary between
|
||||
// the agent loop and tool execution. All tool calls pass through the SecureBus
|
||||
// which enforces capability policies, injects secrets, scans output for leaks,
|
||||
// and writes an append-only audit log.
|
||||
//
|
||||
// Transport abstracts the communication channel. The same SecureBus code runs
|
||||
// on all backends:
|
||||
// - ChannelTransport: in-process Go channels (default, zero overhead)
|
||||
// - SocketTransport: Unix domain socket (daemon mode, Layer 4)
|
||||
// - WasmTransport: wazero isolate (untrusted tools, Layer 5)
|
||||
package securebus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/itr"
|
||||
)
|
||||
|
||||
// Transport abstracts the communication channel between the agent loop (caller)
|
||||
// and the SecureBus worker (executor). Implementations must be safe for
|
||||
// concurrent use by multiple goroutines.
|
||||
type Transport interface {
|
||||
// Send submits a request and blocks until the response is available or ctx
|
||||
// is cancelled.
|
||||
Send(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error)
|
||||
|
||||
// Close shuts down the transport and releases any resources.
|
||||
Close() error
|
||||
}
|
||||
|
||||
// ChannelTransport implements Transport using a pair of buffered Go channels.
|
||||
// It runs the SecureBus in the same goroutine that handles the request, making
|
||||
// it suitable for in-process use with near-zero overhead.
|
||||
type ChannelTransport struct {
|
||||
reqCh chan channelEnvelope
|
||||
closed chan struct{}
|
||||
}
|
||||
|
||||
type channelEnvelope struct {
|
||||
req itr.ToolRequest
|
||||
respCh chan channelResult
|
||||
}
|
||||
|
||||
type channelResult struct {
|
||||
resp itr.ToolResponse
|
||||
err error
|
||||
}
|
||||
|
||||
// NewChannelTransport creates a ChannelTransport with the given buffer depth.
|
||||
// depth 0 makes Send synchronous (no buffering). Typical value: 16.
|
||||
func NewChannelTransport(depth int) *ChannelTransport {
|
||||
if depth < 0 {
|
||||
depth = 0
|
||||
}
|
||||
ct := &ChannelTransport{
|
||||
reqCh: make(chan channelEnvelope, depth),
|
||||
closed: make(chan struct{}),
|
||||
}
|
||||
return ct
|
||||
}
|
||||
|
||||
// Requests returns the channel that the SecureBus worker should read from.
|
||||
// This is the server-side handle — the SecureBus calls this once to start
|
||||
// processing.
|
||||
func (ct *ChannelTransport) Requests() <-chan channelEnvelope {
|
||||
return ct.reqCh
|
||||
}
|
||||
|
||||
// Send enqueues req and blocks until the response arrives or ctx is cancelled.
|
||||
func (ct *ChannelTransport) Send(ctx context.Context, req itr.ToolRequest) (itr.ToolResponse, error) {
|
||||
respCh := make(chan channelResult, 1)
|
||||
env := channelEnvelope{req: req, respCh: respCh}
|
||||
|
||||
select {
|
||||
case ct.reqCh <- env:
|
||||
case <-ct.closed:
|
||||
return itr.ToolResponse{}, fmt.Errorf("transport closed")
|
||||
case <-ctx.Done():
|
||||
return itr.ToolResponse{}, ctx.Err()
|
||||
}
|
||||
|
||||
select {
|
||||
case result := <-respCh:
|
||||
return result.resp, result.err
|
||||
case <-ct.closed:
|
||||
return itr.ToolResponse{}, fmt.Errorf("transport closed while waiting for response")
|
||||
case <-ctx.Done():
|
||||
return itr.ToolResponse{}, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// Close drains the request channel and signals that no new requests will be accepted.
|
||||
func (ct *ChannelTransport) Close() error {
|
||||
select {
|
||||
case <-ct.closed:
|
||||
default:
|
||||
close(ct.closed)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// reply sends a response back to the caller of Send.
|
||||
func (env channelEnvelope) reply(resp itr.ToolResponse, err error) {
|
||||
env.respCh <- channelResult{resp: resp, err: err}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue