From 11304190191ba735b29147a3e1919c0ee0f47754 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 13:25:05 +0200 Subject: [PATCH 1/8] fix(session): lock-safe Messages access in AddMessage/branch/fork --- pkg/session/branch.go | 17 +++-- pkg/session/session.go | 34 +++++++++- pkg/session/session_race_test.go | 107 +++++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+), 6 deletions(-) diff --git a/pkg/session/branch.go b/pkg/session/branch.go index b4d9fcb0d2..8a47a503ab 100644 --- a/pkg/session/branch.go +++ b/pkg/session/branch.go @@ -36,7 +36,12 @@ func branchSessionWithTitle(parent *Session, branchAtPosition int, titleFn func( if parent == nil { return nil, errors.New("parent session is nil") } - if branchAtPosition < 0 || branchAtPosition > len(parent.Messages) { + + // Snapshot under parent.mu (like Clone) so a concurrent AddMessage/ + // ApplyCompaction — e.g. from a live HTTP stream on the same session — + // cannot race with this read or shift branchAtPosition mid-copy. + items := parent.snapshotItems() + if branchAtPosition < 0 || branchAtPosition > len(items) { return nil, fmt.Errorf("branch position %d out of range", branchAtPosition) } @@ -45,7 +50,7 @@ func branchSessionWithTitle(parent *Session, branchAtPosition int, titleFn func( branched.Messages = make([]Item, 0, branchAtPosition) for i := range branchAtPosition { - cloned, err := cloneSessionItem(parent.Messages[i]) + cloned, err := cloneSessionItem(items[i]) if err != nil { return nil, err } @@ -159,8 +164,12 @@ func cloneSubSession(src *Session) (*Session, error) { copySessionMetadata(cloned, src, src.Title) cloned.CreatedAt = src.CreatedAt - cloned.Messages = make([]Item, 0, len(src.Messages)) - for _, item := range src.Messages { + // Snapshot under src.mu: a sub-session created by a background agent + // task can still be actively appended to while the top-level session is + // being branched/forked. + items := src.snapshotItems() + cloned.Messages = make([]Item, 0, len(items)) + for _, item := range items { clonedItem, err := cloneSessionItem(item) if err != nil { return nil, err diff --git a/pkg/session/session.go b/pkg/session/session.go index 6351ac7015..bd5242ee2e 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -579,14 +579,21 @@ func cloneSchemaValue(v any) any { // Session helper methods -// AddMessage adds a message to the session -func (s *Session) AddMessage(msg *Message) { +// AddMessage adds a message to the session and returns the index the new +// item occupies in s.Messages. Callers that need to stamp an event with the +// message's position (e.g. UserMessageEvent.SessionPosition) must use this +// return value rather than a separate len(sess.Messages)-1 read: the latter +// races with concurrent AddMessage/ApplyCompaction calls (e.g. from a live +// HTTP AddMessage while a stream is running) and can also observe a later, +// larger length than the one that matched this append. +func (s *Session) AddMessage(msg *Message) int { s.mu.Lock() defer s.mu.Unlock() if msg != nil { capToolResultContent(&msg.Message, s.MaxToolResultTokens) } s.Messages = append(s.Messages, NewMessageItem(msg)) + return len(s.Messages) - 1 } // SetUsage records cumulative input/output token counts under s.mu. @@ -1025,6 +1032,29 @@ func (s *Session) MessageCount() int { return n } +// ItemCount returns the total number of items in s.Messages — messages, +// sub-sessions, summaries, and recorded errors alike. Unlike MessageCount, +// it counts every item, matching what len(s.Messages) would return outside +// the lock. Hot paths that need "the index the next appended item will +// occupy" (e.g. before calling AddSubSession) should use this instead of +// reading len(sess.Messages) directly, which races with concurrent +// AddMessage/ApplyCompaction. +func (s *Session) ItemCount() int { + s.mu.RLock() + defer s.mu.RUnlock() + return len(s.Messages) +} + +// MessagesSnapshot returns a lock-safe copy of the session's items, deep- +// copying each Message so the result cannot alias a concurrent AddMessage / +// UpdateMessage mutation. It is the exported counterpart of snapshotItems for +// callers outside this package (e.g. pkg/server's ForkSession) that need to +// iterate Messages without racing session.mu; in-package callers should keep +// using snapshotItems directly. +func (s *Session) MessagesSnapshot() []Item { + return s.snapshotItems() +} + // TotalCost computes the total cost of a session by walking all messages, // sub-sessions, and summary items. It does not use the session-level Cost // field, which exists only for backward-compatible persistence. diff --git a/pkg/session/session_race_test.go b/pkg/session/session_race_test.go index 2a504bb621..167ca54170 100644 --- a/pkg/session/session_race_test.go +++ b/pkg/session/session_race_test.go @@ -49,3 +49,110 @@ func TestCompactionInputConcurrent(t *testing.T) { }) wg.Wait() } + +// TestAddMessageReturnsAppendedIndex pins AddMessage's return-value +// contract: the index of the item it just appended. Hot paths (e.g. +// pkg/runtime/loop.go's UserMessageEvent emission) rely on this instead of +// a separate len(sess.Messages)-1 read, which would race with a concurrent +// AddMessage/ApplyCompaction and could observe a later, larger length. +func TestAddMessageReturnsAppendedIndex(t *testing.T) { + t.Parallel() + + s := New() + if got := s.AddMessage(UserMessage("a")); got != 0 { + t.Errorf("expected index 0, got %d", got) + } + if got := s.AddMessage(UserMessage("b")); got != 1 { + t.Errorf("expected index 1, got %d", got) + } +} + +// TestAddMessageConcurrentReturnsUniqueIndices pins the atomicity of +// AddMessage's returned index under concurrent callers: every call must +// observe the position of its own append (under s.mu), never a stale or +// duplicate one, which is what makes the returned value safe to stamp an +// event's SessionPosition with instead of reading len(sess.Messages)-1 +// after the fact. +func TestAddMessageConcurrentReturnsUniqueIndices(t *testing.T) { + t.Parallel() + + s := New() + const n = 200 + indices := make(chan int, n) + var wg sync.WaitGroup + for range n { + wg.Go(func() { + indices <- s.AddMessage(&Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u"}}) + }) + } + wg.Wait() + close(indices) + + seen := make(map[int]bool, n) + for idx := range indices { + if seen[idx] { + t.Fatalf("duplicate index %d returned by concurrent AddMessage calls", idx) + } + seen[idx] = true + } + if len(seen) != n { + t.Errorf("expected %d unique indices, got %d", n, len(seen)) + } +} + +// TestBranchSessionConcurrent pins the data-race fix for branch/fork: +// branchSessionWithTitle (BranchSession/ForkSession) must read +// parent.Messages under parent.mu (via snapshotItems), mirroring Clone, +// so it stays safe against a concurrent AddMessage on the same live +// session — e.g. the HTTP AddMessage path racing a TUI branch/fork action. +// Run with -race; without the lock the branchAtPosition bounds check and +// the parent.Messages[i] reads alias the live backing array and the race +// detector flags the AddMessage append. +func TestBranchSessionConcurrent(t *testing.T) { + t.Parallel() + + parent := New(WithUserMessage("seed")) + var wg sync.WaitGroup + for range 50 { + wg.Go(func() { + parent.AddMessage(&Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u"}}) + }) + wg.Go(func() { + if _, err := BranchSession(parent, 1); err != nil { + t.Errorf("BranchSession: %v", err) + } + }) + wg.Go(func() { + if _, err := ForkSession(parent, 1); err != nil { + t.Errorf("ForkSession: %v", err) + } + }) + } + wg.Wait() +} + +// TestForkSessionConcurrentWithSubSessionMutation pins the same fix for +// cloneSubSession: forking a session that contains a sub-session (e.g. a +// background agent task) must snapshot the sub-session's own Messages +// under its own mu, since that sub-session can still be actively appended +// to while the top-level session is being forked. +func TestForkSessionConcurrentWithSubSessionMutation(t *testing.T) { + t.Parallel() + + sub := New(WithUserMessage("sub-seed")) + parent := New(WithUserMessage("seed")) + parent.AddSubSession(sub) + + var wg sync.WaitGroup + for range 50 { + wg.Go(func() { + sub.AddMessage(&Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u"}}) + }) + wg.Go(func() { + if _, err := ForkSession(parent, 2); err != nil { + t.Errorf("ForkSession: %v", err) + } + }) + } + wg.Wait() +} From ada937e98bf16e2090a41232bd439e2ed35fe9f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 13:25:12 +0200 Subject: [PATCH 2/8] fix(runtime): route hot-path Messages reads through locked accessors pkg/runtime/loop.go, runtime.go, and compactor.go read sess.Messages (or len(sess.Messages)) directly in several hot paths, racing a concurrent HTTP AddMessage or compaction on the same live session and risking a wrong UserMessageEvent.SessionPosition or an index-out-of-range panic: - appendSteerAndEmit and the follow-up injection path in loop.go now use AddMessage's returned index instead of a separate len(sess.Messages)-1 read. - The initial-turn UserMessage emission in loop.go and firstKeptSessionIndex in compactor.go now use sess.ItemCount() instead of len(sess.Messages). - EmitStartupInfo's session-restore LastMessage reconstruction in runtime.go now iterates a MessagesSnapshot() instead of sess.Messages. Refs #3590 --- pkg/runtime/compactor/compactor.go | 7 +- pkg/runtime/compactor/compactor_test.go | 25 +++++++ pkg/runtime/loop.go | 10 +-- pkg/runtime/runtime.go | 9 ++- pkg/runtime/session_messages_race_test.go | 88 +++++++++++++++++++++++ 5 files changed, 129 insertions(+), 10 deletions(-) create mode 100644 pkg/runtime/session_messages_race_test.go diff --git a/pkg/runtime/compactor/compactor.go b/pkg/runtime/compactor/compactor.go index 6e8758f6c9..bdff8ed962 100644 --- a/pkg/runtime/compactor/compactor.go +++ b/pkg/runtime/compactor/compactor.go @@ -337,13 +337,14 @@ func extractMessages(sess *session.Session, _ *agent.Agent, contextLimit int64, // firstKeptSessionIndex translates a split index produced against the // chat-message list returned by [gatherCompactionInput] back to an // index in sess.Messages, suitable for the new summary's -// FirstKeptEntry. Out-of-range splits map to len(sess.Messages), +// FirstKeptEntry. Out-of-range splits map to sess.ItemCount(), // matching the "compact everything; keep nothing of the tail" // sentinel that session.buildSessionSummaryMessages handles by -// skipping the conversation loop. +// skipping the conversation loop. ItemCount takes sess.mu so this +// stays race-free against a concurrent AddMessage/ApplyCompaction. func firstKeptSessionIndex(sess *session.Session, sessIndices []int, splitIdx int) int { if splitIdx >= len(sessIndices) { - return len(sess.Messages) + return sess.ItemCount() } return sessIndices[splitIdx] } diff --git a/pkg/runtime/compactor/compactor_test.go b/pkg/runtime/compactor/compactor_test.go index c089f6ea89..e33325e6ad 100644 --- a/pkg/runtime/compactor/compactor_test.go +++ b/pkg/runtime/compactor/compactor_test.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "strings" + "sync" "testing" "github.com/stretchr/testify/assert" @@ -391,6 +392,30 @@ func TestFirstKeptSessionIndex_SplitZeroOnEmptyInputUsesSafeSentinel(t *testing. assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sess, sessIndices, 0)) } +// TestFirstKeptSessionIndexConcurrent pins the data-race fix for issue +// #3590: firstKeptSessionIndex's out-of-range sentinel must read the +// session's item count through the locked ItemCount accessor, not +// len(sess.Messages) directly, so it stays race-free against a concurrent +// AddMessage/ApplyCompaction on the same live session (e.g. a live HTTP +// AddMessage arriving mid-compaction). Run with -race; before the fix, +// len(sess.Messages) aliases the live backing array and races the +// concurrent AddMessage goroutine below. +func TestFirstKeptSessionIndexConcurrent(t *testing.T) { + t.Parallel() + + sess := session.New() + var wg sync.WaitGroup + for range 100 { + wg.Go(func() { + sess.AddMessage(session.UserMessage("u")) + }) + wg.Go(func() { + _ = firstKeptSessionIndex(sess, nil, 0) + }) + } + wg.Wait() +} + // TestGatherCompactionInput_PriorSummaryWithoutFirstKeptEntry covers // the case where a prior summary was applied as "compact everything, // keep nothing" (FirstKeptEntry left at zero): the iteration must diff --git a/pkg/runtime/loop.go b/pkg/runtime/loop.go index 994ac86f25..e9de2eeb7c 100644 --- a/pkg/runtime/loop.go +++ b/pkg/runtime/loop.go @@ -58,8 +58,8 @@ func (r *LocalRuntime) registerDefaultTools() { // appendSteerAndEmit adds a steer message to the session and emits the corresponding event. func (r *LocalRuntime) appendSteerAndEmit(sess *session.Session, sm QueuedMessage, events EventSink) { - sess.AddMessage(session.UserMessage(sm.Content, sm.MultiContent...)) - events.Emit(UserMessage(sm.Content, sess.ID, sm.MultiContent, len(sess.Messages)-1)) + pos := sess.AddMessage(session.UserMessage(sm.Content, sm.MultiContent...)) + events.Emit(UserMessage(sm.Content, sess.ID, sm.MultiContent, pos)) } // drainAndEmitSteered drains all messages from the steer queue and injects @@ -372,7 +372,7 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session, // signal here too: "a real user prompt is at the tail of the session". if sess.SendUserMessage && len(messages) > 0 { lastMsg := messages[len(messages)-1] - sink.Emit(UserMessage(lastMsg.Content, sess.ID, lastMsg.MultiContent, len(sess.Messages)-1)) + sink.Emit(UserMessage(lastMsg.Content, sess.ID, lastMsg.MultiContent, sess.ItemCount()-1)) // user_prompt_submit fires once per real user message, after // session_start and before the first model call. @@ -959,8 +959,8 @@ func (r *LocalRuntime) runTurn( // undivided agent turn. if followUp, ok := r.followUpQueue.Dequeue(ctx); ok { userMsg := session.UserMessage(followUp.Content, followUp.MultiContent...) - sess.AddMessage(userMsg) - events.Emit(UserMessage(followUp.Content, sess.ID, followUp.MultiContent, len(sess.Messages)-1)) + pos := sess.AddMessage(userMsg) + events.Emit(UserMessage(followUp.Content, sess.ID, followUp.MultiContent, pos)) stop, msg, ctxMsgs := r.executeUserFollowupSubmitHooks(ctx, sess, a, followUp.Content, events) if stop { slog.WarnContext(ctx, "user_followup_submit hook signalled run termination", diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index fa6a5365f1..13a7cdba66 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -1570,8 +1570,13 @@ func (r *LocalRuntime) EmitStartupInfo(ctx context.Context, sess *session.Sessio // parent agent's state: this event carries the parent session_id, // and sub-agents emit their own token_usage events with their own // session_id during live streaming. - for i := range slices.Backward(sess.Messages) { - item := &sess.Messages[i] + // + // MessagesSnapshot takes sess.mu so this cannot race a concurrent + // AddMessage/ApplyCompaction (e.g. a live HTTP AddMessage arriving + // while startup info is being emitted for a restored session). + items := sess.MessagesSnapshot() + for i := range slices.Backward(items) { + item := &items[i] if !item.IsMessage() || item.Message.Message.Role != chat.MessageRoleAssistant { continue } diff --git a/pkg/runtime/session_messages_race_test.go b/pkg/runtime/session_messages_race_test.go new file mode 100644 index 0000000000..cb7bf01fb5 --- /dev/null +++ b/pkg/runtime/session_messages_race_test.go @@ -0,0 +1,88 @@ +package runtime + +import ( + "sync" + "testing" + + "github.com/docker/docker-agent/pkg/chat" + "github.com/docker/docker-agent/pkg/session" +) + +// TestAppendSteerAndEmitConcurrentWithAddMessage pins the data-race fix for +// loop.go's appendSteerAndEmit (issue #3590): the emitted +// UserMessageEvent.SessionPosition must come from AddMessage's own atomic +// return value, not a separate len(sess.Messages)-1 read taken after the +// fact. Run with -race: without the fix, reading len(sess.Messages) +// unlocked races the concurrent AddMessage goroutine simulating a live HTTP +// AddMessage/compaction on the same session. Beyond race-freedom, this also +// asserts the emitted position always resolves back to the message that +// call actually appended, which a racy read cannot guarantee. +func TestAppendSteerAndEmitConcurrentWithAddMessage(t *testing.T) { + t.Parallel() + + rt, _ := newTestRuntime(t) + sess := session.New() + + const n = 100 + events := make(chan Event, n) + sink := NewChannelSink(events) + + var wg sync.WaitGroup + for range n { + wg.Go(func() { + rt.appendSteerAndEmit(sess, QueuedMessage{Content: "steer"}, sink) + }) + wg.Go(func() { + sess.AddMessage(session.UserMessage("concurrent-http-add")) + }) + } + wg.Wait() + close(events) + + items := sess.MessagesSnapshot() + for ev := range events { + um, ok := ev.(*UserMessageEvent) + if !ok { + t.Fatalf("unexpected event type %T", ev) + } + if um.SessionPosition < 0 || um.SessionPosition >= len(items) { + t.Fatalf("SessionPosition %d out of range [0,%d)", um.SessionPosition, len(items)) + } + item := items[um.SessionPosition] + if !item.IsMessage() || item.Message.Message.Content != um.Message { + t.Fatalf("SessionPosition %d does not point at the emitted message %q", um.SessionPosition, um.Message) + } + } +} + +// TestEmitStartupInfoConcurrentWithAddMessage pins the data-race fix for +// runtime.go's session-restore LastMessage reconstruction (issue #3590): it +// must iterate a MessagesSnapshot rather than sess.Messages directly, so +// restoring startup info for a session cannot race a concurrent +// AddMessage/ApplyCompaction (e.g. a live HTTP AddMessage or the runtime's +// own compaction). Run with -race; before the fix, slices.Backward over the +// live sess.Messages slice races the concurrent appends below. +func TestEmitStartupInfoConcurrentWithAddMessage(t *testing.T) { + t.Parallel() + + rt, _ := newTestRuntime(t) + sess := session.New() + sess.SetUsage(10, 20) + sess.AddMessage(session.NewAgentMessage("root", &chat.Message{ + Role: chat.MessageRoleAssistant, + Content: "hi", + })) + + sink := EventSinkFunc(func(Event) {}) + + done := make(chan struct{}) + go func() { + defer close(done) + for range 200 { + sess.AddMessage(session.UserMessage("concurrent-restore-add")) + } + }() + + rt.EmitStartupInfo(t.Context(), sess, sink) + <-done +} From 4ac052896f583be2967326bd49a8fba6704678bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 13:25:19 +0200 Subject: [PATCH 3/8] fix(server): ForkSession must not iterate a live session's Messages InMemorySessionStore.GetSession returns the live, shared *Session pointer, not a copy. userMessageOrdinalToItemIndex iterated s.Messages directly to translate a user-message ordinal into an item index, racing a concurrent AddMessage on that same session (e.g. the HTTP AddMessage handler racing a TUI fork action) and risking an index computed against a torn read. It now walks s.MessagesSnapshot() instead. Refs #3590 --- pkg/server/session_manager.go | 8 +++++- pkg/server/session_manager_test.go | 40 ++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index d6d77c05e0..068ec9f388 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -460,12 +460,18 @@ func (sm *SessionManager) ForkSession(ctx context.Context, sessionID string, use // userMessageOrdinalToItemIndex maps a 0-based user-message ordinal // into an index in the parent's Session.Messages Item slice. Returns // ErrForkOutOfRange or ErrForkInSubSession on invalid input. +// +// It walks a MessagesSnapshot rather than s.Messages directly: s is the +// live, shared session pointer returned by InMemorySessionStore.GetSession, +// which a concurrent HTTP AddMessage or the runtime's own compaction can +// still be mutating while ForkSession runs. func userMessageOrdinalToItemIndex(s *session.Session, ordinal int) (int, error) { if ordinal < 0 { return 0, fmt.Errorf("%w: %d", ErrForkOutOfRange, ordinal) } + items := s.MessagesSnapshot() seen := 0 - for i, item := range s.Messages { + for i, item := range items { switch { case item.IsMessage(): // Mirror GetAllMessages: system messages don't count. diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index f0e2247913..1a4da0f096 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -443,6 +443,46 @@ func TestForkSession_CopiesHistoryBeforeUserMessage(t *testing.T) { assert.Equal(t, forked.ID, loaded.ID) } +// TestForkSession_ConcurrentWithLiveSessionMutation pins the data-race fix +// for issue #3590: InMemorySessionStore.GetSession returns the live, shared +// *Session pointer (not a copy), so ForkSession's index computation +// (userMessageOrdinalToItemIndex) and session.ForkSession's own copy must +// both go through locked snapshots to stay safe against a concurrent +// AddMessage on that same live session — e.g. the HTTP AddMessage handler +// racing a TUI fork action. Run with -race; before the fix, iterating +// s.Messages directly races the concurrent AddMessage goroutine below. +func TestForkSession_ConcurrentWithLiveSessionMutation(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + parent := session.New() + parent.Title = "Parent Title" + parent.Messages = []session.Item{ + session.NewMessageItem(session.UserMessage("first user")), + session.NewMessageItem(session.NewAgentMessage("root", &chat.Message{ + Role: chat.MessageRoleAssistant, + Content: "first answer", + })), + } + require.NoError(t, store.AddSession(ctx, parent)) + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + + var wg sync.WaitGroup + for range 50 { + wg.Go(func() { + parent.AddMessage(session.UserMessage("concurrent")) + }) + wg.Go(func() { + if _, err := sm.ForkSession(ctx, parent.ID, 0); err != nil { + t.Errorf("ForkSession: %v", err) + } + }) + } + wg.Wait() +} + // Regression: repeated forks of the same parent must pick (fork 1), // (fork 2), (fork 3) rather than three copies of (fork 1). func TestForkSession_TitleIncrementsAcrossSiblings(t *testing.T) { From 011431a84144709b8cd969344b74f8f5911b5c90 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 13:25:29 +0200 Subject: [PATCH 4/8] fix(server): reject AddMessage/UpdateMessage with 409 while streaming session.Session.mu now makes AddMessage/branch/fork race-free against a live RunStream, but a message injected or edited mid-stream (mid-tool-call in particular) can still desynchronize the in-flight turn from what the model/tools expect. As defense in depth on top of the locking, reject the mutation outright while the session has an active stream. The check reuses the same activeRuntimes.streaming lock RunSession already uses to serialize concurrent RunSession calls (TryLock/Unlock), so it can never race a stream that starts between the check and the mutation: both require sm.mux, which AddMessage/UpdateMessage hold for their entire body. The existing ErrSessionBusy sentinel (already mapped to 409 for runAgent) is reused and now also returned by AddMessage/UpdateMessage; server.go maps it to 409 Conflict for both endpoints. Documents the two message endpoints (previously missing from the API reference) and their new 409 behavior. Refs #3590 --- docs/features/api-server/index.md | 2 + pkg/server/server.go | 6 ++ pkg/server/session_manager.go | 26 ++++++ pkg/server/session_manager_test.go | 142 +++++++++++++++++++++++++++++ 4 files changed, 176 insertions(+) diff --git a/docs/features/api-server/index.md b/docs/features/api-server/index.md index 2a1ba6432d..a7e452c744 100644 --- a/docs/features/api-server/index.md +++ b/docs/features/api-server/index.md @@ -51,6 +51,8 @@ All endpoints are under the `/api` prefix. | `PATCH` | `/api/sessions/:id/title` | Update session title | | `PATCH` | `/api/sessions/:id/permissions` | Update session permissions | | `POST` | `/api/sessions/:id/fork` | Fork a session at a user message — creates a new session with messages `[0, message_index)` of the parent (see [Session Forking](#session-forking)) | +| `POST` | `/api/sessions/:id/messages` | Append a message directly to a session's history (bypasses the model). Returns `409 Conflict` while the session has an active run (see [Agent Execution](#agent-execution)). | +| `PATCH` | `/api/sessions/:id/messages/:msg_id` | Update an existing message by ID. Returns `409 Conflict` while the session has an active run. | | `POST` | `/api/sessions/:id/resume` | Resume a paused session (after tool confirmation) | | `POST` | `/api/sessions/:id/tools/toggle` | Toggle auto-approve (YOLO) mode | | `POST` | `/api/sessions/:id/elicitation` | Respond to an MCP tool elicitation request | diff --git a/pkg/server/server.go b/pkg/server/server.go index 6f1c7e3b74..204f8f00ee 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -730,6 +730,9 @@ func (s *Server) addMessage(c echo.Context) error { } if err := s.sm.AddMessage(c.Request().Context(), sessionID, req.Message); err != nil { + if errors.Is(err, ErrSessionBusy) { + return echo.NewHTTPError(http.StatusConflict, err.Error()) + } return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to add message: %v", err)) } @@ -749,6 +752,9 @@ func (s *Server) updateMessage(c echo.Context) error { } if err := s.sm.UpdateMessage(c.Request().Context(), sessionID, msgID, req.Message); err != nil { + if errors.Is(err, ErrSessionBusy) { + return echo.NewHTTPError(http.StatusConflict, err.Error()) + } return echo.NewHTTPError(http.StatusInternalServerError, fmt.Sprintf("failed to update message: %v", err)) } diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 068ec9f388..475e76ea36 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -1271,10 +1271,26 @@ func (sm *SessionManager) GetAgentToolCount(ctx context.Context, agentFilename, } // AddMessage adds a message to a session. +// +// It rejects the mutation with ErrSessionBusy while the session has an +// active RunStream: session.Session.mu makes the append itself race-free, +// but a message added mid-stream (mid-tool-call in particular) can still +// desynchronize the in-flight turn from what the model/tools expect, so we +// also reject at the API boundary. The check reuses the same activeRuntimes +// streaming lock RunSession uses, so it can never race a stream that starts +// between the check and the mutation below: both require sm.mux, which this +// method holds for its entire body. func (sm *SessionManager) AddMessage(ctx context.Context, sessionID string, msg *session.Message) error { sm.mux.Lock() defer sm.mux.Unlock() + if rt, ok := sm.runtimeSessions.Load(sessionID); ok { + if !rt.streaming.TryLock() { + return ErrSessionBusy + } + rt.streaming.Unlock() + } + _, err := sm.sessionStore.AddMessage(ctx, sessionID, msg) if err != nil { return err @@ -1289,10 +1305,20 @@ func (sm *SessionManager) AddMessage(ctx context.Context, sessionID string, msg } // UpdateMessage updates a message in a session. +// +// Rejected with ErrSessionBusy while the session has an active RunStream; +// see AddMessage's comment for why and how the check is race-free. func (sm *SessionManager) UpdateMessage(ctx context.Context, sessionID, msgID string, msg *session.Message) error { sm.mux.Lock() defer sm.mux.Unlock() + if rt, ok := sm.runtimeSessions.Load(sessionID); ok { + if !rt.streaming.TryLock() { + return ErrSessionBusy + } + rt.streaming.Unlock() + } + // Parse msgID as int64 var msgPos int64 _, err := fmt.Sscanf(msgID, "%d", &msgPos) diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 1a4da0f096..15dad1d91e 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -1,14 +1,20 @@ package server import ( + "bytes" "context" + "encoding/json" + "net/http" + "net/http/httptest" "path/filepath" + "strconv" "strings" "sync" "sync/atomic" "testing" "time" + "github.com/labstack/echo/v4" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -191,6 +197,142 @@ func TestRunSession_MessagesNotAddedWhenBusy(t *testing.T) { } } +// TestAddMessage_RejectsWhileSessionStreaming verifies the 409-busy guard +// added for issue #3590: AddMessage must reject with ErrSessionBusy while +// the session has an active RunStream. session.Session.mu already makes the +// append itself race-free, but a message injected mid-stream (mid-tool-call +// in particular) can still desynchronize the turn from what the model/tools +// expect, so the API layer also rejects it outright. +func TestAddMessage_RejectsWhileSessionStreaming(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + release := make(chan struct{}) + fake := &fakeRuntime{release: release} + sm := newTestSessionManager(t, sess, fake) + + ch, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "hi"}}, "") + require.NoError(t, err) + + err = sm.AddMessage(ctx, sess.ID, session.UserMessage("should be rejected")) + require.ErrorIs(t, err, ErrSessionBusy) + + close(release) + for range ch { + } + + // After the stream ends, AddMessage must succeed normally. + require.NoError(t, sm.AddMessage(ctx, sess.ID, session.UserMessage("accepted"))) +} + +// TestUpdateMessage_RejectsWhileSessionStreaming mirrors +// TestAddMessage_RejectsWhileSessionStreaming for UpdateMessage. +func TestUpdateMessage_RejectsWhileSessionStreaming(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + release := make(chan struct{}) + fake := &fakeRuntime{release: release} + sm := newTestSessionManager(t, sess, fake) + + msgID, err := sm.sessionStore.AddMessage(ctx, sess.ID, session.UserMessage("original")) + require.NoError(t, err) + msgIDStr := strconv.FormatInt(msgID, 10) + + ch, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "hi"}}, "") + require.NoError(t, err) + + err = sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("should be rejected")) + require.ErrorIs(t, err, ErrSessionBusy) + + close(release) + for range ch { + } + + // After the stream ends, UpdateMessage must succeed normally. + require.NoError(t, sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("accepted"))) +} + +// TestServer_AddMessage_Returns409WhileSessionStreaming and +// TestServer_UpdateMessage_Returns409WhileSessionStreaming drive the actual +// HTTP handlers (not just SessionManager) to pin the 409-busy guard added +// for issue #3590 end to end: ErrSessionBusy from the manager must surface +// as echo.NewHTTPError(http.StatusConflict, ...), mirroring how runAgent +// already maps it. +func TestServer_AddMessage_Returns409WhileSessionStreaming(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + release := make(chan struct{}) + fake := &fakeRuntime{release: release} + sm := newTestSessionManager(t, sess, fake) + srv := NewWithManager(sm, "") + + ch, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "hi"}}, "") + require.NoError(t, err) + + body, err := json.Marshal(api.AddMessageRequest{Message: session.UserMessage("should be rejected")}) + require.NoError(t, err) + + e := echo.New() + req := httptest.NewRequestWithContext(ctx, http.MethodPost, "/api/sessions/"+sess.ID+"/messages", bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(sess.ID) + + err = srv.addMessage(c) + var httpErr *echo.HTTPError + require.ErrorAs(t, err, &httpErr) + assert.Equal(t, http.StatusConflict, httpErr.Code) + + close(release) + for range ch { + } +} + +func TestServer_UpdateMessage_Returns409WhileSessionStreaming(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + release := make(chan struct{}) + fake := &fakeRuntime{release: release} + sm := newTestSessionManager(t, sess, fake) + srv := NewWithManager(sm, "") + + msgID, err := sm.sessionStore.AddMessage(ctx, sess.ID, session.UserMessage("original")) + require.NoError(t, err) + + ch, err := sm.RunSession(ctx, sess.ID, "agent", "root", []api.Message{{Content: "hi"}}, "") + require.NoError(t, err) + + body, err := json.Marshal(api.UpdateMessageRequest{Message: session.UserMessage("should be rejected")}) + require.NoError(t, err) + + e := echo.New() + msgIDStr := strconv.FormatInt(msgID, 10) + req := httptest.NewRequestWithContext(ctx, http.MethodPatch, "/api/sessions/"+sess.ID+"/messages/"+msgIDStr, bytes.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id", "msg_id") + c.SetParamValues(sess.ID, msgIDStr) + + err = srv.updateMessage(c) + var httpErr *echo.HTTPError + require.ErrorAs(t, err, &httpErr) + assert.Equal(t, http.StatusConflict, httpErr.Code) + + close(release) + for range ch { + } +} + // TestRunSession_SequentialRequestsSucceed verifies that sequential // (non-overlapping) requests on the same session work normally. func TestRunSession_SequentialRequestsSucceed(t *testing.T) { From 565eaa8deea14d19e258f4a3e8fee060d66add51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 14:14:40 +0200 Subject: [PATCH 5/8] fix(server,app): extend 409-busy guard to attached/TUI RunStream ownership AddMessage/UpdateMessage/RunSession detect an active stream via activeRuntimes.streaming, but that lock was only ever acquired by RunSession itself. Runtimes registered via AttachRuntime (the TUI's --listen control plane and local recall coordinator) stream directly through pkg/app.App.Run/Retry/RunWithMessage, which called runtime.RunStream without ever touching that lock, so a REST AddMessage/UpdateMessage (or a second RunSession) during a genuinely active attached/TUI stream wrongly succeeded instead of 409ing. AttachRuntime now returns the same lock (sync.Locker) RunSession and AddMessage/UpdateMessage already TryLock. A new app.WithStreamGuard option lets the App hold it for the duration of every direct RunStream call, acquired before mutating the session and released only once the stream ends, mirroring RunSession's own ordering. cmd/root wires this option into both the --listen and local-recall paths. Adds SessionManager and App-level regression tests: driving a real attached stream through the returned lock (not through RunSession) now gets AddMessage/UpdateMessage to correctly return ErrSessionBusy, and a blocking-runtime App test pins that the lock is held for the stream's actual duration, not just while scheduling it. Refs #3590 --- cmd/root/run.go | 7 ++- cmd/root/run_listen.go | 14 ++++-- pkg/app/app.go | 41 +++++++++++++++++ pkg/app/app_test.go | 74 ++++++++++++++++++++++++++++++ pkg/server/session_manager.go | 19 ++++++-- pkg/server/session_manager_test.go | 51 ++++++++++++++++++++ 6 files changed, 199 insertions(+), 7 deletions(-) diff --git a/cmd/root/run.go b/cmd/root/run.go index ac56bb7b6c..017f3c5f54 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -463,6 +463,10 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s return f.handleExecMode(ctx, out, rt, sess, args) } + // startSessionCoordinator wires the App's own RunStream calls into the + // same busy guard RunSession/AddMessage/UpdateMessage use, so a + // concurrent REST call is correctly rejected with 409 instead of racing + // the TUI's live stream (#3590); see app.WithStreamGuard. coordinatorOpt, err := f.startSessionCoordinator(ctx, out, rt, sess) if err != nil { return err @@ -1104,7 +1108,8 @@ func (f *runExecFlags) createSessionSpawner(agentSource config.Source, sessStore if ctrl != nil { appOpts = append(appOpts, app.WithSnapshotController(ctrl)) } - if coordinatorOpt := f.recallCoordinatorOpt(spawnCtx, localRt, newSess); coordinatorOpt != nil { + coordinatorOpt := f.recallCoordinatorOpt(spawnCtx, localRt, newSess) + if coordinatorOpt != nil { appOpts = append(appOpts, coordinatorOpt) } diff --git a/cmd/root/run_listen.go b/cmd/root/run_listen.go index f6344c45c2..d420f68ce5 100644 --- a/cmd/root/run_listen.go +++ b/cmd/root/run_listen.go @@ -21,10 +21,16 @@ import ( // in-process runtime even when the HTTP control plane is disabled. That lets a // background tool wake an idle local TUI by routing through the same injector // used by control-plane follow-ups. +// +// It also wires app.WithStreamGuard so the App's own direct RunStream calls +// (Run/Retry/RunWithMessage) hold the same lock RunSession/AddMessage/ +// UpdateMessage use to detect an active stream, closing the gap where a +// concurrent REST mutation could slip in during an attached/TUI stream (#3590). func (f *runExecFlags) recallCoordinatorOpt(ctx context.Context, rt runtime.Runtime, sess *session.Session) app.Opt { sm := server.NewSessionManager(ctx, nil, rt.SessionStore(), 0, &f.runConfig) - sm.AttachRuntime(ctx, sess.ID, rt, sess) + guard := sm.AttachRuntime(ctx, sess.ID, rt, sess) return func(a *app.App) { + app.WithStreamGuard(guard)(a) sm.RegisterFollowUpInjector(sess.ID, a.InjectUserMessage) } } @@ -32,14 +38,15 @@ func (f *runExecFlags) recallCoordinatorOpt(ctx context.Context, rt runtime.Runt // startSessionCoordinator wires local recall delivery for the in-process // runtime and, when --listen is set, exposes that runtime over HTTP so // external processes can drive the running TUI (steer, followup, resume, ...). -// It returns an app.Opt that registers the App as the attached session owner. +// It returns an app.Opt that registers the App as the attached session owner +// (see recallCoordinatorOpt for the stream-guard wiring shared by both paths). func (f *runExecFlags) startSessionCoordinator(ctx context.Context, out *cli.Printer, rt runtime.Runtime, sess *session.Session) (app.Opt, error) { if f.listenAddr == "" { return f.recallCoordinatorOpt(ctx, rt, sess), nil } sm := server.NewSessionManager(ctx, nil, rt.SessionStore(), 0, &f.runConfig) - sm.AttachRuntime(ctx, sess.ID, rt, sess) + guard := sm.AttachRuntime(ctx, sess.ID, rt, sess) ln, err := server.Listen(ctx, f.listenAddr) if err != nil { @@ -71,6 +78,7 @@ func (f *runExecFlags) startSessionCoordinator(ctx context.Context, out *cli.Pri }() return func(a *app.App) { + app.WithStreamGuard(guard)(a) sm.RegisterEventSource(sess.ID, func(ctx context.Context, send func(any)) { a.SubscribeWith(ctx, func(msg tea.Msg) { if ev, ok := msg.(runtime.Event); ok { diff --git a/pkg/app/app.go b/pkg/app/app.go index 4ff34bb41d..9838d32647 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -52,6 +52,7 @@ type App struct { titleGenerating atomic.Bool // True when title generation is in progress titleGen *sessiontitle.Generator // Title generator for local runtime (nil for remote) snapshotController builtins.SnapshotController // Drives /undo, /snapshots, /reset; nil for runtimes that don't capture snapshots + streamGuard sync.Locker // Held for the duration of every direct RunStream call; nil when not attached to a SessionManager (see WithStreamGuard) startOnce sync.Once subsMu sync.Mutex @@ -121,6 +122,21 @@ func WithSnapshotController(c builtins.SnapshotController) Opt { } } +// WithStreamGuard makes the App hold l for the duration of every direct +// RunStream call (Run, Retry, RunWithMessage). When the App is attached to a +// SessionManager (the --listen control plane or the local recall +// coordinator), l is the SAME lock RunSession and AddMessage/UpdateMessage +// already use (via TryLock) to detect a session's active stream, so a +// concurrent REST mutation correctly sees the App's own stream as busy +// instead of racing it (#3590). Without this option (a bare App with no +// attached SessionManager) there is nothing to race against, so omitting it +// is safe. +func WithStreamGuard(l sync.Locker) Opt { + return func(a *App) { + a.streamGuard = l + } +} + func New(ctx context.Context, rt runtime.Runtime, sess *session.Session, opts ...Opt) *App { app := &App{ ctx: func() context.Context { return context.WithoutCancel(ctx) }, @@ -474,6 +490,9 @@ func (a *App) Run(ctx context.Context, cancel context.CancelFunc, message string } go func() { + release := a.acquireStreamGuard() + defer release() + if len(attachments) > 0 { multiContent := a.buildUserMultiContent(ctx, message, attachments) a.session.AddMessage(session.UserMessage(message, multiContent...)) @@ -666,6 +685,22 @@ func (a *App) sendEvent(ctx context.Context, event tea.Msg) { } } +// acquireStreamGuard locks a.streamGuard (set via WithStreamGuard) and +// returns the matching release func, or a no-op release when no guard is +// attached (a bare App with no SessionManager to race against). Callers must +// hold the lock for the entire direct RunStream call — acquire it before +// mutating the session and defer the release only after the stream ends — +// mirroring the order SessionManager.RunSession already uses so a concurrent +// AddMessage/UpdateMessage/RunSession sees this stream as busy the same way +// (#3590). +func (a *App) acquireStreamGuard() func() { + if a.streamGuard == nil { + return func() {} + } + a.streamGuard.Lock() + return a.streamGuard.Unlock +} + // processInlineAttachment handles content that is already in memory (e.g. pasted // text). The content is appended to textBuilder wrapped in an XML tag for context. func (a *App) processInlineAttachment(att messages.Attachment, textBuilder *strings.Builder) { @@ -690,6 +725,9 @@ func (a *App) Retry(ctx context.Context, cancel context.CancelFunc) { a.cancel = cancel go func() { + release := a.acquireStreamGuard() + defer release() + streamStarted := false for event := range a.runtime.RunStream(ctx, a.session) { // If context is cancelled, continue draining but don't forward events @@ -742,6 +780,9 @@ func (a *App) RunWithMessage(ctx context.Context, cancel context.CancelFunc, msg } go func() { + release := a.acquireStreamGuard() + defer release() + a.session.AddMessage(msg) for event := range a.runtime.RunStream(ctx, a.session) { // If context is cancelled, continue draining but don't forward events diff --git a/pkg/app/app_test.go b/pkg/app/app_test.go index 1ce4a5331f..2e1946e2fa 100644 --- a/pkg/app/app_test.go +++ b/pkg/app/app_test.go @@ -134,6 +134,80 @@ func (m *retryMockRuntime) RunStream(_ context.Context, sess *session.Session) < return ch } +// blockingRunStreamRuntime's RunStream blocks until release is closed (or +// ctx is cancelled), letting tests observe a streamGuard held for the +// stream's actual duration instead of racing a fake runtime that returns +// immediately. +type blockingRunStreamRuntime struct { + mockRuntime + + release chan struct{} +} + +func (r *blockingRunStreamRuntime) RunStream(ctx context.Context, _ *session.Session) <-chan runtime.Event { + ch := make(chan runtime.Event) + go func() { + defer close(ch) + select { + case <-r.release: + case <-ctx.Done(): + } + }() + return ch +} + +// TestApp_Run_HoldsStreamGuardForStreamDuration pins the #3590 fix: Run must +// hold the WithStreamGuard lock for the entire direct RunStream call, not +// just while scheduling it, so a concurrent SessionManager.AddMessage/ +// UpdateMessage/RunSession sharing the same lock (see AttachRuntime) sees +// the attached stream as busy for as long as it is genuinely active. +func TestApp_Run_HoldsStreamGuardForStreamDuration(t *testing.T) { + t.Parallel() + + release := make(chan struct{}) + rt := &blockingRunStreamRuntime{release: release} + var guard sync.Mutex + + app := &App{ + runtime: rt, + session: session.New(), + events: make(chan tea.Msg, 16), + streamGuard: &guard, + } + + ctx, cancel := context.WithCancel(t.Context()) + defer cancel() + app.Run(ctx, cancel, "hello", nil) + + // isLocked probes guard without leaking a spurious lock acquisition: + // TryLock succeeding would otherwise leave the test goroutine holding + // the mutex, making the later "released" check hang forever. + isLocked := func() bool { + if guard.TryLock() { + guard.Unlock() + return false + } + return true + } + + require.Eventually(t, isLocked, time.Second, time.Millisecond, "streamGuard should be held while RunStream is active") + + close(release) + + require.Eventually(t, func() bool { return !isLocked() }, time.Second, time.Millisecond, "streamGuard should be released once RunStream ends") +} + +// TestApp_AcquireStreamGuard_NoopWhenUnset verifies that a bare App with no +// WithStreamGuard option (the common case: no attached SessionManager) never +// blocks on a nil lock. +func TestApp_AcquireStreamGuard_NoopWhenUnset(t *testing.T) { + t.Parallel() + + app := &App{} + release := app.acquireStreamGuard() + require.NotPanics(t, release) +} + func TestApp_Retry_SuppressesReEmittedUserMessage(t *testing.T) { t.Parallel() diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index 475e76ea36..b7dc817fc7 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -219,16 +219,29 @@ func (sm *SessionManager) StreamEvents(ctx context.Context, sessionID string, si // The internal cancellation signal is fired by [SessionManager.DeleteSession]; // SSE streams and other lifetime-bound consumers use it (via // [SessionManager.StreamEvents]) to terminate when the session is detached. -func (sm *SessionManager) AttachRuntime(ctx context.Context, sessionID string, rt runtime.Runtime, sess *session.Session) { +// +// It returns the same lock RunSession and AddMessage/UpdateMessage already +// use (via TryLock) to detect and reject concurrent mutations while a stream +// is active. Callers that stream the attached runtime directly — bypassing +// RunSession entirely, e.g. the TUI's App.Run/Retry/RunWithMessage calling +// rt.RunStream itself — previously left that lock unheld for the whole +// attached/TUI stream, so a concurrent AddMessage/UpdateMessage or +// RunSession wrongly observed the session as idle instead of 409ing (#3590). +// The caller must hold this lock for the duration of every direct RunStream +// call (see the pkg/app WithStreamGuard option) so the busy check sees +// attached streams too. +func (sm *SessionManager) AttachRuntime(ctx context.Context, sessionID string, rt runtime.Runtime, sess *session.Session) sync.Locker { ctx, cancel := context.WithCancel(context.WithoutCancel(ctx)) - sm.runtimeSessions.Store(sessionID, &activeRuntimes{ + rs := &activeRuntimes{ runtime: rt, done: ctx.Done(), cancel: cancel, session: sess, - }) + } + sm.runtimeSessions.Store(sessionID, rs) sm.registerRecallHandler(sessionID, rt) sm.markReady() + return &rs.streaming } // GetSession retrieves a session by ID. diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 15dad1d91e..754ae56888 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -255,6 +255,57 @@ func TestUpdateMessage_RejectsWhileSessionStreaming(t *testing.T) { require.NoError(t, sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("accepted"))) } +// TestAttachedStream_AddMessageAndUpdateMessageRejectWhileStreaming pins the +// fix for the other #3590 blocker: runtimes attached via AttachRuntime +// stream directly through RunStream (see pkg/app.App.Run/Retry/ +// RunWithMessage), never going through RunSession, which is the only place +// that used to acquire activeRuntimes.streaming. Before the fix nothing held +// that lock for an attached stream, so AddMessage/UpdateMessage wrongly +// succeeded during a genuinely active attached stream instead of returning +// ErrSessionBusy. AttachRuntime now returns the same lock RunSession uses; +// the App holds it for the duration of every direct RunStream call (see +// app.WithStreamGuard/acquireStreamGuard). This test drives a real stream +// through that lock exactly the way the App does — NOT through +// sm.RunSession — to prove the guard covers the attached path too. +func TestAttachedStream_AddMessageAndUpdateMessageRejectWhileStreaming(t *testing.T) { + t.Parallel() + + ctx := t.Context() + store := session.NewInMemorySessionStore() + sess := session.New() + require.NoError(t, store.AddSession(ctx, sess)) + + msgID, err := store.AddMessage(ctx, sess.ID, session.UserMessage("original")) + require.NoError(t, err) + msgIDStr := strconv.FormatInt(msgID, 10) + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + release := make(chan struct{}) + fake := &fakeRuntime{release: release} + guard := sm.AttachRuntime(ctx, sess.ID, fake, sess) + + // Simulate the TUI/attached owner streaming directly through the + // runtime, exactly like pkg/app.App.acquireStreamGuard + Run do — NOT + // through sm.RunSession. + guard.Lock() + ch := fake.RunStream(ctx, sess) + + err = sm.AddMessage(ctx, sess.ID, session.UserMessage("should be rejected")) + require.ErrorIs(t, err, ErrSessionBusy) + + err = sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("should be rejected")) + require.ErrorIs(t, err, ErrSessionBusy) + + close(release) + for range ch { + } + guard.Unlock() + + // After the attached stream ends, both must succeed normally. + require.NoError(t, sm.AddMessage(ctx, sess.ID, session.UserMessage("accepted"))) + require.NoError(t, sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("accepted"))) +} + // TestServer_AddMessage_Returns409WhileSessionStreaming and // TestServer_UpdateMessage_Returns409WhileSessionStreaming drive the actual // HTTP handlers (not just SessionManager) to pin the 409-busy guard added From 1f7c030ecbf5d5b979d042061587dca24ca3626e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 14:14:52 +0200 Subject: [PATCH 6/8] fix(session,runtime): compaction sentinel must use CompactionInput's own snapshot count gatherCompactionInput took a locked snapshot via Session.CompactionInput, but firstKeptSessionIndex's out-of-range sentinel (the 'compact everything, keep nothing' boundary) called sess.ItemCount() again instead of reusing that snapshot's own length. ItemCount is lock-safe on its own, but it can observe an AddMessage/ApplyCompaction that landed after CompactionInput's snapshot was taken, so the sentinel could describe a longer session than the one messages/sessIndices actually cover, breaking the invariant that messages appended after the snapshot stay in the kept tail. CompactionInput now returns the snapshot's own item count alongside messages and sessIndices; gatherCompactionInput and firstKeptSessionIndex thread that count through instead of re-querying the live session, so the sentinel always describes the SAME snapshot the split index was computed against. Adds a deterministic test (no goroutines/-race needed: this was a logic error, not a data race) that gathers the compaction input, lets a message race in afterwards, and asserts the out-of-range sentinel still matches the original snapshot count rather than the now-longer live session. Also replaces the now-obsolete TestFirstKeptSessionIndexConcurrent (firstKeptSessionIndex no longer touches the session) with TestGatherCompactionInputConcurrent, which keeps -race coverage on gatherCompactionInput's own snapshot read. Refs #3590 --- benchmarks/compaction/main.go | 2 +- pkg/runtime/compactor/compactor.go | 39 +++++++----- pkg/runtime/compactor/compactor_test.go | 82 ++++++++++++++++++------- pkg/runtime/compactor/scenarios_test.go | 6 +- pkg/session/session.go | 12 +++- pkg/session/session_race_test.go | 2 +- pkg/session/session_test.go | 14 +++-- 7 files changed, 109 insertions(+), 48 deletions(-) diff --git a/benchmarks/compaction/main.go b/benchmarks/compaction/main.go index a38fe74276..f347af8f56 100644 --- a/benchmarks/compaction/main.go +++ b/benchmarks/compaction/main.go @@ -233,7 +233,7 @@ func run(ctx context.Context) error { // compaction input trimmed of its keep-tail, truncated to the context budget, // wrapped in the canonical compaction system/user prompts. func buildCompactionMessages(sess *session.Session) []chat.Message { - messages, _ := sess.CompactionInput() + messages, _, _ := sess.CompactionInput() for i := range messages { messages[i].Cost = 0 messages[i].CacheControl = false diff --git a/pkg/runtime/compactor/compactor.go b/pkg/runtime/compactor/compactor.go index bdff8ed962..00ba5ecb3a 100644 --- a/pkg/runtime/compactor/compactor.go +++ b/pkg/runtime/compactor/compactor.go @@ -251,8 +251,8 @@ func RunLLM(ctx context.Context, args LLMArgs) (result *Result, err error) { // a hook supplies its own summary so the kept-tail policy stays // consistent across the two strategies. func ComputeFirstKeptEntry(sess *session.Session, contextLimit int64) int { - messages, sessIndices := gatherCompactionInput(sess) - return firstKeptSessionIndex(sess, sessIndices, compaction.SplitIndexForKeep(messages, keepTokenBudget(contextLimit))) + messages, sessIndices, itemCount := gatherCompactionInput(sess) + return firstKeptSessionIndex(sessIndices, itemCount, compaction.SplitIndexForKeep(messages, keepTokenBudget(contextLimit))) } // gatherCompactionInput is a thin wrapper around @@ -272,13 +272,19 @@ func ComputeFirstKeptEntry(sess *session.Session, contextLimit int64) int { // past the prior summary, and tracking origin indices in sess.Messages // — lives on Session itself so it can run under sess.mu.RLock and stay // race-safe against concurrent AddMessage / ApplyCompaction calls. -func gatherCompactionInput(sess *session.Session) ([]chat.Message, []int) { - messages, sessIndices := sess.CompactionInput() +// +// The returned itemCount is the SAME snapshot's total item count +// (session.Session.CompactionInput's third return value); callers must use +// it — not a fresh sess.ItemCount() call — for any out-of-range sentinel +// derived from messages/sessIndices, or the boundary can describe a later +// session length than the snapshot it is supposed to describe (#3590). +func gatherCompactionInput(sess *session.Session) (messages []chat.Message, sessIndices []int, itemCount int) { + messages, sessIndices, itemCount = sess.CompactionInput() for i := range messages { messages[i].Cost = 0 messages[i].CacheControl = false } - return messages, sessIndices + return messages, sessIndices, itemCount } // extractMessages returns the messages to send to the compaction @@ -297,10 +303,10 @@ func gatherCompactionInput(sess *session.Session) ([]chat.Message, []int) { // (contextLimit − summary budget − prompt-overhead), older messages // are dropped from the front of the to-compact list to make room. func extractMessages(sess *session.Session, _ *agent.Agent, contextLimit int64, additionalPrompt string) ([]chat.Message, int) { - messages, sessIndices := gatherCompactionInput(sess) + messages, sessIndices, itemCount := gatherCompactionInput(sess) splitIdx := compaction.SplitIndexForKeep(messages, keepTokenBudget(contextLimit)) - firstKeptEntry := firstKeptSessionIndex(sess, sessIndices, splitIdx) + firstKeptEntry := firstKeptSessionIndex(sessIndices, itemCount, splitIdx) messages = messages[:splitIdx] systemPromptMessage := chat.Message{ @@ -337,14 +343,19 @@ func extractMessages(sess *session.Session, _ *agent.Agent, contextLimit int64, // firstKeptSessionIndex translates a split index produced against the // chat-message list returned by [gatherCompactionInput] back to an // index in sess.Messages, suitable for the new summary's -// FirstKeptEntry. Out-of-range splits map to sess.ItemCount(), -// matching the "compact everything; keep nothing of the tail" -// sentinel that session.buildSessionSummaryMessages handles by -// skipping the conversation loop. ItemCount takes sess.mu so this -// stays race-free against a concurrent AddMessage/ApplyCompaction. -func firstKeptSessionIndex(sess *session.Session, sessIndices []int, splitIdx int) int { +// FirstKeptEntry. Out-of-range splits map to itemCount, matching the +// "compact everything; keep nothing of the tail" sentinel that +// session.buildSessionSummaryMessages handles by skipping the +// conversation loop. +// +// itemCount MUST come from the same [session.Session.CompactionInput] +// snapshot as sessIndices (i.e. via gatherCompactionInput), not a fresh +// sess.ItemCount() call: the live session can already hold an append +// that landed after the snapshot was taken, which would describe a +// boundary the snapshot never had (#3590). +func firstKeptSessionIndex(sessIndices []int, itemCount, splitIdx int) int { if splitIdx >= len(sessIndices) { - return sess.ItemCount() + return itemCount } return sessIndices[splitIdx] } diff --git a/pkg/runtime/compactor/compactor_test.go b/pkg/runtime/compactor/compactor_test.go index e33325e6ad..781ee060a4 100644 --- a/pkg/runtime/compactor/compactor_test.go +++ b/pkg/runtime/compactor/compactor_test.go @@ -285,15 +285,15 @@ func TestGatherCompactionInput_NoPriorSummary(t *testing.T) { session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u2"}}), })) - messages, sessIndices := gatherCompactionInput(sess) + messages, sessIndices, itemCount := gatherCompactionInput(sess) require.Len(t, messages, 3) assert.Equal(t, []int{1, 2, 4}, sessIndices) - assert.Equal(t, 1, firstKeptSessionIndex(sess, sessIndices, 0)) - assert.Equal(t, 2, firstKeptSessionIndex(sess, sessIndices, 1)) - assert.Equal(t, 4, firstKeptSessionIndex(sess, sessIndices, 2)) + assert.Equal(t, 1, firstKeptSessionIndex(sessIndices, itemCount, 0)) + assert.Equal(t, 2, firstKeptSessionIndex(sessIndices, itemCount, 1)) + assert.Equal(t, 4, firstKeptSessionIndex(sessIndices, itemCount, 2)) // Past the end: returns len(sess.Messages) (compact-everything sentinel). - assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sess, sessIndices, 3)) + assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sessIndices, itemCount, 3)) } // TestGatherCompactionInput_WithPriorSummary pins the regression where @@ -333,7 +333,7 @@ func TestGatherCompactionInput_WithPriorSummary(t *testing.T) { } sess := session.New(session.WithMessages(items)) - messages, sessIndices := gatherCompactionInput(sess) + messages, sessIndices, itemCount := gatherCompactionInput(sess) // Expected filtered list: // [0]: synthetic Session Summary user message (origin: prior summary at idx 10) @@ -351,16 +351,16 @@ func TestGatherCompactionInput_WithPriorSummary(t *testing.T) { // A split that keeps the last two messages should map to items[13] // (the user message at idx 13), not to items[5] which is what the // old count-from-zero implementation produced. - assert.Equal(t, 13, firstKeptSessionIndex(sess, sessIndices, 5)) + assert.Equal(t, 13, firstKeptSessionIndex(sessIndices, itemCount, 5)) // A split that keeps the entire post-summary tail (everything from // items[8] onwards including the prior summary) maps the synthetic // message back to its originating summary index so the prior // summary item is preserved across the new compaction. - assert.Equal(t, 10, firstKeptSessionIndex(sess, sessIndices, 0)) + assert.Equal(t, 10, firstKeptSessionIndex(sessIndices, itemCount, 0)) // Out-of-range split: compact everything, keep nothing. - assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sess, sessIndices, len(messages))) + assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sessIndices, itemCount, len(messages))) } // TestFirstKeptSessionIndex_SplitZeroOnEmptyInputUsesSafeSentinel @@ -385,22 +385,25 @@ func TestFirstKeptSessionIndex_SplitZeroOnEmptyInputUsesSafeSentinel(t *testing. sess := session.New() var sessIndices []int + itemCount := sess.ItemCount() // Empty input is the only legitimate way splitIdx==0 reaches // firstKeptSessionIndex. Both branches (>= len(sessIndices) and - // the indexed lookup) must yield len(sess.Messages) here. - assert.Equal(t, len(sess.Messages), firstKeptSessionIndex(sess, sessIndices, 0)) + // the indexed lookup) must yield itemCount here. + assert.Equal(t, itemCount, firstKeptSessionIndex(sessIndices, itemCount, 0)) } -// TestFirstKeptSessionIndexConcurrent pins the data-race fix for issue -// #3590: firstKeptSessionIndex's out-of-range sentinel must read the -// session's item count through the locked ItemCount accessor, not -// len(sess.Messages) directly, so it stays race-free against a concurrent -// AddMessage/ApplyCompaction on the same live session (e.g. a live HTTP -// AddMessage arriving mid-compaction). Run with -race; before the fix, -// len(sess.Messages) aliases the live backing array and races the -// concurrent AddMessage goroutine below. -func TestFirstKeptSessionIndexConcurrent(t *testing.T) { +// TestGatherCompactionInputConcurrent extends session's own +// TestCompactionInputConcurrent (pkg/session/session_race_test.go) to the +// compactor's own boundary computation: gatherCompactionInput plus +// firstKeptSessionIndex must stay race-free when called concurrently with +// AddMessage on the same live session. firstKeptSessionIndex itself no +// longer touches the session (it now takes the snapshot's own itemCount +// instead of calling sess.ItemCount() — see +// TestGatherCompactionInput_OutOfRangeSentinelMatchesSnapshotCount below for +// why), so the only thing left to prove race-free here is +// gatherCompactionInput's snapshot read itself. +func TestGatherCompactionInputConcurrent(t *testing.T) { t.Parallel() sess := session.New() @@ -410,12 +413,47 @@ func TestFirstKeptSessionIndexConcurrent(t *testing.T) { sess.AddMessage(session.UserMessage("u")) }) wg.Go(func() { - _ = firstKeptSessionIndex(sess, nil, 0) + _, sessIndices, itemCount := gatherCompactionInput(sess) + _ = firstKeptSessionIndex(sessIndices, itemCount, 0) }) } wg.Wait() } +// TestGatherCompactionInput_OutOfRangeSentinelMatchesSnapshotCount pins the +// other #3590 blocker: the out-of-range sentinel returned by +// firstKeptSessionIndex must describe the SAME snapshot gatherCompactionInput +// produced sessIndices from, not whatever sess.ItemCount() returns when read +// later. A reviewer probe demonstrated the bug directly: the sentinel +// described a session length one longer (200001) than the snapshot it was +// supposedly derived from (200000) once a message was appended in between. +// This reproduces that shape deterministically — no goroutines or -race +// needed, since the bug was a logic error (reading two different sources of +// truth), not a data race. +func TestGatherCompactionInput_OutOfRangeSentinelMatchesSnapshotCount(t *testing.T) { + t.Parallel() + + sess := session.New(session.WithMessages([]session.Item{ + session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u1"}}), + session.NewMessageItem(&session.Message{Message: chat.Message{Role: chat.MessageRoleAssistant, Content: "a1"}}), + })) + + _, sessIndices, itemCount := gatherCompactionInput(sess) + snapshotCount := len(sess.Messages) + require.Equal(t, snapshotCount, itemCount) + + // A message "races in" after the snapshot was taken, e.g. a concurrent + // HTTP AddMessage landing mid-compaction. + sess.AddMessage(session.UserMessage("late arrival")) + require.Equal(t, snapshotCount+1, sess.ItemCount(), "sanity: the live session grew past the snapshot") + + // The out-of-range sentinel must still describe the snapshot count, not + // the now-longer live session. + got := firstKeptSessionIndex(sessIndices, itemCount, len(sessIndices)) + assert.Equal(t, snapshotCount, got) + assert.NotEqual(t, sess.ItemCount(), got) +} + // TestGatherCompactionInput_PriorSummaryWithoutFirstKeptEntry covers // the case where a prior summary was applied as "compact everything, // keep nothing" (FirstKeptEntry left at zero): the iteration must @@ -437,7 +475,7 @@ func TestGatherCompactionInput_PriorSummaryWithoutFirstKeptEntry(t *testing.T) { } sess := session.New(session.WithMessages(items)) - messages, sessIndices := gatherCompactionInput(sess) + messages, sessIndices, _ := gatherCompactionInput(sess) // Filtered list: synthetic-summary, items[3], items[4]. // items[0..1] are excluded because they were compacted into the diff --git a/pkg/runtime/compactor/scenarios_test.go b/pkg/runtime/compactor/scenarios_test.go index e1e1dd61b9..beeb48910f 100644 --- a/pkg/runtime/compactor/scenarios_test.go +++ b/pkg/runtime/compactor/scenarios_test.go @@ -38,7 +38,7 @@ func TestScenario_SecondRoundCompaction_PartitionIsExact(t *testing.T) { } sess := session.New(session.WithMessages(items)) - messages, sessIndices := gatherCompactionInput(sess) + messages, sessIndices, _ := gatherCompactionInput(sess) // synthetic summary(->4), m2(->2), m3(->3), m5..m8(->5..8) require.Equal(t, []int{4, 2, 3, 5, 6, 7, 8}, sessIndices) require.Contains(t, messages[0].Content, "Session Summary: old summary") @@ -54,7 +54,7 @@ func TestScenario_SecondRoundCompaction_PartitionIsExact(t *testing.T) { // Reconstruct what the next prompt will contain via a third // CompactionInput round (mirrors buildSessionSummaryMessages). - after, afterIdx := gatherCompactionInput(sess) + after, afterIdx, _ := gatherCompactionInput(sess) require.Contains(t, after[0].Content, "Session Summary: new summary") // Every original message must be either folded (index < firstKept in @@ -112,7 +112,7 @@ func TestScenario_MessagesAppendedDuringCompactionAreKept(t *testing.T) { sess.AddMessage(&session.Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "raced steer"}}) sess.ApplyCompaction(10, 0, session.Item{Summary: "sum", FirstKeptEntry: firstKept}) - msgs, _ := gatherCompactionInput(sess) + msgs, _, _ := gatherCompactionInput(sess) require.Len(t, msgs, 2, "summary + raced message") assert.Contains(t, msgs[0].Content, "Session Summary: sum") assert.Equal(t, "raced steer", msgs[1].Content, "the raced message must be preserved verbatim") diff --git a/pkg/session/session.go b/pkg/session/session.go index bd5242ee2e..251628f21d 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -1360,10 +1360,18 @@ func (s *Session) buildSessionSummaryMessages(items []Item) ([]chat.Message, int // have hidden), supplies its own system/user prompt, and runs through // a sub-runtime that re-applies sanitization on its own session. // +// The third return value is the snapshot's total item count +// (len(s.Messages) at the instant the snapshot was taken). Callers that +// need an out-of-range sentinel for a split computed against messages/ +// sessIndices (i.e. "keep nothing of the tail") must use this value rather +// than a fresh call to ItemCount(): the live count can already include an +// append that happened after this snapshot, which would describe a longer +// session than the one messages/sessIndices actually cover. +// // All work is performed under s.mu.RLock via snapshotItems, so this // method is safe to call concurrently with AddMessage / ApplyCompaction // on the same session. -func (s *Session) CompactionInput() ([]chat.Message, []int) { +func (s *Session) CompactionInput() ([]chat.Message, []int, int) { items := s.snapshotItems() lastSummaryIndex := -1 @@ -1411,7 +1419,7 @@ func (s *Session) CompactionInput() ([]chat.Message, []int) { messages = append(messages, msg) sessIndices = append(sessIndices, i) } - return messages, sessIndices + return messages, sessIndices, len(items) } func (s *Session) GetMessages(a *agent.Agent, extraSystemMessages ...chat.Message) []chat.Message { diff --git a/pkg/session/session_race_test.go b/pkg/session/session_race_test.go index 167ca54170..ab1880894d 100644 --- a/pkg/session/session_race_test.go +++ b/pkg/session/session_race_test.go @@ -39,7 +39,7 @@ func TestCompactionInputConcurrent(t *testing.T) { s.AddMessage(&Message{Message: chat.Message{Role: chat.MessageRoleUser, Content: "u"}}) }) wg.Go(func() { - _, _ = s.CompactionInput() + _, _, _ = s.CompactionInput() }) } // One concurrent ApplyCompaction-shaped write to exercise the same diff --git a/pkg/session/session_test.go b/pkg/session/session_test.go index 3a22aeb6ef..f7c6d82188 100644 --- a/pkg/session/session_test.go +++ b/pkg/session/session_test.go @@ -860,9 +860,10 @@ func TestCompactionInput(t *testing.T) { t.Run("empty session returns empty", func(t *testing.T) { t.Parallel() sess := New() - messages, sessIndices := sess.CompactionInput() + messages, sessIndices, itemCount := sess.CompactionInput() assert.Empty(t, messages) assert.Empty(t, sessIndices) + assert.Zero(t, itemCount) }) t.Run("system messages on the session are filtered out", func(t *testing.T) { @@ -875,12 +876,13 @@ func TestCompactionInput(t *testing.T) { newMsg(chat.MessageRoleUser, "u2"), })) - messages, sessIndices := sess.CompactionInput() + messages, sessIndices, itemCount := sess.CompactionInput() require.Len(t, messages, 3) assert.Equal(t, []int{1, 2, 4}, sessIndices) assert.Equal(t, "u1", messages[0].Content) assert.Equal(t, "a1", messages[1].Content) assert.Equal(t, "u2", messages[2].Content) + assert.Equal(t, len(sess.Messages), itemCount) }) t.Run("prior summary surfaces synthetic message and starts at FirstKeptEntry", func(t *testing.T) { @@ -896,7 +898,7 @@ func TestCompactionInput(t *testing.T) { } sess := New(WithMessages(items)) - messages, sessIndices := sess.CompactionInput() + messages, sessIndices, itemCount := sess.CompactionInput() require.Len(t, messages, 5) assert.Equal(t, chat.MessageRoleUser, messages[0].Role) @@ -905,6 +907,7 @@ func TestCompactionInput(t *testing.T) { // kept-tail then resumes at the prior FirstKeptEntry, skipping // the (non-message) summary item itself. assert.Equal(t, []int{4, 2, 3, 5, 6}, sessIndices) + assert.Equal(t, len(items), itemCount) }) t.Run("prior summary without FirstKeptEntry starts strictly after the summary", func(t *testing.T) { @@ -918,10 +921,11 @@ func TestCompactionInput(t *testing.T) { } sess := New(WithMessages(items)) - messages, sessIndices := sess.CompactionInput() + messages, sessIndices, itemCount := sess.CompactionInput() require.Len(t, messages, 3) assert.Equal(t, []int{2, 3, 4}, sessIndices) + assert.Equal(t, len(items), itemCount) }) t.Run("returned messages are independent copies safe to mutate", func(t *testing.T) { @@ -935,7 +939,7 @@ func TestCompactionInput(t *testing.T) { }}), })) - messages, _ := sess.CompactionInput() + messages, _, _ := sess.CompactionInput() require.Len(t, messages, 1) messages[0].Cost = 0 From a99b2593ea64db5338d6f9114ac11c20c8841eb8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 14:59:30 +0200 Subject: [PATCH 7/8] fix(server): hold streaming lock across AddMessage/UpdateMessage mutation AddMessage/UpdateMessage TryLock activeRuntimes.streaming to reject a mutation while a stream is active, but released the lock immediately after the check, before performing the append/update. Holding sm.mux for the rest of the method does not close the gap for attached runtimes (see AttachRuntime): App.acquireStreamGuard only ever acquires streaming, never sm.mux, so an attached stream could start after the busy check passed but before/during the REST mutation. Both methods now hold streaming via defer until the mutation (store write, plus the in-memory session update for AddMessage) has fully completed. TryLock stays non-blocking, so a busy REST call still returns immediately instead of waiting on an in-progress attached stream, and the attached guard's plain Lock() cannot deadlock against it. Adds a deterministic blocking-store regression test (mirroring the reviewer's TestReview_AttachedGuardCannotStartBetweenBusyCheckAndMutation probe) for both AddMessage and UpdateMessage: a fake store blocks mid-write so the test can assert the attached-stream guard stays blocked until the mutation returns. Confirmed both tests fail without the defer fix and pass with it. Refs #3590 --- pkg/server/session_manager.go | 24 +++--- pkg/server/session_manager_test.go | 120 +++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+), 9 deletions(-) diff --git a/pkg/server/session_manager.go b/pkg/server/session_manager.go index b7dc817fc7..b69bc6b2f4 100644 --- a/pkg/server/session_manager.go +++ b/pkg/server/session_manager.go @@ -1289,19 +1289,24 @@ func (sm *SessionManager) GetAgentToolCount(ctx context.Context, agentFilename, // active RunStream: session.Session.mu makes the append itself race-free, // but a message added mid-stream (mid-tool-call in particular) can still // desynchronize the in-flight turn from what the model/tools expect, so we -// also reject at the API boundary. The check reuses the same activeRuntimes -// streaming lock RunSession uses, so it can never race a stream that starts -// between the check and the mutation below: both require sm.mux, which this -// method holds for its entire body. +// also reject at the API boundary. The busy check TryLocks the same +// activeRuntimes.streaming lock RunSession/AttachRuntime use, and — unlike +// a bare check-then-release probe — HOLDS it across the entire mutation +// below, releasing only via defer once AddMessage returns. sm.mux alone +// cannot close this gap: an attached runtime's stream (see AttachRuntime, +// pkg/app's WithStreamGuard) only ever acquires streaming, never sm.mux, so +// a stream that starts the instant after the TryLock check but before the +// store write completes would otherwise interleave with it (#3590). func (sm *SessionManager) AddMessage(ctx context.Context, sessionID string, msg *session.Message) error { sm.mux.Lock() defer sm.mux.Unlock() - if rt, ok := sm.runtimeSessions.Load(sessionID); ok { + rt, ok := sm.runtimeSessions.Load(sessionID) + if ok { if !rt.streaming.TryLock() { return ErrSessionBusy } - rt.streaming.Unlock() + defer rt.streaming.Unlock() } _, err := sm.sessionStore.AddMessage(ctx, sessionID, msg) @@ -1310,7 +1315,7 @@ func (sm *SessionManager) AddMessage(ctx context.Context, sessionID string, msg } // If the session is actively running, update the in-memory session - if rt, ok := sm.runtimeSessions.Load(sessionID); ok && rt.session != nil { + if ok && rt.session != nil { rt.session.AddMessage(msg) } @@ -1320,7 +1325,8 @@ func (sm *SessionManager) AddMessage(ctx context.Context, sessionID string, msg // UpdateMessage updates a message in a session. // // Rejected with ErrSessionBusy while the session has an active RunStream; -// see AddMessage's comment for why and how the check is race-free. +// see AddMessage's comment for why the busy check holds streaming across +// the whole mutation instead of releasing it right after the check. func (sm *SessionManager) UpdateMessage(ctx context.Context, sessionID, msgID string, msg *session.Message) error { sm.mux.Lock() defer sm.mux.Unlock() @@ -1329,7 +1335,7 @@ func (sm *SessionManager) UpdateMessage(ctx context.Context, sessionID, msgID st if !rt.streaming.TryLock() { return ErrSessionBusy } - rt.streaming.Unlock() + defer rt.streaming.Unlock() } // Parse msgID as int64 diff --git a/pkg/server/session_manager_test.go b/pkg/server/session_manager_test.go index 754ae56888..5289463f89 100644 --- a/pkg/server/session_manager_test.go +++ b/pkg/server/session_manager_test.go @@ -306,6 +306,126 @@ func TestAttachedStream_AddMessageAndUpdateMessageRejectWhileStreaming(t *testin require.NoError(t, sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("accepted"))) } +// blockingStore wraps a session.Store and blocks inside AddMessage/ +// UpdateMessage until release is closed, letting a test pause the manager +// mid-mutation — after the busy check has already passed — to observe +// whether a concurrent attached stream can slip in before the mutation +// actually completes. entered is closed the instant the blocked call is +// reached, so the test can synchronize on it instead of sleeping. +type blockingStore struct { + session.Store + + release chan struct{} + entered chan struct{} +} + +func (s *blockingStore) AddMessage(ctx context.Context, sessionID string, msg *session.Message) (int64, error) { + close(s.entered) + <-s.release + return s.Store.AddMessage(ctx, sessionID, msg) +} + +func (s *blockingStore) UpdateMessage(ctx context.Context, messageID int64, msg *session.Message) error { + close(s.entered) + <-s.release + return s.Store.UpdateMessage(ctx, messageID, msg) +} + +// assertAttachedGuardBlockedDuringMutation drives the reviewer's +// deterministic "blocking-store" probe (#3590 finding A1): it starts +// mutate (an AddMessage or UpdateMessage call) against a store that blocks +// mid-write, waits for the busy check inside mutate to have already passed +// (store.entered closes), and then tries to acquire the attached-stream +// guard exactly the way pkg/app.App.acquireStreamGuard does (a plain +// Lock(), not TryLock()). Before the #3590 fix, AddMessage/UpdateMessage +// released activeRuntimes.streaming immediately after the busy check, so +// the guard acquisition below would succeed while mutate was still +// blocked inside the store write — an attached stream starting between +// the busy check and the mutation completing. With the fix (the streaming +// lock held via defer until mutate returns), the guard acquisition must +// stay blocked until mutate has fully returned. +func assertAttachedGuardBlockedDuringMutation(t *testing.T, guard sync.Locker, store *blockingStore, mutate func() error) { + t.Helper() + + mutateErrCh := make(chan error, 1) + go func() { mutateErrCh <- mutate() }() + + select { + case <-store.entered: + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for the mutation to reach the blocking store") + } + + guardAcquired := make(chan struct{}) + go func() { + guard.Lock() + close(guardAcquired) + }() + + select { + case <-guardAcquired: + t.Fatal("attached stream guard acquired before the REST mutation completed") + case <-time.After(100 * time.Millisecond): + // Expected: the guard stays held across the in-flight mutation. + } + + close(store.release) + require.NoError(t, <-mutateErrCh) + + select { + case <-guardAcquired: + guard.Unlock() + case <-time.After(2 * time.Second): + t.Fatal("attached stream guard never acquired after the REST mutation completed") + } +} + +// TestReview_AttachedGuardCannotStartBetweenBusyCheckAndMutation_AddMessage +// is the reviewer's deterministic regression probe for #3590 finding A1: +// AddMessage must hold activeRuntimes.streaming across its entire mutation, +// not just across the busy check, otherwise an attached stream (the only +// consumer of that lock outside RunSession — see AttachRuntime) can start +// in the gap between the check passing and the store write completing. +func TestReview_AttachedGuardCannotStartBetweenBusyCheckAndMutation_AddMessage(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(ctx, sess)) + store := &blockingStore{Store: inner, release: make(chan struct{}), entered: make(chan struct{})} + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + guard := sm.AttachRuntime(ctx, sess.ID, &fakeRuntime{}, sess) + + assertAttachedGuardBlockedDuringMutation(t, guard, store, func() error { + return sm.AddMessage(ctx, sess.ID, session.UserMessage("mutating")) + }) +} + +// TestReview_AttachedGuardCannotStartBetweenBusyCheckAndMutation_UpdateMessage +// mirrors the AddMessage probe above for UpdateMessage. +func TestReview_AttachedGuardCannotStartBetweenBusyCheckAndMutation_UpdateMessage(t *testing.T) { + t.Parallel() + + ctx := t.Context() + sess := session.New() + inner := session.NewInMemorySessionStore() + require.NoError(t, inner.AddSession(ctx, sess)) + msgID, err := inner.AddMessage(ctx, sess.ID, session.UserMessage("original")) + require.NoError(t, err) + msgIDStr := strconv.FormatInt(msgID, 10) + + store := &blockingStore{Store: inner, release: make(chan struct{}), entered: make(chan struct{})} + + sm := NewSessionManager(ctx, config.Sources{}, store, 0, &config.RuntimeConfig{}) + guard := sm.AttachRuntime(ctx, sess.ID, &fakeRuntime{}, sess) + + assertAttachedGuardBlockedDuringMutation(t, guard, store, func() error { + return sm.UpdateMessage(ctx, sess.ID, msgIDStr, session.UserMessage("mutating")) + }) +} + // TestServer_AddMessage_Returns409WhileSessionStreaming and // TestServer_UpdateMessage_Returns409WhileSessionStreaming drive the actual // HTTP handlers (not just SessionManager) to pin the 409-busy guard added From 0d074335a10507b0e69b90797c08d903453bb8c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= Date: Sat, 11 Jul 2026 14:59:36 +0200 Subject: [PATCH 8/8] fix(app): add missing sync import in app_test.go TestApp_Run_HoldsStreamGuardForStreamDuration declares 'var guard sync.Mutex' but the file never imported sync, so the package failed to compile (pkg/app/app_test.go:169:12: undefined: sync). task dev and the -race suite could not even build. Refs #3590 --- pkg/app/app_test.go | 1 + 1 file changed, 1 insertion(+) diff --git a/pkg/app/app_test.go b/pkg/app/app_test.go index 2e1946e2fa..24a965a2e0 100644 --- a/pkg/app/app_test.go +++ b/pkg/app/app_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "path/filepath" + "sync" "testing" "time"