From a5eec36f4d2683708a6006d8567e4b98a56b91d5 Mon Sep 17 00:00:00 2001 From: Anabella Buckvar Date: Tue, 18 Aug 2026 14:29:01 +0000 Subject: [PATCH] Persist turn closeout messages during shutdown Co-authored-by: Shelley --- loop/loop.go | 41 +++++++++--- loop/loop_test.go | 117 ++++++++++++++++++++++++++++++++++ server/convo.go | 86 ++++++++++++++++++------- server/convo_shutdown_test.go | 84 ++++++++++++++++++++++++ server/server.go | 8 +-- 5 files changed, 299 insertions(+), 37 deletions(-) create mode 100644 server/convo_shutdown_test.go diff --git a/loop/loop.go b/loop/loop.go index 3b061dcef..0f77c8060 100644 --- a/loop/loop.go +++ b/loop/loop.go @@ -23,7 +23,14 @@ import ( // heartbeats/keepalives (which reset the idle timer) from hanging a turn // indefinitely. It is deliberately far larger than the idle window so that // genuinely long, steadily-streaming turns are unaffected. -const maxTurnDuration = 15 * time.Minute +const ( + maxTurnDuration = 15 * time.Minute + durableMessageTimeout = 5 * time.Second +) + +// ErrShutdown marks loop cancellation caused by graceful server shutdown. +// Final tool results and closeout errors use a detached context for this cause. +var ErrShutdown = errors.New("loop stopped by server shutdown") // MessageRecordFunc is called to record new messages to persistent storage. // otherUsage carries the usage of indirect LLM calls affiliated with the @@ -138,6 +145,21 @@ func NewLoop(config Config) *Loop { } } +// recordMessageDurably detaches a bounded write from the turn context while +// preserving its values. +func (l *Loop) recordMessageDurably(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { + durableCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), durableMessageTimeout) + defer cancel() + return l.recordMessage(durableCtx, message, usage, otherUsage) +} + +func (l *Loop) recordToolResult(ctx context.Context, message llm.Message, otherUsage []llm.PurposedUsage) error { + if ctx.Err() == nil || (errors.Is(ctx.Err(), context.Canceled) && !errors.Is(context.Cause(ctx), ErrShutdown)) { + return l.recordMessage(ctx, message, llm.Usage{}, otherUsage) + } + return l.recordMessageDurably(ctx, message, llm.Usage{}, otherUsage) +} + // Retry signals the loop to re-attempt the next LLM request without queueing // a new user message. The loop's in-memory history is unchanged (failed // requests don't append anything to history, and error messages are persisted @@ -267,6 +289,9 @@ func (l *Loop) Go(ctx context.Context) error { l.logger.Debug("processing queued messages", "count", 1) if err := l.processLLMRequest(ctx); err != nil { l.logger.Error("failed to process LLM request", "error", err) + if ctx.Err() != nil { + return ctx.Err() + } time.Sleep(time.Second) // Wait before retrying continue } @@ -458,17 +483,15 @@ func (l *Loop) processLLMRequest(ctx context.Context) error { // duplicate that — and, because ctx is already dead, the write // failed anyway ("failed to create message: Tx: context canceled"), // producing a scary log line on every user cancel. Skip it. - if errors.Is(ctx.Err(), context.Canceled) { + if errors.Is(ctx.Err(), context.Canceled) && !errors.Is(context.Cause(ctx), ErrShutdown) { l.logger.Info("LLM request aborted by loop cancellation", "error", err) return fmt.Errorf("LLM request failed: %w", err) } // Record the error as a message so it can be displayed in the UI. // EndOfTurn must be true so the agent working state is properly - // updated. WithoutCancel: this row's MarkAgentDone is what clears - // the persisted agent_working flag; losing the write to a context - // that expired (e.g. the 12-hour conversation-loop ceiling) or was - // cancelled between the failure above and this write would wedge - // the conversation in "Agent working..." forever. + // updated. recordMessageDurably detaches this write when the turn + // expired or graceful shutdown cancelled it; losing the row would + // wedge the conversation in "Agent working..." forever. errorMessage := llm.Message{ Role: llm.MessageRoleAssistant, Content: []llm.Content{ @@ -481,7 +504,7 @@ func (l *Loop) processLLMRequest(ctx context.Context) error { ErrorType: llm.ErrorTypeLLMRequest, ErrorRetryable: IsRetryableLLMError(err), } - if recordErr := l.recordMessage(context.WithoutCancel(ctx), errorMessage, llm.Usage{}, nil); recordErr != nil { + if recordErr := l.recordMessageDurably(ctx, errorMessage, llm.Usage{}, nil); recordErr != nil { l.logger.Error("failed to record error message", "error", recordErr) } return fmt.Errorf("LLM request failed: %w", err) @@ -870,7 +893,7 @@ func (l *Loop) executeToolCalls(ctx context.Context, content []llm.Content) erro l.mu.Unlock() // Record tool result message - if err := l.recordMessage(ctx, toolMessage, llm.Usage{}, otherUsage.Take()); err != nil { + if err := l.recordToolResult(ctx, toolMessage, otherUsage.Take()); err != nil { l.logger.Error("failed to record tool result message", "error", err) } } diff --git a/loop/loop_test.go b/loop/loop_test.go index b513676a0..671e1699e 100644 --- a/loop/loop_test.go +++ b/loop/loop_test.go @@ -3,6 +3,7 @@ package loop import ( "context" "encoding/json" + "errors" "fmt" "io" "net/http" @@ -1494,6 +1495,49 @@ func TestProcessLLMRequestError(t *testing.T) { } } +func TestProcessLLMRequestShutdownPersistsError(t *testing.T) { + ctx, cancel := context.WithCancelCause(context.WithValue(context.Background(), "test-key", "test-value")) + service := &cancelingLLMService{cancel: cancel} + + var recorded llm.Message + var recordCount int + var recordCtxErr error + var recordCtxValue any + var recordCtxHasDeadline bool + agentLoop := NewLoop(Config{ + LLM: service, + RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { + recordCount++ + recorded = message + recordCtxErr = ctx.Err() + recordCtxValue = ctx.Value("test-key") + _, recordCtxHasDeadline = ctx.Deadline() + return nil + }, + }) + agentLoop.QueueUserMessage(llm.UserStringMessage("cancel this request")) + + err := agentLoop.ProcessOneTurn(ctx) + if !errors.Is(err, context.Canceled) { + t.Fatalf("ProcessOneTurn error = %v, want context.Canceled", err) + } + if recordCount != 1 { + t.Fatalf("record count = %d, want 1", recordCount) + } + if recorded.Role != llm.MessageRoleAssistant || !recorded.EndOfTurn { + t.Fatalf("recorded message = %+v, want end-of-turn assistant error", recorded) + } + if recordCtxErr != nil { + t.Fatalf("record context error = %v, want nil", recordCtxErr) + } + if recordCtxValue != "test-value" { + t.Fatalf("record context value = %v, want test-value", recordCtxValue) + } + if !recordCtxHasDeadline { + t.Fatal("record context has no durability deadline") + } +} + // errorLLMService is a test LLM service that always returns an error type errorLLMService struct { err error @@ -1517,6 +1561,16 @@ func (e *errorLLMService) MaxImageBytes() int { return 5 * 1024 * 1024 } +type cancelingLLMService struct { + errorLLMService + cancel context.CancelCauseFunc +} + +func (s *cancelingLLMService) Do(ctx context.Context, req *llm.Request) (*llm.Response, error) { + s.cancel(ErrShutdown) + return nil, ctx.Err() +} + // retryableLLMService fails with a retryable error a specified number of times, then succeeds type retryableLLMService struct { failuresRemaining int @@ -1872,6 +1926,69 @@ func TestExecuteToolCallsWithMissingTool(t *testing.T) { } } +func TestExecuteToolCallsShutdownPersistsResult(t *testing.T) { + var recorded llm.Message + var recordCount int + agentLoop := NewLoop(Config{ + Tools: []*llm.Tool{}, + RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { + if err := ctx.Err(); err != nil { + return err + } + recordCount++ + recorded = message + return nil + }, + }) + + ctx, cancel := context.WithCancelCause(context.Background()) + cancel(ErrShutdown) + err := agentLoop.executeToolCalls(ctx, []llm.Content{{ + ID: "cancelled-tool", + Type: llm.ContentTypeToolUse, + ToolName: "missing-tool", + }}) + if err != nil { + t.Fatalf("executeToolCalls error = %v", err) + } + if recordCount != 1 { + t.Fatalf("record count = %d, want 1", recordCount) + } + if recorded.Role != llm.MessageRoleUser || len(recorded.Content) != 1 { + t.Fatalf("recorded message = %+v, want one user tool result", recorded) + } + if got := recorded.Content[0]; got.Type != llm.ContentTypeToolResult || got.ToolUseID != "cancelled-tool" { + t.Fatalf("recorded content = %+v, want cancelled-tool result", got) + } +} + +func TestExecuteToolCallsUserCancellationDoesNotPersistLateResult(t *testing.T) { + var persisted bool + agentLoop := NewLoop(Config{ + Tools: []*llm.Tool{}, + RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { + if err := ctx.Err(); err != nil { + return err + } + persisted = true + return nil + }, + }) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if err := agentLoop.executeToolCalls(ctx, []llm.Content{{ + ID: "cancelled-tool", + Type: llm.ContentTypeToolUse, + ToolName: "missing-tool", + }}); err != nil { + t.Fatalf("executeToolCalls error = %v", err) + } + if persisted { + t.Fatal("late tool result persisted after user cancellation") + } +} + func TestExecuteToolCallsWithErrorTool(t *testing.T) { var recordedMessages []llm.Message recordFunc := func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { diff --git a/server/convo.go b/server/convo.go index 5c2fe9c5f..b1bf78e35 100644 --- a/server/convo.go +++ b/server/convo.go @@ -87,7 +87,9 @@ type ConversationManager struct { conversationOptions db.ConversationOptions db *db.DB loop *loop.Loop - loopCancel context.CancelFunc + loopCancel context.CancelCauseFunc + loopTimeoutCancel context.CancelFunc + loopDone chan struct{} loopCtx context.Context mu sync.Mutex lastActivity time.Time @@ -1871,9 +1873,13 @@ func (cm *ConversationManager) ensureLoop(service llm.Service, modelID string) e // The list patch stream refreshes from the Pool commit hook. } - // Create a context with the conversation ID for LLM request recording/prefix dedup + // Create a context with the conversation ID for LLM request recording/prefix dedup. + // The cancel cause distinguishes graceful shutdown (whose final writes must + // survive cancellation) from user cancellation (whose synthetic closeout is + // recorded by CancelConversation). baseCtx := llmhttp.WithConversationID(context.Background(), conversationID) - processCtx, cancel := context.WithTimeout(baseCtx, 12*time.Hour) + cancelCtx, cancel := context.WithCancelCause(baseCtx) + processCtx, timeoutCancel := context.WithTimeout(cancelCtx, 12*time.Hour) toolSetConfig.ToolOverrides = conversationOpts.ToolOverrides toolSetConfig.DisableAllTools = conversationOpts.DisableAllTools @@ -1912,11 +1918,13 @@ func (cm *ConversationManager) ensureLoop(service llm.Service, modelID string) e return cm.takeInjectableSubagentDone(ctx) }, }) + loopDone := make(chan struct{}) cm.mu.Lock() if cm.loop != nil { cm.mu.Unlock() - cancel() + cancel(context.Canceled) + timeoutCancel() toolSet.Cleanup() existingModel := cm.modelID if existingModel != "" && modelID != "" && existingModel != modelID { @@ -1928,6 +1936,8 @@ func (cm *ConversationManager) ensureLoop(service llm.Service, modelID string) e needsPersist := cm.modelID == "" && modelID != "" cm.loop = loopInstance cm.loopCancel = cancel + cm.loopTimeoutCancel = timeoutCancel + cm.loopDone = loopDone cm.loopCtx = processCtx cm.modelID = modelID cm.toolSet = toolSet @@ -1953,6 +1963,7 @@ func (cm *ConversationManager) ensureLoop(service llm.Service, modelID string) e } go func() { + defer close(loopDone) if err := loopInstance.Go(processCtx); err != nil && err != context.DeadlineExceeded && err != context.Canceled { if logger != nil { logger.Error("Conversation loop stopped", "error", err) @@ -1966,19 +1977,37 @@ func (cm *ConversationManager) ensureLoop(service llm.Service, modelID string) e } func (cm *ConversationManager) stopLoop() { - cm.resetLoop(false) + cm.resetLoop(false, loop.ErrShutdown) +} + +func (cm *ConversationManager) stopLoopAndWait(ctx context.Context) { + loopDone := cm.resetLoop(false, loop.ErrShutdown) + if loopDone == nil { + return + } + select { + case <-loopDone: + case <-ctx.Done(): + } } // ResetLoop drops the in-memory LLM loop so the next turn hydrates from the DB. func (cm *ConversationManager) ResetLoop() { - cm.resetLoop(true) + cm.resetLoop(true, context.Canceled) } -func (cm *ConversationManager) resetLoop(markUnhydrated bool) { +func (cm *ConversationManager) resetLoop(markUnhydrated bool, cause error) <-chan struct{} { cm.mu.Lock() cancel := cm.loopCancel + timeoutCancel := cm.loopTimeoutCancel + loopDone := cm.loopDone toolSet := cm.toolSet + if cm.cancelling { + cause = context.Canceled + } cm.loopCancel = nil + cm.loopTimeoutCancel = nil + cm.loopDone = nil cm.loopCtx = nil cm.loop = nil cm.modelID = "" @@ -1990,11 +2019,15 @@ func (cm *ConversationManager) resetLoop(markUnhydrated bool) { cm.mu.Unlock() if cancel != nil { - cancel() + cancel(cause) + } + if timeoutCancel != nil { + timeoutCancel() } if toolSet != nil { toolSet.Cleanup() } + return loopDone } // CancelConversation cancels the current conversation loop and records a cancelled tool result if a tool was in progress @@ -2003,6 +2036,10 @@ func (cm *ConversationManager) CancelConversation(ctx context.Context) error { loopInstance := cm.loop loopCtx := cm.loopCtx cancel := cm.loopCancel + timeoutCancel := cm.loopTimeoutCancel + if loopInstance != nil { + cm.cancelling = true + } cm.mu.Unlock() if loopInstance == nil { @@ -2010,13 +2047,7 @@ func (cm *ConversationManager) CancelConversation(ctx context.Context) error { return nil } - // Mark the manager as cancelling so the synthetic "[Operation cancelled]" - // end-of-turn message recorded below does not fire onDone — a cancellation - // is not a subagent completion and must not notify the parent. Cleared - // once teardown finishes. - cm.mu.Lock() - cm.cancelling = true - cm.mu.Unlock() + // Suppress onDone while the synthetic cancellation closeout is recorded. defer func() { cm.mu.Lock() cm.cancelling = false @@ -2082,9 +2113,14 @@ func (cm *ConversationManager) CancelConversation(ctx context.Context) error { } } - // Cancel the context + // Cancel the context. Ordinary user cancellation deliberately does not use + // loop.ErrShutdown: CancelConversation owns the synthetic tool result and + // end-of-turn rows, so late loop writes must remain cancelled. if cancel != nil { - cancel() + cancel(context.Canceled) + } + if timeoutCancel != nil { + timeoutCancel() } // Wait briefly for the loop to stop @@ -2164,12 +2200,16 @@ func (cm *ConversationManager) CancelConversation(ctx context.Context) error { cm.SetAgentWorking(false) cm.mu.Lock() - cm.loopCancel = nil - cm.loopCtx = nil - cm.loop = nil - cm.modelID = "" - // Reset hydrated so that the next AcceptUserMessage will reload history from the database - cm.hydrated = false + if cm.loop == loopInstance { + cm.loopCancel = nil + cm.loopTimeoutCancel = nil + cm.loopDone = nil + cm.loopCtx = nil + cm.loop = nil + cm.modelID = "" + // Reset hydrated so that the next AcceptUserMessage reloads history. + cm.hydrated = false + } cm.mu.Unlock() return nil diff --git a/server/convo_shutdown_test.go b/server/convo_shutdown_test.go new file mode 100644 index 000000000..c0694a64b --- /dev/null +++ b/server/convo_shutdown_test.go @@ -0,0 +1,84 @@ +package server + +import ( + "context" + "errors" + "testing" + + "shelley.exe.dev/llm" + "shelley.exe.dev/loop" +) + +type shutdownBlockingService struct { + started chan struct{} +} + +func (s *shutdownBlockingService) Do(ctx context.Context, req *llm.Request) (*llm.Response, error) { + close(s.started) + <-ctx.Done() + return nil, ctx.Err() +} + +func (*shutdownBlockingService) Provider() string { return "test" } +func (*shutdownBlockingService) TokenContextWindow() int { return 200000 } +func (*shutdownBlockingService) MaxImageDimension() int { return 2000 } +func (*shutdownBlockingService) MaxImageBytes() int { return 5 * 1024 * 1024 } +func (*shutdownBlockingService) SupportsImages() bool { return false } + +func TestStopLoopWaitsForDurableCloseout(t *testing.T) { + service := &shutdownBlockingService{started: make(chan struct{})} + recordStarted := make(chan struct{}) + releaseRecord := make(chan struct{}) + recordCount := 0 + agentLoop := loop.NewLoop(loop.Config{ + LLM: service, + RecordMessage: func(ctx context.Context, message llm.Message, usage llm.Usage, otherUsage []llm.PurposedUsage) error { + if err := ctx.Err(); err != nil { + return err + } + recordCount++ + close(recordStarted) + <-releaseRecord + return nil + }, + }) + agentLoop.QueueUserMessage(llm.UserStringMessage("start")) + + cancelCtx, cancel := context.WithCancelCause(context.Background()) + processCtx, timeoutCancel := context.WithCancel(cancelCtx) + loopDone := make(chan struct{}) + go func() { + defer close(loopDone) + _ = agentLoop.Go(processCtx) + }() + <-service.started + + manager := &ConversationManager{ + loop: agentLoop, + loopCancel: cancel, + loopTimeoutCancel: timeoutCancel, + loopDone: loopDone, + loopCtx: processCtx, + } + stopped := make(chan struct{}) + go func() { + manager.stopLoopAndWait(context.Background()) + close(stopped) + }() + + <-recordStarted + if !errors.Is(context.Cause(processCtx), loop.ErrShutdown) { + t.Fatalf("cancellation cause = %v, want loop.ErrShutdown", context.Cause(processCtx)) + } + select { + case <-stopped: + t.Fatal("stopLoop returned before the closeout message was recorded") + default: + } + + close(releaseRecord) + <-stopped + if recordCount != 1 { + t.Fatalf("record count = %d, want 1", recordCount) + } +} diff --git a/server/server.go b/server/server.go index fe3d0a8be..78531d719 100644 --- a/server/server.go +++ b/server/server.go @@ -1721,10 +1721,8 @@ func (s *Server) IsAgentWorking(conversationID string) bool { // stopAllConversations stops every active conversation loop and cleans up // their tool sets. Used on graceful shutdown to ensure browser subprocesses // (headless-shell and its descendants) are killed. Returns when every loop -// has stopped or ctx is done. On timeout, lingering stopLoop goroutines keep -// running in the background; if the process exits before they finish their -// browser groups will be orphaned, but that's a strict improvement over the -// previous unbounded behavior. +// has stopped or ctx is done. On timeout, lingering loop goroutines may keep +// running until process exit, but shutdown itself remains bounded. func (s *Server) stopAllConversations(ctx context.Context) { s.mu.Lock() managers := make([]*ConversationManager, 0, len(s.activeConversations)) @@ -1744,7 +1742,7 @@ func (s *Server) stopAllConversations(ctx context.Context) { wg.Add(1) go func(m *ConversationManager) { defer wg.Done() - m.stopLoop() + m.stopLoopAndWait(ctx) }(manager) } wg.Wait()