From f3fa1337281386b03659078dc559fca7501361c7 Mon Sep 17 00:00:00 2001 From: Dwin Gharibi Date: Fri, 7 Aug 2026 08:49:44 +0330 Subject: [PATCH 1/3] feat(pkg/runtime/toolexec/dispatcher.go): adding up readonly tool caching for improving costs and tool calls count --- pkg/hooks/builtins/builtins.go | 6 ++++++ pkg/hooks/types.go | 7 +++++++ pkg/runtime/toolexec/dispatcher.go | 3 +++ 3 files changed, 16 insertions(+) diff --git a/pkg/hooks/builtins/builtins.go b/pkg/hooks/builtins/builtins.go index 3a3c512bcc..7a4f6e1921 100644 --- a/pkg/hooks/builtins/builtins.go +++ b/pkg/hooks/builtins/builtins.go @@ -30,6 +30,11 @@ // - limit_large_tool_results // (tool_response_transform) — store oversized tool output in a temp file // and replace it with a bounded tail plus notice +// - elide_repeated_tool_results +// (tool_response_transform, session_end) — replace a read-only tool's +// output with a short marker when it is byte-for-byte identical to what +// the model already saw for the same arguments in this session. The tool +// always runs, so this cannot serve stale data; it saves tokens, not I/O. // - safer_shell (pre_tool_use) — deprecated labeller // shim. The runtime classifies shell commands // natively via pkg/safety; pinned entries only @@ -112,6 +117,7 @@ func Register(r *hooks.Registry, opts ...Option) error { r.RegisterBuiltin(MaxIterations, maxIterations), r.RegisterBuiltin(RedactSecrets, redactSecrets), r.RegisterBuiltin(LimitLargeToolResults, limitLargeToolResults), + r.RegisterBuiltin(ElideRepeatedToolResults, elideRepeatedToolResults), r.RegisterBuiltin(SaferShell, saferShell), r.RegisterBuiltin(HTTPPost, newHTTPPost(o.httpPostClient)), r.RegisterBuiltin(Unload, unload), diff --git a/pkg/hooks/types.go b/pkg/hooks/types.go index 1a7c2990dd..7b22afdfce 100644 --- a/pkg/hooks/types.go +++ b/pkg/hooks/types.go @@ -274,6 +274,13 @@ type Input struct { ToolUseID string `json:"tool_use_id,omitempty"` ToolInput map[string]any `json:"tool_input,omitempty"` + // ToolReadOnly mirrors the dispatching tool's ReadOnlyHint annotation, so a + // hook can tell a tool that only observes from one that mutates something. + // False whenever the hint is absent or the tool is unknown to the agent, + // which keeps consumers fail-safe: a hook that acts only on read-only tools + // does nothing when the declaration is missing. + ToolReadOnly bool `json:"tool_read_only,omitempty"` + // SafetyPolicy mirrors the session's effective safety mode // (strict / balanced / autonomous, empty for the legacy default; // see [github.com/docker/docker-agent/pkg/session.SafetyPolicy]) diff --git a/pkg/runtime/toolexec/dispatcher.go b/pkg/runtime/toolexec/dispatcher.go index 4a2f154b98..88efc9eb30 100644 --- a/pkg/runtime/toolexec/dispatcher.go +++ b/pkg/runtime/toolexec/dispatcher.go @@ -1079,6 +1079,9 @@ func (c *call) applyToolResponseTransform(ctx context.Context, payload string, i } in := NewPostToolHooksInput(c.sess, c.tc, &tools.ToolCallResult{Output: payload, IsError: isError}) in.ToolCategory = c.tool.Category + // Zero when !c.available (the tool isn't in the agent's toolset), which is + // the fail-safe direction: consumers keyed on read-only-ness stay inert. + in.ToolReadOnly = c.tool.Annotations.ReadOnlyHint result := c.d.Hooks.Dispatch(ctx, c.a, hooks.EventToolResponseTransform, in) if result == nil || result.UpdatedToolResponse == nil { return payload From 7f5659f5f0b2a5aa962df503aa480217589ca1cd Mon Sep 17 00:00:00 2001 From: Dwin Gharibi Date: Fri, 7 Aug 2026 08:51:24 +0330 Subject: [PATCH 2/3] feat(pkg/hooks/builtins/elide_repeated_tool_results.go): adding up the elide hook to make sure cache consistency --- .../builtins/elide_repeated_tool_results.go | 181 ++++++++++++++++++ 1 file changed, 181 insertions(+) create mode 100644 pkg/hooks/builtins/elide_repeated_tool_results.go diff --git a/pkg/hooks/builtins/elide_repeated_tool_results.go b/pkg/hooks/builtins/elide_repeated_tool_results.go new file mode 100644 index 0000000000..76e94d7c65 --- /dev/null +++ b/pkg/hooks/builtins/elide_repeated_tool_results.go @@ -0,0 +1,181 @@ +package builtins + +import ( + "context" + "crypto/sha256" + "encoding/json" + "fmt" + "log/slog" + "sync" + + "github.com/docker/docker-agent/pkg/hooks" +) + +// ElideRepeatedToolResults is the registered name of the builtin +// tool_response_transform hook that stops re-sending a read-only tool's output +// when it is byte-for-byte identical to what the model already saw earlier in +// the same session. +// +// # Why this cannot serve stale data +// +// This is deliberately NOT a cache. The tool always executes and its fresh +// output is always what gets hashed; the hook only decides whether to repeat +// bytes the model has already been shown. There is no stored payload to go +// stale, no expiry to tune, and no invalidation to get wrong: if the file (or +// whatever the tool reads) changed by even one byte, the hashes differ and the +// full new output is passed through untouched. +// +// The saving is in tokens, not in I/O — a repeated 40 KiB read_file result +// becomes a one-line marker. Latency is unchanged because the tool still runs. +const ElideRepeatedToolResults = "elide_repeated_tool_results" + +const ( + // minElidableBytes is the payload size below which eliding is a net loss: + // the marker itself costs tokens, so replacing a short result with it would + // make the conversation bigger, not smaller. + minElidableBytes = 256 + + // maxElideKeysPerSession bounds per-session memory. A session that calls + // read-only tools with thousands of distinct argument sets stops recording + // new fingerprints rather than growing without limit; already-recorded keys + // keep working. Each entry is a 32-byte hash plus a map key. + maxElideKeysPerSession = 4096 +) + +// elideState remembers, per session, the fingerprint of the most recent output +// seen for each (tool, arguments) pair. +// +// Package-level state mirrors the limit_large_tool_results builtin, which keeps +// per-session scratch state for the same reason: builtins are registered as +// plain functions and have nowhere else to live. Entries are dropped on +// session_end. +type elideState struct { + mu sync.Mutex + // seen maps session ID -> call key -> sha256 of the last output. + seen map[string]map[string][sha256.Size]byte +} + +var elideStore = &elideState{seen: make(map[string]map[string][sha256.Size]byte)} + +// elideRepeatedToolResults is the [hooks.BuiltinFunc] registered under +// [ElideRepeatedToolResults]. It dispatches on the event so one YAML entry can +// cover both the transform leg and the session_end cleanup. +func elideRepeatedToolResults(_ context.Context, in *hooks.Input, _ []string) (*hooks.Output, error) { + if in == nil { + return nil, nil + } + switch in.HookEventName { + case hooks.EventToolResponseTransform: + return elideRepeatedToolResponse(in), nil + case hooks.EventSessionEnd: + elideStore.forget(in.SessionID) + return nil, nil + default: + // Lenient on misconfiguration, matching redact_secrets: log the typo + // but never fail the run loop over a misplaced hook entry. + slog.Warn("elide_repeated_tool_results builtin invoked under unsupported event; no-op", + "event", in.HookEventName) + return nil, nil + } +} + +// elideRepeatedToolResponse returns a marker in place of payloads that repeat +// an earlier identical result, or nil to leave the response untouched. +func elideRepeatedToolResponse(in *hooks.Input) *hooks.Output { + // Only tools the author declared read-only are eligible. A tool with side + // effects may legitimately return identical output for two calls that each + // did something (e.g. an append that was then undone), so eliding the + // second would hide a real event. + if !in.ToolReadOnly { + return nil + } + // An error result is diagnostic: the model needs it every time, and a + // repeated identical failure is itself information. + if in.ToolError { + return nil + } + // Without a session there is nothing to scope the state to. + if in.SessionID == "" { + return nil + } + + payload, ok := in.ToolResponse.(string) + if !ok || len(payload) < minElidableBytes { + return nil + } + + key, ok := elideCallKey(in.ToolName, in.ToolInput) + if !ok { + return nil + } + + if !elideStore.observe(in.SessionID, key, sha256.Sum256([]byte(payload))) { + return nil + } + + marker := fmt.Sprintf( + "[docker-agent] The %s tool ran and returned output byte-for-byte identical to its "+ + "earlier result for these same arguments in this session, so the %d-byte payload is "+ + "not repeated here. Nothing has changed since you last saw it.", + in.ToolName, len(payload)) + + return &hooks.Output{ + HookSpecificOutput: &hooks.HookSpecificOutput{ + HookEventName: hooks.EventToolResponseTransform, + UpdatedToolResponse: &marker, + }, + } +} + +// elideCallKey fingerprints a call as its tool name plus its arguments. +// [encoding/json] sorts map keys, so the result does not depend on Go's +// randomized map iteration order. Arguments that cannot be marshalled yield +// ok=false, which makes the caller leave the response untouched. +func elideCallKey(tool string, args map[string]any) (string, bool) { + encoded, err := json.Marshal(args) + if err != nil { + return "", false + } + h := sha256.New() + h.Write([]byte(tool)) + h.Write([]byte{0}) + h.Write(encoded) + return string(h.Sum(nil)), true +} + +// observe records sum as the latest output fingerprint for (session, key) and +// reports whether it repeats what was already recorded. A mismatch overwrites +// the stored fingerprint, so the *next* identical call elides. +func (s *elideState) observe(sessionID, key string, sum [sha256.Size]byte) bool { + s.mu.Lock() + defer s.mu.Unlock() + + perSession, ok := s.seen[sessionID] + if !ok { + perSession = make(map[string][sha256.Size]byte, 1) + s.seen[sessionID] = perSession + } + + previous, seen := perSession[key] + if seen { + if previous == sum { + return true + } + perSession[key] = sum + return false + } + + // New key: respect the per-session cap. Declining to record simply means + // this call is never elided — correctness is unaffected. + if len(perSession) < maxElideKeysPerSession { + perSession[key] = sum + } + return false +} + +// forget drops all state for a session. +func (s *elideState) forget(sessionID string) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.seen, sessionID) +} From eaf0fa380946ba0575a02a7169fb7f62a67e9f47 Mon Sep 17 00:00:00 2001 From: Dwin Gharibi Date: Fri, 7 Aug 2026 08:51:51 +0330 Subject: [PATCH 3/3] test(pkg/hooks/builtins/elide_repeated_tool_results_test.go): adding some edge case tests for new elide hook added --- .../elide_repeated_tool_results_test.go | 256 ++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 pkg/hooks/builtins/elide_repeated_tool_results_test.go diff --git a/pkg/hooks/builtins/elide_repeated_tool_results_test.go b/pkg/hooks/builtins/elide_repeated_tool_results_test.go new file mode 100644 index 0000000000..2c0e0e864b --- /dev/null +++ b/pkg/hooks/builtins/elide_repeated_tool_results_test.go @@ -0,0 +1,256 @@ +package builtins + +import ( + "strings" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/docker/docker-agent/pkg/hooks" +) + +// bigPayload returns a payload comfortably above minElidableBytes. +func bigPayload(marker string) string { + return marker + strings.Repeat("x", minElidableBytes*2) +} + +// forgetAllElideState resets the package-level store between tests. These tests +// deliberately do not run in parallel with each other: they share that store, +// which is the same state the runtime shares across a process. +func forgetAllElideState() { + elideStore.mu.Lock() + defer elideStore.mu.Unlock() + elideStore.seen = make(map[string]map[string][32]byte) +} + +// elideStoreLen reports how many call keys are recorded for a session. +func elideStoreLen(sessionID string) int { + elideStore.mu.Lock() + defer elideStore.mu.Unlock() + return len(elideStore.seen[sessionID]) +} + +func transformInput(sessionID, tool, payload string, args map[string]any) *hooks.Input { + return &hooks.Input{ + HookEventName: hooks.EventToolResponseTransform, + SessionID: sessionID, + ToolName: tool, + ToolReadOnly: true, + ToolInput: args, + ToolResponse: payload, + } +} + +func elide(t *testing.T, in *hooks.Input) *hooks.Output { + t.Helper() + out, err := elideRepeatedToolResults(t.Context(), in, nil) + require.NoError(t, err) + return out +} + +func TestElideRepeatedToolResults_FirstCallPassesThrough(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + + out := elide(t, transformInput("s1", "read_file", payload, map[string]any{"path": "a.txt"})) + assert.Nil(t, out, "the first result must reach the model in full") +} + +func TestElideRepeatedToolResults_IdenticalRepeatIsElided(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + args := map[string]any{"path": "a.txt"} + + require.Nil(t, elide(t, transformInput("s1", "read_file", payload, args))) + + out := elide(t, transformInput("s1", "read_file", payload, args)) + require.NotNil(t, out) + require.NotNil(t, out.HookSpecificOutput) + require.NotNil(t, out.HookSpecificOutput.UpdatedToolResponse) + + got := *out.HookSpecificOutput.UpdatedToolResponse + assert.NotEqual(t, payload, got) + assert.Less(t, len(got), len(payload), "the marker must be smaller than the payload it replaces") + assert.Contains(t, got, "read_file") + assert.Contains(t, got, "identical") +} + +// THE consistency property: the payload is only ever elided when the tool's +// fresh output is byte-for-byte identical to what the model already saw. The +// tool always executes, so a changed file can never be served from cache. +func TestElideRepeatedToolResults_ChangedOutputIsNeverElided(t *testing.T) { + forgetAllElideState() + args := map[string]any{"path": "a.txt"} + first := bigPayload("version-one") + second := bigPayload("version-two") + + require.Nil(t, elide(t, transformInput("s1", "read_file", first, args))) + + out := elide(t, transformInput("s1", "read_file", second, args)) + assert.Nil(t, out, "changed output must always reach the model in full") + + // And the new output becomes the baseline, so a repeat of *it* elides + // while a return to the old content does not. + require.NotNil(t, elide(t, transformInput("s1", "read_file", second, args))) + assert.Nil(t, elide(t, transformInput("s1", "read_file", first, args)), + "reverting to earlier content must reach the model in full") +} + +func TestElideRepeatedToolResults_NonReadOnlyToolIsNeverElided(t *testing.T) { + forgetAllElideState() + payload := bigPayload("side effects") + args := map[string]any{"cmd": "date"} + + in := transformInput("s1", "shell", payload, args) + in.ToolReadOnly = false + require.Nil(t, elide(t, in)) + + in2 := transformInput("s1", "shell", payload, args) + in2.ToolReadOnly = false + assert.Nil(t, elide(t, in2), "a tool with side effects must never be elided") +} + +func TestElideRepeatedToolResults_ErrorResultIsNeverElided(t *testing.T) { + forgetAllElideState() + payload := bigPayload("boom") + args := map[string]any{"path": "a.txt"} + + in := transformInput("s1", "read_file", payload, args) + in.ToolError = true + require.Nil(t, elide(t, in)) + + in2 := transformInput("s1", "read_file", payload, args) + in2.ToolError = true + assert.Nil(t, elide(t, in2)) +} + +func TestElideRepeatedToolResults_DifferentArgsAreDistinct(t *testing.T) { + forgetAllElideState() + payload := bigPayload("same bytes") + + require.Nil(t, elide(t, transformInput("s1", "read_file", payload, map[string]any{"path": "a.txt"}))) + assert.Nil(t, elide(t, transformInput("s1", "read_file", payload, map[string]any{"path": "b.txt"})), + "a different argument set is a different call") +} + +// Key building must not depend on Go's randomized map iteration order. +func TestElideRepeatedToolResults_ArgOrderIsIrrelevant(t *testing.T) { + forgetAllElideState() + payload := bigPayload("stable") + + require.Nil(t, elide(t, transformInput("s1", "read_file", payload, + map[string]any{"path": "a.txt", "line": 1, "limit": 20}))) + + for range 20 { + out := elide(t, transformInput("s1", "read_file", payload, + map[string]any{"limit": 20, "line": 1, "path": "a.txt"})) + require.NotNil(t, out, "identical args in any map order must be the same key") + } +} + +func TestElideRepeatedToolResults_SessionsAreIsolated(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + args := map[string]any{"path": "a.txt"} + + require.Nil(t, elide(t, transformInput("s1", "read_file", payload, args))) + assert.Nil(t, elide(t, transformInput("s2", "read_file", payload, args)), + "another session has not seen this output") +} + +func TestElideRepeatedToolResults_SmallPayloadNotWorthEliding(t *testing.T) { + forgetAllElideState() + small := "tiny" + args := map[string]any{"path": "a.txt"} + + require.Nil(t, elide(t, transformInput("s1", "read_file", small, args))) + assert.Nil(t, elide(t, transformInput("s1", "read_file", small, args)), + "eliding a payload smaller than the marker would cost tokens, not save them") +} + +func TestElideRepeatedToolResults_SessionEndForgetsState(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + args := map[string]any{"path": "a.txt"} + + require.Nil(t, elide(t, transformInput("s1", "read_file", payload, args))) + require.NotNil(t, elide(t, transformInput("s1", "read_file", payload, args))) + + _, err := elideRepeatedToolResults(t.Context(), &hooks.Input{ + HookEventName: hooks.EventSessionEnd, + SessionID: "s1", + }, nil) + require.NoError(t, err) + + assert.Nil(t, elide(t, transformInput("s1", "read_file", payload, args)), + "state must be dropped when the session ends") +} + +func TestElideRepeatedToolResults_PerSessionKeyCapIsBounded(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + + // Fill past the cap with distinct argument sets. + for i := range maxElideKeysPerSession + 50 { + elide(t, transformInput("s1", "read_file", payload, map[string]any{"path": i})) + } + assert.LessOrEqual(t, elideStoreLen("s1"), maxElideKeysPerSession, + "per-session key count must stay bounded") +} + +// Parallel tool calls dispatch this hook concurrently; run under -race. +func TestElideRepeatedToolResults_ConcurrentDispatch(t *testing.T) { + forgetAllElideState() + payload := bigPayload("contents") + + var wg sync.WaitGroup + for i := range 32 { + wg.Go(func() { + for range 8 { + _, err := elideRepeatedToolResults(t.Context(), + transformInput("s1", "read_file", payload, map[string]any{"path": i % 4}), nil) + assert.NoError(t, err) + } + }) + } + wg.Wait() +} + +func TestElideRepeatedToolResults_IsRegistered(t *testing.T) { + forgetAllElideState() + reg := hooks.NewRegistry() + require.NoError(t, Register(reg)) + + handler, ok := reg.LookupBuiltin(ElideRepeatedToolResults) + require.Truef(t, ok, "builtin %q must be registered", ElideRepeatedToolResults) + + payload := bigPayload("contents") + args := map[string]any{"path": "a.txt"} + + first, err := handler(t.Context(), transformInput("s9", "read_file", payload, args), nil) + require.NoError(t, err) + require.Nil(t, first) + + second, err := handler(t.Context(), transformInput("s9", "read_file", payload, args), nil) + require.NoError(t, err) + require.NotNil(t, second) +} + +func TestElideRepeatedToolResults_UnsupportedEventIsNoOp(t *testing.T) { + forgetAllElideState() + out, err := elideRepeatedToolResults(t.Context(), &hooks.Input{ + HookEventName: hooks.EventTurnStart, + SessionID: "s1", + }, nil) + require.NoError(t, err) + assert.Nil(t, out) +} + +func TestElideRepeatedToolResults_NilInput(t *testing.T) { + forgetAllElideState() + out, err := elideRepeatedToolResults(t.Context(), nil, nil) + require.NoError(t, err) + assert.Nil(t, out) +}