Vectrum
Sign In
CodeyPublic
Like0 likesFork0 forks

agent: report each message as it joins the history; tools: Registry.Exclude

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Browse Files
Jjonicommitted 1 parent 6158e7bcommit3916e1d

Showing 4 changed files with 123 additions and 1 deletion.

+25−1internal/agent/loop.goView File
@@ -304,6 +304,11 @@ type Options struct {
304304 // the model is done, the window is full, or the context is cancelled.
305305 MaxTurns int
306306 System string
307+ // SystemAppend is added to the end of the system prompt, whichever one
308+ // that is — `--append-system-prompt`, for a script that wants to add a
309+ // house rule without replacing the instructions on how to work. It is
310+ // constant for the session, so it costs no cache hit.
311+ SystemAppend string
307312 // Effort maps to output_config.effort. Empty means the API default (high).
308313 Effort string
309314 // ContextWindow, CompactAt and CompactFloor control when the conversation
@@ -385,6 +390,14 @@ const (
385390 // EventResumed when it is let go again.
386391 EventPaused EventKind = "paused"
387392 EventResumed EventKind = "resumed"
393+ // EventMessage is a message joining the history: each complete response
394+ // from the model, with Response set, and each batch of tool results that
395+ // answers one, with Message set. It is for a front end that reports the
396+ // conversation itself rather than drawing it — `-p --output-format
397+ // stream-json` prints one line per message — and everything else can
398+ // ignore it, because the same content has already arrived as the events
399+ // above.
400+ EventMessage EventKind = "message"
388401 )
389402
390403 // Event is what the loop reports to an observer. It carries plain data only —
@@ -404,6 +417,11 @@ type Event struct {
404417 Result string
405418 IsError bool
406419 Usage llm.Usage
420+ // Response and Message are EventMessage's: the model's whole response,
421+ // or the user message of tool results the loop sent back. Exactly one is
422+ // set.
423+ Response *llm.Response
424+ Message *llm.Message
407425 }
408426
409427 // Approver decides whether a tool call may proceed. It is an interface here so
@@ -555,6 +573,9 @@ func New(client *llm.Client, reg *tools.Registry, state *State, log *metrics.Log
555573 }
556574 opts.System += skillListing(opts.Skills)
557575 opts.System += modelListing(opts.Models)
576+ if opts.SystemAppend != "" {
577+ opts.System += "\n\n" + opts.SystemAppend
578+ }
558579 return &Agent{client: client, tools: reg, state: state, log: log, opts: opts}
559580 }
560581
@@ -686,6 +707,7 @@ func (a *Agent) drive(ctx context.Context) (string, error) {
686707 // The assistant message goes into history verbatim, thinking blocks and
687708 // all: their signatures must come back unchanged on the next request.
688709 a.state.Append(resp.Message())
710+ a.emit(Event{Kind: EventMessage, Turn: turn, Response: resp})
689711
690712 // Hosted tool calls are shown here, once the response is complete,
691713 // on both the streaming and the plain path. They never reach dispatch
@@ -727,7 +749,9 @@ func (a *Agent) drive(ctx context.Context) (string, error) {
727749 // message. Splitting them across messages is an API error, and even
728750 // where it is not it teaches the model to stop calling tools in
729751 // parallel.
730- a.state.Append(llm.Message{Role: llm.RoleUser, Content: results})
752+ answer := llm.Message{Role: llm.RoleUser, Content: results}
753+ a.state.Append(answer)
754+ a.emit(Event{Kind: EventMessage, Turn: turn, Message: &answer})
731755
732756 case llm.StopPauseTurn:
733757 // A server-side tool needs another round trip; resend as-is. Hosted
+47−0internal/agent/loop_test.goView File
@@ -1129,3 +1129,50 @@ func TestRedactedReasoningIsNotReported(t *testing.T) {
11291129 t.Error("the redacted block was dropped from the conversation")
11301130 }
11311131 }
1132+
1133+ // Every message that joins the history is reported whole, in order: the
1134+ // model's response with its id and usage, then the tool results that answer
1135+ // it. A front end printing the conversation as JSON reads nothing else.
1136+ func TestEveryMessageIsReportedAsItJoinsTheHistory(t *testing.T) {
1137+ root := t.TempDir()
1138+ os.WriteFile(filepath.Join(root, "main.go"), []byte("package main\n"), 0o644)
1139+ f := newFakeAPI(t,
1140+ msg("tool_use", textBlock("Let me look."), toolUseBlock("tu_1", "read_file", `{"path":"main.go"}`)),
1141+ msg("end_turn", textBlock("It declares package main.")),
1142+ )
1143+ a := newTestAgent(t, f, root, &bytes.Buffer{})
1144+ var got []Event
1145+ a.Observe = func(e Event) {
1146+ if e.Kind == EventMessage {
1147+ got = append(got, e)
1148+ }
1149+ }
1150+ if _, err := a.Run(context.Background(), "what does main.go do?"); err != nil {
1151+ t.Fatalf("Run: %v", err)
1152+ }
1153+ if len(got) != 3 {
1154+ t.Fatalf("got %d message events, want 3 (response, results, response)", len(got))
1155+ }
1156+ if r := got[0].Response; r == nil || r.ID != "msg" || len(r.ToolUses()) != 1 || r.Usage.InputTokens != 100 || got[0].Message != nil {
1157+ t.Errorf("first event should be the whole first response, got %+v", got[0])
1158+ }
1159+ if m := got[1].Message; m == nil || m.Role != llm.RoleUser || len(m.Content) != 1 || m.Content[0].ToolUseID != "tu_1" || got[1].Response != nil {
1160+ t.Errorf("second event should be the tool results, got %+v", got[1])
1161+ }
1162+ if r := got[2].Response; r == nil || r.Text() != "It declares package main." {
1163+ t.Errorf("third event should be the answer, got %+v", got[2])
1164+ }
1165+ }
1166+
1167+ func TestSystemAppendGoesAtTheEnd(t *testing.T) {
1168+ a := New(nil, tools.NewRegistry(nil), NewState(), nil, Options{SystemAppend: "Answer in French."})
1169+ blocks := a.systemBlocks()
1170+ text := blocks[len(blocks)-1].Text
1171+ if !strings.HasPrefix(text, defaultSystem) || !strings.HasSuffix(text, "\n\nAnswer in French.") {
1172+ t.Errorf("the appended text should follow the default prompt, got the tail %q", text[len(text)-40:])
1173+ }
1174+ replaced := New(nil, tools.NewRegistry(nil), NewState(), nil, Options{System: "Be terse.", SystemAppend: "Answer in French."})
1175+ if got := replaced.systemBlocks()[0].Text; got != "Be terse.\n\nAnswer in French." {
1176+ t.Errorf("an append goes after a replaced prompt too, got %q", got)
1177+ }
1178+ }
+28−0internal/tools/registry.goView File
@@ -117,6 +117,8 @@ type Registry struct {
117117 byName map[string]Tool
118118 order []string
119119 log *metrics.Logger
120+ // drop is what Exclude was told to keep out.
121+ drop func(name string) bool
120122 }
121123
122124 // NewRegistry returns an empty registry. log may be nil.
@@ -129,6 +131,9 @@ func (r *Registry) Register(ts ...Tool) *Registry {
129131 r.mu.Lock()
130132 defer r.mu.Unlock()
131133 for _, t := range ts {
134+ if r.drop != nil && r.drop(t.Name()) {
135+ continue
136+ }
132137 if _, dup := r.byName[t.Name()]; dup {
133138 panic(fmt.Sprintf("tools: duplicate registration of %q", t.Name()))
134139 }
@@ -138,6 +143,29 @@ func (r *Registry) Register(ts ...Tool) *Registry {
138143 return r
139144 }
140145
146+ // Exclude keeps every tool drop says yes to out of the registry: the ones
147+ // already in it are taken out, and the ones registered later are skipped. It
148+ // is how --disallowedTools works — a tool the model is never shown is one it
149+ // cannot call, which beats refusing it every time it tries — and it works on
150+ // the registry rather than on a list of names because the tools of a session
151+ // are registered from several places, and every one of them has to hear it.
152+ func (r *Registry) Exclude(drop func(name string) bool) *Registry {
153+ r.mu.Lock()
154+ defer r.mu.Unlock()
155+ prev := r.drop
156+ r.drop = func(name string) bool { return (prev != nil && prev(name)) || drop(name) }
157+ kept := r.order[:0]
158+ for _, n := range r.order {
159+ if drop(n) {
160+ delete(r.byName, n)
161+ continue
162+ }
163+ kept = append(kept, n)
164+ }
165+ r.order = kept
166+ return r
167+ }
168+
141169 // Get returns the named tool.
142170 func (r *Registry) Get(name string) (Tool, bool) {
143171 r.mu.RLock()
+23−0internal/tools/registry_test.goView File
@@ -214,3 +214,26 @@ func TestRegistryLogsEveryDispatch(t *testing.T) {
214214 t.Errorf("unknown-tool dispatch should log a failure, got %s", lines[1])
215215 }
216216 }
217+
218+ func TestExcludeKeepsToolsOutBeforeAndAfter(t *testing.T) {
219+ named := func(name string) Tool {
220+ return New(name, "x", func(context.Context, struct{}) (string, error) { return "", nil })
221+ }
222+ r := NewRegistry(nil).Register(named("bash"), named("read_file"))
223+ r.Exclude(func(n string) bool { return n == "bash" || n == "web_fetch" })
224+ r.Exclude(func(n string) bool { return strings.HasPrefix(n, "mcp__db") })
225+ r.Register(named("web_fetch"), named("mcp__db__query"), named("grep"))
226+
227+ var got []string
228+ for _, tl := range r.List() {
229+ got = append(got, tl.Name())
230+ }
231+ if strings.Join(got, ",") != "grep,read_file" {
232+ t.Errorf("tools = %v, want only grep and read_file", got)
233+ }
234+ if _, err := r.Execute(context.Background(), "bash", nil); err == nil {
235+ t.Error("an excluded tool cannot be executed")
236+ }
237+ // Registering an excluded name again is not a duplicate: it was never in.
238+ r.Register(named("bash"))
239+ }