feat(extensions): phase 3 — event subscriptions + tool-call interception
Two new capabilities, both ride on the existing subprocess
protocol with a couple of new frame types.
Event subscriptions (one-way notifications):
ext -> host: subscribe {events: [...], intercept: [...]}
host -> ext: event {event, ...payload}
Recognised events: session_start, turn_start, turn_end,
tool_call, assistant_message. Subscribers get fire-and-forget
notifications on each. Useful for telemetry, audit logs, custom
state widgets that follow live agent activity.
Tool-call interception (round-trip, can refuse):
host -> ext: event_intercept {id, event:"tool_call", tool_name, tool_args}
ext -> host: event_intercept_response {id, block?, reason?}
When at least one extension subscribed to "tool_call" intercept,
zot asks each one in turn before running every tool call. First
blocker wins; reason becomes the tool-result error text the model
sees. Per-extension 5s timeout treats unresponsive interceptors
as "allow" so a wedged extension never stalls the agent.
Wire format additions (internal/extproto):
ext -> host: SubscribeFromExt, EventInterceptResponseFromExt
host -> ext: EventFromHost, EventInterceptFromHost
Manager (internal/agent/extensions):
- per-extension eventSubs / interceptSubs sets, populated by the
subscribe frame
- EmitEvent fans out to every subscribed extension on its own
goroutine (won't block the agent on slow stdin writes)
- InterceptToolCall walks subscribers serially, returning the
first refusal; 5s timeout per subscriber (allow on timeout)
- readLoop handles event_intercept_response correlations the
same way it handles command/tool responses
Core (internal/core/agent.go):
- Agent.BeforeToolExecute hook called from runOneTool right
before tool.Execute. Returning (allowed=false, reason)
short-circuits with an IsError tool result containing reason.
- Agent.OnEvent observer fires for every emitted AgentEvent;
composed transparently with the per-Prompt sink via wrapSink
so neither the existing TUI nor the rpc loop need changes.
Wiring (internal/agent/cli.go, rpc.go):
- wireAgentExt sets BeforeToolExecute -> InterceptToolCall and
OnEvent -> fanoutAgentEvent for every freshly-built agent
(initial, login rebuild, model swap)
- fanoutAgentEvent translates core AgentEvent kinds into
extproto.EventFromHost. Internal-only events (text_delta,
tool_progress) are dropped to keep the per-extension stream
sane.
- session_start emitted once after extensions come up
SDK (pkg/zotext):
- On(name, EventHandler) registers per-event observers
- InterceptToolCall(InterceptHandler) registers a single
intercept callback
- Run() now also sends a subscribe frame before the ready
sentinel, with the union of subscribed events + intercept
- Frame loop handles "event" and "event_intercept" frames,
runs the handlers (intercepts on a goroutine to avoid
head-of-line blocking)
- Capabilities advertised: commands + tools + events
Example (examples/extensions/guard):
- subscribes to session_start / turn_start / tool_call / turn_end
and writes one-line audit entries
- intercepts every bash call; refuses commands matching
rm -rf, sudo, dd of=/, mkfs, the fork bomb, chmod -R 777
- end-to-end verified live: agent -> bash("rm -rf /tmp/foo")
-> guard refuses -> model sees the refusal text and surfaces
it in its reply ("the guard blocked it, as expected — the
pattern \brm\s+-rf\b matched")
Docs/extensions.md updated with all five new frame types and the
guard example.- commit
- 83ae236571c98405a5a59c630807b870400972fd
- author
- patriceckhart · · 2026-04-19 12:57 UTC
- parent
- 74709a0
- signature
- unsigned
- refs
- tag: v0.0.24
11 files · +709 -36
M docs/extensions.md +67 -8
| @@ -6,13 +6,12 @@ its stdin/stdout. Extensions can be written in **any language** that | |||
| 6 | 6 | can read and write JSON lines from stdio — Go, TypeScript, Python, | |
| 7 | 7 | Rust, shell with `jq`, anything. | |
| 8 | 8 | ||
| 9 | - | Two phases shipped so far: | |
| 9 | + | Three phases shipped so far: | |
| 10 | 10 | ||
| 11 | 11 | - **Phase 1**: slash commands + chat notifications. | |
| 12 | 12 | - **Phase 2**: tools the LLM can call. | |
| 13 | - | ||
| 14 | - | Phase 3 (lifecycle event subscriptions + tool-call interception for | |
| 15 | - | guardrail extensions) is on the roadmap below. | |
| 13 | + | - **Phase 3**: lifecycle event subscriptions + tool-call interception | |
| 14 | + | for guardrail extensions. | |
| 16 | 15 | ||
| 17 | 16 | ## Quick start | |
| 18 | 17 | ||
| @@ -186,6 +185,33 @@ message blocks; each block is `{"type":"text","text":"..."}` or | |||
| 186 | 185 | "content":[{"type":"text","text":"Berlin: 16°C, fog"}]} | |
| 187 | 186 | ``` | |
| 188 | 187 | ||
| 188 | + | #### `subscribe` | |
| 189 | + | ||
| 190 | + | Declares which lifecycle events the extension wants to observe and | |
| 191 | + | which it wants to intercept. Send once after `hello`, before `ready`. | |
| 192 | + | ||
| 193 | + | ```json | |
| 194 | + | {"type":"subscribe", | |
| 195 | + | "events":["session_start","turn_start","tool_call","turn_end","assistant_message"], | |
| 196 | + | "intercept":["tool_call"]} | |
| 197 | + | ``` | |
| 198 | + | ||
| 199 | + | Recognised event names: `session_start`, `turn_start`, `turn_end`, | |
| 200 | + | `tool_call`, `assistant_message`. Only `tool_call` is interceptable | |
| 201 | + | in this version; other names listed under `intercept` are ignored. | |
| 202 | + | ||
| 203 | + | #### `event_intercept_response` | |
| 204 | + | ||
| 205 | + | Reply to an `event_intercept` from the host. `block: true` refuses | |
| 206 | + | the action; `reason` is shown to the model as the tool error text. | |
| 207 | + | Missing the response within 5s is treated as "allow" (i.e. an | |
| 208 | + | unresponsive extension never stalls the agent). | |
| 209 | + | ||
| 210 | + | ```json | |
| 211 | + | {"type":"event_intercept_response","id":"...", | |
| 212 | + | "block":true,"reason":"refused: matches danger pattern \"rm -rf\""} | |
| 213 | + | ``` | |
| 214 | + | ||
| 189 | 215 | #### `command_response` (reply to `command_invoked`) | |
| 190 | 216 | ||
| 191 | 217 | ```json | |
| @@ -260,6 +286,33 @@ Reply with `tool_result` within the host's tool timeout (default 60s). | |||
| 260 | 286 | Missing the timeout surfaces an error to the model and the call is | |
| 261 | 287 | marked as failed. | |
| 262 | 288 | ||
| 289 | + | #### `event` | |
| 290 | + | ||
| 291 | + | Lifecycle notification for events the extension subscribed to via | |
| 292 | + | `subscribe`. One-way — no response expected. | |
| 293 | + | ||
| 294 | + | ```json | |
| 295 | + | {"type":"event","event":"turn_start","step":1} | |
| 296 | + | {"type":"event","event":"tool_call", | |
| 297 | + | "tool_id":"...","tool_name":"read","tool_args":{"path":"foo.go"}} | |
| 298 | + | {"type":"event","event":"turn_end","stop":"end_turn"} | |
| 299 | + | ``` | |
| 300 | + | ||
| 301 | + | #### `event_intercept` | |
| 302 | + | ||
| 303 | + | Sent when zot wants to give the extension a chance to block a | |
| 304 | + | lifecycle event before it happens. Same payload shape as `event`. | |
| 305 | + | Reply with `event_intercept_response` within 5s; missing the deadline | |
| 306 | + | is treated as "allow". | |
| 307 | + | ||
| 308 | + | Only `event: "tool_call"` is sent in this version. | |
| 309 | + | ||
| 310 | + | ```json | |
| 311 | + | {"type":"event_intercept","id":"...","event":"tool_call", | |
| 312 | + | "tool_id":"...","tool_name":"bash", | |
| 313 | + | "tool_args":{"command":"rm -rf /tmp/foo"}} | |
| 314 | + | ``` | |
| 315 | + | ||
| 263 | 316 | #### `shutdown` | |
| 264 | 317 | ||
| 265 | 318 | Sent during graceful zot exit (or `/reload-ext` once that lands). | |
| @@ -323,6 +376,8 @@ See: | |||
| 323 | 376 | - `examples/extensions/hello/` — slash commands | |
| 324 | 377 | - `examples/extensions/clock/` — slash commands in plain Node, no SDK | |
| 325 | 378 | - `examples/extensions/weather/` — LLM-callable tool | |
| 379 | + | - `examples/extensions/guard/` — event subscriptions + tool-call | |
| 380 | + | interception (refuses dangerous bash patterns) | |
| 326 | 381 | ||
| 327 | 382 | ### TypeScript / Python | |
| 328 | 383 | ||
| @@ -355,8 +410,12 @@ Phase 2 (shipped): | |||
| 355 | 410 | - [x] `ready` sentinel for safe agent-registry build timing | |
| 356 | 411 | - [x] tool result attribution surfaces extension name in details | |
| 357 | 412 | ||
| 358 | - | Phase 3: | |
| 359 | - | - [ ] event subscriptions (`turn_start`, `turn_end`, `tool_call_*`, | |
| 360 | - | `session_start`, etc.) | |
| 361 | - | - [ ] tool-call interception (block / modify before execution) | |
| 413 | + | Phase 3 (shipped): | |
| 414 | + | - [x] event subscriptions (`session_start`, `turn_start`, `turn_end`, | |
| 415 | + | `tool_call`, `assistant_message`) | |
| 416 | + | - [x] tool-call interception (block before execution) | |
| 417 | + | ||
| 418 | + | Future (no firm timeline): | |
| 419 | + | - [ ] interception for additional events beyond `tool_call` | |
| 420 | + | - [ ] modify (not just block) tool args mid-flight | |
| 362 | 421 | - [ ] `/reload-ext` slash command (hot-reload without restarting zot) | |
| @@ -6,13 +6,12 @@ its stdin/stdout. Extensions can be written in **any language** that | |||||
| 6 | can read and write JSON lines from stdio — Go, TypeScript, Python, | 6 | can read and write JSON lines from stdio — Go, TypeScript, Python, | ||
| 7 | Rust, shell with `jq`, anything. | 7 | Rust, shell with `jq`, anything. | ||
| 8 | 8 | ||||
| 9 | - | Two phases shipped so far: | 9 | + | Three phases shipped so far: |
| 10 | 10 | ||||
| 11 | - **Phase 1**: slash commands + chat notifications. | 11 | - **Phase 1**: slash commands + chat notifications. | ||
| 12 | - **Phase 2**: tools the LLM can call. | 12 | - **Phase 2**: tools the LLM can call. | ||
| 13 | - | 13 | + | - **Phase 3**: lifecycle event subscriptions + tool-call interception | |
| 14 | - | Phase 3 (lifecycle event subscriptions + tool-call interception for | 14 | + | for guardrail extensions. |
| 15 | - | guardrail extensions) is on the roadmap below. | |||
| 16 | 15 | ||||
| 17 | ## Quick start | 16 | ## Quick start | ||
| 18 | 17 | ||||
| @@ -186,6 +185,33 @@ message blocks; each block is `{"type":"text","text":"..."}` or | |||||
| 186 | "content":[{"type":"text","text":"Berlin: 16°C, fog"}]} | 185 | "content":[{"type":"text","text":"Berlin: 16°C, fog"}]} | ||
| 187 | ``` | 186 | ``` | ||
| 188 | 187 | ||||
| 188 | + | #### `subscribe` | |||
| 189 | + | ||||
| 190 | + | Declares which lifecycle events the extension wants to observe and | |||
| 191 | + | which it wants to intercept. Send once after `hello`, before `ready`. | |||
| 192 | + | ||||
| 193 | + | ```json | |||
| 194 | + | {"type":"subscribe", | |||
| 195 | + | "events":["session_start","turn_start","tool_call","turn_end","assistant_message"], | |||
| 196 | + | "intercept":["tool_call"]} | |||
| 197 | + | ``` | |||
| 198 | + | ||||
| 199 | + | Recognised event names: `session_start`, `turn_start`, `turn_end`, | |||
| 200 | + | `tool_call`, `assistant_message`. Only `tool_call` is interceptable | |||
| 201 | + | in this version; other names listed under `intercept` are ignored. | |||
| 202 | + | ||||
| 203 | + | #### `event_intercept_response` | |||
| 204 | + | ||||
| 205 | + | Reply to an `event_intercept` from the host. `block: true` refuses | |||
| 206 | + | the action; `reason` is shown to the model as the tool error text. | |||
| 207 | + | Missing the response within 5s is treated as "allow" (i.e. an | |||
| 208 | + | unresponsive extension never stalls the agent). | |||
| 209 | + | ||||
| 210 | + | ```json | |||
| 211 | + | {"type":"event_intercept_response","id":"...", | |||
| 212 | + | "block":true,"reason":"refused: matches danger pattern \"rm -rf\""} | |||
| 213 | + | ``` | |||
| 214 | + | ||||
| 189 | #### `command_response` (reply to `command_invoked`) | 215 | #### `command_response` (reply to `command_invoked`) | ||
| 190 | 216 | ||||
| 191 | ```json | 217 | ```json | ||
| @@ -260,6 +286,33 @@ Reply with `tool_result` within the host's tool timeout (default 60s). | |||||
| 260 | Missing the timeout surfaces an error to the model and the call is | 286 | Missing the timeout surfaces an error to the model and the call is | ||
| 261 | marked as failed. | 287 | marked as failed. | ||
| 262 | 288 | ||||
| 289 | + | #### `event` | |||
| 290 | + | ||||
| 291 | + | Lifecycle notification for events the extension subscribed to via | |||
| 292 | + | `subscribe`. One-way — no response expected. | |||
| 293 | + | ||||
| 294 | + | ```json | |||
| 295 | + | {"type":"event","event":"turn_start","step":1} | |||
| 296 | + | {"type":"event","event":"tool_call", | |||
| 297 | + | "tool_id":"...","tool_name":"read","tool_args":{"path":"foo.go"}} | |||
| 298 | + | {"type":"event","event":"turn_end","stop":"end_turn"} | |||
| 299 | + | ``` | |||
| 300 | + | ||||
| 301 | + | #### `event_intercept` | |||
| 302 | + | ||||
| 303 | + | Sent when zot wants to give the extension a chance to block a | |||
| 304 | + | lifecycle event before it happens. Same payload shape as `event`. | |||
| 305 | + | Reply with `event_intercept_response` within 5s; missing the deadline | |||
| 306 | + | is treated as "allow". | |||
| 307 | + | ||||
| 308 | + | Only `event: "tool_call"` is sent in this version. | |||
| 309 | + | ||||
| 310 | + | ```json | |||
| 311 | + | {"type":"event_intercept","id":"...","event":"tool_call", | |||
| 312 | + | "tool_id":"...","tool_name":"bash", | |||
| 313 | + | "tool_args":{"command":"rm -rf /tmp/foo"}} | |||
| 314 | + | ``` | |||
| 315 | + | ||||
| 263 | #### `shutdown` | 316 | #### `shutdown` | ||
| 264 | 317 | ||||
| 265 | Sent during graceful zot exit (or `/reload-ext` once that lands). | 318 | Sent during graceful zot exit (or `/reload-ext` once that lands). | ||
| @@ -323,6 +376,8 @@ See: | |||||
| 323 | - `examples/extensions/hello/` — slash commands | 376 | - `examples/extensions/hello/` — slash commands | ||
| 324 | - `examples/extensions/clock/` — slash commands in plain Node, no SDK | 377 | - `examples/extensions/clock/` — slash commands in plain Node, no SDK | ||
| 325 | - `examples/extensions/weather/` — LLM-callable tool | 378 | - `examples/extensions/weather/` — LLM-callable tool | ||
| 379 | + | - `examples/extensions/guard/` — event subscriptions + tool-call | |||
| 380 | + | interception (refuses dangerous bash patterns) | |||
| 326 | 381 | ||||
| 327 | ### TypeScript / Python | 382 | ### TypeScript / Python | ||
| 328 | 383 | ||||
| @@ -355,8 +410,12 @@ Phase 2 (shipped): | |||||
| 355 | - [x] `ready` sentinel for safe agent-registry build timing | 410 | - [x] `ready` sentinel for safe agent-registry build timing | ||
| 356 | - [x] tool result attribution surfaces extension name in details | 411 | - [x] tool result attribution surfaces extension name in details | ||
| 357 | 412 | ||||
| 358 | - | Phase 3: | 413 | + | Phase 3 (shipped): |
| 359 | - | - [ ] event subscriptions (`turn_start`, `turn_end`, `tool_call_*`, | 414 | + | - [x] event subscriptions (`session_start`, `turn_start`, `turn_end`, |
| 360 | - | `session_start`, etc.) | 415 | + | `tool_call`, `assistant_message`) |
| 361 | - | - [ ] tool-call interception (block / modify before execution) | 416 | + | - [x] tool-call interception (block before execution) |
| 417 | + | ||||
| 418 | + | Future (no firm timeline): | |||
| 419 | + | - [ ] interception for additional events beyond `tool_call` | |||
| 420 | + | - [ ] modify (not just block) tool args mid-flight | |||
| 362 | - [ ] `/reload-ext` slash command (hot-reload without restarting zot) | 421 | - [ ] `/reload-ext` slash command (hot-reload without restarting zot) | ||
A examples/extensions/guard/README.md +58 -0
| @@ -0,0 +1,58 @@ | |||
| 1 | + | # guard — example zot extension (Go, phase 3) | |
| 2 | + | ||
| 3 | + | Demonstrates the **event subscription** and **tool-call interception** | |
| 4 | + | half of the extension protocol (phase 3). | |
| 5 | + | ||
| 6 | + | What it does: | |
| 7 | + | ||
| 8 | + | - Subscribes to `session_start`, `turn_start`, `tool_call`, `turn_end` | |
| 9 | + | and appends a line to `/tmp/zot-guard-audit.log` for each. | |
| 10 | + | - Intercepts every `bash` tool call. If the command matches a danger | |
| 11 | + | regex (`rm -rf`, `sudo`, `dd of=/`, `mkfs`, the fork bomb, `chmod -R | |
| 12 | + | 777`), the call is **refused**. The model sees the refusal as the | |
| 13 | + | tool error and (typically) proposes something safer or asks for | |
| 14 | + | confirmation. | |
| 15 | + | ||
| 16 | + | ## Build | |
| 17 | + | ||
| 18 | + | ```bash | |
| 19 | + | cd examples/extensions/guard | |
| 20 | + | go build -o guard . | |
| 21 | + | ``` | |
| 22 | + | ||
| 23 | + | ## Install | |
| 24 | + | ||
| 25 | + | ```bash | |
| 26 | + | zot ext install . | |
| 27 | + | ``` | |
| 28 | + | ||
| 29 | + | ## Try it | |
| 30 | + | ||
| 31 | + | In zot, ask: | |
| 32 | + | ||
| 33 | + | > Run `rm -rf /tmp/foo` | |
| 34 | + | ||
| 35 | + | The model's `bash` call is intercepted and refused; the model | |
| 36 | + | explains the refusal in its reply. No file is touched. | |
| 37 | + | ||
| 38 | + | > Run `ls /tmp` | |
| 39 | + | ||
| 40 | + | Allowed; the audit log records the call. | |
| 41 | + | ||
| 42 | + | Tail the audit log: | |
| 43 | + | ||
| 44 | + | ```bash | |
| 45 | + | tail -f /tmp/zot-guard-audit.log | |
| 46 | + | ``` | |
| 47 | + | ||
| 48 | + | ## Extending the danger list | |
| 49 | + | ||
| 50 | + | Edit `dangerPatterns` in `main.go`. Each entry is a Go regexp; the | |
| 51 | + | match is case-insensitive. Rebuild and reinstall. | |
| 52 | + | ||
| 53 | + | ## See also | |
| 54 | + | ||
| 55 | + | - `examples/extensions/hello` — slash command (phase 1) | |
| 56 | + | - `examples/extensions/clock` — slash command in plain Node (phase 1) | |
| 57 | + | - `examples/extensions/weather` — LLM-callable tool (phase 2) | |
| 58 | + | - `docs/extensions.md` — full protocol reference | |
| @@ -0,0 +1,58 @@ | |||||
| 1 | + | # guard — example zot extension (Go, phase 3) | |||
| 2 | + | ||||
| 3 | + | Demonstrates the **event subscription** and **tool-call interception** | |||
| 4 | + | half of the extension protocol (phase 3). | |||
| 5 | + | ||||
| 6 | + | What it does: | |||
| 7 | + | ||||
| 8 | + | - Subscribes to `session_start`, `turn_start`, `tool_call`, `turn_end` | |||
| 9 | + | and appends a line to `/tmp/zot-guard-audit.log` for each. | |||
| 10 | + | - Intercepts every `bash` tool call. If the command matches a danger | |||
| 11 | + | regex (`rm -rf`, `sudo`, `dd of=/`, `mkfs`, the fork bomb, `chmod -R | |||
| 12 | + | 777`), the call is **refused**. The model sees the refusal as the | |||
| 13 | + | tool error and (typically) proposes something safer or asks for | |||
| 14 | + | confirmation. | |||
| 15 | + | ||||
| 16 | + | ## Build | |||
| 17 | + | ||||
| 18 | + | ```bash | |||
| 19 | + | cd examples/extensions/guard | |||
| 20 | + | go build -o guard . | |||
| 21 | + | ``` | |||
| 22 | + | ||||
| 23 | + | ## Install | |||
| 24 | + | ||||
| 25 | + | ```bash | |||
| 26 | + | zot ext install . | |||
| 27 | + | ``` | |||
| 28 | + | ||||
| 29 | + | ## Try it | |||
| 30 | + | ||||
| 31 | + | In zot, ask: | |||
| 32 | + | ||||
| 33 | + | > Run `rm -rf /tmp/foo` | |||
| 34 | + | ||||
| 35 | + | The model's `bash` call is intercepted and refused; the model | |||
| 36 | + | explains the refusal in its reply. No file is touched. | |||
| 37 | + | ||||
| 38 | + | > Run `ls /tmp` | |||
| 39 | + | ||||
| 40 | + | Allowed; the audit log records the call. | |||
| 41 | + | ||||
| 42 | + | Tail the audit log: | |||
| 43 | + | ||||
| 44 | + | ```bash | |||
| 45 | + | tail -f /tmp/zot-guard-audit.log | |||
| 46 | + | ``` | |||
| 47 | + | ||||
| 48 | + | ## Extending the danger list | |||
| 49 | + | ||||
| 50 | + | Edit `dangerPatterns` in `main.go`. Each entry is a Go regexp; the | |||
| 51 | + | match is case-insensitive. Rebuild and reinstall. | |||
| 52 | + | ||||
| 53 | + | ## See also | |||
| 54 | + | ||||
| 55 | + | - `examples/extensions/hello` — slash command (phase 1) | |||
| 56 | + | - `examples/extensions/clock` — slash command in plain Node (phase 1) | |||
| 57 | + | - `examples/extensions/weather` — LLM-callable tool (phase 2) | |||
| 58 | + | - `docs/extensions.md` — full protocol reference | |||
A examples/extensions/guard/extension.json +8 -0
| @@ -0,0 +1,8 @@ | |||
| 1 | + | { | |
| 2 | + | "name": "guard", | |
| 3 | + | "version": "1.0.0", | |
| 4 | + | "exec": "./guard", | |
| 5 | + | "language": "go", | |
| 6 | + | "description": "audits lifecycle events and refuses dangerous bash patterns", | |
| 7 | + | "enabled": true | |
| 8 | + | } | |
| @@ -0,0 +1,8 @@ | |||||
| 1 | + | { | |||
| 2 | + | "name": "guard", | |||
| 3 | + | "version": "1.0.0", | |||
| 4 | + | "exec": "./guard", | |||
| 5 | + | "language": "go", | |||
| 6 | + | "description": "audits lifecycle events and refuses dangerous bash patterns", | |||
| 7 | + | "enabled": true | |||
| 8 | + | } | |||
A examples/extensions/guard/main.go +91 -0
| @@ -0,0 +1,91 @@ | |||
| 1 | + | // guard — phase 3 example: event subscriptions + tool-call interception. | |
| 2 | + | // | |
| 3 | + | // What it does: | |
| 4 | + | // | |
| 5 | + | // - Subscribes to turn_start, tool_call, and turn_end events and | |
| 6 | + | // writes a one-line audit entry per event to its own log file. | |
| 7 | + | // | |
| 8 | + | // - Intercepts every tool call. If the tool is `bash` and the | |
| 9 | + | // command matches a danger pattern (rm -rf, sudo, dd of=/, etc.) | |
| 10 | + | // the call is refused and the model gets a one-line refusal it | |
| 11 | + | // can react to. Everything else passes through. | |
| 12 | + | // | |
| 13 | + | // Build it: | |
| 14 | + | // | |
| 15 | + | // cd examples/extensions/guard | |
| 16 | + | // go build -o guard . | |
| 17 | + | // | |
| 18 | + | // Drop it next to its extension.json under $ZOT_HOME/extensions/guard/ | |
| 19 | + | // (or `zot ext install ./guard` from this directory) and the next zot | |
| 20 | + | // session will load it automatically. | |
| 21 | + | package main | |
| 22 | + | ||
| 23 | + | import ( | |
| 24 | + | "encoding/json" | |
| 25 | + | "fmt" | |
| 26 | + | "os" | |
| 27 | + | "path/filepath" | |
| 28 | + | "regexp" | |
| 29 | + | "strings" | |
| 30 | + | "time" | |
| 31 | + | ||
| 32 | + | "github.com/patriceckhart/zot/pkg/zotext" | |
| 33 | + | ) | |
| 34 | + | ||
| 35 | + | var dangerPatterns = []*regexp.Regexp{ | |
| 36 | + | regexp.MustCompile(`(?i)\brm\s+-rf\b`), | |
| 37 | + | regexp.MustCompile(`(?i)\bsudo\b`), | |
| 38 | + | regexp.MustCompile(`(?i)\bdd\s+of=/`), | |
| 39 | + | regexp.MustCompile(`(?i)\bmkfs\b`), | |
| 40 | + | regexp.MustCompile(`(?i)\b:\s*\(\s*\)\s*\{\s*:\|:`), // fork bomb | |
| 41 | + | regexp.MustCompile(`(?i)\bchmod\s+-R\s+777\b`), | |
| 42 | + | } | |
| 43 | + | ||
| 44 | + | func main() { | |
| 45 | + | ext := zotext.New("guard", "1.0.0") | |
| 46 | + | ||
| 47 | + | auditPath := filepath.Join(os.TempDir(), "zot-guard-audit.log") | |
| 48 | + | auditFile, _ := os.OpenFile(auditPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) | |
| 49 | + | if auditFile != nil { | |
| 50 | + | defer auditFile.Close() | |
| 51 | + | } | |
| 52 | + | audit := func(format string, args ...any) { | |
| 53 | + | line := fmt.Sprintf("%s "+format+"\n", append([]any{time.Now().Format(time.RFC3339)}, args...)...) | |
| 54 | + | if auditFile != nil { | |
| 55 | + | auditFile.WriteString(line) | |
| 56 | + | } | |
| 57 | + | ext.Logf("%s", strings.TrimRight(line, "\n")) | |
| 58 | + | } | |
| 59 | + | ||
| 60 | + | // Lifecycle observers. | |
| 61 | + | ext.On("session_start", func(ev zotext.Event) { audit("[guard] session_start") }) | |
| 62 | + | ext.On("turn_start", func(ev zotext.Event) { audit("[guard] turn_start step=%d", ev.Step) }) | |
| 63 | + | ext.On("tool_call", func(ev zotext.Event) { | |
| 64 | + | audit("[guard] tool_call %s args=%s", ev.ToolName, string(ev.ToolArgs)) | |
| 65 | + | }) | |
| 66 | + | ext.On("turn_end", func(ev zotext.Event) { audit("[guard] turn_end stop=%s err=%s", ev.Stop, ev.Error) }) | |
| 67 | + | ||
| 68 | + | // Tool-call gate. | |
| 69 | + | ext.InterceptToolCall(func(toolName string, args json.RawMessage) (bool, string) { | |
| 70 | + | if toolName != "bash" { | |
| 71 | + | return true, "" // only guard bash | |
| 72 | + | } | |
| 73 | + | var in struct { | |
| 74 | + | Command string `json:"command"` | |
| 75 | + | } | |
| 76 | + | if err := json.Unmarshal(args, &in); err != nil { | |
| 77 | + | return true, "" | |
| 78 | + | } | |
| 79 | + | for _, pat := range dangerPatterns { | |
| 80 | + | if pat.MatchString(in.Command) { | |
| 81 | + | audit("[guard] BLOCKED bash: %s (matched %s)", in.Command, pat.String()) | |
| 82 | + | return false, fmt.Sprintf("guard refused: command matches the danger pattern %q. Ask the user before running this.", pat.String()) | |
| 83 | + | } | |
| 84 | + | } | |
| 85 | + | return true, "" | |
| 86 | + | }) | |
| 87 | + | ||
| 88 | + | if err := ext.Run(); err != nil { | |
| 89 | + | ext.Logf("fatal: %v", err) | |
| 90 | + | } | |
| 91 | + | } | |
| @@ -0,0 +1,91 @@ | |||||
| 1 | + | // guard — phase 3 example: event subscriptions + tool-call interception. | |||
| 2 | + | // | |||
| 3 | + | // What it does: | |||
| 4 | + | // | |||
| 5 | + | // - Subscribes to turn_start, tool_call, and turn_end events and | |||
| 6 | + | // writes a one-line audit entry per event to its own log file. | |||
| 7 | + | // | |||
| 8 | + | // - Intercepts every tool call. If the tool is `bash` and the | |||
| 9 | + | // command matches a danger pattern (rm -rf, sudo, dd of=/, etc.) | |||
| 10 | + | // the call is refused and the model gets a one-line refusal it | |||
| 11 | + | // can react to. Everything else passes through. | |||
| 12 | + | // | |||
| 13 | + | // Build it: | |||
| 14 | + | // | |||
| 15 | + | // cd examples/extensions/guard | |||
| 16 | + | // go build -o guard . | |||
| 17 | + | // | |||
| 18 | + | // Drop it next to its extension.json under $ZOT_HOME/extensions/guard/ | |||
| 19 | + | // (or `zot ext install ./guard` from this directory) and the next zot | |||
| 20 | + | // session will load it automatically. | |||
| 21 | + | package main | |||
| 22 | + | ||||
| 23 | + | import ( | |||
| 24 | + | "encoding/json" | |||
| 25 | + | "fmt" | |||
| 26 | + | "os" | |||
| 27 | + | "path/filepath" | |||
| 28 | + | "regexp" | |||
| 29 | + | "strings" | |||
| 30 | + | "time" | |||
| 31 | + | ||||
| 32 | + | "github.com/patriceckhart/zot/pkg/zotext" | |||
| 33 | + | ) | |||
| 34 | + | ||||
| 35 | + | var dangerPatterns = []*regexp.Regexp{ | |||
| 36 | + | regexp.MustCompile(`(?i)\brm\s+-rf\b`), | |||
| 37 | + | regexp.MustCompile(`(?i)\bsudo\b`), | |||
| 38 | + | regexp.MustCompile(`(?i)\bdd\s+of=/`), | |||
| 39 | + | regexp.MustCompile(`(?i)\bmkfs\b`), | |||
| 40 | + | regexp.MustCompile(`(?i)\b:\s*\(\s*\)\s*\{\s*:\|:`), // fork bomb | |||
| 41 | + | regexp.MustCompile(`(?i)\bchmod\s+-R\s+777\b`), | |||
| 42 | + | } | |||
| 43 | + | ||||
| 44 | + | func main() { | |||
| 45 | + | ext := zotext.New("guard", "1.0.0") | |||
| 46 | + | ||||
| 47 | + | auditPath := filepath.Join(os.TempDir(), "zot-guard-audit.log") | |||
| 48 | + | auditFile, _ := os.OpenFile(auditPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o644) | |||
| 49 | + | if auditFile != nil { | |||
| 50 | + | defer auditFile.Close() | |||
| 51 | + | } | |||
| 52 | + | audit := func(format string, args ...any) { | |||
| 53 | + | line := fmt.Sprintf("%s "+format+"\n", append([]any{time.Now().Format(time.RFC3339)}, args...)...) | |||
| 54 | + | if auditFile != nil { | |||
| 55 | + | auditFile.WriteString(line) | |||
| 56 | + | } | |||
| 57 | + | ext.Logf("%s", strings.TrimRight(line, "\n")) | |||
| 58 | + | } | |||
| 59 | + | ||||
| 60 | + | // Lifecycle observers. | |||
| 61 | + | ext.On("session_start", func(ev zotext.Event) { audit("[guard] session_start") }) | |||
| 62 | + | ext.On("turn_start", func(ev zotext.Event) { audit("[guard] turn_start step=%d", ev.Step) }) | |||
| 63 | + | ext.On("tool_call", func(ev zotext.Event) { | |||
| 64 | + | audit("[guard] tool_call %s args=%s", ev.ToolName, string(ev.ToolArgs)) | |||
| 65 | + | }) | |||
| 66 | + | ext.On("turn_end", func(ev zotext.Event) { audit("[guard] turn_end stop=%s err=%s", ev.Stop, ev.Error) }) | |||
| 67 | + | ||||
| 68 | + | // Tool-call gate. | |||
| 69 | + | ext.InterceptToolCall(func(toolName string, args json.RawMessage) (bool, string) { | |||
| 70 | + | if toolName != "bash" { | |||
| 71 | + | return true, "" // only guard bash | |||
| 72 | + | } | |||
| 73 | + | var in struct { | |||
| 74 | + | Command string `json:"command"` | |||
| 75 | + | } | |||
| 76 | + | if err := json.Unmarshal(args, &in); err != nil { | |||
| 77 | + | return true, "" | |||
| 78 | + | } | |||
| 79 | + | for _, pat := range dangerPatterns { | |||
| 80 | + | if pat.MatchString(in.Command) { | |||
| 81 | + | audit("[guard] BLOCKED bash: %s (matched %s)", in.Command, pat.String()) | |||
| 82 | + | return false, fmt.Sprintf("guard refused: command matches the danger pattern %q. Ask the user before running this.", pat.String()) | |||
| 83 | + | } | |||
| 84 | + | } | |||
| 85 | + | return true, "" | |||
| 86 | + | }) | |||
| 87 | + | ||||
| 88 | + | if err := ext.Run(); err != nil { | |||
| 89 | + | ext.Logf("fatal: %v", err) | |||
| 90 | + | } | |||
| 91 | + | } | |||
M internal/agent/cli.go +54 -3
| @@ -13,6 +13,7 @@ import ( | |||
| 13 | 13 | "github.com/patriceckhart/zot/internal/agent/modes" | |
| 14 | 14 | "github.com/patriceckhart/zot/internal/auth" | |
| 15 | 15 | "github.com/patriceckhart/zot/internal/core" | |
| 16 | + | "github.com/patriceckhart/zot/internal/extproto" | |
| 16 | 17 | "github.com/patriceckhart/zot/internal/provider" | |
| 17 | 18 | "github.com/patriceckhart/zot/internal/skills" | |
| 18 | 19 | "github.com/patriceckhart/zot/internal/tui" | |
| @@ -86,6 +87,42 @@ func (a *extToolAdapter) NewExtensionTool(info ExtensionToolInfo) core.Tool { | |||
| 86 | 87 | }) | |
| 87 | 88 | } | |
| 88 | 89 | ||
| 90 | + | // fanoutAgentEvent translates a core.AgentEvent into the wire-format | |
| 91 | + | // EventFromHost and pushes it through the extension manager. Only | |
| 92 | + | // the events that have a clear extension-facing meaning are | |
| 93 | + | // forwarded; internal-only ones (text_delta, tool_progress) are | |
| 94 | + | // dropped to keep the per-extension stream sane. | |
| 95 | + | func fanoutAgentEvent(mgr *extensions.Manager, ev core.AgentEvent) { | |
| 96 | + | if mgr == nil { | |
| 97 | + | return | |
| 98 | + | } | |
| 99 | + | switch e := ev.(type) { | |
| 100 | + | case core.EvTurnStart: | |
| 101 | + | mgr.EmitEvent(extproto.EventFromHost{Event: "turn_start", Step: e.Step}) | |
| 102 | + | case core.EvToolCall: | |
| 103 | + | mgr.EmitEvent(extproto.EventFromHost{ | |
| 104 | + | Event: "tool_call", ToolID: e.ID, ToolName: e.Name, ToolArgs: e.Args, | |
| 105 | + | }) | |
| 106 | + | case core.EvAssistantMessage: | |
| 107 | + | // Concat the visible text portions of the message; binary | |
| 108 | + | // blocks (tool_use, etc.) are skipped because subscribers | |
| 109 | + | // usually want a string they can grep / display. | |
| 110 | + | var text string | |
| 111 | + | for _, c := range e.Message.Content { | |
| 112 | + | if tb, ok := c.(provider.TextBlock); ok { | |
| 113 | + | text += tb.Text | |
| 114 | + | } | |
| 115 | + | } | |
| 116 | + | mgr.EmitEvent(extproto.EventFromHost{Event: "assistant_message", Text: text}) | |
| 117 | + | case core.EvTurnEnd: | |
| 118 | + | ev := extproto.EventFromHost{Event: "turn_end", Stop: string(e.Stop)} | |
| 119 | + | if e.Err != nil { | |
| 120 | + | ev.Error = e.Err.Error() | |
| 121 | + | } | |
| 122 | + | mgr.EmitEvent(ev) | |
| 123 | + | } | |
| 124 | + | } | |
| 125 | + | ||
| 89 | 126 | // Run is the top-level entrypoint for the zot binary. | |
| 90 | 127 | func Run(rawArgs []string, version string) error { | |
| 91 | 128 | // Subcommand router: `zot bot ...` is handled separately so the | |
| @@ -233,6 +270,17 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 233 | 270 | ||
| 234 | 271 | // Capture current args in a closure so BuildAgent can re-resolve | |
| 235 | 272 | // after a successful login (picks up the newly stored credential). | |
| 273 | + | wireAgentExt := func(a *core.Agent) *core.Agent { | |
| 274 | + | if a == nil { | |
| 275 | + | return a | |
| 276 | + | } | |
| 277 | + | a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string) { | |
| 278 | + | return extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments) | |
| 279 | + | } | |
| 280 | + | a.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) } | |
| 281 | + | return a | |
| 282 | + | } | |
| 283 | + | ||
| 236 | 284 | buildAgent := func() (*core.Agent, string, string, error) { | |
| 237 | 285 | resolved, err := Resolve(args, true) | |
| 238 | 286 | if err != nil { | |
| @@ -240,7 +288,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 240 | 288 | } | |
| 241 | 289 | resolved.UseSandbox(sharedSandbox) | |
| 242 | 290 | resolved.MergeExtensionTools(extToolAdapter) | |
| 243 | - | return resolved.NewAgent(), resolved.Provider, resolved.Model, nil | |
| 291 | + | return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil | |
| 244 | 292 | } | |
| 245 | 293 | ||
| 246 | 294 | // Rebuild agent with an explicit provider/model override. | |
| @@ -258,14 +306,17 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||
| 258 | 306 | } | |
| 259 | 307 | resolved.UseSandbox(sharedSandbox) | |
| 260 | 308 | resolved.MergeExtensionTools(extToolAdapter) | |
| 261 | - | return resolved.NewAgent(), resolved.Provider, resolved.Model, nil | |
| 309 | + | return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil | |
| 262 | 310 | } | |
| 263 | 311 | ||
| 264 | 312 | var ag *core.Agent | |
| 265 | 313 | if r.HasCredential() { | |
| 266 | - | ag = r.NewAgent() | |
| 314 | + | ag = wireAgentExt(r.NewAgent()) | |
| 267 | 315 | } | |
| 268 | 316 | ||
| 317 | + | // Fire session_start once we know the manager's running. | |
| 318 | + | extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"}) | |
| 319 | + | ||
| 269 | 320 | var sess *core.Session | |
| 270 | 321 | var sessBaselineMsgs int // messages already on disk when current session opened | |
| 271 | 322 | if !args.NoSess && ag != nil { | |
| @@ -13,6 +13,7 @@ import ( | |||||
| 13 | "github.com/patriceckhart/zot/internal/agent/modes" | 13 | "github.com/patriceckhart/zot/internal/agent/modes" | ||
| 14 | "github.com/patriceckhart/zot/internal/auth" | 14 | "github.com/patriceckhart/zot/internal/auth" | ||
| 15 | "github.com/patriceckhart/zot/internal/core" | 15 | "github.com/patriceckhart/zot/internal/core" | ||
| 16 | + | "github.com/patriceckhart/zot/internal/extproto" | |||
| 16 | "github.com/patriceckhart/zot/internal/provider" | 17 | "github.com/patriceckhart/zot/internal/provider" | ||
| 17 | "github.com/patriceckhart/zot/internal/skills" | 18 | "github.com/patriceckhart/zot/internal/skills" | ||
| 18 | "github.com/patriceckhart/zot/internal/tui" | 19 | "github.com/patriceckhart/zot/internal/tui" | ||
| @@ -86,6 +87,42 @@ func (a *extToolAdapter) NewExtensionTool(info ExtensionToolInfo) core.Tool { | |||||
| 86 | }) | 87 | }) | ||
| 87 | } | 88 | } | ||
| 88 | 89 | ||||
| 90 | + | // fanoutAgentEvent translates a core.AgentEvent into the wire-format | |||
| 91 | + | // EventFromHost and pushes it through the extension manager. Only | |||
| 92 | + | // the events that have a clear extension-facing meaning are | |||
| 93 | + | // forwarded; internal-only ones (text_delta, tool_progress) are | |||
| 94 | + | // dropped to keep the per-extension stream sane. | |||
| 95 | + | func fanoutAgentEvent(mgr *extensions.Manager, ev core.AgentEvent) { | |||
| 96 | + | if mgr == nil { | |||
| 97 | + | return | |||
| 98 | + | } | |||
| 99 | + | switch e := ev.(type) { | |||
| 100 | + | case core.EvTurnStart: | |||
| 101 | + | mgr.EmitEvent(extproto.EventFromHost{Event: "turn_start", Step: e.Step}) | |||
| 102 | + | case core.EvToolCall: | |||
| 103 | + | mgr.EmitEvent(extproto.EventFromHost{ | |||
| 104 | + | Event: "tool_call", ToolID: e.ID, ToolName: e.Name, ToolArgs: e.Args, | |||
| 105 | + | }) | |||
| 106 | + | case core.EvAssistantMessage: | |||
| 107 | + | // Concat the visible text portions of the message; binary | |||
| 108 | + | // blocks (tool_use, etc.) are skipped because subscribers | |||
| 109 | + | // usually want a string they can grep / display. | |||
| 110 | + | var text string | |||
| 111 | + | for _, c := range e.Message.Content { | |||
| 112 | + | if tb, ok := c.(provider.TextBlock); ok { | |||
| 113 | + | text += tb.Text | |||
| 114 | + | } | |||
| 115 | + | } | |||
| 116 | + | mgr.EmitEvent(extproto.EventFromHost{Event: "assistant_message", Text: text}) | |||
| 117 | + | case core.EvTurnEnd: | |||
| 118 | + | ev := extproto.EventFromHost{Event: "turn_end", Stop: string(e.Stop)} | |||
| 119 | + | if e.Err != nil { | |||
| 120 | + | ev.Error = e.Err.Error() | |||
| 121 | + | } | |||
| 122 | + | mgr.EmitEvent(ev) | |||
| 123 | + | } | |||
| 124 | + | } | |||
| 125 | + | ||||
| 89 | // Run is the top-level entrypoint for the zot binary. | 126 | // Run is the top-level entrypoint for the zot binary. | ||
| 90 | func Run(rawArgs []string, version string) error { | 127 | func Run(rawArgs []string, version string) error { | ||
| 91 | // Subcommand router: `zot bot ...` is handled separately so the | 128 | // Subcommand router: `zot bot ...` is handled separately so the | ||
| @@ -233,6 +270,17 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 233 | 270 | ||||
| 234 | // Capture current args in a closure so BuildAgent can re-resolve | 271 | // Capture current args in a closure so BuildAgent can re-resolve | ||
| 235 | // after a successful login (picks up the newly stored credential). | 272 | // after a successful login (picks up the newly stored credential). | ||
| 273 | + | wireAgentExt := func(a *core.Agent) *core.Agent { | |||
| 274 | + | if a == nil { | |||
| 275 | + | return a | |||
| 276 | + | } | |||
| 277 | + | a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string) { | |||
| 278 | + | return extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments) | |||
| 279 | + | } | |||
| 280 | + | a.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) } | |||
| 281 | + | return a | |||
| 282 | + | } | |||
| 283 | + | ||||
| 236 | buildAgent := func() (*core.Agent, string, string, error) { | 284 | buildAgent := func() (*core.Agent, string, string, error) { | ||
| 237 | resolved, err := Resolve(args, true) | 285 | resolved, err := Resolve(args, true) | ||
| 238 | if err != nil { | 286 | if err != nil { | ||
| @@ -240,7 +288,7 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 240 | } | 288 | } | ||
| 241 | resolved.UseSandbox(sharedSandbox) | 289 | resolved.UseSandbox(sharedSandbox) | ||
| 242 | resolved.MergeExtensionTools(extToolAdapter) | 290 | resolved.MergeExtensionTools(extToolAdapter) | ||
| 243 | - | return resolved.NewAgent(), resolved.Provider, resolved.Model, nil | 291 | + | return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil |
| 244 | } | 292 | } | ||
| 245 | 293 | ||||
| 246 | // Rebuild agent with an explicit provider/model override. | 294 | // Rebuild agent with an explicit provider/model override. | ||
| @@ -258,14 +306,17 @@ func runInteractive(ctx context.Context, args Args, version string) error { | |||||
| 258 | } | 306 | } | ||
| 259 | resolved.UseSandbox(sharedSandbox) | 307 | resolved.UseSandbox(sharedSandbox) | ||
| 260 | resolved.MergeExtensionTools(extToolAdapter) | 308 | resolved.MergeExtensionTools(extToolAdapter) | ||
| 261 | - | return resolved.NewAgent(), resolved.Provider, resolved.Model, nil | 309 | + | return wireAgentExt(resolved.NewAgent()), resolved.Provider, resolved.Model, nil |
| 262 | } | 310 | } | ||
| 263 | 311 | ||||
| 264 | var ag *core.Agent | 312 | var ag *core.Agent | ||
| 265 | if r.HasCredential() { | 313 | if r.HasCredential() { | ||
| 266 | - | ag = r.NewAgent() | 314 | + | ag = wireAgentExt(r.NewAgent()) |
| 267 | } | 315 | } | ||
| 268 | 316 | ||||
| 317 | + | // Fire session_start once we know the manager's running. | |||
| 318 | + | extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"}) | |||
| 319 | + | ||||
| 269 | var sess *core.Session | 320 | var sess *core.Session | ||
| 270 | var sessBaselineMsgs int // messages already on disk when current session opened | 321 | var sessBaselineMsgs int // messages already on disk when current session opened | ||
| 271 | if !args.NoSess && ag != nil { | 322 | if !args.NoSess && ag != nil { | ||
A internal/agent/extensions/events.go +122 -0
| @@ -0,0 +1,122 @@ | |||
| 1 | + | package extensions | |
| 2 | + | ||
| 3 | + | import ( | |
| 4 | + | "context" | |
| 5 | + | "encoding/json" | |
| 6 | + | "fmt" | |
| 7 | + | "time" | |
| 8 | + | ||
| 9 | + | "github.com/patriceckhart/zot/internal/extproto" | |
| 10 | + | ) | |
| 11 | + | ||
| 12 | + | // EmitEvent fires a one-way lifecycle event to every extension that | |
| 13 | + | // subscribed to it via SubscribeFromExt.events. Non-blocking: each | |
| 14 | + | // extension's pipe write happens on a per-call goroutine so a slow | |
| 15 | + | // extension can't stall the agent loop. | |
| 16 | + | // | |
| 17 | + | // Event names are documented on extproto.EventFromHost. Unknown event | |
| 18 | + | // names are still routed (subscribers can use any string they want). | |
| 19 | + | func (m *Manager) EmitEvent(ev extproto.EventFromHost) { | |
| 20 | + | ev.Type = "event" | |
| 21 | + | frame, err := extproto.Encode(ev) | |
| 22 | + | if err != nil { | |
| 23 | + | return | |
| 24 | + | } | |
| 25 | + | m.mu.RLock() | |
| 26 | + | subs := make([]*Extension, 0, len(m.ext)) | |
| 27 | + | for _, ext := range m.ext { | |
| 28 | + | ext.mu.Lock() | |
| 29 | + | _, subscribed := ext.eventSubs[ev.Event] | |
| 30 | + | ext.mu.Unlock() | |
| 31 | + | if subscribed { | |
| 32 | + | subs = append(subs, ext) | |
| 33 | + | } | |
| 34 | + | } | |
| 35 | + | m.mu.RUnlock() | |
| 36 | + | ||
| 37 | + | for _, ext := range subs { | |
| 38 | + | go func(ext *Extension) { | |
| 39 | + | defer func() { | |
| 40 | + | // A panicking write to a closed pipe shouldn't kill | |
| 41 | + | // the calling goroutine. | |
| 42 | + | _ = recover() | |
| 43 | + | }() | |
| 44 | + | _, _ = ext.stdin.Write(frame) | |
| 45 | + | }(ext) | |
| 46 | + | } | |
| 47 | + | } | |
| 48 | + | ||
| 49 | + | // InterceptToolCall asks every extension that subscribed to tool_call | |
| 50 | + | // interception whether the call may proceed. Subscribers are invoked | |
| 51 | + | // serially; the first one to return Block: true wins, with its Reason | |
| 52 | + | // surfaced as the refusal text. allowed=true means proceed. | |
| 53 | + | // | |
| 54 | + | // Each subscriber gets up to interceptTimeout to reply; missing the | |
| 55 | + | // deadline counts as "allow" (i.e. an unresponsive interceptor never | |
| 56 | + | // stalls the agent). | |
| 57 | + | func (m *Manager) InterceptToolCall(ctx context.Context, toolID, toolName string, args json.RawMessage) (allowed bool, reason string) { | |
| 58 | + | m.mu.RLock() | |
| 59 | + | subs := make([]*Extension, 0, len(m.ext)) | |
| 60 | + | for _, ext := range m.ext { | |
| 61 | + | ext.mu.Lock() | |
| 62 | + | _, subscribed := ext.interceptSubs["tool_call"] | |
| 63 | + | ext.mu.Unlock() | |
| 64 | + | if subscribed { | |
| 65 | + | subs = append(subs, ext) | |
| 66 | + | } | |
| 67 | + | } | |
| 68 | + | m.mu.RUnlock() | |
| 69 | + | if len(subs) == 0 { | |
| 70 | + | return true, "" | |
| 71 | + | } | |
| 72 | + | ||
| 73 | + | for _, ext := range subs { | |
| 74 | + | ok, reason := m.askIntercept(ctx, ext, toolID, toolName, args) | |
| 75 | + | if !ok { | |
| 76 | + | return false, reason | |
| 77 | + | } | |
| 78 | + | } | |
| 79 | + | return true, "" | |
| 80 | + | } | |
| 81 | + | ||
| 82 | + | const interceptTimeout = 5 * time.Second | |
| 83 | + | ||
| 84 | + | func (m *Manager) askIntercept(ctx context.Context, ext *Extension, toolID, toolName string, args json.RawMessage) (allowed bool, reason string) { | |
| 85 | + | id := newCorrelationID() | |
| 86 | + | ch := make(chan extproto.EventInterceptResponseFromExt, 1) | |
| 87 | + | ext.mu.Lock() | |
| 88 | + | ext.pendingIntercept[id] = ch | |
| 89 | + | ext.mu.Unlock() | |
| 90 | + | ||
| 91 | + | frame, _ := extproto.Encode(extproto.EventInterceptFromHost{ | |
| 92 | + | Type: "event_intercept", | |
| 93 | + | ID: id, | |
| 94 | + | Event: "tool_call", | |
| 95 | + | ToolID: toolID, | |
| 96 | + | ToolName: toolName, | |
| 97 | + | ToolArgs: args, | |
| 98 | + | }) | |
| 99 | + | if _, err := ext.stdin.Write(frame); err != nil { | |
| 100 | + | ext.mu.Lock() | |
| 101 | + | delete(ext.pendingIntercept, id) | |
| 102 | + | ext.mu.Unlock() | |
| 103 | + | fmt.Fprintf(ext.logFile, "[zot] intercept write failed: %v\n", err) | |
| 104 | + | return true, "" | |
| 105 | + | } | |
| 106 | + | ||
| 107 | + | select { | |
| 108 | + | case resp := <-ch: | |
| 109 | + | return !resp.Block, resp.Reason | |
| 110 | + | case <-time.After(interceptTimeout): | |
| 111 | + | ext.mu.Lock() | |
| 112 | + | delete(ext.pendingIntercept, id) | |
| 113 | + | ext.mu.Unlock() | |
| 114 | + | fmt.Fprintf(ext.logFile, "[zot] intercept %s timed out; allowing\n", toolName) | |
| 115 | + | return true, "" | |
| 116 | + | case <-ctx.Done(): | |
| 117 | + | ext.mu.Lock() | |
| 118 | + | delete(ext.pendingIntercept, id) | |
| 119 | + | ext.mu.Unlock() | |
| 120 | + | return true, "" | |
| 121 | + | } | |
| 122 | + | } | |
| @@ -0,0 +1,122 @@ | |||||
| 1 | + | package extensions | |||
| 2 | + | ||||
| 3 | + | import ( | |||
| 4 | + | "context" | |||
| 5 | + | "encoding/json" | |||
| 6 | + | "fmt" | |||
| 7 | + | "time" | |||
| 8 | + | ||||
| 9 | + | "github.com/patriceckhart/zot/internal/extproto" | |||
| 10 | + | ) | |||
| 11 | + | ||||
| 12 | + | // EmitEvent fires a one-way lifecycle event to every extension that | |||
| 13 | + | // subscribed to it via SubscribeFromExt.events. Non-blocking: each | |||
| 14 | + | // extension's pipe write happens on a per-call goroutine so a slow | |||
| 15 | + | // extension can't stall the agent loop. | |||
| 16 | + | // | |||
| 17 | + | // Event names are documented on extproto.EventFromHost. Unknown event | |||
| 18 | + | // names are still routed (subscribers can use any string they want). | |||
| 19 | + | func (m *Manager) EmitEvent(ev extproto.EventFromHost) { | |||
| 20 | + | ev.Type = "event" | |||
| 21 | + | frame, err := extproto.Encode(ev) | |||
| 22 | + | if err != nil { | |||
| 23 | + | return | |||
| 24 | + | } | |||
| 25 | + | m.mu.RLock() | |||
| 26 | + | subs := make([]*Extension, 0, len(m.ext)) | |||
| 27 | + | for _, ext := range m.ext { | |||
| 28 | + | ext.mu.Lock() | |||
| 29 | + | _, subscribed := ext.eventSubs[ev.Event] | |||
| 30 | + | ext.mu.Unlock() | |||
| 31 | + | if subscribed { | |||
| 32 | + | subs = append(subs, ext) | |||
| 33 | + | } | |||
| 34 | + | } | |||
| 35 | + | m.mu.RUnlock() | |||
| 36 | + | ||||
| 37 | + | for _, ext := range subs { | |||
| 38 | + | go func(ext *Extension) { | |||
| 39 | + | defer func() { | |||
| 40 | + | // A panicking write to a closed pipe shouldn't kill | |||
| 41 | + | // the calling goroutine. | |||
| 42 | + | _ = recover() | |||
| 43 | + | }() | |||
| 44 | + | _, _ = ext.stdin.Write(frame) | |||
| 45 | + | }(ext) | |||
| 46 | + | } | |||
| 47 | + | } | |||
| 48 | + | ||||
| 49 | + | // InterceptToolCall asks every extension that subscribed to tool_call | |||
| 50 | + | // interception whether the call may proceed. Subscribers are invoked | |||
| 51 | + | // serially; the first one to return Block: true wins, with its Reason | |||
| 52 | + | // surfaced as the refusal text. allowed=true means proceed. | |||
| 53 | + | // | |||
| 54 | + | // Each subscriber gets up to interceptTimeout to reply; missing the | |||
| 55 | + | // deadline counts as "allow" (i.e. an unresponsive interceptor never | |||
| 56 | + | // stalls the agent). | |||
| 57 | + | func (m *Manager) InterceptToolCall(ctx context.Context, toolID, toolName string, args json.RawMessage) (allowed bool, reason string) { | |||
| 58 | + | m.mu.RLock() | |||
| 59 | + | subs := make([]*Extension, 0, len(m.ext)) | |||
| 60 | + | for _, ext := range m.ext { | |||
| 61 | + | ext.mu.Lock() | |||
| 62 | + | _, subscribed := ext.interceptSubs["tool_call"] | |||
| 63 | + | ext.mu.Unlock() | |||
| 64 | + | if subscribed { | |||
| 65 | + | subs = append(subs, ext) | |||
| 66 | + | } | |||
| 67 | + | } | |||
| 68 | + | m.mu.RUnlock() | |||
| 69 | + | if len(subs) == 0 { | |||
| 70 | + | return true, "" | |||
| 71 | + | } | |||
| 72 | + | ||||
| 73 | + | for _, ext := range subs { | |||
| 74 | + | ok, reason := m.askIntercept(ctx, ext, toolID, toolName, args) | |||
| 75 | + | if !ok { | |||
| 76 | + | return false, reason | |||
| 77 | + | } | |||
| 78 | + | } | |||
| 79 | + | return true, "" | |||
| 80 | + | } | |||
| 81 | + | ||||
| 82 | + | const interceptTimeout = 5 * time.Second | |||
| 83 | + | ||||
| 84 | + | func (m *Manager) askIntercept(ctx context.Context, ext *Extension, toolID, toolName string, args json.RawMessage) (allowed bool, reason string) { | |||
| 85 | + | id := newCorrelationID() | |||
| 86 | + | ch := make(chan extproto.EventInterceptResponseFromExt, 1) | |||
| 87 | + | ext.mu.Lock() | |||
| 88 | + | ext.pendingIntercept[id] = ch | |||
| 89 | + | ext.mu.Unlock() | |||
| 90 | + | ||||
| 91 | + | frame, _ := extproto.Encode(extproto.EventInterceptFromHost{ | |||
| 92 | + | Type: "event_intercept", | |||
| 93 | + | ID: id, | |||
| 94 | + | Event: "tool_call", | |||
| 95 | + | ToolID: toolID, | |||
| 96 | + | ToolName: toolName, | |||
| 97 | + | ToolArgs: args, | |||
| 98 | + | }) | |||
| 99 | + | if _, err := ext.stdin.Write(frame); err != nil { | |||
| 100 | + | ext.mu.Lock() | |||
| 101 | + | delete(ext.pendingIntercept, id) | |||
| 102 | + | ext.mu.Unlock() | |||
| 103 | + | fmt.Fprintf(ext.logFile, "[zot] intercept write failed: %v\n", err) | |||
| 104 | + | return true, "" | |||
| 105 | + | } | |||
| 106 | + | ||||
| 107 | + | select { | |||
| 108 | + | case resp := <-ch: | |||
| 109 | + | return !resp.Block, resp.Reason | |||
| 110 | + | case <-time.After(interceptTimeout): | |||
| 111 | + | ext.mu.Lock() | |||
| 112 | + | delete(ext.pendingIntercept, id) | |||
| 113 | + | ext.mu.Unlock() | |||
| 114 | + | fmt.Fprintf(ext.logFile, "[zot] intercept %s timed out; allowing\n", toolName) | |||
| 115 | + | return true, "" | |||
| 116 | + | case <-ctx.Done(): | |||
| 117 | + | ext.mu.Lock() | |||
| 118 | + | delete(ext.pendingIntercept, id) | |||
| 119 | + | ext.mu.Unlock() | |||
| 120 | + | return true, "" | |||
| 121 | + | } | |||
| 122 | + | } | |||
M internal/agent/extensions/manager.go +49 -8
| @@ -73,9 +73,17 @@ type Extension struct { | |||
| 73 | 73 | // pending command invocations waiting on a CommandResponseFromExt | |
| 74 | 74 | // keyed by the id we sent in CommandInvokedFromHost. | |
| 75 | 75 | // pendingTool is the same idea for tool calls. | |
| 76 | - | mu sync.Mutex | |
| 77 | - | pending map[string]chan extproto.CommandResponseFromExt | |
| 78 | - | pendingTool map[string]chan extproto.ToolResultFromExt | |
| 76 | + | // pendingIntercept is the same idea for event_intercept calls. | |
| 77 | + | mu sync.Mutex | |
| 78 | + | pending map[string]chan extproto.CommandResponseFromExt | |
| 79 | + | pendingTool map[string]chan extproto.ToolResultFromExt | |
| 80 | + | pendingIntercept map[string]chan extproto.EventInterceptResponseFromExt | |
| 81 | + | ||
| 82 | + | // eventSubs and interceptSubs are the sets of event names this | |
| 83 | + | // extension subscribed to via SubscribeFromExt. Used by | |
| 84 | + | // EmitEvent / InterceptToolCall to filter recipients. | |
| 85 | + | eventSubs map[string]struct{} | |
| 86 | + | interceptSubs map[string]struct{} | |
| 79 | 87 | } | |
| 80 | 88 | ||
| 81 | 89 | // HostHooks is the small interface the manager calls back into the | |
| @@ -206,11 +214,14 @@ func (m *Manager) loadOne(ctx context.Context, dir string) error { | |||
| 206 | 214 | } | |
| 207 | 215 | ||
| 208 | 216 | ext := &Extension{ | |
| 209 | - | Manifest: mf, | |
| 210 | - | Dir: dir, | |
| 211 | - | readyCh: make(chan struct{}), | |
| 212 | - | pending: map[string]chan extproto.CommandResponseFromExt{}, | |
| 213 | - | pendingTool: map[string]chan extproto.ToolResultFromExt{}, | |
| 217 | + | Manifest: mf, | |
| 218 | + | Dir: dir, | |
| 219 | + | readyCh: make(chan struct{}), | |
| 220 | + | pending: map[string]chan extproto.CommandResponseFromExt{}, | |
| 221 | + | pendingTool: map[string]chan extproto.ToolResultFromExt{}, | |
| 222 | + | pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{}, | |
| 223 | + | eventSubs: map[string]struct{}{}, | |
| 224 | + | interceptSubs: map[string]struct{}{}, | |
| 214 | 225 | } | |
| 215 | 226 | if err := m.spawn(ctx, ext); err != nil { | |
| 216 | 227 | return err | |
| @@ -383,6 +394,36 @@ func (m *Manager) readLoop(ext *Extension, scanner *bufio.Scanner) { | |||
| 383 | 394 | m.mu.Unlock() | |
| 384 | 395 | case "ready": | |
| 385 | 396 | ext.readyOnce.Do(func() { close(ext.readyCh) }) | |
| 397 | + | case "subscribe": | |
| 398 | + | var sub extproto.SubscribeFromExt | |
| 399 | + | if err := json.Unmarshal(line, &sub); err == nil { | |
| 400 | + | ext.mu.Lock() | |
| 401 | + | for _, ev := range sub.Events { | |
| 402 | + | ext.eventSubs[ev] = struct{}{} | |
| 403 | + | } | |
| 404 | + | for _, ev := range sub.Intercept { | |
| 405 | + | if ev == "tool_call" { // only kind supported in v1 | |
| 406 | + | ext.interceptSubs[ev] = struct{}{} | |
| 407 | + | } | |
| 408 | + | } | |
| 409 | + | ext.mu.Unlock() | |
| 410 | + | } | |
| 411 | + | case "event_intercept_response": | |
| 412 | + | var er extproto.EventInterceptResponseFromExt | |
| 413 | + | if err := json.Unmarshal(line, &er); err == nil { | |
| 414 | + | ext.mu.Lock() | |
| 415 | + | ch, ok := ext.pendingIntercept[er.ID] | |
| 416 | + | if ok { | |
| 417 | + | delete(ext.pendingIntercept, er.ID) | |
| 418 | + | } | |
| 419 | + | ext.mu.Unlock() | |
| 420 | + | if ok { | |
| 421 | + | select { | |
| 422 | + | case ch <- er: | |
| 423 | + | default: | |
| 424 | + | } | |
| 425 | + | } | |
| 426 | + | } | |
| 386 | 427 | case "tool_result": | |
| 387 | 428 | var tr extproto.ToolResultFromExt | |
| 388 | 429 | if err := json.Unmarshal(line, &tr); err == nil { | |
| @@ -73,9 +73,17 @@ type Extension struct { | |||||
| 73 | // pending command invocations waiting on a CommandResponseFromExt | 73 | // pending command invocations waiting on a CommandResponseFromExt | ||
| 74 | // keyed by the id we sent in CommandInvokedFromHost. | 74 | // keyed by the id we sent in CommandInvokedFromHost. | ||
| 75 | // pendingTool is the same idea for tool calls. | 75 | // pendingTool is the same idea for tool calls. | ||
| 76 | - | mu sync.Mutex | 76 | + | // pendingIntercept is the same idea for event_intercept calls. |
| 77 | - | pending map[string]chan extproto.CommandResponseFromExt | 77 | + | mu sync.Mutex |
| 78 | - | pendingTool map[string]chan extproto.ToolResultFromExt | 78 | + | pending map[string]chan extproto.CommandResponseFromExt |
| 79 | + | pendingTool map[string]chan extproto.ToolResultFromExt | |||
| 80 | + | pendingIntercept map[string]chan extproto.EventInterceptResponseFromExt | |||
| 81 | + | ||||
| 82 | + | // eventSubs and interceptSubs are the sets of event names this | |||
| 83 | + | // extension subscribed to via SubscribeFromExt. Used by | |||
| 84 | + | // EmitEvent / InterceptToolCall to filter recipients. | |||
| 85 | + | eventSubs map[string]struct{} | |||
| 86 | + | interceptSubs map[string]struct{} | |||
| 79 | } | 87 | } | ||
| 80 | 88 | ||||
| 81 | // HostHooks is the small interface the manager calls back into the | 89 | // HostHooks is the small interface the manager calls back into the | ||
| @@ -206,11 +214,14 @@ func (m *Manager) loadOne(ctx context.Context, dir string) error { | |||||
| 206 | } | 214 | } | ||
| 207 | 215 | ||||
| 208 | ext := &Extension{ | 216 | ext := &Extension{ | ||
| 209 | - | Manifest: mf, | 217 | + | Manifest: mf, |
| 210 | - | Dir: dir, | 218 | + | Dir: dir, |
| 211 | - | readyCh: make(chan struct{}), | 219 | + | readyCh: make(chan struct{}), |
| 212 | - | pending: map[string]chan extproto.CommandResponseFromExt{}, | 220 | + | pending: map[string]chan extproto.CommandResponseFromExt{}, |
| 213 | - | pendingTool: map[string]chan extproto.ToolResultFromExt{}, | 221 | + | pendingTool: map[string]chan extproto.ToolResultFromExt{}, |
| 222 | + | pendingIntercept: map[string]chan extproto.EventInterceptResponseFromExt{}, | |||
| 223 | + | eventSubs: map[string]struct{}{}, | |||
| 224 | + | interceptSubs: map[string]struct{}{}, | |||
| 214 | } | 225 | } | ||
| 215 | if err := m.spawn(ctx, ext); err != nil { | 226 | if err := m.spawn(ctx, ext); err != nil { | ||
| 216 | return err | 227 | return err | ||
| @@ -383,6 +394,36 @@ func (m *Manager) readLoop(ext *Extension, scanner *bufio.Scanner) { | |||||
| 383 | m.mu.Unlock() | 394 | m.mu.Unlock() | ||
| 384 | case "ready": | 395 | case "ready": | ||
| 385 | ext.readyOnce.Do(func() { close(ext.readyCh) }) | 396 | ext.readyOnce.Do(func() { close(ext.readyCh) }) | ||
| 397 | + | case "subscribe": | |||
| 398 | + | var sub extproto.SubscribeFromExt | |||
| 399 | + | if err := json.Unmarshal(line, &sub); err == nil { | |||
| 400 | + | ext.mu.Lock() | |||
| 401 | + | for _, ev := range sub.Events { | |||
| 402 | + | ext.eventSubs[ev] = struct{}{} | |||
| 403 | + | } | |||
| 404 | + | for _, ev := range sub.Intercept { | |||
| 405 | + | if ev == "tool_call" { // only kind supported in v1 | |||
| 406 | + | ext.interceptSubs[ev] = struct{}{} | |||
| 407 | + | } | |||
| 408 | + | } | |||
| 409 | + | ext.mu.Unlock() | |||
| 410 | + | } | |||
| 411 | + | case "event_intercept_response": | |||
| 412 | + | var er extproto.EventInterceptResponseFromExt | |||
| 413 | + | if err := json.Unmarshal(line, &er); err == nil { | |||
| 414 | + | ext.mu.Lock() | |||
| 415 | + | ch, ok := ext.pendingIntercept[er.ID] | |||
| 416 | + | if ok { | |||
| 417 | + | delete(ext.pendingIntercept, er.ID) | |||
| 418 | + | } | |||
| 419 | + | ext.mu.Unlock() | |||
| 420 | + | if ok { | |||
| 421 | + | select { | |||
| 422 | + | case ch <- er: | |||
| 423 | + | default: | |||
| 424 | + | } | |||
| 425 | + | } | |||
| 426 | + | } | |||
| 386 | case "tool_result": | 427 | case "tool_result": | ||
| 387 | var tr extproto.ToolResultFromExt | 428 | var tr extproto.ToolResultFromExt | ||
| 388 | if err := json.Unmarshal(line, &tr); err == nil { | 429 | if err := json.Unmarshal(line, &tr); err == nil { | ||
M internal/agent/rpc.go +6 -0
| @@ -15,6 +15,7 @@ import ( | |||
| 15 | 15 | "github.com/patriceckhart/zot/internal/agent/extensions" | |
| 16 | 16 | "github.com/patriceckhart/zot/internal/agent/modes" | |
| 17 | 17 | "github.com/patriceckhart/zot/internal/core" | |
| 18 | + | "github.com/patriceckhart/zot/internal/extproto" | |
| 18 | 19 | "github.com/patriceckhart/zot/internal/provider" | |
| 19 | 20 | ) | |
| 20 | 21 | ||
| @@ -58,6 +59,11 @@ func runRPCMode(ctx context.Context, args Args, version string) error { | |||
| 58 | 59 | r.MergeExtensionTools(&extToolAdapter{mgr: extMgr}) | |
| 59 | 60 | ||
| 60 | 61 | ag := r.NewAgent() | |
| 62 | + | ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string) { | |
| 63 | + | return extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments) | |
| 64 | + | } | |
| 65 | + | ag.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) } | |
| 66 | + | extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"}) | |
| 61 | 67 | ||
| 62 | 68 | server := &rpcServer{ | |
| 63 | 69 | ctx: ctx, | |
| @@ -15,6 +15,7 @@ import ( | |||||
| 15 | "github.com/patriceckhart/zot/internal/agent/extensions" | 15 | "github.com/patriceckhart/zot/internal/agent/extensions" | ||
| 16 | "github.com/patriceckhart/zot/internal/agent/modes" | 16 | "github.com/patriceckhart/zot/internal/agent/modes" | ||
| 17 | "github.com/patriceckhart/zot/internal/core" | 17 | "github.com/patriceckhart/zot/internal/core" | ||
| 18 | + | "github.com/patriceckhart/zot/internal/extproto" | |||
| 18 | "github.com/patriceckhart/zot/internal/provider" | 19 | "github.com/patriceckhart/zot/internal/provider" | ||
| 19 | ) | 20 | ) | ||
| 20 | 21 | ||||
| @@ -58,6 +59,11 @@ func runRPCMode(ctx context.Context, args Args, version string) error { | |||||
| 58 | r.MergeExtensionTools(&extToolAdapter{mgr: extMgr}) | 59 | r.MergeExtensionTools(&extToolAdapter{mgr: extMgr}) | ||
| 59 | 60 | ||||
| 60 | ag := r.NewAgent() | 61 | ag := r.NewAgent() | ||
| 62 | + | ag.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string) { | |||
| 63 | + | return extMgr.InterceptToolCall(ctx, call.ID, call.Name, call.Arguments) | |||
| 64 | + | } | |||
| 65 | + | ag.OnEvent = func(ev core.AgentEvent) { fanoutAgentEvent(extMgr, ev) } | |||
| 66 | + | extMgr.EmitEvent(extproto.EventFromHost{Event: "session_start"}) | |||
| 61 | 67 | ||||
| 62 | server := &rpcServer{ | 68 | server := &rpcServer{ | ||
| 63 | ctx: ctx, | 69 | ctx: ctx, | ||
M internal/core/agent.go +45 -0
| @@ -21,6 +21,19 @@ type Agent struct { | |||
| 21 | 21 | MaxSteps int | |
| 22 | 22 | Reasoning string | |
| 23 | 23 | ||
| 24 | + | // BeforeToolExecute, if set, is called immediately before each | |
| 25 | + | // tool runs. Returning (allowed=false, reason) short-circuits | |
| 26 | + | // the call with an error result containing reason. Used by the | |
| 27 | + | // extension manager's tool-call interception. Always treated as | |
| 28 | + | // (allowed=true, "") when nil. | |
| 29 | + | BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string) | |
| 30 | + | ||
| 31 | + | // OnEvent, if set, mirrors every AgentEvent the loop emits to | |
| 32 | + | // this callback in addition to the per-Prompt sink. Used by the | |
| 33 | + | // extension manager to fan events out to subscribed extensions | |
| 34 | + | // without each caller having to compose sinks manually. | |
| 35 | + | OnEvent func(AgentEvent) | |
| 36 | + | ||
| 24 | 37 | mu sync.Mutex | |
| 25 | 38 | messages []provider.Message | |
| 26 | 39 | cost CostTracker | |
| @@ -76,6 +89,7 @@ func (a *Agent) Prompt(ctx context.Context, text string, images []provider.Image | |||
| 76 | 89 | if sink == nil { | |
| 77 | 90 | sink = func(AgentEvent) {} | |
| 78 | 91 | } | |
| 92 | + | sink = a.wrapSink(sink) | |
| 79 | 93 | content := []provider.Content{} | |
| 80 | 94 | if text != "" { | |
| 81 | 95 | content = append(content, provider.TextBlock{Text: text}) | |
| @@ -99,9 +113,24 @@ func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error { | |||
| 99 | 113 | if sink == nil { | |
| 100 | 114 | sink = func(AgentEvent) {} | |
| 101 | 115 | } | |
| 116 | + | sink = a.wrapSink(sink) | |
| 102 | 117 | return a.runLoop(ctx, sink) | |
| 103 | 118 | } | |
| 104 | 119 | ||
| 120 | + | // wrapSink composes the per-call sink with a.OnEvent (if set) so the | |
| 121 | + | // extension manager (or any other observer) sees every AgentEvent | |
| 122 | + | // without having to thread itself through every Prompt callsite. | |
| 123 | + | func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) { | |
| 124 | + | if a.OnEvent == nil { | |
| 125 | + | return sink | |
| 126 | + | } | |
| 127 | + | obs := a.OnEvent | |
| 128 | + | return func(ev AgentEvent) { | |
| 129 | + | obs(ev) | |
| 130 | + | sink(ev) | |
| 131 | + | } | |
| 132 | + | } | |
| 133 | + | ||
| 105 | 134 | func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error { | |
| 106 | 135 | for step := 1; step <= a.MaxSteps; step++ { | |
| 107 | 136 | sink(EvTurnStart{Step: step}) | |
| @@ -236,6 +265,22 @@ func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink | |||
| 236 | 265 | } | |
| 237 | 266 | } | |
| 238 | 267 | ||
| 268 | + | // Intercept hook: an extension or other guard can refuse the | |
| 269 | + | // call before any side effect happens. The model sees the | |
| 270 | + | // reason as the tool error, learns from it, and (typically) | |
| 271 | + | // proposes a different action. | |
| 272 | + | if a.BeforeToolExecute != nil { | |
| 273 | + | if allowed, reason := a.BeforeToolExecute(tc); !allowed { | |
| 274 | + | if reason == "" { | |
| 275 | + | reason = "tool call refused by extension guard" | |
| 276 | + | } | |
| 277 | + | return ToolResult{ | |
| 278 | + | Content: []provider.Content{provider.TextBlock{Text: reason}}, | |
| 279 | + | IsError: true, | |
| 280 | + | } | |
| 281 | + | } | |
| 282 | + | } | |
| 283 | + | ||
| 239 | 284 | args := tc.Arguments | |
| 240 | 285 | if len(args) == 0 { | |
| 241 | 286 | args = json.RawMessage("{}") | |
| @@ -21,6 +21,19 @@ type Agent struct { | |||||
| 21 | MaxSteps int | 21 | MaxSteps int | ||
| 22 | Reasoning string | 22 | Reasoning string | ||
| 23 | 23 | ||||
| 24 | + | // BeforeToolExecute, if set, is called immediately before each | |||
| 25 | + | // tool runs. Returning (allowed=false, reason) short-circuits | |||
| 26 | + | // the call with an error result containing reason. Used by the | |||
| 27 | + | // extension manager's tool-call interception. Always treated as | |||
| 28 | + | // (allowed=true, "") when nil. | |||
| 29 | + | BeforeToolExecute func(call provider.ToolCallBlock) (allowed bool, reason string) | |||
| 30 | + | ||||
| 31 | + | // OnEvent, if set, mirrors every AgentEvent the loop emits to | |||
| 32 | + | // this callback in addition to the per-Prompt sink. Used by the | |||
| 33 | + | // extension manager to fan events out to subscribed extensions | |||
| 34 | + | // without each caller having to compose sinks manually. | |||
| 35 | + | OnEvent func(AgentEvent) | |||
| 36 | + | ||||
| 24 | mu sync.Mutex | 37 | mu sync.Mutex | ||
| 25 | messages []provider.Message | 38 | messages []provider.Message | ||
| 26 | cost CostTracker | 39 | cost CostTracker | ||
| @@ -76,6 +89,7 @@ func (a *Agent) Prompt(ctx context.Context, text string, images []provider.Image | |||||
| 76 | if sink == nil { | 89 | if sink == nil { | ||
| 77 | sink = func(AgentEvent) {} | 90 | sink = func(AgentEvent) {} | ||
| 78 | } | 91 | } | ||
| 92 | + | sink = a.wrapSink(sink) | |||
| 79 | content := []provider.Content{} | 93 | content := []provider.Content{} | ||
| 80 | if text != "" { | 94 | if text != "" { | ||
| 81 | content = append(content, provider.TextBlock{Text: text}) | 95 | content = append(content, provider.TextBlock{Text: text}) | ||
| @@ -99,9 +113,24 @@ func (a *Agent) Continue(ctx context.Context, sink func(AgentEvent)) error { | |||||
| 99 | if sink == nil { | 113 | if sink == nil { | ||
| 100 | sink = func(AgentEvent) {} | 114 | sink = func(AgentEvent) {} | ||
| 101 | } | 115 | } | ||
| 116 | + | sink = a.wrapSink(sink) | |||
| 102 | return a.runLoop(ctx, sink) | 117 | return a.runLoop(ctx, sink) | ||
| 103 | } | 118 | } | ||
| 104 | 119 | ||||
| 120 | + | // wrapSink composes the per-call sink with a.OnEvent (if set) so the | |||
| 121 | + | // extension manager (or any other observer) sees every AgentEvent | |||
| 122 | + | // without having to thread itself through every Prompt callsite. | |||
| 123 | + | func (a *Agent) wrapSink(sink func(AgentEvent)) func(AgentEvent) { | |||
| 124 | + | if a.OnEvent == nil { | |||
| 125 | + | return sink | |||
| 126 | + | } | |||
| 127 | + | obs := a.OnEvent | |||
| 128 | + | return func(ev AgentEvent) { | |||
| 129 | + | obs(ev) | |||
| 130 | + | sink(ev) | |||
| 131 | + | } | |||
| 132 | + | } | |||
| 133 | + | ||||
| 105 | func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error { | 134 | func (a *Agent) runLoop(ctx context.Context, sink func(AgentEvent)) error { | ||
| 106 | for step := 1; step <= a.MaxSteps; step++ { | 135 | for step := 1; step <= a.MaxSteps; step++ { | ||
| 107 | sink(EvTurnStart{Step: step}) | 136 | sink(EvTurnStart{Step: step}) | ||
| @@ -236,6 +265,22 @@ func (a *Agent) runOneTool(ctx context.Context, tc provider.ToolCallBlock, sink | |||||
| 236 | } | 265 | } | ||
| 237 | } | 266 | } | ||
| 238 | 267 | ||||
| 268 | + | // Intercept hook: an extension or other guard can refuse the | |||
| 269 | + | // call before any side effect happens. The model sees the | |||
| 270 | + | // reason as the tool error, learns from it, and (typically) | |||
| 271 | + | // proposes a different action. | |||
| 272 | + | if a.BeforeToolExecute != nil { | |||
| 273 | + | if allowed, reason := a.BeforeToolExecute(tc); !allowed { | |||
| 274 | + | if reason == "" { | |||
| 275 | + | reason = "tool call refused by extension guard" | |||
| 276 | + | } | |||
| 277 | + | return ToolResult{ | |||
| 278 | + | Content: []provider.Content{provider.TextBlock{Text: reason}}, | |||
| 279 | + | IsError: true, | |||
| 280 | + | } | |||
| 281 | + | } | |||
| 282 | + | } | |||
| 283 | + | ||||
| 239 | args := tc.Arguments | 284 | args := tc.Arguments | ||
| 240 | if len(args) == 0 { | 285 | if len(args) == 0 { | ||
| 241 | args = json.RawMessage("{}") | 286 | args = json.RawMessage("{}") | ||
M internal/extproto/extproto.go +67 -0
| @@ -73,6 +73,33 @@ type ReadyFromExt struct { | |||
| 73 | 73 | Type string `json:"type"` // "ready" | |
| 74 | 74 | } | |
| 75 | 75 | ||
| 76 | + | // SubscribeFromExt declares which lifecycle events the extension | |
| 77 | + | // wants to observe (one-way `event` frames) and which it wants to | |
| 78 | + | // intercept (round-trip `event_intercept` frames). Send once after | |
| 79 | + | // hello, before ready. | |
| 80 | + | // | |
| 81 | + | // Recognised event names: "session_start", "turn_start", | |
| 82 | + | // "turn_end", "tool_call", "assistant_message". | |
| 83 | + | // | |
| 84 | + | // Only "tool_call" supports interception in this version; values | |
| 85 | + | // listed in Intercept that aren't "tool_call" are ignored. | |
| 86 | + | type SubscribeFromExt struct { | |
| 87 | + | Type string `json:"type"` // "subscribe" | |
| 88 | + | Events []string `json:"events,omitempty"` | |
| 89 | + | Intercept []string `json:"intercept,omitempty"` | |
| 90 | + | } | |
| 91 | + | ||
| 92 | + | // EventInterceptResponseFromExt is the extension's reply to an | |
| 93 | + | // EventInterceptFromHost. block=true refuses the underlying action; | |
| 94 | + | // reason is shown to the model as the tool error message. Both | |
| 95 | + | // fields default to (false, "") meaning "allow". | |
| 96 | + | type EventInterceptResponseFromExt struct { | |
| 97 | + | Type string `json:"type"` // "event_intercept_response" | |
| 98 | + | ID string `json:"id"` | |
| 99 | + | Block bool `json:"block,omitempty"` | |
| 100 | + | Reason string `json:"reason,omitempty"` | |
| 101 | + | } | |
| 102 | + | ||
| 76 | 103 | // ToolResultFromExt is the extension's reply to a ToolCallFromHost. | |
| 77 | 104 | // Content[] follows the same shape as elsewhere in zot: | |
| 78 | 105 | // | |
| @@ -164,6 +191,46 @@ type ToolCallFromHost struct { | |||
| 164 | 191 | Args json.RawMessage `json:"args"` | |
| 165 | 192 | } | |
| 166 | 193 | ||
| 194 | + | // EventFromHost is a one-way lifecycle notification. The payload | |
| 195 | + | // fields populated depend on Event: | |
| 196 | + | // | |
| 197 | + | // session_start : (no extra fields) | |
| 198 | + | // turn_start : Step | |
| 199 | + | // turn_end : Stop, optional Error | |
| 200 | + | // tool_call : ToolID, ToolName, ToolArgs | |
| 201 | + | // assistant_message: Text | |
| 202 | + | type EventFromHost struct { | |
| 203 | + | Type string `json:"type"` // "event" | |
| 204 | + | Event string `json:"event"` | |
| 205 | + | ||
| 206 | + | Step int `json:"step,omitempty"` | |
| 207 | + | Stop string `json:"stop,omitempty"` | |
| 208 | + | Error string `json:"error,omitempty"` | |
| 209 | + | ||
| 210 | + | ToolID string `json:"tool_id,omitempty"` | |
| 211 | + | ToolName string `json:"tool_name,omitempty"` | |
| 212 | + | ToolArgs json.RawMessage `json:"tool_args,omitempty"` | |
| 213 | + | ||
| 214 | + | Text string `json:"text,omitempty"` | |
| 215 | + | } | |
| 216 | + | ||
| 217 | + | // EventInterceptFromHost is sent when zot wants to give the | |
| 218 | + | // extension a chance to block / annotate a lifecycle event before | |
| 219 | + | // it happens. Same payload shape as EventFromHost. Reply with | |
| 220 | + | // EventInterceptResponseFromExt within the host's intercept timeout | |
| 221 | + | // (default 5s); missing the deadline is treated as "allow". | |
| 222 | + | // | |
| 223 | + | // Only Event="tool_call" is sent in this version. | |
| 224 | + | type EventInterceptFromHost struct { | |
| 225 | + | Type string `json:"type"` // "event_intercept" | |
| 226 | + | ID string `json:"id"` | |
| 227 | + | Event string `json:"event"` | |
| 228 | + | ||
| 229 | + | ToolID string `json:"tool_id,omitempty"` | |
| 230 | + | ToolName string `json:"tool_name,omitempty"` | |
| 231 | + | ToolArgs json.RawMessage `json:"tool_args,omitempty"` | |
| 232 | + | } | |
| 233 | + | ||
| 167 | 234 | // ShutdownFromHost asks the extension to clean up and exit. Zot | |
| 168 | 235 | // sends this when the user runs /reload-ext or zot itself is exiting | |
| 169 | 236 | // gracefully. Extensions that don't reply within a few seconds get | |
| @@ -73,6 +73,33 @@ type ReadyFromExt struct { | |||||
| 73 | Type string `json:"type"` // "ready" | 73 | Type string `json:"type"` // "ready" | ||
| 74 | } | 74 | } | ||
| 75 | 75 | ||||
| 76 | + | // SubscribeFromExt declares which lifecycle events the extension | |||
| 77 | + | // wants to observe (one-way `event` frames) and which it wants to | |||
| 78 | + | // intercept (round-trip `event_intercept` frames). Send once after | |||
| 79 | + | // hello, before ready. | |||
| 80 | + | // | |||
| 81 | + | // Recognised event names: "session_start", "turn_start", | |||
| 82 | + | // "turn_end", "tool_call", "assistant_message". | |||
| 83 | + | // | |||
| 84 | + | // Only "tool_call" supports interception in this version; values | |||
| 85 | + | // listed in Intercept that aren't "tool_call" are ignored. | |||
| 86 | + | type SubscribeFromExt struct { | |||
| 87 | + | Type string `json:"type"` // "subscribe" | |||
| 88 | + | Events []string `json:"events,omitempty"` | |||
| 89 | + | Intercept []string `json:"intercept,omitempty"` | |||
| 90 | + | } | |||
| 91 | + | ||||
| 92 | + | // EventInterceptResponseFromExt is the extension's reply to an | |||
| 93 | + | // EventInterceptFromHost. block=true refuses the underlying action; | |||
| 94 | + | // reason is shown to the model as the tool error message. Both | |||
| 95 | + | // fields default to (false, "") meaning "allow". | |||
| 96 | + | type EventInterceptResponseFromExt struct { | |||
| 97 | + | Type string `json:"type"` // "event_intercept_response" | |||
| 98 | + | ID string `json:"id"` | |||
| 99 | + | Block bool `json:"block,omitempty"` | |||
| 100 | + | Reason string `json:"reason,omitempty"` | |||
| 101 | + | } | |||
| 102 | + | ||||
| 76 | // ToolResultFromExt is the extension's reply to a ToolCallFromHost. | 103 | // ToolResultFromExt is the extension's reply to a ToolCallFromHost. | ||
| 77 | // Content[] follows the same shape as elsewhere in zot: | 104 | // Content[] follows the same shape as elsewhere in zot: | ||
| 78 | // | 105 | // | ||
| @@ -164,6 +191,46 @@ type ToolCallFromHost struct { | |||||
| 164 | Args json.RawMessage `json:"args"` | 191 | Args json.RawMessage `json:"args"` | ||
| 165 | } | 192 | } | ||
| 166 | 193 | ||||
| 194 | + | // EventFromHost is a one-way lifecycle notification. The payload | |||
| 195 | + | // fields populated depend on Event: | |||
| 196 | + | // | |||
| 197 | + | // session_start : (no extra fields) | |||
| 198 | + | // turn_start : Step | |||
| 199 | + | // turn_end : Stop, optional Error | |||
| 200 | + | // tool_call : ToolID, ToolName, ToolArgs | |||
| 201 | + | // assistant_message: Text | |||
| 202 | + | type EventFromHost struct { | |||
| 203 | + | Type string `json:"type"` // "event" | |||
| 204 | + | Event string `json:"event"` | |||
| 205 | + | ||||
| 206 | + | Step int `json:"step,omitempty"` | |||
| 207 | + | Stop string `json:"stop,omitempty"` | |||
| 208 | + | Error string `json:"error,omitempty"` | |||
| 209 | + | ||||
| 210 | + | ToolID string `json:"tool_id,omitempty"` | |||
| 211 | + | ToolName string `json:"tool_name,omitempty"` | |||
| 212 | + | ToolArgs json.RawMessage `json:"tool_args,omitempty"` | |||
| 213 | + | ||||
| 214 | + | Text string `json:"text,omitempty"` | |||
| 215 | + | } | |||
| 216 | + | ||||
| 217 | + | // EventInterceptFromHost is sent when zot wants to give the | |||
| 218 | + | // extension a chance to block / annotate a lifecycle event before | |||
| 219 | + | // it happens. Same payload shape as EventFromHost. Reply with | |||
| 220 | + | // EventInterceptResponseFromExt within the host's intercept timeout | |||
| 221 | + | // (default 5s); missing the deadline is treated as "allow". | |||
| 222 | + | // | |||
| 223 | + | // Only Event="tool_call" is sent in this version. | |||
| 224 | + | type EventInterceptFromHost struct { | |||
| 225 | + | Type string `json:"type"` // "event_intercept" | |||
| 226 | + | ID string `json:"id"` | |||
| 227 | + | Event string `json:"event"` | |||
| 228 | + | ||||
| 229 | + | ToolID string `json:"tool_id,omitempty"` | |||
| 230 | + | ToolName string `json:"tool_name,omitempty"` | |||
| 231 | + | ToolArgs json.RawMessage `json:"tool_args,omitempty"` | |||
| 232 | + | } | |||
| 233 | + | ||||
| 167 | // ShutdownFromHost asks the extension to clean up and exit. Zot | 234 | // ShutdownFromHost asks the extension to clean up and exit. Zot | ||
| 168 | // sends this when the user runs /reload-ext or zot itself is exiting | 235 | // sends this when the user runs /reload-ext or zot itself is exiting | ||
| 169 | // gracefully. Extensions that don't reply within a few seconds get | 236 | // gracefully. Extensions that don't reply within a few seconds get | ||
M pkg/zotext/zotext.go +142 -17
| @@ -52,6 +52,39 @@ type CommandHandler func(args string) Response | |||
| 52 | 52 | // ToolResult describing what zot should send back to the model. | |
| 53 | 53 | type ToolHandler func(args json.RawMessage) ToolResult | |
| 54 | 54 | ||
| 55 | + | // EventHandler is called for each lifecycle event the extension | |
| 56 | + | // subscribed to via Subscribe. The handler is invoked synchronously | |
| 57 | + | // on the read goroutine; keep it quick or hand off to your own | |
| 58 | + | // worker. | |
| 59 | + | type EventHandler func(ev Event) | |
| 60 | + | ||
| 61 | + | // Event is a lifecycle notification from zot. The fields populated | |
| 62 | + | // depend on Name (the host's event_name string): | |
| 63 | + | // | |
| 64 | + | // session_start : (no extra fields) | |
| 65 | + | // turn_start : Step | |
| 66 | + | // turn_end : Stop, optional Error | |
| 67 | + | // tool_call : ToolID, ToolName, ToolArgs | |
| 68 | + | // assistant_message: Text | |
| 69 | + | type Event struct { | |
| 70 | + | Name string | |
| 71 | + | ||
| 72 | + | Step int | |
| 73 | + | Stop string | |
| 74 | + | Error string | |
| 75 | + | ||
| 76 | + | ToolID string | |
| 77 | + | ToolName string | |
| 78 | + | ToolArgs json.RawMessage | |
| 79 | + | ||
| 80 | + | Text string | |
| 81 | + | } | |
| 82 | + | ||
| 83 | + | // InterceptHandler decides whether a tool call may proceed. Return | |
| 84 | + | // (allow=true) to permit, (allow=false, reason) to refuse. The | |
| 85 | + | // reason is shown to the model as the tool's error text. | |
| 86 | + | type InterceptHandler func(toolName string, args json.RawMessage) (allow bool, reason string) | |
| 87 | + | ||
| 55 | 88 | // ToolResult is the extension's reply to a tool invocation. Build | |
| 56 | 89 | // one with TextResult, ImageResult, or directly when you need to | |
| 57 | 90 | // combine multiple blocks. | |
| @@ -138,15 +171,17 @@ type Extension struct { | |||
| 138 | 171 | stderr io.Writer | |
| 139 | 172 | writeMu sync.Mutex | |
| 140 | 173 | ||
| 141 | - | mu sync.Mutex | |
| 142 | - | commands map[string]CommandHandler | |
| 143 | - | descriptions []descTuple // ordered so register frames arrive in registration order | |
| 144 | - | tools map[string]ToolHandler | |
| 145 | - | toolDefs []toolDef // ordered so register frames arrive in registration order | |
| 146 | - | ||
| 147 | - | // Caps reported in the hello frame. The SDK enables "commands" | |
| 148 | - | // and "tools" automatically; future capabilities (events) will | |
| 149 | - | // live here when those phases land. | |
| 174 | + | mu sync.Mutex | |
| 175 | + | commands map[string]CommandHandler | |
| 176 | + | descriptions []descTuple // ordered so register frames arrive in registration order | |
| 177 | + | tools map[string]ToolHandler | |
| 178 | + | toolDefs []toolDef // ordered so register frames arrive in registration order | |
| 179 | + | eventHandlers map[string]EventHandler | |
| 180 | + | eventNames []string // declared subscription order | |
| 181 | + | interceptTool InterceptHandler | |
| 182 | + | interceptOn bool | |
| 183 | + | ||
| 184 | + | // Caps reported in the hello frame. | |
| 150 | 185 | caps []string | |
| 151 | 186 | ||
| 152 | 187 | // hostInfo is filled in once HelloAck arrives. | |
| @@ -177,14 +212,15 @@ type HostInfo struct { | |||
| 177 | 212 | // match the name field in extension.json. | |
| 178 | 213 | func New(name, version string) *Extension { | |
| 179 | 214 | return &Extension{ | |
| 180 | - | name: name, | |
| 181 | - | version: version, | |
| 182 | - | in: os.Stdin, | |
| 183 | - | out: os.Stdout, | |
| 184 | - | stderr: os.Stderr, | |
| 185 | - | commands: map[string]CommandHandler{}, | |
| 186 | - | tools: map[string]ToolHandler{}, | |
| 187 | - | caps: []string{"commands", "tools"}, | |
| 215 | + | name: name, | |
| 216 | + | version: version, | |
| 217 | + | in: os.Stdin, | |
| 218 | + | out: os.Stdout, | |
| 219 | + | stderr: os.Stderr, | |
| 220 | + | commands: map[string]CommandHandler{}, | |
| 221 | + | tools: map[string]ToolHandler{}, | |
| 222 | + | eventHandlers: map[string]EventHandler{}, | |
| 223 | + | caps: []string{"commands", "tools", "events"}, | |
| 188 | 224 | } | |
| 189 | 225 | } | |
| 190 | 226 | ||
| @@ -227,6 +263,35 @@ func (e *Extension) Tool(name, description string, schema json.RawMessage, fn To | |||
| 227 | 263 | e.mu.Unlock() | |
| 228 | 264 | } | |
| 229 | 265 | ||
| 266 | + | // On subscribes to a lifecycle event. fn is called for each | |
| 267 | + | // notification; the same name can only have one handler (later | |
| 268 | + | // registrations replace earlier ones). Recognised names: | |
| 269 | + | // session_start, turn_start, turn_end, tool_call, | |
| 270 | + | // assistant_message. | |
| 271 | + | func (e *Extension) On(name string, fn EventHandler) { | |
| 272 | + | e.mu.Lock() | |
| 273 | + | if _, exists := e.eventHandlers[name]; !exists { | |
| 274 | + | e.eventNames = append(e.eventNames, name) | |
| 275 | + | } | |
| 276 | + | e.eventHandlers[name] = fn | |
| 277 | + | e.mu.Unlock() | |
| 278 | + | } | |
| 279 | + | ||
| 280 | + | // InterceptToolCall registers a guard that runs immediately before | |
| 281 | + | // each tool call. Returning (false, reason) refuses the call; the | |
| 282 | + | // model sees reason as the tool error text. Multiple extensions can | |
| 283 | + | // install interceptors; if any one refuses, the call is blocked. | |
| 284 | + | // | |
| 285 | + | // Use this to build permission gates: refuse `bash` calls containing | |
| 286 | + | // `rm -rf`, ask the user for confirmation on dangerous patterns, | |
| 287 | + | // audit-log every call, etc. | |
| 288 | + | func (e *Extension) InterceptToolCall(fn InterceptHandler) { | |
| 289 | + | e.mu.Lock() | |
| 290 | + | e.interceptTool = fn | |
| 291 | + | e.interceptOn = true | |
| 292 | + | e.mu.Unlock() | |
| 293 | + | } | |
| 294 | + | ||
| 230 | 295 | // Notify pushes an info-level status note into zot's chat without | |
| 231 | 296 | // requiring a slash command from the user. | |
| 232 | 297 | func (e *Extension) Notify(level, message string) { | |
| @@ -254,6 +319,8 @@ func (e *Extension) Run() error { | |||
| 254 | 319 | e.mu.Lock() | |
| 255 | 320 | descs := append([]descTuple(nil), e.descriptions...) | |
| 256 | 321 | toolDefs := append([]toolDef(nil), e.toolDefs...) | |
| 322 | + | eventNames := append([]string(nil), e.eventNames...) | |
| 323 | + | interceptOn := e.interceptOn | |
| 257 | 324 | e.mu.Unlock() | |
| 258 | 325 | for _, d := range descs { | |
| 259 | 326 | _ = e.send(extproto.RegisterCommandFromExt{ | |
| @@ -270,6 +337,13 @@ func (e *Extension) Run() error { | |||
| 270 | 337 | Schema: td.schema, | |
| 271 | 338 | }) | |
| 272 | 339 | } | |
| 340 | + | if len(eventNames) > 0 || interceptOn { | |
| 341 | + | sub := extproto.SubscribeFromExt{Type: "subscribe", Events: eventNames} | |
| 342 | + | if interceptOn { | |
| 343 | + | sub.Intercept = []string{"tool_call"} | |
| 344 | + | } | |
| 345 | + | _ = e.send(sub) | |
| 346 | + | } | |
| 273 | 347 | // Sentinel: tells the host that all initial registrations have | |
| 274 | 348 | // been flushed and the agent registry can be built. Never block | |
| 275 | 349 | // on this; if the host doesn't act on it, registrations still | |
| @@ -341,6 +415,57 @@ func (e *Extension) Run() error { | |||
| 341 | 415 | res := fn(args) | |
| 342 | 416 | e.respondTool(id, res) | |
| 343 | 417 | }(tc.ID, fn, tc.Args) | |
| 418 | + | case "event": | |
| 419 | + | var ef extproto.EventFromHost | |
| 420 | + | if err := json.Unmarshal(line, &ef); err != nil { | |
| 421 | + | continue | |
| 422 | + | } | |
| 423 | + | e.mu.Lock() | |
| 424 | + | handler := e.eventHandlers[ef.Event] | |
| 425 | + | e.mu.Unlock() | |
| 426 | + | if handler != nil { | |
| 427 | + | func() { | |
| 428 | + | defer func() { | |
| 429 | + | if r := recover(); r != nil { | |
| 430 | + | e.Logf("event %s handler panicked: %v", ef.Event, r) | |
| 431 | + | } | |
| 432 | + | }() | |
| 433 | + | handler(Event{ | |
| 434 | + | Name: ef.Event, Step: ef.Step, Stop: ef.Stop, | |
| 435 | + | Error: ef.Error, ToolID: ef.ToolID, ToolName: ef.ToolName, | |
| 436 | + | ToolArgs: ef.ToolArgs, Text: ef.Text, | |
| 437 | + | }) | |
| 438 | + | }() | |
| 439 | + | } | |
| 440 | + | case "event_intercept": | |
| 441 | + | var ei extproto.EventInterceptFromHost | |
| 442 | + | if err := json.Unmarshal(line, &ei); err != nil { | |
| 443 | + | continue | |
| 444 | + | } | |
| 445 | + | e.mu.Lock() | |
| 446 | + | fn := e.interceptTool | |
| 447 | + | e.mu.Unlock() | |
| 448 | + | if fn == nil || ei.Event != "tool_call" { | |
| 449 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |
| 450 | + | Type: "event_intercept_response", ID: ei.ID, | |
| 451 | + | }) | |
| 452 | + | continue | |
| 453 | + | } | |
| 454 | + | go func(id, name string, args json.RawMessage) { | |
| 455 | + | defer func() { | |
| 456 | + | if r := recover(); r != nil { | |
| 457 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |
| 458 | + | Type: "event_intercept_response", ID: id, | |
| 459 | + | Block: true, Reason: fmt.Sprintf("intercept panic: %v", r), | |
| 460 | + | }) | |
| 461 | + | } | |
| 462 | + | }() | |
| 463 | + | allow, reason := fn(name, args) | |
| 464 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |
| 465 | + | Type: "event_intercept_response", ID: id, | |
| 466 | + | Block: !allow, Reason: reason, | |
| 467 | + | }) | |
| 468 | + | }(ei.ID, ei.ToolName, ei.ToolArgs) | |
| 344 | 469 | case "shutdown": | |
| 345 | 470 | _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"}) | |
| 346 | 471 | return nil | |
| @@ -52,6 +52,39 @@ type CommandHandler func(args string) Response | |||||
| 52 | // ToolResult describing what zot should send back to the model. | 52 | // ToolResult describing what zot should send back to the model. | ||
| 53 | type ToolHandler func(args json.RawMessage) ToolResult | 53 | type ToolHandler func(args json.RawMessage) ToolResult | ||
| 54 | 54 | ||||
| 55 | + | // EventHandler is called for each lifecycle event the extension | |||
| 56 | + | // subscribed to via Subscribe. The handler is invoked synchronously | |||
| 57 | + | // on the read goroutine; keep it quick or hand off to your own | |||
| 58 | + | // worker. | |||
| 59 | + | type EventHandler func(ev Event) | |||
| 60 | + | ||||
| 61 | + | // Event is a lifecycle notification from zot. The fields populated | |||
| 62 | + | // depend on Name (the host's event_name string): | |||
| 63 | + | // | |||
| 64 | + | // session_start : (no extra fields) | |||
| 65 | + | // turn_start : Step | |||
| 66 | + | // turn_end : Stop, optional Error | |||
| 67 | + | // tool_call : ToolID, ToolName, ToolArgs | |||
| 68 | + | // assistant_message: Text | |||
| 69 | + | type Event struct { | |||
| 70 | + | Name string | |||
| 71 | + | ||||
| 72 | + | Step int | |||
| 73 | + | Stop string | |||
| 74 | + | Error string | |||
| 75 | + | ||||
| 76 | + | ToolID string | |||
| 77 | + | ToolName string | |||
| 78 | + | ToolArgs json.RawMessage | |||
| 79 | + | ||||
| 80 | + | Text string | |||
| 81 | + | } | |||
| 82 | + | ||||
| 83 | + | // InterceptHandler decides whether a tool call may proceed. Return | |||
| 84 | + | // (allow=true) to permit, (allow=false, reason) to refuse. The | |||
| 85 | + | // reason is shown to the model as the tool's error text. | |||
| 86 | + | type InterceptHandler func(toolName string, args json.RawMessage) (allow bool, reason string) | |||
| 87 | + | ||||
| 55 | // ToolResult is the extension's reply to a tool invocation. Build | 88 | // ToolResult is the extension's reply to a tool invocation. Build | ||
| 56 | // one with TextResult, ImageResult, or directly when you need to | 89 | // one with TextResult, ImageResult, or directly when you need to | ||
| 57 | // combine multiple blocks. | 90 | // combine multiple blocks. | ||
| @@ -138,15 +171,17 @@ type Extension struct { | |||||
| 138 | stderr io.Writer | 171 | stderr io.Writer | ||
| 139 | writeMu sync.Mutex | 172 | writeMu sync.Mutex | ||
| 140 | 173 | ||||
| 141 | - | mu sync.Mutex | 174 | + | mu sync.Mutex |
| 142 | - | commands map[string]CommandHandler | 175 | + | commands map[string]CommandHandler |
| 143 | - | descriptions []descTuple // ordered so register frames arrive in registration order | 176 | + | descriptions []descTuple // ordered so register frames arrive in registration order |
| 144 | - | tools map[string]ToolHandler | 177 | + | tools map[string]ToolHandler |
| 145 | - | toolDefs []toolDef // ordered so register frames arrive in registration order | 178 | + | toolDefs []toolDef // ordered so register frames arrive in registration order |
| 146 | - | 179 | + | eventHandlers map[string]EventHandler | |
| 147 | - | // Caps reported in the hello frame. The SDK enables "commands" | 180 | + | eventNames []string // declared subscription order |
| 148 | - | // and "tools" automatically; future capabilities (events) will | 181 | + | interceptTool InterceptHandler |
| 149 | - | // live here when those phases land. | 182 | + | interceptOn bool |
| 183 | + | ||||
| 184 | + | // Caps reported in the hello frame. | |||
| 150 | caps []string | 185 | caps []string | ||
| 151 | 186 | ||||
| 152 | // hostInfo is filled in once HelloAck arrives. | 187 | // hostInfo is filled in once HelloAck arrives. | ||
| @@ -177,14 +212,15 @@ type HostInfo struct { | |||||
| 177 | // match the name field in extension.json. | 212 | // match the name field in extension.json. | ||
| 178 | func New(name, version string) *Extension { | 213 | func New(name, version string) *Extension { | ||
| 179 | return &Extension{ | 214 | return &Extension{ | ||
| 180 | - | name: name, | 215 | + | name: name, |
| 181 | - | version: version, | 216 | + | version: version, |
| 182 | - | in: os.Stdin, | 217 | + | in: os.Stdin, |
| 183 | - | out: os.Stdout, | 218 | + | out: os.Stdout, |
| 184 | - | stderr: os.Stderr, | 219 | + | stderr: os.Stderr, |
| 185 | - | commands: map[string]CommandHandler{}, | 220 | + | commands: map[string]CommandHandler{}, |
| 186 | - | tools: map[string]ToolHandler{}, | 221 | + | tools: map[string]ToolHandler{}, |
| 187 | - | caps: []string{"commands", "tools"}, | 222 | + | eventHandlers: map[string]EventHandler{}, |
| 223 | + | caps: []string{"commands", "tools", "events"}, | |||
| 188 | } | 224 | } | ||
| 189 | } | 225 | } | ||
| 190 | 226 | ||||
| @@ -227,6 +263,35 @@ func (e *Extension) Tool(name, description string, schema json.RawMessage, fn To | |||||
| 227 | e.mu.Unlock() | 263 | e.mu.Unlock() | ||
| 228 | } | 264 | } | ||
| 229 | 265 | ||||
| 266 | + | // On subscribes to a lifecycle event. fn is called for each | |||
| 267 | + | // notification; the same name can only have one handler (later | |||
| 268 | + | // registrations replace earlier ones). Recognised names: | |||
| 269 | + | // session_start, turn_start, turn_end, tool_call, | |||
| 270 | + | // assistant_message. | |||
| 271 | + | func (e *Extension) On(name string, fn EventHandler) { | |||
| 272 | + | e.mu.Lock() | |||
| 273 | + | if _, exists := e.eventHandlers[name]; !exists { | |||
| 274 | + | e.eventNames = append(e.eventNames, name) | |||
| 275 | + | } | |||
| 276 | + | e.eventHandlers[name] = fn | |||
| 277 | + | e.mu.Unlock() | |||
| 278 | + | } | |||
| 279 | + | ||||
| 280 | + | // InterceptToolCall registers a guard that runs immediately before | |||
| 281 | + | // each tool call. Returning (false, reason) refuses the call; the | |||
| 282 | + | // model sees reason as the tool error text. Multiple extensions can | |||
| 283 | + | // install interceptors; if any one refuses, the call is blocked. | |||
| 284 | + | // | |||
| 285 | + | // Use this to build permission gates: refuse `bash` calls containing | |||
| 286 | + | // `rm -rf`, ask the user for confirmation on dangerous patterns, | |||
| 287 | + | // audit-log every call, etc. | |||
| 288 | + | func (e *Extension) InterceptToolCall(fn InterceptHandler) { | |||
| 289 | + | e.mu.Lock() | |||
| 290 | + | e.interceptTool = fn | |||
| 291 | + | e.interceptOn = true | |||
| 292 | + | e.mu.Unlock() | |||
| 293 | + | } | |||
| 294 | + | ||||
| 230 | // Notify pushes an info-level status note into zot's chat without | 295 | // Notify pushes an info-level status note into zot's chat without | ||
| 231 | // requiring a slash command from the user. | 296 | // requiring a slash command from the user. | ||
| 232 | func (e *Extension) Notify(level, message string) { | 297 | func (e *Extension) Notify(level, message string) { | ||
| @@ -254,6 +319,8 @@ func (e *Extension) Run() error { | |||||
| 254 | e.mu.Lock() | 319 | e.mu.Lock() | ||
| 255 | descs := append([]descTuple(nil), e.descriptions...) | 320 | descs := append([]descTuple(nil), e.descriptions...) | ||
| 256 | toolDefs := append([]toolDef(nil), e.toolDefs...) | 321 | toolDefs := append([]toolDef(nil), e.toolDefs...) | ||
| 322 | + | eventNames := append([]string(nil), e.eventNames...) | |||
| 323 | + | interceptOn := e.interceptOn | |||
| 257 | e.mu.Unlock() | 324 | e.mu.Unlock() | ||
| 258 | for _, d := range descs { | 325 | for _, d := range descs { | ||
| 259 | _ = e.send(extproto.RegisterCommandFromExt{ | 326 | _ = e.send(extproto.RegisterCommandFromExt{ | ||
| @@ -270,6 +337,13 @@ func (e *Extension) Run() error { | |||||
| 270 | Schema: td.schema, | 337 | Schema: td.schema, | ||
| 271 | }) | 338 | }) | ||
| 272 | } | 339 | } | ||
| 340 | + | if len(eventNames) > 0 || interceptOn { | |||
| 341 | + | sub := extproto.SubscribeFromExt{Type: "subscribe", Events: eventNames} | |||
| 342 | + | if interceptOn { | |||
| 343 | + | sub.Intercept = []string{"tool_call"} | |||
| 344 | + | } | |||
| 345 | + | _ = e.send(sub) | |||
| 346 | + | } | |||
| 273 | // Sentinel: tells the host that all initial registrations have | 347 | // Sentinel: tells the host that all initial registrations have | ||
| 274 | // been flushed and the agent registry can be built. Never block | 348 | // been flushed and the agent registry can be built. Never block | ||
| 275 | // on this; if the host doesn't act on it, registrations still | 349 | // on this; if the host doesn't act on it, registrations still | ||
| @@ -341,6 +415,57 @@ func (e *Extension) Run() error { | |||||
| 341 | res := fn(args) | 415 | res := fn(args) | ||
| 342 | e.respondTool(id, res) | 416 | e.respondTool(id, res) | ||
| 343 | }(tc.ID, fn, tc.Args) | 417 | }(tc.ID, fn, tc.Args) | ||
| 418 | + | case "event": | |||
| 419 | + | var ef extproto.EventFromHost | |||
| 420 | + | if err := json.Unmarshal(line, &ef); err != nil { | |||
| 421 | + | continue | |||
| 422 | + | } | |||
| 423 | + | e.mu.Lock() | |||
| 424 | + | handler := e.eventHandlers[ef.Event] | |||
| 425 | + | e.mu.Unlock() | |||
| 426 | + | if handler != nil { | |||
| 427 | + | func() { | |||
| 428 | + | defer func() { | |||
| 429 | + | if r := recover(); r != nil { | |||
| 430 | + | e.Logf("event %s handler panicked: %v", ef.Event, r) | |||
| 431 | + | } | |||
| 432 | + | }() | |||
| 433 | + | handler(Event{ | |||
| 434 | + | Name: ef.Event, Step: ef.Step, Stop: ef.Stop, | |||
| 435 | + | Error: ef.Error, ToolID: ef.ToolID, ToolName: ef.ToolName, | |||
| 436 | + | ToolArgs: ef.ToolArgs, Text: ef.Text, | |||
| 437 | + | }) | |||
| 438 | + | }() | |||
| 439 | + | } | |||
| 440 | + | case "event_intercept": | |||
| 441 | + | var ei extproto.EventInterceptFromHost | |||
| 442 | + | if err := json.Unmarshal(line, &ei); err != nil { | |||
| 443 | + | continue | |||
| 444 | + | } | |||
| 445 | + | e.mu.Lock() | |||
| 446 | + | fn := e.interceptTool | |||
| 447 | + | e.mu.Unlock() | |||
| 448 | + | if fn == nil || ei.Event != "tool_call" { | |||
| 449 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |||
| 450 | + | Type: "event_intercept_response", ID: ei.ID, | |||
| 451 | + | }) | |||
| 452 | + | continue | |||
| 453 | + | } | |||
| 454 | + | go func(id, name string, args json.RawMessage) { | |||
| 455 | + | defer func() { | |||
| 456 | + | if r := recover(); r != nil { | |||
| 457 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |||
| 458 | + | Type: "event_intercept_response", ID: id, | |||
| 459 | + | Block: true, Reason: fmt.Sprintf("intercept panic: %v", r), | |||
| 460 | + | }) | |||
| 461 | + | } | |||
| 462 | + | }() | |||
| 463 | + | allow, reason := fn(name, args) | |||
| 464 | + | _ = e.send(extproto.EventInterceptResponseFromExt{ | |||
| 465 | + | Type: "event_intercept_response", ID: id, | |||
| 466 | + | Block: !allow, Reason: reason, | |||
| 467 | + | }) | |||
| 468 | + | }(ei.ID, ei.ToolName, ei.ToolArgs) | |||
| 344 | case "shutdown": | 469 | case "shutdown": | ||
| 345 | _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"}) | 470 | _ = e.send(extproto.ShutdownAckFromExt{Type: "shutdown_ack"}) | ||
| 346 | return nil | 471 | return nil | ||