opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-10-16 14:48:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(chatloop): simplify step persistence and interruption

Browse Source
This commit is contained in:
Kyle Carberry
2026-02-24 17:47:13 -05:00
parent cc8b7f90e1
commit 560b46987a
3 changed files with 116 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/chatd/chatd.go
+63 -73
View File
@@ -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))
},
coderd/chatd/chatloop/chatloop.go
+30 -51
View File
@@ -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
}
coderd/chatd/chatloop/chatloop_test.go
+23 -39
View File
@@ -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)