mirror of
https://github.com/patriceckhart/zot.git
synced 2026-06-27 05:46:34 +02:00
99c9ba8062
Clears every deferred extension todo in one push:
1) Interception expands to three events: tool_call (already shipped),
turn_start (gate the turn before the model call, e.g. rate-limit /
business-hour), and assistant_message (suppress or rewrite the
user-visible text while keeping the model's original output in
the transcript).
2) Tool-call args can now be rewritten mid-flight. An interceptor
returning modified_args replaces the JSON the tool actually
receives, without the model seeing the rewrite. Chains: each
subscriber sees the previous one's output, letting guards
successively redact / patch / augment. Invalid JSON is dropped
safely.
3) /reload-ext hot-reloads every extension without restarting zot.
The manager gracefully shuts down all running subprocesses,
re-reads extension.json from disk, respawns (including --ext
paths remembered from startup), and the host rebuilds the agent's
tool registry in-place so freshly-registered tools are callable
immediately.
Wire-format changes (extproto):
- EventInterceptResponseFromExt gains modified_args and replace_text
fields (both optional, ignored when block=true).
- EventInterceptFromHost gains Step (for turn_start) and Text (for
assistant_message) alongside the existing tool_call payload.
Core agent changes:
- BeforeToolExecute signature now returns (allowed, reason,
modifiedArgs json.RawMessage). Non-nil+valid JSON args replace
tc.Arguments before Tool.Execute runs.
- New BeforeTurn hook, invoked in runLoop before oneTurn. Blocking
cancels the turn with an EvTurnEnd{StopError} carrying the reason.
- New BeforeAssistantMessage hook, invoked after finalMsg is
assembled but before the EvAssistantMessage emit. Supports
suppress (block=true) and text rewrite (replace_text). Transcript
always gets the original; UI gets the rewritten text.
- New SetTools(reg) so /reload-ext can swap the registry on the
live agent under the agent mutex.
Manager changes:
- InterceptToolCall now returns InterceptResult (Block, Reason,
ModifiedArgs, ReplaceText), with a chain that folds rewrites.
- New InterceptTurnStart and InterceptAssistantMessage.
- New Reload(ctx, grace) tears down and respawns everything,
returning ReloadStats{Stopped, Loaded, Ready, Errors}.
- New SetOnReload(fn) callback the host uses to rebuild the agent
tool registry after a reload.
- LoadExplicit remembers --ext paths so Reload respawns them.
- subscribe accepts "tool_call", "turn_start", "assistant_message"
under "intercept".
SDK (pkg/zotext):
- New handler types: ToolCallHandler, TurnStartHandler,
AssistantMessageHandler, and their decision structs
(ToolCallDecision with ModifiedArgs, AssistantMessageDecision
with ReplaceText).
- New registration methods: InterceptToolCallX (rich variant of
the existing InterceptToolCall), InterceptTurnStart,
InterceptAssistantMessage.
- dispatchIntercept routes per-event with panic recovery and
always emits exactly one event_intercept_response.
TUI:
- /reload-ext slash command registered in slashCatalog and
runSlash. Added to slashCancelsTurn so it waits for idle like
/compact does.
- runReloadExt shows a "reloading extensions..." status, runs the
Manager.Reload on a goroutine, and reports the resulting stats.
Tests:
- internal/core/intercept_test.go: verifies args are actually
rewritten on the way to Tool.Execute, malformed JSON is ignored,
and block surfaces the reason as an error ToolResult.
- internal/agent/extensions/intercept_test.go: end-to-end with a
bash extension subprocess that blocks rm -rf, rewrites other bash
args to "echo GUARDED:", passes through read calls, allows
turn_start, and redacts SECRET in assistant messages. Second test
verifies Reload respawns the subprocess, re-registers its command,
and fires the onReload callback.
Docs:
- docs/extensions.md: rewrote the intercept section to cover all
three events, added a table of event_intercept_response fields,
documented the /reload-ext hot-reload command, expanded the SDK
section with examples of every handler, moved the old "future"
items into a shipped Phase 4.
- README.md: extensions summary mentions intercept beyond tool_call,
/reload-ext added to the slash-commands table and to the
turn-cancel list in "Queued messages".
109 lines
3.3 KiB
Go
109 lines
3.3 KiB
Go
package core
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"github.com/patriceckhart/zot/internal/provider"
|
|
)
|
|
|
|
// recordingTool captures the args it was invoked with so the test
|
|
// can verify the interceptor-rewritten args reached execution.
|
|
type recordingTool struct {
|
|
lastArgs json.RawMessage
|
|
}
|
|
|
|
func (r *recordingTool) Name() string { return "echo" }
|
|
func (r *recordingTool) Description() string { return "echoes" }
|
|
func (r *recordingTool) Schema() json.RawMessage { return json.RawMessage(`{"type":"object"}`) }
|
|
|
|
func (r *recordingTool) Execute(_ context.Context, args json.RawMessage, _ func(string)) (ToolResult, error) {
|
|
r.lastArgs = append(json.RawMessage(nil), args...)
|
|
return ToolResult{
|
|
Content: []provider.Content{provider.TextBlock{Text: "ok"}},
|
|
}, nil
|
|
}
|
|
|
|
// TestBeforeToolExecuteModifiesArgs verifies that a non-nil
|
|
// modifiedArgs returned from BeforeToolExecute is what the tool
|
|
// actually sees.
|
|
func TestBeforeToolExecuteModifiesArgs(t *testing.T) {
|
|
rec := &recordingTool{}
|
|
reg := Registry{"echo": rec}
|
|
a := NewAgent(nil, "test", "", reg)
|
|
|
|
newArgs := json.RawMessage(`{"command":"echo GUARDED: ls"}`)
|
|
a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
|
|
return true, "", newArgs
|
|
}
|
|
|
|
ctx := context.Background()
|
|
res := a.runOneTool(ctx, provider.ToolCallBlock{
|
|
ID: "T1",
|
|
Name: "echo",
|
|
Arguments: json.RawMessage(`{"command":"ls"}`),
|
|
}, func(AgentEvent) {})
|
|
if res.IsError {
|
|
t.Fatalf("unexpected error result: %v", res.Content)
|
|
}
|
|
if string(rec.lastArgs) != string(newArgs) {
|
|
t.Errorf("tool saw %s, want %s", string(rec.lastArgs), string(newArgs))
|
|
}
|
|
}
|
|
|
|
// TestBeforeToolExecuteInvalidJSONIgnored verifies that returning
|
|
// malformed JSON as modifiedArgs leaves the original args intact
|
|
// (safety: a buggy interceptor can't corrupt the call).
|
|
func TestBeforeToolExecuteInvalidJSONIgnored(t *testing.T) {
|
|
rec := &recordingTool{}
|
|
reg := Registry{"echo": rec}
|
|
a := NewAgent(nil, "test", "", reg)
|
|
|
|
a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
|
|
return true, "", json.RawMessage(`{not json`)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
orig := json.RawMessage(`{"command":"ls"}`)
|
|
a.runOneTool(ctx, provider.ToolCallBlock{
|
|
ID: "T1",
|
|
Name: "echo",
|
|
Arguments: orig,
|
|
}, func(AgentEvent) {})
|
|
if string(rec.lastArgs) != string(orig) {
|
|
t.Errorf("tool saw %s, want original %s", string(rec.lastArgs), string(orig))
|
|
}
|
|
}
|
|
|
|
// TestBeforeToolExecuteBlockSurfacesReason verifies a refusal from
|
|
// the interceptor returns an error ToolResult with the reason text.
|
|
func TestBeforeToolExecuteBlockSurfacesReason(t *testing.T) {
|
|
rec := &recordingTool{}
|
|
reg := Registry{"echo": rec}
|
|
a := NewAgent(nil, "test", "", reg)
|
|
|
|
a.BeforeToolExecute = func(call provider.ToolCallBlock) (bool, string, json.RawMessage) {
|
|
return false, "nope", nil
|
|
}
|
|
|
|
ctx := context.Background()
|
|
res := a.runOneTool(ctx, provider.ToolCallBlock{
|
|
ID: "T1",
|
|
Name: "echo",
|
|
Arguments: json.RawMessage(`{"command":"ls"}`),
|
|
}, func(AgentEvent) {})
|
|
if !res.IsError {
|
|
t.Fatal("want error result, got success")
|
|
}
|
|
if len(res.Content) == 0 {
|
|
t.Fatal("no content")
|
|
}
|
|
tb, ok := res.Content[0].(provider.TextBlock)
|
|
if !ok || tb.Text != "nope" {
|
|
t.Errorf("want reason 'nope', got %v", res.Content[0])
|
|
}
|
|
if rec.lastArgs != nil {
|
|
t.Error("tool ran despite block")
|
|
}
|
|
}
|