mirror of
https://github.com/coder/coder.git
synced 2026-09-23 14:03:57 +08:00
Closes CODAGT-835 ## Summary `chat_messages.runtime_ms` becomes the billing source of truth for Coder Agents runtime (summed hourly by #27312), but it was built for debugging: the June refactor (#26270) silently stopped recording tool-step runtime, compaction was never measured, and interrupted turns lost their partial runtime entirely. This PR defines the billable metric, closes the paths that dropped it, and documents the definition where the data lives. ## The billable definition **`runtime_ms` is the wall-clock duration of the model invocation that produced the persisted message content**, measured from just before the provider stream opens until it is fully consumed. What counts: - Assistant generation steps, in top-level and sub-agent chats (sub-agents are ordinary chats on the same generation path). - Compaction summarization calls, persisted on the compaction assistant message (**new**). - Interrupted attempts: the message-part episode's lifetime is persisted on the partial assistant message committed by `FinishInterruption`, so partial generation time survives interruption (**new**; measured via a new `Buffer.EpisodeDuration`, which works even though the generation goroutine and the interrupt task are different tasks). What deliberately does not count (each is documented in code and docs): - **Local tool execution.** Tool wall time includes idle waits, most importantly `wait_agent` polling a sub-agent chat that already bills its own model invocations; billing the batch would double count, and excluding one tool from a concurrent batch's wall time is ill-defined. Pre-refactor instrumentation did include tool time; this makes the exclusion an explicit product definition instead of a silent regression. - **Failed model calls whose output is discarded** (retried attempts, terminal errors, content-filter refusals). They persist no content, so they bill nothing; billing errs toward undercounting. Notably a stream-silence timeout can burn 10 idle minutes before a retry, which should not be billable "active generation". If product later wants failed attempts billed, that needs a place to persist runtime on error turns (`FinishError` inserts no rows today) and is a deliberate follow-up, not instrumentation drift. - **Ancillary calls that produce no chat messages** (title generation, advisor, turn summaries) and all idle/parked time (`requires_action`, queueing). The definition is documented as `COMMENT ON COLUMN chat_messages.runtime_ms` (migration 000551, surfacing as a Go doc comment on `ChatMessage.RuntimeMs`), on `chatloop.PersistedStep.Runtime`, in the chatd architecture doc, and in the Spend Management docs page. ## Index for the hourly scan None needed: `GetTotalChatMessageRuntimeMsInRange` (#27312) filters an hour-wide `created_at` range, which the existing `idx_chat_messages_created_at` b-tree already serves; the residual `runtime_ms IS NOT NULL` filter applies to one hour of rows. A partial index would add permanent write amplification for a query that runs once an hour. > [!NOTE] > Migration 000551 is also claimed by #27312; whichever merges second renumbers via `fix_migration_numbers.sh`. ## Tests - End-to-end: the existing full-server generation test now asserts `RuntimeMs.Valid` on the committed assistant row (it previously read `.Int64` without checking `.Valid`, so it passed on NULL). - Interrupted turn: full task-level test (real DB, mock clock) asserting the partial assistant message persists the attempt's runtime. - Errored stream: asserts a failed invocation yields no step and no runtime. - Tool-using turn: asserts runtime lands on the assistant row only and tool rows stay NULL. - Compaction: asserts the summarization call duration is recorded and lands on the compaction assistant message only. - `messagepartbuffer.EpisodeDuration` unit coverage. Blocks: CODAGT-843 (B3), CODAGT-838 (D8). --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Hugo Dutka <hugo@coder.com>
725 lines
23 KiB
Go
725 lines
23 KiB
Go
package chatd
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"golang.org/x/xerrors"
|
|
|
|
"cdr.dev/slog/v3"
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatdebug"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
"github.com/coder/quartz"
|
|
)
|
|
|
|
const (
|
|
postCommitWatchPublishTimeout = 10 * time.Second
|
|
// defaultTaskTimeout must exceed chatloop's stream-silence guard so
|
|
// silent provider streams fail through chat-specific retry handling
|
|
// before the runner retries the whole task.
|
|
defaultTaskTimeout = 15 * time.Minute
|
|
)
|
|
|
|
var (
|
|
errTaskExpectedExit = xerrors.New("chatworker task expected exit")
|
|
errTaskRetryable = xerrors.New("chatworker task retryable error")
|
|
errTaskTimeout = xerrors.New("chatworker task timeout")
|
|
)
|
|
|
|
type taskRetryableError struct {
|
|
err error
|
|
}
|
|
|
|
func (e taskRetryableError) Error() string {
|
|
if e.err == nil {
|
|
return errTaskRetryable.Error()
|
|
}
|
|
return e.err.Error()
|
|
}
|
|
|
|
func (e taskRetryableError) Unwrap() error {
|
|
if e.err == nil {
|
|
return errTaskRetryable
|
|
}
|
|
return errors.Join(errTaskRetryable, e.err)
|
|
}
|
|
|
|
type retryWrapperOptions struct {
|
|
clock quartz.Clock
|
|
logger slog.Logger
|
|
initialDelay time.Duration
|
|
maxDelay time.Duration
|
|
}
|
|
|
|
type retryWrapperTaskInfo struct {
|
|
ChatID uuid.UUID
|
|
WorkerID uuid.UUID
|
|
RunnerID uuid.UUID
|
|
}
|
|
|
|
// runTaskWithRetry ensures that a task doesn't exit until it completes
|
|
// successfully or gets canceled. It retries the task in case of any ephemeral errors.
|
|
// It's critical for the correct operation of the chat runner:
|
|
// this function is THE place that ensures task liveness within the runner.
|
|
func runTaskWithRetry(
|
|
ctx context.Context,
|
|
opts retryWrapperOptions,
|
|
kind taskKind,
|
|
info retryWrapperTaskInfo,
|
|
fn func(context.Context) error,
|
|
) error {
|
|
if opts.clock == nil {
|
|
opts.clock = quartz.NewReal()
|
|
}
|
|
if opts.initialDelay <= 0 {
|
|
opts.initialDelay = defaultTaskRetryInitialBackoff
|
|
}
|
|
if opts.maxDelay <= 0 {
|
|
opts.maxDelay = defaultTaskRetryMaxBackoff
|
|
}
|
|
if opts.maxDelay < opts.initialDelay {
|
|
opts.maxDelay = opts.initialDelay
|
|
}
|
|
delay := opts.initialDelay
|
|
for {
|
|
attemptCtx, cancelAttempt := taskAttemptContext(ctx, opts.clock, kind)
|
|
err := executeTaskSafely(attemptCtx, fn)
|
|
timedOut := errors.Is(context.Cause(attemptCtx), errTaskTimeout)
|
|
cancelAttempt()
|
|
if timedOut && err != nil {
|
|
if !errors.Is(err, errTaskExpectedExit) ||
|
|
errors.Is(err, context.Canceled) ||
|
|
errors.Is(err, context.DeadlineExceeded) ||
|
|
errors.Is(err, errTaskTimeout) {
|
|
err = taskRetryableError{err: errors.Join(errTaskTimeout, err)}
|
|
}
|
|
}
|
|
if err == nil {
|
|
// no log on success to avoid noise
|
|
return nil
|
|
}
|
|
|
|
exitReason := ""
|
|
switch {
|
|
case ctx.Err() != nil:
|
|
exitReason = "context_canceled"
|
|
case errors.Is(err, errTaskExpectedExit) && !errors.Is(err, errTaskRetryable):
|
|
exitReason = "expected_non_retryable_exit"
|
|
}
|
|
if exitReason != "" {
|
|
opts.logger.Debug(ctx, "chatworker task exited",
|
|
slog.F("task_kind", kind),
|
|
slog.F("reason", exitReason),
|
|
slog.F("chat_id", info.ChatID),
|
|
slog.F("worker_id", info.WorkerID),
|
|
slog.F("runner_id", info.RunnerID),
|
|
slogError(err),
|
|
)
|
|
return nil
|
|
}
|
|
|
|
opts.logger.Warn(ctx, "chatworker task retrying",
|
|
slog.F("task_kind", kind),
|
|
slog.F("delay", delay),
|
|
slog.F("chat_id", info.ChatID),
|
|
slog.F("worker_id", info.WorkerID),
|
|
slog.F("runner_id", info.RunnerID),
|
|
slogError(err),
|
|
)
|
|
timer := opts.clock.NewTimer(delay, "chatworker", "task-retry-"+string(kind))
|
|
select {
|
|
case <-timer.C:
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return nil
|
|
}
|
|
timer.Stop()
|
|
if delay < opts.maxDelay {
|
|
delay *= 2
|
|
if delay > opts.maxDelay {
|
|
delay = opts.maxDelay
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func taskAttemptContext(ctx context.Context, clock quartz.Clock, kind taskKind) (context.Context, func()) {
|
|
attemptCtx, cancelCause := context.WithCancelCause(ctx)
|
|
timer := clock.AfterFunc(defaultTaskTimeout, func() {
|
|
cancelCause(errTaskTimeout)
|
|
}, "chatworker", "task-timeout-"+string(kind))
|
|
return attemptCtx, func() {
|
|
timer.Stop()
|
|
cancelCause(nil)
|
|
}
|
|
}
|
|
|
|
func executeTaskSafely(ctx context.Context, fn func(context.Context) error) (err error) {
|
|
defer func() {
|
|
if recovered := recover(); recovered != nil {
|
|
err = xerrors.Errorf("chatworker task panic: %v", recovered)
|
|
}
|
|
}()
|
|
return fn(ctx)
|
|
}
|
|
|
|
type interruptionOutcome struct {
|
|
Chat database.Chat
|
|
Kind runnerActionKind
|
|
WatchEventKind codersdk.ChatWatchEventKind
|
|
}
|
|
|
|
type taskStarter struct {
|
|
server *Server
|
|
opts chatWorkerOptions
|
|
routeStateHint func(context.Context, runnerStateUpdate)
|
|
requestCleanup func(context.Context, runnerKey)
|
|
afterInterruptionOutcome func(context.Context, interruptionOutcome) error
|
|
}
|
|
|
|
func newTaskStarter(
|
|
server *Server,
|
|
opts chatWorkerOptions,
|
|
routeStateHint func(context.Context, runnerStateUpdate),
|
|
requestCleanup func(context.Context, runnerKey),
|
|
) (*taskStarter, error) {
|
|
if server == nil {
|
|
return nil, xerrors.New("chatworker: server is required")
|
|
}
|
|
if opts.Store == nil {
|
|
return nil, xerrors.New("chatworker: task store is required")
|
|
}
|
|
if opts.Pubsub == nil {
|
|
return nil, xerrors.New("chatworker: task pubsub is required")
|
|
}
|
|
if opts.MessagePartBuffer == nil {
|
|
return nil, xerrors.New("chatworker: message part buffer is required")
|
|
}
|
|
if opts.Clock == nil {
|
|
opts.Clock = quartz.NewReal()
|
|
}
|
|
if opts.TaskRetryInitialBackoff <= 0 {
|
|
opts.TaskRetryInitialBackoff = defaultTaskRetryInitialBackoff
|
|
}
|
|
if opts.TaskRetryMaxBackoff <= 0 {
|
|
opts.TaskRetryMaxBackoff = defaultTaskRetryMaxBackoff
|
|
}
|
|
if opts.TaskRetryMaxBackoff < opts.TaskRetryInitialBackoff {
|
|
opts.TaskRetryMaxBackoff = opts.TaskRetryInitialBackoff
|
|
}
|
|
if routeStateHint == nil {
|
|
return nil, xerrors.New("chatworker: route state hint callback is required")
|
|
}
|
|
if requestCleanup == nil {
|
|
return nil, xerrors.New("chatworker: cleanup callback is required")
|
|
}
|
|
return &taskStarter{
|
|
server: server,
|
|
opts: opts,
|
|
routeStateHint: routeStateHint,
|
|
requestCleanup: requestCleanup,
|
|
}, nil
|
|
}
|
|
|
|
func (o chatWorkerOptions) retryOptions() retryWrapperOptions {
|
|
return retryWrapperOptions{
|
|
clock: o.Clock,
|
|
logger: o.Logger,
|
|
initialDelay: o.TaskRetryInitialBackoff,
|
|
maxDelay: o.TaskRetryMaxBackoff,
|
|
}
|
|
}
|
|
|
|
func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskStartInput) error {
|
|
machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
|
|
var chat database.Chat
|
|
err := machine.ReadLock(ctx, func(store database.Store) error {
|
|
loadedChat, err := loadChatForTask(ctx, store, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
|
|
if err != nil {
|
|
return xerrors.Errorf("load chat for task: %w", err)
|
|
}
|
|
chat = loadedChat
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return normalizeTaskInfrastructureError(err, "lock chat for interrupt")
|
|
}
|
|
|
|
key := messagepartbuffer.Key{
|
|
ChatID: input.ChatID,
|
|
HistoryVersion: input.HistoryVersion,
|
|
GenerationAttempt: chat.GenerationAttempt,
|
|
}
|
|
modelInvokedAt := s.opts.MessagePartBuffer.ModelInvokedAt(key)
|
|
if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil {
|
|
if ctx.Err() != nil {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("close message part episode: %w", err), ctx.Err())
|
|
}
|
|
return taskRetryableError{err: xerrors.Errorf("close message part episode: %w", err)}
|
|
}
|
|
parts, err := s.opts.MessagePartBuffer.GetParts(key)
|
|
if errors.Is(err, messagepartbuffer.ErrEpisodeNotFound) {
|
|
parts = nil
|
|
err = nil
|
|
}
|
|
if err != nil {
|
|
if ctx.Err() != nil {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("get message part episode: %w", err), ctx.Err())
|
|
}
|
|
return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)}
|
|
}
|
|
interruptedAt := s.opts.Clock.Now("chatworker", "interrupt")
|
|
var attemptRuntime time.Duration
|
|
if !modelInvokedAt.IsZero() {
|
|
attemptRuntime = interruptedAt.Sub(modelInvokedAt)
|
|
}
|
|
partialMessages, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{
|
|
parts: parts,
|
|
modelConfigID: chat.LastModelConfigID,
|
|
contentVersion: chatprompt.CurrentContentVersion,
|
|
logger: s.opts.Logger,
|
|
interruptedAt: interruptedAt,
|
|
attemptRuntime: attemptRuntime,
|
|
})
|
|
if err != nil {
|
|
return xerrors.Errorf("convert buffered parts: %w", err)
|
|
}
|
|
|
|
var committed database.Chat
|
|
err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
chat, err := loadChatForTask(ctx, store, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true})
|
|
if err != nil {
|
|
return xerrors.Errorf("load chat for task: %w", err)
|
|
}
|
|
messages := partialMessages
|
|
committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, chat, s.opts.Clock.Now("chatworker", "interrupt"))
|
|
if err != nil {
|
|
return xerrors.Errorf("committed pending local tool cancellation messages: %w", err)
|
|
}
|
|
if len(committedCancels) > 0 {
|
|
messages = append(append([]chatstate.Message{}, partialMessages...), committedCancels...)
|
|
}
|
|
if _, err := tx.FinishInterruption(chatstate.FinishInterruptionInput{PartialMessages: messages}); err != nil {
|
|
return xerrors.Errorf("finish interruption: %w", err)
|
|
}
|
|
committed, err = store.GetChatByID(ctx, input.ChatID)
|
|
if err != nil {
|
|
return xerrors.Errorf("load committed chat: %w", err)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
if current, ok := s.committedStateAfterUpdateError(ctx, committed); ok {
|
|
return s.publishWatchAndRoute(ctx, current, codersdk.ChatWatchEventKindStatusChange)
|
|
}
|
|
return normalizeTaskTransitionError(err, "finish interruption")
|
|
}
|
|
input.DebugTurn.RecordOutcome(chatdebug.StatusInterrupted)
|
|
if err := s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindStatusChange); err != nil {
|
|
return xerrors.Errorf("publish watch and route: %w", err)
|
|
}
|
|
return s.runAfterInterruptionOutcome(ctx, interruptionOutcome{
|
|
Chat: committed,
|
|
Kind: runnerActionKindFinishInterruption,
|
|
WatchEventKind: codersdk.ChatWatchEventKindStatusChange,
|
|
})
|
|
}
|
|
|
|
func (s *taskStarter) runAfterInterruptionOutcome(ctx context.Context, outcome interruptionOutcome) error {
|
|
afterOutcome := s.afterInterruptionOutcome
|
|
if afterOutcome == nil {
|
|
afterOutcome = s.server.afterInterruptionOutcome
|
|
}
|
|
if afterOutcome == nil {
|
|
return nil
|
|
}
|
|
if err := afterOutcome(ctx, outcome); err != nil {
|
|
return taskRetryableError{err: xerrors.Errorf("interruption post-outcome side effects: %w", err)}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *taskStarter) StartRequiresActionTimeout(ctx context.Context, input chatWorkerTaskStartInput) error {
|
|
machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
|
|
for {
|
|
decision, err := decideRequiresActionTimeout(ctx, machine, input)
|
|
if err != nil {
|
|
return xerrors.Errorf("decide requires action timeout: %w", err)
|
|
}
|
|
if decision.cancel {
|
|
return s.cancelRequiresAction(ctx, machine, input, decision.reason)
|
|
}
|
|
if !decision.waitUntil.Valid {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("requires action deadline is missing"))
|
|
}
|
|
if err := s.waitUntil(ctx, decision.waitUntil.Time); err != nil {
|
|
return xerrors.Errorf("wait until: %w", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
type requiresActionTimeoutDecision struct {
|
|
cancel bool
|
|
reason string
|
|
waitUntil sql.NullTime
|
|
}
|
|
|
|
func decideRequiresActionTimeout(
|
|
ctx context.Context,
|
|
machine *chatstate.ChatMachine,
|
|
input chatWorkerTaskStartInput,
|
|
) (requiresActionTimeoutDecision, error) {
|
|
var decision requiresActionTimeoutDecision
|
|
err := machine.ReadLock(ctx, func(store database.Store) error {
|
|
chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
|
|
if err != nil {
|
|
return xerrors.Errorf("load chat for task: %w", err)
|
|
}
|
|
if !chat.RequiresActionDeadlineAt.Valid {
|
|
decision.cancel = true
|
|
decision.reason = "Tool execution canceled because the action deadline was missing"
|
|
return nil
|
|
}
|
|
now, err := store.GetDatabaseNow(ctx)
|
|
if err != nil {
|
|
return xerrors.Errorf("get database time: %w", err)
|
|
}
|
|
if now.Before(chat.RequiresActionDeadlineAt.Time) {
|
|
decision.waitUntil = chat.RequiresActionDeadlineAt
|
|
return nil
|
|
}
|
|
decision.cancel = true
|
|
decision.reason = "Tool execution timed out"
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return requiresActionTimeoutDecision{}, normalizeTaskInfrastructureError(err, "lock chat for requires action timeout")
|
|
}
|
|
return decision, nil
|
|
}
|
|
|
|
func (s *taskStarter) waitUntil(ctx context.Context, deadline time.Time) error {
|
|
now := s.opts.Clock.Now("chatworker", "requires-action-timeout")
|
|
if !now.Before(deadline) {
|
|
return nil
|
|
}
|
|
timer := s.opts.Clock.NewTimer(deadline.Sub(now), "chatworker", "requires-action-timeout")
|
|
defer timer.Stop()
|
|
select {
|
|
case <-timer.C:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("wait until: %w", ctx.Err()))
|
|
}
|
|
}
|
|
|
|
func (s *taskStarter) cancelRequiresAction(
|
|
ctx context.Context,
|
|
machine *chatstate.ChatMachine,
|
|
input chatWorkerTaskStartInput,
|
|
reason string,
|
|
) error {
|
|
var committed database.Chat
|
|
err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
chat, err := loadChatForTask(ctx, store, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true})
|
|
if err != nil {
|
|
return xerrors.Errorf("load chat for task: %w", err)
|
|
}
|
|
if chat.RequiresActionDeadlineAt.Valid {
|
|
now, err := store.GetDatabaseNow(ctx)
|
|
if err != nil {
|
|
return xerrors.Errorf("get database time: %w", err)
|
|
}
|
|
if now.Before(chat.RequiresActionDeadlineAt.Time) {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("requires action deadline is in the future"))
|
|
}
|
|
}
|
|
if _, err := tx.CancelRequiresAction(chatstate.CancelRequiresActionInput{Reason: reason}); err != nil {
|
|
return xerrors.Errorf("cancel requires action: %w", err)
|
|
}
|
|
committed, err = store.GetChatByID(ctx, input.ChatID)
|
|
if err != nil {
|
|
return xerrors.Errorf("load committed chat: %w", err)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
if current, ok := s.committedStateAfterUpdateError(ctx, committed); ok {
|
|
return s.publishWatchAndRoute(ctx, current, codersdk.ChatWatchEventKindStatusChange)
|
|
}
|
|
return normalizeTaskTransitionError(err, "cancel requires action")
|
|
}
|
|
return s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindStatusChange)
|
|
}
|
|
|
|
func (s *taskStarter) StartAbandon(ctx context.Context, input chatWorkerTaskStartInput) error {
|
|
machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID)
|
|
mismatch := false
|
|
err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
chat, err := store.GetChatByID(ctx, input.ChatID)
|
|
if err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
mismatch = true
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
|
|
}
|
|
return xerrors.Errorf("load chat: %w", err)
|
|
}
|
|
if !ownedByTask(chat, input) {
|
|
mismatch = true
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat not owned by task"))
|
|
}
|
|
if err := verifyTaskFence(chat, input, input.Status, taskFenceOptions{requireHistory: true, allowArchived: true}); err != nil {
|
|
return xerrors.Errorf("task fence mismatch: %w", err)
|
|
}
|
|
if _, err := tx.Abandon(chatstate.AbandonInput{}); err != nil {
|
|
return xerrors.Errorf("abandon chat: %w", err)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
if errors.Is(err, errTaskExpectedExit) && mismatch {
|
|
s.requestCleanup(ctx, runnerKey{ChatID: input.ChatID, RunnerID: input.RunnerID})
|
|
return nil
|
|
}
|
|
return normalizeTaskTransitionError(err, "abandon chat")
|
|
}
|
|
s.requestCleanup(ctx, runnerKey{ChatID: input.ChatID, RunnerID: input.RunnerID})
|
|
return nil
|
|
}
|
|
|
|
func (s *taskStarter) committedStateAfterUpdateError(ctx context.Context, committed database.Chat) (database.Chat, bool) {
|
|
if committed.ID == uuid.Nil {
|
|
return database.Chat{}, false
|
|
}
|
|
current, err := s.opts.Store.GetChatByID(ctx, committed.ID)
|
|
if err != nil {
|
|
return database.Chat{}, false
|
|
}
|
|
if current.SnapshotVersion != committed.SnapshotVersion ||
|
|
current.HistoryVersion != committed.HistoryVersion ||
|
|
current.QueueVersion != committed.QueueVersion ||
|
|
current.GenerationAttempt != committed.GenerationAttempt ||
|
|
current.Status != committed.Status ||
|
|
current.Archived != committed.Archived ||
|
|
current.WorkerID != committed.WorkerID ||
|
|
current.RunnerID != committed.RunnerID {
|
|
return database.Chat{}, false
|
|
}
|
|
return current, true
|
|
}
|
|
|
|
func (s *taskStarter) publishWatchAndRoute(
|
|
ctx context.Context,
|
|
chat database.Chat,
|
|
kind codersdk.ChatWatchEventKind,
|
|
) error {
|
|
watchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postCommitWatchPublishTimeout)
|
|
defer cancel()
|
|
if err := s.publishWatchWithRetry(watchCtx, chat, kind); err != nil {
|
|
return xerrors.Errorf("publish watch with retry: %w", err)
|
|
}
|
|
s.routeStateHint(ctx, stateUpdateFromChat(chat))
|
|
return nil
|
|
}
|
|
|
|
func (s *taskStarter) publishWatchWithRetry(
|
|
ctx context.Context,
|
|
chat database.Chat,
|
|
kind codersdk.ChatWatchEventKind,
|
|
) error {
|
|
delay := s.opts.TaskRetryInitialBackoff
|
|
for {
|
|
if err := publishChatWatchEvent(s.opts.Pubsub, chat, kind); err == nil {
|
|
return nil
|
|
} else if ctx.Err() != nil {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("publishChatWatchEvent: %w", ctx.Err()))
|
|
}
|
|
timer := s.opts.Clock.NewTimer(delay, "chatworker", "watch-publish-retry")
|
|
select {
|
|
case <-timer.C:
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("watch publish retry context done: %w", ctx.Err()))
|
|
}
|
|
timer.Stop()
|
|
if delay < s.opts.TaskRetryMaxBackoff {
|
|
delay *= 2
|
|
if delay > s.opts.TaskRetryMaxBackoff {
|
|
delay = s.opts.TaskRetryMaxBackoff
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func publishChatWatchEvent(pubsub chatWorkerPubsub, chat database.Chat, kind codersdk.ChatWatchEventKind) error {
|
|
event := codersdk.ChatWatchEvent{
|
|
Kind: kind,
|
|
Chat: chatWatchEventSDKChat(chat, nil),
|
|
}
|
|
payload, err := json.Marshal(event)
|
|
if err != nil {
|
|
return xerrors.Errorf("marshal chat watch event: %w", err)
|
|
}
|
|
if err := pubsub.Publish(coderdpubsub.ChatWatchEventChannel(chat.OwnerID), payload); err != nil {
|
|
return xerrors.Errorf("publish chat watch event: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type taskFenceOptions struct {
|
|
requireHistory bool
|
|
allowArchived bool
|
|
}
|
|
|
|
// loadChatForTask loads the chat row and verifies the task fence in one
|
|
// step so call sites cannot skip the fence check. It returns an error
|
|
// wrapping errTaskExpectedExit when the chat no longer exists or the fence
|
|
// no longer matches.
|
|
func loadChatForTask(
|
|
ctx context.Context,
|
|
store database.Store,
|
|
input chatWorkerTaskStartInput,
|
|
status database.ChatStatus,
|
|
opts taskFenceOptions,
|
|
) (database.Chat, error) {
|
|
chat, err := store.GetChatByID(ctx, input.ChatID)
|
|
if err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return database.Chat{}, errors.Join(errTaskExpectedExit, xerrors.Errorf("load chat: %w", err))
|
|
}
|
|
return database.Chat{}, xerrors.Errorf("load chat: %w", err)
|
|
}
|
|
if err := verifyTaskFence(chat, input, status, opts); err != nil {
|
|
return database.Chat{}, xerrors.Errorf("verifyTaskFence: %w", err)
|
|
}
|
|
return chat, nil
|
|
}
|
|
|
|
func verifyTaskFence(
|
|
chat database.Chat,
|
|
input chatWorkerTaskStartInput,
|
|
status database.ChatStatus,
|
|
opts taskFenceOptions,
|
|
) error {
|
|
if !ownedByTask(chat, input) {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat not owned by task"))
|
|
}
|
|
if chat.Status != status {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat status mismatch: %s != %s", chat.Status, status))
|
|
}
|
|
if !opts.allowArchived && chat.Archived {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat archived"))
|
|
}
|
|
if opts.requireHistory && chat.HistoryVersion != input.HistoryVersion {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("chat history version mismatch: %d != %d", chat.HistoryVersion, input.HistoryVersion))
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func ownedByTask(chat database.Chat, input chatWorkerTaskStartInput) bool {
|
|
return chat.WorkerID.Valid && chat.WorkerID.UUID == input.WorkerID &&
|
|
chat.RunnerID.Valid && chat.RunnerID.UUID == input.RunnerID
|
|
}
|
|
|
|
func normalizeTaskInfrastructureError(err error, action string) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
if errors.Is(err, errTaskExpectedExit) {
|
|
return err
|
|
}
|
|
if errors.Is(err, chatstate.ErrChatNotFound) || errors.Is(err, sql.ErrNoRows) || errors.Is(err, context.Canceled) {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("%s: %w", action, err))
|
|
}
|
|
return taskRetryableError{err: xerrors.Errorf("%s: %w", action, err)}
|
|
}
|
|
|
|
func normalizeTaskTransitionError(err error, action string) error {
|
|
if err == nil {
|
|
return nil
|
|
}
|
|
if errors.Is(err, errTaskExpectedExit) {
|
|
return err
|
|
}
|
|
if errors.Is(err, chatstate.ErrChatNotFound) || errors.Is(err, sql.ErrNoRows) || errors.Is(err, context.Canceled) {
|
|
return errors.Join(errTaskExpectedExit, xerrors.Errorf("%s: %w", action, err))
|
|
}
|
|
if errors.Is(err, chatstate.ErrTransitionNotAllowed) || errors.Is(err, chatstate.ErrInvalidState) {
|
|
return xerrors.Errorf("%s: %w", action, err)
|
|
}
|
|
return taskRetryableError{err: xerrors.Errorf("%s: %w", action, err)}
|
|
}
|
|
|
|
func dynamicToolNamesFromChat(chat database.Chat) map[string]bool {
|
|
if !chat.DynamicTools.Valid || len(chat.DynamicTools.RawMessage) == 0 {
|
|
return nil
|
|
}
|
|
var tools []codersdk.DynamicTool
|
|
if err := json.Unmarshal(chat.DynamicTools.RawMessage, &tools); err != nil {
|
|
return nil
|
|
}
|
|
names := make(map[string]bool, len(tools))
|
|
for _, tool := range tools {
|
|
name := strings.TrimSpace(tool.Name)
|
|
if name != "" {
|
|
names[name] = true
|
|
}
|
|
}
|
|
return names
|
|
}
|
|
|
|
func committedPendingLocalToolCancellationMessages(
|
|
ctx context.Context,
|
|
store database.Store,
|
|
chat database.Chat,
|
|
interruptedAt time.Time,
|
|
) ([]chatstate.Message, error) {
|
|
messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
|
|
ChatID: chat.ID,
|
|
AfterID: 0,
|
|
})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("load committed messages for interruption: %w", err)
|
|
}
|
|
localCalls, _, err := unresolvedToolCallsFromHistory(messages, dynamicToolNamesFromChat(chat))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(localCalls) == 0 {
|
|
return nil, nil
|
|
}
|
|
result := make([]chatstate.Message, 0, len(localCalls))
|
|
for _, call := range localCalls {
|
|
payload, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("marshal interrupted tool result: %w", err)
|
|
}
|
|
part := codersdk.ChatMessageToolResult(call.ToolCallID, call.ToolName, payload, true, false)
|
|
if !interruptedAt.IsZero() {
|
|
part.CreatedAt = &interruptedAt
|
|
}
|
|
content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part})
|
|
if err != nil {
|
|
return nil, xerrors.Errorf("marshal interrupted tool result part: %w", err)
|
|
}
|
|
result = append(result, chatstate.Message{
|
|
Role: database.ChatMessageRoleTool,
|
|
Content: content,
|
|
Visibility: database.ChatMessageVisibilityBoth,
|
|
ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: chat.LastModelConfigID != uuid.Nil},
|
|
ContentVersion: chatprompt.CurrentContentVersion,
|
|
})
|
|
}
|
|
return result, nil
|
|
}
|