mirror of
https://github.com/coder/coder.git
synced 2026-11-01 04:47:17 +08:00
## Summary Scope chat model configurations and runtime model resolution to organizations. Each model configuration now belongs to exactly one organization, each organization can have its own default model, and Chatd only uses models that belong to the chat's organization. Depends on #27968. ## Problem Chat model configurations are currently deployment-wide even though chats belong to organizations. A default model, explicitly selected model, or configured model override can therefore reference the same model configuration from every organization. This prevents model catalogs and defaults from being isolated by organization and allows a chat to retain or select a model that its organization does not own. Existing model configurations also have no organization ownership or organization-scoped ACL data, and the database enforces a single live default across the entire deployment. ## Fix Add required organization ownership and ACL storage to chat model configurations. The migration assigns every existing configuration to the default organization, seeds organization-wide read access, and replaces the deployment-wide default constraint with a per-organization constraint. Scope Chatd lookup, admission, fallback, and queue-promotion behavior to the chat's organization: - New chats use the usable default model from their organization. - Explicit model selections from another organization are rejected. - Deployment or personal overrides that point to a foreign model are not used for the chat. - Historical cross-organization model references remain unchanged in the database, but runtime generation falls back to a usable local default. - Queue promotion resolves unavailable, disabled, missing, or foreign stored models to the usable local default and fails when no usable local default exists. The existing model configuration management and availability routes remain compatibility routes for the default organization in this layer. A later layer in the stack adds organization-scoped management routes and updates clients to use them. ## Breaking change This changes chat model behavior from deployment-wide to organization-scoped. Existing configurations are migrated only to the default organization; they are not copied into every organization. As a result: - A non-default organization no longer inherits the deployment's existing default model. - Creating a chat in an organization without a usable local default fails with `400: No default chat model config is configured.` - Explicitly selecting a configuration owned by another organization fails with `400: Model config does not belong to the specified organization.` - Existing chats whose stored model belongs to another organization can generate only if a usable local default is available; otherwise generation fails. Deployments using Agents in non-default organizations must configure a usable default chat model separately for each organization after the organization-scoped management layers in this stack are available. _This pull request description was generated by Coder Agents._
727 lines
23 KiB
Go
727 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
|
|
}
|