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/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/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/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..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" @@ -134,6 +135,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/runtime/compactor/compactor.go b/pkg/runtime/compactor/compactor.go index 6e8758f6c9..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,13 +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 len(sess.Messages), -// matching the "compact everything; keep nothing of the tail" -// sentinel that session.buildSessionSummaryMessages handles by -// skipping the conversation loop. -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 len(sess.Messages) + return itemCount } return sessIndices[splitIdx] } diff --git a/pkg/runtime/compactor/compactor_test.go b/pkg/runtime/compactor/compactor_test.go index c089f6ea89..781ee060a4 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" @@ -284,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 @@ -332,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) @@ -350,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 @@ -384,11 +385,73 @@ 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)) +} + +// 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() + var wg sync.WaitGroup + for range 100 { + wg.Go(func() { + sess.AddMessage(session.UserMessage("u")) + }) + wg.Go(func() { + _, 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 @@ -412,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/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 +} 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 d6d77c05e0..b69bc6b2f4 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. @@ -460,12 +473,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. @@ -1265,17 +1284,38 @@ 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 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() + rt, ok := sm.runtimeSessions.Load(sessionID) + if ok { + if !rt.streaming.TryLock() { + return ErrSessionBusy + } + defer rt.streaming.Unlock() + } + _, err := sm.sessionStore.AddMessage(ctx, sessionID, msg) if err != nil { return err } // 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) } @@ -1283,10 +1323,21 @@ 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 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() + if rt, ok := sm.runtimeSessions.Load(sessionID); ok { + if !rt.streaming.TryLock() { + return ErrSessionBusy + } + defer 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 f0e2247913..5289463f89 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,313 @@ 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"))) +} + +// 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"))) +} + +// 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 +// 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) { @@ -443,6 +756,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) { 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..251628f21d 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. @@ -1330,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 @@ -1381,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 2a504bb621..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 @@ -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() +} 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