From ddfe630757a0561b9cdc70ebf63f42b39d685579 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Mon, 2 Mar 2026 18:51:57 -0500 Subject: [PATCH] refactor(chatd): replace fantasy.Agent with custom agent loop (#22507) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Replaces fantasy's `Agent` abstraction with a direct step loop calling `LanguageModel.Stream()`. Fantasy is retained as the provider abstraction layer (streaming parsers, types, tool schema) but we no longer use `fantasy.Agent`, `AgentStreamCall`, `AgentResult`, or `StepResult`. ## Problems solved | Problem | Before | After | |---|---|---| | **Sentinel prompt hack** | fantasy.Agent requires non-empty Prompt → UUID sentinel generated and stripped in PrepareStep | Messages passed directly to `model.Stream()` | | **Discarded PersistStep errors** | `_ = opts.OnStepFinish(result)` silently swallows errors | Errors propagate directly from `PersistStep()` | | **Shadow draft state** | ~160 LOC tracking content in parallel because fantasy doesn't expose in-progress content on interruption | `stepResult` owns content directly; `flushActiveState()` is trivial | | **Nested retry layers** | fantasy's 2-attempt retry nested inside chatretry's indefinite retry | Single `chatretry.Retry` layer | | **Callback-mediated compaction** | Mutex + boolean flag + coordination between OnStepFinish/PrepareStep callbacks | Inline `if` statement between steps | | **Duplicate compaction paths** | `compactStep()` + `maybeCompact()` sharing ~80% logic | Single `tryCompact()` function | ## Changes ### `coderd/chatd/chatloop/chatloop.go` — Rewritten - **Removed**: `fantasy.NewAgent()`, `AgentStreamCall`, sentinel prompt, shadow draft state (~160 LOC of closures), `compactedMu`/`compacted` flag, `PrepareStepResult` - **Added**: `stepResult` struct, `processStepStream()` (stream consumer), `executeTools()` (sequential tool execution), `flushActiveState()` (interrupt handling), `buildToolDefinitions()`, `toResponseMessages()` - **Changed**: `Run()` return type from `(*fantasy.AgentResult, error)` to `error` (callers already discarded the result) - **Preserved**: Anthropic prompt caching, reasoning title extraction, `extractContextLimit()`, `ErrInterrupted` semantics ### `coderd/chatd/chatloop/compaction.go` — Simplified - Merged `compactStep()` + `maybeCompact()` → single `tryCompact()` - Removed `[]StepResult` parameter from `generateCompactionSummary()` (caller provides complete message list) - Kept helper functions: `normalizedCompactionConfig`, `contextTokensFromUsage`, `resolveContextLimit`, `shouldCompact` ### `coderd/chatd/chatd.go` — Caller updates - Removed `AgentStreamCall` construction - Changed `_, err = chatloop.Run(...)` to `err = chatloop.Run(...)` - Model parameters moved from `AgentStreamCall` fields to `RunOptions` fields ### Tests — 4 new tests - `MidLoopCompactionReloadsMessages` — compaction fires mid-loop, messages reloaded - `PostRunCompactionSkippedAfterMidLoop` — no double compaction - `MultiStepToolExecution` — tools execute between steps, results feed next step - `PersistStepErrorPropagates` — persistence errors propagate (was silently discarded) --- coderd/chatd/chatd.go | 64 +- coderd/chatd/chatloop/chatloop.go | 1107 +++++++++++++--------- coderd/chatd/chatloop/chatloop_test.go | 135 ++- coderd/chatd/chatloop/compaction.go | 322 ++++--- coderd/chatd/chatloop/compaction_test.go | 258 ++++- 5 files changed, 1286 insertions(+), 600 deletions(-) diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 1e18ab8322..94d4a9902d 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -2036,28 +2036,19 @@ func (p *Server) runChat( return nil } - streamCall := fantasy.AgentStreamCall{ - MaxOutputTokens: callConfig.MaxOutputTokens, - Temperature: callConfig.Temperature, - TopP: callConfig.TopP, - TopK: callConfig.TopK, - PresencePenalty: callConfig.PresencePenalty, - FrequencyPenalty: callConfig.FrequencyPenalty, - ProviderOptions: chatprovider.ProviderOptionsFromChatModelConfig(model, callConfig.ProviderOptions), - } - - if streamCall.MaxOutputTokens == nil { + // Apply the default MaxOutputTokens if the model config + // does not specify one. + if callConfig.MaxOutputTokens == nil { maxOutputTokens := int64(32_000) - streamCall.MaxOutputTokens = &maxOutputTokens + callConfig.MaxOutputTokens = &maxOutputTokens } - // Generate the tool call ID up front so that the OnStart - // streaming part and the Persist durable messages share - // the same identifier. Without this the client cannot - // correlate the "Summarizing..." tool call with the - // "Summarized" tool result. + // Generate the tool call ID up front so that the streaming + // parts and durable messages share the same identifier. + // Without this the client cannot correlate the + // "Summarizing..." tool call with the "Summarized" tool + // result. compactionToolCallID := "chat_summarized_" + uuid.NewString() - compactionOptions := &chatloop.CompactionOptions{ ThresholdPercent: modelConfig.CompressionThreshold, ContextLimit: modelConfig.ContextLimit, @@ -2083,15 +2074,10 @@ func (p *Server) runChat( ) return nil }, - OnStart: func() { - // Publish a streaming tool-call part immediately so - // connected clients see "Summarizing..." while the - // LLM generates the summary. - p.publishMessagePart(chat.ID, string(fantasy.MessageRoleAssistant), codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: compactionToolCallID, - ToolName: "chat_summarized", - }) + ToolCallID: compactionToolCallID, + ToolName: "chat_summarized", + PublishMessagePart: func(role fantasy.MessageRole, part codersdk.ChatMessagePart) { + p.publishMessagePart(chat.ID, string(role), part) }, OnError: func(err error) { logger.Warn(ctx, "failed to compact chat context", slog.Error(err)) @@ -2147,12 +2133,14 @@ func (p *Server) runChat( })...) } - _, err = chatloop.Run(ctx, chatloop.RunOptions{ - Model: model, - Messages: prompt, - Tools: tools, - StreamCall: streamCall, - MaxSteps: maxChatSteps, + err = chatloop.Run(ctx, chatloop.RunOptions{ + Model: model, + Messages: prompt, + Tools: tools, + MaxSteps: maxChatSteps, + + ModelConfig: callConfig, + ProviderOptions: chatprovider.ProviderOptionsFromChatModelConfig(model, callConfig.ProviderOptions), ContextLimitFallback: modelConfigContextLimit, @@ -2315,16 +2303,6 @@ func (p *Server) persistChatContextSummary( return xerrors.Errorf("insert summary tool result message: %w", err) } - // Publish a streaming tool-result part so connected clients - // transition from "Summarizing..." to "Summarized" before the - // durable messages and status change arrive. - p.publishMessagePart(chatID, string(fantasy.MessageRoleTool), codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: toolCallID, - ToolName: "chat_summarized", - Result: summaryResult, - }) - p.publishMessage(chatID, assistantMessage) p.publishMessage(chatID, toolMessage) return nil diff --git a/coderd/chatd/chatloop/chatloop.go b/coderd/chatd/chatloop/chatloop.go index 53f340eea7..6a6d8633c2 100644 --- a/coderd/chatd/chatloop/chatloop.go +++ b/coderd/chatd/chatloop/chatloop.go @@ -5,14 +5,14 @@ import ( "database/sql" "encoding/json" "errors" + "slices" "strconv" "strings" - "sync" "time" "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" - "github.com/google/uuid" + "charm.land/fantasy/schema" "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/chatd/chatprompt" @@ -28,10 +28,9 @@ var ErrInterrupted = xerrors.New("chat interrupted") // PersistedStep contains the full content of a completed or // interrupted agent step. Content includes both assistant blocks -// (text, reasoning, tool calls) and tool result blocks, mirroring -// what fantasy provides in StepResult.Content. The persistence -// layer is responsible for splitting these into separate database -// messages by role. +// (text, reasoning, tool calls) and tool result blocks. The +// persistence layer is responsible for splitting these into +// separate database messages by role. type PersistedStep struct { Content []fantasy.Content Usage fantasy.Usage @@ -40,21 +39,30 @@ type PersistedStep struct { // RunOptions configures a single streaming chat loop run. type RunOptions struct { - Model fantasy.LanguageModel - Messages []fantasy.Message - Tools []fantasy.AgentTool - StreamCall fantasy.AgentStreamCall - MaxSteps int + Model fantasy.LanguageModel + Messages []fantasy.Message + Tools []fantasy.AgentTool + MaxSteps int ActiveTools []string ContextLimitFallback int64 + // ModelConfig holds per-call LLM parameters (temperature, + // max tokens, etc.) read from the chat model configuration. + ModelConfig codersdk.ChatModelCallConfig + // ProviderOptions are provider-specific call options + // converted from ModelConfig.ProviderOptions. This is a + // separate field because the conversion requires knowledge + // of the provider, which lives in chatd, not chatloop. + ProviderOptions fantasy.ProviderOptions + PersistStep func(context.Context, PersistedStep) error PublishMessagePart func( role fantasy.MessageRole, part codersdk.ChatMessagePart, ) - Compaction *CompactionOptions + Compaction *CompactionOptions + ReloadMessages func(context.Context) ([]fantasy.Message, error) // OnRetry is called before each retry attempt when the LLM // stream fails with a retryable error. It provides the attempt @@ -65,13 +73,117 @@ type RunOptions struct { OnInterruptedPersistError func(error) } -// Run executes the chat step-stream loop and delegates persistence/publishing to callbacks. -func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { +// stepResult holds the accumulated output of a single streaming +// step. Since we own the stream consumer, all content is tracked +// directly here — no shadow draft state needed. +type stepResult struct { + content []fantasy.Content + usage fantasy.Usage + providerMetadata fantasy.ProviderMetadata + finishReason fantasy.FinishReason + toolCalls []fantasy.ToolCallContent + shouldContinue bool +} + +// toResponseMessages converts step content into messages suitable +// for appending to the conversation. Mirrors fantasy's +// toResponseMessages logic. +func (r stepResult) toResponseMessages() []fantasy.Message { + var assistantParts []fantasy.MessagePart + var toolParts []fantasy.MessagePart + + for _, c := range r.content { + switch c.GetType() { + case fantasy.ContentTypeText: + text, ok := fantasy.AsContentType[fantasy.TextContent](c) + if !ok { + continue + } + assistantParts = append(assistantParts, fantasy.TextPart{ + Text: text.Text, + ProviderOptions: fantasy.ProviderOptions(text.ProviderMetadata), + }) + case fantasy.ContentTypeReasoning: + reasoning, ok := fantasy.AsContentType[fantasy.ReasoningContent](c) + if !ok { + continue + } + assistantParts = append(assistantParts, fantasy.ReasoningPart{ + Text: reasoning.Text, + ProviderOptions: fantasy.ProviderOptions(reasoning.ProviderMetadata), + }) + case fantasy.ContentTypeToolCall: + toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](c) + if !ok { + continue + } + assistantParts = append(assistantParts, fantasy.ToolCallPart{ + ToolCallID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + Input: toolCall.Input, + ProviderExecuted: toolCall.ProviderExecuted, + ProviderOptions: fantasy.ProviderOptions(toolCall.ProviderMetadata), + }) + case fantasy.ContentTypeFile: + file, ok := fantasy.AsContentType[fantasy.FileContent](c) + if !ok { + continue + } + assistantParts = append(assistantParts, fantasy.FilePart{ + Data: file.Data, + MediaType: file.MediaType, + ProviderOptions: fantasy.ProviderOptions(file.ProviderMetadata), + }) + case fantasy.ContentTypeSource: + // Sources are metadata about references; they don't + // need to be included in conversation messages. + continue + case fantasy.ContentTypeToolResult: + result, ok := fantasy.AsContentType[fantasy.ToolResultContent](c) + if !ok { + continue + } + toolParts = append(toolParts, fantasy.ToolResultPart{ + ToolCallID: result.ToolCallID, + Output: result.Result, + ProviderOptions: fantasy.ProviderOptions(result.ProviderMetadata), + }) + default: + continue + } + } + + var messages []fantasy.Message + if len(assistantParts) > 0 { + messages = append(messages, fantasy.Message{ + Role: fantasy.MessageRoleAssistant, + Content: assistantParts, + }) + } + if len(toolParts) > 0 { + messages = append(messages, fantasy.Message{ + Role: fantasy.MessageRoleTool, + Content: toolParts, + }) + } + return messages +} + +// reasoningState accumulates reasoning content and provider +// metadata while the stream is in flight. +type reasoningState struct { + text string + options fantasy.ProviderMetadata +} + +// Run executes the chat step-stream loop and delegates +// persistence/publishing to callbacks. +func Run(ctx context.Context, opts RunOptions) error { if opts.Model == nil { - return nil, xerrors.New("chat model is required") + return xerrors.New("chat model is required") } if opts.PersistStep == nil { - return nil, xerrors.New("persist step callback is required") + return xerrors.New("persist step callback is required") } if opts.MaxSteps <= 0 { opts.MaxSteps = 1 @@ -84,359 +196,94 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { opts.PublishMessagePart(role, part) } - var ( - stepStateMu sync.Mutex - streamToolNames map[string]string - streamReasoningTitles map[string]string - streamReasoningText map[string]string - // stepToolResultContents tracks tool results received during - // streaming. These are needed for the interrupted-step path - // where OnStepFinish never fires. - stepToolResultContents []fantasy.ToolResultContent - stepAssistantDraft []fantasy.Content - stepToolCallIndexByID map[string]int - ) - - resetStepState := func() { - stepStateMu.Lock() - streamToolNames = make(map[string]string) - streamReasoningTitles = make(map[string]string) - streamReasoningText = make(map[string]string) - stepToolResultContents = nil - stepAssistantDraft = nil - stepToolCallIndexByID = make(map[string]int) - stepStateMu.Unlock() - } - - setReasoningTitleFromText := func(id string, text string) { - if id == "" || strings.TrimSpace(text) == "" { - return - } - - stepStateMu.Lock() - defer stepStateMu.Unlock() - - if streamReasoningTitles[id] != "" { - return - } - - streamReasoningText[id] += text - if !strings.ContainsAny(streamReasoningText[id], "\r\n") { - return - } - title := chatprompt.ReasoningTitleFromFirstLine(streamReasoningText[id]) - if title == "" { - return - } - - streamReasoningTitles[id] = title - } - - appendDraftText := func(text string) { - if text == "" { - return - } - - stepStateMu.Lock() - defer stepStateMu.Unlock() - - if len(stepAssistantDraft) > 0 { - lastIndex := len(stepAssistantDraft) - 1 - switch last := stepAssistantDraft[lastIndex].(type) { - case fantasy.TextContent: - last.Text += text - stepAssistantDraft[lastIndex] = last - return - case *fantasy.TextContent: - last.Text += text - stepAssistantDraft[lastIndex] = fantasy.TextContent{Text: last.Text} - return - } - } - stepAssistantDraft = append(stepAssistantDraft, fantasy.TextContent{Text: text}) - } - - appendDraftReasoning := func(text string) { - if text == "" { - return - } - - stepStateMu.Lock() - defer stepStateMu.Unlock() - - if len(stepAssistantDraft) > 0 { - lastIndex := len(stepAssistantDraft) - 1 - switch last := stepAssistantDraft[lastIndex].(type) { - case fantasy.ReasoningContent: - last.Text += text - stepAssistantDraft[lastIndex] = last - return - case *fantasy.ReasoningContent: - last.Text += text - stepAssistantDraft[lastIndex] = fantasy.ReasoningContent{Text: last.Text} - return - } - } - stepAssistantDraft = append(stepAssistantDraft, fantasy.ReasoningContent{Text: text}) - } - - upsertDraftToolCall := func(toolCallID, toolName, input string, appendInput bool) { - if toolCallID == "" { - return - } - - stepStateMu.Lock() - defer stepStateMu.Unlock() - - if strings.TrimSpace(toolName) != "" { - streamToolNames[toolCallID] = toolName - } - - index, exists := stepToolCallIndexByID[toolCallID] - if !exists { - stepToolCallIndexByID[toolCallID] = len(stepAssistantDraft) - stepAssistantDraft = append(stepAssistantDraft, fantasy.ToolCallContent{ - ToolCallID: toolCallID, - ToolName: toolName, - Input: input, - }) - return - } - - if index < 0 || index >= len(stepAssistantDraft) { - stepToolCallIndexByID[toolCallID] = len(stepAssistantDraft) - stepAssistantDraft = append(stepAssistantDraft, fantasy.ToolCallContent{ - ToolCallID: toolCallID, - ToolName: toolName, - Input: input, - }) - return - } - - existingCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](stepAssistantDraft[index]) - if !ok { - if ptrCall, ptrOK := fantasy.AsContentType[*fantasy.ToolCallContent](stepAssistantDraft[index]); ptrOK && ptrCall != nil { - existingCall = *ptrCall - ok = true - } - } - if !ok { - stepToolCallIndexByID[toolCallID] = len(stepAssistantDraft) - stepAssistantDraft = append(stepAssistantDraft, fantasy.ToolCallContent{ - ToolCallID: toolCallID, - ToolName: toolName, - Input: input, - }) - return - } - - if strings.TrimSpace(toolName) != "" { - existingCall.ToolName = toolName - } - if appendInput { - existingCall.Input += input - } else if input != "" || existingCall.Input == "" { - existingCall.Input = input - } - stepAssistantDraft[index] = existingCall - } - - appendDraftSource := func(source fantasy.SourceContent) { - stepStateMu.Lock() - stepAssistantDraft = append(stepAssistantDraft, source) - stepStateMu.Unlock() - } - - persistInterruptedStep := func() error { - stepStateMu.Lock() - draft := append([]fantasy.Content(nil), stepAssistantDraft...) - toolResults := append([]fantasy.ToolResultContent(nil), stepToolResultContents...) - toolNameByCallID := make(map[string]string, len(streamToolNames)) - for id, name := range streamToolNames { - toolNameByCallID[id] = name - } - stepStateMu.Unlock() - - if len(draft) == 0 && len(toolResults) == 0 { - return nil - } - - // Track which tool calls already have results. - answeredToolCalls := make(map[string]struct{}, len(toolResults)) - for _, tr := range toolResults { - if tr.ToolCallID != "" { - answeredToolCalls[tr.ToolCallID] = struct{}{} - } - } - - // Build the combined content: draft + received tool results - // + synthetic interrupted results for unanswered tool calls. - content := make([]fantasy.Content, 0, len(draft)+len(toolResults)) - content = append(content, draft...) - for _, tr := range toolResults { - content = append(content, tr) - } - - for _, block := range draft { - toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block) - if !ok { - if ptrCall, ptrOK := fantasy.AsContentType[*fantasy.ToolCallContent](block); ptrOK && ptrCall != nil { - toolCall = *ptrCall - ok = true - } - } - if !ok || toolCall.ToolCallID == "" { - continue - } - if _, exists := answeredToolCalls[toolCall.ToolCallID]; exists { - continue - } - - toolName := strings.TrimSpace(toolCall.ToolName) - if toolName == "" { - toolName = strings.TrimSpace(toolNameByCallID[toolCall.ToolCallID]) - } - - content = append(content, fantasy.ToolResultContent{ - ToolCallID: toolCall.ToolCallID, - ToolName: toolName, - Result: fantasy.ToolResultOutputContentError{ - Error: xerrors.New(interruptedToolResultErrorMessage), - }, - }) - answeredToolCalls[toolCall.ToolCallID] = struct{}{} - } - - persistCtx := context.WithoutCancel(ctx) - return opts.PersistStep(persistCtx, PersistedStep{ - Content: content, - }) - } - - resetStepState() - - agent := fantasy.NewAgent( - opts.Model, - fantasy.WithTools(opts.Tools...), - fantasy.WithStopConditions(fantasy.StepCountIs(opts.MaxSteps)), - ) + tools := buildToolDefinitions(opts.Tools, opts.ActiveTools) applyAnthropicCaching := shouldApplyAnthropicPromptCaching(opts.Model) - // Fantasy's AgentStreamCall currently requires a non-empty Prompt and always - // appends it as a user message. chatd already supplies the full history in - // Messages, so we pass and then strip a sentinel user message in PrepareStep. - sentinelPrompt := "__chatd_agent_prompt_sentinel_" + uuid.NewString() - streamCall := opts.StreamCall - streamCall.Prompt = sentinelPrompt - streamCall.Messages = opts.Messages - streamCall.PrepareStep = func( - stepCtx context.Context, - options fantasy.PrepareStepFunctionOptions, - ) (context.Context, fantasy.PrepareStepResult, error) { - return stepCtx, prepareStepResult( - options.Messages, - sentinelPrompt, - opts.ActiveTools, - applyAnthropicCaching, - ), nil - } - streamCall.OnStepStart = func(_ int) error { - resetStepState() - return nil - } - streamCall.OnTextDelta = func(_ string, text string) error { - appendDraftText(text) - publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeText, - Text: text, - }) - return nil - } - streamCall.OnReasoningDelta = func(id string, text string) error { - appendDraftReasoning(text) - setReasoningTitleFromText(id, text) - stepStateMu.Lock() - title := streamReasoningTitles[id] - stepStateMu.Unlock() - publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeReasoning, - Text: text, - Title: title, - }) - return nil - } - streamCall.OnReasoningEnd = func(id string, _ fantasy.ReasoningContent) error { - stepStateMu.Lock() - if streamReasoningTitles[id] == "" { - // At the end of reasoning we have the full text, so we can - // safely evaluate first-line title format even if no newline - // ever arrived in deltas. - streamReasoningTitles[id] = chatprompt.ReasoningTitleFromFirstLine( - streamReasoningText[id], - ) + messages := opts.Messages + alreadyCompacted := false + var lastUsage fantasy.Usage + var lastProviderMetadata fantasy.ProviderMetadata + + for step := 0; step < opts.MaxSteps; step++ { + // Copy messages so that provider-specific caching + // mutations don't leak back to the caller's slice. + // copy copies Message structs by value, so field + // reassignments in addAnthropicPromptCaching only + // affect the prepared slice. + prepared := make([]fantasy.Message, len(messages)) + copy(prepared, messages) + if applyAnthropicCaching { + addAnthropicPromptCaching(prepared) } - title := streamReasoningTitles[id] - stepStateMu.Unlock() - if title != "" { - // Publish a title-only reasoning part so clients can update the - // reasoning header when metadata arrives at the end of streaming. - publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeReasoning, - Title: title, + + call := fantasy.Call{ + Prompt: prepared, + Tools: tools, + MaxOutputTokens: opts.ModelConfig.MaxOutputTokens, + Temperature: opts.ModelConfig.Temperature, + TopP: opts.ModelConfig.TopP, + TopK: opts.ModelConfig.TopK, + PresencePenalty: opts.ModelConfig.PresencePenalty, + FrequencyPenalty: opts.ModelConfig.FrequencyPenalty, + ProviderOptions: opts.ProviderOptions, + } + + var result stepResult + err := chatretry.Retry(ctx, func(retryCtx context.Context) error { + stream, streamErr := opts.Model.Stream(retryCtx, call) + if streamErr != nil { + return streamErr + } + var processErr error + result, processErr = processStepStream(retryCtx, stream, publishMessagePart) + return processErr + }, func(attempt int, retryErr error, delay time.Duration) { + // Reset result from the failed attempt so the next + // attempt starts clean. + result = stepResult{} + if opts.OnRetry != nil { + opts.OnRetry(attempt, retryErr, delay) + } + }) + if err != nil { + if errors.Is(err, ErrInterrupted) { + persistInterruptedStep(ctx, opts, &result) + return ErrInterrupted + } + return xerrors.Errorf("stream response: %w", err) + } + + // Execute tools before persisting so that tool results + // are included in the persisted step content. The + // persistence layer splits assistant and tool-result + // blocks into separate database messages by role. + var toolResults []fantasy.ToolResultContent + if result.shouldContinue { + // Check for context cancellation before starting + // tool execution. If the chat was interrupted + // between stream completion and here, persist + // what we have and bail out. + if ctx.Err() != nil { + if errors.Is(context.Cause(ctx), ErrInterrupted) { + persistInterruptedStep(ctx, opts, &result) + return ErrInterrupted + } + return ctx.Err() + } + + toolResults = executeTools(ctx, opts.Tools, result.toolCalls, func(tr fantasy.ToolResultContent) { + publishMessagePart( + fantasy.MessageRoleTool, + chatprompt.PartFromContent(tr), + ) }) + for _, tr := range toolResults { + result.content = append(result.content, tr) + } } - return nil - } - streamCall.OnToolInputStart = func(id, toolName string) error { - upsertDraftToolCall(id, toolName, "", false) - return nil - } - streamCall.OnToolInputDelta = func(id, delta string) error { - stepStateMu.Lock() - toolName := streamToolNames[id] - stepStateMu.Unlock() - upsertDraftToolCall(id, toolName, delta, true) - publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: id, - ToolName: toolName, - ArgsDelta: delta, - }) - return nil - } - streamCall.OnToolCall = func(toolCall fantasy.ToolCallContent) error { - upsertDraftToolCall(toolCall.ToolCallID, toolCall.ToolName, toolCall.Input, false) - publishMessagePart( - fantasy.MessageRoleAssistant, - chatprompt.PartFromContent(toolCall), - ) - return nil - } - streamCall.OnSource = func(source fantasy.SourceContent) error { - appendDraftSource(source) - publishMessagePart( - fantasy.MessageRoleAssistant, - chatprompt.PartFromContent(source), - ) - return nil - } - streamCall.OnToolResult = func(result fantasy.ToolResultContent) error { - publishMessagePart( - fantasy.MessageRoleTool, - chatprompt.PartFromContent(result), - ) - stepStateMu.Lock() - if result.ToolCallID != "" && strings.TrimSpace(result.ToolName) != "" { - streamToolNames[result.ToolCallID] = result.ToolName - } - stepToolResultContents = append(stepToolResultContents, result) - stepStateMu.Unlock() - - return nil - } - streamCall.OnStepFinish = func(stepResult fantasy.StepResult) error { - contextLimit := extractContextLimit(stepResult.ProviderMetadata) + // Extract context limit from provider metadata. + contextLimit := extractContextLimit(result.providerMetadata) if !contextLimit.Valid && opts.ContextLimitFallback > 0 { contextLimit = sql.NullInt64{ Int64: opts.ContextLimitFallback, @@ -444,95 +291,494 @@ func Run(ctx context.Context, opts RunOptions) (*fantasy.AgentResult, error) { } } - return opts.PersistStep(ctx, PersistedStep{ - Content: stepResult.Content, - Usage: stepResult.Usage, + // Persist the step — errors propagate directly. + if err := opts.PersistStep(ctx, PersistedStep{ + Content: result.content, + Usage: result.usage, ContextLimit: contextLimit, - }) - } + }); err != nil { + return xerrors.Errorf("persist step: %w", err) + } - var result *fantasy.AgentResult - err := chatretry.Retry(ctx, func(retryCtx context.Context) error { - var streamErr error - result, streamErr = agent.Stream(retryCtx, streamCall) - if streamErr != nil { - // Interrupts are not retryable — propagate them - // immediately so processChat can set the correct - // status. - if errors.Is(streamErr, context.Canceled) && - errors.Is(context.Cause(retryCtx), ErrInterrupted) { - if persistErr := persistInterruptedStep(); persistErr != nil { - if opts.OnInterruptedPersistError != nil { - opts.OnInterruptedPersistError(persistErr) - } + lastUsage = result.usage + lastProviderMetadata = result.providerMetadata + + // Inline compaction. + if opts.Compaction != nil && opts.ReloadMessages != nil { + did, compactErr := tryCompact( + ctx, + opts.Model, + opts.Compaction, + opts.ContextLimitFallback, + result.usage, + result.providerMetadata, + messages, + ) + if compactErr != nil && opts.Compaction.OnError != nil { + opts.Compaction.OnError(compactErr) + } + if did { + alreadyCompacted = true + reloaded, reloadErr := opts.ReloadMessages(ctx) + if reloadErr != nil { + return xerrors.Errorf("reload messages after compaction: %w", reloadErr) } - // Return ErrInterrupted directly so the retry - // loop sees a non-retryable error and stops. - return ErrInterrupted + messages = reloaded } - return streamErr } - return nil - }, func(attempt int, retryErr error, delay time.Duration) { - // Reset accumulated draft state from the failed attempt - // so the next attempt starts clean. - resetStepState() - if opts.OnRetry != nil { - opts.OnRetry(attempt, retryErr, delay) + if !result.shouldContinue { + break } - }) - if err != nil { - if errors.Is(err, ErrInterrupted) { - return nil, ErrInterrupted - } - return nil, xerrors.Errorf("stream response: %w", err) + + // Build messages from the step for the next iteration. + // toResponseMessages produces assistant-role content + // (text, reasoning, tool calls) and tool-result content. + stepMessages := result.toResponseMessages() + messages = append(messages, stepMessages...) } - if opts.Compaction != nil { - if err := maybeCompact(ctx, opts, result); err != nil { + + // Post-run compaction safety net: if we never compacted + // during the loop, try once at the end. + if !alreadyCompacted && opts.Compaction != nil { + if _, err := tryCompact( + ctx, + opts.Model, + opts.Compaction, + opts.ContextLimitFallback, + lastUsage, + lastProviderMetadata, + messages, + ); err != nil { if opts.Compaction.OnError != nil { opts.Compaction.OnError(err) } } } + return nil +} + +// processStepStream consumes a fantasy StreamResponse and +// accumulates all content into a stepResult. Callbacks fire +// inline and their errors propagate directly. +func processStepStream( + ctx context.Context, + stream fantasy.StreamResponse, + publishMessagePart func(fantasy.MessageRole, codersdk.ChatMessagePart), +) (stepResult, error) { + var result stepResult + + activeToolCalls := make(map[string]*fantasy.ToolCallContent) + activeTextContent := make(map[string]string) + activeReasoningContent := make(map[string]reasoningState) + // Track tool names by ID for input delta publishing. + toolNames := make(map[string]string) + // Track reasoning text/titles for title extraction. + reasoningTitles := make(map[string]string) + reasoningText := make(map[string]string) + + setReasoningTitleFromText := func(id string, text string) { + if id == "" || strings.TrimSpace(text) == "" { + return + } + if reasoningTitles[id] != "" { + return + } + reasoningText[id] += text + if !strings.ContainsAny(reasoningText[id], "\r\n") { + return + } + title := chatprompt.ReasoningTitleFromFirstLine(reasoningText[id]) + if title == "" { + return + } + reasoningTitles[id] = title + } + + for part := range stream { + switch part.Type { + case fantasy.StreamPartTypeTextStart: + activeTextContent[part.ID] = "" + + case fantasy.StreamPartTypeTextDelta: + if _, exists := activeTextContent[part.ID]; exists { + activeTextContent[part.ID] += part.Delta + } + publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeText, + Text: part.Delta, + }) + + case fantasy.StreamPartTypeTextEnd: + if text, exists := activeTextContent[part.ID]; exists { + result.content = append(result.content, fantasy.TextContent{ + Text: text, + ProviderMetadata: part.ProviderMetadata, + }) + delete(activeTextContent, part.ID) + } + + case fantasy.StreamPartTypeReasoningStart: + activeReasoningContent[part.ID] = reasoningState{ + text: part.Delta, + options: part.ProviderMetadata, + } + + case fantasy.StreamPartTypeReasoningDelta: + if active, exists := activeReasoningContent[part.ID]; exists { + active.text += part.Delta + active.options = part.ProviderMetadata + activeReasoningContent[part.ID] = active + } + setReasoningTitleFromText(part.ID, part.Delta) + title := reasoningTitles[part.ID] + publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeReasoning, + Text: part.Delta, + Title: title, + }) + + case fantasy.StreamPartTypeReasoningEnd: + if active, exists := activeReasoningContent[part.ID]; exists { + if part.ProviderMetadata != nil { + active.options = part.ProviderMetadata + } + content := fantasy.ReasoningContent{ + Text: active.text, + ProviderMetadata: active.options, + } + result.content = append(result.content, content) + delete(activeReasoningContent, part.ID) + + // Derive reasoning title at end of reasoning + // block if we haven't yet. + if reasoningTitles[part.ID] == "" { + reasoningTitles[part.ID] = chatprompt.ReasoningTitleFromFirstLine( + reasoningText[part.ID], + ) + } + title := reasoningTitles[part.ID] + if title != "" { + publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeReasoning, + Title: title, + }) + } + } + + case fantasy.StreamPartTypeToolInputStart: + activeToolCalls[part.ID] = &fantasy.ToolCallContent{ + ToolCallID: part.ID, + ToolName: part.ToolCallName, + Input: "", + ProviderExecuted: part.ProviderExecuted, + } + if strings.TrimSpace(part.ToolCallName) != "" { + toolNames[part.ID] = part.ToolCallName + } + + case fantasy.StreamPartTypeToolInputDelta: + if toolCall, exists := activeToolCalls[part.ID]; exists { + toolCall.Input += part.Delta + } + toolName := toolNames[part.ID] + publishMessagePart(fantasy.MessageRoleAssistant, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: part.ID, + ToolName: toolName, + ArgsDelta: part.Delta, + }) + + case fantasy.StreamPartTypeToolInputEnd: + // No callback needed; the full tool call arrives in + // StreamPartTypeToolCall. + + case fantasy.StreamPartTypeToolCall: + tc := fantasy.ToolCallContent{ + ToolCallID: part.ID, + ToolName: part.ToolCallName, + Input: part.ToolCallInput, + ProviderExecuted: part.ProviderExecuted, + ProviderMetadata: part.ProviderMetadata, + } + result.toolCalls = append(result.toolCalls, tc) + result.content = append(result.content, tc) + if strings.TrimSpace(part.ToolCallName) != "" { + toolNames[part.ID] = part.ToolCallName + } + // Clean up active tool call tracking. + delete(activeToolCalls, part.ID) + + publishMessagePart( + fantasy.MessageRoleAssistant, + chatprompt.PartFromContent(tc), + ) + + case fantasy.StreamPartTypeSource: + sourceContent := fantasy.SourceContent{ + SourceType: part.SourceType, + ID: part.ID, + URL: part.URL, + Title: part.Title, + ProviderMetadata: part.ProviderMetadata, + } + result.content = append(result.content, sourceContent) + publishMessagePart( + fantasy.MessageRoleAssistant, + chatprompt.PartFromContent(sourceContent), + ) + + case fantasy.StreamPartTypeFinish: + result.usage = part.Usage + result.finishReason = part.FinishReason + result.providerMetadata = part.ProviderMetadata + + case fantasy.StreamPartTypeError: + // Detect interruption: context canceled with + // ErrInterrupted as the cause. + if errors.Is(part.Error, context.Canceled) && + errors.Is(context.Cause(ctx), ErrInterrupted) { + // Flush in-progress content so that + // persistInterruptedStep has access to partial + // text, reasoning, and tool calls that were + // still streaming when the interrupt arrived. + flushActiveState( + &result, + activeTextContent, + activeReasoningContent, + activeToolCalls, + toolNames, + ) + return result, ErrInterrupted + } + return result, part.Error + } + } + + result.shouldContinue = len(result.toolCalls) > 0 && + result.finishReason == fantasy.FinishReasonToolCalls return result, nil } -//nolint:revive // Boolean controls Anthropic-specific caching behavior. -func prepareStepResult( - messages []fantasy.Message, - sentinel string, - activeTools []string, - anthropicCaching bool, -) fantasy.PrepareStepResult { - filtered := make([]fantasy.Message, 0, len(messages)) - removed := false - for _, message := range messages { - if !removed && - message.Role == fantasy.MessageRoleUser && - len(message.Content) == 1 { - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](message.Content[0]) - if ok && textPart.Text == sentinel { - removed = true - continue - } - } - filtered = append(filtered, message) +// executeTools runs each tool call sequentially after the stream +// completes. Results are published via onResult as each tool +// finishes. +func executeTools( + ctx context.Context, + allTools []fantasy.AgentTool, + toolCalls []fantasy.ToolCallContent, + onResult func(fantasy.ToolResultContent), +) []fantasy.ToolResultContent { + if len(toolCalls) == 0 { + return nil } - result := fantasy.PrepareStepResult{ - Messages: filtered, + toolMap := make(map[string]fantasy.AgentTool, len(allTools)) + for _, t := range allTools { + toolMap[t.Info().Name] = t } - if anthropicCaching { - result.Messages = addAnthropicPromptCaching(result.Messages) + + results := make([]fantasy.ToolResultContent, 0, len(toolCalls)) + for _, tc := range toolCalls { + tr := executeSingleTool(ctx, toolMap, tc) + results = append(results, tr) + if onResult != nil { + onResult(tr) + } } - if len(activeTools) > 0 { - result.ActiveTools = append([]string(nil), activeTools...) + return results +} + +// executeSingleTool executes one tool call and converts the +// response into a ToolResultContent. +func executeSingleTool( + ctx context.Context, + toolMap map[string]fantasy.AgentTool, + tc fantasy.ToolCallContent, +) fantasy.ToolResultContent { + result := fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + ProviderExecuted: false, + } + + tool, exists := toolMap[tc.ToolName] + if !exists { + result.Result = fantasy.ToolResultOutputContentError{ + Error: xerrors.New("Tool not found: " + tc.ToolName), + } + return result + } + + resp, err := tool.Run(ctx, fantasy.ToolCall{ + ID: tc.ToolCallID, + Name: tc.ToolName, + Input: tc.Input, + }) + if err != nil { + result.Result = fantasy.ToolResultOutputContentError{ + Error: err, + } + result.ClientMetadata = resp.Metadata + return result + } + + result.ClientMetadata = resp.Metadata + switch { + case resp.IsError: + result.Result = fantasy.ToolResultOutputContentError{ + Error: xerrors.New(resp.Content), + } + case resp.Type == "image" || resp.Type == "media": + result.Result = fantasy.ToolResultOutputContentMedia{ + Data: string(resp.Data), + MediaType: resp.MediaType, + Text: resp.Content, + } + default: + result.Result = fantasy.ToolResultOutputContentText{ + Text: resp.Content, + } } return result } +// flushActiveState moves any in-progress text, reasoning, and +// tool calls from the active tracking maps into result.content +// and result.toolCalls. This is called on interruption so that +// partial content from an incomplete stream is available for +// persistence. +func flushActiveState( + result *stepResult, + activeText map[string]string, + activeReasoning map[string]reasoningState, + activeToolCalls map[string]*fantasy.ToolCallContent, + toolNames map[string]string, +) { + // Flush partial text content. + for _, text := range activeText { + if text != "" { + result.content = append(result.content, fantasy.TextContent{Text: text}) + } + } + + // Flush partial reasoning content. + for _, rs := range activeReasoning { + if rs.text != "" { + result.content = append(result.content, fantasy.ReasoningContent{ + Text: rs.text, + ProviderMetadata: rs.options, + }) + } + } + + // Flush in-progress tool calls. These haven't received a + // StreamPartTypeToolCall yet, so they only exist in + // activeToolCalls. We add them to both content and toolCalls + // so persistInterruptedStep can generate synthetic error + // results for them. + for id, tc := range activeToolCalls { + if tc == nil { + continue + } + // Prefer the tool name from the toolNames map since + // ToolInputStart may provide a cleaner name. + toolName := tc.ToolName + if name, ok := toolNames[id]; ok && strings.TrimSpace(name) != "" { + toolName = name + } + flushed := fantasy.ToolCallContent{ + ToolCallID: tc.ToolCallID, + ToolName: toolName, + Input: tc.Input, + ProviderExecuted: tc.ProviderExecuted, + } + result.content = append(result.content, flushed) + result.toolCalls = append(result.toolCalls, flushed) + } +} + +// persistInterruptedStep saves all accumulated content from a +// partial stream. Since we own the stepResult directly, no shadow +// state is needed. +func persistInterruptedStep( + ctx context.Context, + opts RunOptions, + result *stepResult, +) { + if result == nil || (len(result.content) == 0 && len(result.toolCalls) == 0) { + return + } + + // Track which tool calls already have results in the content. + answeredToolCalls := make(map[string]struct{}) + for _, c := range result.content { + tr, ok := fantasy.AsContentType[fantasy.ToolResultContent](c) + if ok && tr.ToolCallID != "" { + answeredToolCalls[tr.ToolCallID] = struct{}{} + } + } + + // Build combined content: all accumulated content + synthetic + // interrupted results for any unanswered tool calls. + content := make([]fantasy.Content, 0, len(result.content)) + content = append(content, result.content...) + + for _, tc := range result.toolCalls { + if tc.ToolCallID == "" { + continue + } + if _, exists := answeredToolCalls[tc.ToolCallID]; exists { + continue + } + content = append(content, fantasy.ToolResultContent{ + ToolCallID: tc.ToolCallID, + ToolName: tc.ToolName, + Result: fantasy.ToolResultOutputContentError{ + Error: xerrors.New(interruptedToolResultErrorMessage), + }, + }) + answeredToolCalls[tc.ToolCallID] = struct{}{} + } + + persistCtx := context.WithoutCancel(ctx) + if err := opts.PersistStep(persistCtx, PersistedStep{ + Content: content, + }); err != nil { + if opts.OnInterruptedPersistError != nil { + opts.OnInterruptedPersistError(err) + } + } +} + +// buildToolDefinitions converts AgentTool definitions into the +// fantasy.Tool slice expected by fantasy.Call. When activeTools +// is non-empty, only tools whose name appears in the list are +// included. This mirrors fantasy's agent.prepareTools filtering. +func buildToolDefinitions(tools []fantasy.AgentTool, activeTools []string) []fantasy.Tool { + prepared := make([]fantasy.Tool, 0, len(tools)) + for _, tool := range tools { + info := tool.Info() + if len(activeTools) > 0 && !slices.Contains(activeTools, info.Name) { + continue + } + inputSchema := map[string]any{ + "type": "object", + "properties": info.Parameters, + "required": info.Required, + } + schema.Normalize(inputSchema) + prepared = append(prepared, fantasy.FunctionTool{ + Name: info.Name, + Description: info.Description, + InputSchema: inputSchema, + ProviderOptions: tool.ProviderOptions(), + }) + } + return prepared +} + func shouldApplyAnthropicPromptCaching(model fantasy.LanguageModel) bool { if model == nil { return false @@ -540,7 +786,10 @@ func shouldApplyAnthropicPromptCaching(model fantasy.LanguageModel) bool { return model.Provider() == fantasyanthropic.Name } -func addAnthropicPromptCaching(messages []fantasy.Message) []fantasy.Message { +// addAnthropicPromptCaching mutates messages in-place, setting +// ProviderOptions for Anthropic prompt caching on the last system +// message and the final two messages. +func addAnthropicPromptCaching(messages []fantasy.Message) { for i := range messages { messages[i].ProviderOptions = nil } @@ -564,8 +813,6 @@ func addAnthropicPromptCaching(messages []fantasy.Message) []fantasy.Message { messages[i].ProviderOptions = providerOption } } - - return messages } func extractContextLimit(metadata fantasy.ProviderMetadata) sql.NullInt64 { diff --git a/coderd/chatd/chatloop/chatloop_test.go b/coderd/chatd/chatloop/chatloop_test.go index c6ae4309e1..aac9ee20b4 100644 --- a/coderd/chatd/chatloop/chatloop_test.go +++ b/coderd/chatd/chatloop/chatloop_test.go @@ -4,6 +4,7 @@ import ( "context" "iter" "strings" + "sync" "testing" "charm.land/fantasy" @@ -34,7 +35,7 @@ func TestRun_ActiveToolsPrepareBehavior(t *testing.T) { persistStepCalls := 0 var persistedStep PersistedStep - _, err := Run(context.Background(), RunOptions{ + err := Run(context.Background(), RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleSystem, "sys-1"), @@ -130,7 +131,7 @@ func TestRun_InterruptedStepPersistsSyntheticToolResult(t *testing.T) { persistedAssistantCtxErr := xerrors.New("unset") var persistedContent []fantasy.Content - _, err := Run(ctx, RunOptions{ + err := Run(ctx, RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleUser, "hello"), @@ -274,6 +275,136 @@ func containsPromptSentinel(prompt []fantasy.Message) bool { return false } +func TestRun_MultiStepToolExecution(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var streamCalls int + var secondCallPrompt []fantasy.Message + + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { + mu.Lock() + step := streamCalls + streamCalls++ + mu.Unlock() + + switch step { + case 0: + // Step 0: produce a tool call. + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "read_file"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{"path":"main.go"}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "read_file", + ToolCallInput: `{"path":"main.go"}`, + }, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, + }), nil + default: + // Step 1: capture the prompt the loop sent us, + // then return plain text. + mu.Lock() + secondCallPrompt = append([]fantasy.Message(nil), call.Prompt...) + mu.Unlock() + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "all done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + } + }, + } + + var persistStepCalls int + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "please read main.go"), + }, + Tools: []fantasy.AgentTool{ + newNoopTool("read_file"), + }, + MaxSteps: 5, + PersistStep: func(_ context.Context, _ PersistedStep) error { + persistStepCalls++ + return nil + }, + }) + require.NoError(t, err) + + // Stream was called twice: once for the tool-call step, + // once for the follow-up text step. + require.Equal(t, 2, streamCalls) + + // PersistStep is called once per step. + require.Equal(t, 2, persistStepCalls) + + // The second call's prompt must contain the assistant message + // from step 0 (with the tool call) and a tool-result message. + require.NotEmpty(t, secondCallPrompt) + + var foundAssistantToolCall bool + var foundToolResult bool + for _, msg := range secondCallPrompt { + if msg.Role == fantasy.MessageRoleAssistant { + for _, part := range msg.Content { + if tc, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](part); ok { + if tc.ToolCallID == "tc-1" && tc.ToolName == "read_file" { + foundAssistantToolCall = true + } + } + } + } + if msg.Role == fantasy.MessageRoleTool { + for _, part := range msg.Content { + if tr, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part); ok { + if tr.ToolCallID == "tc-1" { + foundToolResult = true + } + } + } + } + } + require.True(t, foundAssistantToolCall, "second call prompt should contain assistant tool call from step 0") + require.True(t, foundToolResult, "second call prompt should contain tool result message") +} + +func TestRun_PersistStepErrorPropagates(t *testing.T) { + t.Parallel() + + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "hello"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, + }), nil + }, + } + + persistErr := xerrors.New("database write failed") + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "hello"), + }, + MaxSteps: 1, + PersistStep: func(_ context.Context, _ PersistedStep) error { + return persistErr + }, + }) + require.Error(t, err) + require.ErrorContains(t, err, "database write failed") +} + func hasAnthropicEphemeralCacheControl(message fantasy.Message) bool { if len(message.ProviderOptions) == 0 { return false diff --git a/coderd/chatd/chatloop/compaction.go b/coderd/chatd/chatloop/compaction.go index 610682169f..2238c92544 100644 --- a/coderd/chatd/chatloop/compaction.go +++ b/coderd/chatd/chatloop/compaction.go @@ -2,11 +2,14 @@ package chatloop import ( "context" + "encoding/json" "strings" "time" "charm.land/fantasy" "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk" ) const ( @@ -30,8 +33,18 @@ type CompactionOptions struct { SystemSummaryPrefix string Timeout time.Duration Persist func(context.Context, CompactionResult) error - OnStart func() - OnError func(error) + + // ToolCallID and ToolName identify the synthetic tool call + // used to represent compaction in the message stream. + ToolCallID string + ToolName string + + // PublishMessagePart publishes streaming parts to connected + // clients so they see "Summarizing..." / "Summarized" UI + // transitions during compaction. + PublishMessagePart func(fantasy.MessageRole, codersdk.ChatMessagePart) + + OnError func(error) } type CompactionResult struct { @@ -43,18 +56,145 @@ type CompactionResult struct { ContextLimit int64 } -func maybeCompact( +// tryCompact checks whether context usage exceeds the compaction +// threshold and, if so, generates and persists a summary. Returns +// (true, nil) when compaction was performed, (false, nil) when not +// needed, and (false, err) on failure. +func tryCompact( ctx context.Context, - runOpts RunOptions, - runResult *fantasy.AgentResult, -) error { - if runResult == nil || runOpts.Compaction == nil { - return nil + model fantasy.LanguageModel, + compaction *CompactionOptions, + contextLimitFallback int64, + stepUsage fantasy.Usage, + stepMetadata fantasy.ProviderMetadata, + allMessages []fantasy.Message, +) (bool, error) { + config, ok := normalizedCompactionConfig(compaction) + if !ok { + return false, nil } - config := *runOpts.Compaction + contextTokens := contextTokensFromUsage(stepUsage) + if contextTokens <= 0 { + return false, nil + } + + metadataLimit := extractContextLimit(stepMetadata) + contextLimit := resolveContextLimit( + metadataLimit.Int64, + config.ContextLimit, + contextLimitFallback, + ) + + usagePercent, compact := shouldCompact( + contextTokens, contextLimit, config.ThresholdPercent, + ) + if !compact { + return false, nil + } + + // Publish the "Summarizing..." tool-call indicator so + // connected clients see activity during summary generation. + if config.PublishMessagePart != nil && config.ToolCallID != "" { + config.PublishMessagePart( + fantasy.MessageRoleAssistant, + codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: config.ToolCallID, + ToolName: config.ToolName, + }, + ) + } + + summary, err := generateCompactionSummary( + ctx, model, allMessages, config, + ) + if err != nil { + return false, err + } + if summary == "" { + // Publish a tool-result error so connected clients + // see the compaction failure. + publishCompactionError(config, "compaction produced an empty summary") + return false, xerrors.New("compaction produced an empty summary") + } + + systemSummary := strings.TrimSpace( + config.SystemSummaryPrefix + "\n\n" + summary, + ) + + err = config.Persist(ctx, CompactionResult{ + SystemSummary: systemSummary, + SummaryReport: summary, + ThresholdPercent: config.ThresholdPercent, + UsagePercent: usagePercent, + ContextTokens: contextTokens, + ContextLimit: contextLimit, + }) + if err != nil { + publishCompactionError(config, "failed to persist compaction result") + return false, xerrors.Errorf("persist compaction: %w", err) + } + + // Publish the "Summarized" tool-result part so the client + // transitions from the in-progress indicator to the final + // state. + if config.PublishMessagePart != nil && config.ToolCallID != "" { + resultJSON, _ := json.Marshal(map[string]any{ + "summary": summary, + "source": "automatic", + "threshold_percent": config.ThresholdPercent, + "usage_percent": usagePercent, + "context_tokens": contextTokens, + "context_limit_tokens": contextLimit, + }) + config.PublishMessagePart( + fantasy.MessageRoleTool, + codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: config.ToolCallID, + ToolName: config.ToolName, + Result: resultJSON, + }, + ) + } + + return true, nil +} + +// publishCompactionError sends a tool-result error part so +// connected clients see that compaction failed. +func publishCompactionError(config CompactionOptions, msg string) { + if config.PublishMessagePart == nil || config.ToolCallID == "" { + return + } + errJSON, _ := json.Marshal(map[string]any{ + "error": msg, + }) + config.PublishMessagePart( + fantasy.MessageRoleTool, + codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: config.ToolCallID, + ToolName: config.ToolName, + Result: errJSON, + IsError: true, + }, + ) +} + +// normalizedCompactionConfig returns a copy of the compaction options +// with defaults applied. The bool is false when compaction is +// disabled (nil options, missing Persist callback, or threshold at +// 100%). +func normalizedCompactionConfig(opts *CompactionOptions) (CompactionOptions, bool) { + if opts == nil { + return CompactionOptions{}, false + } + + config := *opts if config.Persist == nil { - return xerrors.New("compaction persist callback is required") + return CompactionOptions{}, false } if strings.TrimSpace(config.SummaryPrompt) == "" { config.SummaryPrompt = defaultCompactionSummaryPrompt @@ -69,116 +209,80 @@ func maybeCompact( config.ThresholdPercent > maxCompactionThresholdPercent { config.ThresholdPercent = defaultCompactionThresholdPercent } - - if config.ThresholdPercent >= maxCompactionThresholdPercent { - return nil - } - if runOpts.MaxSteps > 0 && len(runResult.Steps) >= runOpts.MaxSteps { - lastStep := runResult.Steps[len(runResult.Steps)-1] - if lastStep.FinishReason == fantasy.FinishReasonToolCalls && - len(lastStep.Content.ToolCalls()) > 0 { - return nil - } + if config.ThresholdPercent == maxCompactionThresholdPercent { + return CompactionOptions{}, false } - contextTokens := int64(0) - contextLimitFromMetadata := int64(0) - for i := len(runResult.Steps) - 1; i >= 0; i-- { - usage := runResult.Steps[i].Usage - total := int64(0) - hasContextTokens := false - - if usage.InputTokens > 0 { - total += usage.InputTokens - hasContextTokens = true - } - if usage.CacheReadTokens > 0 { - total += usage.CacheReadTokens - hasContextTokens = true - } - if usage.CacheCreationTokens > 0 { - total += usage.CacheCreationTokens - hasContextTokens = true - } - if !hasContextTokens && usage.TotalTokens > 0 { - total = usage.TotalTokens - hasContextTokens = true - } - if !hasContextTokens || total <= 0 { - continue - } - - contextTokens = total - metadataLimit := extractContextLimit(runResult.Steps[i].ProviderMetadata) - if metadataLimit.Valid && metadataLimit.Int64 > 0 { - contextLimitFromMetadata = metadataLimit.Int64 - } - break - } - if contextTokens <= 0 { - return nil - } - - contextLimit := contextLimitFromMetadata - if contextLimit <= 0 && config.ContextLimit > 0 { - contextLimit = config.ContextLimit - } - if contextLimit <= 0 && runOpts.ContextLimitFallback > 0 { - contextLimit = runOpts.ContextLimitFallback - } - if contextLimit <= 0 { - return nil - } - - usagePercent := (float64(contextTokens) / float64(contextLimit)) * 100 - if usagePercent < float64(config.ThresholdPercent) { - return nil - } - - if config.OnStart != nil { - config.OnStart() - } - - summary, err := generateCompactionSummary( - ctx, - runOpts.Model, - runOpts.Messages, - runResult.Steps, - config, - ) - if err != nil { - return err - } - if summary == "" { - return nil - } - - systemSummary := strings.TrimSpace( - config.SystemSummaryPrefix + "\n\n" + summary, - ) - - return config.Persist(ctx, CompactionResult{ - SystemSummary: systemSummary, - SummaryReport: summary, - ThresholdPercent: config.ThresholdPercent, - UsagePercent: usagePercent, - ContextTokens: contextTokens, - ContextLimit: contextLimit, - }) + return config, true } +// contextTokensFromUsage returns the total context token count from +// a step's usage report. It sums input, cache-read, and +// cache-creation tokens when available, falling back to TotalTokens +// if none of the granular fields are set. +func contextTokensFromUsage(usage fantasy.Usage) int64 { + total := int64(0) + hasContextTokens := false + + if usage.InputTokens > 0 { + total += usage.InputTokens + hasContextTokens = true + } + if usage.CacheReadTokens > 0 { + total += usage.CacheReadTokens + hasContextTokens = true + } + if usage.CacheCreationTokens > 0 { + total += usage.CacheCreationTokens + hasContextTokens = true + } + if !hasContextTokens && usage.TotalTokens > 0 { + total = usage.TotalTokens + } + + return total +} + +// resolveContextLimit picks the first positive value from metadata, +// configured limit, and fallback — in that priority order. Returns +// 0 when none are positive. +func resolveContextLimit(metadataLimit, configLimit, fallback int64) int64 { + if metadataLimit > 0 { + return metadataLimit + } + if configLimit > 0 { + return configLimit + } + if fallback > 0 { + return fallback + } + return 0 +} + +// shouldCompact returns the usage percentage and whether it exceeds +// the threshold. Returns (0, false) when contextLimit is +// non-positive. +func shouldCompact(contextTokens, contextLimit int64, thresholdPercent int32) (float64, bool) { + if contextLimit <= 0 { + return 0, false + } + usagePercent := (float64(contextTokens) / float64(contextLimit)) * 100 + return usagePercent, usagePercent >= float64(thresholdPercent) +} + +// generateCompactionSummary asks the model to summarize the +// conversation so far. The provided messages should contain the +// complete history (system prompt, user/assistant turns, tool +// results). A final user message with the summary prompt is appended +// before calling the model. func generateCompactionSummary( ctx context.Context, model fantasy.LanguageModel, messages []fantasy.Message, - steps []fantasy.StepResult, options CompactionOptions, ) (string, error) { - summaryPrompt := make([]fantasy.Message, 0, len(messages)+len(steps)+1) + summaryPrompt := make([]fantasy.Message, 0, len(messages)+1) summaryPrompt = append(summaryPrompt, messages...) - for _, step := range steps { - summaryPrompt = append(summaryPrompt, step.Messages...) - } summaryPrompt = append(summaryPrompt, fantasy.Message{ Role: fantasy.MessageRoleUser, Content: []fantasy.MessagePart{ diff --git a/coderd/chatd/chatloop/compaction_test.go b/coderd/chatd/chatloop/compaction_test.go index f2f4df18dc..c0a2739b9a 100644 --- a/coderd/chatd/chatloop/compaction_test.go +++ b/coderd/chatd/chatloop/compaction_test.go @@ -2,11 +2,14 @@ package chatloop //nolint:testpackage // Uses internal symbols. import ( "context" + "sync" "testing" "charm.land/fantasy" "github.com/stretchr/testify/require" "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk" ) func TestRun_Compaction(t *testing.T) { @@ -54,7 +57,7 @@ func TestRun_Compaction(t *testing.T) { }, } - _, err := Run(context.Background(), RunOptions{ + err := Run(context.Background(), RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleUser, "hello"), @@ -83,14 +86,14 @@ func TestRun_Compaction(t *testing.T) { require.InDelta(t, 80.0, persistedCompaction.UsagePercent, 0.0001) }) - t.Run("OnStartFiresBeforePersist", func(t *testing.T) { + t.Run("PublishesPartsBeforeAndAfterPersist", func(t *testing.T) { t.Parallel() const summaryText = "compaction summary for ordering test" - // Track the order of callbacks to verify OnStart fires - // before the Generate call (summary generation) and - // before Persist. + // Track the order of callbacks to verify the tool-call + // part publishes before Generate (summary generation) + // and the tool-result part publishes after Persist. var callOrder []string model := &loopTestModel{ @@ -120,7 +123,7 @@ func TestRun_Compaction(t *testing.T) { }, } - _, err := Run(context.Background(), RunOptions{ + err := Run(context.Background(), RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleUser, "hello"), @@ -133,8 +136,15 @@ func TestRun_Compaction(t *testing.T) { Compaction: &CompactionOptions{ ThresholdPercent: 70, SummaryPrompt: "summarize now", - OnStart: func() { - callOrder = append(callOrder, "on_start") + ToolCallID: "test-tool-call-id", + ToolName: "chat_summarized", + PublishMessagePart: func(role fantasy.MessageRole, part codersdk.ChatMessagePart) { + switch part.Type { + case codersdk.ChatMessagePartTypeToolCall: + callOrder = append(callOrder, "publish_tool_call") + case codersdk.ChatMessagePartTypeToolResult: + callOrder = append(callOrder, "publish_tool_result") + } }, Persist: func(_ context.Context, _ CompactionResult) error { callOrder = append(callOrder, "persist") @@ -143,13 +153,18 @@ func TestRun_Compaction(t *testing.T) { }, }) require.NoError(t, err) - require.Equal(t, []string{"on_start", "generate", "persist"}, callOrder) + require.Equal(t, []string{ + "publish_tool_call", + "generate", + "persist", + "publish_tool_result", + }, callOrder) }) - t.Run("OnStartNotCalledBelowThreshold", func(t *testing.T) { + t.Run("PublishNotCalledBelowThreshold", func(t *testing.T) { t.Parallel() - onStartCalled := false + publishCalled := false model := &loopTestModel{ provider: "fake", @@ -166,7 +181,7 @@ func TestRun_Compaction(t *testing.T) { }, } - _, err := Run(context.Background(), RunOptions{ + err := Run(context.Background(), RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleUser, "hello"), @@ -178,8 +193,10 @@ func TestRun_Compaction(t *testing.T) { ContextLimitFallback: 100, Compaction: &CompactionOptions{ ThresholdPercent: 70, - OnStart: func() { - onStartCalled = true + ToolCallID: "test-tool-call-id", + ToolName: "chat_summarized", + PublishMessagePart: func(_ fantasy.MessageRole, _ codersdk.ChatMessagePart) { + publishCalled = true }, Persist: func(_ context.Context, _ CompactionResult) error { return nil @@ -187,7 +204,216 @@ func TestRun_Compaction(t *testing.T) { }, }) require.NoError(t, err) - require.False(t, onStartCalled, "OnStart should not fire when usage is below threshold") + require.False(t, publishCalled, "PublishMessagePart should not fire when usage is below threshold") + }) + + t.Run("MidLoopCompactionReloadsMessages", func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var streamCallCount int + persistCompactionCalls := 0 + reloadCalls := 0 + + const summaryText = "compacted summary" + + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + mu.Lock() + step := streamCallCount + streamCallCount++ + mu.Unlock() + + switch step { + case 0: + // Step 0: tool call with high usage (80/100 = 80% > 70%). + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "read_file"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "read_file", + ToolCallInput: `{}`, + }, + { + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonToolCalls, + Usage: fantasy.Usage{ + InputTokens: 80, + TotalTokens: 85, + }, + }, + }), nil + default: + // Step 1: text with low usage (30/100 = 30% < 70%). + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + { + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + Usage: fantasy.Usage{ + InputTokens: 30, + TotalTokens: 35, + }, + }, + }), nil + } + }, + generateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: []fantasy.Content{ + fantasy.TextContent{Text: summaryText}, + }, + }, nil + }, + } + + compactedMessages := []fantasy.Message{ + textMessage(fantasy.MessageRoleSystem, "compacted system"), + textMessage(fantasy.MessageRoleUser, "compacted user"), + } + + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "hello"), + }, + Tools: []fantasy.AgentTool{ + newNoopTool("read_file"), + }, + MaxSteps: 5, + PersistStep: func(_ context.Context, _ PersistedStep) error { + return nil + }, + ContextLimitFallback: 100, + Compaction: &CompactionOptions{ + ThresholdPercent: 70, + SummaryPrompt: "summarize now", + Persist: func(_ context.Context, _ CompactionResult) error { + persistCompactionCalls++ + return nil + }, + }, + ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { + reloadCalls++ + return compactedMessages, nil + }, + }) + require.NoError(t, err) + + // Compaction fired after step 0 (above threshold). + require.GreaterOrEqual(t, persistCompactionCalls, 1) + // ReloadMessages was called after mid-loop compaction. + require.GreaterOrEqual(t, reloadCalls, 1) + // Both steps ran (tool-call step + follow-up text step). + require.Equal(t, 2, streamCallCount) + }) + + t.Run("PostRunCompactionSkippedAfterMidLoop", func(t *testing.T) { + t.Parallel() + + var mu sync.Mutex + var streamCallCount int + persistCompactionCalls := 0 + + const summaryText = "compacted summary for skip test" + + model := &loopTestModel{ + provider: "fake", + streamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + mu.Lock() + step := streamCallCount + streamCallCount++ + mu.Unlock() + + switch step { + case 0: + // Step 0: tool call with high usage (80/100 = 80% > 70%). + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "read_file"}, + {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, + {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, + { + Type: fantasy.StreamPartTypeToolCall, + ID: "tc-1", + ToolCallName: "read_file", + ToolCallInput: `{}`, + }, + { + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonToolCalls, + Usage: fantasy.Usage{ + InputTokens: 80, + TotalTokens: 85, + }, + }, + }), nil + default: + // Step 1: text with low usage (20/100 = 20% < 70%). + return streamFromParts([]fantasy.StreamPart{ + {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, + {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"}, + {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, + { + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + Usage: fantasy.Usage{ + InputTokens: 20, + TotalTokens: 25, + }, + }, + }), nil + } + }, + generateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { + return &fantasy.Response{ + Content: []fantasy.Content{ + fantasy.TextContent{Text: summaryText}, + }, + }, nil + }, + } + + compactedMessages := []fantasy.Message{ + textMessage(fantasy.MessageRoleSystem, "compacted system"), + textMessage(fantasy.MessageRoleUser, "compacted user"), + } + + err := Run(context.Background(), RunOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "hello"), + }, + Tools: []fantasy.AgentTool{ + newNoopTool("read_file"), + }, + MaxSteps: 5, + PersistStep: func(_ context.Context, _ PersistedStep) error { + return nil + }, + ContextLimitFallback: 100, + Compaction: &CompactionOptions{ + ThresholdPercent: 70, + SummaryPrompt: "summarize now", + Persist: func(_ context.Context, _ CompactionResult) error { + persistCompactionCalls++ + return nil + }, + }, + ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { + return compactedMessages, nil + }, + }) + require.NoError(t, err) + + // Only mid-loop compaction fires after step 0. The post-run + // safety net is skipped because alreadyCompacted is true. + require.Equal(t, 1, persistCompactionCalls) }) t.Run("ErrorsAreReported", func(t *testing.T) { @@ -212,7 +438,7 @@ func TestRun_Compaction(t *testing.T) { } compactionErr := xerrors.New("unset") - _, err := Run(context.Background(), RunOptions{ + err := Run(context.Background(), RunOptions{ Model: model, Messages: []fantasy.Message{ textMessage(fantasy.MessageRoleUser, "hello"),