From 560b46987a92c5273a2603fdd3fb19c936ff43aa Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Tue, 24 Feb 2026 17:47:13 -0500 Subject: [PATCH] refactor(chatloop): simplify step persistence and interruption --- coderd/chatd/chatd.go | 136 ++++++++++++------------- coderd/chatd/chatloop/chatloop.go | 81 ++++++--------- coderd/chatd/chatloop/chatloop_test.go | 62 +++++------ 3 files changed, 116 insertions(+), 163 deletions(-) diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index de517de549..91cbdcdf91 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -85,10 +85,6 @@ const ( externalAuthWaitCompletedStatus = "completed" ) -var ( - ErrChatInterrupted = xerrors.New("chat interrupted") -) - // Processor handles background processing of pending chats. type Processor struct { cancel context.CancelFunc @@ -430,7 +426,7 @@ func (p *Processor) InterruptChat(chatID uuid.UUID) bool { if !ok { return false } - cancel(ErrChatInterrupted) + cancel(chatloop.ErrInterrupted) if p.streamManager != nil { p.streamManager.StopStream(chatID) } @@ -869,7 +865,7 @@ func (p *Processor) processChat(ctx context.Context, chat database.Chat) { }() if err := p.runChat(chatCtx, chat, logger, reportOnly, activeSubagentRequestID); err != nil { - if errors.Is(err, ErrChatInterrupted) || errors.Is(context.Cause(chatCtx), ErrChatInterrupted) { + if errors.Is(err, chatloop.ErrInterrupted) || errors.Is(context.Cause(chatCtx), chatloop.ErrInterrupted) { logger.Info(ctx, "chat interrupted") status = database.ChatStatusWaiting return @@ -1528,77 +1524,73 @@ func (p *Processor) runChatWithAgent( modelConfigContextLimit = compressionConfig.ContextLimit } - persistAssistant := func( - persistCtx context.Context, - content []fantasy.Content, - usage fantasy.Usage, - contextLimit sql.NullInt64, - ) error { - if len(content) == 0 { - return nil + persistStep := func(persistCtx context.Context, step chatloop.PersistedStep) error { + if len(step.AssistantContent) > 0 { + assistantContent, err := chatprompt.MarshalContent(step.AssistantContent) + if err != nil { + return err + } + + hasUsage := step.Usage != (fantasy.Usage{}) + assistantMessage, err := p.db.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ + ChatID: chat.ID, + Role: string(fantasy.MessageRoleAssistant), + Content: assistantContent, + ToolCallID: sql.NullString{ + Valid: false, + }, + Thinking: sql.NullString{ + Valid: false, + }, + Hidden: false, + SubagentRequestID: subagentRequest, + InputTokens: usageNullInt64(step.Usage.InputTokens, hasUsage), + OutputTokens: usageNullInt64(step.Usage.OutputTokens, hasUsage), + TotalTokens: usageNullInt64(step.Usage.TotalTokens, hasUsage), + ReasoningTokens: usageNullInt64(step.Usage.ReasoningTokens, hasUsage), + CacheCreationTokens: usageNullInt64( + step.Usage.CacheCreationTokens, + hasUsage, + ), + CacheReadTokens: usageNullInt64(step.Usage.CacheReadTokens, hasUsage), + ContextLimit: step.ContextLimit, + }) + if err != nil { + return xerrors.Errorf("insert assistant message: %w", err) + } + p.publishMessage(chat.ID, assistantMessage) } - assistantContent, err := chatprompt.MarshalContent(content) - if err != nil { - return err - } + for _, result := range step.ToolResults { + resultContent, err := chatprompt.MarshalResults([]chatprompt.ToolResultBlock{result}) + if err != nil { + return err + } - hasUsage := usage != (fantasy.Usage{}) - assistantMessage, err := p.db.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - Role: string(fantasy.MessageRoleAssistant), - Content: assistantContent, - ToolCallID: sql.NullString{ - Valid: false, - }, - Thinking: sql.NullString{ - Valid: false, - }, - Hidden: false, - SubagentRequestID: subagentRequest, - InputTokens: usageNullInt64(usage.InputTokens, hasUsage), - OutputTokens: usageNullInt64(usage.OutputTokens, hasUsage), - TotalTokens: usageNullInt64(usage.TotalTokens, hasUsage), - ReasoningTokens: usageNullInt64(usage.ReasoningTokens, hasUsage), - CacheCreationTokens: usageNullInt64( - usage.CacheCreationTokens, - hasUsage, - ), - CacheReadTokens: usageNullInt64(usage.CacheReadTokens, hasUsage), - ContextLimit: contextLimit, - }) - if err != nil { - return xerrors.Errorf("insert assistant message: %w", err) - } - p.publishMessage(chat.ID, assistantMessage) - return nil - } + toolMessage, err := p.db.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ + ChatID: chat.ID, + Role: string(fantasy.MessageRoleTool), + Content: resultContent, + ToolCallID: sql.NullString{ + String: result.ToolCallID, + Valid: result.ToolCallID != "", + }, + Hidden: false, + SubagentRequestID: subagentRequest, + }) + if err != nil { + return xerrors.Errorf("insert tool result: %w", err) + } - persistToolResult := func(persistCtx context.Context, result chatprompt.ToolResultBlock) error { - resultContent, err := chatprompt.MarshalResults([]chatprompt.ToolResultBlock{result}) - if err != nil { - return err + p.publishMessage(chat.ID, toolMessage) } - - toolMessage, err := p.db.InsertChatMessage(persistCtx, database.InsertChatMessageParams{ - ChatID: chat.ID, - Role: string(fantasy.MessageRoleTool), - Content: resultContent, - ToolCallID: sql.NullString{ - String: result.ToolCallID, - Valid: result.ToolCallID != "", - }, - Hidden: false, - SubagentRequestID: subagentRequest, - }) - if err != nil { - return xerrors.Errorf("insert tool result: %w", err) - } - - p.publishMessage(chat.ID, toolMessage) return nil } streamCall := streamCallOptionsFromChatModelConfig(model, modelConfig) + var activeTools []string + if reportOnly { + activeTools = append(activeTools, "subagent_report") + } return chatloop.Run(ctx, chatloop.RunOptions{ Model: model, @@ -1607,11 +1599,10 @@ func (p *Processor) runChatWithAgent( StreamCall: streamCall, MaxSteps: maxChatSteps, - ReportOnly: reportOnly, + ActiveTools: activeTools, ContextLimitFallback: modelConfigContextLimit, - PersistAssistant: persistAssistant, - PersistToolResult: persistToolResult, + PersistStep: persistStep, PublishMessagePart: func( role fantasy.MessageRole, part codersdk.ChatMessagePart, @@ -1619,7 +1610,6 @@ func (p *Processor) runChatWithAgent( p.publishMessagePart(chat.ID, string(role), part) }, - InterruptedError: ErrChatInterrupted, OnInterruptedPersistError: func(err error) { p.logger.Warn(ctx, "failed to persist interrupted chat step", slog.Error(err)) }, diff --git a/coderd/chatd/chatloop/chatloop.go b/coderd/chatd/chatloop/chatloop.go index 646b525112..090aeb63e5 100644 --- a/coderd/chatd/chatloop/chatloop.go +++ b/coderd/chatd/chatloop/chatloop.go @@ -20,9 +20,17 @@ import ( const ( interruptedToolResultErrorMessage = "tool call was interrupted before it produced a result" - reportOnlyToolName = "subagent_report" ) +var ErrInterrupted = xerrors.New("chat interrupted") + +type PersistedStep struct { + AssistantContent []fantasy.Content + ToolResults []chatprompt.ToolResultBlock + Usage fantasy.Usage + ContextLimit sql.NullInt64 +} + // RunOptions configures a single streaming chat loop run. type RunOptions struct { Model fantasy.LanguageModel @@ -31,22 +39,15 @@ type RunOptions struct { StreamCall fantasy.AgentStreamCall MaxSteps int - ReportOnly bool + ActiveTools []string ContextLimitFallback int64 - PersistAssistant func( - persistCtx context.Context, - content []fantasy.Content, - usage fantasy.Usage, - contextLimit sql.NullInt64, - ) error - PersistToolResult func(context.Context, chatprompt.ToolResultBlock) error + PersistStep func(context.Context, PersistedStep) error PublishMessagePart func( role fantasy.MessageRole, part codersdk.ChatMessagePart, ) - InterruptedError error OnInterruptedPersistError func(error) } @@ -55,11 +56,8 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { if opts.Model == nil { return nil, xerrors.New("chat model is required") } - if opts.PersistAssistant == nil { - return nil, xerrors.New("persist assistant callback is required") - } - if opts.PersistToolResult == nil { - return nil, xerrors.New("persist tool result callback is required") + if opts.PersistStep == nil { + return nil, xerrors.New("persist step callback is required") } if opts.MaxSteps <= 0 { opts.MaxSteps = 1 @@ -290,20 +288,10 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { toolResults = append(toolResults, missingResults...) persistCtx := context.WithoutCancel(ctx) - if err := opts.PersistAssistant( - persistCtx, - stepAssistantContent(draft, toolResults), - fantasy.Usage{}, - sql.NullInt64{}, - ); err != nil { - return err - } - for _, toolResult := range toolResults { - if err := opts.PersistToolResult(persistCtx, toolResult); err != nil { - return err - } - } - return nil + return opts.PersistStep(persistCtx, PersistedStep{ + AssistantContent: stepAssistantContent(draft, toolResults), + ToolResults: toolResults, + }) } resetStepState() @@ -326,7 +314,7 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { return stepCtx, prepareStepResult( options.Messages, sentinelPrompt, - opts.ReportOnly, + opts.ActiveTools, applyAnthropicCaching, ), nil } @@ -448,33 +436,24 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { } } - if err := opts.PersistAssistant( - ctx, - stepAssistantContent(stepResult.Content, toolResults), - stepResult.Usage, - contextLimit, - ); err != nil { - return err - } - for _, toolResult := range toolResults { - if err := opts.PersistToolResult(ctx, toolResult); err != nil { - return err - } - } - return nil + return opts.PersistStep(ctx, PersistedStep{ + AssistantContent: stepAssistantContent(stepResult.Content, toolResults), + ToolResults: toolResults, + Usage: stepResult.Usage, + ContextLimit: contextLimit, + }) } result, err := agent.Stream(ctx, streamCall) if err != nil { - if opts.InterruptedError != nil && - errors.Is(err, context.Canceled) && - errors.Is(context.Cause(ctx), opts.InterruptedError) { + if errors.Is(err, context.Canceled) && + errors.Is(context.Cause(ctx), ErrInterrupted) { if persistErr := persistInterruptedStep(); persistErr != nil { if opts.OnInterruptedPersistError != nil { opts.OnInterruptedPersistError(persistErr) } } - return nil, opts.InterruptedError + return nil, ErrInterrupted } return nil, xerrors.Errorf("stream response: %w", err) } @@ -485,7 +464,7 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { func prepareStepResult( messages []fantasy.Message, sentinel string, - reportOnly bool, + activeTools []string, anthropicCaching bool, ) fantasy.PrepareStepResult { filtered := make([]fantasy.Message, 0, len(messages)) @@ -509,8 +488,8 @@ func prepareStepResult( if anthropicCaching { result.Messages = addAnthropicPromptCaching(result.Messages) } - if reportOnly { - result.ActiveTools = []string{reportOnlyToolName} + if len(activeTools) > 0 { + result.ActiveTools = append([]string(nil), activeTools...) } return result } diff --git a/coderd/chatd/chatloop/chatloop_test.go b/coderd/chatd/chatloop/chatloop_test.go index 446e30439d..3e055efde1 100644 --- a/coderd/chatd/chatloop/chatloop_test.go +++ b/coderd/chatd/chatloop/chatloop_test.go @@ -2,7 +2,6 @@ package chatloop import ( "context" - "database/sql" "iter" "strings" "testing" @@ -15,7 +14,9 @@ import ( "github.com/coder/coder/v2/coderd/chatd/chatprompt" ) -func TestRun_ReportOnlyPrepareBehavior(t *testing.T) { +const subagentReportToolName = "subagent_report" + +func TestRun_ActiveToolsPrepareBehavior(t *testing.T) { t.Parallel() var capturedCall fantasy.Call @@ -32,9 +33,8 @@ func TestRun_ReportOnlyPrepareBehavior(t *testing.T) { }, } - persistAssistantCalls := 0 - persistToolResultCalls := 0 - persistedContextLimit := sql.NullInt64{} + persistStepCalls := 0 + var persistedStep PersistedStep _, err := Run(context.Background(), RunOptions{ Model: model, @@ -46,38 +46,29 @@ func TestRun_ReportOnlyPrepareBehavior(t *testing.T) { textMessage(fantasy.MessageRoleUser, "continue"), }, Tools: []fantasy.AgentTool{ - newNoopTool(reportOnlyToolName), + newNoopTool(subagentReportToolName), newNoopTool("read_file"), }, MaxSteps: 3, - ReportOnly: true, + ActiveTools: []string{subagentReportToolName}, ContextLimitFallback: 4096, - PersistAssistant: func( - _ context.Context, - _ []fantasy.Content, - _ fantasy.Usage, - contextLimit sql.NullInt64, - ) error { - persistAssistantCalls++ - persistedContextLimit = contextLimit - return nil - }, - PersistToolResult: func(_ context.Context, _ chatprompt.ToolResultBlock) error { - persistToolResultCalls++ + PersistStep: func(_ context.Context, step PersistedStep) error { + persistStepCalls++ + persistedStep = step return nil }, }) require.NoError(t, err) - require.Equal(t, 1, persistAssistantCalls) - require.Equal(t, 0, persistToolResultCalls) - require.True(t, persistedContextLimit.Valid) - require.Equal(t, int64(4096), persistedContextLimit.Int64) + require.Equal(t, 1, persistStepCalls) + require.Empty(t, persistedStep.ToolResults) + require.True(t, persistedStep.ContextLimit.Valid) + require.Equal(t, int64(4096), persistedStep.ContextLimit.Int64) require.NotEmpty(t, capturedCall.Prompt) require.False(t, containsPromptSentinel(capturedCall.Prompt)) require.Len(t, capturedCall.Tools, 1) - require.Equal(t, reportOnlyToolName, capturedCall.Tools[0].GetName()) + require.Equal(t, subagentReportToolName, capturedCall.Tools[0].GetName()) require.Len(t, capturedCall.Prompt, 5) require.False(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[0])) @@ -131,13 +122,12 @@ func TestRun_InterruptedStepPersistsSyntheticToolResult(t *testing.T) { }, } - interruptedErr := xerrors.New("chat interrupted") ctx, cancel := context.WithCancelCause(context.Background()) defer cancel(nil) go func() { <-started - cancel(interruptedErr) + cancel(ErrInterrupted) }() persistedAssistantCtxErr := xerrors.New("unset") @@ -153,23 +143,17 @@ func TestRun_InterruptedStepPersistsSyntheticToolResult(t *testing.T) { newNoopTool("read_file"), }, MaxSteps: 3, - PersistAssistant: func( - persistCtx context.Context, - content []fantasy.Content, - _ fantasy.Usage, - _ sql.NullInt64, - ) error { + PersistStep: func(persistCtx context.Context, step PersistedStep) error { persistedAssistantCtxErr = persistCtx.Err() - persistedAssistantContent = append([]fantasy.Content(nil), content...) + persistedAssistantContent = append([]fantasy.Content(nil), step.AssistantContent...) + persistedToolResults = append( + persistedToolResults, + step.ToolResults..., + ) return nil }, - PersistToolResult: func(_ context.Context, result chatprompt.ToolResultBlock) error { - persistedToolResults = append(persistedToolResults, result) - return nil - }, - InterruptedError: interruptedErr, }) - require.ErrorIs(t, err, interruptedErr) + require.ErrorIs(t, err, ErrInterrupted) require.NoError(t, persistedAssistantCtxErr) require.NotEmpty(t, persistedAssistantContent)