package chatloop

import (

"context"

"database/sql"

"encoding/base64"

"encoding/json"

"errors"

"maps"

"slices"

"strconv"

"strings"

"sync"

"time"

"unicode"

"charm.land/fantasy"

fantasyanthropic "charm.land/fantasy/providers/anthropic"

fantasyopenai "charm.land/fantasy/providers/openai"

"charm.land/fantasy/schema"

"golang.org/x/xerrors"

"cdr.dev/slog/v3"

"github.com/coder/coder/v2/coderd/database/dbtime"

"github.com/coder/coder/v2/coderd/x/chatd/chatdebug"

"github.com/coder/coder/v2/coderd/x/chatd/chaterror"

"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"

"github.com/coder/coder/v2/coderd/x/chatd/chatretry"

"github.com/coder/coder/v2/coderd/x/chatd/chatsanitize"

"github.com/coder/coder/v2/coderd/x/chatd/chattool"

"github.com/coder/coder/v2/codersdk"

"github.com/coder/quartz"

)

const (

interruptedToolResultErrorMessage = "tool call was interrupted before it produced a result"

// maxCompactionRetries limits how many times the post-run

// compaction safety net can re-enter the step loop. This

// prevents infinite compaction loops when the model keeps

// hitting the context limit after summarization.

maxCompactionRetries = 3

// defaultStartupTimeout bounds how long an individual

// model attempt may spend starting to respond before

// the attempt is canceled and retried.

defaultStartupTimeout = 60 * time.Second

)

var (

ErrInterrupted = xerrors.New("chat interrupted")

ErrDynamicToolCall = xerrors.New("dynamic tool call")

// ErrStopAfterTool is returned when a tool listed in

// StopAfterTools produces a successful result, indicating

// the run should terminate cleanly after persistence.

ErrStopAfterTool = xerrors.New("stop after tool")

errStartupTimeout = xerrors.New(

"chat response did not start before the startup timeout",

)

)

// PendingToolCall describes a tool call that targets a dynamic

// tool. These calls are not executed by the chatloop; instead

// they are persisted so the caller can fulfill them externally.

type PendingToolCall struct {

ToolCallID string

ToolName string

Args string

}

// 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. The

// persistence layer is responsible for splitting these into

// separate database messages by role.

type PersistedStep struct {

Content []fantasy.Content

Usage fantasy.Usage

ContextLimit sql.NullInt64

ProviderResponseID string

// Runtime is the wall-clock duration of this step,

// covering LLM streaming, tool execution, and retries.

// Zero indicates the duration was not measured (e.g.

// interrupted steps).

Runtime time.Duration

// PendingDynamicToolCalls lists tool calls that target

// dynamic tools. When non-empty the chatloop exits with

// ErrDynamicToolCall so the caller can execute them

// externally and resume the loop.

PendingDynamicToolCalls []PendingToolCall

// ToolCallCreatedAt maps tool-call IDs to the time

// the model emitted each tool call. Applied by the

// persistence layer to set CreatedAt on persisted

// tool-call ChatMessageParts.

ToolCallCreatedAt map[string]time.Time

// ToolResultCreatedAt maps tool-call IDs to the time

// each tool result was produced (or interrupted).

// Applied by the persistence layer to set CreatedAt

// on persisted tool-result ChatMessageParts.

ToolResultCreatedAt map[string]time.Time

}

// RunOptions configures a single streaming chat loop run.

type RunOptions struct {

Model fantasy.LanguageModel

Messages []fantasy.Message

Tools []fantasy.AgentTool

MaxSteps int

// StartupTimeout bounds how long each model attempt may

// spend opening the provider stream and waiting for its

// first stream part before the attempt is canceled and

// retried. Zero uses the production default.

StartupTimeout time.Duration

// Clock creates startup guard timers. In production use a

// real clock; tests can inject quartz.NewMock(t) to make

// startup timeout behavior deterministic.

Clock quartz.Clock

ActiveTools []string

ContextLimitFallback int64

// DynamicToolNames lists tool names that are handled

// externally. When the model invokes one of these tools

// the chatloop persists partial results and exits with

// ErrDynamicToolCall instead of executing the tool.

DynamicToolNames map[string]bool

// StopAfterTools lists tool names that, when they produce a

// successful result, cause the run to stop after persisting

// the current step. This is used for plan turns where

// propose_plan should terminate the run on success.

StopAfterTools map[string]struct{}

// 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

// ProviderTools are provider-native tools (like web search

// and computer use) whose definitions are passed directly

// to the provider API. When a ProviderTool has a non-nil

// Runner, tool calls are executed locally; otherwise the

// provider handles execution (e.g. web search).

ProviderTools []ProviderTool

PersistStep func(context.Context, PersistedStep) error

PublishMessagePart func(

role codersdk.ChatMessageRole,

part codersdk.ChatMessagePart,

)

// Callers should attach correlation fields (chat_id, owner_id, etc.)

// using Logger.With before passing the logger in.

Logger slog.Logger

Compaction *CompactionOptions

ReloadMessages func(context.Context) ([]fantasy.Message, error)

DisableChainMode func()

// PrepareMessages is called before each LLM step with the

// current message history. If it returns non-nil, the returned

// slice replaces messages for this and all subsequent steps.

// Used to inject system context that becomes available mid-loop

// (e.g. AGENTS.md after create_workspace).

PrepareMessages func([]fantasy.Message) []fantasy.Message

// OnRetry is called before each retry attempt when the LLM

// stream fails with a retryable error. It provides the attempt

// number, raw error, normalized classification, and backoff

// delay so callers can publish status events to connected

// clients. Callers should also clear any buffered stream state

// from the failed attempt in this callback to avoid sending

// duplicated content.

OnRetry chatretry.OnRetryFn

OnInterruptedPersistError func(error)

// Metrics records Prometheus metrics for the chatd subsystem.

// When nil, no metrics are recorded.

Metrics *Metrics

// BuiltinToolNames lists tool names that are built into chatd.

BuiltinToolNames map[string]bool

}

// ProviderTool pairs a provider-native tool definition with an

// optional local executor. When Runner is nil the tool is fully

// provider-executed (e.g. web search). When Runner is non-nil

// the definition is sent to the API but execution is handled

// locally (e.g. computer use).

type ProviderTool struct {

Definition fantasy.Tool

Runner fantasy.AgentTool

}

// 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

toolCallCreatedAt map[string]time.Time

toolResultCreatedAt map[string]time.Time

}

// 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 || strings.TrimSpace(text.Text) == "" {

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 || strings.TrimSpace(reasoning.Text) == "" {

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

}

part := fantasy.ToolResultPart{

ToolCallID: result.ToolCallID,

Output: result.Result,

ProviderExecuted: result.ProviderExecuted,

ProviderOptions: fantasy.ProviderOptions(result.ProviderMetadata),

}

// Provider-executed tool results (e.g. web_search)

// must stay in the assistant message so the result

// block appears inline after the corresponding

// server_tool_use block. This matches the persistence

// layer in chatd.go which keeps them in

// assistantBlocks.

if result.ProviderExecuted {

assistantParts = append(assistantParts, part)

} else {

toolParts = append(toolParts, part)

}

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 xerrors.New("chat model is required")

}

if opts.PersistStep == nil {

return xerrors.New("persist step callback is required")

}

if opts.MaxSteps <= 0 {

opts.MaxSteps = 1

}

if opts.StartupTimeout <= 0 {

opts.StartupTimeout = defaultStartupTimeout

}

if opts.Clock == nil {

opts.Clock = quartz.NewReal()

}

if opts.Metrics == nil {

opts.Metrics = NopMetrics()

}

publishMessagePart := func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) {

if opts.PublishMessagePart == nil {

return

}

opts.PublishMessagePart(role, part)

}

tools := buildToolDefinitions(opts.Tools, opts.ActiveTools, opts.ProviderTools)

applyAnthropicCaching := shouldApplyAnthropicPromptCaching(opts.Model)

messages := opts.Messages

var lastUsage fantasy.Usage

var lastProviderMetadata fantasy.ProviderMetadata

needsFullHistoryReload := false

reloadFullHistory := func(stage string) error {

if opts.ReloadMessages == nil {

return nil

}

reloaded, err := opts.ReloadMessages(ctx)

if err != nil {

return xerrors.Errorf("reload messages %s: %w", stage, err)

}

messages = reloaded

return nil

}

totalSteps := 0

// When totalSteps reaches MaxSteps the inner loop exits immediately

// (its condition is false), stoppedByModel stays false, and the

// post-loop guard breaks the outer compaction loop.

for compactionAttempt := 0; ; compactionAttempt++ {

alreadyCompacted := false

// stoppedByModel is true when the inner step loop

// exited because the model produced no tool calls

// (shouldContinue was false). This distinguishes a

// natural stop from hitting MaxSteps.

stoppedByModel := false

// compactedOnFinalStep tracks whether compaction

// occurred on the very step where the model stopped.

// Only in that case should we re-enter, because the

// agent never had a chance to use the compacted context.

compactedOnFinalStep := false

for step := 0; totalSteps < opts.MaxSteps; step++ {

totalSteps++

provider := opts.Model.Provider()

modelName := opts.Model.Model()

opts.Metrics.StepsTotal.WithLabelValues(provider, modelName).Inc()

stepStart := time.Now()

// 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.

if opts.PrepareMessages != nil {

if updated := opts.PrepareMessages(messages); updated != nil {

messages = updated

}

}

prepared := make([]fantasy.Message, len(messages))

copy(prepared, messages)

prepared, sanitizeStats := chatsanitize.SanitizeAnthropicProviderToolHistory(provider, prepared)

chatsanitize.LogAnthropicProviderToolSanitization(

ctx, opts.Logger, "pre_request", provider, modelName, sanitizeStats,

slog.F("step_index", step),

slog.F("total_steps", totalSteps),

)

prepared = chatsanitize.ApplyAnthropicProviderToolGuard(

ctx, opts.Logger, provider, modelName, prepared,

)

if applyAnthropicCaching {

addAnthropicPromptCaching(prepared)

}

opts.Metrics.MessageCount.WithLabelValues(provider, modelName).Observe(float64(len(prepared)))

opts.Metrics.PromptSizeBytes.WithLabelValues(provider, modelName).Observe(float64(EstimatePromptSize(prepared)))

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

stepCtx := chatdebug.ReuseStep(ctx)

err := chatretry.Retry(stepCtx, func(retryCtx context.Context) error {

attempt, streamErr := guardedStream(

retryCtx,

provider,

modelName,

opts.Clock,

opts.StartupTimeout,

func(attemptCtx context.Context) (fantasy.StreamResponse, error) {

return opts.Model.Stream(attemptCtx, call)

},

opts.Metrics,

)

if streamErr != nil {

return streamErr

}

defer attempt.release()

var processErr error

result, processErr = processStepStream(

attempt.ctx,

attempt.stream,

publishMessagePart,

)

return attempt.finish(processErr)

}, func(

attempt int,

retryErr error,

classified chatretry.ClassifiedError,

delay time.Duration,

) {

// Reset result from the failed attempt so the next

// attempt starts clean.

result = stepResult{}

// Record before OnRetry so a panicking callback can't

// drop the sample. The metric's provider label comes

// from the outer local; WithProvider only affects the

// classified payload handed to OnRetry.

classified = classified.WithProvider(provider)

opts.Metrics.RecordStreamRetry(provider, modelName, classified)

if opts.OnRetry != nil {

opts.OnRetry(attempt, retryErr, classified, 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()

}

// Partition tool calls into built-in and dynamic.

var builtinCalls, dynamicCalls []fantasy.ToolCallContent

if len(opts.DynamicToolNames) > 0 {

for _, tc := range result.toolCalls {

if opts.DynamicToolNames[tc.ToolName] {

dynamicCalls = append(dynamicCalls, tc)

} else {

builtinCalls = append(builtinCalls, tc)

}

}

} else {

builtinCalls = result.toolCalls

}

// Execute only built-in tools.

toolResults = executeTools(ctx, opts.Tools, opts.ActiveTools, opts.ProviderTools, builtinCalls, opts.Metrics, opts.Logger, provider, modelName, opts.BuiltinToolNames, func(tr fantasy.ToolResultContent, completedAt time.Time) {

recordToolResultTimestamp(&result, tr.ToolCallID, completedAt)

publishToolAttachments(ctx, opts.Logger, tr, completedAt, publishMessagePart)

ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr)

ssePart.CreatedAt = &completedAt

publishMessagePart(codersdk.ChatMessageRoleTool, ssePart)

})

for _, tr := range toolResults {

result.content = append(result.content, tr)

}

// If dynamic tools were called, persist what we

// have (assistant + built-in results) and exit so

// the caller can execute them externally.

if len(dynamicCalls) > 0 {

pending := make([]PendingToolCall, 0, len(dynamicCalls))

for _, dc := range dynamicCalls {

pending = append(pending, PendingToolCall{

ToolCallID: dc.ToolCallID,

ToolName: dc.ToolName,

Args: dc.Input,

})

}

contextLimit := extractContextLimitWithFallback(

result.providerMetadata,

opts.ContextLimitFallback,

)

result.content = chatsanitize.SanitizeAnthropicProviderToolStepContent(

ctx, opts.Logger, provider, modelName,

"dynamic_tool_persist", step, result.finishReason, result.content,

)

if len(result.content) == 0 && len(pending) == 0 {

tryCompactOnExit(ctx, opts, result.usage, result.providerMetadata)

return ErrDynamicToolCall

}

if err := opts.PersistStep(ctx, PersistedStep{

Content: result.content,

Usage: result.usage,

ContextLimit: contextLimit,

ProviderResponseID: extractOpenAIResponseIDIfStored(opts.ProviderOptions, result.providerMetadata),

Runtime: time.Since(stepStart),

PendingDynamicToolCalls: pending,

}); err != nil {

if errors.Is(err, ErrInterrupted) {

persistInterruptedStep(ctx, opts, &result)

return ErrInterrupted

}

return xerrors.Errorf("persist step: %w", err)

}

tryCompactOnExit(ctx, opts, result.usage, result.providerMetadata)

return ErrDynamicToolCall

}

// Check for interruption after tool execution.

// Tools that were canceled mid-flight produce error

// results via ctx cancellation. Persist the full

// step (assistant blocks + tool results) through

// the interrupt-safe path so nothing is lost.

if ctx.Err() != nil {

if errors.Is(context.Cause(ctx), ErrInterrupted) {

persistInterruptedStep(ctx, opts, &result)

return ErrInterrupted

}

return ctx.Err()

}

}

// Extract context limit from provider metadata.

contextLimit := extractContextLimitWithFallback(

result.providerMetadata,

opts.ContextLimitFallback,

)

result.content = chatsanitize.SanitizeAnthropicProviderToolStepContent(

ctx, opts.Logger, provider, modelName,

"normal_persist", step, result.finishReason, result.content,

)

if len(result.content) == 0 {

lastUsage = result.usage

lastProviderMetadata = result.providerMetadata

stoppedByModel = true

break

}

// Persist the step. If persistence fails because

// the chat was interrupted between the previous

// check and here, fall back to the interrupt-safe

// path so partial content is not lost.

if err := opts.PersistStep(ctx, PersistedStep{

Content: result.content,

Usage: result.usage,

ContextLimit: contextLimit,

ProviderResponseID: extractOpenAIResponseIDIfStored(opts.ProviderOptions, result.providerMetadata),

Runtime: time.Since(stepStart),

ToolCallCreatedAt: result.toolCallCreatedAt,

ToolResultCreatedAt: result.toolResultCreatedAt,

}); err != nil {

if errors.Is(err, ErrInterrupted) {

persistInterruptedStep(ctx, opts, &result)

return ErrInterrupted

}

return xerrors.Errorf("persist step: %w", err)

}

lastUsage = result.usage

lastProviderMetadata = result.providerMetadata

// Check if any executed tool triggers an early stop.

if shouldStopAfterTools(opts.StopAfterTools, toolResults) {

tryCompactOnExit(ctx, opts, result.usage, result.providerMetadata)

return ErrStopAfterTool

}

// When chain mode is active (PreviousResponseID set), exit

// it after persisting the first chained step. Continuation

// steps include tool-result messages, which fantasy rejects

// when previous_response_id is set, so we must leave chain

// mode and reload the full history before the next call.

stepMessages := result.toResponseMessages()

if hasPreviousResponseID(opts.ProviderOptions) {

clearPreviousResponseID(opts.ProviderOptions)

if opts.DisableChainMode != nil {

opts.DisableChainMode()

}

switch {

case opts.ReloadMessages != nil:

if err := reloadFullHistory("after chain mode exit"); err != nil {

return err

}

needsFullHistoryReload = false

default:

messages = append(messages, stepMessages...)

needsFullHistoryReload = false

}

} else {

messages = append(messages, stepMessages...)

}

if needsFullHistoryReload && !result.shouldContinue &&

opts.ReloadMessages != nil {

if err := reloadFullHistory("before final compaction after chain mode exit"); err != nil {

return err

}

needsFullHistoryReload = false

}

// Inline compaction.

if !needsFullHistoryReload && opts.Compaction != nil && opts.ReloadMessages != nil {

did, compactErr := tryCompact(

ctx,

opts.Model,

opts.Compaction,

opts.ContextLimitFallback,

result.usage,

result.providerMetadata,

messages,

)

opts.Metrics.RecordCompaction(provider, modelName, did, compactErr)

if compactErr != nil && opts.Compaction.OnError != nil {

opts.Compaction.OnError(compactErr)

}

if did {

alreadyCompacted = true

compactedOnFinalStep = true

if err := reloadFullHistory("after compaction"); err != nil {

return err

}

}

}

if !result.shouldContinue {

stoppedByModel = true

break

}

// The agent is continuing with tool calls, so any

// prior compaction has already been consumed.

compactedOnFinalStep = false

}

if needsFullHistoryReload && stoppedByModel && opts.ReloadMessages != nil {

if err := reloadFullHistory("before post-run compaction after chain mode exit"); err != nil {

return err

}

needsFullHistoryReload = false

}

// Post-run compaction safety net: if we never compacted

// during the loop, try once at the end.

if !needsFullHistoryReload && !alreadyCompacted && opts.Compaction != nil && opts.ReloadMessages != nil {

did, err := tryCompact(

ctx,

opts.Model,

opts.Compaction,

opts.ContextLimitFallback,

lastUsage,

lastProviderMetadata,

messages,

)

opts.Metrics.RecordCompaction(opts.Model.Provider(), opts.Model.Model(), did, err)

if err != nil {

if opts.Compaction.OnError != nil {

opts.Compaction.OnError(err)

}

}

if did {

compactedOnFinalStep = true

}

}

// Re-enter the step loop when compaction fired on the

// model's final step. This lets the agent continue

// working with fresh summarized context instead of

// stopping. When the inner loop continued after inline

// compaction (tool-call steps kept going), the agent

// already used the compacted context, so no re-entry

// is needed. Limit retries to prevent infinite loops.

if compactedOnFinalStep && stoppedByModel &&

opts.ReloadMessages != nil &&

compactionAttempt < maxCompactionRetries {

reloaded, reloadErr := opts.ReloadMessages(ctx)

if reloadErr != nil {

return xerrors.Errorf("reload messages after compaction: %w", reloadErr)

}

messages = reloaded

continue

}

break

}

return nil

}

// guardedAttempt owns an attempt-scoped context and startup guard

// around a provider stream. release is idempotent and frees the

// attempt-scoped timer/context. finish canonicalizes startup timeout

// errors before the retry loop classifies them.

type guardedAttempt struct {

ctx context.Context

stream fantasy.StreamResponse

release func()

finish func(error) error

}

// startupGuard arbitrates whether an attempt times out during

// stream startup. Exactly one outcome wins: the timer cancels

// the attempt, or the first-part path disarms the timer.

type startupGuard struct {

timer *quartz.Timer

cancel context.CancelCauseFunc

once sync.Once

}

func newStartupGuard(

clock quartz.Clock,

timeout time.Duration,

cancel context.CancelCauseFunc,

) *startupGuard {

guard := &startupGuard{cancel: cancel}

guard.timer = clock.AfterFunc(timeout, guard.onTimeout, "startupGuard")

return guard

}

func (g *startupGuard) onTimeout() {

g.once.Do(func() {

g.cancel(errStartupTimeout)

})

}

func (g *startupGuard) Disarm() {

g.once.Do(func() {

g.timer.Stop()

})

}

func classifyStartupTimeout(

attemptCtx context.Context,

provider string,

err error,

) error {

if !errors.Is(context.Cause(attemptCtx), errStartupTimeout) {

return err

}

if err == nil {

err = errStartupTimeout

}

return chaterror.WithClassification(err, chaterror.ClassifiedError{

Kind: chaterror.KindStartupTimeout,

Provider: provider,

Retryable: true,

})

}

func guardedStream(

parent context.Context,

provider, model string,

clock quartz.Clock,

timeout time.Duration,

openStream func(context.Context) (fantasy.StreamResponse, error),

metrics *Metrics,

) (guardedAttempt, error) {

attemptCtx, cancelAttempt := context.WithCancelCause(parent)

guard := newStartupGuard(clock, timeout, cancelAttempt)

var releaseOnce sync.Once

release := func() {

releaseOnce.Do(func() {

guard.Disarm()

cancelAttempt(nil)

})

}

streamStart := clock.Now()

stream, err := openStream(attemptCtx)

if err != nil {

err = classifyStartupTimeout(attemptCtx, provider, err)

release()

return guardedAttempt{}, err

}

recordTTFT := sync.OnceFunc(func() {

metrics.TTFTSeconds.WithLabelValues(provider, model).Observe(

clock.Since(streamStart).Seconds(),

)

})

return guardedAttempt{

ctx: attemptCtx,

stream: fantasy.StreamResponse(func(yield func(fantasy.StreamPart) bool) {

for part := range stream {

guard.Disarm()

recordTTFT()

if !yield(part) {

return

}

}

}),

release: release,

finish: func(err error) error {

return classifyStartupTimeout(attemptCtx, provider, err)

},

}, 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(codersdk.ChatMessageRole, 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)

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(codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText(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

}

publishMessagePart(codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageReasoning(part.Delta))

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)

}

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:

var providerExecuted bool

if toolCall, exists := activeToolCalls[part.ID]; exists {

toolCall.Input += part.Delta

providerExecuted = toolCall.ProviderExecuted

}

toolName := toolNames[part.ID]

publishMessagePart(codersdk.ChatMessageRoleAssistant, codersdk.ChatMessagePart{

Type: codersdk.ChatMessagePartTypeToolCall,

ToolCallID: part.ID,

ToolName: toolName,

ArgsDelta: part.Delta,

ProviderExecuted: providerExecuted,

})

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)

// Record when the model emitted this tool call

// so the persisted part carries an accurate

// timestamp for duration computation.

now := dbtime.Now()

if result.toolCallCreatedAt == nil {

result.toolCallCreatedAt = make(map[string]time.Time)

}

result.toolCallCreatedAt[part.ID] = now

ssePart := chatprompt.PartFromContent(tc)

ssePart.CreatedAt = &now

publishMessagePart(

codersdk.ChatMessageRoleAssistant,

ssePart,

)

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(

codersdk.ChatMessageRoleAssistant,

chatprompt.PartFromContent(sourceContent),

)

case fantasy.StreamPartTypeToolResult:

// Provider-executed tool results (e.g. web search)

// are emitted by the provider and added directly

// to the step content for multi-turn round-tripping.

// This mirrors fantasy's agent.go accumulation logic.

if part.ProviderExecuted {

tr := fantasy.ToolResultContent{

ToolCallID: part.ID,

ToolName: part.ToolCallName,

ProviderExecuted: part.ProviderExecuted,

ProviderMetadata: part.ProviderMetadata,

}

result.content = append(result.content, tr)