From 39edf8b77db4f7c9df210b3b7db38058e2c820c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Wed, 14 Oct 2026 08:27:50 +0000 Subject: [PATCH 1/2] fix: isolate think scratchpads by conversation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Persist caller- and instance-scoped thoughts in the session store, bind tool execution to the durable conversation root, and cover reload, concurrent isolation, and cleanup. Preserve Origin in same-ID session clones so retained thoughts remain addressable. Signed-off-by: Arnaud Héritier --- docs/tools/think/index.md | 30 +- pkg/runtime/loop.go | 5 +- pkg/runtime/think_test.go | 562 ++++++++++++++++++++++++ pkg/runtime/tool_dispatch.go | 107 ++++- pkg/runtime/toolexec/dispatcher.go | 10 + pkg/runtime/toolexec/dispatcher_test.go | 38 ++ pkg/session/branch.go | 1 + pkg/session/clone_test.go | 2 + pkg/session/migrations.go | 18 + pkg/session/migrations_pinned_test.go | 2 +- pkg/session/store.go | 21 +- pkg/session/think_scratchpad.go | 129 ++++++ pkg/session/think_scratchpad_test.go | 238 ++++++++++ pkg/tools/builtin/think/think.go | 28 +- pkg/tools/builtin/think/think_test.go | 75 +++- pkg/tools/runtime.go | 6 + 16 files changed, 1235 insertions(+), 37 deletions(-) create mode 100644 pkg/runtime/think_test.go create mode 100644 pkg/session/think_scratchpad.go create mode 100644 pkg/session/think_scratchpad_test.go diff --git a/docs/tools/think/index.md b/docs/tools/think/index.md index 24873fcacf..7ed1e0ed0e 100644 --- a/docs/tools/think/index.md +++ b/docs/tools/think/index.md @@ -13,7 +13,7 @@ _Step-by-step reasoning scratchpad for planning and decision-making._ The think tool is a reasoning scratchpad that lets agents think step-by-step before acting. The agent can write its thoughts without producing visible output to the user — ideal for planning complex tasks, breaking down problems, and reasoning through multi-step solutions. -This is a lightweight tool with no side effects. It is most useful for models that lack built-in reasoning or thinking capabilities (e.g., smaller or older models). For models that already support native thinking — such as Claude with extended thinking, OpenAI o-series, or Gemini with a thinking budget — this tool is unnecessary since the model can reason internally. +This lightweight tool changes only its internal scratchpad, not external data. It is most useful for models that lack built-in reasoning or thinking capabilities (e.g., smaller or older models). For models that already support native thinking — such as Claude with extended thinking, OpenAI o-series, or Gemini with a thinking budget — this tool is unnecessary since the model can reason internally. ## Configuration @@ -28,3 +28,31 @@ No configuration options. > **When to use** > > Use the think tool with models that don't have native reasoning capabilities. If your model already supports a [thinking budget](../../configuration/models/index.md#thinking-budget), you likely don't need this tool. + +## Scope and lifetime + +Thoughts are isolated by conversation ID, calling agent and think toolset instance. +They accumulate across turns, runtime closure and compaction. The built-in SQLite +session store retains raw thoughts across reloads and restarts; the in-memory +store retains them only for its lifetime. This internal state is stored before +response-output redaction, so session-database and backup privacy policies apply. + +A branch or fork with a new conversation ID starts empty. A same-ID session clone +addresses the same durable scratchpad; cloning messages does not provide a +transactional copy of hidden thoughts. JSON export/import does not transport the +scratchpad, and existing conversations start with empty dedicated state on upgrade. + +Toolset namespaces follow configuration order (configured and extra toolsets are +separate). Resuming with reordered toolsets or another configuration with the same +agent name can reassign these namespaces. Stable identity requires unchanged topology. + +Root deletion removes all root-owned thoughts, including those from aborted +sub-sessions. Deleting a persisted child removes known subtree thoughts; +never-persisted descendants remain root-owned until root deletion. Stop writers +before deleting sessions. Deletion is logical, not secure erasure or backup removal. + +There is no TTL, eviction, quota or output omission: every call returns the full +scratchpad, which can grow beyond model context limits. Storage and output memory +scale with retained thoughts. Persistence, ownership or unsupported-store failures +are tool errors, never silent resets. Only genuinely unscoped hosts, such as +pre-run prompt expansion, return an ephemeral single thought. diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 4d9a39ca8a..3dcd9d73ec 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -283,13 +283,15 @@ func (r *LocalRuntime) runStream(ctx context.Context, sess *session.Session, idl // the whole lifetime of its stream. entry := r.registerLiveSessionWithIdleRetry(sess, idleRetryPolicy) + // Persist the root before tools can append root-owned state. + out := r.observe(ctx, sess, events) go func() { if rootStream { defer r.activeRootStreams.Add(-1) } r.runStreamLoop(ctx, sess, entry, events) }() - return r.observe(ctx, sess, events) + return out } // runStreamLoop is the body of RunStream. Pulled out of the anonymous @@ -297,6 +299,7 @@ func (r *LocalRuntime) runStream(ctx context.Context, sess *session.Session, idl // in editors. func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, liveEntry *liveSessionEntry, events chan Event) { sink := &channelSink{ch: events} + ctx = context.WithValue(ctx, scratchpadRootKey{}, liveEntry.treeRootID) // Seed the cagent session ID at the run-loop boundary so any // gateway-bound HTTP call originating from this loop can correlate diff --git a/pkg/runtime/think_test.go b/pkg/runtime/think_test.go new file mode 100644 index 0000000000..fc1fec8a2b --- /dev/null +++ b/pkg/runtime/think_test.go @@ -0,0 +1,562 @@ +package runtime + +import ( + "context" + "encoding/json" + "fmt" + "path/filepath" + "strings" + "sync" + "testing" + "testing/synctest" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "go.opentelemetry.io/otel/baggage" + + "github.com/docker/docker-agent/pkg/agent" + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/httpclient" + "github.com/docker/docker-agent/pkg/session" + "github.com/docker/docker-agent/pkg/session/sqlitestore" + "github.com/docker/docker-agent/pkg/team" + "github.com/docker/docker-agent/pkg/tools" + "github.com/docker/docker-agent/pkg/tools/builtin/think" + "github.com/docker/docker-agent/pkg/tools/builtin/transfertask" + "github.com/docker/docker-agent/pkg/tools/codemode" +) + +func thinkStream(thought string) chat.MessageStream { + return newStreamBuilder().AddToolCallName("think-call", "think"). + AddToolCallArguments("think-call", fmt.Sprintf(`{"thought":%q}`, thought)).AddToolCallStopWithUsage(1, 1).Build() +} + +func thinkStop() chat.MessageStream { + return newStreamBuilder().AddContent("done").AddStopWithUsage(1, 1).Build() +} + +func thinkRuntime(t *testing.T, tm *team.Team, store session.Store, current string) *LocalRuntime { + t.Helper() + rt, err := NewLocalRuntime(t.Context(), tm, WithCurrentAgent(current), WithSessionStore(store), WithSessionCompaction(false), WithModelStore(mockModelStore{})) + require.NoError(t, err) + return rt +} + +func drainThink(t *testing.T, rt *LocalRuntime, sess *session.Session) []string { + t.Helper() + var outputs []string + for ev := range rt.RunStream(t.Context(), sess) { + switch ev := ev.(type) { + case *ToolCallResponseEvent: + if ev.ToolDefinition.Name == "think" { + assert.False(t, ev.Result.IsError, ev.Response) + outputs = append(outputs, ev.Response) + } + case *ErrorEvent: + t.Errorf("runtime error: %v", ev) + } + } + return outputs +} + +func TestThinkConversationIsolationAfterClose(t *testing.T) { + t.Parallel() + prov := &queueProvider{streams: []chat.MessageStream{thinkStream("first"), thinkStop(), thinkStream("second"), thinkStop()}} + root := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(think.New())) + tm := team.New(team.WithAgents(root)) + store := session.NewInMemorySessionStore() + rt := thinkRuntime(t, tm, store, "root") + one := session.New(session.WithUserMessage("start")) + one.ToolsApproved = true + assert.Equal(t, []string{"Thoughts:\nfirst"}, drainThink(t, rt, one)) + require.NoError(t, rt.Close()) + rt = thinkRuntime(t, tm, store, "root") + two := session.New(session.WithUserMessage("start")) + two.ToolsApproved = true + assert.Equal(t, []string{"Thoughts:\nsecond"}, drainThink(t, rt, two)) + require.NoError(t, rt.Close()) +} + +func TestThinkRetentionReloadCloneFork(t *testing.T) { + t.Parallel() + for _, kind := range []string{"memory", "sqlite"} { + t.Run(kind, func(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "session.db") + store := session.NewInMemorySessionStore() + if kind == "sqlite" { + var err error + store, err = sqlitestore.New(t.Context(), path) + require.NoError(t, err) + } + prov := &queueProvider{streams: []chat.MessageStream{ + thinkStream("first"), thinkStop(), thinkStream("second"), thinkStop(), + thinkStream("third"), thinkStop(), thinkStream("fork"), thinkStop(), + }} + root := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(think.New())) + tm := team.New(team.WithAgents(root)) + rt := thinkRuntime(t, tm, store, "root") + sess := session.New(session.WithUserMessage("start")) + sess.ToolsApproved = true + stale := sess.Clone() + stale.Origin = sess.Origin + assert.Equal(t, []string{"Thoughts:\nfirst"}, drainThink(t, rt, sess)) + assert.Equal(t, []string{"Thoughts:\nfirst\nsecond"}, drainThink(t, rt, sess)) + require.NoError(t, rt.Close()) + require.NoError(t, store.UpdateSession(t.Context(), stale)) + if kind == "sqlite" { + require.NoError(t, store.Close()) + var err error + store, err = sqlitestore.New(t.Context(), path) + require.NoError(t, err) + } + t.Cleanup(func() { require.NoError(t, store.Close()) }) + reloaded, err := store.GetSession(t.Context(), sess.ID) + require.NoError(t, err) + reloaded.ToolsApproved = true + rt = thinkRuntime(t, tm, store, "root") + clone := reloaded.Clone() + assert.Equal(t, []string{"Thoughts:\nfirst\nsecond\nthird"}, drainThink(t, rt, clone)) + fork, err := session.ForkSession(clone, 0) + require.NoError(t, err) + fork.ToolsApproved = true + assert.Equal(t, []string{"Thoughts:\nfork"}, drainThink(t, rt, fork)) + require.NoError(t, rt.Close()) + }) + } +} + +func TestThinkParallelBatchAndAgents(t *testing.T) { + t.Parallel() + batch := newStreamBuilder().AddToolCallName("a", "think").AddToolCallArguments("a", `{"thought":"a"}`). + AddToolCallName("b", "think").AddToolCallArguments("b", `{"thought":"b"}`).AddToolCallStopWithUsage(1, 1).Build() + oneProvider := &queueProvider{streams: []chat.MessageStream{batch, thinkStop(), thinkStream("final"), thinkStop()}} + twoProvider := &queueProvider{streams: []chat.MessageStream{thinkStream("worker"), thinkStop()}} + shared := think.New() + one := agent.New("root", "test", agent.WithModel(oneProvider), agent.WithToolSets(shared)) + two := agent.New("worker", "test", agent.WithModel(twoProvider), agent.WithToolSets(shared)) + tm := team.New(team.WithAgents(one, two)) + store := session.NewInMemorySessionStore() + rt := thinkRuntime(t, tm, store, "root") + sess := session.New(session.WithUserMessage("start")) + sess.ToolsApproved = true + outputs := drainThink(t, rt, sess) + require.Len(t, outputs, 2) + assert.ElementsMatch(t, []int{2, 3}, []int{len(strings.Split(outputs[0], "\n")), len(strings.Split(outputs[1], "\n"))}) + outputs = drainThink(t, rt, sess) + require.Len(t, outputs, 1) + lines := strings.Split(outputs[0], "\n") + require.Len(t, lines, 4) + assert.ElementsMatch(t, []string{"a", "b"}, lines[1:3]) + assert.Equal(t, "final", lines[3]) + require.NoError(t, rt.Close()) + rt = thinkRuntime(t, tm, store, "worker") + sess.AgentName = "worker" + assert.Equal(t, []string{"Thoughts:\nworker"}, drainThink(t, rt, sess)) + require.NoError(t, rt.Close()) +} + +func dispatchThink(t *testing.T, rt *LocalRuntime, ctx context.Context, sess *session.Session, caller *agent.Agent, ts tools.ToolSet, thought string) *tools.ToolCallResult { + t.Helper() + available, err := ts.Tools(t.Context()) + require.NoError(t, err) + events := make(chan Event, 64) + rt.processToolCalls(ctx, sess, caller, []tools.ToolCall{{ID: "call", Type: "function", Function: tools.FunctionCall{Name: "think", Arguments: fmt.Sprintf(`{"thought":%q}`, thought)}}}, available, NewChannelSink(events)) + close(events) + var result *tools.ToolCallResult + for ev := range events { + if ev, ok := ev.(*ToolCallResponseEvent); ok { + result = ev.Result + } + } + require.NotNil(t, result) + return result +} + +func TestThinkInstanceExtraAndCallerBinding(t *testing.T) { + t.Parallel() + first, second, extra := think.New(), think.New(), think.New() + caller := agent.New("root", "test", agent.WithModel(&mockProvider{stream: &mockStream{}}), agent.WithToolSets(codemode.Wrap(first), second)) + other := agent.New("worker", "test", agent.WithModel(&mockProvider{stream: &mockStream{}}), agent.WithToolSets(first)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller, other)), session.NewInMemorySessionStore(), "worker") + sess := session.New(session.WithExtraToolSets([]tools.ToolSet{first, extra})) + sess.ToolsApproved = true + require.NoError(t, rt.sessionStore.UpdateSession(t.Context(), sess)) + assert.Equal(t, map[*think.ToolSet]string{first: "cfg:0", second: "cfg:1", extra: "extra:0"}, thinkOwners(caller, sess.ExtraToolSets)) + for _, ts := range []*think.ToolSet{first, second, extra} { + res := dispatchThink(t, rt, t.Context(), sess, caller, ts, "one") + assert.False(t, res.IsError) + assert.Equal(t, "Thoughts:\none", res.Output) + } + res := dispatchThink(t, rt, t.Context(), sess, caller, first, "two") + assert.Equal(t, "Thoughts:\none\ntwo", res.Output) + res = dispatchThink(t, rt, t.Context(), sess, other, first, "worker") + assert.Equal(t, "Thoughts:\nworker", res.Output) + res = dispatchThink(t, rt, t.Context(), sess, caller, think.New(), "unknown") + assert.True(t, res.IsError) + assert.Contains(t, res.Output, "unknown") +} + +func TestThinkCodeModeSharesDirectScope(t *testing.T) { + t.Parallel() + ts := think.New() + wrapped := codemode.Wrap(ts) + script, err := json.Marshal(codemode.RunToolsWithJavascriptArgs{Script: `return await think({thought: "nested"});`}) + require.NoError(t, err) + prov := &queueProvider{streams: []chat.MessageStream{ + newStreamBuilder().AddToolCallName("outer", "run_tools_with_javascript").AddToolCallArguments("outer", string(script)).AddToolCallStopWithUsage(1, 1).Build(), thinkStop(), + }} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(wrapped)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), session.NewInMemorySessionStore(), "root") + sess := session.New(session.WithUserMessage("start")) + sess.ToolsApproved = true + assert.Equal(t, []string{"Thoughts:\nnested"}, drainThink(t, rt, sess)) + res := dispatchThink(t, rt, t.Context(), sess, caller, ts, "direct") + assert.Equal(t, "Thoughts:\nnested\ndirect", res.Output) + require.NoError(t, rt.Close()) +} + +func TestThinkChildRootAndMissingAncestors(t *testing.T) { + t.Parallel() + ts := think.New() + prov := &queueProvider{streams: []chat.MessageStream{thinkStream("child"), thinkStop()}} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(ts)) + store := session.NewInMemorySessionStore() + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), store, "root") + root := session.New(session.WithOrigin("a2a")) + parent := session.New() + parent.ParentID = root.ID + for _, sess := range []*session.Session{root, parent} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + child := session.New(session.WithUserMessage("start")) + child.ParentID = parent.ID + child.ToolsApproved = true + assert.Equal(t, []string{"Thoughts:\nchild"}, drainThink(t, rt, child)) + res := dispatchThink(t, rt, t.Context(), child, caller, ts, "fallback") + assert.False(t, res.IsError) + assert.Equal(t, "Thoughts:\nchild\nfallback", res.Output) + require.NoError(t, store.DeleteSession(t.Context(), root.ID)) + res = dispatchThink(t, rt, t.Context(), child, caller, ts, "missing") + assert.True(t, res.IsError) + require.NoError(t, store.UpdateSession(t.Context(), root)) + require.NoError(t, store.UpdateSession(t.Context(), parent)) + res = dispatchThink(t, rt, t.Context(), child, caller, ts, "after deletion") + assert.Equal(t, "Thoughts:\nafter deletion", res.Output) + parent.ParentID = parent.ID + require.NoError(t, store.UpdateSession(t.Context(), parent)) + res = dispatchThink(t, rt, t.Context(), child, caller, ts, "cycle") + assert.True(t, res.IsError) + require.NoError(t, rt.Close()) +} + +type noThinkStore struct{ session.Store } + +func TestThinkMissingStoreCapabilityAndRoot(t *testing.T) { + t.Parallel() + ts := think.New() + caller := agent.New("root", "test", agent.WithModel(&mockProvider{stream: &mockStream{}}), agent.WithToolSets(ts)) + tm := team.New(team.WithAgents(caller)) + sess := session.New() + sess.ToolsApproved = true + rt := thinkRuntime(t, tm, noThinkStore{session.NewInMemorySessionStore()}, "root") + res := dispatchThink(t, rt, t.Context(), sess, caller, ts, "unsupported") + assert.True(t, res.IsError) + assert.Contains(t, res.Output, "does not support") + rt = thinkRuntime(t, tm, session.NewInMemorySessionStore(), "root") + res = dispatchThink(t, rt, t.Context(), sess, caller, ts, "missing root") + assert.True(t, res.IsError) + assert.Contains(t, res.Output, "think root unavailable") +} + +type renamedThink struct { + tools.ToolSet + + name string +} + +func (r renamedThink) Unwrap() tools.ToolSet { return r.ToolSet } + +func (r renamedThink) Tools(ctx context.Context) ([]tools.Tool, error) { + list, err := r.ToolSet.Tools(ctx) + if err == nil { + list[0].Name = r.name + } + return list, err +} + +func TestThinkCodeModeDistinctInstances(t *testing.T) { + t.Parallel() + wrapped := codemode.Wrap(renamedThink{ToolSet: think.New(), name: "first_think"}, renamedThink{ToolSet: think.New(), name: "second_think"}) + script, err := json.Marshal(codemode.RunToolsWithJavascriptArgs{Script: `await first_think({thought:"one"}); await second_think({thought:"two"}); return await first_think({thought:"three"});`}) + require.NoError(t, err) + prov := &queueProvider{streams: []chat.MessageStream{ + newStreamBuilder().AddToolCallName("outer", "run_tools_with_javascript").AddToolCallArguments("outer", string(script)).AddToolCallStopWithUsage(1, 1).Build(), thinkStop(), + }} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(wrapped)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), session.NewInMemorySessionStore(), "root") + sess := session.New(session.WithUserMessage("start")) + sess.ToolsApproved = true + var results []string + for ev := range rt.RunStream(t.Context(), sess) { + if ev, ok := ev.(*ToolCallResponseEvent); ok && (ev.ToolDefinition.Name == "first_think" || ev.ToolDefinition.Name == "second_think") { + assert.False(t, ev.Result.IsError, ev.Response) + results = append(results, ev.Response) + } + } + assert.Equal(t, []string{"Thoughts:\none", "Thoughts:\ntwo", "Thoughts:\none\nthree"}, results) + require.NoError(t, rt.Close()) +} + +func TestThinkParallelReloadedConversations(t *testing.T) { + t.Parallel() + store, err := sqlitestore.New(t.Context(), filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, store.Close()) }) + root := session.New(session.WithOrigin("a2a"), session.WithUserMessage("start")) + root.ToolsApproved = true + require.NoError(t, store.UpdateSession(t.Context(), root)) + shared := think.New() + var wg sync.WaitGroup + outputs := make([][]string, 2) + for i := range outputs { + loaded, err := store.GetSession(t.Context(), root.ID) + require.NoError(t, err) + prov := &queueProvider{streams: []chat.MessageStream{thinkStream(fmt.Sprintf("parallel-%d", i)), thinkStop()}} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(shared)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), store, "root") + wg.Go(func() { outputs[i] = drainThink(t, rt, loaded); assert.NoError(t, rt.Close()) }) + } + wg.Wait() + for _, output := range outputs { + require.Len(t, output, 1) + } + assert.ElementsMatch(t, []int{2, 3}, []int{len(strings.Split(outputs[0][0], "\n")), len(strings.Split(outputs[1][0], "\n"))}) + loaded, err := store.GetSession(t.Context(), root.ID) + require.NoError(t, err) + prov := &queueProvider{streams: []chat.MessageStream{thinkStream("final"), thinkStop()}} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(shared)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), store, "root") + result := drainThink(t, rt, loaded) + require.Len(t, result, 1) + assert.ElementsMatch(t, []string{"parallel-0", "parallel-1", "final"}, strings.Split(result[0], "\n")[1:]) + require.NoError(t, rt.Close()) +} + +type thinkStartedProvider struct { + queueProvider + + started chan struct{} + once sync.Once +} + +func (p *thinkStartedProvider) CreateChatCompletionStream(ctx context.Context, messages []chat.Message, available []tools.Tool) (chat.MessageStream, error) { + p.once.Do(func() { close(p.started) }) + return p.queueProvider.CreateChatCompletionStream(ctx, messages, available) +} + +func TestThinkRootObserverPrecedesProducer(t *testing.T) { + t.Parallel() + synctest.Test(t, func(t *testing.T) { + prov := &thinkStartedProvider{queueProvider: queueProvider{streams: []chat.MessageStream{thinkStream("ready"), thinkStop()}}, started: make(chan struct{})} + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(think.New())) + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), session.NewInMemorySessionStore(), "root") + release, entered, finished := make(chan struct{}), make(chan struct{}), make(chan struct{}) + rt.observers = append(rt.observers, &fnObserver{onStart: func(context.Context, *session.Session) { close(entered); <-release }}) + sess := session.New(session.WithUserMessage("start")) + sess.ToolsApproved = true + go func() { defer close(finished); assert.Equal(t, []string{"Thoughts:\nready"}, drainThink(t, rt, sess)) }() + <-entered + synctest.Wait() + select { + case <-prov.started: + t.Error("producer started before synchronous root observers completed") + default: + } + close(release) + <-finished + require.NoError(t, rt.Close()) + }) +} + +func TestThinkTelemetryBaggageDoesNotOwnState(t *testing.T) { + t.Parallel() + ts := think.New() + caller := agent.New("root", "test", agent.WithModel(&mockProvider{stream: &mockStream{}}), agent.WithToolSets(ts)) + store := session.NewInMemorySessionStore() + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), store, "root") + member, err := baggage.NewMember("gen_ai.conversation.id", "attacker-selected") + require.NoError(t, err) + bag, err := baggage.New(member) + require.NoError(t, err) + ctx := baggage.ContextWithBaggage(t.Context(), bag) + for _, thought := range []string{"one", "two"} { + sess := session.New() + sess.ToolsApproved = true + require.NoError(t, store.UpdateSession(t.Context(), sess)) + res := dispatchThink(t, rt, ctx, sess, caller, ts, thought) + assert.False(t, res.IsError) + assert.Equal(t, "Thoughts:\n"+thought, res.Output) + } +} + +type thinkRecordingStore struct { + session.Store + + mu sync.Mutex + roots map[string]string +} + +func (s *thinkRecordingStore) AppendThought(ctx context.Context, rootID string, key session.ScratchpadKey, thought string) ([]string, error) { + out, err := s.Store.(session.ThinkScratchpadStore).AppendThought(ctx, rootID, key, thought) + if err == nil { + s.mu.Lock() + s.roots[key.SessionID] = rootID + s.mu.Unlock() + } + return out, err +} + +func TestThinkResumedChildDelegates(t *testing.T) { + t.Parallel() + for _, kind := range []string{"memory", "sqlite"} { + t.Run(kind, func(t *testing.T) { + t.Parallel() + store := session.NewInMemorySessionStore() + if kind == "sqlite" { + var err error + store, err = sqlitestore.New(t.Context(), filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + } + t.Cleanup(func() { require.NoError(t, store.Close()) }) + recorded := &thinkRecordingStore{Store: store, roots: make(map[string]string)} + shared := think.New() + worker := agent.New("worker", "test", agent.WithModel(&queueProvider{streams: []chat.MessageStream{thinkStream("delegated"), thinkStop()}}), agent.WithToolSets(shared)) + transfer := newStreamBuilder().AddToolCallName("transfer", "transfer_task"). + AddToolCallArguments("transfer", transferToolCall("worker").Function.Arguments).AddToolCallStopWithUsage(1, 1).Build() + caller := agent.New("root", "test", agent.WithModel(&queueProvider{streams: []chat.MessageStream{thinkStream("resumed"), transfer, thinkStop()}}), + agent.WithToolSets(shared, transfertask.New()), agent.WithSubAgents(worker)) + rt := thinkRuntime(t, team.New(team.WithAgents(caller, worker)), recorded, "root") + root := session.New(session.WithOrigin("a2a")) + intermediate := session.New(session.WithParentID(root.ID)) + for _, sess := range []*session.Session{root, intermediate} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + // Resumed streams need not re-persist their own row before delegation. + resumed := session.New(session.WithParentID(intermediate.ID), session.WithUserMessage("start"), session.WithToolsApproved(true)) + assert.Equal(t, []string{"Thoughts:\nresumed", "Thoughts:\ndelegated"}, drainThink(t, rt, resumed)) + delegated := firstSubSession(resumed) + require.NotNil(t, delegated) + assert.Equal(t, map[string]string{resumed.ID: root.ID, delegated.ID: root.ID}, recorded.roots) + require.NoError(t, rt.Close()) + require.NoError(t, store.DeleteSession(t.Context(), root.ID)) + for _, id := range []string{resumed.ID, delegated.ID} { + _, err := recorded.AppendThought(t.Context(), root.ID, session.ScratchpadKey{SessionID: id, AgentName: "root", ToolsetKey: "cfg:0"}, "deleted") + require.ErrorIs(t, err, session.ErrScratchpadOwner) + } + require.NoError(t, store.UpdateSession(t.Context(), root)) + for name, sess := range map[string]*session.Session{"root": resumed, "worker": delegated} { + out, err := recorded.AppendThought(t.Context(), root.ID, session.ScratchpadKey{SessionID: sess.ID, AgentName: name, ToolsetKey: "cfg:0"}, "after deletion") + require.NoError(t, err) + assert.Equal(t, []string{"after deletion"}, out) + } + }) + } +} + +type thinkOverlappingProvider struct { + queueProvider + + mu sync.Mutex + calls map[string]int + arrived [2]chan string + release [2]chan struct{} +} + +func (p *thinkOverlappingProvider) CreateChatCompletionStream(ctx context.Context, _ []chat.Message, _ []tools.Tool) (chat.MessageStream, error) { + id := httpclient.SessionIDFromContext(ctx) + p.mu.Lock() + call := p.calls[id] + p.calls[id]++ + p.mu.Unlock() + if call >= 2 { + return thinkStop(), nil + } + p.arrived[call] <- id + select { + case <-p.release[call]: + return thinkStream(fmt.Sprintf("%s:%d", id, call)), nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func TestThinkOverlappingConversationIsolation(t *testing.T) { + t.Parallel() + for _, kind := range []string{"memory", "sqlite"} { + t.Run(kind, func(t *testing.T) { + t.Parallel() + store := session.NewInMemorySessionStore() + if kind == "sqlite" { + var err error + store, err = sqlitestore.New(t.Context(), filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + } + t.Cleanup(func() { require.NoError(t, store.Close()) }) + recorded := &thinkRecordingStore{Store: store, roots: make(map[string]string)} + prov := &thinkOverlappingProvider{calls: make(map[string]int)} + for i := range prov.arrived { + prov.arrived[i] = make(chan string, 2) + prov.release[i] = make(chan struct{}) + } + shared := think.New() + available, err := shared.Tools(t.Context()) + require.NoError(t, err) + // Static advertisement avoids collectTools' concurrent TryIsStarted skips. + caller := agent.New("root", "test", agent.WithModel(prov), agent.WithToolSets(shared), agent.WithTools(available...)) + tm := team.New(team.WithAgents(caller)) + sessions := []*session.Session{ + session.New(session.WithUserMessage("one"), session.WithToolsApproved(true)), + session.New(session.WithUserMessage("two"), session.WithToolsApproved(true)), + } + var wg sync.WaitGroup + outputs := make([][]string, len(sessions)) + for i, sess := range sessions { + rt := thinkRuntime(t, tm, recorded, "root") + wg.Go(func() { outputs[i] = drainThink(t, rt, sess); assert.NoError(t, rt.Close()) }) + } + for i := range prov.arrived { + assert.ElementsMatch(t, []string{sessions[0].ID, sessions[1].ID}, []string{<-prov.arrived[i], <-prov.arrived[i]}) + close(prov.release[i]) + } + wg.Wait() + for i, sess := range sessions { + assert.Equal(t, []string{"Thoughts:\n" + sess.ID + ":0", "Thoughts:\n" + sess.ID + ":0\n" + sess.ID + ":1"}, outputs[i]) + out, err := recorded.AppendThought(t.Context(), sess.ID, session.ScratchpadKey{SessionID: sess.ID, AgentName: caller.Name(), ToolsetKey: "cfg:0"}, "verify") + require.NoError(t, err) + assert.Equal(t, []string{sess.ID + ":0", sess.ID + ":1", "verify"}, out) + } + assert.Equal(t, map[string]string{sessions[0].ID: sessions[0].ID, sessions[1].ID: sessions[1].ID}, recorded.roots) + }) + } +} + +func TestThinkRootIgnoresInheritedCandidate(t *testing.T) { + t.Parallel() + ts := think.New() + caller := agent.New("root", "test", agent.WithModel(&mockProvider{stream: &mockStream{}}), agent.WithToolSets(ts)) + store := session.NewInMemorySessionStore() + rt := thinkRuntime(t, team.New(team.WithAgents(caller)), store, "root") + root, other := session.New(session.WithToolsApproved(true)), session.New() + for _, sess := range []*session.Session{root, other} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + ctx := context.WithValue(t.Context(), scratchpadRootKey{}, other.ID) + res := dispatchThink(t, rt, ctx, root, caller, ts, "own root") + assert.False(t, res.IsError, res.Output) + require.NoError(t, store.DeleteSession(t.Context(), other.ID)) + res = dispatchThink(t, rt, t.Context(), root, caller, ts, "retained") + assert.False(t, res.IsError, res.Output) + assert.Equal(t, "Thoughts:\nown root\nretained", res.Output) + require.NoError(t, rt.Close()) +} diff --git a/pkg/runtime/tool_dispatch.go b/pkg/runtime/tool_dispatch.go index 8120ad9788..fb81d0a294 100644 --- a/pkg/runtime/tool_dispatch.go +++ b/pkg/runtime/tool_dispatch.go @@ -2,6 +2,8 @@ package runtime import ( "context" + "errors" + "fmt" "uuid" "github.com/docker/docker-agent/pkg/agent" @@ -15,6 +17,7 @@ import ( "github.com/docker/docker-agent/pkg/tools/builtin/handoff" "github.com/docker/docker-agent/pkg/tools/builtin/modelpicker" "github.com/docker/docker-agent/pkg/tools/builtin/skills" + "github.com/docker/docker-agent/pkg/tools/builtin/think" "github.com/docker/docker-agent/pkg/tools/builtin/transfertask" ) @@ -66,12 +69,13 @@ func (r *LocalRuntime) processToolCalls(ctx context.Context, sess *session.Sessi } d := &toolexec.Dispatcher{ - Tracer: r.tracer, - Hooks: &hookDispatcher{r: r, events: events}, - Resume: r.resumeChan, - AgentFor: func(*session.Session) *agent.Agent { return caller }, - Permissions: r.permissionCheckers, - Handlers: handlers, + AppendThought: r.thoughtAppender(ctx, sess, caller), + Tracer: r.tracer, + Hooks: &hookDispatcher{r: r, events: events}, + Resume: r.resumeChan, + AgentFor: func(*session.Session) *agent.Agent { return caller }, + Permissions: r.permissionCheckers, + Handlers: handlers, Recall: func(ctx context.Context, _ *session.Session, _ *agent.Agent, message string) error { return r.recall(ctx, QueuedMessage{Content: message}) }, @@ -201,3 +205,94 @@ func addAgentMessage(sess *session.Session, a *agent.Agent, msg *chat.Message, e } events.Emit(MessageAdded(sess.ID, agentMsg, a.Name())) } + +// scratchpadRootKey carries this stream's immutable lineage, never telemetry baggage. +type scratchpadRootKey struct{} + +func thinkOwners(caller *agent.Agent, extras []tools.ToolSet) map[*think.ToolSet]string { + owners := make(map[*think.ToolSet]string) + for category, sets := range [][]tools.ToolSet{caller.ToolSets(), extras} { + prefix := "cfg" + if category == 1 { + prefix = "extra" + } + ordinal := 0 + for _, ts := range sets { + tools.Walk(ts, func(ts tools.ToolSet) bool { + if owner, ok := ts.(*think.ToolSet); ok { + if _, seen := owners[owner]; !seen { + owners[owner] = fmt.Sprintf("%s:%d", prefix, ordinal) + ordinal++ + } + } + return true + }) + } + } + return owners +} + +func (r *LocalRuntime) thoughtAppender(streamCtx context.Context, sess *session.Session, caller *agent.Agent) func(context.Context, any, string) ([]string, error) { + owners := thinkOwners(caller, sess.ExtraToolSets) + store, supported := r.sessionStore.(session.ThinkScratchpadStore) + var root string + var rootErr error + if len(owners) != 0 && supported { + root, rootErr = r.thoughtRoot(streamCtx, sess) + } + return func(ctx context.Context, owner any, thought string) ([]string, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + ts, ok := owner.(*think.ToolSet) + if !ok { + return nil, errors.New("unknown think scratchpad owner") + } + key, ok := owners[ts] + if !ok { + return nil, errors.New("unknown think scratchpad owner") + } + if !supported { + return nil, errors.New("session store does not support think scratchpads") + } + if rootErr != nil { + return nil, rootErr + } + return store.AppendThought(ctx, root, session.ScratchpadKey{SessionID: sess.ID, AgentName: caller.Name(), ToolsetKey: key}, thought) + } +} + +func (r *LocalRuntime) thoughtRoot(ctx context.Context, sess *session.Session) (string, error) { + root, _ := ctx.Value(scratchpadRootKey{}).(string) + // Resumed children and their descendants may inherit an intermediate ancestor. + // Resolve stored ancestry once per batch, never from a mutable live registry. + if sess.ParentID == "" { + root = sess.ID + } else { + if root == "" { + root = sess.ParentID + } + seen := make(map[string]bool) + for { + if seen[root] { + return "", session.ErrScratchpadOwner + } + seen[root] = true + parent, err := r.sessionStore.GetSession(ctx, root) + if err != nil { + return "", fmt.Errorf("resolve think root: %w", err) + } + if parent.ParentID == "" { + break + } + root = parent.ParentID + } + } + // Delegated children need not have the protocol origin of their root. + if root == sess.ID { + if _, err := r.sessionStore.GetSessionByOrigin(ctx, root, sess.Origin); err != nil { + return "", fmt.Errorf("think root unavailable: %w", err) + } + } + return root, nil +} diff --git a/pkg/runtime/toolexec/dispatcher.go b/pkg/runtime/toolexec/dispatcher.go index 6856ccecf6..c545b75808 100644 --- a/pkg/runtime/toolexec/dispatcher.go +++ b/pkg/runtime/toolexec/dispatcher.go @@ -228,6 +228,9 @@ type Dispatcher struct { // [tools.ErrRecallNotSupported]. Recall func(ctx context.Context, sess *session.Session, a *agent.Agent, message string) error + // AppendThought binds internal scratchpad mutation to this batch's caller. + AppendThought func(ctx context.Context, owner any, thought string) ([]string, error) + confirmationMu *sync.Mutex fallbackConfirmationMu sync.Mutex } @@ -1064,6 +1067,13 @@ func (r callRuntime) RunTool(ctx context.Context, tc tools.ToolCall, tool tools. return result, err } +func (r callRuntime) AppendThought(ctx context.Context, owner any, thought string) ([]string, error) { + if r.c.d.AppendThought == nil { + return nil, errors.New("think scratchpad unavailable in this runtime") + } + return r.c.d.AppendThought(ctx, owner, thought) +} + func (r callRuntime) Recall(ctx context.Context, message string) error { if r.c.d.Recall == nil { return tools.ErrRecallNotSupported diff --git a/pkg/runtime/toolexec/dispatcher_test.go b/pkg/runtime/toolexec/dispatcher_test.go index e173f45b80..f544848f1e 100644 --- a/pkg/runtime/toolexec/dispatcher_test.go +++ b/pkg/runtime/toolexec/dispatcher_test.go @@ -2132,3 +2132,41 @@ func TestDispatcherBoundsResultsAfterHooks(t *testing.T) { } } } + +func TestDispatcher_ThinkAppender(t *testing.T) { + t.Parallel() + for _, available := range []bool{true, false} { + t.Run(fmt.Sprintf("available=%t", available), func(t *testing.T) { + t.Parallel() + sess := session.New() + sess.ToolsApproved = true + a := newAgent() + owner := new(byte) + d := &toolexec.Dispatcher{AgentFor: func(*session.Session) *agent.Agent { return a }} + if available { + d.AppendThought = func(_ context.Context, actual any, thought string) ([]string, error) { + assert.Same(t, owner, actual) + return []string{"retained", thought}, nil + } + } + tool := tools.Tool{Name: "think", Handler: func(ctx context.Context, _ tools.ToolCall, rt tools.Runtime) (*tools.ToolCallResult, error) { + appender, ok := rt.(tools.ScratchpadAppender) + require.True(t, ok) + out, err := appender.AppendThought(ctx, owner, "current") + if err != nil { + return tools.ResultError(err.Error()), nil + } + return tools.ResultSuccess(strings.Join(out, "\n")), nil + }} + em := &captureEmitter{} + d.Process(t.Context(), sess, []tools.ToolCall{{ID: "call", Function: tools.FunctionCall{Name: "think"}}}, []tools.Tool{tool}, em) + require.Len(t, em.responses, 1) + assert.Equal(t, !available, em.responses[0].IsError) + if available { + assert.Equal(t, "retained\ncurrent", em.responses[0].Output) + } else { + assert.Contains(t, em.responses[0].Output, "unavailable") + } + }) + } +} diff --git a/pkg/session/branch.go b/pkg/session/branch.go index ea1573d157..77bb3a0d41 100644 --- a/pkg/session/branch.go +++ b/pkg/session/branch.go @@ -79,6 +79,7 @@ func (s *Session) Clone() *Session { clone := &Session{ ID: s.ID, + Origin: s.Origin, InputID: s.InputID, Title: s.Title, Evals: cloneEvalCriteria(s.Evals), diff --git a/pkg/session/clone_test.go b/pkg/session/clone_test.go index e375e73831..ca5a760167 100644 --- a/pkg/session/clone_test.go +++ b/pkg/session/clone_test.go @@ -20,6 +20,7 @@ func TestClone_CopiesScalarFields(t *testing.T) { t.Parallel() orig := &Session{ ID: "sess-1", + Origin: "run", Title: "title", ToolsApproved: true, NonInteractive: true, @@ -51,6 +52,7 @@ func TestClone_CopiesScalarFields(t *testing.T) { // Unlike BranchSession, Clone keeps the original identity and history. assert.Equal(t, "sess-1", clone.ID) + assert.Equal(t, "run", clone.Origin) assert.Equal(t, "title", clone.Title) assert.True(t, clone.ToolsApproved) assert.True(t, clone.NonInteractive) diff --git a/pkg/session/migrations.go b/pkg/session/migrations.go index 632fc1b1ec..18cba97b8a 100644 --- a/pkg/session/migrations.go +++ b/pkg/session/migrations.go @@ -482,6 +482,24 @@ func getAllMigrations() []Migration { Description: "Add root_kind to generated_media_manifest so user-confirmed out-of-workspace generated files stay manifest-gated alongside workspace-relative ones", UpSQL: `ALTER TABLE generated_media_manifest ADD COLUMN root_kind TEXT NOT NULL DEFAULT 'workspace'`, }, + { + ID: 31, + Name: "031_add_think_scratchpad", + Description: "Retain conversation-scoped think scratchpads under their root session", + UpSQL: ` + CREATE TABLE think_scratchpad ( + root_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE, + session_id TEXT NOT NULL, + agent_name TEXT NOT NULL, + toolset_key TEXT NOT NULL, + seq INTEGER NOT NULL, + thought TEXT NOT NULL, + PRIMARY KEY (session_id, agent_name, toolset_key, seq) + ); + CREATE INDEX idx_think_scratchpad_root ON think_scratchpad(root_id); + `, + DownSQL: `DROP TABLE think_scratchpad`, + }, } } diff --git a/pkg/session/migrations_pinned_test.go b/pkg/session/migrations_pinned_test.go index 5e0cbaa738..b7359b30c6 100644 --- a/pkg/session/migrations_pinned_test.go +++ b/pkg/session/migrations_pinned_test.go @@ -39,7 +39,7 @@ func TestMigrationCatalogIsContentPinned(t *testing.T) { got := digestMigrationCatalog(getAllMigrations()) - const wantDigest = "e2d764e159e9c952596cdac552c860df2df3c90d5cf00fdac0ff4d9eda5b15ba" + const wantDigest = "403f21f32ebfb6de873517a8849ce5518125ab273941c3045ce938bf030a05b2" if got != wantDigest { t.Fatalf(`migration catalogue content has changed. diff --git a/pkg/session/store.go b/pkg/session/store.go index 65132ed123..c0e7bbb630 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -11,6 +11,7 @@ import ( "slices" "strconv" "strings" + "sync" "sync/atomic" "time" @@ -157,7 +158,9 @@ type InMemorySessionStore struct { sessions *concurrent.Map[string, *Session] generatedFiles *concurrent.Map[string, GeneratedFile] // keyed by generatedFileKey generatedBlobs *concurrent.Map[string, []byte] // keyed by generatedFileKey - messageID atomic.Int64 // counter for message IDs, incremented via Add(1) + scratchpadMu sync.Mutex + scratchpads map[string]map[ScratchpadKey][]string + messageID atomic.Int64 // counter for message IDs, incremented via Add(1) } func NewInMemorySessionStore() Store { @@ -269,6 +272,8 @@ func (s *InMemorySessionStore) DeleteSession(ctx context.Context, id string) err if err := ctx.Err(); err != nil { return err } + s.scratchpadMu.Lock() + defer s.scratchpadMu.Unlock() pending := []string{id} seen := make(map[string]struct{}) for len(pending) > 0 { @@ -283,6 +288,17 @@ func (s *InMemorySessionStore) DeleteSession(ctx context.Context, id string) err s.deleteGeneratedFiles(id) s.deleteGeneratedBlobs(id) } + for rootID, entries := range s.scratchpads { + if _, deleted := seen[rootID]; deleted { + delete(s.scratchpads, rootID) + continue + } + for key := range entries { + if _, deleted := seen[key.SessionID]; deleted { + delete(entries, key) + } + } + } return nil } @@ -1121,6 +1137,9 @@ func (s *SQLiteSessionStore) DeleteSession(ctx context.Context, id string) error if _, err := tx.ExecContext(ctx, subtree+"DELETE FROM generated_media_manifest WHERE session_id IN (SELECT id FROM subtree)", id); err != nil { return err } + if _, err := tx.ExecContext(ctx, subtree+"DELETE FROM think_scratchpad WHERE session_id IN (SELECT id FROM subtree)", id); err != nil { + return err + } result, err := tx.ExecContext(ctx, "DELETE FROM sessions WHERE id = ?", id) if err != nil { return err diff --git a/pkg/session/think_scratchpad.go b/pkg/session/think_scratchpad.go new file mode 100644 index 0000000000..f04630d540 --- /dev/null +++ b/pkg/session/think_scratchpad.go @@ -0,0 +1,129 @@ +package session + +import ( + "context" + "errors" + "fmt" + "slices" +) + +// ScratchpadKey separates conversations, callers and configured think instances. +type ScratchpadKey struct { + SessionID string + AgentName string + ToolsetKey string +} + +// ThinkScratchpadStore atomically appends a thought and returns an independent +// ordered snapshot. rootID owns cleanup, not the scratchpad's sharing identity. +type ThinkScratchpadStore interface { + AppendThought(ctx context.Context, rootID string, key ScratchpadKey, thought string) ([]string, error) +} + +var ErrScratchpadOwner = errors.New("think scratchpad root missing or ownership inconsistent") + +func validateScratchpadKey(rootID string, key ScratchpadKey) error { + if rootID == "" || key.SessionID == "" { + return ErrEmptyID + } + if key.AgentName == "" || key.ToolsetKey == "" { + return ErrScratchpadOwner + } + return nil +} + +func (s *InMemorySessionStore) AppendThought(ctx context.Context, rootID string, key ScratchpadKey, thought string) ([]string, error) { + if err := validateScratchpadKey(rootID, key); err != nil { + return nil, err + } + s.scratchpadMu.Lock() + defer s.scratchpadMu.Unlock() + if err := ctx.Err(); err != nil { + return nil, err + } + root, exists := s.sessions.Load(rootID) + if !exists { + return nil, ErrScratchpadOwner + } + root.mu.RLock() + isRoot := root.ParentID == "" + root.mu.RUnlock() + if !isRoot { + return nil, ErrScratchpadOwner + } + for owner, entries := range s.scratchpads { + if owner != rootID { + for existing := range entries { + if existing.SessionID == key.SessionID { + return nil, ErrScratchpadOwner + } + } + } + } + if s.scratchpads == nil { + s.scratchpads = make(map[string]map[ScratchpadKey][]string) + } + if s.scratchpads[rootID] == nil { + s.scratchpads[rootID] = make(map[ScratchpadKey][]string) + } + thoughts := append(s.scratchpads[rootID][key], thought) + s.scratchpads[rootID][key] = thoughts + return slices.Clone(thoughts), nil +} + +func (s *SQLiteSessionStore) AppendThought(ctx context.Context, rootID string, key ScratchpadKey, thought string) ([]string, error) { + if err := validateScratchpadKey(rootID, key); err != nil { + return nil, err + } + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return nil, fmt.Errorf("begin think append: %w", err) + } + defer func() { _ = tx.Rollback() }() + // Acquire the write lock before reading: independently loaded sessions must + // not upgrade stale read transactions or allocate the same sequence. + result, err := tx.ExecContext(ctx, ` + INSERT INTO think_scratchpad (root_id, session_id, agent_name, toolset_key, seq, thought) + SELECT ?, ?, ?, ?, COALESCE(( + SELECT MAX(seq) FROM think_scratchpad WHERE session_id = ? AND agent_name = ? AND toolset_key = ? + ), 0) + 1, ? + WHERE EXISTS (SELECT 1 FROM sessions WHERE id = ? AND (parent_id IS NULL OR parent_id = '')) + AND NOT EXISTS ( + SELECT 1 FROM think_scratchpad WHERE session_id = ? AND root_id != ? + )`, rootID, key.SessionID, key.AgentName, key.ToolsetKey, + key.SessionID, key.AgentName, key.ToolsetKey, thought, rootID, + key.SessionID, rootID) + if err != nil { + return nil, fmt.Errorf("append think thought: %w", err) + } + n, err := result.RowsAffected() + if err != nil { + return nil, err + } + if n != 1 { + return nil, ErrScratchpadOwner + } + rows, err := tx.QueryContext(ctx, `SELECT thought FROM think_scratchpad + WHERE session_id = ? AND agent_name = ? AND toolset_key = ? ORDER BY seq`, + key.SessionID, key.AgentName, key.ToolsetKey) + if err != nil { + return nil, err + } + defer rows.Close() + var thoughts []string + for rows.Next() { + var thought string + if err := rows.Scan(&thought); err != nil { + return nil, err + } + thoughts = append(thoughts, thought) + } + err = errors.Join(rows.Err(), rows.Close()) + if err != nil { + return nil, err + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("commit think thought: %w", err) + } + return thoughts, nil +} diff --git a/pkg/session/think_scratchpad_test.go b/pkg/session/think_scratchpad_test.go new file mode 100644 index 0000000000..75617bed4e --- /dev/null +++ b/pkg/session/think_scratchpad_test.go @@ -0,0 +1,238 @@ +package session + +import ( + "context" + "fmt" + "path/filepath" + "strconv" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func scratchpadStores(t *testing.T, check func(*testing.T, Store)) { + t.Helper() + for _, kind := range []string{"memory", "sqlite"} { + t.Run(kind, func(t *testing.T) { + t.Parallel() + store := NewInMemorySessionStore() + if kind == "sqlite" { + var err error + store, err = newSQLiteStoreForTest(t, filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + } + t.Cleanup(func() { require.NoError(t, store.Close()) }) + check(t, store) + }) + } +} + +func TestThinkScratchpadStore(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + ctx := t.Context() + appendStore := store.(ThinkScratchpadStore) + root := New() + key := ScratchpadKey{SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"} + _, err := appendStore.AppendThought(ctx, root.ID, key, "before persistence") + require.ErrorIs(t, err, ErrScratchpadOwner) + require.NoError(t, store.UpdateSession(ctx, root)) + stale, err := store.GetSession(ctx, root.ID) + require.NoError(t, err) + stale = stale.Clone() + stale.Origin = root.Origin + out, err := appendStore.AppendThought(ctx, root.ID, key, "first\nline") + require.NoError(t, err) + require.Equal(t, []string{"first\nline"}, out) + out[0] = "caller mutation" + require.NoError(t, store.UpdateSession(ctx, stale)) + require.NoError(t, store.PersistCompaction(ctx, stale, 0, 0, Item{Summary: "compacted"})) + out, err = appendStore.AppendThought(ctx, root.ID, key, "") + require.NoError(t, err) + assert.Equal(t, []string{"first\nline", ""}, out) + for _, other := range []ScratchpadKey{ + {SessionID: root.ID, AgentName: "worker", ToolsetKey: "cfg:0"}, + {SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:1"}, + {SessionID: "unpersisted-child", AgentName: "root", ToolsetKey: "cfg:0"}, + } { + out, err := appendStore.AppendThought(ctx, root.ID, other, "isolated") + require.NoError(t, err) + assert.Equal(t, []string{"isolated"}, out) + } + otherRoot := New() + require.NoError(t, store.UpdateSession(ctx, otherRoot)) + _, err = appendStore.AppendThought(ctx, otherRoot.ID, key, "wrong owner") + require.ErrorIs(t, err, ErrScratchpadOwner) + wrongKey := key + wrongKey.AgentName = "different-agent" + _, err = appendStore.AppendThought(ctx, otherRoot.ID, wrongKey, "wrong root with fresh key") + require.ErrorIs(t, err, ErrScratchpadOwner) + cancelled, cancel := context.WithCancel(ctx) + cancel() + _, err = appendStore.AppendThought(cancelled, root.ID, key, "cancelled") + require.ErrorIs(t, err, context.Canceled) + out, err = appendStore.AppendThought(ctx, root.ID, key, "last") + require.NoError(t, err) + assert.Equal(t, []string{"first\nline", "", "last"}, out) + }) +} + +func TestThinkScratchpadConcurrent(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + root := New() + require.NoError(t, store.UpdateSession(t.Context(), root)) + appender := store.(ThinkScratchpadStore) + key := ScratchpadKey{SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"} + const n = 64 + outputs := make([][]string, n) + errs := make([]error, n) + thoughts := make([]string, n) + var wg sync.WaitGroup + for i := range n { + thoughts[i] = fmt.Sprintf("thought-%d", i) + wg.Go(func() { outputs[i], errs[i] = appender.AppendThought(t.Context(), root.ID, key, thoughts[i]) }) + } + wg.Wait() + lengths := make(map[int]bool) + for i := range n { + require.NoError(t, errs[i]) + require.Equal(t, thoughts[i], outputs[i][len(outputs[i])-1]) + lengths[len(outputs[i])] = true + } + require.Len(t, lengths, n) + out, err := appender.AppendThought(t.Context(), root.ID, key, "final") + require.NoError(t, err) + require.Len(t, out, n+1) + assert.ElementsMatch(t, thoughts, out[:n]) + }) +} + +func TestThinkScratchpadDeletion(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + root := New() + child := New() + child.ParentID = root.ID + grandchild := New() + grandchild.ParentID = child.ID + for _, sess := range []*Session{root, child, grandchild} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + appender := store.(ThinkScratchpadStore) + keys := []ScratchpadKey{ + {SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"}, + {SessionID: child.ID, AgentName: "worker", ToolsetKey: "cfg:0"}, + {SessionID: grandchild.ID, AgentName: "worker", ToolsetKey: "cfg:0"}, + {SessionID: "aborted", AgentName: "worker", ToolsetKey: "cfg:0"}, + } + for _, key := range keys { + _, err := appender.AppendThought(t.Context(), root.ID, key, "initial") + require.NoError(t, err) + } + require.NoError(t, store.DeleteSession(t.Context(), child.ID)) + for _, key := range keys[1:3] { + out, err := appender.AppendThought(t.Context(), root.ID, key, "after child deletion") + require.NoError(t, err) + assert.Equal(t, []string{"after child deletion"}, out) + } + out, err := appender.AppendThought(t.Context(), root.ID, keys[3], "still root owned") + require.NoError(t, err) + assert.Equal(t, []string{"initial", "still root owned"}, out) + require.NoError(t, store.DeleteSession(t.Context(), root.ID)) + _, err = appender.AppendThought(t.Context(), root.ID, keys[3], "no resurrection") + require.ErrorIs(t, err, ErrScratchpadOwner) + require.NoError(t, store.UpdateSession(t.Context(), root)) + for _, key := range keys { + out, err := appender.AppendThought(t.Context(), root.ID, key, "after root deletion") + require.NoError(t, err) + assert.Equal(t, []string{"after root deletion"}, out) + } + }) +} + +func TestThinkScratchpadSQLiteReopenAndMultipleHandles(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "sessions.db") + store, err := newSQLiteStoreForTest(t, path) + require.NoError(t, err) + root := New() + require.NoError(t, store.UpdateSession(t.Context(), root)) + key := ScratchpadKey{SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"} + _, err = store.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, key, "before restart") + require.NoError(t, err) + require.NoError(t, store.Close()) + store, err = newSQLiteStoreForTest(t, path) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, store.Close()) }) + second, err := newSQLiteStoreForTest(t, path) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, second.Close()) }) + var wg sync.WaitGroup + errs := make([]error, 20) + for i := range errs { + wg.Go(func() { + writer := store + if i%2 == 0 { + writer = second + } + _, errs[i] = writer.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, key, strconv.Itoa(i)) + }) + } + wg.Wait() + for _, err := range errs { + require.NoError(t, err) + } + out, err := store.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, key, "end") + require.NoError(t, err) + require.Len(t, out, 22) + assert.Equal(t, "before restart", out[0]) + assert.Equal(t, "end", out[21]) +} + +func TestThinkScratchpadMigrationUpgrade(t *testing.T) { + t.Parallel() + path := filepath.Join(t.TempDir(), "sessions.db") + store, err := newSQLiteStoreForTest(t, path) + require.NoError(t, err) + sqlite := store.(*SQLiteSessionStore) + root := New() + require.NoError(t, store.UpdateSession(t.Context(), root)) + _, err = sqlite.db.ExecContext(t.Context(), `DROP TABLE think_scratchpad; DELETE FROM migrations WHERE id = 31`) + require.NoError(t, err) + require.NoError(t, store.Close()) + store, err = newSQLiteStoreForTest(t, path) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, store.Close()) }) + out, err := store.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, ScratchpadKey{SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"}, "upgraded") + require.NoError(t, err) + assert.Equal(t, []string{"upgraded"}, out) + mgr := NewMigrationManager(store.(*SQLiteSessionStore).db) + require.NoError(t, mgr.InitializeMigrations(t.Context())) +} + +func TestThinkScratchpadDeleteRollback(t *testing.T) { + t.Parallel() + store, err := newSQLiteStoreForTest(t, filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, store.Close()) }) + root, child := New(), New() + child.ParentID = root.ID + for _, sess := range []*Session{root, child} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + key := ScratchpadKey{SessionID: child.ID, AgentName: "worker", ToolsetKey: "cfg:0"} + _, err = store.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, key, "retained") + require.NoError(t, err) + _, err = store.(*SQLiteSessionStore).db.ExecContext(t.Context(), `CREATE TRIGGER prevent_session_delete BEFORE DELETE ON sessions BEGIN SELECT RAISE(ABORT, 'delete blocked'); END;`) + require.NoError(t, err) + require.Error(t, store.DeleteSession(t.Context(), child.ID)) + out, err := store.(ThinkScratchpadStore).AppendThought(t.Context(), root.ID, key, "after failed delete") + require.NoError(t, err) + assert.Equal(t, []string{"retained", "after failed delete"}, out) +} diff --git a/pkg/tools/builtin/think/think.go b/pkg/tools/builtin/think/think.go index aaf370b65a..b3548eae5f 100644 --- a/pkg/tools/builtin/think/think.go +++ b/pkg/tools/builtin/think/think.go @@ -3,7 +3,6 @@ package think import ( "context" "strings" - "sync" "github.com/docker/docker-agent/pkg/tools" ) @@ -15,20 +14,25 @@ func CreateToolSet() (tools.ToolSet, error) { return New(), nil } -type ToolSet struct { - mu sync.Mutex - thoughts []string -} +// ToolSet has nonzero size so distinct instances have distinct pointer identities. +type ToolSet struct{ _ byte } type Args struct { Thought string `json:"thought" jsonschema:"The thought to think about"` } -func (t *ToolSet) callTool(_ context.Context, params Args) (*tools.ToolCallResult, error) { - t.mu.Lock() - defer t.mu.Unlock() - t.thoughts = append(t.thoughts, params.Thought) - return tools.ResultSuccess("Thoughts:\n" + strings.Join(t.thoughts, "\n")), nil +func (t *ToolSet) callTool(ctx context.Context, params Args, rt tools.Runtime) (*tools.ToolCallResult, error) { + thoughts := []string{params.Thought} + if appender, ok := rt.(tools.ScratchpadAppender); ok { + var err error + thoughts, err = appender.AppendThought(ctx, t, params.Thought) + if err != nil { + return tools.ResultError(err.Error()), nil + } + } else if _, unscoped := rt.(tools.NopRuntime); !unscoped { + return tools.ResultError("think scratchpad unavailable in this runtime"), nil + } + return tools.ResultSuccess("Thoughts:\n" + strings.Join(thoughts, "\n")), nil } func New() *ToolSet { @@ -50,10 +54,10 @@ func (t *ToolSet) Tools(context.Context) ([]tools.Tool, error) { { Name: ToolNameThink, Category: "think", - Description: "Use the tool to think about something. It will not obtain new information or change the database, but just append the thought to the log. Use it when complex reasoning or some cache memory is needed.", + Description: "Use the tool to think about something. It will not obtain new information or change external data; it appends to an internal conversation scratchpad. Use it when complex reasoning or some cache memory is needed.", Parameters: tools.MustSchemaFor[Args](), OutputSchema: tools.MustSchemaFor[string](), - Handler: tools.NewHandler(t.callTool), + Handler: tools.NewRuntimeHandler(t.callTool), Annotations: tools.ToolAnnotations{ ReadOnlyHint: true, Title: "Think", diff --git a/pkg/tools/builtin/think/think_test.go b/pkg/tools/builtin/think/think_test.go index 64baa3ca7f..047e6197c5 100644 --- a/pkg/tools/builtin/think/think_test.go +++ b/pkg/tools/builtin/think/think_test.go @@ -3,6 +3,7 @@ package think import ( "context" "fmt" + "slices" "strings" "sync" "testing" @@ -13,19 +14,33 @@ import ( "github.com/docker/docker-agent/pkg/tools" ) +type scratchpadRuntime struct { + tools.NopRuntime + + mu sync.Mutex + thoughts []string +} + +func (r *scratchpadRuntime) AppendThought(ctx context.Context, _ any, thought string) ([]string, error) { + r.mu.Lock() + defer r.mu.Unlock() + if err := ctx.Err(); err != nil { + return nil, err + } + r.thoughts = append(r.thoughts, thought) + return slices.Clone(r.thoughts), nil +} + func TestThinkTool_Handler(t *testing.T) { t.Parallel() - tool := New() - - result, err := tool.callTool(t.Context(), Args{Thought: "This is a test thought"}) + handler := thinkHandler(t) + rt := &scratchpadRuntime{} + output, err := callThink(t.Context(), handler, "This is a test thought", rt) require.NoError(t, err) - assert.Contains(t, result.Output, "This is a test thought") - - result, err = tool.callTool(t.Context(), Args{Thought: "Another thought"}) + assert.Equal(t, "Thoughts:\nThis is a test thought", output) + output, err = callThink(t.Context(), handler, "Another thought", rt) require.NoError(t, err) - - assert.Contains(t, result.Output, "This is a test thought") - assert.Contains(t, result.Output, "Another thought") + assert.Equal(t, "Thoughts:\nThis is a test thought\nAnother thought", output) } func thinkHandler(t *testing.T) tools.ToolHandler { @@ -36,25 +51,29 @@ func thinkHandler(t *testing.T) tools.ToolHandler { return allTools[0].Handler } -func callThink(ctx context.Context, handler tools.ToolHandler, thought string) (string, error) { +func callThink(ctx context.Context, handler tools.ToolHandler, thought string, rt tools.Runtime) (string, error) { result, err := handler(ctx, tools.ToolCall{ Function: tools.FunctionCall{Name: ToolNameThink, Arguments: fmt.Sprintf(`{"thought":%q}`, thought)}, - }, tools.NopRuntime{}) + }, rt) if err != nil { return "", err } + if result.IsError { + return "", fmt.Errorf("%s", result.Output) + } return result.Output, nil } func TestThinkTool_SequentialOutput(t *testing.T) { t.Parallel() handler := thinkHandler(t) + rt := &scratchpadRuntime{} - output, err := callThink(t.Context(), handler, "first") + output, err := callThink(t.Context(), handler, "first", rt) require.NoError(t, err) assert.Equal(t, "Thoughts:\nfirst", output) - output, err = callThink(t.Context(), handler, "second") + output, err = callThink(t.Context(), handler, "second", rt) require.NoError(t, err) assert.Equal(t, "Thoughts:\nfirst\nsecond", output) } @@ -62,6 +81,7 @@ func TestThinkTool_SequentialOutput(t *testing.T) { func TestThinkTool_ConcurrentCalls(t *testing.T) { t.Parallel() handler := thinkHandler(t) + rt := &scratchpadRuntime{} const n = 64 thoughts := make([]string, n) @@ -71,7 +91,7 @@ func TestThinkTool_ConcurrentCalls(t *testing.T) { for i := range n { thoughts[i] = fmt.Sprintf("thought-%02d", i) wg.Go(func() { - outputs[i], errs[i] = callThink(t.Context(), handler, thoughts[i]) + outputs[i], errs[i] = callThink(t.Context(), handler, thoughts[i], rt) }) } wg.Wait() @@ -83,7 +103,7 @@ func TestThinkTool_ConcurrentCalls(t *testing.T) { assert.Contains(t, lines[1:], thoughts[i]) } - output, err := callThink(t.Context(), handler, "final") + output, err := callThink(t.Context(), handler, "final", rt) require.NoError(t, err) lines := strings.Split(output, "\n") require.Len(t, lines, n+2) @@ -120,3 +140,28 @@ func TestThinkTool_ParametersAreObjects(t *testing.T) { assert.Equal(t, "object", m["type"]) } } + +func TestThinkTool_UnscopedIsEphemeral(t *testing.T) { + t.Parallel() + handler := thinkHandler(t) + for _, thought := range []string{"first\nline", "second", ""} { + output, err := callThink(t.Context(), handler, thought, tools.NopRuntime{}) + require.NoError(t, err) + assert.Equal(t, "Thoughts:\n"+thought, output) + } +} + +type scopedWithoutScratchpad struct{ tools.NopRuntime } + +func TestThinkTool_Failures(t *testing.T) { + t.Parallel() + handler := thinkHandler(t) + _, err := callThink(t.Context(), handler, "hidden", scopedWithoutScratchpad{}) + require.ErrorContains(t, err, "unavailable") + ctx, cancel := context.WithCancel(t.Context()) + cancel() + rt := &scratchpadRuntime{} + _, err = callThink(ctx, handler, "cancelled", rt) + require.ErrorContains(t, err, "canceled") + assert.Empty(t, rt.thoughts) +} diff --git a/pkg/tools/runtime.go b/pkg/tools/runtime.go index 3d3ef13924..c5ada7f586 100644 --- a/pkg/tools/runtime.go +++ b/pkg/tools/runtime.go @@ -76,6 +76,12 @@ type ToolRunner interface { RunTool(ctx context.Context, toolCall ToolCall, tool Tool) (*ToolCallResult, error) } +// ScratchpadAppender is an optional host-owned, conversation-scoped think store. +// owner identifies the calling toolset instance, not a caller-supplied namespace. +type ScratchpadAppender interface { + AppendThought(ctx context.Context, owner any, thought string) ([]string, error) +} + // ErrRecallNotSupported is returned by [Runtime.Recall] implementations that // cannot steer an agent loop. var ErrRecallNotSupported = errors.New("recall is not supported by this host") From d63148c9c5a4a8dde773091e28ef95a16c7c3b99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20H=C3=A9ritier?= Date: Wed, 14 Oct 2026 09:04:40 +0000 Subject: [PATCH 2/2] fix: clean up think scratchpad roots and index ownership MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Arnaud Héritier --- pkg/session/store.go | 22 ++-- pkg/session/think_scratchpad.go | 14 +-- pkg/session/think_scratchpad_test.go | 158 +++++++++++++++++++++++++++ 3 files changed, 179 insertions(+), 15 deletions(-) diff --git a/pkg/session/store.go b/pkg/session/store.go index c0e7bbb630..b10d35a83d 100644 --- a/pkg/session/store.go +++ b/pkg/session/store.go @@ -155,12 +155,13 @@ type Store interface { } type InMemorySessionStore struct { - sessions *concurrent.Map[string, *Session] - generatedFiles *concurrent.Map[string, GeneratedFile] // keyed by generatedFileKey - generatedBlobs *concurrent.Map[string, []byte] // keyed by generatedFileKey - scratchpadMu sync.Mutex - scratchpads map[string]map[ScratchpadKey][]string - messageID atomic.Int64 // counter for message IDs, incremented via Add(1) + sessions *concurrent.Map[string, *Session] + generatedFiles *concurrent.Map[string, GeneratedFile] // keyed by generatedFileKey + generatedBlobs *concurrent.Map[string, []byte] // keyed by generatedFileKey + scratchpadMu sync.Mutex + scratchpads map[string]map[ScratchpadKey][]string + scratchpadOwners map[string]string + messageID atomic.Int64 // counter for message IDs, incremented via Add(1) } func NewInMemorySessionStore() Store { @@ -299,6 +300,13 @@ func (s *InMemorySessionStore) DeleteSession(ctx context.Context, id string) err } } } + for sessionID, rootID := range s.scratchpadOwners { + _, sessionDeleted := seen[sessionID] + _, rootDeleted := seen[rootID] + if sessionDeleted || rootDeleted { + delete(s.scratchpadOwners, sessionID) + } + } return nil } @@ -1137,7 +1145,7 @@ func (s *SQLiteSessionStore) DeleteSession(ctx context.Context, id string) error if _, err := tx.ExecContext(ctx, subtree+"DELETE FROM generated_media_manifest WHERE session_id IN (SELECT id FROM subtree)", id); err != nil { return err } - if _, err := tx.ExecContext(ctx, subtree+"DELETE FROM think_scratchpad WHERE session_id IN (SELECT id FROM subtree)", id); err != nil { + if _, err := tx.ExecContext(ctx, subtree+"DELETE FROM think_scratchpad WHERE session_id IN (SELECT id FROM subtree) OR root_id IN (SELECT id FROM subtree)", id); err != nil { return err } result, err := tx.ExecContext(ctx, "DELETE FROM sessions WHERE id = ?", id) diff --git a/pkg/session/think_scratchpad.go b/pkg/session/think_scratchpad.go index f04630d540..2582219d93 100644 --- a/pkg/session/think_scratchpad.go +++ b/pkg/session/think_scratchpad.go @@ -51,14 +51,11 @@ func (s *InMemorySessionStore) AppendThought(ctx context.Context, rootID string, if !isRoot { return nil, ErrScratchpadOwner } - for owner, entries := range s.scratchpads { - if owner != rootID { - for existing := range entries { - if existing.SessionID == key.SessionID { - return nil, ErrScratchpadOwner - } - } - } + if owner, exists := s.scratchpadOwners[key.SessionID]; exists && owner != rootID { + return nil, ErrScratchpadOwner + } + if s.scratchpadOwners == nil { + s.scratchpadOwners = make(map[string]string) } if s.scratchpads == nil { s.scratchpads = make(map[string]map[ScratchpadKey][]string) @@ -68,6 +65,7 @@ func (s *InMemorySessionStore) AppendThought(ctx context.Context, rootID string, } thoughts := append(s.scratchpads[rootID][key], thought) s.scratchpads[rootID][key] = thoughts + s.scratchpadOwners[key.SessionID] = rootID return slices.Clone(thoughts), nil } diff --git a/pkg/session/think_scratchpad_test.go b/pkg/session/think_scratchpad_test.go index 75617bed4e..6b5620d5bf 100644 --- a/pkg/session/think_scratchpad_test.go +++ b/pkg/session/think_scratchpad_test.go @@ -2,6 +2,7 @@ package session import ( "context" + "database/sql" "fmt" "path/filepath" "strconv" @@ -236,3 +237,160 @@ func TestThinkScratchpadDeleteRollback(t *testing.T) { require.NoError(t, err) assert.Equal(t, []string{"retained", "after failed delete"}, out) } + +func TestThinkScratchpadSQLiteDeletionWithoutForeignKeys(t *testing.T) { + t.Parallel() + db, err := sql.Open("sqlite", filepath.Join(t.TempDir(), "sessions.db")) + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, db.Close()) }) + store, err := NewSQLiteSessionStoreFromDB(t.Context(), db) + require.NoError(t, err) + var foreignKeys int + require.NoError(t, db.QueryRowContext(t.Context(), "PRAGMA foreign_keys").Scan(&foreignKeys)) + require.Zero(t, foreignKeys) + + roots := []*Session{New(), New()} + keys := make([][]ScratchpadKey, len(roots)) + for i, root := range roots { + require.NoError(t, store.UpdateSession(t.Context(), root)) + keys[i] = []ScratchpadKey{ + {SessionID: root.ID, AgentName: "root", ToolsetKey: "cfg:0"}, + {SessionID: root.ID + "-unpersisted", AgentName: "worker", ToolsetKey: "cfg:0"}, + } + for _, key := range keys[i] { + _, err := store.AppendThought(t.Context(), root.ID, key, "initial") + require.NoError(t, err) + } + } + require.NoError(t, store.DeleteSession(t.Context(), roots[0].ID)) + var count int + require.NoError(t, db.QueryRowContext(t.Context(), "SELECT COUNT(*) FROM think_scratchpad WHERE root_id = ?", roots[0].ID).Scan(&count)) + require.Zero(t, count) + require.NoError(t, db.QueryRowContext(t.Context(), "SELECT COUNT(*) FROM think_scratchpad").Scan(&count)) + require.Equal(t, 2, count) + for _, key := range keys[1] { + out, err := store.AppendThought(t.Context(), roots[1].ID, key, "unrelated survives") + require.NoError(t, err) + require.Equal(t, []string{"initial", "unrelated survives"}, out) + } + require.NoError(t, store.DeleteSession(t.Context(), roots[1].ID)) + require.NoError(t, db.QueryRowContext(t.Context(), "SELECT COUNT(*) FROM think_scratchpad").Scan(&count)) + require.Zero(t, count) + require.NoError(t, store.UpdateSession(t.Context(), roots[0])) + for _, key := range keys[0] { + out, err := store.AppendThought(t.Context(), roots[0].ID, key, "recreated") + require.NoError(t, err) + require.Equal(t, []string{"recreated"}, out) + } +} + +func TestThinkScratchpadOwnershipAcrossKeys(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + roots := []*Session{New(), New()} + for _, root := range roots { + require.NoError(t, store.UpdateSession(t.Context(), root)) + } + appender := store.(ThinkScratchpadStore) + for _, agent := range []string{"root", "worker"} { + for _, toolset := range []string{"cfg:0", "cfg:1"} { + key := ScratchpadKey{SessionID: "unpersisted", AgentName: agent, ToolsetKey: toolset} + out, err := appender.AppendThought(t.Context(), roots[0].ID, key, "first") + require.NoError(t, err) + require.Equal(t, []string{"first"}, out) + _, err = appender.AppendThought(t.Context(), roots[1].ID, key, "wrong owner") + require.ErrorIs(t, err, ErrScratchpadOwner) + } + } + }) +} + +func TestThinkScratchpadOwnershipDeletion(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + root, otherRoot, child := New(), New(), New() + for _, sess := range []*Session{root, otherRoot} { + require.NoError(t, store.UpdateSession(t.Context(), sess)) + } + require.NoError(t, store.AddSubSession(t.Context(), root.ID, child)) + appender := store.(ThinkScratchpadStore) + keys := []ScratchpadKey{ + {SessionID: child.ID, AgentName: "worker", ToolsetKey: "cfg:0"}, + {SessionID: "unpersisted", AgentName: "worker", ToolsetKey: "cfg:0"}, + } + unrelated := ScratchpadKey{SessionID: otherRoot.ID, AgentName: "root", ToolsetKey: "cfg:0"} + for _, key := range append([]ScratchpadKey{unrelated}, keys...) { + owner := root.ID + if key == unrelated { + owner = otherRoot.ID + } + _, err := appender.AppendThought(t.Context(), owner, key, "initial") + require.NoError(t, err) + } + if memory, ok := store.(*InMemorySessionStore); ok { + require.Equal(t, map[string]string{child.ID: root.ID, "unpersisted": root.ID, otherRoot.ID: otherRoot.ID}, memory.scratchpadOwners) + } + require.NoError(t, store.DeleteSession(t.Context(), child.ID)) + if memory, ok := store.(*InMemorySessionStore); ok { + require.Equal(t, map[string]string{"unpersisted": root.ID, otherRoot.ID: otherRoot.ID}, memory.scratchpadOwners) + } + out, err := appender.AppendThought(t.Context(), otherRoot.ID, keys[0], "new owner after child deletion") + require.NoError(t, err) + require.Equal(t, []string{"new owner after child deletion"}, out) + _, err = appender.AppendThought(t.Context(), otherRoot.ID, keys[1], "still root owned") + require.ErrorIs(t, err, ErrScratchpadOwner) + require.NoError(t, store.DeleteSession(t.Context(), root.ID)) + if memory, ok := store.(*InMemorySessionStore); ok { + require.Equal(t, map[string]string{otherRoot.ID: otherRoot.ID}, memory.scratchpadOwners) + } + out, err = appender.AppendThought(t.Context(), otherRoot.ID, keys[1], "new owner after root deletion") + require.NoError(t, err) + require.Equal(t, []string{"new owner after root deletion"}, out) + out, err = appender.AppendThought(t.Context(), otherRoot.ID, unrelated, "unrelated survives") + require.NoError(t, err) + require.Equal(t, []string{"initial", "unrelated survives"}, out) + }) +} + +func TestThinkScratchpadRejectedAppendDoesNotClaimOwnership(t *testing.T) { + t.Parallel() + scratchpadStores(t, func(t *testing.T, store Store) { + t.Helper() + root, child := New(), New() + require.NoError(t, store.UpdateSession(t.Context(), root)) + require.NoError(t, store.AddSubSession(t.Context(), root.ID, child)) + cancelled, cancel := context.WithCancel(t.Context()) + cancel() + appender := store.(ThinkScratchpadStore) + for _, tc := range []struct { + name string + cancelled bool + owner string + err error + }{ + {"missing", false, "missing-root", ErrScratchpadOwner}, + {"non-root", false, child.ID, ErrScratchpadOwner}, + {"cancelled", true, root.ID, context.Canceled}, + } { + t.Run(tc.name, func(t *testing.T) { + key := ScratchpadKey{SessionID: tc.name, AgentName: "worker", ToolsetKey: "cfg:0"} + ctx := t.Context() + if tc.cancelled { + ctx = cancelled + } + _, err := appender.AppendThought(ctx, tc.owner, key, "rejected") + require.ErrorIs(t, err, tc.err) + if memory, ok := store.(*InMemorySessionStore); ok { + require.NotContains(t, memory.scratchpadOwners, key.SessionID) + } + otherRoot := New() + require.NoError(t, store.UpdateSession(t.Context(), otherRoot)) + out, err := appender.AppendThought(t.Context(), otherRoot.ID, key, "accepted") + require.NoError(t, err) + require.Equal(t, []string{"accepted"}, out) + }) + } + }) +}