diff --git a/cli/server.go b/cli/server.go index 70db62e6b0..1dbdc5a152 100644 --- a/cli/server.go +++ b/cli/server.go @@ -1158,7 +1158,7 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd. defer shutdownConns() // Ensures that old database entries are cleaned up over time! - purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, options.DeploymentValues, options.PrometheusRegistry, &coderAPI.Auditor, dbpurge.WithNotificationsEnqueuer(options.NotificationsEnqueuer)) + purger := dbpurge.New(ctx, logger.Named("dbpurge"), options.Database, options.DeploymentValues, options.PrometheusRegistry, &coderAPI.Auditor) defer purger.Close() // Updates workspace usage diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 6ed25d6046..55661ec6f2 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -818,6 +818,42 @@ const docTemplate = `{ ] } }, + "/api/experimental/chats/{chat}/reconcile-invalid": { + "post": { + "description": "Experimental: this endpoint is subject to change.", + "produces": [ + "application/json" + ], + "tags": [ + "Chats" + ], + "summary": "Reconcile invalid chat state", + "operationId": "reconcile-invalid-chat-state", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.Chat" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/experimental/chats/{chat}/stream": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -923,6 +959,45 @@ const docTemplate = `{ ] } }, + "/api/experimental/chats/{chat}/stream/parts": { + "get": { + "description": "Experimental: this endpoint is subject to change.", + "produces": [ + "application/json" + ], + "tags": [ + "Chats" + ], + "summary": "Stream chat parts via WebSockets", + "operationId": "stream-chat-parts-via-websockets", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.ChatStreamEvent" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + } + }, "/api/experimental/chats/{chat}/title/regenerate": { "post": { "description": "Experimental: this endpoint is subject to change.", @@ -17087,7 +17162,8 @@ const docTemplate = `{ "paused", "completed", "error", - "requires_action" + "requires_action", + "interrupting" ], "x-enum-varnames": [ "ChatStatusWaiting", @@ -17096,7 +17172,8 @@ const docTemplate = `{ "ChatStatusPaused", "ChatStatusCompleted", "ChatStatusError", - "ChatStatusRequiresAction" + "ChatStatusRequiresAction", + "ChatStatusInterrupting" ] }, "codersdk.ChatStreamActionRequired": { @@ -17155,7 +17232,9 @@ const docTemplate = `{ "error", "queue_update", "retry", - "action_required" + "action_required", + "preview_reset", + "history_reset" ], "x-enum-varnames": [ "ChatStreamEventTypeMessagePart", @@ -17164,17 +17243,28 @@ const docTemplate = `{ "ChatStreamEventTypeError", "ChatStreamEventTypeQueueUpdate", "ChatStreamEventTypeRetry", - "ChatStreamEventTypeActionRequired" + "ChatStreamEventTypeActionRequired", + "ChatStreamEventTypePreviewReset", + "ChatStreamEventTypeHistoryReset" ] }, "codersdk.ChatStreamMessagePart": { "type": "object", "properties": { + "generation_attempt": { + "type": "integer" + }, + "history_version": { + "type": "integer" + }, "part": { "$ref": "#/definitions/codersdk.ChatMessagePart" }, "role": { "$ref": "#/definitions/codersdk.ChatMessageRole" + }, + "seq": { + "type": "integer" } } }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 4440fe155b..8a542d8024 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -723,6 +723,38 @@ ] } }, + "/api/experimental/chats/{chat}/reconcile-invalid": { + "post": { + "description": "Experimental: this endpoint is subject to change.", + "produces": ["application/json"], + "tags": ["Chats"], + "summary": "Reconcile invalid chat state", + "operationId": "reconcile-invalid-chat-state", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.Chat" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/api/experimental/chats/{chat}/stream": { "get": { "description": "Experimental: this endpoint is subject to change.", @@ -816,6 +848,41 @@ ] } }, + "/api/experimental/chats/{chat}/stream/parts": { + "get": { + "description": "Experimental: this endpoint is subject to change.", + "produces": ["application/json"], + "tags": ["Chats"], + "summary": "Stream chat parts via WebSockets", + "operationId": "stream-chat-parts-via-websockets", + "parameters": [ + { + "type": "string", + "format": "uuid", + "description": "Chat ID", + "name": "chat", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.ChatStreamEvent" + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ], + "x-apidocgen": { + "skip": true + } + } + }, "/api/experimental/chats/{chat}/title/regenerate": { "post": { "description": "Experimental: this endpoint is subject to change.", @@ -15392,7 +15459,8 @@ "paused", "completed", "error", - "requires_action" + "requires_action", + "interrupting" ], "x-enum-varnames": [ "ChatStatusWaiting", @@ -15401,7 +15469,8 @@ "ChatStatusPaused", "ChatStatusCompleted", "ChatStatusError", - "ChatStatusRequiresAction" + "ChatStatusRequiresAction", + "ChatStatusInterrupting" ] }, "codersdk.ChatStreamActionRequired": { @@ -15460,7 +15529,9 @@ "error", "queue_update", "retry", - "action_required" + "action_required", + "preview_reset", + "history_reset" ], "x-enum-varnames": [ "ChatStreamEventTypeMessagePart", @@ -15469,17 +15540,28 @@ "ChatStreamEventTypeError", "ChatStreamEventTypeQueueUpdate", "ChatStreamEventTypeRetry", - "ChatStreamEventTypeActionRequired" + "ChatStreamEventTypeActionRequired", + "ChatStreamEventTypePreviewReset", + "ChatStreamEventTypeHistoryReset" ] }, "codersdk.ChatStreamMessagePart": { "type": "object", "properties": { + "generation_attempt": { + "type": "integer" + }, + "history_version": { + "type": "integer" + }, "part": { "$ref": "#/definitions/codersdk.ChatMessagePart" }, "role": { "$ref": "#/definitions/codersdk.ChatMessageRole" + }, + "seq": { + "type": "integer" } } }, diff --git a/coderd/coderd.go b/coderd/coderd.go index c5e5a12793..0605e57829 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -251,9 +251,10 @@ type Options struct { SSHConfig codersdk.SSHConfigResponse HTTPClient *http.Client - // ChatSubscribeFn provides cross-replica subscription merging. - // Set by enterprise for HA deployments. Nil in AGPL single-replica. - ChatSubscribeFn chatd.SubscribeFn + // ChatStreamPartsDialer dials remote chat stream parts. + // Set by enterprise for HA deployments. Nil uses chatd's local + // in-process channel dialer. + ChatStreamPartsDialer chatd.StreamPartsDialer // ChatProviderAPIKeys overrides deployment-derived provider keys. // Test harnesses use this to route chat models to local providers. ChatProviderAPIKeys *chatprovider.ProviderAPIKeys @@ -814,11 +815,11 @@ func New(options *Options) *API { chatAIGatewayRoutingEnabled := options.DeploymentValues.AI.BridgeConfig.Enabled.Value() && options.DeploymentValues.AI.Chat.AIGatewayRoutingEnabled.Value() - api.chatDaemon = chatd.New(chatd.Config{ + api.chatDaemon = chatd.New(options.Pubsub, chatd.Config{ Logger: options.Logger.Named("chatd"), Database: options.Database, ReplicaID: api.ID, - SubscribeFn: options.ChatSubscribeFn, + StreamPartsDialer: options.ChatStreamPartsDialer, MaxChatsPerAcquire: int32(maxChatsPerAcquire), //nolint:gosec // maxChatsPerAcquire is clamped to int32 range above. ProviderAPIKeys: providerAPIKeys, AllowBYOK: options.DeploymentValues.AI.BridgeConfig.AllowBYOK.Value(), @@ -832,11 +833,12 @@ func New(options *Options) *API { CreateWorkspace: api.chatCreateWorkspace, StartWorkspace: api.chatStartWorkspace, StopWorkspace: api.chatStopWorkspace, - Pubsub: options.Pubsub, WebpushDispatcher: options.WebPushDispatcher, UsageTracker: options.WorkspaceUsageTracker, PrometheusRegistry: options.PrometheusRegistry, OIDCTokenSource: oidcMCPSrc, + NotificationsEnqueuer: options.NotificationsEnqueuer, + Auditor: &api.Auditor, }).Start() gitSyncLogger := options.Logger.Named("gitsync") refresher := gitsync.NewRefresher( @@ -1342,10 +1344,12 @@ func New(options *Options) *API { r.Get("/prompts", api.getChatUserPrompts) r.Route("/stream", func(r chi.Router) { r.Get("/", api.streamChat) + r.Get("/parts", api.streamChatParts) r.Get("/desktop", api.watchChatDesktop) r.Get("/git", api.watchChatGit) }) r.Post("/interrupt", api.interruptChat) + r.Post("/reconcile-invalid", api.reconcileInvalidChatState) r.Post("/tool-results", api.postChatToolResults) r.Post("/title/regenerate", api.regenerateChatTitle) r.Post("/title/propose", api.proposeChatTitle) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 4bdbdcf6cf..d27bcbc518 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -1692,9 +1692,6 @@ func (q *querier) ArchiveUnusedTemplateVersions(ctx context.Context, arg databas } func (q *querier) AutoArchiveInactiveChats(ctx context.Context, arg database.AutoArchiveInactiveChatsParams) ([]database.AutoArchiveInactiveChatsRow, error) { - // Background write by dbpurge. The LATERAL read of chat_messages rows - // happens below the RBAC boundary; only the chat row itself requires - // authorization. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { return nil, err } @@ -1718,6 +1715,13 @@ func (q *querier) BackoffChatDiffStatus(ctx context.Context, arg database.Backof return q.db.BackoffChatDiffStatus(ctx, arg) } +func (q *querier) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return 0, err + } + return q.db.BatchDeleteChatHeartbeats(ctx, arg) +} + func (q *querier) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error { // Could be any workspace agent and checking auth to each workspace agent is overkill for // the purpose of this function. @@ -1743,6 +1747,13 @@ func (q *querier) BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg datab return q.db.BatchUpdateWorkspaceNextStartAt(ctx, arg) } +func (q *querier) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return err + } + return q.db.BatchUpsertChatHeartbeats(ctx, arg) +} + func (q *querier) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil { return err @@ -1857,6 +1868,14 @@ func (q *querier) CountAuditLogs(ctx context.Context, arg database.CountAuditLog return q.db.CountAuthorizedAuditLogs(ctx, arg, prep) } +func (q *querier) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { + _, err := q.GetChatByID(ctx, chatID) + if err != nil { + return 0, err + } + return q.db.CountChatQueuedMessages(ctx, chatID) +} + func (q *querier) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) { // Just like the actual query, shortcut if the user is an owner. err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceConnectionLog) @@ -1954,6 +1973,18 @@ func (q *querier) DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) e return q.db.DeleteAPIKeysByUserID(ctx, userID) } +func (q *querier) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error { + chat, err := q.db.GetChatByID(ctx, chatID) + if err != nil { + return err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return err + } + _ = chat + return q.db.DeleteAllChatHeartbeats(ctx, chatID) +} + func (q *querier) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error { chat, err := q.db.GetChatByID(ctx, chatID) if err != nil { @@ -1965,6 +1996,18 @@ func (q *querier) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.U return q.db.DeleteAllChatQueuedMessages(ctx, chatID) } +func (q *querier) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { + chat, err := q.db.GetChatByID(ctx, chatID) + if err != nil { + return 0, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return 0, err + } + _ = chat + return q.db.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID) +} + func (q *querier) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil { return nil, err @@ -2043,6 +2086,18 @@ func (q *querier) DeleteChatQueuedMessage(ctx context.Context, arg database.Dele return q.db.DeleteChatQueuedMessage(ctx, arg) } +func (q *querier) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return 0, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return 0, err + } + _ = chat + return q.db.DeleteChatQueuedMessageReturningCount(ctx, arg) +} + func (q *querier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err @@ -2315,6 +2370,13 @@ func (q *querier) DeleteRuntimeConfig(ctx context.Context, key string) error { return q.db.DeleteRuntimeConfig(ctx, key) } +func (q *querier) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return 0, err + } + return q.db.DeleteStaleChatHeartbeats(ctx, staleSeconds) +} + func (q *querier) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceTailnetCoordinator); err != nil { return database.DeleteTailnetPeerRow{}, err @@ -2826,6 +2888,13 @@ func (q *querier) GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUI return q.db.GetAuthorizationUserRoles(ctx, userID) } +func (q *querier) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return nil, err + } + return q.db.GetAutoArchiveInactiveChatCandidates(ctx, arg) +} + func (q *querier) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceBoundaryLog); err != nil { return database.BoundaryLog{}, err @@ -2877,6 +2946,10 @@ func (q *querier) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat, return fetch(q.log, q.auth, q.db.GetChatByID)(ctx, id) } +func (q *querier) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) { + return fetch(q.log, q.auth, q.db.GetChatByIDForShare)(ctx, id) +} + func (q *querier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) { return fetch(q.log, q.auth, q.db.GetChatByIDForUpdate)(ctx, id) } @@ -3048,6 +3121,18 @@ func (q *querier) GetChatExploreModelOverride(ctx context.Context) (string, erro return q.db.GetChatExploreModelOverride(ctx) } +func (q *querier) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + // This is a read-only query: it returns the chat IDs that belong + // to a family. Authorize as Read against the root chat. The + // individual SetArchived (or other) transitions that consume + // these IDs run their own per-row authorization, so we do not + // gate the listing itself on Update permission. + if _, err := q.GetChatByID(ctx, id); err != nil { + return nil, err + } + return q.db.GetChatFamilyIDsByRootID(ctx, id) +} + func (q *querier) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) { file, err := q.db.GetChatFileByID(ctx, id) if err != nil { @@ -3114,6 +3199,14 @@ func (q *querier) GetChatGeneralModelOverride(ctx context.Context) (string, erro return q.db.GetChatGeneralModelOverride(ctx) } +func (q *querier) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return database.ChatHeartbeat{}, err + } + return q.db.GetChatHeartbeat(ctx, arg) +} + func (q *querier) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) { // The include-default-system-prompt flag is a deployment-wide setting read // during chat creation by every authenticated user, so no RBAC policy @@ -3174,6 +3267,14 @@ func (q *querier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg return q.db.GetChatMessagesByChatIDDescPaginated(ctx, arg) } +func (q *querier) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return nil, err + } + return q.db.GetChatMessagesByRevisionForStream(ctx, arg) +} + func (q *querier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { // Authorize read on the parent chat. _, err := q.GetChatByID(ctx, chatID) @@ -3222,6 +3323,22 @@ func (q *querier) GetChatPlanModeInstructions(ctx context.Context) (string, erro return q.db.GetChatPlanModeInstructions(ctx) } +func (q *querier) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return database.ChatQueuedMessage{}, err + } + return q.db.GetChatQueuedMessageByID(ctx, arg) +} + +func (q *querier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { + _, err := q.GetChatByID(ctx, chatID) + if err != nil { + return database.ChatQueuedMessage{}, err + } + return q.db.GetChatQueuedMessageHead(ctx, chatID) +} + func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { _, err := q.GetChatByID(ctx, chatID) if err != nil { @@ -3230,6 +3347,14 @@ func (q *querier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ( return q.db.GetChatQueuedMessages(ctx, chatID) } +func (q *querier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { + _, err := q.GetChatByID(ctx, chatID) + if err != nil { + return nil, err + } + return q.db.GetChatQueuedMessagesByPosition(ctx, chatID) +} + func (q *querier) GetChatRetentionDays(ctx context.Context) (int32, error) { // Chat retention is a deployment-wide config read by dbpurge. // Only requires a valid actor in context. @@ -3239,6 +3364,13 @@ func (q *querier) GetChatRetentionDays(ctx context.Context) (int32, error) { return q.db.GetChatRetentionDays(ctx) } +func (q *querier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat); err != nil { + return nil, err + } + return q.db.GetChatStreamSyncRows(ctx, ids) +} + func (q *querier) GetChatSystemPrompt(ctx context.Context) (string, error) { // The system prompt is a deployment-wide setting read during chat // creation by every authenticated user, so no RBAC policy check @@ -3311,6 +3443,13 @@ func (q *querier) GetChatUserPromptsByChatID(ctx context.Context, arg database.G return q.db.GetChatUserPromptsByChatID(ctx, arg) } +func (q *querier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return nil, err + } + return q.db.GetChatWorkerAcquisitionCandidates(ctx, arg) +} + func (q *querier) GetChatWorkspaceTTL(ctx context.Context) (string, error) { // The workspace-TTL setting is a deployment-wide value read by any // authenticated chat user. We only require that an explicit actor is @@ -3333,6 +3472,13 @@ func (q *querier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) ([ return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByChatFileID)(ctx, fileID) } +func (q *querier) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { + return nil, err + } + return q.db.GetChatsByIDsForRunnerSync(ctx, ids) +} + func (q *querier) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByWorkspaceIDs)(ctx, ids) } @@ -3403,6 +3549,10 @@ func (q *querier) GetDERPMeshKey(ctx context.Context) (string, error) { return q.db.GetDERPMeshKey(ctx) } +func (q *querier) GetDatabaseNow(ctx context.Context) (time.Time, error) { + return q.db.GetDatabaseNow(ctx) +} + func (q *querier) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { // Reading the default model config is needed for chat creation. // TODO(CODAGT-161): scope this check when org context is available. @@ -5468,6 +5618,18 @@ func (q *querier) GetWorkspacesForWorkspaceMetrics(ctx context.Context) ([]datab return q.db.GetWorkspacesForWorkspaceMetrics(ctx) } +func (q *querier) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) { + chat, err := q.db.GetChatByID(ctx, id) + if err != nil { + return 0, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return 0, err + } + _ = chat + return q.db.IncrementChatGenerationAttempt(ctx, id) +} + func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg) } @@ -5638,6 +5800,18 @@ func (q *querier) InsertChatQueuedMessage(ctx context.Context, arg database.Inse return q.db.InsertChatQueuedMessage(ctx, arg) } +func (q *querier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return database.ChatQueuedMessage{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.ChatQueuedMessage{}, err + } + _ = chat + return q.db.InsertChatQueuedMessageWithCreator(ctx, arg) +} + func (q *querier) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) { if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceCryptoKey); err != nil { return database.CryptoKey{}, err @@ -6210,6 +6384,14 @@ func (q *querier) InsertWorkspaceResourceMetadata(ctx context.Context, arg datab return q.db.InsertWorkspaceResourceMetadata(ctx, arg) } +func (q *querier) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) { + _, err := q.GetChatByID(ctx, arg.ChatID) + if err != nil { + return false, err + } + return q.db.IsChatHeartbeatStale(ctx, arg) +} + func (q *querier) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) { chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { @@ -6394,6 +6576,18 @@ func (q *querier) ListWorkspaceAgentPortShares(ctx context.Context, workspaceID return q.db.ListWorkspaceAgentPortShares(ctx, workspaceID) } +func (q *querier) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { + chat, err := q.db.GetChatByID(ctx, id) + if err != nil { + return database.Chat{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.Chat{}, err + } + _ = chat + return q.db.LockChatAndBumpSnapshotVersion(ctx, id) +} + func (q *querier) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error { resource := rbac.ResourceInboxNotification.WithOwner(arg.UserID.String()) @@ -6500,6 +6694,18 @@ func (q *querier) ReorderChatQueuedMessageToFront(ctx context.Context, arg datab return q.db.ReorderChatQueuedMessageToFront(ctx, arg) } +func (q *querier) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return 0, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return 0, err + } + _ = chat + return q.db.ReorderChatQueuedMessageToHead(ctx, arg) +} + func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil { return database.ResolveUserChatSpendLimitRow{}, err @@ -6742,6 +6948,18 @@ func (q *querier) UpdateChatDebugStep(ctx context.Context, arg database.UpdateCh return q.db.UpdateChatDebugStep(ctx, arg) } +func (q *querier) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) { + chat, err := q.db.GetChatByID(ctx, arg.ID) + if err != nil { + return database.Chat{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.Chat{}, err + } + _ = chat + return q.db.UpdateChatExecutionState(ctx, arg) +} + func (q *querier) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { // The batch heartbeat is a system-level operation filtered by // worker_id. Authorization is enforced by the AsChatd context @@ -6864,6 +7082,19 @@ func (q *querier) UpdateChatPlanModeByID(ctx context.Context, arg database.Updat return q.db.UpdateChatPlanModeByID(ctx, arg) } +func (q *querier) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) { + // UpdateChatRetryState is used by the chat processor to publish + // transient retry state. It should be called with system context. + chat, err := q.db.GetChatByID(ctx, arg.ID) + if err != nil { + return database.Chat{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.Chat{}, err + } + return q.db.UpdateChatRetryState(ctx, arg) +} + func (q *querier) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) { // UpdateChatStatus is used by the chat processor to change chat status. // It should be called with system context. @@ -8243,6 +8474,18 @@ func (q *querier) UpsertChatGeneralModelOverride(ctx context.Context, value stri return q.db.UpsertChatGeneralModelOverride(ctx, value) } +func (q *querier) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return err + } + _ = chat + return q.db.UpsertChatHeartbeat(ctx, arg) +} + func (q *querier) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 4cbd1a4fdc..e9c7a8dad2 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -537,6 +537,21 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().AcquireChats(gomock.Any(), arg).Return([]database.Chat{chat}, nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat}) })) + s.Run("GetChatWorkerAcquisitionCandidates", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + arg := database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: 30, + LimitCount: 100, + } + row := testutil.Fake(s.T(), faker, database.GetChatWorkerAcquisitionCandidatesRow{}) + dbm.EXPECT().GetChatWorkerAcquisitionCandidates(gomock.Any(), arg).Return([]database.GetChatWorkerAcquisitionCandidatesRow{row}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.GetChatWorkerAcquisitionCandidatesRow{row}) + })) + s.Run("GetChatsByIDsForRunnerSync", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{uuid.New(), uuid.New()} + chat := testutil.Fake(s.T(), faker, database.Chat{ID: ids[0]}) + dbm.EXPECT().GetChatsByIDsForRunnerSync(gomock.Any(), ids).Return([]database.Chat{chat}, nil).AnyTimes() + check.Args(ids).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.Chat{chat}) + })) s.Run("DeleteAllChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() @@ -763,6 +778,24 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().GetChatByIDForUpdate(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(chat) })) + s.Run("GetChatByIDForShare", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByIDForShare(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(chat) + })) + s.Run("GetChatStreamSyncRows", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{uuid.New(), uuid.New()} + rows := []database.GetChatStreamSyncRowsRow{{ID: ids[0]}} + dbm.EXPECT().GetChatStreamSyncRows(gomock.Any(), ids).Return(rows, nil).AnyTimes() + check.Args(ids).Asserts(rbac.ResourceChat, policy.ActionRead).Returns(rows) + })) + s.Run("GetChatFamilyIDsByRootID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + ids := []uuid.UUID{chat.ID} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatFamilyIDsByRootID(gomock.Any(), chat.ID).Return(ids, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(ids) + })) s.Run("GetChatsByWorkspaceIDs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chatA := testutil.Fake(s.T(), faker, database.Chat{}) chatB := testutil.Fake(s.T(), faker, database.Chat{}) @@ -950,9 +983,10 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().UpsertChatAutoArchiveDays(gomock.Any(), int32(90)).Return(nil).AnyTimes() check.Args(int32(90)).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate) })) - s.Run("AutoArchiveInactiveChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - dbm.EXPECT().AutoArchiveInactiveChats(gomock.Any(), database.AutoArchiveInactiveChatsParams{}).Return([]database.AutoArchiveInactiveChatsRow{}, nil).AnyTimes() - check.Args(database.AutoArchiveInactiveChatsParams{}).Asserts(rbac.ResourceChat, policy.ActionUpdate) + s.Run("GetAutoArchiveInactiveChatCandidates", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.GetAutoArchiveInactiveChatCandidatesParams{LimitCount: 100} + dbm.EXPECT().GetAutoArchiveInactiveChatCandidates(gomock.Any(), arg).Return([]database.GetAutoArchiveInactiveChatCandidatesRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.GetAutoArchiveInactiveChatCandidatesRow{}) })) s.Run("GetChatMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) @@ -985,6 +1019,14 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().GetChatMessagesByChatIDDescPaginated(gomock.Any(), arg).Return(msgs, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) })) + s.Run("GetChatMessagesByRevisionForStream", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + msgs := []database.ChatMessage{testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID})} + arg := database.GetChatMessagesByRevisionForStreamParams{ChatID: chat.ID, AfterRevision: 1} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatMessagesByRevisionForStream(gomock.Any(), arg).Return(msgs, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(msgs) + })) s.Run("GetChatUserPromptsByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) rows := []database.GetChatUserPromptsByChatIDRow{{ID: 1, Text: "hello"}} @@ -1214,6 +1256,136 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().UpdateChatACLByID(gomock.Any(), arg).Return(nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionShare).Returns() })) + s.Run("LockChatAndBumpSnapshotVersion", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().LockChatAndBumpSnapshotVersion(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(chat) + })) + s.Run("UpdateChatExecutionState", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.UpdateChatExecutionStateParams{ID: chat.ID, Status: database.ChatStatusRunning} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().UpdateChatExecutionState(gomock.Any(), arg).Return(chat, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat) + })) + s.Run("IncrementChatGenerationAttempt", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().IncrementChatGenerationAttempt(gomock.Any(), chat.ID).Return(int64(7), nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(int64(7)) + })) + s.Run("UpdateChatRetryState", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.UpdateChatRetryStateParams{ID: chat.ID, RetryState: []byte(`{"attempt":1}`)} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().UpdateChatRetryState(gomock.Any(), arg).Return(chat, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(chat) + })) + s.Run("GetDatabaseNow", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + now := time.Now() + dbm.EXPECT().GetDatabaseNow(gomock.Any()).Return(now, nil).AnyTimes() + check.Args().Asserts().Returns(now) + })) + s.Run("InsertChatQueuedMessageWithCreator", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := testutil.Fake(s.T(), faker, database.InsertChatQueuedMessageWithCreatorParams{ChatID: chat.ID}) + qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().InsertChatQueuedMessageWithCreator(gomock.Any(), arg).Return(qm, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(qm) + })) + s.Run("GetChatQueuedMessagesByPosition", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + qms := []database.ChatQueuedMessage{} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatQueuedMessagesByPosition(gomock.Any(), chat.ID).Return(qms, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qms) + })) + s.Run("CountChatQueuedMessages", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().CountChatQueuedMessages(gomock.Any(), chat.ID).Return(int64(3), nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(int64(3)) + })) + s.Run("GetChatQueuedMessageHead", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatQueuedMessageHead(gomock.Any(), chat.ID).Return(qm, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(qm) + })) + s.Run("GetChatQueuedMessageByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + qm := testutil.Fake(s.T(), faker, database.ChatQueuedMessage{ChatID: chat.ID}) + arg := database.GetChatQueuedMessageByIDParams{ID: qm.ID, ChatID: chat.ID} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatQueuedMessageByID(gomock.Any(), arg).Return(qm, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(qm) + })) + s.Run("DeleteChatQueuedMessageReturningCount", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.DeleteChatQueuedMessageReturningCountParams{ID: 1, ChatID: chat.ID} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().DeleteChatQueuedMessageReturningCount(gomock.Any(), arg).Return(int64(1), nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1)) + })) + s.Run("DeleteAllChatQueuedMessagesReturningCount", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().DeleteAllChatQueuedMessagesReturningCount(gomock.Any(), chat.ID).Return(int64(1), nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns(int64(1)) + })) + s.Run("ReorderChatQueuedMessageToHead", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.ReorderChatQueuedMessageToHeadParams{ChatID: chat.ID, ID: 1} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().ReorderChatQueuedMessageToHead(gomock.Any(), arg).Return(int64(1), nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(int64(1)) + })) + s.Run("UpsertChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.UpsertChatHeartbeatParams{ChatID: chat.ID, RunnerID: uuid.New()} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().UpsertChatHeartbeat(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns() + })) + s.Run("BatchUpsertChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.BatchUpsertChatHeartbeatsParams{ChatIds: []uuid.UUID{uuid.New()}, RunnerIds: []uuid.UUID{uuid.New()}} + dbm.EXPECT().BatchUpsertChatHeartbeats(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns() + })) + s.Run("GetChatHeartbeat", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.GetChatHeartbeatParams{ChatID: chat.ID, RunnerID: uuid.New()} + hb := database.ChatHeartbeat{ChatID: chat.ID, RunnerID: arg.RunnerID} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatHeartbeat(gomock.Any(), arg).Return(hb, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(hb) + })) + s.Run("IsChatHeartbeatStale", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + arg := database.IsChatHeartbeatStaleParams{ChatID: chat.ID, RunnerID: uuid.New(), StaleSeconds: 30} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().IsChatHeartbeatStale(gomock.Any(), arg).Return(false, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(false) + })) + s.Run("DeleteAllChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().DeleteAllChatHeartbeats(gomock.Any(), chat.ID).Return(nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionUpdate).Returns() + })) + s.Run("BatchDeleteChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.BatchDeleteChatHeartbeatsParams{ChatIds: []uuid.UUID{uuid.New()}, RunnerIds: []uuid.UUID{uuid.New()}} + dbm.EXPECT().BatchDeleteChatHeartbeats(gomock.Any(), arg).Return(int64(1), nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(int64(1)) + })) + s.Run("DeleteStaleChatHeartbeats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + const staleSeconds int32 = 30 + dbm.EXPECT().DeleteStaleChatHeartbeats(gomock.Any(), staleSeconds).Return(int64(1), nil).AnyTimes() + check.Args(staleSeconds).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(int64(1)) + })) s.Run("UpdateChatByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) arg := database.UpdateChatByIDParams{ @@ -1408,6 +1580,14 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().BackoffChatDiffStatus(gomock.Any(), arg).Return(nil).AnyTimes() check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns() })) + s.Run("AutoArchiveInactiveChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + arg := database.AutoArchiveInactiveChatsParams{ + ArchiveCutoff: dbtime.Now(), + LimitCount: 100, + } + dbm.EXPECT().AutoArchiveInactiveChats(gomock.Any(), arg).Return([]database.AutoArchiveInactiveChatsRow{}, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns([]database.AutoArchiveInactiveChatsRow{}) + })) s.Run("UpsertChatIncludeDefaultSystemPrompt", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { dbm.EXPECT().UpsertChatIncludeDefaultSystemPrompt(gomock.Any(), false).Return(nil).AnyTimes() check.Args(false).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate) @@ -1697,9 +1877,9 @@ func (s *MethodTestSuite) TestChats() { s.Run("UpdateChatLastTurnSummary", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) arg := database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true}, } dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() dbm.EXPECT().UpdateChatLastTurnSummary(gomock.Any(), arg).Return(int64(1), nil).AnyTimes() diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 9d9e12f118..3955220efe 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -115,13 +115,23 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat if seed.Content.Valid { content = string(seed.Content.RawMessage) } + role := takeFirst(seed.Role, database.ChatMessageRoleUser) + apiKeyID := seed.APIKeyID.String + // Mint a real API key for user turns so the api_key_id foreign key is + // satisfied. Without a creator we leave it empty, which the insert query + // stores as NULL. + if role == database.ChatMessageRoleUser && apiKeyID == "" && + seed.CreatedBy.Valid && seed.CreatedBy.UUID != uuid.Nil { + key, _ := APIKey(t, db, database.APIKey{UserID: seed.CreatedBy.UUID}) + apiKeyID = key.ID + } msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ ChatID: seed.ChatID, CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, - APIKeyID: []string{seed.APIKeyID.String}, + APIKeyID: []string{apiKeyID}, ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, - Role: []database.ChatMessageRole{takeFirst(seed.Role, database.ChatMessageRoleUser)}, + Role: []database.ChatMessageRole{role}, Content: []string{content}, ContentVersion: []int16{takeFirst(seed.ContentVersion, chatprompt.CurrentContentVersion)}, Visibility: []database.ChatMessageVisibility{takeFirst(seed.Visibility, database.ChatMessageVisibilityBoth)}, diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 6c873b3841..6460cc6c4f 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -202,6 +202,14 @@ func (m queryMetricsStore) BackoffChatDiffStatus(ctx context.Context, arg databa return r0 } +func (m queryMetricsStore) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.BatchDeleteChatHeartbeats(ctx, arg) + m.queryLatencies.WithLabelValues("BatchDeleteChatHeartbeats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BatchDeleteChatHeartbeats").Inc() + return r0, r1 +} + func (m queryMetricsStore) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error { start := time.Now() r0 := m.s.BatchUpdateWorkspaceAgentMetadata(ctx, arg) @@ -226,6 +234,14 @@ func (m queryMetricsStore) BatchUpdateWorkspaceNextStartAt(ctx context.Context, return r0 } +func (m queryMetricsStore) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error { + start := time.Now() + r0 := m.s.BatchUpsertChatHeartbeats(ctx, arg) + m.queryLatencies.WithLabelValues("BatchUpsertChatHeartbeats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "BatchUpsertChatHeartbeats").Inc() + return r0 +} + func (m queryMetricsStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error { start := time.Now() r0 := m.s.BatchUpsertConnectionLogs(ctx, arg) @@ -322,6 +338,14 @@ func (m queryMetricsStore) CountAuditLogs(ctx context.Context, arg database.Coun return r0, r1 } +func (m queryMetricsStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { + start := time.Now() + r0, r1 := m.s.CountChatQueuedMessages(ctx, chatID) + m.queryLatencies.WithLabelValues("CountChatQueuedMessages").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "CountChatQueuedMessages").Inc() + return r0, r1 +} + func (m queryMetricsStore) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) { start := time.Now() r0, r1 := m.s.CountConnectionLogs(ctx, arg) @@ -418,6 +442,14 @@ func (m queryMetricsStore) DeleteAPIKeysByUserID(ctx context.Context, userID uui return r0 } +func (m queryMetricsStore) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error { + start := time.Now() + r0 := m.s.DeleteAllChatHeartbeats(ctx, chatID) + m.queryLatencies.WithLabelValues("DeleteAllChatHeartbeats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAllChatHeartbeats").Inc() + return r0 +} + func (m queryMetricsStore) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error { start := time.Now() r0 := m.s.DeleteAllChatQueuedMessages(ctx, chatID) @@ -426,6 +458,14 @@ func (m queryMetricsStore) DeleteAllChatQueuedMessages(ctx context.Context, chat return r0 } +func (m queryMetricsStore) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { + start := time.Now() + r0, r1 := m.s.DeleteAllChatQueuedMessagesReturningCount(ctx, chatID) + m.queryLatencies.WithLabelValues("DeleteAllChatQueuedMessagesReturningCount").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteAllChatQueuedMessagesReturningCount").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) { start := time.Now() r0, r1 := m.s.DeleteAllTailnetTunnels(ctx, arg) @@ -498,6 +538,14 @@ func (m queryMetricsStore) DeleteChatQueuedMessage(ctx context.Context, arg data return r0 } +func (m queryMetricsStore) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.DeleteChatQueuedMessageReturningCount(ctx, arg) + m.queryLatencies.WithLabelValues("DeleteChatQueuedMessageReturningCount").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatQueuedMessageReturningCount").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error { start := time.Now() r0 := m.s.DeleteChatUsageLimitGroupOverride(ctx, groupID) @@ -778,6 +826,14 @@ func (m queryMetricsStore) DeleteRuntimeConfig(ctx context.Context, key string) return r0 } +func (m queryMetricsStore) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) { + start := time.Now() + r0, r1 := m.s.DeleteStaleChatHeartbeats(ctx, staleSeconds) + m.queryLatencies.WithLabelValues("DeleteStaleChatHeartbeats").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteStaleChatHeartbeats").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) { start := time.Now() r0, r1 := m.s.DeleteTailnetPeer(ctx, arg) @@ -1274,6 +1330,14 @@ func (m queryMetricsStore) GetAuthorizationUserRoles(ctx context.Context, userID return r0, r1 } +func (m queryMetricsStore) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) { + start := time.Now() + r0, r1 := m.s.GetAutoArchiveInactiveChatCandidates(ctx, arg) + m.queryLatencies.WithLabelValues("GetAutoArchiveInactiveChatCandidates").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAutoArchiveInactiveChatCandidates").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) { start := time.Now() r0, r1 := m.s.GetBoundaryLogByID(ctx, id) @@ -1322,6 +1386,14 @@ func (m queryMetricsStore) GetChatByID(ctx context.Context, id uuid.UUID) (datab return r0, r1 } +func (m queryMetricsStore) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) { + start := time.Now() + r0, r1 := m.s.GetChatByIDForShare(ctx, id) + m.queryLatencies.WithLabelValues("GetChatByIDForShare").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatByIDForShare").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) { start := time.Now() r0, r1 := m.s.GetChatByIDForUpdate(ctx, id) @@ -1450,6 +1522,14 @@ func (m queryMetricsStore) GetChatExploreModelOverride(ctx context.Context) (str return r0, r1 } +func (m queryMetricsStore) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + start := time.Now() + r0, r1 := m.s.GetChatFamilyIDsByRootID(ctx, id) + m.queryLatencies.WithLabelValues("GetChatFamilyIDsByRootID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatFamilyIDsByRootID").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) { start := time.Now() r0, r1 := m.s.GetChatFileByID(ctx, id) @@ -1482,6 +1562,14 @@ func (m queryMetricsStore) GetChatGeneralModelOverride(ctx context.Context) (str return r0, r1 } +func (m queryMetricsStore) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) { + start := time.Now() + r0, r1 := m.s.GetChatHeartbeat(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatHeartbeat").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatHeartbeat").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) { start := time.Now() r0, r1 := m.s.GetChatIncludeDefaultSystemPrompt(ctx) @@ -1530,6 +1618,14 @@ func (m queryMetricsStore) GetChatMessagesByChatIDDescPaginated(ctx context.Cont return r0, r1 } +func (m queryMetricsStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatMessagesByRevisionForStream(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatMessagesByRevisionForStream").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatMessagesByRevisionForStream").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { start := time.Now() r0, r1 := m.s.GetChatMessagesForPromptByChatID(ctx, chatID) @@ -1578,6 +1674,22 @@ func (m queryMetricsStore) GetChatPlanModeInstructions(ctx context.Context) (str return r0, r1 } +func (m queryMetricsStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatQueuedMessageByID(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatQueuedMessageByID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessageByID").Inc() + return r0, r1 +} + +func (m queryMetricsStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatQueuedMessageHead(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatQueuedMessageHead").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessageHead").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { start := time.Now() r0, r1 := m.s.GetChatQueuedMessages(ctx, chatID) @@ -1586,6 +1698,14 @@ func (m queryMetricsStore) GetChatQueuedMessages(ctx context.Context, chatID uui return r0, r1 } +func (m queryMetricsStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { + start := time.Now() + r0, r1 := m.s.GetChatQueuedMessagesByPosition(ctx, chatID) + m.queryLatencies.WithLabelValues("GetChatQueuedMessagesByPosition").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatQueuedMessagesByPosition").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatRetentionDays(ctx context.Context) (int32, error) { start := time.Now() r0, r1 := m.s.GetChatRetentionDays(ctx) @@ -1594,6 +1714,14 @@ func (m queryMetricsStore) GetChatRetentionDays(ctx context.Context) (int32, err return r0, r1 } +func (m queryMetricsStore) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) { + start := time.Now() + r0, r1 := m.s.GetChatStreamSyncRows(ctx, ids) + m.queryLatencies.WithLabelValues("GetChatStreamSyncRows").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatStreamSyncRows").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatSystemPrompt(ctx context.Context) (string, error) { start := time.Now() r0, r1 := m.s.GetChatSystemPrompt(ctx) @@ -1658,6 +1786,14 @@ func (m queryMetricsStore) GetChatUserPromptsByChatID(ctx context.Context, arg d return r0, r1 } +func (m queryMetricsStore) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) { + start := time.Now() + r0, r1 := m.s.GetChatWorkerAcquisitionCandidates(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatWorkerAcquisitionCandidates").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatWorkerAcquisitionCandidates").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatWorkspaceTTL(ctx context.Context) (string, error) { start := time.Now() r0, r1 := m.s.GetChatWorkspaceTTL(ctx) @@ -1682,6 +1818,14 @@ func (m queryMetricsStore) GetChatsByChatFileID(ctx context.Context, fileID uuid return r0, r1 } +func (m queryMetricsStore) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { + start := time.Now() + r0, r1 := m.s.GetChatsByIDsForRunnerSync(ctx, ids) + m.queryLatencies.WithLabelValues("GetChatsByIDsForRunnerSync").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatsByIDsForRunnerSync").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { start := time.Now() r0, r1 := m.s.GetChatsByWorkspaceIDs(ctx, ids) @@ -1754,6 +1898,14 @@ func (m queryMetricsStore) GetDERPMeshKey(ctx context.Context) (string, error) { return r0, r1 } +func (m queryMetricsStore) GetDatabaseNow(ctx context.Context) (time.Time, error) { + start := time.Now() + r0, r1 := m.s.GetDatabaseNow(ctx) + m.queryLatencies.WithLabelValues("GetDatabaseNow").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetDatabaseNow").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { start := time.Now() r0, r1 := m.s.GetDefaultChatModelConfig(ctx) @@ -3706,6 +3858,14 @@ func (m queryMetricsStore) GetWorkspacesForWorkspaceMetrics(ctx context.Context) return r0, r1 } +func (m queryMetricsStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) { + start := time.Now() + r0, r1 := m.s.IncrementChatGenerationAttempt(ctx, id) + m.queryLatencies.WithLabelValues("IncrementChatGenerationAttempt").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "IncrementChatGenerationAttempt").Inc() + return r0, r1 +} + func (m queryMetricsStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { start := time.Now() r0, r1 := m.s.InsertAIBridgeInterception(ctx, arg) @@ -3866,6 +4026,14 @@ func (m queryMetricsStore) InsertChatQueuedMessage(ctx context.Context, arg data return r0, r1 } +func (m queryMetricsStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { + start := time.Now() + r0, r1 := m.s.InsertChatQueuedMessageWithCreator(ctx, arg) + m.queryLatencies.WithLabelValues("InsertChatQueuedMessageWithCreator").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatQueuedMessageWithCreator").Inc() + return r0, r1 +} + func (m queryMetricsStore) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) { start := time.Now() r0, r1 := m.s.InsertCryptoKey(ctx, arg) @@ -4362,6 +4530,14 @@ func (m queryMetricsStore) InsertWorkspaceResourceMetadata(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) { + start := time.Now() + r0, r1 := m.s.IsChatHeartbeatStale(ctx, arg) + m.queryLatencies.WithLabelValues("IsChatHeartbeatStale").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "IsChatHeartbeatStale").Inc() + return r0, r1 +} + func (m queryMetricsStore) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) { start := time.Now() r0, r1 := m.s.LinkChatFiles(ctx, arg) @@ -4546,6 +4722,14 @@ func (m queryMetricsStore) ListWorkspaceAgentPortShares(ctx context.Context, wor return r0, r1 } +func (m queryMetricsStore) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { + start := time.Now() + r0, r1 := m.s.LockChatAndBumpSnapshotVersion(ctx, id) + m.queryLatencies.WithLabelValues("LockChatAndBumpSnapshotVersion").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "LockChatAndBumpSnapshotVersion").Inc() + return r0, r1 +} + func (m queryMetricsStore) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error { start := time.Now() r0 := m.s.MarkAllInboxNotificationsAsRead(ctx, arg) @@ -4634,6 +4818,14 @@ func (m queryMetricsStore) ReorderChatQueuedMessageToFront(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.ReorderChatQueuedMessageToHead(ctx, arg) + m.queryLatencies.WithLabelValues("ReorderChatQueuedMessageToHead").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ReorderChatQueuedMessageToHead").Inc() + return r0, r1 +} + func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { start := time.Now() r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID) @@ -4826,6 +5018,14 @@ func (m queryMetricsStore) UpdateChatDebugStep(ctx context.Context, arg database return r0, r1 } +func (m queryMetricsStore) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) { + start := time.Now() + r0, r1 := m.s.UpdateChatExecutionState(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateChatExecutionState").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatExecutionState").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { start := time.Now() r0, r1 := m.s.UpdateChatHeartbeats(ctx, arg) @@ -4914,6 +5114,14 @@ func (m queryMetricsStore) UpdateChatPlanModeByID(ctx context.Context, arg datab return r0, r1 } +func (m queryMetricsStore) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) { + start := time.Now() + r0, r1 := m.s.UpdateChatRetryState(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateChatRetryState").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatRetryState").Inc() + return r0, r1 +} + func (m queryMetricsStore) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) { start := time.Now() r0, r1 := m.s.UpdateChatStatus(ctx, arg) @@ -5866,6 +6074,14 @@ func (m queryMetricsStore) UpsertChatGeneralModelOverride(ctx context.Context, v return r0 } +func (m queryMetricsStore) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error { + start := time.Now() + r0 := m.s.UpsertChatHeartbeat(ctx, arg) + m.queryLatencies.WithLabelValues("UpsertChatHeartbeat").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatHeartbeat").Inc() + return r0 +} + func (m queryMetricsStore) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error { start := time.Now() r0 := m.s.UpsertChatIncludeDefaultSystemPrompt(ctx, includeDefaultSystemPrompt) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 60ddd85446..79ae18782d 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -223,6 +223,21 @@ func (mr *MockStoreMockRecorder) BackoffChatDiffStatus(ctx, arg any) *gomock.Cal return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BackoffChatDiffStatus", reflect.TypeOf((*MockStore)(nil).BackoffChatDiffStatus), ctx, arg) } +// BatchDeleteChatHeartbeats mocks base method. +func (m *MockStore) BatchDeleteChatHeartbeats(ctx context.Context, arg database.BatchDeleteChatHeartbeatsParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BatchDeleteChatHeartbeats", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// BatchDeleteChatHeartbeats indicates an expected call of BatchDeleteChatHeartbeats. +func (mr *MockStoreMockRecorder) BatchDeleteChatHeartbeats(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchDeleteChatHeartbeats", reflect.TypeOf((*MockStore)(nil).BatchDeleteChatHeartbeats), ctx, arg) +} + // BatchUpdateWorkspaceAgentMetadata mocks base method. func (m *MockStore) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg database.BatchUpdateWorkspaceAgentMetadataParams) error { m.ctrl.T.Helper() @@ -265,6 +280,20 @@ func (mr *MockStoreMockRecorder) BatchUpdateWorkspaceNextStartAt(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpdateWorkspaceNextStartAt", reflect.TypeOf((*MockStore)(nil).BatchUpdateWorkspaceNextStartAt), ctx, arg) } +// BatchUpsertChatHeartbeats mocks base method. +func (m *MockStore) BatchUpsertChatHeartbeats(ctx context.Context, arg database.BatchUpsertChatHeartbeatsParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "BatchUpsertChatHeartbeats", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// BatchUpsertChatHeartbeats indicates an expected call of BatchUpsertChatHeartbeats. +func (mr *MockStoreMockRecorder) BatchUpsertChatHeartbeats(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BatchUpsertChatHeartbeats", reflect.TypeOf((*MockStore)(nil).BatchUpsertChatHeartbeats), ctx, arg) +} + // BatchUpsertConnectionLogs mocks base method. func (m *MockStore) BatchUpsertConnectionLogs(ctx context.Context, arg database.BatchUpsertConnectionLogsParams) error { m.ctrl.T.Helper() @@ -484,6 +513,21 @@ func (mr *MockStoreMockRecorder) CountAuthorizedConnectionLogs(ctx, arg, prepare return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountAuthorizedConnectionLogs", reflect.TypeOf((*MockStore)(nil).CountAuthorizedConnectionLogs), ctx, arg, prepared) } +// CountChatQueuedMessages mocks base method. +func (m *MockStore) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CountChatQueuedMessages", ctx, chatID) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CountChatQueuedMessages indicates an expected call of CountChatQueuedMessages. +func (mr *MockStoreMockRecorder) CountChatQueuedMessages(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).CountChatQueuedMessages), ctx, chatID) +} + // CountConnectionLogs mocks base method. func (m *MockStore) CountConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams) (int64, error) { m.ctrl.T.Helper() @@ -660,6 +704,20 @@ func (mr *MockStoreMockRecorder) DeleteAPIKeysByUserID(ctx, userID any) *gomock. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAPIKeysByUserID", reflect.TypeOf((*MockStore)(nil).DeleteAPIKeysByUserID), ctx, userID) } +// DeleteAllChatHeartbeats mocks base method. +func (m *MockStore) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAllChatHeartbeats", ctx, chatID) + ret0, _ := ret[0].(error) + return ret0 +} + +// DeleteAllChatHeartbeats indicates an expected call of DeleteAllChatHeartbeats. +func (mr *MockStoreMockRecorder) DeleteAllChatHeartbeats(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatHeartbeats", reflect.TypeOf((*MockStore)(nil).DeleteAllChatHeartbeats), ctx, chatID) +} + // DeleteAllChatQueuedMessages mocks base method. func (m *MockStore) DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error { m.ctrl.T.Helper() @@ -674,6 +732,21 @@ func (mr *MockStoreMockRecorder) DeleteAllChatQueuedMessages(ctx, chatID any) *g return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).DeleteAllChatQueuedMessages), ctx, chatID) } +// DeleteAllChatQueuedMessagesReturningCount mocks base method. +func (m *MockStore) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteAllChatQueuedMessagesReturningCount", ctx, chatID) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteAllChatQueuedMessagesReturningCount indicates an expected call of DeleteAllChatQueuedMessagesReturningCount. +func (mr *MockStoreMockRecorder) DeleteAllChatQueuedMessagesReturningCount(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteAllChatQueuedMessagesReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteAllChatQueuedMessagesReturningCount), ctx, chatID) +} + // DeleteAllTailnetTunnels mocks base method. func (m *MockStore) DeleteAllTailnetTunnels(ctx context.Context, arg database.DeleteAllTailnetTunnelsParams) ([]database.DeleteAllTailnetTunnelsRow, error) { m.ctrl.T.Helper() @@ -803,6 +876,21 @@ func (mr *MockStoreMockRecorder) DeleteChatQueuedMessage(ctx, arg any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessage", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessage), ctx, arg) } +// DeleteChatQueuedMessageReturningCount mocks base method. +func (m *MockStore) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg database.DeleteChatQueuedMessageReturningCountParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteChatQueuedMessageReturningCount", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteChatQueuedMessageReturningCount indicates an expected call of DeleteChatQueuedMessageReturningCount. +func (mr *MockStoreMockRecorder) DeleteChatQueuedMessageReturningCount(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessageReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessageReturningCount), ctx, arg) +} + // DeleteChatUsageLimitGroupOverride mocks base method. func (m *MockStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error { m.ctrl.T.Helper() @@ -1305,6 +1393,21 @@ func (mr *MockStoreMockRecorder) DeleteRuntimeConfig(ctx, key any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteRuntimeConfig", reflect.TypeOf((*MockStore)(nil).DeleteRuntimeConfig), ctx, key) } +// DeleteStaleChatHeartbeats mocks base method. +func (m *MockStore) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteStaleChatHeartbeats", ctx, staleSeconds) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteStaleChatHeartbeats indicates an expected call of DeleteStaleChatHeartbeats. +func (mr *MockStoreMockRecorder) DeleteStaleChatHeartbeats(ctx, staleSeconds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteStaleChatHeartbeats", reflect.TypeOf((*MockStore)(nil).DeleteStaleChatHeartbeats), ctx, staleSeconds) +} + // DeleteTailnetPeer mocks base method. func (m *MockStore) DeleteTailnetPeer(ctx context.Context, arg database.DeleteTailnetPeerParams) (database.DeleteTailnetPeerRow, error) { m.ctrl.T.Helper() @@ -2341,6 +2444,21 @@ func (mr *MockStoreMockRecorder) GetAuthorizedWorkspacesAndAgentsByOwnerID(ctx, return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthorizedWorkspacesAndAgentsByOwnerID", reflect.TypeOf((*MockStore)(nil).GetAuthorizedWorkspacesAndAgentsByOwnerID), ctx, ownerID, prepared) } +// GetAutoArchiveInactiveChatCandidates mocks base method. +func (m *MockStore) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAutoArchiveInactiveChatCandidates", ctx, arg) + ret0, _ := ret[0].([]database.GetAutoArchiveInactiveChatCandidatesRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAutoArchiveInactiveChatCandidates indicates an expected call of GetAutoArchiveInactiveChatCandidates. +func (mr *MockStoreMockRecorder) GetAutoArchiveInactiveChatCandidates(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAutoArchiveInactiveChatCandidates", reflect.TypeOf((*MockStore)(nil).GetAutoArchiveInactiveChatCandidates), ctx, arg) +} + // GetBoundaryLogByID mocks base method. func (m *MockStore) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (database.BoundaryLog, error) { m.ctrl.T.Helper() @@ -2431,6 +2549,21 @@ func (mr *MockStoreMockRecorder) GetChatByID(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatByID", reflect.TypeOf((*MockStore)(nil).GetChatByID), ctx, id) } +// GetChatByIDForShare mocks base method. +func (m *MockStore) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (database.Chat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatByIDForShare", ctx, id) + ret0, _ := ret[0].(database.Chat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatByIDForShare indicates an expected call of GetChatByIDForShare. +func (mr *MockStoreMockRecorder) GetChatByIDForShare(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatByIDForShare", reflect.TypeOf((*MockStore)(nil).GetChatByIDForShare), ctx, id) +} + // GetChatByIDForUpdate mocks base method. func (m *MockStore) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (database.Chat, error) { m.ctrl.T.Helper() @@ -2671,6 +2804,21 @@ func (mr *MockStoreMockRecorder) GetChatExploreModelOverride(ctx any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatExploreModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatExploreModelOverride), ctx) } +// GetChatFamilyIDsByRootID mocks base method. +func (m *MockStore) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatFamilyIDsByRootID", ctx, id) + ret0, _ := ret[0].([]uuid.UUID) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatFamilyIDsByRootID indicates an expected call of GetChatFamilyIDsByRootID. +func (mr *MockStoreMockRecorder) GetChatFamilyIDsByRootID(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatFamilyIDsByRootID", reflect.TypeOf((*MockStore)(nil).GetChatFamilyIDsByRootID), ctx, id) +} + // GetChatFileByID mocks base method. func (m *MockStore) GetChatFileByID(ctx context.Context, id uuid.UUID) (database.ChatFile, error) { m.ctrl.T.Helper() @@ -2731,6 +2879,21 @@ func (mr *MockStoreMockRecorder) GetChatGeneralModelOverride(ctx any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatGeneralModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatGeneralModelOverride), ctx) } +// GetChatHeartbeat mocks base method. +func (m *MockStore) GetChatHeartbeat(ctx context.Context, arg database.GetChatHeartbeatParams) (database.ChatHeartbeat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatHeartbeat", ctx, arg) + ret0, _ := ret[0].(database.ChatHeartbeat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatHeartbeat indicates an expected call of GetChatHeartbeat. +func (mr *MockStoreMockRecorder) GetChatHeartbeat(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatHeartbeat", reflect.TypeOf((*MockStore)(nil).GetChatHeartbeat), ctx, arg) +} + // GetChatIncludeDefaultSystemPrompt mocks base method. func (m *MockStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) { m.ctrl.T.Helper() @@ -2821,6 +2984,21 @@ func (mr *MockStoreMockRecorder) GetChatMessagesByChatIDDescPaginated(ctx, arg a return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByChatIDDescPaginated", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByChatIDDescPaginated), ctx, arg) } +// GetChatMessagesByRevisionForStream mocks base method. +func (m *MockStore) GetChatMessagesByRevisionForStream(ctx context.Context, arg database.GetChatMessagesByRevisionForStreamParams) ([]database.ChatMessage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatMessagesByRevisionForStream", ctx, arg) + ret0, _ := ret[0].([]database.ChatMessage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatMessagesByRevisionForStream indicates an expected call of GetChatMessagesByRevisionForStream. +func (mr *MockStoreMockRecorder) GetChatMessagesByRevisionForStream(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatMessagesByRevisionForStream", reflect.TypeOf((*MockStore)(nil).GetChatMessagesByRevisionForStream), ctx, arg) +} + // GetChatMessagesForPromptByChatID mocks base method. func (m *MockStore) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatMessage, error) { m.ctrl.T.Helper() @@ -2911,6 +3089,36 @@ func (mr *MockStoreMockRecorder) GetChatPlanModeInstructions(ctx any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatPlanModeInstructions", reflect.TypeOf((*MockStore)(nil).GetChatPlanModeInstructions), ctx) } +// GetChatQueuedMessageByID mocks base method. +func (m *MockStore) GetChatQueuedMessageByID(ctx context.Context, arg database.GetChatQueuedMessageByIDParams) (database.ChatQueuedMessage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatQueuedMessageByID", ctx, arg) + ret0, _ := ret[0].(database.ChatQueuedMessage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatQueuedMessageByID indicates an expected call of GetChatQueuedMessageByID. +func (mr *MockStoreMockRecorder) GetChatQueuedMessageByID(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessageByID", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessageByID), ctx, arg) +} + +// GetChatQueuedMessageHead mocks base method. +func (m *MockStore) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (database.ChatQueuedMessage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatQueuedMessageHead", ctx, chatID) + ret0, _ := ret[0].(database.ChatQueuedMessage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatQueuedMessageHead indicates an expected call of GetChatQueuedMessageHead. +func (mr *MockStoreMockRecorder) GetChatQueuedMessageHead(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessageHead", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessageHead), ctx, chatID) +} + // GetChatQueuedMessages mocks base method. func (m *MockStore) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { m.ctrl.T.Helper() @@ -2926,6 +3134,21 @@ func (mr *MockStoreMockRecorder) GetChatQueuedMessages(ctx, chatID any) *gomock. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessages", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessages), ctx, chatID) } +// GetChatQueuedMessagesByPosition mocks base method. +func (m *MockStore) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]database.ChatQueuedMessage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatQueuedMessagesByPosition", ctx, chatID) + ret0, _ := ret[0].([]database.ChatQueuedMessage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatQueuedMessagesByPosition indicates an expected call of GetChatQueuedMessagesByPosition. +func (mr *MockStoreMockRecorder) GetChatQueuedMessagesByPosition(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatQueuedMessagesByPosition", reflect.TypeOf((*MockStore)(nil).GetChatQueuedMessagesByPosition), ctx, chatID) +} + // GetChatRetentionDays mocks base method. func (m *MockStore) GetChatRetentionDays(ctx context.Context) (int32, error) { m.ctrl.T.Helper() @@ -2941,6 +3164,21 @@ func (mr *MockStoreMockRecorder) GetChatRetentionDays(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatRetentionDays", reflect.TypeOf((*MockStore)(nil).GetChatRetentionDays), ctx) } +// GetChatStreamSyncRows mocks base method. +func (m *MockStore) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]database.GetChatStreamSyncRowsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatStreamSyncRows", ctx, ids) + ret0, _ := ret[0].([]database.GetChatStreamSyncRowsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatStreamSyncRows indicates an expected call of GetChatStreamSyncRows. +func (mr *MockStoreMockRecorder) GetChatStreamSyncRows(ctx, ids any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatStreamSyncRows", reflect.TypeOf((*MockStore)(nil).GetChatStreamSyncRows), ctx, ids) +} + // GetChatSystemPrompt mocks base method. func (m *MockStore) GetChatSystemPrompt(ctx context.Context) (string, error) { m.ctrl.T.Helper() @@ -3061,6 +3299,21 @@ func (mr *MockStoreMockRecorder) GetChatUserPromptsByChatID(ctx, arg any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUserPromptsByChatID", reflect.TypeOf((*MockStore)(nil).GetChatUserPromptsByChatID), ctx, arg) } +// GetChatWorkerAcquisitionCandidates mocks base method. +func (m *MockStore) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg database.GetChatWorkerAcquisitionCandidatesParams) ([]database.GetChatWorkerAcquisitionCandidatesRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatWorkerAcquisitionCandidates", ctx, arg) + ret0, _ := ret[0].([]database.GetChatWorkerAcquisitionCandidatesRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatWorkerAcquisitionCandidates indicates an expected call of GetChatWorkerAcquisitionCandidates. +func (mr *MockStoreMockRecorder) GetChatWorkerAcquisitionCandidates(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatWorkerAcquisitionCandidates", reflect.TypeOf((*MockStore)(nil).GetChatWorkerAcquisitionCandidates), ctx, arg) +} + // GetChatWorkspaceTTL mocks base method. func (m *MockStore) GetChatWorkspaceTTL(ctx context.Context) (string, error) { m.ctrl.T.Helper() @@ -3106,6 +3359,21 @@ func (mr *MockStoreMockRecorder) GetChatsByChatFileID(ctx, fileID any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatsByChatFileID", reflect.TypeOf((*MockStore)(nil).GetChatsByChatFileID), ctx, fileID) } +// GetChatsByIDsForRunnerSync mocks base method. +func (m *MockStore) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatsByIDsForRunnerSync", ctx, ids) + ret0, _ := ret[0].([]database.Chat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatsByIDsForRunnerSync indicates an expected call of GetChatsByIDsForRunnerSync. +func (mr *MockStoreMockRecorder) GetChatsByIDsForRunnerSync(ctx, ids any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatsByIDsForRunnerSync", reflect.TypeOf((*MockStore)(nil).GetChatsByIDsForRunnerSync), ctx, ids) +} + // GetChatsByWorkspaceIDs mocks base method. func (m *MockStore) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]database.Chat, error) { m.ctrl.T.Helper() @@ -3241,6 +3509,21 @@ func (mr *MockStoreMockRecorder) GetDERPMeshKey(ctx any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDERPMeshKey", reflect.TypeOf((*MockStore)(nil).GetDERPMeshKey), ctx) } +// GetDatabaseNow mocks base method. +func (m *MockStore) GetDatabaseNow(ctx context.Context) (time.Time, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetDatabaseNow", ctx) + ret0, _ := ret[0].(time.Time) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetDatabaseNow indicates an expected call of GetDatabaseNow. +func (mr *MockStoreMockRecorder) GetDatabaseNow(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDatabaseNow", reflect.TypeOf((*MockStore)(nil).GetDatabaseNow), ctx) +} + // GetDefaultChatModelConfig mocks base method. func (m *MockStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { m.ctrl.T.Helper() @@ -6945,6 +7228,21 @@ func (mr *MockStoreMockRecorder) InTx(arg0, arg1 any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InTx", reflect.TypeOf((*MockStore)(nil).InTx), arg0, arg1) } +// IncrementChatGenerationAttempt mocks base method. +func (m *MockStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IncrementChatGenerationAttempt", ctx, id) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// IncrementChatGenerationAttempt indicates an expected call of IncrementChatGenerationAttempt. +func (mr *MockStoreMockRecorder) IncrementChatGenerationAttempt(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IncrementChatGenerationAttempt", reflect.TypeOf((*MockStore)(nil).IncrementChatGenerationAttempt), ctx, id) +} + // InsertAIBridgeInterception mocks base method. func (m *MockStore) InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) { m.ctrl.T.Helper() @@ -7245,6 +7543,21 @@ func (mr *MockStoreMockRecorder) InsertChatQueuedMessage(ctx, arg any) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatQueuedMessage", reflect.TypeOf((*MockStore)(nil).InsertChatQueuedMessage), ctx, arg) } +// InsertChatQueuedMessageWithCreator mocks base method. +func (m *MockStore) InsertChatQueuedMessageWithCreator(ctx context.Context, arg database.InsertChatQueuedMessageWithCreatorParams) (database.ChatQueuedMessage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertChatQueuedMessageWithCreator", ctx, arg) + ret0, _ := ret[0].(database.ChatQueuedMessage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// InsertChatQueuedMessageWithCreator indicates an expected call of InsertChatQueuedMessageWithCreator. +func (mr *MockStoreMockRecorder) InsertChatQueuedMessageWithCreator(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatQueuedMessageWithCreator", reflect.TypeOf((*MockStore)(nil).InsertChatQueuedMessageWithCreator), ctx, arg) +} + // InsertCryptoKey mocks base method. func (m *MockStore) InsertCryptoKey(ctx context.Context, arg database.InsertCryptoKeyParams) (database.CryptoKey, error) { m.ctrl.T.Helper() @@ -8160,6 +8473,21 @@ func (mr *MockStoreMockRecorder) InsertWorkspaceResourceMetadata(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceResourceMetadata", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceResourceMetadata), ctx, arg) } +// IsChatHeartbeatStale mocks base method. +func (m *MockStore) IsChatHeartbeatStale(ctx context.Context, arg database.IsChatHeartbeatStaleParams) (bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "IsChatHeartbeatStale", ctx, arg) + ret0, _ := ret[0].(bool) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// IsChatHeartbeatStale indicates an expected call of IsChatHeartbeatStale. +func (mr *MockStoreMockRecorder) IsChatHeartbeatStale(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "IsChatHeartbeatStale", reflect.TypeOf((*MockStore)(nil).IsChatHeartbeatStale), ctx, arg) +} + // LinkChatFiles mocks base method. func (m *MockStore) LinkChatFiles(ctx context.Context, arg database.LinkChatFilesParams) (int32, error) { m.ctrl.T.Helper() @@ -8565,6 +8893,21 @@ func (mr *MockStoreMockRecorder) ListWorkspaceAgentPortShares(ctx, workspaceID a return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListWorkspaceAgentPortShares", reflect.TypeOf((*MockStore)(nil).ListWorkspaceAgentPortShares), ctx, workspaceID) } +// LockChatAndBumpSnapshotVersion mocks base method. +func (m *MockStore) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "LockChatAndBumpSnapshotVersion", ctx, id) + ret0, _ := ret[0].(database.Chat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// LockChatAndBumpSnapshotVersion indicates an expected call of LockChatAndBumpSnapshotVersion. +func (mr *MockStoreMockRecorder) LockChatAndBumpSnapshotVersion(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "LockChatAndBumpSnapshotVersion", reflect.TypeOf((*MockStore)(nil).LockChatAndBumpSnapshotVersion), ctx, id) +} + // MarkAllInboxNotificationsAsRead mocks base method. func (m *MockStore) MarkAllInboxNotificationsAsRead(ctx context.Context, arg database.MarkAllInboxNotificationsAsReadParams) error { m.ctrl.T.Helper() @@ -8757,6 +9100,21 @@ func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToFront(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToFront", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToFront), ctx, arg) } +// ReorderChatQueuedMessageToHead mocks base method. +func (m *MockStore) ReorderChatQueuedMessageToHead(ctx context.Context, arg database.ReorderChatQueuedMessageToHeadParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ReorderChatQueuedMessageToHead", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ReorderChatQueuedMessageToHead indicates an expected call of ReorderChatQueuedMessageToHead. +func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToHead(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToHead", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToHead), ctx, arg) +} + // ResolveUserChatSpendLimit mocks base method. func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) { m.ctrl.T.Helper() @@ -9103,6 +9461,21 @@ func (mr *MockStoreMockRecorder) UpdateChatDebugStep(ctx, arg any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatDebugStep", reflect.TypeOf((*MockStore)(nil).UpdateChatDebugStep), ctx, arg) } +// UpdateChatExecutionState mocks base method. +func (m *MockStore) UpdateChatExecutionState(ctx context.Context, arg database.UpdateChatExecutionStateParams) (database.Chat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateChatExecutionState", ctx, arg) + ret0, _ := ret[0].(database.Chat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpdateChatExecutionState indicates an expected call of UpdateChatExecutionState. +func (mr *MockStoreMockRecorder) UpdateChatExecutionState(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatExecutionState", reflect.TypeOf((*MockStore)(nil).UpdateChatExecutionState), ctx, arg) +} + // UpdateChatHeartbeats mocks base method. func (m *MockStore) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { m.ctrl.T.Helper() @@ -9266,6 +9639,21 @@ func (mr *MockStoreMockRecorder) UpdateChatPlanModeByID(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatPlanModeByID", reflect.TypeOf((*MockStore)(nil).UpdateChatPlanModeByID), ctx, arg) } +// UpdateChatRetryState mocks base method. +func (m *MockStore) UpdateChatRetryState(ctx context.Context, arg database.UpdateChatRetryStateParams) (database.Chat, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateChatRetryState", ctx, arg) + ret0, _ := ret[0].(database.Chat) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// UpdateChatRetryState indicates an expected call of UpdateChatRetryState. +func (mr *MockStoreMockRecorder) UpdateChatRetryState(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatRetryState", reflect.TypeOf((*MockStore)(nil).UpdateChatRetryState), ctx, arg) +} + // UpdateChatStatus mocks base method. func (m *MockStore) UpdateChatStatus(ctx context.Context, arg database.UpdateChatStatusParams) (database.Chat, error) { m.ctrl.T.Helper() @@ -10990,6 +11378,20 @@ func (mr *MockStoreMockRecorder) UpsertChatGeneralModelOverride(ctx, value any) return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatGeneralModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatGeneralModelOverride), ctx, value) } +// UpsertChatHeartbeat mocks base method. +func (m *MockStore) UpsertChatHeartbeat(ctx context.Context, arg database.UpsertChatHeartbeatParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpsertChatHeartbeat", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpsertChatHeartbeat indicates an expected call of UpsertChatHeartbeat. +func (mr *MockStoreMockRecorder) UpsertChatHeartbeat(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatHeartbeat", reflect.TypeOf((*MockStore)(nil).UpsertChatHeartbeat), ctx, arg) +} + // UpsertChatIncludeDefaultSystemPrompt mocks base method. func (m *MockStore) UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error { m.ctrl.T.Helper() diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index c87bc5a8df..646bd7edd9 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/coderd/database/dbpurge/dbpurge.go @@ -1,18 +1,12 @@ package dbpurge import ( - "cmp" "context" "errors" "io" - "net/http" - "slices" - "strconv" "sync/atomic" "time" - "github.com/dustin/go-humanize" - "github.com/google/uuid" "github.com/prometheus/client_golang/prometheus" "golang.org/x/xerrors" @@ -21,9 +15,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbtime" - "github.com/coder/coder/v2/coderd/notifications" "github.com/coder/coder/v2/coderd/pproflabel" - "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/codersdk" "github.com/coder/quartz" ) @@ -52,18 +44,8 @@ const ( // Chat debug run deletions can cascade into steps with large JSONB // payloads, so they use the same conservative batch size. chatDebugRunsBatchSize = 1000 - // chatAutoArchiveDigestMaxChats bounds how many chat titles a - // single digest body lists. Past the cap, surplus titles are - // summarized as "...and N more". 25 is a readable email-friendly - // length; the cap is unrelated to chatAutoArchiveBatchSize, which - // bounds work per tick. - chatAutoArchiveDigestMaxChats = 25 ) -// defaultChatAutoArchiveBatchSize bounds how many root chats one -// tick will archive by default. -const defaultChatAutoArchiveBatchSize int32 = 1000 - type Option func(*instance) // WithClock overrides the clock used by the purger. Defaults to @@ -72,34 +54,12 @@ func WithClock(clk quartz.Clock) Option { return func(i *instance) { i.clk = clk } } -// WithChatAutoArchiveBatchSize overrides how many root chats a -// single tick will auto-archive. Defaults to -// defaultChatAutoArchiveBatchSize (1000). -func WithChatAutoArchiveBatchSize(n int32) Option { - return func(i *instance) { i.chatAutoArchiveBatchSize = n } -} - -// WithNotificationsEnqueuer sets the enqueuer used for digest -// notifications. Defaults to notifications.NewNoopEnqueuer(). Panics -// if e is nil: a nil enqueuer would NPE on the first dispatch tick, -// and failing fast at option-apply time surfaces the misuse at -// startup rather than minutes later. -func WithNotificationsEnqueuer(e notifications.Enqueuer) Option { - if e == nil { - panic("developer error: WithNotificationsEnqueuer called with nil enqueuer") - } - return func(i *instance) { i.enqueuer = e } -} - // New creates a new periodically purging database instance. // Callers must Close the returned instance. // -// The auditor pointer is loaded on each dispatch tick so runtime -// entitlement changes (e.g. toggling the audit-log feature) take -// effect without restarting the process. Notifications enqueuer -// defaults to no-op. Use WithNotificationsEnqueuer to pass a real -// one. -func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, auditor *atomic.Pointer[audit.Auditor], opts ...Option) io.Closer { +// The auditor pointer is accepted for compatibility with other background +// services. Dbpurge does not emit audit logs directly. +func New(ctx context.Context, logger slog.Logger, db database.Store, vals *codersdk.DeploymentValues, reg prometheus.Registerer, _ *atomic.Pointer[audit.Auditor], opts ...Option) io.Closer { closed := make(chan struct{}) ctx, cancelFunc := context.WithCancel(ctx) @@ -123,26 +83,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder }, []string{"record_type"}) reg.MustRegister(recordsPurged) - chatAutoArchiveRecords := prometheus.NewCounter(prometheus.CounterOpts{ - Namespace: "coderd", - Subsystem: "chat_auto_archive", - Name: "records_archived_total", - Help: "Total number of chats archived by the auto-archive job (counting both roots and cascaded children).", - }) - reg.MustRegister(chatAutoArchiveRecords) - inst := &instance{ - cancel: cancelFunc, - closed: closed, - logger: logger, - vals: vals, - clk: quartz.NewReal(), - auditor: auditor, - enqueuer: notifications.NewNoopEnqueuer(), - iterationDuration: iterationDuration, - recordsPurged: recordsPurged, - chatAutoArchiveRecords: chatAutoArchiveRecords, - chatAutoArchiveBatchSize: defaultChatAutoArchiveBatchSize, + cancel: cancelFunc, + closed: closed, + logger: logger, + vals: vals, + clk: quartz.NewReal(), + iterationDuration: iterationDuration, + recordsPurged: recordsPurged, } for _, opt := range opts { opt(inst) @@ -185,19 +133,14 @@ func New(ctx context.Context, logger slog.Logger, db database.Store, vals *coder func (i *instance) purgeTick(ctx context.Context, db database.Store, start time.Time) error { // Read chat configs outside the tx so a corrupt value can't // poison subsequent queries. On config read errors, log and stash - // the error, then run unrelated purges best-effort. Retention and - // auto-archive errors skip only the conversation purge and - // auto-archive work. Debug retention errors skip only the debug - // purge. purgeTick returns chatConfigErr after the tx so the failed - // iteration is operator-visible via metric and logs. + // the error, then run unrelated purges best-effort. Retention + // errors skip only the conversation purge. Debug retention errors + // skip only the debug purge. purgeTick returns chatConfigErr after + // the tx so the failed iteration is operator-visible via metric and + // logs. chatRetentionDays, chatRetentionErr := db.GetChatRetentionDays(ctx) if chatRetentionErr != nil { - i.logger.Error(ctx, "failed to read chat retention config: skipping chat purge and auto-archive this tick", slog.Error(chatRetentionErr)) - } - - chatAutoArchiveDays, chatAutoArchiveErr := db.GetChatAutoArchiveDays(ctx, codersdk.DefaultChatAutoArchiveDays) - if chatAutoArchiveErr != nil { - i.logger.Error(ctx, "failed to read chat auto-archive config: skipping chat purge and auto-archive this tick", slog.Error(chatAutoArchiveErr)) + i.logger.Error(ctx, "failed to read chat retention config: skipping chat purge this tick", slog.Error(chatRetentionErr)) } chatDebugRetentionDays, chatDebugRetentionErr := db.GetChatDebugRetentionDays(ctx, codersdk.DefaultChatDebugRetentionDays) @@ -205,11 +148,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. i.logger.Error(ctx, "failed to read chat debug retention config: skipping chat debug purge this tick", slog.Error(chatDebugRetentionErr)) } - chatRetentionConfigErr := errors.Join(chatRetentionErr, chatAutoArchiveErr) - chatConfigErr := errors.Join(chatRetentionConfigErr, chatDebugRetentionErr) - - // Populated inside the tx; dispatched post-commit. - var archivedChats []database.AutoArchiveInactiveChatsRow + chatConfigErr := errors.Join(chatRetentionErr, chatDebugRetentionErr) // Start a transaction to grab advisory lock, we don't want to run // multiple purges at the same time (multiple replicas). @@ -316,8 +255,8 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. } var purgedChats, purgedChatFiles, purgedChatDebugRuns int64 - if chatRetentionConfigErr == nil { - purgedChats, purgedChatFiles, archivedChats, err = i.purgeChatsInTx(ctx, tx, start, chatRetentionDays, chatAutoArchiveDays) + if chatRetentionErr == nil { + purgedChats, purgedChatFiles, err = i.purgeChatsInTx(ctx, tx, start, chatRetentionDays) if err != nil { return xerrors.Errorf("failed to purge chats: %w", err) } @@ -345,7 +284,6 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. slog.F("chats", purgedChats), slog.F("chat_files", purgedChatFiles), slog.F("chat_debug_runs", purgedChatDebugRuns), - slog.F("auto_archived_chats", len(archivedChats)), slog.F("duration", i.clk.Since(start)), ) @@ -379,35 +317,17 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. return xerrors.Errorf("chat config read failed this tick: %w", chatConfigErr) } - // Dispatch audits and digests post-commit. Detached context for audit - // so that ticker cancellation cannot truncate the audit trail. - // Notification enqueue uses the cancellable parent context to avoid - // stalling shutdown. - // Owners with more eligible chats than batch size will get a - // notification per tick until their backlog drains. - // If this is deemed too noisy, users can disable the - // "Chats Auto-Archived" template from their notification preferences. - if len(archivedChats) > 0 { - i.chatAutoArchiveRecords.Add(float64(len(archivedChats))) - auditCtx := context.WithoutCancel(ctx) - i.dispatchChatAutoArchive(auditCtx, ctx, start, chatAutoArchiveDays, chatRetentionDays, archivedChats) - } - return nil } type instance struct { - cancel context.CancelFunc - closed chan struct{} - logger slog.Logger - vals *codersdk.DeploymentValues - clk quartz.Clock - auditor *atomic.Pointer[audit.Auditor] - enqueuer notifications.Enqueuer - iterationDuration *prometheus.HistogramVec - recordsPurged *prometheus.CounterVec - chatAutoArchiveRecords prometheus.Counter - chatAutoArchiveBatchSize int32 + cancel context.CancelFunc + closed chan struct{} + logger slog.Logger + vals *codersdk.DeploymentValues + clk quartz.Clock + iterationDuration *prometheus.HistogramVec + recordsPurged *prometheus.CounterVec } func (i *instance) Close() error { @@ -416,73 +336,8 @@ func (i *instance) Close() error { return nil } -// chatFromAutoArchiveRow reshapes the query row into a database.Chat for -// audit.Auditable[database.Chat]. -func chatFromAutoArchiveRow(logger slog.Logger, r database.AutoArchiveInactiveChatsRow) database.Chat { - var labels database.StringMap - // sqlc's StringMap override doesn't reach CTE-aliased columns, so Labels - // arrives as raw JSON bytes. StringMap.Scan handles []byte and nil. - if err := labels.Scan([]byte(r.Labels)); err != nil { - logger.Warn(context.Background(), "failed to parse chat labels from auto-archive row", - slog.F("chat_id", r.ID), - slog.F("raw_labels", string(r.Labels)), - slog.Error(err), - ) - } - - var userACL database.ChatACL - if err := userACL.Scan([]byte(r.UserACL)); err != nil { - logger.Warn(context.Background(), "failed to parse chat user ACL from auto-archive row", - slog.F("chat_id", r.ID), - slog.F("raw_user_acl", string(r.UserACL)), - slog.Error(err), - ) - } - - var groupACL database.ChatACL - if err := groupACL.Scan([]byte(r.GroupACL)); err != nil { - logger.Warn(context.Background(), "failed to parse chat group ACL from auto-archive row", - slog.F("chat_id", r.ID), - slog.F("raw_group_acl", string(r.GroupACL)), - slog.Error(err), - ) - } - - return database.Chat{ - ID: r.ID, - OwnerID: r.OwnerID, - OrganizationID: r.OrganizationID, - WorkspaceID: r.WorkspaceID, - BuildID: r.BuildID, - AgentID: r.AgentID, - Title: r.Title, - Status: r.Status, - WorkerID: r.WorkerID, - StartedAt: r.StartedAt, - HeartbeatAt: r.HeartbeatAt, - CreatedAt: r.CreatedAt, - UpdatedAt: r.UpdatedAt, - ParentChatID: r.ParentChatID, - RootChatID: r.RootChatID, - LastModelConfigID: r.LastModelConfigID, - Archived: r.Archived, - LastError: r.LastError, - Mode: r.Mode, - MCPServerIDs: r.MCPServerIDs, - Labels: labels, - UserACL: userACL, - GroupACL: groupACL, - PinOrder: r.PinOrder, - LastReadMessageID: r.LastReadMessageID, - LastInjectedContext: r.LastInjectedContext, - DynamicTools: r.DynamicTools, - PlanMode: r.PlanMode, - ClientType: r.ClientType, - } -} - // purgeChatsInTx MUST BE CALLED WITH A TRANSACTION -func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start time.Time, chatRetentionDays, chatAutoArchiveDays int32) (purgedChats, purgedChatFiles int64, archivedChats []database.AutoArchiveInactiveChatsRow, err error) { +func (*instance) purgeChatsInTx(ctx context.Context, tx database.Store, start time.Time, chatRetentionDays int32) (purgedChats, purgedChatFiles int64, err error) { // Delete old archived chats first, then orphaned files // (cascade clears chat_file_links but not chat_files). if chatRetentionDays > 0 { @@ -492,7 +347,7 @@ func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start LimitCount: chatsBatchSize, }) if err != nil { - return 0, 0, nil, xerrors.Errorf("failed to delete old chats: %w", err) + return 0, 0, xerrors.Errorf("failed to delete old chats: %w", err) } purgedChatFiles, err = tx.DeleteOldChatFiles(ctx, database.DeleteOldChatFilesParams{ @@ -500,149 +355,9 @@ func (i *instance) purgeChatsInTx(ctx context.Context, tx database.Store, start LimitCount: chatFilesBatchSize, }) if err != nil { - return 0, 0, nil, xerrors.Errorf("failed to delete old chat files: %w", err) + return 0, 0, xerrors.Errorf("failed to delete old chat files: %w", err) } } - // Auto-archive runs after the delete pass so newly - // archived chats aren't eligible for deletion this tick. - // Eligibility uses UTC day boundaries: a chat is archived on the - // start of the UTC day after its inactivity period has elapsed. - if chatAutoArchiveDays > 0 { - today := dbtime.StartOfDay(start) - archiveCutoff := today.Add(-time.Duration(chatAutoArchiveDays) * 24 * time.Hour) - archivedChats, err = tx.AutoArchiveInactiveChats(ctx, database.AutoArchiveInactiveChatsParams{ - ArchiveCutoff: archiveCutoff, - LimitCount: i.chatAutoArchiveBatchSize, - }) - if err != nil { - return 0, 0, nil, xerrors.Errorf("failed to auto-archive inactive chats: %w", err) - } - } - return purgedChats, purgedChatFiles, archivedChats, nil -} - -// dispatchChatAutoArchive audits every archived root chat and enqueues one -// notification per owner covering the roots archived in this tick. Children -// inherit their root's archival decision and are skipped for audit, matching -// the manual archive path (patchChat audits the root only). Enqueue is -// per-tick: owners whose backlog spans multiple ticks receive multiple -// notifications; notification_messages dedupe does not collapse them because -// each tick's payload differs. -// -// auditCtx is detached from the ticker so audits always complete. enqueueCtx -// is the cancellable parent: on shutdown we abandon any remaining digests -// rather than blocking Close. -func (i *instance) dispatchChatAutoArchive(auditCtx, enqueueCtx context.Context, tickStart time.Time, autoArchiveDays, retentionDays int32, archived []database.AutoArchiveInactiveChatsRow) { - // Children inherit their root's archival decision and are skipped - // for both audit and digest. Partition once so the two loops - // cannot drift apart if the cascade shape ever changes. - roots := slice.Filter(archived, func(r database.AutoArchiveInactiveChatsRow) bool { - return !r.ParentChatID.Valid - }) - - auditor := *i.auditor.Load() - for _, row := range roots { - after := chatFromAutoArchiveRow(i.logger, row) - before := after - before.Archived = false - audit.BackgroundAudit(auditCtx, &audit.BackgroundAuditParams[database.Chat]{ - Audit: auditor, - Log: i.logger, - UserID: row.OwnerID, - OrganizationID: row.OrganizationID, - Action: database.AuditActionWrite, - Old: before, - New: after, - Status: http.StatusOK, - AdditionalFields: audit.BackgroundTaskFieldsBytes(auditCtx, i.logger, audit.BackgroundSubsystemChatAutoArchive), - }) - } - - // Group archived roots by owner. Inline because this is the - // only call site and the loop body is self-explanatory. - rootsByOwner := make(map[uuid.UUID][]database.AutoArchiveInactiveChatsRow, len(roots)) - for _, row := range roots { - rootsByOwner[row.OwnerID] = append(rootsByOwner[row.OwnerID], row) - } - - // Sort owner IDs so shutdown abandons a deterministic tail of the dispatch list. - ownerIDs := make([]uuid.UUID, 0, len(rootsByOwner)) - for id := range rootsByOwner { - ownerIDs = append(ownerIDs, id) - } - slices.SortFunc(ownerIDs, func(a, b uuid.UUID) int { - return cmp.Compare(a.String(), b.String()) - }) - - dispatched := 0 - for _, ownerID := range ownerIDs { - // Check between iterations so shutdown unblocks promptly. A - // hung in-flight enqueue is unblocked by enqueueCtx propagating - // cancellation into the DB call. Skipped owners are not - // re-notified on the next tick because AutoArchiveInactiveChats - // only returns rows with archived = false; we accept that - // tradeoff over hanging shutdown. - if err := enqueueCtx.Err(); err != nil { - i.logger.Warn(enqueueCtx, "chat auto-archive digest dispatch canceled", - slog.F("remaining_owners", len(ownerIDs)-dispatched), - slog.Error(err)) - return - } - dispatched++ - - ownerRoots := rootsByOwner[ownerID] - data := buildDigestData(ownerRoots, autoArchiveDays, retentionDays, tickStart) - - // nolint:gocritic // Background digest runs as the notifier subject. - if _, err := i.enqueuer.EnqueueWithData( - dbauthz.AsNotifier(enqueueCtx), - ownerID, - notifications.TemplateChatAutoArchiveDigest, - map[string]string{}, - data, - string(audit.BackgroundSubsystemChatAutoArchive), - ); err != nil { - i.logger.Warn(enqueueCtx, "failed to enqueue chat auto-archive digest", - slog.F("owner_id", ownerID), - slog.Error(err)) - } - } -} - -// buildDigestData builds the notification payload; shape mirrors the -// golden fixtures in coderd/notifications/testdata. Truncation keeps -// the oldest archived roots (created_at ASC from the query) to -// preserve index-driven ordering; revisit if the digest becomes the -// primary surface for reviewing archived chats. -func buildDigestData(rows []database.AutoArchiveInactiveChatsRow, autoArchiveDays, retentionDays int32, tickStart time.Time) map[string]any { - // Cap titles; overflow surfaces as "...and N more" via the template. - overflow := 0 - if len(rows) > chatAutoArchiveDigestMaxChats { - overflow = len(rows) - chatAutoArchiveDigestMaxChats - rows = rows[:chatAutoArchiveDigestMaxChats] - } - - chats := make([]map[string]any, 0, len(rows)) - for _, r := range rows { - chats = append(chats, map[string]any{ - "title": r.Title, - "last_activity_humanized": humanize.RelTime(r.LastActivityAt, tickStart, "ago", "from now"), - }) - } - - // Stringify the int32 config values: the template's - // {{if eq .Data.retention_days "0"}} branch requires both - // operands to share a type, and Go templates do not coerce - // numeric ↔ string. Storing a raw int here would silently - // take the deletion-warning branch on every notification. - data := map[string]any{ - "auto_archive_days": strconv.Itoa(int(autoArchiveDays)), - "retention_days": strconv.Itoa(int(retentionDays)), - "archived_chats": chats, - } - if overflow > 0 { - data["additional_archived_count"] = strconv.Itoa(overflow) - } - return data + return purgedChats, purgedChatFiles, nil } diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 4ebd645a7a..c0e784f538 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -32,9 +32,6 @@ import ( "github.com/coder/coder/v2/coderd/database/dbrollup" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" - "github.com/coder/coder/v2/coderd/notifications" - "github.com/coder/coder/v2/coderd/notifications/notificationsmock" - "github.com/coder/coder/v2/coderd/notifications/notificationstest" "github.com/coder/coder/v2/coderd/provisionerdserver" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/provisionerd/proto" @@ -60,7 +57,6 @@ func TestPurge(t *testing.T) { done := awaitDoTick(ctx, t, clk) mDB := dbmock.NewMockStore(gomock.NewController(t)) mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(0), nil).AnyTimes() - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays).Return(int32(0), nil).AnyTimes() mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays).Return(int32(0), nil).AnyTimes() mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")).Return(nil).Times(2) purger := dbpurge.New(context.Background(), testutil.Logger(t), mDB, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), nopAuditorPtr(t), dbpurge.WithClock(clk)) @@ -163,8 +159,6 @@ func TestMetrics(t *testing.T) { ctrl := gomock.NewController(t) mDB := dbmock.NewMockStore(ctrl) mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(0), nil).AnyTimes() - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays). - Return(int32(0), nil).AnyTimes() mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays). Return(int32(0), nil).AnyTimes() mDB.EXPECT().TryAcquireLock(gomock.Any(), int64(database.LockIDDBPurge)).Return(false, nil).AnyTimes() @@ -203,7 +197,6 @@ func TestMetrics(t *testing.T) { ctrl := gomock.NewController(t) mDB := dbmock.NewMockStore(ctrl) mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(0), nil).AnyTimes() - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays).Return(int32(0), nil).AnyTimes() mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays). Return(int32(0), nil).AnyTimes() mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")). @@ -230,8 +223,8 @@ func TestMetrics(t *testing.T) { }) // A failed retention read must not block unrelated or chat debug - // purges, but must skip the conversation purge and auto-archive - // passes and surface as a failed iteration via the metric. + // purges, but must skip the conversation purge and surface as a + // failed iteration via the metric. t.Run("FailedChatRetentionRead", func(t *testing.T) { ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort) defer cancel() @@ -248,8 +241,6 @@ func TestMetrics(t *testing.T) { MinTimes(1) // All reads happen before the bail; InTx still runs so unrelated // purges and chat debug purge commit best-effort. - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays). - Return(int32(0), nil).AnyTimes() mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays). Return(int32(7), nil).AnyTimes() mDB.EXPECT().TryAcquireLock(gomock.Any(), int64(database.LockIDDBPurge)).Return(true, nil).AnyTimes() @@ -285,50 +276,6 @@ func TestMetrics(t *testing.T) { require.Nil(t, successHist, "should not have success=true metric on retention read failure") }) - // Same contract as FailedChatRetentionRead, but the - // auto-archive read is the half that fails. - t.Run("FailedChatAutoArchiveRead", func(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort) - defer cancel() - - reg := prometheus.NewRegistry() - clk := quartz.NewMock(t) - now := clk.Now() - clk.Set(now).MustWait(ctx) - - ctrl := gomock.NewController(t) - mDB := dbmock.NewMockStore(ctrl) - mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(30), nil).AnyTimes() - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays). - Return(int32(0), xerrors.New("simulated auto-archive read error")). - MinTimes(1) - mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays). - Return(int32(0), nil).AnyTimes() - // InTx still runs so unrelated purges commit; chat - // passes inside the tx are skipped. - mDB.EXPECT().InTx(gomock.Any(), database.DefaultTXOptions().WithID("db_purge")). - Return(nil).MinTimes(1) - - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, mDB, &codersdk.DeploymentValues{}, reg, nopAuditorPtr(t), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - hist := promhelp.HistogramValue(t, reg, "coderd_dbpurge_iteration_duration_seconds", prometheus.Labels{ - "success": "false", - }) - require.NotNil(t, hist) - require.Greater(t, hist.GetSampleCount(), uint64(0), - "failed auto-archive read must record a failed iteration") - - successHist := promhelp.MetricValue(t, reg, "coderd_dbpurge_iteration_duration_seconds", prometheus.Labels{ - "success": "true", - }) - require.Nil(t, successHist, "should not have success=true metric on auto-archive read failure") - }) - // Same contract as the other chat config reads, but debug retention // read failures skip only debug purging. t.Run("FailedChatDebugRetentionRead", func(t *testing.T) { @@ -343,8 +290,6 @@ func TestMetrics(t *testing.T) { ctrl := gomock.NewController(t) mDB := dbmock.NewMockStore(ctrl) mDB.EXPECT().GetChatRetentionDays(gomock.Any()).Return(int32(30), nil).AnyTimes() - mDB.EXPECT().GetChatAutoArchiveDays(gomock.Any(), codersdk.DefaultChatAutoArchiveDays). - Return(int32(0), nil).AnyTimes() mDB.EXPECT().GetChatDebugRetentionDays(gomock.Any(), codersdk.DefaultChatDebugRetentionDays). Return(int32(0), xerrors.New("simulated chat debug retention read error")). MinTimes(1) @@ -632,63 +577,6 @@ func awaitDoTick(ctx context.Context, t *testing.T, clk *quartz.Mock) chan struc return ch } -// tickDriver drives one or more dbpurge ticks against a single -// dbpurge.New instance. Unlike awaitDoTick it must be constructed -// *before* dbpurge.New so its traps are installed when the forced -// initial tick fires. awaitInitial waits for the forced tick's -// doTick to complete without advancing the clock, so no loop -// iteration has yet run; awaitNext then explicitly drives each -// subsequent iteration. This keeps each tick's observable state -// isolated and deterministic, which matters for tests where -// per-tick work differs (e.g. batch-size pagination). -type tickDriver struct { - clk *quartz.Mock - trapNow *quartz.Trap - trapStop *quartz.Trap - trapReset *quartz.Trap -} - -func newTickDriver(t *testing.T, clk *quartz.Mock) *tickDriver { - t.Helper() - d := &tickDriver{ - clk: clk, - trapNow: clk.Trap().Now(), - trapStop: clk.Trap().TickerStop(), - trapReset: clk.Trap().TickerReset(), - } - return d -} - -// close releases all traps. Call this via defer *after* the defer -// that closes the dbpurge instance so trap closure releases the -// shutdown ticker.Stop() rather than blocking on it. -func (d *tickDriver) close() { - d.trapReset.Close() - d.trapStop.Close() - d.trapNow.Close() -} - -// awaitInitial waits for the forced initial tick's doTick to -// complete. No loop iteration runs because the clock has not been -// advanced. -func (d *tickDriver) awaitInitial(ctx context.Context, t *testing.T) { - t.Helper() - d.trapNow.MustWait(ctx).MustRelease(ctx) - d.trapReset.MustWait(ctx).MustRelease(ctx) -} - -// awaitNext advances the clock by the tick interval, lets the loop -// receive the tick and run doTick, and waits for the ensuing -// ticker.Reset so the driver is ready for another awaitNext. -func (d *tickDriver) awaitNext(ctx context.Context, t *testing.T) { - t.Helper() - dur, w := d.clk.AdvanceNext() - require.Equal(t, 10*time.Minute, dur) - w.MustWait(ctx) - d.trapStop.MustWait(ctx).MustRelease(ctx) - d.trapReset.MustWait(ctx).MustRelease(ctx) -} - func assertNoWorkspaceAgentLogs(ctx context.Context, t *testing.T, db database.Store, agentID uuid.UUID) { t.Helper() agentLogs, err := db.GetWorkspaceAgentLogsAfter(ctx, database.GetWorkspaceAgentLogsAfterParams{ @@ -1922,14 +1810,6 @@ func nopAuditorPtr(t *testing.T) *atomic.Pointer[audit.Auditor] { return &p } -// mockAuditorPtr wraps a *MockAuditor in an atomic pointer for tests. -func mockAuditorPtr(m *audit.MockAuditor) *atomic.Pointer[audit.Auditor] { - a := audit.Auditor(m) - var p atomic.Pointer[audit.Auditor] - p.Store(&a) - return &p -} - //nolint:paralleltest // It uses LockIDDBPurge. func TestPurgeChatDebugRuns(t *testing.T) { now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) @@ -2585,943 +2465,3 @@ func TestDeleteOldChatFiles(t *testing.T) { }) } } - -// helpers for TestAutoArchiveInactiveChats. Kept scoped to the -// test so they don't leak into the package surface area. -func archiveTestDeps(t *testing.T, db database.Store) chatAutoArchiveDeps { - t.Helper() - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) - _ = dbgen.ChatProvider(t, db, database.ChatProvider{ - Provider: "openai", - DisplayName: "OpenAI", - }) - mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Provider: "openai", - Model: "test-model", - ContextLimit: 8192, - }) - return chatAutoArchiveDeps{user: user, org: org, modelConfig: mc} -} - -type chatAutoArchiveDeps struct { - user database.User - org database.Organization - modelConfig database.ChatModelConfig -} - -// archiveHarness bundles the per-subtest setup shared by every -// TestAutoArchiveInactiveChats case. Subtests read fields off the -// harness directly instead of repeating six lines of identical -// plumbing. -type archiveHarness struct { - ctx context.Context - clk *quartz.Mock - db database.Store - rawDB *sql.DB - logger slog.Logger - deps chatAutoArchiveDeps -} - -func newArchiveHarness(t *testing.T, now time.Time) *archiveHarness { - t.Helper() - ctx := testutil.Context(t, testutil.WaitLong) - clk := quartz.NewMock(t) - clk.Set(now).MustWait(ctx) - db, _, rawDB := dbtestutil.NewDBWithSQLDB(t, dbtestutil.WithDumpOnFailure()) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - return &archiveHarness{ - ctx: ctx, - clk: clk, - db: db, - rawDB: rawDB, - logger: logger, - deps: archiveTestDeps(t, db), - } -} - -// createArchiveChat inserts a chat with an optional backdated -// created_at. Title is propagated through so tests can assert on -// digest contents. -func createArchiveChat(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, deps chatAutoArchiveDeps, title string, createdAt time.Time) database.Chat { - t.Helper() - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: deps.org.ID, - OwnerID: deps.user.ID, - LastModelConfigID: deps.modelConfig.ID, - Title: title, - }) - _, err := rawDB.ExecContext(ctx, "UPDATE chats SET created_at = $1, updated_at = $1 WHERE id = $2", createdAt, chat.ID) - require.NoError(t, err) - return chat -} - -// insertTextMessage appends a non-deleted user message with a -// backdated created_at. Used to establish "last activity" for the -// auto-archive query's LATERAL subquery. -func insertTextMessage(ctx context.Context, t *testing.T, db database.Store, rawDB *sql.DB, chatID, userID, modelConfigID uuid.UUID, createdAt time.Time) { - t.Helper() - msg := dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: chatID, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, - Role: database.ChatMessageRoleUser, - }) - _, err := rawDB.ExecContext(ctx, "UPDATE chat_messages SET created_at = $1 WHERE id = $2", createdAt, msg.ID) - require.NoError(t, err) -} - -//nolint:paralleltest // It uses LockIDDBPurge. -func TestAutoArchiveInactiveChats(t *testing.T) { - now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) - - tests := []struct { - name string - run func(t *testing.T) - }{ - { - name: "AutoArchiveDisabled", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.Zero(t, codersdk.DefaultChatAutoArchiveDays) - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, codersdk.DefaultChatAutoArchiveDays)) - - // Chat older than any reasonable cutoff. - staleChat := createArchiveChat(ctx, t, db, rawDB, deps, "stale-chat", now.Add(-365*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - // Not archived, no audits, no digests. - refreshed, err := db.GetChatByID(ctx, staleChat.ID) - require.NoError(t, err) - require.False(t, refreshed.Archived, "chat should stay active when auto-archive is disabled") - - require.Empty(t, auditor.AuditLogs(), "no audit log entries expected") - require.Empty(t, enqueuer.Sent(), "no digest notifications expected") - }, - }, - { - name: "ArchivesInactiveRoot", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - // Regression guard: ensure that both auto-archive and retention - // are both set to a distinct non-zero value. - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - require.NoError(t, db.UpsertChatRetentionDays(ctx, int32(30))) - - // Inactive root: newest message 100 days old. - staleChat := createArchiveChat(ctx, t, db, rawDB, deps, "stale-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, staleChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - - // Active root: message 10 days old, within cutoff. - activeChat := createArchiveChat(ctx, t, db, rawDB, deps, "active-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, activeChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-10*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - refreshedStale, err := db.GetChatByID(ctx, staleChat.ID) - require.NoError(t, err) - require.True(t, refreshedStale.Archived, "stale chat should be auto-archived") - - refreshedActive, err := db.GetChatByID(ctx, activeChat.ID) - require.NoError(t, err) - require.False(t, refreshedActive.Archived, "active chat should stay live") - - // Exactly one audit entry, for the stale root. - logs := auditor.AuditLogs() - require.Len(t, logs, 1, "expected one audit entry") - require.Equal(t, staleChat.ID, logs[0].ResourceID) - require.Equal(t, database.ResourceTypeChat, logs[0].ResourceType) - require.Equal(t, database.AuditActionWrite, logs[0].Action) - require.Contains(t, string(logs[0].AdditionalFields), "chat_auto_archive", - "audit entry must carry the auto-archive subsystem tag") - - // Exactly one digest, addressed to the owner. - sent := enqueuer.Sent() - require.Len(t, sent, 1, "expected one digest notification") - require.Equal(t, notifications.TemplateChatAutoArchiveDigest, sent[0].TemplateID) - require.Equal(t, deps.user.ID, sent[0].UserID) - // Ensure that config-derived fields flow through to payload. - require.Equal(t, "90", sent[0].Data["auto_archive_days"]) - require.Equal(t, "30", sent[0].Data["retention_days"]) - }, - }, - { - name: "DateBoundary", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // With now = 2025-06-15 12:00 UTC, the Go code - // truncates to today = 2025-06-15 00:00 UTC, then - // subtracts 90 days -> cutoff = 2025-03-17 00:00 UTC. - // A chat's last-activity UTC date must be strictly < - // 2025-03-17 to be archived. - - // Activity on the cutoff date (2025-03-17): must survive. - onDate := createArchiveChat(ctx, t, db, rawDB, deps, "on-date", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, onDate.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 17, 15, 30, 0, 0, time.UTC)) - - // Activity day before cutoff date (2025-03-16): must be archived. - beforeDate := createArchiveChat(ctx, t, db, rawDB, deps, "before-date", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, beforeDate.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 16, 23, 59, 59, 0, time.UTC)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - refreshedOn, err := db.GetChatByID(ctx, onDate.ID) - require.NoError(t, err) - require.False(t, refreshedOn.Archived, "chat with activity on cutoff date must survive") - - refreshedBefore, err := db.GetChatByID(ctx, beforeDate.ID) - require.NoError(t, err) - require.True(t, refreshedBefore.Archived, "chat with activity day before cutoff must be archived") - - require.Len(t, auditor.AuditLogs(), 1, "only the before-date chat should produce an audit entry") - }, - }, - { - name: "DayBoundaryLateActivity", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Activity at 23:59:59 UTC on 2025-03-17 (cutoff date). - // The UTC date is still 2025-03-17, NOT < cutoff date, - // so it must NOT be archived. - lateChat := createArchiveChat(ctx, t, db, rawDB, deps, "late-activity", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, lateChat.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 17, 23, 59, 59, 0, time.UTC)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - refreshed, err := db.GetChatByID(ctx, lateChat.ID) - require.NoError(t, err) - require.False(t, refreshed.Archived, "activity at 23:59:59 UTC on cutoff date must not be archived") - require.Empty(t, auditor.AuditLogs()) - }, - }, - { - name: "SameDayActivityNotArchived", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Activity at 00:00:01 UTC on the cutoff date - // (2025-03-17). Same date as cutoff, NOT strictly <, - // so must NOT be archived. - earlyChat := createArchiveChat(ctx, t, db, rawDB, deps, "early-same-day", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, earlyChat.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 17, 0, 0, 1, 0, time.UTC)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - refreshed, err := db.GetChatByID(ctx, earlyChat.ID) - require.NoError(t, err) - require.False(t, refreshed.Archived, "activity at start of cutoff date must not be archived") - require.Empty(t, auditor.AuditLogs()) - }, - }, - { - name: "SameDayBatch", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Three chats all with last activity on 2025-03-16 - // (one day before cutoff) but at different times. - // All should be archived in the same batch. - chat1 := createArchiveChat(ctx, t, db, rawDB, deps, "batch-1", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, chat1.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 16, 1, 0, 0, 0, time.UTC)) - - chat2 := createArchiveChat(ctx, t, db, rawDB, deps, "batch-2", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, chat2.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 16, 12, 0, 0, 0, time.UTC)) - - chat3 := createArchiveChat(ctx, t, db, rawDB, deps, "batch-3", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, chat3.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 16, 23, 59, 0, 0, time.UTC)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - for _, tc := range []struct { - name string - id uuid.UUID - }{ - {"batch-1", chat1.ID}, - {"batch-2", chat2.ID}, - {"batch-3", chat3.ID}, - } { - refreshed, err := db.GetChatByID(ctx, tc.id) - require.NoError(t, err) - require.True(t, refreshed.Archived, "%s should be archived", tc.name) - } - - require.Len(t, auditor.AuditLogs(), 3, "all three chats should produce audit entries") - }, - }, - { - // CutoffStableAcrossSameDayTicks verifies that the archive - // cutoff is derived from the UTC day, not from the wall-clock - // time. Advancing the clock within the same UTC day must not - // change the archival decision ("no trickle" property). The - // chat is only archived once the clock crosses into the next - // UTC day and the cutoff date advances. - name: "CutoffStableAcrossSameDayTicks", - run: func(t *testing.T) { - // Start close to midnight so exactly two awaitNext calls - // cross the UTC day boundary: tick 1 at 23:49, tick 2 at - // 23:59 (still June 15, cutoff unchanged), tick 3 at - // 00:09 June 16 (new day, cutoff advances). - nearMidnight := time.Date(2025, 6, 15, 23, 49, 0, 0, time.UTC) - h := newArchiveHarness(t, nearMidnight) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Chat last active on 2025-03-17, which equals the cutoff - // for any tick on 2025-06-15: truncate(today) - 90d = - // 2025-03-17. The query requires last-activity < cutoff - // (strict), so the chat must survive all June-15 ticks. - chat := createArchiveChat(ctx, t, db, rawDB, deps, "boundary-chat", nearMidnight.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, chat.ID, deps.user.ID, deps.modelConfig.ID, - time.Date(2025, 3, 17, 12, 0, 0, 0, time.UTC)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - defer driver.close() - - // Tick 1 (23:49 UTC June 15): cutoff = 2025-03-17. - // Activity on the cutoff date is not strictly less than - // the cutoff, so the chat must not be archived. - driver.awaitInitial(ctx, t) - - refreshed, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.False(t, refreshed.Archived, "tick 1: chat on cutoff date must not be archived") - require.Empty(t, auditor.AuditLogs(), "tick 1: no audit entries expected") - - // Tick 2 (23:59 UTC June 15): still the same UTC day. - // The cutoff is unchanged (still 2025-03-17), so advancing - // the wall clock within the same day must not archive the - // chat. - driver.awaitNext(ctx, t) - - refreshed, err = db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.False(t, refreshed.Archived, "tick 2: same UTC day, cutoff unchanged, chat must still survive") - require.Empty(t, auditor.AuditLogs(), "tick 2: no audit entries expected") - - // Tick 3 (00:09 UTC June 16): new UTC day. The cutoff - // advances to 2025-03-18, so activity on 2025-03-17 is - // now strictly less than the cutoff and the chat must be - // archived. - driver.awaitNext(ctx, t) - - refreshed, err = db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.True(t, refreshed.Archived, "tick 3: cutoff advanced to 2025-03-18, chat must now be archived") - require.Len(t, auditor.AuditLogs(), 1, "tick 3: exactly one audit entry expected") - }, - }, - - { - name: "DeletedMessagesIgnored", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Chat created 120 days ago with a recent message - // (10 days old) that is then soft-deleted. The - // LATERAL subquery filters cm.deleted = false, so - // the chat should fall back to created_at and be - // archived. - chat := createArchiveChat(ctx, t, db, rawDB, deps, "deleted-msg", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, chat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-10*24*time.Hour)) - // Soft-delete all messages on this chat. - _, err := rawDB.ExecContext(ctx, "UPDATE chat_messages SET deleted = true WHERE chat_id = $1", chat.ID) - require.NoError(t, err) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - refreshed, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.True(t, refreshed.Archived, "chat with only deleted messages should be archived") - require.Len(t, auditor.AuditLogs(), 1) - }, - }, - { - name: "ChildActivityKeepsRootAlive", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Stale root with no messages of its own. - root := createArchiveChat(ctx, t, db, rawDB, deps, "stale-root", now.Add(-120*24*time.Hour)) - - // Child linked to root with a recent message (10 days old, - // well within the 90-day cutoff). - child := createArchiveChat(ctx, t, db, rawDB, deps, "active-child", now.Add(-120*24*time.Hour)) - _, err := rawDB.ExecContext(ctx, "UPDATE chats SET parent_chat_id = $1, root_chat_id = $1 WHERE id = $2", root.ID, child.ID) - require.NoError(t, err) - insertTextMessage(ctx, t, db, rawDB, child.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-10*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - refreshedRoot, err := db.GetChatByID(ctx, root.ID) - require.NoError(t, err) - require.False(t, refreshedRoot.Archived, "root must stay active because child has recent activity") - - refreshedChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.False(t, refreshedChild.Archived, "child must stay active") - - require.Empty(t, auditor.AuditLogs(), "no chats should be archived") - require.Empty(t, enqueuer.Sent(), "no notifications should be sent") - }, - }, - { - name: "SkipsActiveStatusChats", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(90))) - - // Stale chats whose status prevents archiving. - runningChat := createArchiveChat(ctx, t, db, rawDB, deps, "running-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, runningChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - _, err := rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusRunning, runningChat.ID) - require.NoError(t, err) - - requiresActionChat := createArchiveChat(ctx, t, db, rawDB, deps, "requires-action-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, requiresActionChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusRequiresAction, requiresActionChat.ID) - require.NoError(t, err) - - pendingChat := createArchiveChat(ctx, t, db, rawDB, deps, "pending-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, pendingChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusPending, pendingChat.ID) - require.NoError(t, err) - - pausedChat := createArchiveChat(ctx, t, db, rawDB, deps, "paused-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, pausedChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusPaused, pausedChat.ID) - require.NoError(t, err) - - // Control: a stale chat with archivable status that - // should be archived. - completedChat := createArchiveChat(ctx, t, db, rawDB, deps, "completed-chat", now.Add(-120*24*time.Hour)) - insertTextMessage(ctx, t, db, rawDB, completedChat.ID, deps.user.ID, deps.modelConfig.ID, now.Add(-100*24*time.Hour)) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusCompleted, completedChat.ID) - require.NoError(t, err) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - refreshedRunning, err := db.GetChatByID(ctx, runningChat.ID) - require.NoError(t, err) - require.False(t, refreshedRunning.Archived, "running chat must not be archived") - - refreshedRA, err := db.GetChatByID(ctx, requiresActionChat.ID) - require.NoError(t, err) - require.False(t, refreshedRA.Archived, "requires_action chat must not be archived") - - refreshedPending, err := db.GetChatByID(ctx, pendingChat.ID) - require.NoError(t, err) - require.False(t, refreshedPending.Archived, "pending chat must not be archived") - - refreshedPaused, err := db.GetChatByID(ctx, pausedChat.ID) - require.NoError(t, err) - require.False(t, refreshedPaused.Archived, "paused chat must not be archived") - - refreshedCompleted, err := db.GetChatByID(ctx, completedChat.ID) - require.NoError(t, err) - require.True(t, refreshedCompleted.Archived, "completed stale chat should be archived") - - logs := auditor.AuditLogs() - require.Len(t, logs, 1, "only the completed chat should produce an audit entry") - require.Equal(t, completedChat.ID, logs[0].ResourceID) - - // Assert number of sent notifications to catch dispatch regressions. - sent := enqueuer.Sent() - require.Len(t, sent, 1, "expected one digest notification for the completed chat") - require.Equal(t, notifications.TemplateChatAutoArchiveDigest, sent[0].TemplateID) - require.Equal(t, deps.user.ID, sent[0].UserID) - }, - }, - { - name: "SkipsPinnedAndChildren", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - // Pinned stale chat: should be skipped. - pinnedChat := createArchiveChat(ctx, t, db, rawDB, deps, "pinned-chat", now.Add(-90*24*time.Hour)) - _, err := rawDB.ExecContext(ctx, "UPDATE chats SET pin_order = 1 WHERE id = $1", pinnedChat.ID) - require.NoError(t, err) - - // Stale root with a child. - root := createArchiveChat(ctx, t, db, rawDB, deps, "root-chat", now.Add(-90*24*time.Hour)) - child := createArchiveChat(ctx, t, db, rawDB, deps, "child-chat", now.Add(-90*24*time.Hour)) - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET parent_chat_id = $1, root_chat_id = $1 WHERE id = $2", root.ID, child.ID) - require.NoError(t, err) - // Give the child an active status to prove the cascade is - // status-blind by design. If someone adds a status filter - // to the cascade CTE, this assertion will catch it. - _, err = rawDB.ExecContext(ctx, "UPDATE chats SET status = $1 WHERE id = $2", database.ChatStatusRunning, child.ID) - require.NoError(t, err) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - refreshedPinned, err := db.GetChatByID(ctx, pinnedChat.ID) - require.NoError(t, err) - require.False(t, refreshedPinned.Archived, "pinned chat must be skipped") - - refreshedRoot, err := db.GetChatByID(ctx, root.ID) - require.NoError(t, err) - require.True(t, refreshedRoot.Archived, "root should be archived") - - refreshedChild, err := db.GetChatByID(ctx, child.ID) - require.NoError(t, err) - require.True(t, refreshedChild.Archived, "child should be cascade-archived") - - // One audit entry for the root; the cascaded child is - // not audited individually. - require.Len(t, auditor.AuditLogs(), 1) - - // Digest should list only the root (one row). - sent := enqueuer.Sent() - require.Len(t, sent, 1) - data := sent[0].Data - require.NotNil(t, data) - chats, ok := data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - require.Len(t, chats, 1, "digest should only list the root") - require.Equal(t, "root-chat", chats[0]["title"]) - }, - }, - { - name: "DigestOverflowCap", - run: func(t *testing.T) { - // 27 inactive roots exceed chatAutoArchiveDigestMaxChats - // (25). All 27 should archive, but the digest payload - // lists at most 25 titles and surfaces the rest via - // additional_archived_count so the template can render - // "...and N more". - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - const total = 27 - for i := range total { - createArchiveChat(ctx, t, db, rawDB, deps, - fmt.Sprintf("stale-%02d", i), - now.Add(-60*24*time.Hour)) - } - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - // All 27 roots archived (one audit each). - require.Len(t, auditor.AuditLogs(), total) - - sent := enqueuer.Sent() - require.Len(t, sent, 1, "one digest per owner") - chats, ok := sent[0].Data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - require.Len(t, chats, 25, "digest caps titles at 25") - require.Equal(t, "2", sent[0].Data["additional_archived_count"], - "overflow count is total - cap") - // Humanized timestamp is computed from LastActivityAt - // and the tick-start time, not a static fixture, so we - // only assert the suffix the humanizer emits. - humanized, _ := chats[0]["last_activity_humanized"].(string) - require.Contains(t, humanized, "ago", - "last_activity_humanized should be a past relative time") - }, - }, - { - name: "MultipleOwners", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - user2 := dbgen.User(t, db, database.User{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user2.ID, OrganizationID: deps.org.ID}) - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - // Two stale roots per owner, backdated well past - // the 30-day cutoff. - u1Deps := deps - u2Deps := chatAutoArchiveDeps{user: user2, org: deps.org, modelConfig: deps.modelConfig} - createArchiveChat(ctx, t, db, rawDB, u1Deps, "u1-a", now.Add(-60*24*time.Hour)) - createArchiveChat(ctx, t, db, rawDB, u1Deps, "u1-b", now.Add(-60*24*time.Hour)) - createArchiveChat(ctx, t, db, rawDB, u2Deps, "u2-a", now.Add(-60*24*time.Hour)) - createArchiveChat(ctx, t, db, rawDB, u2Deps, "u2-b", now.Add(-60*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - // Four audit rows, one per archived root, attributed - // to the owning user so downstream consumers can - // correlate per-owner activity. - logs := auditor.AuditLogs() - require.Len(t, logs, 4) - auditsByUser := map[uuid.UUID]int{} - for _, l := range logs { - auditsByUser[l.UserID]++ - } - require.Equal(t, 2, auditsByUser[deps.user.ID]) - require.Equal(t, 2, auditsByUser[user2.ID]) - - // One digest per owner, each listing only that owner's - // two chats. - sent := enqueuer.Sent() - require.Len(t, sent, 2, "expected one digest per owner") - - byUser := map[uuid.UUID][]string{} - for _, s := range sent { - require.Equal(t, notifications.TemplateChatAutoArchiveDigest, s.TemplateID) - chats, ok := s.Data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - for _, c := range chats { - title, _ := c["title"].(string) - byUser[s.UserID] = append(byUser[s.UserID], title) - } - } - require.Contains(t, byUser, deps.user.ID) - require.Contains(t, byUser, user2.ID) - slices.Sort(byUser[deps.user.ID]) - slices.Sort(byUser[user2.ID]) - require.Equal(t, []string{"u1-a", "u1-b"}, byUser[deps.user.ID]) - require.Equal(t, []string{"u2-a", "u2-b"}, byUser[user2.ID]) - }, - }, - { - name: "SecondTickIdempotent", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - // Two stale roots seeded before the first tick. - firstA := createArchiveChat(ctx, t, db, rawDB, deps, "first-a", now.Add(-60*24*time.Hour)) - firstB := createArchiveChat(ctx, t, db, rawDB, deps, "first-b", now.Add(-60*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk)) - // Defer driver.close() after closer.Close(): defers - // run LIFO, so this frees shutdown's ticker.Stop() - // before the dbpurge goroutine blocks on it. - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - // Tick 1: both archived, one digest. - require.Len(t, auditor.AuditLogs(), 2, "tick 1 audits") - require.Len(t, enqueuer.Sent(), 1, "tick 1 digests") - - // Seed a third stale root between ticks so tick 2 has - // genuine work and we can distinguish "ignored already - // archived" from "ignored everything". - third := createArchiveChat(ctx, t, db, rawDB, deps, "second-c", now.Add(-60*24*time.Hour)) - - driver.awaitNext(ctx, t) - - // Tick 2: exactly one new audit + one new digest for - // the third chat; tick 1's rows must not be re-archived. - require.Len(t, auditor.AuditLogs(), 3, "tick 2 cumulative audits") - sent := enqueuer.Sent() - require.Len(t, sent, 2, "tick 2 cumulative digests") - chats, ok := sent[1].Data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - require.Len(t, chats, 1, "tick 2 digest lists only the new chat") - require.Equal(t, "second-c", chats[0]["title"]) - - // First-tick chats stayed archived. - for _, id := range []uuid.UUID{firstA.ID, firstB.ID, third.ID} { - refreshed, err := db.GetChatByID(ctx, id) - require.NoError(t, err) - require.True(t, refreshed.Archived, "chat %s should remain archived", id) - } - }, - }, - { - name: "BatchSizePagination", - run: func(t *testing.T) { - // With 27 stale roots and batch size 20, tick 1 - // archives 20, tick 2 archives the remaining 7, and - // tick 3 archives none. We assert the dispatch side - // effects (audits, digests) follow the same pattern: - // dispatch only runs when rows > 0, so tick 3 emits - // no new audits or digests. - // - // The two-digest count asserted here is a consequence - // of the per-tick enqueue model, not a product - // invariant. notification_messages dedupe does not - // collapse these because each tick's payload differs. - // If enqueue is ever restructured to one notification - // per owner per day, this assertion changes with it. - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - const total = 27 - for i := range total { - createArchiveChat(ctx, t, db, rawDB, deps, - fmt.Sprintf("page-%02d", i), - now.Add(-60*24*time.Hour)) - } - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - enqueuer := notificationstest.NewFakeEnqueuer() - driver := newTickDriver(t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(enqueuer), dbpurge.WithClock(clk), dbpurge.WithChatAutoArchiveBatchSize(20)) - // Defer driver.close() after closer.Close() so trap - // cleanup frees shutdown's ticker.Stop() before the - // dbpurge goroutine blocks on it. - defer closer.Close() - defer driver.close() - driver.awaitInitial(ctx, t) - - // Tick 1: first batch (20) archived. - require.Len(t, auditor.AuditLogs(), 20, "tick 1 audits") - sent := enqueuer.Sent() - require.Len(t, sent, 1, "tick 1 digests") - chats1, ok := sent[0].Data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - require.Len(t, chats1, 20, "tick 1 digest lists all 20 titles") - require.NotContains(t, sent[0].Data, "additional_archived_count", - "no overflow when batch <= digest cap; 20 <= 25") - - driver.awaitNext(ctx, t) - - // Tick 2: remaining 7 archived. - require.Len(t, auditor.AuditLogs(), 27, "tick 2 cumulative audits") - sent = enqueuer.Sent() - require.Len(t, sent, 2, "tick 2 cumulative digests") - chats2, ok := sent[1].Data["archived_chats"].([]map[string]any) - require.True(t, ok, "archived_chats should be []map[string]any") - require.Len(t, chats2, 7, "tick 2 digest lists remaining 7") - - driver.awaitNext(ctx, t) - - // Tick 3: nothing left to archive. The dispatch is - // gated on len(archivedChats) > 0, so no new audits - // or digests are produced. If that gate is ever - // removed, update this assertion intentionally. - require.Len(t, auditor.AuditLogs(), 27, "tick 3 cumulative audits unchanged") - require.Len(t, enqueuer.Sent(), 2, "tick 3 cumulative digests unchanged") - }, - }, - { - name: "ShutdownCancelsDigestDispatch", - run: func(t *testing.T) { - // Two owners with one stale root each. The first - // EnqueueWithData call blocks until ctx is canceled. - // Closing the purger must propagate cancellation - // into the in-flight call and short-circuit the - // rest of the loop, so Close returns promptly - // instead of hanging on dispatch. - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - user2 := dbgen.User(t, db, database.User{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user2.ID, OrganizationID: deps.org.ID}) - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - u1Deps := deps - u2Deps := chatAutoArchiveDeps{user: user2, org: deps.org, modelConfig: deps.modelConfig} - createArchiveChat(ctx, t, db, rawDB, u1Deps, "u1-stale", now.Add(-60*24*time.Hour)) - createArchiveChat(ctx, t, db, rawDB, u2Deps, "u2-stale", now.Add(-60*24*time.Hour)) - - // Dispatch iterates owner IDs in ascending UUID order (convention). - expectedFirst := deps.user.ID - if user2.ID.String() < deps.user.ID.String() { - expectedFirst = user2.ID - } - - ctrl := gomock.NewController(t) - mockEnq := notificationsmock.NewMockEnqueuer(ctrl) - started := make(chan struct{}) - mockEnq.EXPECT().EnqueueWithData(gomock.Any(), gomock.Eq(expectedFirst), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(ctx context.Context, _, _ uuid.UUID, _ map[string]string, _ map[string]any, _ string, _ ...uuid.UUID) ([]uuid.UUID, error) { - close(started) - <-ctx.Done() - return nil, ctx.Err() - }).Times(1) - - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), nopAuditorPtr(t), dbpurge.WithNotificationsEnqueuer(mockEnq), dbpurge.WithClock(clk)) - - // Wait for the forced initial tick to reach the first - // enqueue, which then blocks on ctx.Done(). - testutil.TryReceive(ctx, t, started) - - // Blocked enqueue receives ctx cancellation via the parent context. - // Loop-head check abandons the remaining owner instead of trying to enqueue. - done := make(chan error) - go func() { done <- closer.Close() }() - testutil.RequireReceive(ctx, t, done) - }, - }, - { - // A transient enqueue failure for one owner must not abort the dispatch loop. - name: "TransientEnqueueFailureDoesNotAbortLoop", - run: func(t *testing.T) { - h := newArchiveHarness(t, now) - ctx, clk, db, rawDB, logger, deps := h.ctx, h.clk, h.db, h.rawDB, h.logger, h.deps - user2 := dbgen.User(t, db, database.User{}) - _ = dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user2.ID, OrganizationID: deps.org.ID}) - - require.NoError(t, db.UpsertChatAutoArchiveDays(ctx, int32(30))) - - u1Deps := deps - u2Deps := chatAutoArchiveDeps{user: user2, org: deps.org, modelConfig: deps.modelConfig} - createArchiveChat(ctx, t, db, rawDB, u1Deps, "u1-stale", now.Add(-60*24*time.Hour)) - createArchiveChat(ctx, t, db, rawDB, u2Deps, "u2-stale", now.Add(-60*24*time.Hour)) - - auditor := audit.NewMock() - auditorPtr := mockAuditorPtr(auditor) - - ctrl := gomock.NewController(t) - mockEnq := notificationsmock.NewMockEnqueuer(ctrl) - var calls atomic.Int32 - var successUserID uuid.UUID - mockEnq.EXPECT().EnqueueWithData(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). - DoAndReturn(func(_ context.Context, userID, _ uuid.UUID, _ map[string]string, _ map[string]any, _ string, _ ...uuid.UUID) ([]uuid.UUID, error) { - if calls.Add(1) == 1 { - return nil, xerrors.New("simulated transient enqueue failure") - } - successUserID = userID - return nil, nil - }).Times(2) - - done := awaitDoTick(ctx, t, clk) - closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), auditorPtr, dbpurge.WithNotificationsEnqueuer(mockEnq), dbpurge.WithClock(clk)) - defer closer.Close() - testutil.TryReceive(ctx, t, done) - - // Both owners must have been audited regardless of - // digest enqueue outcomes; the audit and digest - // paths are independent. - require.Len(t, auditor.AuditLogs(), 2, "both archived roots must be audited") - - // gomock's .Times(2) already enforces both calls - // happened; this assertion makes the contract - // explicit at the test site. - require.Equal(t, int32(2), calls.Load(), - "loop must attempt every owner even when one fails") - - // The second attempt succeeded for one of the two owners. - require.Contains(t, []uuid.UUID{deps.user.ID, user2.ID}, successUserID, - "successful digest must belong to one of the two owners") - }, - }, - } - - for _, tc := range tests { - t.Run(tc.name, func(t *testing.T) { - tc.run(t) - }) - } -} diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 3a7b27894c..6f4d588d67 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -341,7 +341,8 @@ CREATE TYPE chat_status AS ENUM ( 'paused', 'completed', 'error', - 'requires_action' + 'requires_action', + 'interrupting' ); CREATE TYPE connection_status AS ENUM ( @@ -716,6 +717,29 @@ BEGIN END; $$; +CREATE FUNCTION bump_chat_queue_version_on_queued_message_change() RETURNS trigger + LANGUAGE plpgsql + AS $$ +DECLARE + changed_chat_id uuid; +BEGIN + IF TG_OP = 'DELETE' THEN + changed_chat_id = OLD.chat_id; + ELSE + changed_chat_id = NEW.chat_id; + END IF; + + UPDATE chats + SET queue_version = snapshot_version + WHERE id = changed_chat_id; + + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END; +$$; + CREATE FUNCTION check_workspace_agent_name_unique() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -1293,6 +1317,103 @@ BEGIN END; $$; +CREATE FUNCTION set_chat_message_revision_before() RETURNS trigger + LANGUAGE plpgsql + AS $$ +DECLARE + chat_snapshot_version bigint; +BEGIN + IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN + RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger'; + END IF; + + IF TG_OP = 'UPDATE' THEN + IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN + RAISE EXCEPTION 'chat_messages.chat_id is immutable'; + END IF; + + IF OLD.revision IS DISTINCT FROM NEW.revision THEN + RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger'; + END IF; + + IF OLD IS NOT DISTINCT FROM NEW THEN + RETURN NEW; + END IF; + END IF; + + SELECT snapshot_version INTO chat_snapshot_version + FROM chats WHERE id = NEW.chat_id; + + IF chat_snapshot_version IS NULL THEN + RAISE EXCEPTION 'chat % does not exist', NEW.chat_id; + END IF; + + NEW.revision = chat_snapshot_version; + RETURN NEW; +END; +$$; + +CREATE FUNCTION sync_chat_retry_state() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + IF OLD.retry_state_version IS DISTINCT FROM NEW.retry_state_version THEN + RAISE EXCEPTION 'chats.retry_state_version must be assigned by trigger'; + END IF; + + IF NEW.generation_attempt IS DISTINCT FROM OLD.generation_attempt THEN + NEW.retry_state = NULL; + END IF; + + IF NEW.retry_state IS DISTINCT FROM OLD.retry_state THEN + NEW.retry_state_version = NEW.snapshot_version; + END IF; + + RETURN NEW; +END; +$$; + +CREATE FUNCTION update_chat_history_after_message_insert() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + UPDATE chats c + SET history_version = c.snapshot_version, + generation_attempt = 0 + FROM ( + SELECT DISTINCT chat_id FROM chat_message_history_new_rows + ) AS affected + WHERE c.id = affected.chat_id + AND ( + c.history_version IS DISTINCT FROM c.snapshot_version + OR c.generation_attempt <> 0 + ); + RETURN NULL; +END; +$$; + +CREATE FUNCTION update_chat_history_after_message_update() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + UPDATE chats c + SET history_version = c.snapshot_version, + generation_attempt = 0 + FROM ( + SELECT DISTINCT n.chat_id + FROM chat_message_history_new_rows n + JOIN chat_message_history_old_rows o ON o.id = n.id + WHERE o IS DISTINCT FROM n + ) AS affected + WHERE c.id = affected.chat_id + AND ( + c.history_version IS DISTINCT FROM c.snapshot_version + OR c.generation_attempt <> 0 + ); + RETURN NULL; +END; +$$; + CREATE TABLE ai_gateway_keys ( id uuid NOT NULL, created_at timestamp with time zone NOT NULL, @@ -1670,6 +1791,14 @@ CREATE TABLE chat_files ( data bytea NOT NULL ); +CREATE UNLOGGED TABLE chat_heartbeats ( + chat_id uuid NOT NULL, + runner_id uuid NOT NULL, + heartbeat_at timestamp with time zone NOT NULL +); + +COMMENT ON TABLE chat_heartbeats IS 'Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats.'; + CREATE TABLE chat_messages ( id bigint NOT NULL, chat_id uuid NOT NULL, @@ -1692,7 +1821,8 @@ CREATE TABLE chat_messages ( runtime_ms bigint, deleted boolean DEFAULT false NOT NULL, provider_response_id text, - api_key_id text + api_key_id text, + revision bigint NOT NULL ); CREATE SEQUENCE chat_messages_id_seq @@ -1726,13 +1856,22 @@ CREATE TABLE chat_model_configs ( CONSTRAINT chat_model_configs_context_limit_check CHECK ((context_limit > 0)) ); +CREATE SEQUENCE chat_queued_messages_position_seq + START WITH 1 + INCREMENT BY 1 + NO MINVALUE + NO MAXVALUE + CACHE 1; + CREATE TABLE chat_queued_messages ( id bigint NOT NULL, chat_id uuid NOT NULL, content jsonb NOT NULL, created_at timestamp with time zone DEFAULT now() NOT NULL, model_config_id uuid, - api_key_id text + api_key_id text, + "position" bigint DEFAULT nextval('chat_queued_messages_position_seq'::regclass) NOT NULL, + created_by uuid NOT NULL ); CREATE SEQUENCE chat_queued_messages_id_seq @@ -1797,6 +1936,14 @@ CREATE TABLE chats ( last_turn_summary text, user_acl jsonb DEFAULT '{}'::jsonb NOT NULL, group_acl jsonb DEFAULT '{}'::jsonb NOT NULL, + snapshot_version bigint DEFAULT 1 NOT NULL, + history_version bigint DEFAULT 0 NOT NULL, + queue_version bigint DEFAULT 0 NOT NULL, + generation_attempt bigint DEFAULT 0 NOT NULL, + retry_state jsonb, + retry_state_version bigint DEFAULT 0 NOT NULL, + runner_id uuid, + requires_action_deadline_at timestamp with time zone, CONSTRAINT chat_acl_only_on_root_chats CHECK ((((parent_chat_id IS NULL) AND (root_chat_id IS NULL)) OR ((user_acl = '{}'::jsonb) AND (group_acl = '{}'::jsonb)))), CONSTRAINT chat_group_acl_not_null_jsonb CHECK (((group_acl IS NOT NULL) AND (jsonb_typeof(group_acl) = 'object'::text))), CONSTRAINT chat_user_acl_not_null_jsonb CHECK (((user_acl IS NOT NULL) AND (jsonb_typeof(user_acl) = 'object'::text))), @@ -1804,6 +1951,12 @@ CREATE TABLE chats ( CONSTRAINT chats_pin_order_parent_check CHECK (((pin_order = 0) OR (parent_chat_id IS NULL))) ); +COMMENT ON COLUMN chats.snapshot_version IS 'Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet.'; + +COMMENT ON COLUMN chats.history_version IS 'Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version.'; + +COMMENT ON COLUMN chats.queue_version IS 'Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version.'; + CREATE TABLE users ( id uuid NOT NULL, email text NOT NULL, @@ -1884,6 +2037,14 @@ CREATE VIEW chats_expanded AS c.plan_mode, c.client_type, c.last_turn_summary, + c.snapshot_version, + c.history_version, + c.queue_version, + c.generation_attempt, + c.retry_state, + c.retry_state_version, + c.runner_id, + c.requires_action_deadline_at, COALESCE(root.user_acl, c.user_acl) AS user_acl, COALESCE(root.group_acl, c.group_acl) AS group_acl, owner.username AS owner_username, @@ -3848,6 +4009,9 @@ ALTER TABLE ONLY chat_file_links ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id); +ALTER TABLE ONLY chat_heartbeats + ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id); + ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id); @@ -4190,6 +4354,8 @@ CREATE INDEX api_keys_last_used_idx ON api_keys USING btree (last_used DESC); COMMENT ON INDEX api_keys_last_used_idx IS 'Index for optimizing api_keys queries filtering by last_used'; +CREATE INDEX chat_heartbeats_heartbeat_at_idx ON chat_heartbeats USING btree (heartbeat_at); + CREATE INDEX idx_agent_stats_created_at ON workspace_agent_stats USING btree (created_at); CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_id); @@ -4320,6 +4486,8 @@ CREATE INDEX idx_chats_pending ON chats USING btree (status) WHERE (status = 'pe CREATE INDEX idx_chats_root_chat_id ON chats USING btree (root_chat_id); +CREATE INDEX idx_chats_worker_acquisition_candidates ON chats USING btree (status, updated_at, id) WHERE (archived = false); + CREATE INDEX idx_chats_workspace ON chats USING btree (workspace_id); CREATE INDEX idx_connection_logs_connect_time_desc ON connection_logs USING btree (connect_time DESC); @@ -4550,6 +4718,12 @@ COMMENT ON TRIGGER remove_organization_member_custom_role ON custom_roles IS 'Wh CREATE TRIGGER trigger_aggregate_usage_event AFTER INSERT ON usage_events FOR EACH ROW EXECUTE FUNCTION aggregate_usage_event(); +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_delete AFTER DELETE ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_insert AFTER INSERT ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_update AFTER UPDATE OF content, model_config_id, "position", created_by ON chat_queued_messages FOR EACH ROW EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + CREATE TRIGGER trigger_delete_group_members_on_org_member_delete BEFORE DELETE ON organization_members FOR EACH ROW EXECUTE FUNCTION delete_group_members_on_org_member_delete(); CREATE TRIGGER trigger_delete_oauth2_provider_app_token AFTER DELETE ON oauth2_provider_app_tokens FOR EACH ROW EXECUTE FUNCTION delete_deleted_oauth2_provider_app_token_api_key(); @@ -4566,6 +4740,16 @@ CREATE TRIGGER trigger_insert_organization_system_roles AFTER INSERT ON organiza CREATE TRIGGER trigger_nullify_next_start_at_on_workspace_autostart_modificati AFTER UPDATE ON workspaces FOR EACH ROW EXECUTE FUNCTION nullify_next_start_at_on_workspace_autostart_modification(); +CREATE TRIGGER trigger_set_chat_message_revision_on_insert BEFORE INSERT ON chat_messages FOR EACH ROW EXECUTE FUNCTION set_chat_message_revision_before(); + +CREATE TRIGGER trigger_set_chat_message_revision_on_update BEFORE UPDATE ON chat_messages FOR EACH ROW EXECUTE FUNCTION set_chat_message_revision_before(); + +CREATE TRIGGER trigger_sync_chat_retry_state BEFORE UPDATE OF retry_state, retry_state_version, generation_attempt ON chats FOR EACH ROW EXECUTE FUNCTION sync_chat_retry_state(); + +CREATE TRIGGER trigger_update_chat_history_after_message_insert AFTER INSERT ON chat_messages REFERENCING NEW TABLE AS chat_message_history_new_rows FOR EACH STATEMENT EXECUTE FUNCTION update_chat_history_after_message_insert(); + +CREATE TRIGGER trigger_update_chat_history_after_message_update AFTER UPDATE ON chat_messages REFERENCING OLD TABLE AS chat_message_history_old_rows NEW TABLE AS chat_message_history_new_rows FOR EACH STATEMENT EXECUTE FUNCTION update_chat_history_after_message_update(); + CREATE TRIGGER trigger_update_users AFTER INSERT OR UPDATE ON users FOR EACH ROW WHEN ((new.deleted = true)) EXECUTE FUNCTION delete_deleted_user_resources(); CREATE TRIGGER trigger_upsert_user_links BEFORE INSERT OR UPDATE ON user_links FOR EACH ROW EXECUTE FUNCTION insert_user_links_fail_if_user_deleted(); @@ -4636,6 +4820,9 @@ ALTER TABLE ONLY chat_files ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE; +ALTER TABLE ONLY chat_heartbeats + ADD CONSTRAINT chat_heartbeats_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE; + ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index 8109f2564f..159040d142 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -22,6 +22,7 @@ const ( ForeignKeyChatFileLinksFileID ForeignKeyConstraint = "chat_file_links_file_id_fkey" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_file_id_fkey FOREIGN KEY (file_id) REFERENCES chat_files(id) ON DELETE CASCADE; ForeignKeyChatFilesOrganizationID ForeignKeyConstraint = "chat_files_organization_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; ForeignKeyChatFilesOwnerID ForeignKeyConstraint = "chat_files_owner_id_fkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE; + ForeignKeyChatHeartbeatsChatID ForeignKeyConstraint = "chat_heartbeats_chat_id_fkey" // ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE; ForeignKeyChatMessagesAPIKeyID ForeignKeyConstraint = "chat_messages_api_key_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES api_keys(id) ON DELETE SET NULL; ForeignKeyChatMessagesChatID ForeignKeyConstraint = "chat_messages_chat_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE; ForeignKeyChatMessagesModelConfigID ForeignKeyConstraint = "chat_messages_model_config_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_model_config_id_fkey FOREIGN KEY (model_config_id) REFERENCES chat_model_configs(id); diff --git a/coderd/database/migrations/000519_chatd_core_state_machine.down.sql b/coderd/database/migrations/000519_chatd_core_state_machine.down.sql new file mode 100644 index 0000000000..fd109dc1b6 --- /dev/null +++ b/coderd/database/migrations/000519_chatd_core_state_machine.down.sql @@ -0,0 +1,106 @@ +-- Rollback for the chatd core state machine foundation migration. + +-- 1. Recreate chats_expanded without the new chat fields. We must drop +-- the view first because the subsequent column drops would fail with +-- "view depends on column". +DROP VIEW IF EXISTS chats_expanded; + +-- 2. Drop the worker acquisition candidates index. +DROP INDEX IF EXISTS idx_chats_worker_acquisition_candidates; + +-- 3. Drop the retry state trigger and function. +DROP TRIGGER IF EXISTS trigger_sync_chat_retry_state ON chats; +DROP FUNCTION IF EXISTS sync_chat_retry_state(); + +-- 4. Drop the queue version triggers and function. +DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_delete ON chat_queued_messages; +DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_update ON chat_queued_messages; +DROP TRIGGER IF EXISTS trigger_bump_chat_queue_version_on_queued_message_insert ON chat_queued_messages; +DROP FUNCTION IF EXISTS bump_chat_queue_version_on_queued_message_change(); + +-- 5. Drop the message revision triggers and functions. +DROP TRIGGER IF EXISTS trigger_update_chat_history_after_message_update ON chat_messages; +DROP TRIGGER IF EXISTS trigger_update_chat_history_after_message_insert ON chat_messages; +DROP TRIGGER IF EXISTS trigger_set_chat_message_revision_on_update ON chat_messages; +DROP TRIGGER IF EXISTS trigger_set_chat_message_revision_on_insert ON chat_messages; +DROP FUNCTION IF EXISTS update_chat_history_after_message_update(); +DROP FUNCTION IF EXISTS update_chat_history_after_message_insert(); +-- The pre-split function name is kept here for backward compatibility +-- with environments that may have applied an earlier draft of the up +-- migration. DROP FUNCTION IF EXISTS is a no-op if the function is +-- absent. +DROP FUNCTION IF EXISTS update_chat_history_after_message_changes(); +DROP FUNCTION IF EXISTS set_chat_message_revision_before(); +DROP FUNCTION IF EXISTS set_chat_message_revision(); + +-- 6. Drop chat_heartbeats (and its index by association). +DROP TABLE IF EXISTS chat_heartbeats; + +-- 7. Drop chat_queued_messages.position and its default sequence, plus +-- created_by. +ALTER TABLE chat_queued_messages + ALTER COLUMN position DROP DEFAULT; +ALTER TABLE chat_queued_messages + DROP COLUMN IF EXISTS position, + DROP COLUMN IF EXISTS created_by; +DROP SEQUENCE IF EXISTS chat_queued_messages_position_seq; + +-- 8. Drop chat_messages.revision. +ALTER TABLE chat_messages + DROP COLUMN IF EXISTS revision; + +-- 9. Drop the new chats columns. +ALTER TABLE chats + DROP COLUMN IF EXISTS snapshot_version, + DROP COLUMN IF EXISTS history_version, + DROP COLUMN IF EXISTS queue_version, + DROP COLUMN IF EXISTS generation_attempt, + DROP COLUMN IF EXISTS retry_state, + DROP COLUMN IF EXISTS retry_state_version, + DROP COLUMN IF EXISTS runner_id, + DROP COLUMN IF EXISTS requires_action_deadline_at; + +-- 10. Recreate chats_expanded with the pre-migration field list. +CREATE VIEW chats_expanded AS +SELECT + c.id, + c.owner_id, + c.workspace_id, + c.title, + c.status, + c.worker_id, + c.started_at, + c.heartbeat_at, + c.created_at, + c.updated_at, + c.parent_chat_id, + c.root_chat_id, + c.last_model_config_id, + c.archived, + c.last_error, + c.mode, + c.mcp_server_ids, + c.labels, + c.build_id, + c.agent_id, + c.pin_order, + c.last_read_message_id, + c.last_injected_context, + c.dynamic_tools, + c.organization_id, + c.plan_mode, + c.client_type, + c.last_turn_summary, + COALESCE(root.user_acl, c.user_acl) AS user_acl, + COALESCE(root.group_acl, c.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name +FROM + chats c + LEFT JOIN chats root ON root.id = COALESCE(c.root_chat_id, c.parent_chat_id) + JOIN visible_users owner ON owner.id = c.owner_id; + +-- 11. The `interrupting` chat_status enum value is intentionally left +-- in place. Postgres does not support dropping a single enum value +-- without recreating the entire type, which would require rewriting +-- every chat row and is unsafe inside a transactional rollback. diff --git a/coderd/database/migrations/000519_chatd_core_state_machine.up.sql b/coderd/database/migrations/000519_chatd_core_state_machine.up.sql new file mode 100644 index 0000000000..06277f0c9c --- /dev/null +++ b/coderd/database/migrations/000519_chatd_core_state_machine.up.sql @@ -0,0 +1,358 @@ +-- Adds the core chat state-machine storage model. +-- Adds new versioning fields to chats, a revision column to chat_messages, +-- positional ordering and creator tracking to chat_queued_messages, an +-- unlogged chat_heartbeats table for ownership leases, and Postgres +-- triggers that keep history/queue versioning consistent. + +-- 1. Add `interrupting` to the chat_status enum. +ALTER TYPE chat_status ADD VALUE IF NOT EXISTS 'interrupting'; + +-- 2. Add new versioning, ownership, retry, and pending-action fields to chats. +ALTER TABLE chats + ADD COLUMN snapshot_version bigint NOT NULL DEFAULT 1, + ADD COLUMN history_version bigint NOT NULL DEFAULT 0, + ADD COLUMN queue_version bigint NOT NULL DEFAULT 0, + ADD COLUMN generation_attempt bigint NOT NULL DEFAULT 0, + ADD COLUMN retry_state jsonb, + ADD COLUMN retry_state_version bigint NOT NULL DEFAULT 0, + ADD COLUMN runner_id uuid, + ADD COLUMN requires_action_deadline_at timestamp with time zone; + +COMMENT ON COLUMN chats.snapshot_version IS + 'Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet.'; +COMMENT ON COLUMN chats.history_version IS + 'Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version.'; +COMMENT ON COLUMN chats.queue_version IS + 'Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version.'; + +-- 3. Add `revision` to chat_messages. Adding the column as NOT NULL with +-- a constant default backfills existing rows through catalog metadata +-- only, so the highest-volume table is neither rewritten nor scanned for +-- NOT NULL validation while under ACCESS EXCLUSIVE. The default is +-- dropped immediately because the BEFORE INSERT trigger below rejects +-- inserts that pre-assign revision and assigns it from +-- chats.snapshot_version instead. +ALTER TABLE chat_messages + ADD COLUMN revision bigint NOT NULL DEFAULT 1; +ALTER TABLE chat_messages + ALTER COLUMN revision DROP DEFAULT; + +-- 4. Backfill chats.history_version = 1 for chats that already have at +-- least one message. We avoid recursive trigger fire by performing the +-- backfill before the triggers are created. +UPDATE chats +SET history_version = 1 +WHERE EXISTS ( + SELECT 1 FROM chat_messages WHERE chat_messages.chat_id = chats.id +); + +-- 5. Add `position` and `created_by` to chat_queued_messages. +ALTER TABLE chat_queued_messages + ADD COLUMN position bigint, + ADD COLUMN created_by uuid; + +-- 6. Backfill chat_queued_messages.position per chat using row_number(), +-- ordering by created_at and breaking ties by id. +WITH ordered AS ( + SELECT + id, + row_number() OVER ( + PARTITION BY chat_id + ORDER BY created_at, id + ) AS rn + FROM chat_queued_messages +) +UPDATE chat_queued_messages +SET position = ordered.rn +FROM ordered +WHERE chat_queued_messages.id = ordered.id; + +-- 7. Backfill chat_queued_messages.created_by from chats.owner_id. +UPDATE chat_queued_messages +SET created_by = chats.owner_id +FROM chats +WHERE chat_queued_messages.chat_id = chats.id + AND chat_queued_messages.created_by IS NULL; + +-- 8. Enforce NOT NULL on chat_queued_messages.position and +-- created_by. Legacy queued-message inserts are updated to populate +-- created_by from the chat owner when no explicit creator exists. +ALTER TABLE chat_queued_messages + ALTER COLUMN position SET NOT NULL, + ALTER COLUMN created_by SET NOT NULL; + +-- 9. Default sequence for new queued-message positions. +-- A global sequence is acceptable because ordering only needs to be +-- stable within a chat. +CREATE SEQUENCE IF NOT EXISTS chat_queued_messages_position_seq AS bigint START WITH 1; +SELECT setval( + 'chat_queued_messages_position_seq', + GREATEST((SELECT COALESCE(MAX(position), 0) FROM chat_queued_messages), 1) +); +ALTER TABLE chat_queued_messages + ALTER COLUMN position SET DEFAULT nextval('chat_queued_messages_position_seq'); + +-- 10. Backfill chats.queue_version = 1 for chats that already have queued +-- messages. Same trigger-avoidance reasoning as for history_version. +UPDATE chats +SET queue_version = 1 +WHERE EXISTS ( + SELECT 1 FROM chat_queued_messages WHERE chat_queued_messages.chat_id = chats.id +); + +-- 11. chat_heartbeats: unlogged table for ownership leases. Keyed by +-- (chat_id, runner_id) so a single chat can briefly have entries from +-- multiple runners during failover. +CREATE UNLOGGED TABLE IF NOT EXISTS chat_heartbeats ( + chat_id uuid NOT NULL REFERENCES chats(id) ON DELETE CASCADE, + runner_id uuid NOT NULL, + heartbeat_at timestamp with time zone NOT NULL, + PRIMARY KEY (chat_id, runner_id) +); + +COMMENT ON TABLE chat_heartbeats IS + 'Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats.'; + +CREATE INDEX IF NOT EXISTS chat_heartbeats_heartbeat_at_idx + ON chat_heartbeats (heartbeat_at); + +-- 12. Message revision trigger. +-- The BEFORE-trigger only assigns NEW.revision from chats.snapshot_version +-- and validates immutability. The chats.history_version / +-- generation_attempt update is performed by an AFTER STATEMENT trigger +-- so it doesn't conflict with CTE updates on the chats row in the same +-- command (the legacy InsertChatMessages query updates last_model_config_id +-- in a CTE on chats and then inserts messages). +CREATE FUNCTION set_chat_message_revision_before() +RETURNS trigger AS $$ +DECLARE + chat_snapshot_version bigint; +BEGIN + IF TG_OP = 'INSERT' AND NEW.revision IS NOT NULL THEN + RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger'; + END IF; + + IF TG_OP = 'UPDATE' THEN + IF OLD.chat_id IS DISTINCT FROM NEW.chat_id THEN + RAISE EXCEPTION 'chat_messages.chat_id is immutable'; + END IF; + + IF OLD.revision IS DISTINCT FROM NEW.revision THEN + RAISE EXCEPTION 'chat_messages.revision must be assigned by trigger'; + END IF; + + IF OLD IS NOT DISTINCT FROM NEW THEN + RETURN NEW; + END IF; + END IF; + + SELECT snapshot_version INTO chat_snapshot_version + FROM chats WHERE id = NEW.chat_id; + + IF chat_snapshot_version IS NULL THEN + RAISE EXCEPTION 'chat % does not exist', NEW.chat_id; + END IF; + + NEW.revision = chat_snapshot_version; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +-- AFTER STATEMENT trigger functions. Use the transition tables to +-- update chats.history_version / generation_attempt once per chat per +-- command. Running AFTER row inserts/updates complete lets a CTE +-- update on the same chats row in the same command finalize before +-- this trigger needs to update it. +-- +-- The INSERT and UPDATE variants are split so the UPDATE variant can +-- reference both the OLD and NEW transition tables and skip rows that +-- did not actually change. Without that filter, a no-op UPDATE on a +-- chat_messages row (one whose OLD IS NOT DISTINCT FROM NEW) would +-- still advance chats.history_version whenever the chat's snapshot +-- had previously been bumped. +CREATE FUNCTION update_chat_history_after_message_insert() +RETURNS trigger AS $$ +BEGIN + UPDATE chats c + SET history_version = c.snapshot_version, + generation_attempt = 0 + FROM ( + SELECT DISTINCT chat_id FROM chat_message_history_new_rows + ) AS affected + WHERE c.id = affected.chat_id + AND ( + c.history_version IS DISTINCT FROM c.snapshot_version + OR c.generation_attempt <> 0 + ); + RETURN NULL; +END; +$$ LANGUAGE plpgsql; + +CREATE FUNCTION update_chat_history_after_message_update() +RETURNS trigger AS $$ +BEGIN + UPDATE chats c + SET history_version = c.snapshot_version, + generation_attempt = 0 + FROM ( + SELECT DISTINCT n.chat_id + FROM chat_message_history_new_rows n + JOIN chat_message_history_old_rows o ON o.id = n.id + WHERE o IS DISTINCT FROM n + ) AS affected + WHERE c.id = affected.chat_id + AND ( + c.history_version IS DISTINCT FROM c.snapshot_version + OR c.generation_attempt <> 0 + ); + RETURN NULL; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_set_chat_message_revision_on_insert +BEFORE INSERT ON chat_messages +FOR EACH ROW +EXECUTE FUNCTION set_chat_message_revision_before(); + +CREATE TRIGGER trigger_set_chat_message_revision_on_update +BEFORE UPDATE ON chat_messages +FOR EACH ROW +EXECUTE FUNCTION set_chat_message_revision_before(); + +CREATE TRIGGER trigger_update_chat_history_after_message_insert +AFTER INSERT ON chat_messages +REFERENCING NEW TABLE AS chat_message_history_new_rows +FOR EACH STATEMENT +EXECUTE FUNCTION update_chat_history_after_message_insert(); + +CREATE TRIGGER trigger_update_chat_history_after_message_update +AFTER UPDATE ON chat_messages +REFERENCING OLD TABLE AS chat_message_history_old_rows NEW TABLE AS chat_message_history_new_rows +FOR EACH STATEMENT +EXECUTE FUNCTION update_chat_history_after_message_update(); + +-- 13. Queue version trigger function. +CREATE FUNCTION bump_chat_queue_version_on_queued_message_change() +RETURNS trigger AS $$ +DECLARE + changed_chat_id uuid; +BEGIN + IF TG_OP = 'DELETE' THEN + changed_chat_id = OLD.chat_id; + ELSE + changed_chat_id = NEW.chat_id; + END IF; + + UPDATE chats + SET queue_version = snapshot_version + WHERE id = changed_chat_id; + + IF TG_OP = 'DELETE' THEN + RETURN OLD; + END IF; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_insert +AFTER INSERT ON chat_queued_messages +FOR EACH ROW +EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_update +AFTER UPDATE OF content, model_config_id, position, created_by +ON chat_queued_messages +FOR EACH ROW +EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + +CREATE TRIGGER trigger_bump_chat_queue_version_on_queued_message_delete +AFTER DELETE ON chat_queued_messages +FOR EACH ROW +EXECUTE FUNCTION bump_chat_queue_version_on_queued_message_change(); + +-- 14. Retry state trigger function. +CREATE FUNCTION sync_chat_retry_state() +RETURNS trigger AS $$ +BEGIN + IF OLD.retry_state_version IS DISTINCT FROM NEW.retry_state_version THEN + RAISE EXCEPTION 'chats.retry_state_version must be assigned by trigger'; + END IF; + + IF NEW.generation_attempt IS DISTINCT FROM OLD.generation_attempt THEN + NEW.retry_state = NULL; + END IF; + + IF NEW.retry_state IS DISTINCT FROM OLD.retry_state THEN + NEW.retry_state_version = NEW.snapshot_version; + END IF; + + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +CREATE TRIGGER trigger_sync_chat_retry_state +BEFORE UPDATE OF retry_state, retry_state_version, generation_attempt +ON chats +FOR EACH ROW +EXECUTE FUNCTION sync_chat_retry_state(); + +-- 15. Index for the chat worker acquisition scan, which runs every 30 +-- seconds per replica plus on every worker wake. Leading on status lets +-- the scan touch only rows in the worker-runnable status set instead of +-- sequentially scanning the ever-growing chats table. The status set is +-- intentionally not part of the index predicate: 'interrupting' is added +-- to chat_status above, and Postgres forbids using a new enum value in +-- the same transaction, which all migrations share. +CREATE INDEX idx_chats_worker_acquisition_candidates ON chats + USING btree (status, updated_at, id) + WHERE archived = false; + +-- 16. Refresh chats_expanded to include the new chat fields. Drop and +-- recreate so column ordering is stable. +DROP VIEW IF EXISTS chats_expanded; +CREATE VIEW chats_expanded AS +SELECT + c.id, + c.owner_id, + c.workspace_id, + c.title, + c.status, + c.worker_id, + c.started_at, + c.heartbeat_at, + c.created_at, + c.updated_at, + c.parent_chat_id, + c.root_chat_id, + c.last_model_config_id, + c.archived, + c.last_error, + c.mode, + c.mcp_server_ids, + c.labels, + c.build_id, + c.agent_id, + c.pin_order, + c.last_read_message_id, + c.last_injected_context, + c.dynamic_tools, + c.organization_id, + c.plan_mode, + c.client_type, + c.last_turn_summary, + c.snapshot_version, + c.history_version, + c.queue_version, + c.generation_attempt, + c.retry_state, + c.retry_state_version, + c.runner_id, + c.requires_action_deadline_at, + COALESCE(root.user_acl, c.user_acl) AS user_acl, + COALESCE(root.group_acl, c.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name +FROM + chats c + LEFT JOIN chats root ON root.id = COALESCE(c.root_chat_id, c.parent_chat_id) + JOIN visible_users owner ON owner.id = c.owner_id; diff --git a/coderd/database/migrations/testdata/fixtures/000519_chatd_core_state_machine.up.sql b/coderd/database/migrations/testdata/fixtures/000519_chatd_core_state_machine.up.sql new file mode 100644 index 0000000000..31ce67fd1b --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000519_chatd_core_state_machine.up.sql @@ -0,0 +1,17 @@ +-- Fixture coverage for the chat_heartbeats table introduced in +-- migration 000500. The earlier chat fixtures already insert at least +-- one row into chats; we attach a heartbeat for the first such chat so +-- migration tests see a non-empty chat_heartbeats table without +-- hard-coding a specific chat ID. +INSERT INTO chat_heartbeats ( + chat_id, + runner_id, + heartbeat_at +) +SELECT + chats.id, + '00000000-0000-0000-0000-0000000fea51'::uuid, + '2024-01-01 00:00:00+00' +FROM chats +ORDER BY created_at, id +LIMIT 1; diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index f0645220f4..d89156d654 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -824,6 +824,14 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams, &i.Chat.PlanMode, &i.Chat.ClientType, &i.Chat.LastTurnSummary, + &i.Chat.SnapshotVersion, + &i.Chat.HistoryVersion, + &i.Chat.QueueVersion, + &i.Chat.GenerationAttempt, + &i.Chat.RetryState, + &i.Chat.RetryStateVersion, + &i.Chat.RunnerID, + &i.Chat.RequiresActionDeadlineAt, &i.Chat.UserACL, &i.Chat.GroupACL, &i.Chat.OwnerUsername, @@ -891,6 +899,14 @@ func (q *sqlQuerier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, diff --git a/coderd/database/models.go b/coderd/database/models.go index 10ab50e5dc..e52d5f6c4a 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -1567,6 +1567,7 @@ const ( ChatStatusCompleted ChatStatus = "completed" ChatStatusError ChatStatus = "error" ChatStatusRequiresAction ChatStatus = "requires_action" + ChatStatusInterrupting ChatStatus = "interrupting" ) func (e *ChatStatus) Scan(src interface{}) error { @@ -1612,7 +1613,8 @@ func (e ChatStatus) Valid() bool { ChatStatusPaused, ChatStatusCompleted, ChatStatusError, - ChatStatusRequiresAction: + ChatStatusRequiresAction, + ChatStatusInterrupting: return true } return false @@ -1627,6 +1629,7 @@ func AllChatStatusValues() []ChatStatus { ChatStatusCompleted, ChatStatusError, ChatStatusRequiresAction, + ChatStatusInterrupting, } } @@ -4607,38 +4610,46 @@ type BoundaryUsageStat struct { } type Chat struct { - ID uuid.UUID `db:"id" json:"id"` - OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` - WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` - Title string `db:"title" json:"title"` - Status ChatStatus `db:"status" json:"status"` - WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` - StartedAt sql.NullTime `db:"started_at" json:"started_at"` - HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` - CreatedAt time.Time `db:"created_at" json:"created_at"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` - ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` - RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` - LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` - Archived bool `db:"archived" json:"archived"` - LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` - Mode NullChatMode `db:"mode" json:"mode"` - MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` - Labels StringMap `db:"labels" json:"labels"` - BuildID uuid.NullUUID `db:"build_id" json:"build_id"` - AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` - PinOrder int32 `db:"pin_order" json:"pin_order"` - LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` - LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` - DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` - OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` - PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` - ClientType ChatClientType `db:"client_type" json:"client_type"` - LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` - UserACL ChatACL `db:"user_acl" json:"user_acl"` - GroupACL ChatACL `db:"group_acl" json:"group_acl"` - OwnerUsername string `db:"owner_username" json:"owner_username"` - OwnerName string `db:"owner_name" json:"owner_name"` + ID uuid.UUID `db:"id" json:"id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + Title string `db:"title" json:"title"` + Status ChatStatus `db:"status" json:"status"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` + StartedAt sql.NullTime `db:"started_at" json:"started_at"` + HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` + RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` + LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` + Archived bool `db:"archived" json:"archived"` + LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` + Mode NullChatMode `db:"mode" json:"mode"` + MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` + Labels StringMap `db:"labels" json:"labels"` + BuildID uuid.NullUUID `db:"build_id" json:"build_id"` + AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` + PinOrder int32 `db:"pin_order" json:"pin_order"` + LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` + LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` + DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` + ClientType ChatClientType `db:"client_type" json:"client_type"` + LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + HistoryVersion int64 `db:"history_version" json:"history_version"` + QueueVersion int64 `db:"queue_version" json:"queue_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` + UserACL ChatACL `db:"user_acl" json:"user_acl"` + GroupACL ChatACL `db:"group_acl" json:"group_acl"` + OwnerUsername string `db:"owner_username" json:"owner_username"` + OwnerName string `db:"owner_name" json:"owner_name"` } type ChatDebugRun struct { @@ -4720,6 +4731,13 @@ type ChatFileLink struct { FileID uuid.UUID `db:"file_id" json:"file_id"` } +// Ephemeral runner ownership leases for runnable chats. The table is unlogged because losing heartbeat rows after a crash is safe: missing heartbeats are treated as stale ownership and cause workers to reacquire runnable chats. +type ChatHeartbeat struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + RunnerID uuid.UUID `db:"runner_id" json:"runner_id"` + HeartbeatAt time.Time `db:"heartbeat_at" json:"heartbeat_at"` +} + type ChatMessage struct { ID int64 `db:"id" json:"id"` ChatID uuid.UUID `db:"chat_id" json:"chat_id"` @@ -4743,6 +4761,7 @@ type ChatMessage struct { Deleted bool `db:"deleted" json:"deleted"` ProviderResponseID sql.NullString `db:"provider_response_id" json:"provider_response_id"` APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` + Revision int64 `db:"revision" json:"revision"` } type ChatModelConfig struct { @@ -4771,6 +4790,8 @@ type ChatQueuedMessage struct { CreatedAt time.Time `db:"created_at" json:"created_at"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` + Position int64 `db:"position" json:"position"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` } type ChatTable struct { @@ -4804,6 +4825,17 @@ type ChatTable struct { LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` UserACL ChatACL `db:"user_acl" json:"user_acl"` GroupACL ChatACL `db:"group_acl" json:"group_acl"` + // Monotonic version for the full chat snapshot. Starts at 1 so stream loops and workers can use 0 to mean they have not loaded the chat yet. + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + // Snapshot version of the latest durable history change. Starts at 0 until chat_messages triggers set it to the current snapshot_version. + HistoryVersion int64 `db:"history_version" json:"history_version"` + // Snapshot version of the latest queued-message change. Starts at 0 until chat_queued_messages triggers set it to the current snapshot_version. + QueueVersion int64 `db:"queue_version" json:"queue_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` } type ChatUsageLimitConfig struct { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 63f95141af..39673e7b04 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -77,9 +77,12 @@ type sqlcQuerier interface { // enforces that non-deleted rows always have a provider ID. BackfillChatModelConfigProvider(ctx context.Context, arg BackfillChatModelConfigProviderParams) (sql.Result, error) BackoffChatDiffStatus(ctx context.Context, arg BackoffChatDiffStatusParams) error + // Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs. + BatchDeleteChatHeartbeats(ctx context.Context, arg BatchDeleteChatHeartbeatsParams) (int64, error) BatchUpdateWorkspaceAgentMetadata(ctx context.Context, arg BatchUpdateWorkspaceAgentMetadataParams) error BatchUpdateWorkspaceLastUsedAt(ctx context.Context, arg BatchUpdateWorkspaceLastUsedAtParams) error BatchUpdateWorkspaceNextStartAt(ctx context.Context, arg BatchUpdateWorkspaceNextStartAtParams) error + BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUpsertChatHeartbeatsParams) error BatchUpsertConnectionLogs(ctx context.Context, arg BatchUpsertConnectionLogsParams) error BulkMarkNotificationMessagesFailed(ctx context.Context, arg BulkMarkNotificationMessagesFailedParams) (int64, error) BulkMarkNotificationMessagesSent(ctx context.Context, arg BulkMarkNotificationMessagesSentParams) (int64, error) @@ -94,6 +97,9 @@ type sqlcQuerier interface { ClearChatMessageProviderResponseIDsByChatID(ctx context.Context, chatID uuid.UUID) error CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error) CountAuditLogs(ctx context.Context, arg CountAuditLogsParams) (int64, error) + // Cheap queue-length check used by ChatMachine.Update when deciding + // whether the chat is in a "1" sub-state. + CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) CountConnectionLogs(ctx context.Context, arg CountConnectionLogsParams) (int64, error) // Counts enabled, non-deleted model configs that lack both input and // output pricing in their JSONB options.cost configuration. @@ -111,7 +117,11 @@ type sqlcQuerier interface { DeleteAIProviderKey(ctx context.Context, id uuid.UUID) error DeleteAPIKeyByID(ctx context.Context, id string) error DeleteAPIKeysByUserID(ctx context.Context, userID uuid.UUID) error + // Deletes all heartbeat rows for the chat. Used during ownership + // transitions that abandon a lease. + DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error DeleteAllChatQueuedMessages(ctx context.Context, chatID uuid.UUID) error + DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) DeleteAllTailnetTunnels(ctx context.Context, arg DeleteAllTailnetTunnelsParams) ([]DeleteAllTailnetTunnelsRow, error) // Deletes all existing webpush subscriptions. // This should be called when the VAPID keypair is regenerated, as the old @@ -133,6 +143,10 @@ type sqlcQuerier interface { DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error DeleteChatModelConfigsByProvider(ctx context.Context, provider string) error DeleteChatQueuedMessage(ctx context.Context, arg DeleteChatQueuedMessageParams) error + // Deletes a queued message, scoped to the parent chat. Returns the + // number of affected rows so callers can detect missing rows without + // a follow-up read. + DeleteChatQueuedMessageReturningCount(ctx context.Context, arg DeleteChatQueuedMessageReturningCountParams) (int64, error) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error DeleteCryptoKey(ctx context.Context, arg DeleteCryptoKeyParams) (CryptoKey, error) @@ -201,6 +215,7 @@ type sqlcQuerier interface { DeleteProvisionerKey(ctx context.Context, id uuid.UUID) error DeleteReplicasUpdatedBefore(ctx context.Context, updatedAt time.Time) error DeleteRuntimeConfig(ctx context.Context, key string) error + DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) DeleteTailnetPeer(ctx context.Context, arg DeleteTailnetPeerParams) (DeleteTailnetPeerRow, error) DeleteTailnetTunnel(ctx context.Context, arg DeleteTailnetTunnelParams) (DeleteTailnetTunnelRow, error) DeleteTask(ctx context.Context, arg DeleteTaskParams) (uuid.UUID, error) @@ -323,6 +338,10 @@ type sqlcQuerier interface { // This function returns roles for authorization purposes. Implied member roles // are included. GetAuthorizationUserRoles(ctx context.Context, userID uuid.UUID) (GetAuthorizationUserRolesRow, error) + // Returns read-only root chat candidates for state-machine-backed + // auto-archive. Activity is computed across the root family. The query + // limits roots, not total family members. + GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg GetAutoArchiveInactiveChatCandidatesParams) ([]GetAutoArchiveInactiveChatCandidatesRow, error) GetBoundaryLogByID(ctx context.Context, id uuid.UUID) (BoundaryLog, error) GetBoundarySessionByID(ctx context.Context, id uuid.UUID) (BoundarySession, error) GetChatACLByID(ctx context.Context, id uuid.UUID) (GetChatACLByIDRow, error) @@ -334,6 +353,7 @@ type sqlcQuerier interface { // Auto-archive window in days. 0 disables. GetChatAutoArchiveDays(ctx context.Context, defaultAutoArchiveDays int32) (int32, error) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error) + GetChatByIDForShare(ctx context.Context, id uuid.UUID) (Chat, error) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Chat, error) GetChatComputerUseProvider(ctx context.Context) (string, error) // Per-root-chat cost breakdown for a single user within a date range. @@ -371,6 +391,10 @@ type sqlcQuerier interface { GetChatDiffStatusSummary(ctx context.Context) (GetChatDiffStatusSummaryRow, error) GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds []uuid.UUID) ([]ChatDiffStatus, error) GetChatExploreModelOverride(ctx context.Context) (string, error) + // Returns the chat IDs of every chat in a family (root + all children) + // in deterministic order. The id parameter must be the root id; the + // query does not walk up from a child. + GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) GetChatFileByID(ctx context.Context, id uuid.UUID) (ChatFile, error) // GetChatFileMetadataByChatID returns lightweight file metadata for // all files linked to a chat. The data column is excluded to avoid @@ -378,6 +402,7 @@ type sqlcQuerier interface { GetChatFileMetadataByChatID(ctx context.Context, chatID uuid.UUID) ([]GetChatFileMetadataByChatIDRow, error) GetChatFilesByIDs(ctx context.Context, ids []uuid.UUID) ([]ChatFile, error) GetChatGeneralModelOverride(ctx context.Context) (string, error) + GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatParams) (ChatHeartbeat, error) // GetChatIncludeDefaultSystemPrompt preserves the legacy default // for deployments created before the explicit include-default toggle. // When the toggle is unset, a non-empty custom prompt implies false; @@ -391,6 +416,7 @@ type sqlcQuerier interface { GetChatMessagesByChatID(ctx context.Context, arg GetChatMessagesByChatIDParams) ([]ChatMessage, error) GetChatMessagesByChatIDAscPaginated(ctx context.Context, arg GetChatMessagesByChatIDAscPaginatedParams) ([]ChatMessage, error) GetChatMessagesByChatIDDescPaginated(ctx context.Context, arg GetChatMessagesByChatIDDescPaginatedParams) ([]ChatMessage, error) + GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error) @@ -400,12 +426,18 @@ type sqlcQuerier interface { // personal chat model overrides. It defaults to false when unset. GetChatPersonalModelOverridesEnabled(ctx context.Context) (bool, error) GetChatPlanModeInstructions(ctx context.Context) (string, error) + GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error) + // Returns the queue head (lowest position, then lowest id). + GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) + // Returns queued messages in state-machine order (position ASC, id ASC). + GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) // Returns the chat retention period in days. Chats archived longer // than this and orphaned chat files older than this are purged by // dbpurge. Returns 30 (days) when no value has been configured. // A value of 0 disables chat purging entirely. GetChatRetentionDays(ctx context.Context) (int32, error) + GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]GetChatStreamSyncRowsRow, error) GetChatSystemPrompt(ctx context.Context) (string, error) // GetChatSystemPromptConfig returns both chat system prompt settings in a // single read to avoid torn reads between separate site-config lookups. @@ -430,11 +462,24 @@ type sqlcQuerier interface { // jsonb_array_elements never raises "cannot extract elements from a // scalar". Backed by idx_chat_messages_user_prompts. GetChatUserPromptsByChatID(ctx context.Context, arg GetChatUserPromptsByChatIDParams) ([]GetChatUserPromptsByChatIDRow, error) + // Returns chats that workers may try to acquire. Candidates must be: + // - in a worker-runnable execution status; + // - unarchived; and + // - missing ownership, carrying inconsistent ownership, or lacking a + // fresh heartbeat for the assigned runner. + // + // Missing ownership is worker_id IS NULL. Inconsistent ownership is + // runner_id IS NULL while worker_id is set. Stale ownership is no + // heartbeat row for (chat_id, runner_id), or one older than + // @stale_seconds by database time. Candidates are ordered by oldest + // updated_at first so workers drain stale runnable chats predictably. + GetChatWorkerAcquisitionCandidates(ctx context.Context, arg GetChatWorkerAcquisitionCandidatesParams) ([]GetChatWorkerAcquisitionCandidatesRow, error) // Returns the global TTL for chat workspaces as a Go duration string. // Returns "0s" (disabled) when no value has been configured. GetChatWorkspaceTTL(ctx context.Context) (string, error) GetChats(ctx context.Context, arg GetChatsParams) ([]GetChatsRow, error) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) ([]Chat, error) + GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]Chat, error) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID) ([]Chat, error) // Retrieves chats updated after the given timestamp for telemetry // snapshot collection. Uses updated_at so that long-running chats @@ -451,6 +496,10 @@ type sqlcQuerier interface { GetCryptoKeysByFeature(ctx context.Context, feature CryptoKeyFeature) ([]CryptoKey, error) GetDBCryptKeys(ctx context.Context) ([]DBCryptKey, error) GetDERPMeshKey(ctx context.Context) (string, error) + // Returns the current database timestamp. Used so transitions that + // record deadlines or heartbeats rely on a clock that is consistent + // with the database rather than the caller's local clock. + GetDatabaseNow(ctx context.Context) (time.Time, error) GetDefaultChatModelConfig(ctx context.Context) (ChatModelConfig, error) GetDefaultOrganization(ctx context.Context) (Organization, error) GetDefaultProxyConfig(ctx context.Context) (GetDefaultProxyConfigRow, error) @@ -931,6 +980,8 @@ type sqlcQuerier interface { GetWorkspacesByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceTable, error) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]GetWorkspacesEligibleForTransitionRow, error) GetWorkspacesForWorkspaceMetrics(ctx context.Context) ([]GetWorkspacesForWorkspaceMetricsRow, error) + // Increments generation_attempt and returns the resulting value. + IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) InsertAIBridgeModelThought(ctx context.Context, arg InsertAIBridgeModelThoughtParams) (AIBridgeModelThought, error) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) @@ -961,7 +1012,14 @@ type sqlcQuerier interface { InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) + // Legacy queue insertion path. When no caller-supplied creator exists, + // preserve the created_by invariant by attributing the queued row to the + // chat owner. InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) + // Inserts a queued message that carries a position (from the default + // sequence) and an explicit created_by reference. Use this when the + // queued-message creator differs from the chat owner. + InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error) InsertCryptoKey(ctx context.Context, arg InsertCryptoKeyParams) (CryptoKey, error) InsertCustomRole(ctx context.Context, arg InsertCustomRoleParams) (CustomRole, error) InsertDBCryptKey(ctx context.Context, arg InsertDBCryptKeyParams) error @@ -1041,6 +1099,11 @@ type sqlcQuerier interface { InsertWorkspaceProxy(ctx context.Context, arg InsertWorkspaceProxyParams) (WorkspaceProxy, error) InsertWorkspaceResource(ctx context.Context, arg InsertWorkspaceResourceParams) (WorkspaceResource, error) InsertWorkspaceResourceMetadata(ctx context.Context, arg InsertWorkspaceResourceMetadataParams) ([]WorkspaceResourceMetadatum, error) + // Returns true when there is no heartbeat row for (chat_id, runner_id) + // or the existing row is older than @stale_seconds seconds by database + // time. chatstate calls this in a single query so the staleness check + // is atomic and does not depend on the caller's local clock. + IsChatHeartbeatStale(ctx context.Context, arg IsChatHeartbeatStaleParams) (bool, error) // LinkChatFiles inserts file associations into the chat_file_links // join table with deduplication (ON CONFLICT DO NOTHING). The INSERT // is conditional: it only proceeds when the total number of links @@ -1091,6 +1154,11 @@ type sqlcQuerier interface { ListUserSecretsWithValues(ctx context.Context, userID uuid.UUID) ([]UserSecret, error) ListUserSkillMetadataByUserID(ctx context.Context, userID uuid.UUID) ([]ListUserSkillMetadataByUserIDRow, error) ListWorkspaceAgentPortShares(ctx context.Context, workspaceID uuid.UUID) ([]WorkspaceAgentPortShare, error) + // Locks the chat row with FOR UPDATE and atomically increments its + // snapshot_version, returning the post-bump chat. This is the single + // entry point ChatMachine.Update uses to acquire the row lock and + // allocate a new snapshot version in one round trip. + LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (Chat, error) MarkAllInboxNotificationsAsRead(ctx context.Context, arg MarkAllInboxNotificationsAsReadParams) error OIDCClaimFieldValues(ctx context.Context, arg OIDCClaimFieldValuesParams) ([]string, error) // OIDCClaimFields returns a list of distinct keys in the the merged_claims fields. @@ -1115,6 +1183,9 @@ type sqlcQuerier interface { // Mutates only created_at on the target row; ids are unchanged so // consumers can keep tracking queued messages by id. ReorderChatQueuedMessageToFront(ctx context.Context, arg ReorderChatQueuedMessageToFrontParams) (int64, error) + // Sets the target queued message's position to one less than the + // current minimum position for that chat, moving it to the head. + ReorderChatQueuedMessageToHead(ctx context.Context, arg ReorderChatQueuedMessageToHeadParams) (int64, error) // Resolves the effective spend limit for a user using the hierarchy: // 1. Individual user override (highest priority, applies globally across // all organizations since it lives on the users table) @@ -1216,6 +1287,11 @@ type sqlcQuerier interface { // parameter keeps updated_at under the caller's clock, matching // the injectable quartz.Clock used by FinalizeStale sweeps. UpdateChatDebugStep(ctx context.Context, arg UpdateChatDebugStepParams) (ChatDebugStep, error) + // Atomically updates the execution-state-managed fields on a chat: + // status, archived, last_error, ownership identifiers, and the + // requires-action deadline. Callers compose this with transition + // mutations inside a single ChatMachine.Update transaction. + UpdateChatExecutionState(ctx context.Context, arg UpdateChatExecutionStateParams) (Chat, error) // Bumps the heartbeat timestamp for the given set of chat IDs, // provided they are still running and owned by the specified // worker. Returns the IDs that were actually updated so the @@ -1234,11 +1310,9 @@ type sqlcQuerier interface { // Updates the cached last completed turn summary for sidebar display. // Empty or whitespace-only summaries are stored as NULL here so direct // query callers cannot accidentally persist blank sidebar text. - // This intentionally preserves updated_at. The staleness guard relies on - // every new-turn query, such as UpdateChatStatus and AcquireChats, bumping - // updated_at. Future chat-field updates that do not bump updated_at can let - // stale summaries persist. If this query ever bumps updated_at, later - // goroutine summary writes will be rejected as stale. + // This intentionally preserves updated_at. The staleness guard uses + // history_version so worker lifecycle transitions that do not change the + // active message history cannot reject final turn summary writes. // Two summary workers using the same freshness marker are last-write-wins. UpdateChatLastTurnSummary(ctx context.Context, arg UpdateChatLastTurnSummaryParams) (int64, error) UpdateChatMCPServerIDs(ctx context.Context, arg UpdateChatMCPServerIDsParams) (Chat, error) @@ -1246,6 +1320,9 @@ type sqlcQuerier interface { UpdateChatModelConfig(ctx context.Context, arg UpdateChatModelConfigParams) (ChatModelConfig, error) UpdateChatPinOrder(ctx context.Context, arg UpdateChatPinOrderParams) error UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatPlanModeByIDParams) (Chat, error) + // Stores the client-visible retry payload. retry_state_version is + // assigned by trigger from the current snapshot_version. + UpdateChatRetryState(ctx context.Context, arg UpdateChatRetryStateParams) (Chat, error) UpdateChatStatus(ctx context.Context, arg UpdateChatStatusParams) (Chat, error) UpdateChatStatusPreserveUpdatedAt(ctx context.Context, arg UpdateChatStatusPreserveUpdatedAtParams) (Chat, error) UpdateChatTitleByID(ctx context.Context, arg UpdateChatTitleByIDParams) (Chat, error) @@ -1396,6 +1473,9 @@ type sqlcQuerier interface { UpsertChatDiffStatusReference(ctx context.Context, arg UpsertChatDiffStatusReferenceParams) (ChatDiffStatus, error) UpsertChatExploreModelOverride(ctx context.Context, value string) error UpsertChatGeneralModelOverride(ctx context.Context, value string) error + // Upserts a heartbeat row for the (chat_id, runner_id) lease. Uses + // database time so callers do not depend on a local clock. + UpsertChatHeartbeat(ctx context.Context, arg UpsertChatHeartbeatParams) error UpsertChatIncludeDefaultSystemPrompt(ctx context.Context, includeDefaultSystemPrompt bool) error // UpsertChatPersonalModelOverridesEnabled updates whether users may configure // personal chat model overrides. diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c48088314b..23f170067d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -11063,10 +11063,12 @@ func TestInsertChatMessages(t *testing.T) { insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) { t.Helper() + apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: userID}) _, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chatID, CreatedBy: []uuid.UUID{userID}, + APIKeyID: []string{apiKey.ID}, ModelConfigID: []uuid.UUID{modelConfigID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.CurrentContentVersion}, @@ -11125,10 +11127,12 @@ func TestInsertChatMessages(t *testing.T) { t.Parallel() store, ctx, user, chat, _, modelConfigA := setupChat(t) + apiKey, _ := dbgen.APIKey(t, store, database.APIKey{ID: uuid.NewString(), UserID: user.ID}) msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{user.ID, uuid.Nil, uuid.Nil}, + APIKeyID: []string{apiKey.ID, "", ""}, ModelConfigID: []uuid.UUID{modelConfigA.ID, modelConfigA.ID, modelConfigA.ID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser, database.ChatMessageRoleAssistant, database.ChatMessageRoleTool}, ContentVersion: []int16{chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion, chatprompt.CurrentContentVersion}, @@ -12744,9 +12748,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) { require.NoError(t, err) affected, err := db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: "resolved the issue", Valid: true}, }) require.NoError(t, err) require.EqualValues(t, 1, affected) @@ -12757,9 +12761,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) { require.Equal(t, chat.UpdatedAt, fetched.UpdatedAt) affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: " \n\t ", Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: " \n\t ", Valid: true}, }) require.NoError(t, err) require.EqualValues(t, 1, affected) @@ -12770,9 +12774,9 @@ func TestUpdateChatLastTurnSummary(t *testing.T) { require.Equal(t, chat.UpdatedAt, fetched.UpdatedAt) affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: "fresh summary", Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: "fresh summary", Valid: true}, }) require.NoError(t, err) require.EqualValues(t, 1, affected) @@ -12786,17 +12790,54 @@ func TestUpdateChatLastTurnSummary(t *testing.T) { require.NoError(t, err) affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: "stale summary", Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: "still fresh summary", Valid: true}, + }) + require.NoError(t, err) + require.EqualValues(t, 1, affected) + + fetched, err = db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, sql.NullString{String: "still fresh summary", Valid: true}, fetched.LastTurnSummary) + require.Equal(t, advancedUpdatedAt, fetched.UpdatedAt) + + _, err = db.LockChatAndBumpSnapshotVersion(ctx, chat.ID) + require.NoError(t, err) + _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{owner.ID}, + ModelConfigID: []uuid.UUID{modelCfg.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, + Content: []string{`[{"type":"text","text":"new request"}]`}, + ContentVersion: []int16{chatprompt.CurrentContentVersion}, + Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, + InputTokens: []int64{0}, + OutputTokens: []int64{0}, + TotalTokens: []int64{0}, + ReasoningTokens: []int64{0}, + CacheCreationTokens: []int64{0}, + CacheReadTokens: []int64{0}, + ContextLimit: []int64{0}, + Compressed: []bool{false}, + TotalCostMicros: []int64{0}, + RuntimeMs: []int64{0}, + ProviderResponseID: []string{""}, + }) + require.NoError(t, err) + + affected, err = db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: "stale summary", Valid: true}, }) require.NoError(t, err) require.Zero(t, affected) fetched, err = db.GetChatByID(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, sql.NullString{String: "fresh summary", Valid: true}, fetched.LastTurnSummary) - require.Equal(t, advancedUpdatedAt, fetched.UpdatedAt) + require.Equal(t, sql.NullString{String: "still fresh summary", Valid: true}, fetched.LastTurnSummary) + require.NotEqual(t, chat.HistoryVersion, fetched.HistoryVersion) } func TestDeleteChatDebugDataAfterMessageIDIncludesTriggeredRuns(t *testing.T) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 3cc6884f0b..cb2dfd3ce6 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -5837,7 +5837,7 @@ WHERE LIMIT $3::int ) -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -5869,6 +5869,14 @@ chats_expanded AS ( acquired_chats.plan_mode, acquired_chats.client_type, acquired_chats.last_turn_summary, + acquired_chats.snapshot_version, + acquired_chats.history_version, + acquired_chats.queue_version, + acquired_chats.generation_attempt, + acquired_chats.retry_state, + acquired_chats.retry_state_version, + acquired_chats.runner_id, + acquired_chats.requires_action_deadline_at, COALESCE(root.user_acl, acquired_chats.user_acl) AS user_acl, COALESCE(root.group_acl, acquired_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -5878,7 +5886,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(acquired_chats.root_chat_id, acquired_chats.parent_chat_id) JOIN visible_users owner ON owner.id = acquired_chats.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -5928,6 +5936,14 @@ func (q *sqlQuerier) AcquireChats(ctx context.Context, arg AcquireChatsParams) ( &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -6070,7 +6086,7 @@ WITH updated_chats AS ( UPDATE chats SET archived = true, pin_order = 0, updated_at = NOW() WHERE id = $1::uuid OR root_chat_id = $1::uuid - RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -6102,6 +6118,14 @@ chats_expanded AS ( updated_chats.plan_mode, updated_chats.client_type, updated_chats.last_turn_summary, + updated_chats.snapshot_version, + updated_chats.history_version, + updated_chats.queue_version, + updated_chats.generation_attempt, + updated_chats.retry_state, + updated_chats.retry_state_version, + updated_chats.runner_id, + updated_chats.requires_action_deadline_at, COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -6111,7 +6135,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chats.root_chat_id, updated_chats.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chats.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ORDER BY (chats_expanded.id = $1::uuid) DESC, chats_expanded.created_at ASC, chats_expanded.id ASC ` @@ -6154,6 +6178,14 @@ func (q *sqlQuerier) ArchiveChatByID(ctx context.Context, id uuid.UUID) ([]Chat, &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -6208,10 +6240,10 @@ archived AS ( FROM to_archive t WHERE (c.id = t.id OR c.root_chat_id = t.id) -- cascade to children AND c.archived = false - RETURNING c.id, c.owner_id, c.workspace_id, c.title, c.status, c.worker_id, c.started_at, c.heartbeat_at, c.created_at, c.updated_at, c.parent_chat_id, c.root_chat_id, c.last_model_config_id, c.archived, c.last_error, c.mode, c.mcp_server_ids, c.labels, c.build_id, c.agent_id, c.pin_order, c.last_read_message_id, c.last_injected_context, c.dynamic_tools, c.organization_id, c.plan_mode, c.client_type, c.last_turn_summary, c.user_acl, c.group_acl + RETURNING c.id, c.owner_id, c.workspace_id, c.title, c.status, c.worker_id, c.started_at, c.heartbeat_at, c.created_at, c.updated_at, c.parent_chat_id, c.root_chat_id, c.last_model_config_id, c.archived, c.last_error, c.mode, c.mcp_server_ids, c.labels, c.build_id, c.agent_id, c.pin_order, c.last_read_message_id, c.last_injected_context, c.dynamic_tools, c.organization_id, c.plan_mode, c.client_type, c.last_turn_summary, c.user_acl, c.group_acl, c.snapshot_version, c.history_version, c.queue_version, c.generation_attempt, c.retry_state, c.retry_state_version, c.runner_id, c.requires_action_deadline_at ) SELECT - a.id, a.owner_id, a.workspace_id, a.title, a.status, a.worker_id, a.started_at, a.heartbeat_at, a.created_at, a.updated_at, a.parent_chat_id, a.root_chat_id, a.last_model_config_id, a.archived, a.last_error, a.mode, a.mcp_server_ids, a.labels, a.build_id, a.agent_id, a.pin_order, a.last_read_message_id, a.last_injected_context, a.dynamic_tools, a.organization_id, a.plan_mode, a.client_type, a.last_turn_summary, a.user_acl, a.group_acl, + a.id, a.owner_id, a.workspace_id, a.title, a.status, a.worker_id, a.started_at, a.heartbeat_at, a.created_at, a.updated_at, a.parent_chat_id, a.root_chat_id, a.last_model_config_id, a.archived, a.last_error, a.mode, a.mcp_server_ids, a.labels, a.build_id, a.agent_id, a.pin_order, a.last_read_message_id, a.last_injected_context, a.dynamic_tools, a.organization_id, a.plan_mode, a.client_type, a.last_turn_summary, a.user_acl, a.group_acl, a.snapshot_version, a.history_version, a.queue_version, a.generation_attempt, a.retry_state, a.retry_state_version, a.runner_id, a.requires_action_deadline_at, -- Children inherit their root's activity so last_activity_at is never null. COALESCE( t.last_activity_at, @@ -6229,37 +6261,45 @@ type AutoArchiveInactiveChatsParams struct { } type AutoArchiveInactiveChatsRow struct { - ID uuid.UUID `db:"id" json:"id"` - OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` - WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` - Title string `db:"title" json:"title"` - Status ChatStatus `db:"status" json:"status"` - WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` - StartedAt sql.NullTime `db:"started_at" json:"started_at"` - HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` - CreatedAt time.Time `db:"created_at" json:"created_at"` - UpdatedAt time.Time `db:"updated_at" json:"updated_at"` - ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` - RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` - LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` - Archived bool `db:"archived" json:"archived"` - LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` - Mode NullChatMode `db:"mode" json:"mode"` - MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` - Labels json.RawMessage `db:"labels" json:"labels"` - BuildID uuid.NullUUID `db:"build_id" json:"build_id"` - AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` - PinOrder int32 `db:"pin_order" json:"pin_order"` - LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` - LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` - DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` - OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` - PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` - ClientType ChatClientType `db:"client_type" json:"client_type"` - LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` - UserACL json.RawMessage `db:"user_acl" json:"user_acl"` - GroupACL json.RawMessage `db:"group_acl" json:"group_acl"` - LastActivityAt time.Time `db:"last_activity_at" json:"last_activity_at"` + ID uuid.UUID `db:"id" json:"id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + Title string `db:"title" json:"title"` + Status ChatStatus `db:"status" json:"status"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` + StartedAt sql.NullTime `db:"started_at" json:"started_at"` + HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` + RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` + LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` + Archived bool `db:"archived" json:"archived"` + LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` + Mode NullChatMode `db:"mode" json:"mode"` + MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` + Labels json.RawMessage `db:"labels" json:"labels"` + BuildID uuid.NullUUID `db:"build_id" json:"build_id"` + AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` + PinOrder int32 `db:"pin_order" json:"pin_order"` + LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` + LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` + DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` + ClientType ChatClientType `db:"client_type" json:"client_type"` + LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` + UserACL json.RawMessage `db:"user_acl" json:"user_acl"` + GroupACL json.RawMessage `db:"group_acl" json:"group_acl"` + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + HistoryVersion int64 `db:"history_version" json:"history_version"` + QueueVersion int64 `db:"queue_version" json:"queue_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` + LastActivityAt time.Time `db:"last_activity_at" json:"last_activity_at"` } // Archives inactive root chats (pinned and already-archived chats skipped), @@ -6309,6 +6349,14 @@ func (q *sqlQuerier) AutoArchiveInactiveChats(ctx context.Context, arg AutoArchi &i.LastTurnSummary, &i.UserACL, &i.GroupACL, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.LastActivityAt, ); err != nil { return nil, err @@ -6347,6 +6395,47 @@ func (q *sqlQuerier) BackoffChatDiffStatus(ctx context.Context, arg BackoffChatD return err } +const batchDeleteChatHeartbeats = `-- name: BatchDeleteChatHeartbeats :execrows +DELETE FROM chat_heartbeats +USING unnest($1::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord) +JOIN unnest($2::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord) +WHERE chat_heartbeats.chat_id = chat_ids.chat_id + AND chat_heartbeats.runner_id = runner_ids.runner_id +` + +type BatchDeleteChatHeartbeatsParams struct { + ChatIds []uuid.UUID `db:"chat_ids" json:"chat_ids"` + RunnerIds []uuid.UUID `db:"runner_ids" json:"runner_ids"` +} + +// Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs. +func (q *sqlQuerier) BatchDeleteChatHeartbeats(ctx context.Context, arg BatchDeleteChatHeartbeatsParams) (int64, error) { + result, err := q.db.ExecContext(ctx, batchDeleteChatHeartbeats, pq.Array(arg.ChatIds), pq.Array(arg.RunnerIds)) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +const batchUpsertChatHeartbeats = `-- name: BatchUpsertChatHeartbeats :exec +INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at) +SELECT chat_ids.chat_id, runner_ids.runner_id, NOW() +FROM unnest($1::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord) +JOIN unnest($2::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord) +ON CONFLICT (chat_id, runner_id) DO UPDATE +SET heartbeat_at = EXCLUDED.heartbeat_at +` + +type BatchUpsertChatHeartbeatsParams struct { + ChatIds []uuid.UUID `db:"chat_ids" json:"chat_ids"` + RunnerIds []uuid.UUID `db:"runner_ids" json:"runner_ids"` +} + +func (q *sqlQuerier) BatchUpsertChatHeartbeats(ctx context.Context, arg BatchUpsertChatHeartbeatsParams) error { + _, err := q.db.ExecContext(ctx, batchUpsertChatHeartbeats, pq.Array(arg.ChatIds), pq.Array(arg.RunnerIds)) + return err +} + const clearChatMessageProviderResponseIDsByChatID = `-- name: ClearChatMessageProviderResponseIDsByChatID :exec UPDATE chat_messages SET provider_response_id = NULL @@ -6360,6 +6449,21 @@ func (q *sqlQuerier) ClearChatMessageProviderResponseIDsByChatID(ctx context.Con return err } +const countChatQueuedMessages = `-- name: CountChatQueuedMessages :one +SELECT COUNT(*)::bigint AS count +FROM chat_queued_messages +WHERE chat_id = $1::uuid +` + +// Cheap queue-length check used by ChatMachine.Update when deciding +// whether the chat is in a "1" sub-state. +func (q *sqlQuerier) CountChatQueuedMessages(ctx context.Context, chatID uuid.UUID) (int64, error) { + row := q.db.QueryRowContext(ctx, countChatQueuedMessages, chatID) + var count int64 + err := row.Scan(&count) + return count, err +} + const countEnabledModelsWithoutPricing = `-- name: CountEnabledModelsWithoutPricing :one SELECT COUNT(*)::bigint AS count FROM chat_model_configs @@ -6384,6 +6488,17 @@ func (q *sqlQuerier) CountEnabledModelsWithoutPricing(ctx context.Context) (int6 return count, err } +const deleteAllChatHeartbeats = `-- name: DeleteAllChatHeartbeats :exec +DELETE FROM chat_heartbeats WHERE chat_id = $1::uuid +` + +// Deletes all heartbeat rows for the chat. Used during ownership +// transitions that abandon a lease. +func (q *sqlQuerier) DeleteAllChatHeartbeats(ctx context.Context, chatID uuid.UUID) error { + _, err := q.db.ExecContext(ctx, deleteAllChatHeartbeats, chatID) + return err +} + const deleteAllChatQueuedMessages = `-- name: DeleteAllChatQueuedMessages :exec DELETE FROM chat_queued_messages WHERE chat_id = $1 ` @@ -6393,6 +6508,19 @@ func (q *sqlQuerier) DeleteAllChatQueuedMessages(ctx context.Context, chatID uui return err } +const deleteAllChatQueuedMessagesReturningCount = `-- name: DeleteAllChatQueuedMessagesReturningCount :execrows +DELETE FROM chat_queued_messages +WHERE chat_id = $1::uuid +` + +func (q *sqlQuerier) DeleteAllChatQueuedMessagesReturningCount(ctx context.Context, chatID uuid.UUID) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteAllChatQueuedMessagesReturningCount, chatID) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + const deleteChatQueuedMessage = `-- name: DeleteChatQueuedMessage :exec DELETE FROM chat_queued_messages WHERE id = $1 AND chat_id = $2 ` @@ -6407,6 +6535,27 @@ func (q *sqlQuerier) DeleteChatQueuedMessage(ctx context.Context, arg DeleteChat return err } +const deleteChatQueuedMessageReturningCount = `-- name: DeleteChatQueuedMessageReturningCount :execrows +DELETE FROM chat_queued_messages +WHERE id = $1::bigint AND chat_id = $2::uuid +` + +type DeleteChatQueuedMessageReturningCountParams struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` +} + +// Deletes a queued message, scoped to the parent chat. Returns the +// number of affected rows so callers can detect missing rows without +// a follow-up read. +func (q *sqlQuerier) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg DeleteChatQueuedMessageReturningCountParams) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteChatQueuedMessageReturningCount, arg.ID, arg.ChatID) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + const deleteChatUsageLimitGroupOverride = `-- name: DeleteChatUsageLimitGroupOverride :exec UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = $1::uuid ` @@ -6458,8 +6607,21 @@ func (q *sqlQuerier) DeleteOldChats(ctx context.Context, arg DeleteOldChatsParam return result.RowsAffected() } +const deleteStaleChatHeartbeats = `-- name: DeleteStaleChatHeartbeats :execrows +DELETE FROM chat_heartbeats +WHERE heartbeat_at < NOW() - (INTERVAL '1 second' * $1::int) +` + +func (q *sqlQuerier) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds int32) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteStaleChatHeartbeats, staleSeconds) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + const getActiveChatsByAgentID = `-- name: GetActiveChatsByAgentID :many -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded WHERE agent_id = $1::uuid AND archived = false @@ -6508,6 +6670,14 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -6526,6 +6696,152 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U return items, nil } +const getAutoArchiveInactiveChatCandidates = `-- name: GetAutoArchiveInactiveChatCandidates :many +SELECT + chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.snapshot_version, chats_expanded.history_version, chats_expanded.queue_version, chats_expanded.generation_attempt, chats_expanded.retry_state, chats_expanded.retry_state_version, chats_expanded.runner_id, chats_expanded.requires_action_deadline_at, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, + COALESCE(activity.last_activity_at, chats_expanded.created_at)::timestamptz AS last_activity_at +FROM chats_expanded +LEFT JOIN LATERAL ( + SELECT MAX(chat_messages.created_at) AS last_activity_at + FROM chat_messages + JOIN chats family_chat ON family_chat.id = chat_messages.chat_id + WHERE (family_chat.id = chats_expanded.id OR family_chat.root_chat_id = chats_expanded.id) + AND chat_messages.deleted = false +) activity ON TRUE +WHERE + chats_expanded.archived = false + AND chats_expanded.pin_order = 0 + AND chats_expanded.parent_chat_id IS NULL + AND chats_expanded.created_at < $1::timestamptz + AND chats_expanded.status NOT IN ( + 'running'::chat_status, + 'interrupting'::chat_status, + 'pending'::chat_status, + 'paused'::chat_status, + 'requires_action'::chat_status + ) + AND COALESCE(activity.last_activity_at, chats_expanded.created_at) < $1::timestamptz +ORDER BY chats_expanded.created_at ASC +LIMIT $2::int +` + +type GetAutoArchiveInactiveChatCandidatesParams struct { + ArchiveCutoff time.Time `db:"archive_cutoff" json:"archive_cutoff"` + LimitCount int32 `db:"limit_count" json:"limit_count"` +} + +type GetAutoArchiveInactiveChatCandidatesRow struct { + ID uuid.UUID `db:"id" json:"id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + Title string `db:"title" json:"title"` + Status ChatStatus `db:"status" json:"status"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` + StartedAt sql.NullTime `db:"started_at" json:"started_at"` + HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` + RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` + LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` + Archived bool `db:"archived" json:"archived"` + LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` + Mode NullChatMode `db:"mode" json:"mode"` + MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` + Labels StringMap `db:"labels" json:"labels"` + BuildID uuid.NullUUID `db:"build_id" json:"build_id"` + AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` + PinOrder int32 `db:"pin_order" json:"pin_order"` + LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` + LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` + DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` + ClientType ChatClientType `db:"client_type" json:"client_type"` + LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + HistoryVersion int64 `db:"history_version" json:"history_version"` + QueueVersion int64 `db:"queue_version" json:"queue_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` + UserACL ChatACL `db:"user_acl" json:"user_acl"` + GroupACL ChatACL `db:"group_acl" json:"group_acl"` + OwnerUsername string `db:"owner_username" json:"owner_username"` + OwnerName string `db:"owner_name" json:"owner_name"` + LastActivityAt time.Time `db:"last_activity_at" json:"last_activity_at"` +} + +// Returns read-only root chat candidates for state-machine-backed +// auto-archive. Activity is computed across the root family. The query +// limits roots, not total family members. +func (q *sqlQuerier) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg GetAutoArchiveInactiveChatCandidatesParams) ([]GetAutoArchiveInactiveChatCandidatesRow, error) { + rows, err := q.db.QueryContext(ctx, getAutoArchiveInactiveChatCandidates, arg.ArchiveCutoff, arg.LimitCount) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetAutoArchiveInactiveChatCandidatesRow + for rows.Next() { + var i GetAutoArchiveInactiveChatCandidatesRow + if err := rows.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + &i.LastActivityAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getChatACLByID = `-- name: GetChatACLByID :one SELECT user_acl AS users, @@ -6549,7 +6865,7 @@ func (q *sqlQuerier) GetChatACLByID(ctx context.Context, id uuid.UUID) (GetChatA } const getChatByID = `-- name: GetChatByID :one -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded WHERE id = $1::uuid ` @@ -6586,6 +6902,120 @@ func (q *sqlQuerier) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + ) + return i, err +} + +const getChatByIDForShare = `-- name: GetChatByIDForShare :one +WITH shared_chat AS ( + SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at + FROM chats + WHERE id = $1::uuid + FOR SHARE +), +chats_expanded AS ( + SELECT + shared_chat.id, + shared_chat.owner_id, + shared_chat.workspace_id, + shared_chat.title, + shared_chat.status, + shared_chat.worker_id, + shared_chat.started_at, + shared_chat.heartbeat_at, + shared_chat.created_at, + shared_chat.updated_at, + shared_chat.parent_chat_id, + shared_chat.root_chat_id, + shared_chat.last_model_config_id, + shared_chat.archived, + shared_chat.last_error, + shared_chat.mode, + shared_chat.mcp_server_ids, + shared_chat.labels, + shared_chat.build_id, + shared_chat.agent_id, + shared_chat.pin_order, + shared_chat.last_read_message_id, + shared_chat.last_injected_context, + shared_chat.dynamic_tools, + shared_chat.organization_id, + shared_chat.plan_mode, + shared_chat.client_type, + shared_chat.last_turn_summary, + shared_chat.snapshot_version, + shared_chat.history_version, + shared_chat.queue_version, + shared_chat.generation_attempt, + shared_chat.retry_state, + shared_chat.retry_state_version, + shared_chat.runner_id, + shared_chat.requires_action_deadline_at, + COALESCE(root.user_acl, shared_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, shared_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM + shared_chat + LEFT JOIN chats root ON root.id = COALESCE(shared_chat.root_chat_id, shared_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = shared_chat.owner_id +) +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name +FROM chats_expanded +` + +func (q *sqlQuerier) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (Chat, error) { + row := q.db.QueryRowContext(ctx, getChatByIDForShare, id) + var i Chat + err := row.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -6596,7 +7026,7 @@ func (q *sqlQuerier) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error const getChatByIDForUpdate = `-- name: GetChatByIDForUpdate :one WITH locked_chat AS ( - SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl + SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at FROM chats WHERE id = $1::uuid FOR UPDATE @@ -6631,6 +7061,14 @@ chats_expanded AS ( locked_chat.plan_mode, locked_chat.client_type, locked_chat.last_turn_summary, + locked_chat.snapshot_version, + locked_chat.history_version, + locked_chat.queue_version, + locked_chat.generation_attempt, + locked_chat.retry_state, + locked_chat.retry_state_version, + locked_chat.runner_id, + locked_chat.requires_action_deadline_at, COALESCE(root.user_acl, locked_chat.user_acl) AS user_acl, COALESCE(root.group_acl, locked_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -6640,7 +7078,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(locked_chat.root_chat_id, locked_chat.parent_chat_id) JOIN visible_users owner ON owner.id = locked_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -6676,6 +7114,14 @@ func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Ch &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -7212,9 +7658,59 @@ func (q *sqlQuerier) GetChatDiffStatusesByChatIDs(ctx context.Context, chatIds [ return items, nil } +const getChatFamilyIDsByRootID = `-- name: GetChatFamilyIDsByRootID :many +SELECT id +FROM chats +WHERE id = $1::uuid OR root_chat_id = $1::uuid +ORDER BY (id = $1::uuid) DESC, created_at ASC, id ASC +` + +// Returns the chat IDs of every chat in a family (root + all children) +// in deterministic order. The id parameter must be the root id; the +// query does not walk up from a child. +func (q *sqlQuerier) GetChatFamilyIDsByRootID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + rows, err := q.db.QueryContext(ctx, getChatFamilyIDsByRootID, id) + if err != nil { + return nil, err + } + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getChatHeartbeat = `-- name: GetChatHeartbeat :one +SELECT chat_id, runner_id, heartbeat_at FROM chat_heartbeats +WHERE chat_id = $1::uuid AND runner_id = $2::uuid +` + +type GetChatHeartbeatParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + RunnerID uuid.UUID `db:"runner_id" json:"runner_id"` +} + +func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatParams) (ChatHeartbeat, error) { + row := q.db.QueryRowContext(ctx, getChatHeartbeat, arg.ChatID, arg.RunnerID) + var i ChatHeartbeat + err := row.Scan(&i.ChatID, &i.RunnerID, &i.HeartbeatAt) + return i, err +} + const getChatMessageByID = `-- name: GetChatMessageByID :one SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -7248,6 +7744,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ) return i, err } @@ -7337,7 +7834,7 @@ func (q *sqlQuerier) GetChatMessageSummariesPerChat(ctx context.Context, created const getChatMessagesByChatID = `-- name: GetChatMessagesByChatID :many SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -7386,6 +7883,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ); err != nil { return nil, err } @@ -7402,7 +7900,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes const getChatMessagesByChatIDAscPaginated = `-- name: GetChatMessagesByChatIDAscPaginated :many SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -7454,6 +7952,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ); err != nil { return nil, err } @@ -7470,7 +7969,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar const getChatMessagesByChatIDDescPaginated = `-- name: GetChatMessagesByChatIDDescPaginated :many SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -7535,6 +8034,72 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getChatMessagesByRevisionForStream = `-- name: GetChatMessagesByRevisionForStream :many +SELECT + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision +FROM + chat_messages +WHERE + chat_id = $1::uuid + AND revision > $2::bigint + AND visibility IN ('user', 'both') +ORDER BY + created_at ASC, id ASC +` + +type GetChatMessagesByRevisionForStreamParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + AfterRevision int64 `db:"after_revision" json:"after_revision"` +} + +func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg GetChatMessagesByRevisionForStreamParams) ([]ChatMessage, error) { + rows, err := q.db.QueryContext(ctx, getChatMessagesByRevisionForStream, arg.ChatID, arg.AfterRevision) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ChatMessage + for rows.Next() { + var i ChatMessage + if err := rows.Scan( + &i.ID, + &i.ChatID, + &i.ModelConfigID, + &i.CreatedAt, + &i.Role, + &i.Content, + &i.Visibility, + &i.InputTokens, + &i.OutputTokens, + &i.TotalTokens, + &i.ReasoningTokens, + &i.CacheCreationTokens, + &i.CacheReadTokens, + &i.ContextLimit, + &i.Compressed, + &i.CreatedBy, + &i.ContentVersion, + &i.TotalCostMicros, + &i.RuntimeMs, + &i.Deleted, + &i.ProviderResponseID, + &i.APIKeyID, + &i.Revision, ); err != nil { return nil, err } @@ -7567,7 +8132,7 @@ WITH latest_compressed_summary AS ( 1 ) SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -7640,6 +8205,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ); err != nil { return nil, err } @@ -7700,8 +8266,58 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get return items, nil } +const getChatQueuedMessageByID = `-- name: GetChatQueuedMessageByID :one +SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by FROM chat_queued_messages +WHERE id = $1::bigint AND chat_id = $2::uuid +` + +type GetChatQueuedMessageByIDParams struct { + ID int64 `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` +} + +func (q *sqlQuerier) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQueuedMessageByIDParams) (ChatQueuedMessage, error) { + row := q.db.QueryRowContext(ctx, getChatQueuedMessageByID, arg.ID, arg.ChatID) + var i ChatQueuedMessage + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Content, + &i.CreatedAt, + &i.ModelConfigID, + &i.APIKeyID, + &i.Position, + &i.CreatedBy, + ) + return i, err +} + +const getChatQueuedMessageHead = `-- name: GetChatQueuedMessageHead :one +SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by FROM chat_queued_messages +WHERE chat_id = $1::uuid +ORDER BY position ASC, id ASC +LIMIT 1 +` + +// Returns the queue head (lowest position, then lowest id). +func (q *sqlQuerier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) { + row := q.db.QueryRowContext(ctx, getChatQueuedMessageHead, chatID) + var i ChatQueuedMessage + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Content, + &i.CreatedAt, + &i.ModelConfigID, + &i.APIKeyID, + &i.Position, + &i.CreatedBy, + ) + return i, err +} + const getChatQueuedMessages = `-- name: GetChatQueuedMessages :many -SELECT id, chat_id, content, created_at, model_config_id, api_key_id FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by FROM chat_queued_messages WHERE chat_id = $1 ORDER BY created_at ASC, id ASC ` @@ -7722,6 +8338,105 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID &i.CreatedAt, &i.ModelConfigID, &i.APIKeyID, + &i.Position, + &i.CreatedBy, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getChatQueuedMessagesByPosition = `-- name: GetChatQueuedMessagesByPosition :many +SELECT id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by FROM chat_queued_messages +WHERE chat_id = $1::uuid +ORDER BY position ASC, id ASC +` + +// Returns queued messages in state-machine order (position ASC, id ASC). +func (q *sqlQuerier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error) { + rows, err := q.db.QueryContext(ctx, getChatQueuedMessagesByPosition, chatID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ChatQueuedMessage + for rows.Next() { + var i ChatQueuedMessage + if err := rows.Scan( + &i.ID, + &i.ChatID, + &i.Content, + &i.CreatedAt, + &i.ModelConfigID, + &i.APIKeyID, + &i.Position, + &i.CreatedBy, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getChatStreamSyncRows = `-- name: GetChatStreamSyncRows :many +SELECT + id, + snapshot_version, + history_version, + queue_version, + retry_state_version, + generation_attempt, + status, + worker_id +FROM chats +WHERE id = ANY($1::uuid[]) +ORDER BY id ASC +` + +type GetChatStreamSyncRowsRow struct { + ID uuid.UUID `db:"id" json:"id"` + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + HistoryVersion int64 `db:"history_version" json:"history_version"` + QueueVersion int64 `db:"queue_version" json:"queue_version"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + Status ChatStatus `db:"status" json:"status"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` +} + +func (q *sqlQuerier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID) ([]GetChatStreamSyncRowsRow, error) { + rows, err := q.db.QueryContext(ctx, getChatStreamSyncRows, pq.Array(ids)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetChatStreamSyncRowsRow + for rows.Next() { + var i GetChatStreamSyncRowsRow + if err := rows.Scan( + &i.ID, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.RetryStateVersion, + &i.GenerationAttempt, + &i.Status, + &i.WorkerID, ); err != nil { return nil, err } @@ -7857,6 +8572,166 @@ func (q *sqlQuerier) GetChatUserPromptsByChatID(ctx context.Context, arg GetChat return items, nil } +const getChatWorkerAcquisitionCandidates = `-- name: GetChatWorkerAcquisitionCandidates :many +SELECT + chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.snapshot_version, chats_expanded.history_version, chats_expanded.queue_version, chats_expanded.generation_attempt, chats_expanded.retry_state, chats_expanded.retry_state_version, chats_expanded.runner_id, chats_expanded.requires_action_deadline_at, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, + chat_heartbeats.heartbeat_at AS current_heartbeat_at, + NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats_expanded.id + AND current_lease.runner_id = chats_expanded.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) + ) AS heartbeat_stale +FROM chats_expanded +LEFT JOIN chat_heartbeats + ON chat_heartbeats.chat_id = chats_expanded.id + AND chat_heartbeats.runner_id = chats_expanded.runner_id +WHERE + chats_expanded.status IN ('running'::chat_status, 'interrupting'::chat_status, 'requires_action'::chat_status) + AND chats_expanded.archived = false + AND ( + chats_expanded.worker_id IS NULL + OR chats_expanded.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats_expanded.id + AND current_lease.runner_id = chats_expanded.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * $1::int) + ) + ) +ORDER BY chats_expanded.updated_at ASC, chats_expanded.id ASC +LIMIT $2::int +` + +type GetChatWorkerAcquisitionCandidatesParams struct { + StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` + LimitCount int32 `db:"limit_count" json:"limit_count"` +} + +type GetChatWorkerAcquisitionCandidatesRow struct { + ID uuid.UUID `db:"id" json:"id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + Title string `db:"title" json:"title"` + Status ChatStatus `db:"status" json:"status"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` + StartedAt sql.NullTime `db:"started_at" json:"started_at"` + HeartbeatAt sql.NullTime `db:"heartbeat_at" json:"heartbeat_at"` + CreatedAt time.Time `db:"created_at" json:"created_at"` + UpdatedAt time.Time `db:"updated_at" json:"updated_at"` + ParentChatID uuid.NullUUID `db:"parent_chat_id" json:"parent_chat_id"` + RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"` + LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"` + Archived bool `db:"archived" json:"archived"` + LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` + Mode NullChatMode `db:"mode" json:"mode"` + MCPServerIDs []uuid.UUID `db:"mcp_server_ids" json:"mcp_server_ids"` + Labels StringMap `db:"labels" json:"labels"` + BuildID uuid.NullUUID `db:"build_id" json:"build_id"` + AgentID uuid.NullUUID `db:"agent_id" json:"agent_id"` + PinOrder int32 `db:"pin_order" json:"pin_order"` + LastReadMessageID sql.NullInt64 `db:"last_read_message_id" json:"last_read_message_id"` + LastInjectedContext pqtype.NullRawMessage `db:"last_injected_context" json:"last_injected_context"` + DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + PlanMode NullChatPlanMode `db:"plan_mode" json:"plan_mode"` + ClientType ChatClientType `db:"client_type" json:"client_type"` + LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` + SnapshotVersion int64 `db:"snapshot_version" json:"snapshot_version"` + HistoryVersion int64 `db:"history_version" json:"history_version"` + QueueVersion int64 `db:"queue_version" json:"queue_version"` + GenerationAttempt int64 `db:"generation_attempt" json:"generation_attempt"` + RetryState pqtype.NullRawMessage `db:"retry_state" json:"retry_state"` + RetryStateVersion int64 `db:"retry_state_version" json:"retry_state_version"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` + UserACL ChatACL `db:"user_acl" json:"user_acl"` + GroupACL ChatACL `db:"group_acl" json:"group_acl"` + OwnerUsername string `db:"owner_username" json:"owner_username"` + OwnerName string `db:"owner_name" json:"owner_name"` + CurrentHeartbeatAt sql.NullTime `db:"current_heartbeat_at" json:"current_heartbeat_at"` + HeartbeatStale bool `db:"heartbeat_stale" json:"heartbeat_stale"` +} + +// Returns chats that workers may try to acquire. Candidates must be: +// - in a worker-runnable execution status; +// - unarchived; and +// - missing ownership, carrying inconsistent ownership, or lacking a +// fresh heartbeat for the assigned runner. +// +// Missing ownership is worker_id IS NULL. Inconsistent ownership is +// runner_id IS NULL while worker_id is set. Stale ownership is no +// heartbeat row for (chat_id, runner_id), or one older than +// @stale_seconds by database time. Candidates are ordered by oldest +// updated_at first so workers drain stale runnable chats predictably. +func (q *sqlQuerier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg GetChatWorkerAcquisitionCandidatesParams) ([]GetChatWorkerAcquisitionCandidatesRow, error) { + rows, err := q.db.QueryContext(ctx, getChatWorkerAcquisitionCandidates, arg.StaleSeconds, arg.LimitCount) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetChatWorkerAcquisitionCandidatesRow + for rows.Next() { + var i GetChatWorkerAcquisitionCandidatesRow + if err := rows.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + &i.CurrentHeartbeatAt, + &i.HeartbeatStale, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getChats = `-- name: GetChats :many WITH cursor_chat AS ( SELECT @@ -7867,7 +8742,7 @@ WITH cursor_chat AS ( WHERE id = $7 ) SELECT - chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, + chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.snapshot_version, chats_expanded.history_version, chats_expanded.queue_version, chats_expanded.generation_attempt, chats_expanded.retry_state, chats_expanded.retry_state_version, chats_expanded.runner_id, chats_expanded.requires_action_deadline_at, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, EXISTS ( SELECT 1 FROM chat_messages cm WHERE cm.chat_id = chats_expanded.id @@ -8107,6 +8982,14 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha &i.Chat.PlanMode, &i.Chat.ClientType, &i.Chat.LastTurnSummary, + &i.Chat.SnapshotVersion, + &i.Chat.HistoryVersion, + &i.Chat.QueueVersion, + &i.Chat.GenerationAttempt, + &i.Chat.RetryState, + &i.Chat.RetryStateVersion, + &i.Chat.RunnerID, + &i.Chat.RequiresActionDeadlineAt, &i.Chat.UserACL, &i.Chat.GroupACL, &i.Chat.OwnerUsername, @@ -8128,7 +9011,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha const getChatsByChatFileID = `-- name: GetChatsByChatFileID :many SELECT - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded WHERE @@ -8179,6 +9062,85 @@ func (q *sqlQuerier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getChatsByIDsForRunnerSync = `-- name: GetChatsByIDsForRunnerSync :many +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name +FROM chats_expanded +WHERE id = ANY($1::uuid[]) +ORDER BY id ASC +` + +func (q *sqlQuerier) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid.UUID) ([]Chat, error) { + rows, err := q.db.QueryContext(ctx, getChatsByIDsForRunnerSync, pq.Array(ids)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Chat + for rows.Next() { + var i Chat + if err := rows.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -8198,7 +9160,7 @@ func (q *sqlQuerier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) } const getChatsByWorkspaceIDs = `-- name: GetChatsByWorkspaceIDs :many -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded WHERE archived = false AND workspace_id = ANY($1::uuid[]) @@ -8243,6 +9205,14 @@ func (q *sqlQuerier) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -8331,7 +9301,7 @@ func (q *sqlQuerier) GetChatsUpdatedAfter(ctx context.Context, updatedAfter time const getChildChatsByParentIDs = `-- name: GetChildChatsByParentIDs :many SELECT - chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, + chats_expanded.id, chats_expanded.owner_id, chats_expanded.workspace_id, chats_expanded.title, chats_expanded.status, chats_expanded.worker_id, chats_expanded.started_at, chats_expanded.heartbeat_at, chats_expanded.created_at, chats_expanded.updated_at, chats_expanded.parent_chat_id, chats_expanded.root_chat_id, chats_expanded.last_model_config_id, chats_expanded.archived, chats_expanded.last_error, chats_expanded.mode, chats_expanded.mcp_server_ids, chats_expanded.labels, chats_expanded.build_id, chats_expanded.agent_id, chats_expanded.pin_order, chats_expanded.last_read_message_id, chats_expanded.last_injected_context, chats_expanded.dynamic_tools, chats_expanded.organization_id, chats_expanded.plan_mode, chats_expanded.client_type, chats_expanded.last_turn_summary, chats_expanded.snapshot_version, chats_expanded.history_version, chats_expanded.queue_version, chats_expanded.generation_attempt, chats_expanded.retry_state, chats_expanded.retry_state_version, chats_expanded.runner_id, chats_expanded.requires_action_deadline_at, chats_expanded.user_acl, chats_expanded.group_acl, chats_expanded.owner_username, chats_expanded.owner_name, EXISTS ( SELECT 1 FROM chat_messages cm WHERE cm.chat_id = chats_expanded.id @@ -8404,6 +9374,14 @@ func (q *sqlQuerier) GetChildChatsByParentIDs(ctx context.Context, arg GetChildC &i.Chat.PlanMode, &i.Chat.ClientType, &i.Chat.LastTurnSummary, + &i.Chat.SnapshotVersion, + &i.Chat.HistoryVersion, + &i.Chat.QueueVersion, + &i.Chat.GenerationAttempt, + &i.Chat.RetryState, + &i.Chat.RetryStateVersion, + &i.Chat.RunnerID, + &i.Chat.RequiresActionDeadlineAt, &i.Chat.UserACL, &i.Chat.GroupACL, &i.Chat.OwnerUsername, @@ -8423,9 +9401,23 @@ func (q *sqlQuerier) GetChildChatsByParentIDs(ctx context.Context, arg GetChildC return items, nil } +const getDatabaseNow = `-- name: GetDatabaseNow :one +SELECT NOW()::timestamptz AS now +` + +// Returns the current database timestamp. Used so transitions that +// record deadlines or heartbeats rely on a clock that is consistent +// with the database rather than the caller's local clock. +func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) { + row := q.db.QueryRowContext(ctx, getDatabaseNow) + var now time.Time + err := row.Scan(&now) + return now, err +} + const getLastChatMessageByRole = `-- name: GetLastChatMessageByRole :one SELECT - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision FROM chat_messages WHERE @@ -8469,13 +9461,14 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ) return i, err } const getStaleChats = `-- name: GetStaleChats :many SELECT - id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name + id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded WHERE @@ -8536,6 +9529,14 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -8616,6 +9617,21 @@ func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGrou return limit_micros, err } +const incrementChatGenerationAttempt = `-- name: IncrementChatGenerationAttempt :one +UPDATE chats +SET generation_attempt = generation_attempt + 1, updated_at = NOW() +WHERE id = $1::uuid +RETURNING generation_attempt +` + +// Increments generation_attempt and returns the resulting value. +func (q *sqlQuerier) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) { + row := q.db.QueryRowContext(ctx, incrementChatGenerationAttempt, id) + var generation_attempt int64 + err := row.Scan(&generation_attempt) + return generation_attempt, err +} + const insertChat = `-- name: InsertChat :one WITH inserted_chat AS ( INSERT INTO chats ( @@ -8653,7 +9669,7 @@ INSERT INTO chats ( $15::jsonb, $16::chat_client_type ) -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -8685,6 +9701,14 @@ chats_expanded AS ( inserted_chat.plan_mode, inserted_chat.client_type, inserted_chat.last_turn_summary, + inserted_chat.snapshot_version, + inserted_chat.history_version, + inserted_chat.queue_version, + inserted_chat.generation_attempt, + inserted_chat.retry_state, + inserted_chat.retry_state_version, + inserted_chat.runner_id, + inserted_chat.requires_action_deadline_at, COALESCE(root.user_acl, inserted_chat.user_acl) AS user_acl, COALESCE(root.group_acl, inserted_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -8694,7 +9718,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(inserted_chat.root_chat_id, inserted_chat.parent_chat_id) JOIN visible_users owner ON owner.id = inserted_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -8766,6 +9790,14 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -8845,7 +9877,7 @@ SELECT NULLIF(UNNEST($18::bigint[]), 0), NULLIF(UNNEST($19::text[]), '') RETURNING - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision ` type InsertChatMessagesParams struct { @@ -8922,6 +9954,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ); err != nil { return nil, err } @@ -8937,14 +9970,16 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa } const insertChatQueuedMessage = `-- name: InsertChatQueuedMessage :one -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id) -VALUES ( - $1, - $2, +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by) +SELECT + $1::uuid, + $2::jsonb, $3::uuid, - $4::text -) -RETURNING id, chat_id, content, created_at, model_config_id, api_key_id + $4::text, + chats.owner_id +FROM chats +WHERE chats.id = $1::uuid +RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by ` type InsertChatQueuedMessageParams struct { @@ -8954,6 +9989,9 @@ type InsertChatQueuedMessageParams struct { APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` } +// Legacy queue insertion path. When no caller-supplied creator exists, +// preserve the created_by invariant by attributing the queued row to the +// chat owner. func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) { row := q.db.QueryRowContext(ctx, insertChatQueuedMessage, arg.ChatID, @@ -8969,10 +10007,83 @@ func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChat &i.CreatedAt, &i.ModelConfigID, &i.APIKeyID, + &i.Position, + &i.CreatedBy, ) return i, err } +const insertChatQueuedMessageWithCreator = `-- name: InsertChatQueuedMessageWithCreator :one +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by) +VALUES ( + $1::uuid, + $2::jsonb, + $3::uuid, + $4::text, + $5::uuid +) +RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by +` + +type InsertChatQueuedMessageWithCreatorParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Content json.RawMessage `db:"content" json:"content"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` +} + +// Inserts a queued message that carries a position (from the default +// sequence) and an explicit created_by reference. Use this when the +// queued-message creator differs from the chat owner. +func (q *sqlQuerier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error) { + row := q.db.QueryRowContext(ctx, insertChatQueuedMessageWithCreator, + arg.ChatID, + arg.Content, + arg.ModelConfigID, + arg.APIKeyID, + arg.CreatedBy, + ) + var i ChatQueuedMessage + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Content, + &i.CreatedAt, + &i.ModelConfigID, + &i.APIKeyID, + &i.Position, + &i.CreatedBy, + ) + return i, err +} + +const isChatHeartbeatStale = `-- name: IsChatHeartbeatStale :one +SELECT NOT EXISTS ( + SELECT 1 FROM chat_heartbeats + WHERE chat_id = $1::uuid + AND runner_id = $2::uuid + AND heartbeat_at > NOW() - (INTERVAL '1 second' * $3::int) +) AS stale +` + +type IsChatHeartbeatStaleParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + RunnerID uuid.UUID `db:"runner_id" json:"runner_id"` + StaleSeconds int32 `db:"stale_seconds" json:"stale_seconds"` +} + +// Returns true when there is no heartbeat row for (chat_id, runner_id) +// or the existing row is older than @stale_seconds seconds by database +// time. chatstate calls this in a single query so the staleness check +// is atomic and does not depend on the caller's local clock. +func (q *sqlQuerier) IsChatHeartbeatStale(ctx context.Context, arg IsChatHeartbeatStaleParams) (bool, error) { + row := q.db.QueryRowContext(ctx, isChatHeartbeatStale, arg.ChatID, arg.RunnerID, arg.StaleSeconds) + var stale bool + err := row.Scan(&stale) + return stale, err +} + const linkChatFiles = `-- name: LinkChatFiles :one WITH current AS ( SELECT COUNT(*) AS cnt @@ -9124,6 +10235,119 @@ func (q *sqlQuerier) ListChatUsageLimitOverrides(ctx context.Context) ([]ListCha return items, nil } +const lockChatAndBumpSnapshotVersion = `-- name: LockChatAndBumpSnapshotVersion :one +WITH bumped_chat AS ( + UPDATE chats + SET snapshot_version = snapshot_version + 1 + WHERE id = ( + SELECT id FROM chats + WHERE id = $1::uuid + FOR UPDATE + ) + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at +), +chats_expanded AS ( + SELECT + bumped_chat.id, + bumped_chat.owner_id, + bumped_chat.workspace_id, + bumped_chat.title, + bumped_chat.status, + bumped_chat.worker_id, + bumped_chat.started_at, + bumped_chat.heartbeat_at, + bumped_chat.created_at, + bumped_chat.updated_at, + bumped_chat.parent_chat_id, + bumped_chat.root_chat_id, + bumped_chat.last_model_config_id, + bumped_chat.archived, + bumped_chat.last_error, + bumped_chat.mode, + bumped_chat.mcp_server_ids, + bumped_chat.labels, + bumped_chat.build_id, + bumped_chat.agent_id, + bumped_chat.pin_order, + bumped_chat.last_read_message_id, + bumped_chat.last_injected_context, + bumped_chat.dynamic_tools, + bumped_chat.organization_id, + bumped_chat.plan_mode, + bumped_chat.client_type, + bumped_chat.last_turn_summary, + bumped_chat.snapshot_version, + bumped_chat.history_version, + bumped_chat.queue_version, + bumped_chat.generation_attempt, + bumped_chat.retry_state, + bumped_chat.retry_state_version, + bumped_chat.runner_id, + bumped_chat.requires_action_deadline_at, + COALESCE(root.user_acl, bumped_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, bumped_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM bumped_chat + LEFT JOIN chats root ON root.id = COALESCE(bumped_chat.root_chat_id, bumped_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = bumped_chat.owner_id +) +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name +FROM chats_expanded +` + +// Locks the chat row with FOR UPDATE and atomically increments its +// snapshot_version, returning the post-bump chat. This is the single +// entry point ChatMachine.Update uses to acquire the row lock and +// allocate a new snapshot version in one round trip. +func (q *sqlQuerier) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (Chat, error) { + row := q.db.QueryRowContext(ctx, lockChatAndBumpSnapshotVersion, id) + var i Chat + err := row.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + ) + return i, err +} + const pinChatByID = `-- name: PinChatByID :exec WITH target_chat AS ( SELECT @@ -9193,7 +10417,7 @@ WHERE id = ( ORDER BY cqm.created_at ASC, cqm.id ASC LIMIT 1 ) -RETURNING id, chat_id, content, created_at, model_config_id, api_key_id +RETURNING id, chat_id, content, created_at, model_config_id, api_key_id, position, created_by ` func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) { @@ -9206,6 +10430,8 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) &i.CreatedAt, &i.ModelConfigID, &i.APIKeyID, + &i.Position, + &i.CreatedBy, ) return i, err } @@ -9235,6 +10461,35 @@ func (q *sqlQuerier) ReorderChatQueuedMessageToFront(ctx context.Context, arg Re return result.RowsAffected() } +const reorderChatQueuedMessageToHead = `-- name: ReorderChatQueuedMessageToHead :execrows +UPDATE chat_queued_messages AS target +SET position = COALESCE( + (SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = $1::uuid), + 0 +) - 1 +WHERE target.id = $2::bigint + AND target.chat_id = $1::uuid + AND target.position > COALESCE( + (SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = $1::uuid), + target.position + ) +` + +type ReorderChatQueuedMessageToHeadParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + ID int64 `db:"id" json:"id"` +} + +// Sets the target queued message's position to one less than the +// current minimum position for that chat, moving it to the head. +func (q *sqlQuerier) ReorderChatQueuedMessageToHead(ctx context.Context, arg ReorderChatQueuedMessageToHeadParams) (int64, error) { + result, err := q.db.ExecContext(ctx, reorderChatQueuedMessageToHead, arg.ChatID, arg.ID) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one SELECT CASE WHEN NOT cfg.enabled THEN -1 @@ -9344,7 +10599,7 @@ WITH updated_chats AS ( archived = false, updated_at = NOW() WHERE id = $1::uuid OR root_chat_id = $1::uuid - RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -9376,6 +10631,14 @@ chats_expanded AS ( updated_chats.plan_mode, updated_chats.client_type, updated_chats.last_turn_summary, + updated_chats.snapshot_version, + updated_chats.history_version, + updated_chats.queue_version, + updated_chats.generation_attempt, + updated_chats.retry_state, + updated_chats.retry_state_version, + updated_chats.runner_id, + updated_chats.requires_action_deadline_at, COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -9385,7 +10648,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chats.root_chat_id, updated_chats.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chats.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ORDER BY (chats_expanded.id = $1::uuid) DESC, chats_expanded.created_at ASC, chats_expanded.id ASC ` @@ -9432,6 +10695,14 @@ func (q *sqlQuerier) UnarchiveChatByID(ctx context.Context, id uuid.UUID) ([]Cha &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -9538,7 +10809,7 @@ UPDATE chats SET updated_at = NOW() WHERE id = $3::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -9570,6 +10841,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -9579,7 +10858,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -9621,6 +10900,14 @@ func (q *sqlQuerier) UpdateChatBuildAgentBinding(ctx context.Context, arg Update &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -9638,7 +10925,7 @@ SET updated_at = NOW() WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -9670,6 +10957,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -9679,7 +10974,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -9720,6 +11015,149 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + ) + return i, err +} + +const updateChatExecutionState = `-- name: UpdateChatExecutionState :one +WITH updated_chat AS ( + UPDATE chats + SET + status = $1::chat_status, + archived = $2::boolean, + worker_id = $3::uuid, + runner_id = $4::uuid, + last_error = $5::jsonb, + requires_action_deadline_at = $6::timestamptz, + pin_order = CASE WHEN $2::boolean THEN 0 ELSE pin_order END, + updated_at = NOW() + WHERE id = $7::uuid + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at +), +chats_expanded AS ( + SELECT + updated_chat.id, + updated_chat.owner_id, + updated_chat.workspace_id, + updated_chat.title, + updated_chat.status, + updated_chat.worker_id, + updated_chat.started_at, + updated_chat.heartbeat_at, + updated_chat.created_at, + updated_chat.updated_at, + updated_chat.parent_chat_id, + updated_chat.root_chat_id, + updated_chat.last_model_config_id, + updated_chat.archived, + updated_chat.last_error, + updated_chat.mode, + updated_chat.mcp_server_ids, + updated_chat.labels, + updated_chat.build_id, + updated_chat.agent_id, + updated_chat.pin_order, + updated_chat.last_read_message_id, + updated_chat.last_injected_context, + updated_chat.dynamic_tools, + updated_chat.organization_id, + updated_chat.plan_mode, + updated_chat.client_type, + updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, + COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM updated_chat + LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = updated_chat.owner_id +) +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name +FROM chats_expanded +` + +type UpdateChatExecutionStateParams struct { + Status ChatStatus `db:"status" json:"status"` + Archived bool `db:"archived" json:"archived"` + WorkerID uuid.NullUUID `db:"worker_id" json:"worker_id"` + RunnerID uuid.NullUUID `db:"runner_id" json:"runner_id"` + LastError pqtype.NullRawMessage `db:"last_error" json:"last_error"` + RequiresActionDeadlineAt sql.NullTime `db:"requires_action_deadline_at" json:"requires_action_deadline_at"` + ID uuid.UUID `db:"id" json:"id"` +} + +// Atomically updates the execution-state-managed fields on a chat: +// status, archived, last_error, ownership identifiers, and the +// requires-action deadline. Callers compose this with transition +// mutations inside a single ChatMachine.Update transaction. +func (q *sqlQuerier) UpdateChatExecutionState(ctx context.Context, arg UpdateChatExecutionStateParams) (Chat, error) { + row := q.db.QueryRowContext(ctx, updateChatExecutionState, + arg.Status, + arg.Archived, + arg.WorkerID, + arg.RunnerID, + arg.LastError, + arg.RequiresActionDeadlineAt, + arg.ID, + ) + var i Chat + err := row.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -9782,7 +11220,7 @@ SET updated_at = NOW() WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -9814,6 +11252,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -9823,7 +11269,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -9864,6 +11310,14 @@ func (q *sqlQuerier) UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLab &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -9878,7 +11332,7 @@ UPDATE chats SET last_injected_context = $1::jsonb WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -9910,6 +11364,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -9919,7 +11381,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -9964,6 +11426,14 @@ func (q *sqlQuerier) UpdateChatLastInjectedContext(ctx context.Context, arg Upda &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -9981,7 +11451,7 @@ SET last_model_config_id = $1::uuid WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10013,6 +11483,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10022,7 +11500,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10063,6 +11541,14 @@ func (q *sqlQuerier) UpdateChatLastModelConfigByID(ctx context.Context, arg Upda &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10097,26 +11583,24 @@ SET ), '') WHERE id = $2::uuid - AND updated_at = $3::timestamptz + AND history_version = $3::bigint ` type UpdateChatLastTurnSummaryParams struct { - LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` - ID uuid.UUID `db:"id" json:"id"` - ExpectedUpdatedAt time.Time `db:"expected_updated_at" json:"expected_updated_at"` + LastTurnSummary sql.NullString `db:"last_turn_summary" json:"last_turn_summary"` + ID uuid.UUID `db:"id" json:"id"` + ExpectedHistoryVersion int64 `db:"expected_history_version" json:"expected_history_version"` } // Updates the cached last completed turn summary for sidebar display. // Empty or whitespace-only summaries are stored as NULL here so direct // query callers cannot accidentally persist blank sidebar text. -// This intentionally preserves updated_at. The staleness guard relies on -// every new-turn query, such as UpdateChatStatus and AcquireChats, bumping -// updated_at. Future chat-field updates that do not bump updated_at can let -// stale summaries persist. If this query ever bumps updated_at, later -// goroutine summary writes will be rejected as stale. +// This intentionally preserves updated_at. The staleness guard uses +// history_version so worker lifecycle transitions that do not change the +// active message history cannot reject final turn summary writes. // Two summary workers using the same freshness marker are last-write-wins. func (q *sqlQuerier) UpdateChatLastTurnSummary(ctx context.Context, arg UpdateChatLastTurnSummaryParams) (int64, error) { - result, err := q.db.ExecContext(ctx, updateChatLastTurnSummary, arg.LastTurnSummary, arg.ID, arg.ExpectedUpdatedAt) + result, err := q.db.ExecContext(ctx, updateChatLastTurnSummary, arg.LastTurnSummary, arg.ID, arg.ExpectedHistoryVersion) if err != nil { return 0, err } @@ -10132,7 +11616,7 @@ SET updated_at = NOW() WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10164,6 +11648,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10173,7 +11665,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10214,6 +11706,14 @@ func (q *sqlQuerier) UpdateChatMCPServerIDs(ctx context.Context, arg UpdateChatM &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10231,7 +11731,7 @@ SET WHERE id = $3::bigint RETURNING - id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id + id, chat_id, model_config_id, created_at, role, content, visibility, input_tokens, output_tokens, total_tokens, reasoning_tokens, cache_creation_tokens, cache_read_tokens, context_limit, compressed, created_by, content_version, total_cost_micros, runtime_ms, deleted, provider_response_id, api_key_id, revision ` type UpdateChatMessageByIDParams struct { @@ -10266,6 +11766,7 @@ func (q *sqlQuerier) UpdateChatMessageByID(ctx context.Context, arg UpdateChatMe &i.Deleted, &i.ProviderResponseID, &i.APIKeyID, + &i.Revision, ) return i, err } @@ -10350,7 +11851,7 @@ SET plan_mode = $1::chat_plan_mode WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10382,6 +11883,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10391,7 +11900,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10432,6 +11941,14 @@ func (q *sqlQuerier) UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatP &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10440,20 +11957,14 @@ func (q *sqlQuerier) UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatP return i, err } -const updateChatStatus = `-- name: UpdateChatStatus :one +const updateChatRetryState = `-- name: UpdateChatRetryState :one WITH updated_chat AS ( -UPDATE - chats -SET - status = $1::chat_status, - worker_id = $2::uuid, - started_at = $3::timestamptz, - heartbeat_at = $4::timestamptz, - last_error = $5::jsonb, - updated_at = NOW() -WHERE - id = $6::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl + UPDATE chats + SET + retry_state = $1::jsonb, + updated_at = NOW() + WHERE id = $2::uuid + RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10485,6 +11996,134 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, + COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM updated_chat + LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = updated_chat.owner_id +) +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name +FROM chats_expanded +` + +type UpdateChatRetryStateParams struct { + RetryState json.RawMessage `db:"retry_state" json:"retry_state"` + ID uuid.UUID `db:"id" json:"id"` +} + +// Stores the client-visible retry payload. retry_state_version is +// assigned by trigger from the current snapshot_version. +func (q *sqlQuerier) UpdateChatRetryState(ctx context.Context, arg UpdateChatRetryStateParams) (Chat, error) { + row := q.db.QueryRowContext(ctx, updateChatRetryState, arg.RetryState, arg.ID) + var i Chat + err := row.Scan( + &i.ID, + &i.OwnerID, + &i.WorkspaceID, + &i.Title, + &i.Status, + &i.WorkerID, + &i.StartedAt, + &i.HeartbeatAt, + &i.CreatedAt, + &i.UpdatedAt, + &i.ParentChatID, + &i.RootChatID, + &i.LastModelConfigID, + &i.Archived, + &i.LastError, + &i.Mode, + pq.Array(&i.MCPServerIDs), + &i.Labels, + &i.BuildID, + &i.AgentID, + &i.PinOrder, + &i.LastReadMessageID, + &i.LastInjectedContext, + &i.DynamicTools, + &i.OrganizationID, + &i.PlanMode, + &i.ClientType, + &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, + &i.UserACL, + &i.GroupACL, + &i.OwnerUsername, + &i.OwnerName, + ) + return i, err +} + +const updateChatStatus = `-- name: UpdateChatStatus :one +WITH updated_chat AS ( +UPDATE + chats +SET + status = $1::chat_status, + worker_id = $2::uuid, + started_at = $3::timestamptz, + heartbeat_at = $4::timestamptz, + last_error = $5::jsonb, + updated_at = NOW() +WHERE + id = $6::uuid +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at +), +chats_expanded AS ( + SELECT + updated_chat.id, + updated_chat.owner_id, + updated_chat.workspace_id, + updated_chat.title, + updated_chat.status, + updated_chat.worker_id, + updated_chat.started_at, + updated_chat.heartbeat_at, + updated_chat.created_at, + updated_chat.updated_at, + updated_chat.parent_chat_id, + updated_chat.root_chat_id, + updated_chat.last_model_config_id, + updated_chat.archived, + updated_chat.last_error, + updated_chat.mode, + updated_chat.mcp_server_ids, + updated_chat.labels, + updated_chat.build_id, + updated_chat.agent_id, + updated_chat.pin_order, + updated_chat.last_read_message_id, + updated_chat.last_injected_context, + updated_chat.dynamic_tools, + updated_chat.organization_id, + updated_chat.plan_mode, + updated_chat.client_type, + updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10494,7 +12133,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10546,6 +12185,14 @@ func (q *sqlQuerier) UpdateChatStatus(ctx context.Context, arg UpdateChatStatusP &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10567,7 +12214,7 @@ SET updated_at = $6::timestamptz WHERE id = $7::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10599,6 +12246,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10608,7 +12263,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10662,6 +12317,14 @@ func (q *sqlQuerier) UpdateChatStatusPreserveUpdatedAt(ctx context.Context, arg &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10681,7 +12344,7 @@ SET title = $1::text WHERE id = $2::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10713,6 +12376,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10722,7 +12393,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10763,6 +12434,14 @@ func (q *sqlQuerier) UpdateChatTitleByID(ctx context.Context, arg UpdateChatTitl &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -10779,7 +12458,7 @@ UPDATE chats SET agent_id = $3::uuid, updated_at = NOW() WHERE id = $4::uuid -RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl +RETURNING id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at ), chats_expanded AS ( SELECT @@ -10811,6 +12490,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -10820,7 +12507,7 @@ chats_expanded AS ( LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) JOIN visible_users owner ON owner.id = updated_chat.owner_id ) -SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, user_acl, group_acl, owner_username, owner_name +SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, last_injected_context, dynamic_tools, organization_id, plan_mode, client_type, last_turn_summary, snapshot_version, history_version, queue_version, generation_attempt, retry_state, retry_state_version, runner_id, requires_action_deadline_at, user_acl, group_acl, owner_username, owner_name FROM chats_expanded ` @@ -10868,6 +12555,14 @@ func (q *sqlQuerier) UpdateChatWorkspaceBinding(ctx context.Context, arg UpdateC &i.PlanMode, &i.ClientType, &i.LastTurnSummary, + &i.SnapshotVersion, + &i.HistoryVersion, + &i.QueueVersion, + &i.GenerationAttempt, + &i.RetryState, + &i.RetryStateVersion, + &i.RunnerID, + &i.RequiresActionDeadlineAt, &i.UserACL, &i.GroupACL, &i.OwnerUsername, @@ -11095,6 +12790,25 @@ func (q *sqlQuerier) UpsertChatDiffStatusReference(ctx context.Context, arg Upse return i, err } +const upsertChatHeartbeat = `-- name: UpsertChatHeartbeat :exec +INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at) +VALUES ($1::uuid, $2::uuid, NOW()) +ON CONFLICT (chat_id, runner_id) DO UPDATE +SET heartbeat_at = EXCLUDED.heartbeat_at +` + +type UpsertChatHeartbeatParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + RunnerID uuid.UUID `db:"runner_id" json:"runner_id"` +} + +// Upserts a heartbeat row for the (chat_id, runner_id) lease. Uses +// database time so callers do not depend on a local clock. +func (q *sqlQuerier) UpsertChatHeartbeat(ctx context.Context, arg UpsertChatHeartbeatParams) error { + _, err := q.db.ExecContext(ctx, upsertChatHeartbeat, arg.ChatID, arg.RunnerID) + return err +} + const upsertChatUsageLimitConfig = `-- name: UpsertChatUsageLimitConfig :one INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at) VALUES (TRUE, $1::boolean, $2::bigint, $3::text, NOW()) diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 0d3810acb0..a395429f07 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -35,6 +35,14 @@ chats_expanded AS ( updated_chats.plan_mode, updated_chats.client_type, updated_chats.last_turn_summary, + updated_chats.snapshot_version, + updated_chats.history_version, + updated_chats.queue_version, + updated_chats.generation_attempt, + updated_chats.retry_state, + updated_chats.retry_state_version, + updated_chats.runner_id, + updated_chats.requires_action_deadline_at, COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -90,6 +98,14 @@ chats_expanded AS ( updated_chats.plan_mode, updated_chats.client_type, updated_chats.last_turn_summary, + updated_chats.snapshot_version, + updated_chats.history_version, + updated_chats.queue_version, + updated_chats.generation_attempt, + updated_chats.retry_state, + updated_chats.retry_state_version, + updated_chats.runner_id, + updated_chats.requires_action_deadline_at, COALESCE(root.user_acl, updated_chats.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -293,6 +309,15 @@ SELECT * FROM chats_expanded WHERE id = @id::uuid; +-- name: GetChatFamilyIDsByRootID :many +-- Returns the chat IDs of every chat in a family (root + all children) +-- in deterministic order. The id parameter must be the root id; the +-- query does not walk up from a child. +SELECT id +FROM chats +WHERE id = @id::uuid OR root_chat_id = @id::uuid +ORDER BY (id = @id::uuid) DESC, created_at ASC, id ASC; + -- name: GetChatACLByID :one SELECT user_acl AS users, @@ -333,6 +358,18 @@ WHERE ORDER BY created_at ASC; +-- name: GetChatMessagesByRevisionForStream :many +SELECT + * +FROM + chat_messages +WHERE + chat_id = @chat_id::uuid + AND revision > @after_revision::bigint + AND visibility IN ('user', 'both') +ORDER BY + created_at ASC, id ASC; + -- name: GetChatMessagesByChatIDAscPaginated :many SELECT * @@ -723,6 +760,14 @@ chats_expanded AS ( inserted_chat.plan_mode, inserted_chat.client_type, inserted_chat.last_turn_summary, + inserted_chat.snapshot_version, + inserted_chat.history_version, + inserted_chat.queue_version, + inserted_chat.generation_attempt, + inserted_chat.retry_state, + inserted_chat.retry_state_version, + inserted_chat.runner_id, + inserted_chat.requires_action_deadline_at, COALESCE(root.user_acl, inserted_chat.user_acl) AS user_acl, COALESCE(root.group_acl, inserted_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -860,6 +905,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -915,6 +968,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -968,6 +1029,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1021,6 +1090,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1074,6 +1151,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1126,6 +1211,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1178,6 +1271,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1232,6 +1333,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1248,11 +1357,9 @@ FROM chats_expanded; -- Updates the cached last completed turn summary for sidebar display. -- Empty or whitespace-only summaries are stored as NULL here so direct -- query callers cannot accidentally persist blank sidebar text. --- This intentionally preserves updated_at. The staleness guard relies on --- every new-turn query, such as UpdateChatStatus and AcquireChats, bumping --- updated_at. Future chat-field updates that do not bump updated_at can let --- stale summaries persist. If this query ever bumps updated_at, later --- goroutine summary writes will be rejected as stale. +-- This intentionally preserves updated_at. The staleness guard uses +-- history_version so worker lifecycle transitions that do not change the +-- active message history cannot reject final turn summary writes. -- Two summary workers using the same freshness marker are last-write-wins. UPDATE chats SET @@ -1261,7 +1368,7 @@ SET ), '') WHERE id = @id::uuid - AND updated_at = @expected_updated_at::timestamptz; + AND history_version = @expected_history_version::bigint; -- name: UpdateChatMCPServerIDs :one WITH updated_chat AS ( @@ -1304,6 +1411,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1413,6 +1528,14 @@ chats_expanded AS ( acquired_chats.plan_mode, acquired_chats.client_type, acquired_chats.last_turn_summary, + acquired_chats.snapshot_version, + acquired_chats.history_version, + acquired_chats.queue_version, + acquired_chats.generation_attempt, + acquired_chats.retry_state, + acquired_chats.retry_state_version, + acquired_chats.runner_id, + acquired_chats.requires_action_deadline_at, COALESCE(root.user_acl, acquired_chats.user_acl) AS user_acl, COALESCE(root.group_acl, acquired_chats.group_acl) AS group_acl, owner.username AS owner_username, @@ -1470,6 +1593,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1527,6 +1658,14 @@ chats_expanded AS ( updated_chat.plan_mode, updated_chat.client_type, updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1694,13 +1833,18 @@ RETURNING *; -- name: InsertChatQueuedMessage :one -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id) -VALUES ( - @chat_id, - @content, +-- Legacy queue insertion path. When no caller-supplied creator exists, +-- preserve the created_by invariant by attributing the queued row to the +-- chat owner. +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by) +SELECT + @chat_id::uuid, + @content::jsonb, sqlc.narg('model_config_id')::uuid, - sqlc.narg('api_key_id')::text -) + sqlc.narg('api_key_id')::text, + chats.owner_id +FROM chats +WHERE chats.id = @chat_id::uuid RETURNING *; -- name: GetChatQueuedMessages :many @@ -1786,6 +1930,14 @@ chats_expanded AS ( locked_chat.plan_mode, locked_chat.client_type, locked_chat.last_turn_summary, + locked_chat.snapshot_version, + locked_chat.history_version, + locked_chat.queue_version, + locked_chat.generation_attempt, + locked_chat.retry_state, + locked_chat.retry_state_version, + locked_chat.runner_id, + locked_chat.requires_action_deadline_at, COALESCE(root.user_acl, locked_chat.user_acl) AS user_acl, COALESCE(root.group_acl, locked_chat.group_acl) AS group_acl, owner.username AS owner_username, @@ -1798,6 +1950,63 @@ chats_expanded AS ( SELECT * FROM chats_expanded; +-- name: GetChatByIDForShare :one +WITH shared_chat AS ( + SELECT * + FROM chats + WHERE id = @id::uuid + FOR SHARE +), +chats_expanded AS ( + SELECT + shared_chat.id, + shared_chat.owner_id, + shared_chat.workspace_id, + shared_chat.title, + shared_chat.status, + shared_chat.worker_id, + shared_chat.started_at, + shared_chat.heartbeat_at, + shared_chat.created_at, + shared_chat.updated_at, + shared_chat.parent_chat_id, + shared_chat.root_chat_id, + shared_chat.last_model_config_id, + shared_chat.archived, + shared_chat.last_error, + shared_chat.mode, + shared_chat.mcp_server_ids, + shared_chat.labels, + shared_chat.build_id, + shared_chat.agent_id, + shared_chat.pin_order, + shared_chat.last_read_message_id, + shared_chat.last_injected_context, + shared_chat.dynamic_tools, + shared_chat.organization_id, + shared_chat.plan_mode, + shared_chat.client_type, + shared_chat.last_turn_summary, + shared_chat.snapshot_version, + shared_chat.history_version, + shared_chat.queue_version, + shared_chat.generation_attempt, + shared_chat.retry_state, + shared_chat.retry_state_version, + shared_chat.runner_id, + shared_chat.requires_action_deadline_at, + COALESCE(root.user_acl, shared_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, shared_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM + shared_chat + LEFT JOIN chats root ON root.id = COALESCE(shared_chat.root_chat_id, shared_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = shared_chat.owner_id +) +SELECT * +FROM chats_expanded; + -- name: GetChatsByChatFileID :many SELECT * @@ -2317,6 +2526,420 @@ WHERE chat_id = @chat_id::uuid AND deleted = false AND content::jsonb @> '[{"type": "context-file"}]'; +-- name: GetChatWorkerAcquisitionCandidates :many +-- Returns chats that workers may try to acquire. Candidates must be: +-- - in a worker-runnable execution status; +-- - unarchived; and +-- - missing ownership, carrying inconsistent ownership, or lacking a +-- fresh heartbeat for the assigned runner. +-- +-- Missing ownership is worker_id IS NULL. Inconsistent ownership is +-- runner_id IS NULL while worker_id is set. Stale ownership is no +-- heartbeat row for (chat_id, runner_id), or one older than +-- @stale_seconds by database time. Candidates are ordered by oldest +-- updated_at first so workers drain stale runnable chats predictably. +SELECT + chats_expanded.*, + chat_heartbeats.heartbeat_at AS current_heartbeat_at, + NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats_expanded.id + AND current_lease.runner_id = chats_expanded.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) + ) AS heartbeat_stale +FROM chats_expanded +LEFT JOIN chat_heartbeats + ON chat_heartbeats.chat_id = chats_expanded.id + AND chat_heartbeats.runner_id = chats_expanded.runner_id +WHERE + chats_expanded.status IN ('running'::chat_status, 'interrupting'::chat_status, 'requires_action'::chat_status) + AND chats_expanded.archived = false + AND ( + chats_expanded.worker_id IS NULL + OR chats_expanded.runner_id IS NULL + OR NOT EXISTS ( + SELECT 1 + FROM chat_heartbeats current_lease + WHERE current_lease.chat_id = chats_expanded.id + AND current_lease.runner_id = chats_expanded.runner_id + AND current_lease.heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) + ) + ) +ORDER BY chats_expanded.updated_at ASC, chats_expanded.id ASC +LIMIT @limit_count::int; + +-- name: GetChatsByIDsForRunnerSync :many +SELECT * +FROM chats_expanded +WHERE id = ANY(@ids::uuid[]) +ORDER BY id ASC; + +-- name: BatchUpsertChatHeartbeats :exec +INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at) +SELECT chat_ids.chat_id, runner_ids.runner_id, NOW() +FROM unnest(@chat_ids::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord) +JOIN unnest(@runner_ids::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord) +ON CONFLICT (chat_id, runner_id) DO UPDATE +SET heartbeat_at = EXCLUDED.heartbeat_at; + +-- name: DeleteStaleChatHeartbeats :execrows +DELETE FROM chat_heartbeats +WHERE heartbeat_at < NOW() - (INTERVAL '1 second' * @stale_seconds::int); + +-- name: GetAutoArchiveInactiveChatCandidates :many +-- Returns read-only root chat candidates for state-machine-backed +-- auto-archive. Activity is computed across the root family. The query +-- limits roots, not total family members. +SELECT + chats_expanded.*, + COALESCE(activity.last_activity_at, chats_expanded.created_at)::timestamptz AS last_activity_at +FROM chats_expanded +LEFT JOIN LATERAL ( + SELECT MAX(chat_messages.created_at) AS last_activity_at + FROM chat_messages + JOIN chats family_chat ON family_chat.id = chat_messages.chat_id + WHERE (family_chat.id = chats_expanded.id OR family_chat.root_chat_id = chats_expanded.id) + AND chat_messages.deleted = false +) activity ON TRUE +WHERE + chats_expanded.archived = false + AND chats_expanded.pin_order = 0 + AND chats_expanded.parent_chat_id IS NULL + AND chats_expanded.created_at < @archive_cutoff::timestamptz + AND chats_expanded.status NOT IN ( + 'running'::chat_status, + 'interrupting'::chat_status, + 'pending'::chat_status, + 'paused'::chat_status, + 'requires_action'::chat_status + ) + AND COALESCE(activity.last_activity_at, chats_expanded.created_at) < @archive_cutoff::timestamptz +ORDER BY chats_expanded.created_at ASC +LIMIT @limit_count::int; + + +-- name: LockChatAndBumpSnapshotVersion :one +-- Locks the chat row with FOR UPDATE and atomically increments its +-- snapshot_version, returning the post-bump chat. This is the single +-- entry point ChatMachine.Update uses to acquire the row lock and +-- allocate a new snapshot version in one round trip. +WITH bumped_chat AS ( + UPDATE chats + SET snapshot_version = snapshot_version + 1 + WHERE id = ( + SELECT id FROM chats + WHERE id = @id::uuid + FOR UPDATE + ) + RETURNING * +), +chats_expanded AS ( + SELECT + bumped_chat.id, + bumped_chat.owner_id, + bumped_chat.workspace_id, + bumped_chat.title, + bumped_chat.status, + bumped_chat.worker_id, + bumped_chat.started_at, + bumped_chat.heartbeat_at, + bumped_chat.created_at, + bumped_chat.updated_at, + bumped_chat.parent_chat_id, + bumped_chat.root_chat_id, + bumped_chat.last_model_config_id, + bumped_chat.archived, + bumped_chat.last_error, + bumped_chat.mode, + bumped_chat.mcp_server_ids, + bumped_chat.labels, + bumped_chat.build_id, + bumped_chat.agent_id, + bumped_chat.pin_order, + bumped_chat.last_read_message_id, + bumped_chat.last_injected_context, + bumped_chat.dynamic_tools, + bumped_chat.organization_id, + bumped_chat.plan_mode, + bumped_chat.client_type, + bumped_chat.last_turn_summary, + bumped_chat.snapshot_version, + bumped_chat.history_version, + bumped_chat.queue_version, + bumped_chat.generation_attempt, + bumped_chat.retry_state, + bumped_chat.retry_state_version, + bumped_chat.runner_id, + bumped_chat.requires_action_deadline_at, + COALESCE(root.user_acl, bumped_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, bumped_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM bumped_chat + LEFT JOIN chats root ON root.id = COALESCE(bumped_chat.root_chat_id, bumped_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = bumped_chat.owner_id +) +SELECT * +FROM chats_expanded; + +-- name: UpdateChatExecutionState :one +-- Atomically updates the execution-state-managed fields on a chat: +-- status, archived, last_error, ownership identifiers, and the +-- requires-action deadline. Callers compose this with transition +-- mutations inside a single ChatMachine.Update transaction. +WITH updated_chat AS ( + UPDATE chats + SET + status = @status::chat_status, + archived = @archived::boolean, + worker_id = sqlc.narg('worker_id')::uuid, + runner_id = sqlc.narg('runner_id')::uuid, + last_error = sqlc.narg('last_error')::jsonb, + requires_action_deadline_at = sqlc.narg('requires_action_deadline_at')::timestamptz, + pin_order = CASE WHEN @archived::boolean THEN 0 ELSE pin_order END, + updated_at = NOW() + WHERE id = @id::uuid + RETURNING * +), +chats_expanded AS ( + SELECT + updated_chat.id, + updated_chat.owner_id, + updated_chat.workspace_id, + updated_chat.title, + updated_chat.status, + updated_chat.worker_id, + updated_chat.started_at, + updated_chat.heartbeat_at, + updated_chat.created_at, + updated_chat.updated_at, + updated_chat.parent_chat_id, + updated_chat.root_chat_id, + updated_chat.last_model_config_id, + updated_chat.archived, + updated_chat.last_error, + updated_chat.mode, + updated_chat.mcp_server_ids, + updated_chat.labels, + updated_chat.build_id, + updated_chat.agent_id, + updated_chat.pin_order, + updated_chat.last_read_message_id, + updated_chat.last_injected_context, + updated_chat.dynamic_tools, + updated_chat.organization_id, + updated_chat.plan_mode, + updated_chat.client_type, + updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, + COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM updated_chat + LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = updated_chat.owner_id +) +SELECT * +FROM chats_expanded; + +-- name: UpdateChatRetryState :one +-- Stores the client-visible retry payload. retry_state_version is +-- assigned by trigger from the current snapshot_version. +WITH updated_chat AS ( + UPDATE chats + SET + retry_state = @retry_state::jsonb, + updated_at = NOW() + WHERE id = @id::uuid + RETURNING * +), +chats_expanded AS ( + SELECT + updated_chat.id, + updated_chat.owner_id, + updated_chat.workspace_id, + updated_chat.title, + updated_chat.status, + updated_chat.worker_id, + updated_chat.started_at, + updated_chat.heartbeat_at, + updated_chat.created_at, + updated_chat.updated_at, + updated_chat.parent_chat_id, + updated_chat.root_chat_id, + updated_chat.last_model_config_id, + updated_chat.archived, + updated_chat.last_error, + updated_chat.mode, + updated_chat.mcp_server_ids, + updated_chat.labels, + updated_chat.build_id, + updated_chat.agent_id, + updated_chat.pin_order, + updated_chat.last_read_message_id, + updated_chat.last_injected_context, + updated_chat.dynamic_tools, + updated_chat.organization_id, + updated_chat.plan_mode, + updated_chat.client_type, + updated_chat.last_turn_summary, + updated_chat.snapshot_version, + updated_chat.history_version, + updated_chat.queue_version, + updated_chat.generation_attempt, + updated_chat.retry_state, + updated_chat.retry_state_version, + updated_chat.runner_id, + updated_chat.requires_action_deadline_at, + COALESCE(root.user_acl, updated_chat.user_acl) AS user_acl, + COALESCE(root.group_acl, updated_chat.group_acl) AS group_acl, + owner.username AS owner_username, + owner.name AS owner_name + FROM updated_chat + LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) + JOIN visible_users owner ON owner.id = updated_chat.owner_id +) +SELECT * +FROM chats_expanded; + +-- name: IncrementChatGenerationAttempt :one +-- Increments generation_attempt and returns the resulting value. +UPDATE chats +SET generation_attempt = generation_attempt + 1, updated_at = NOW() +WHERE id = @id::uuid +RETURNING generation_attempt; + +-- name: GetDatabaseNow :one +-- Returns the current database timestamp. Used so transitions that +-- record deadlines or heartbeats rely on a clock that is consistent +-- with the database rather than the caller's local clock. +SELECT NOW()::timestamptz AS now; + +-- name: InsertChatQueuedMessageWithCreator :one +-- Inserts a queued message that carries a position (from the default +-- sequence) and an explicit created_by reference. Use this when the +-- queued-message creator differs from the chat owner. +INSERT INTO chat_queued_messages (chat_id, content, model_config_id, api_key_id, created_by) +VALUES ( + @chat_id::uuid, + @content::jsonb, + sqlc.narg('model_config_id')::uuid, + sqlc.narg('api_key_id')::text, + @created_by::uuid +) +RETURNING *; + +-- name: GetChatQueuedMessagesByPosition :many +-- Returns queued messages in state-machine order (position ASC, id ASC). +SELECT * FROM chat_queued_messages +WHERE chat_id = @chat_id::uuid +ORDER BY position ASC, id ASC; + +-- name: CountChatQueuedMessages :one +-- Cheap queue-length check used by ChatMachine.Update when deciding +-- whether the chat is in a "1" sub-state. +SELECT COUNT(*)::bigint AS count +FROM chat_queued_messages +WHERE chat_id = @chat_id::uuid; + +-- name: GetChatQueuedMessageHead :one +-- Returns the queue head (lowest position, then lowest id). +SELECT * FROM chat_queued_messages +WHERE chat_id = @chat_id::uuid +ORDER BY position ASC, id ASC +LIMIT 1; + +-- name: GetChatQueuedMessageByID :one +SELECT * FROM chat_queued_messages +WHERE id = @id::bigint AND chat_id = @chat_id::uuid; + +-- name: DeleteChatQueuedMessageReturningCount :execrows +-- Deletes a queued message, scoped to the parent chat. Returns the +-- number of affected rows so callers can detect missing rows without +-- a follow-up read. +DELETE FROM chat_queued_messages +WHERE id = @id::bigint AND chat_id = @chat_id::uuid; + +-- name: DeleteAllChatQueuedMessagesReturningCount :execrows +DELETE FROM chat_queued_messages +WHERE chat_id = @chat_id::uuid; + +-- name: ReorderChatQueuedMessageToHead :execrows +-- Sets the target queued message's position to one less than the +-- current minimum position for that chat, moving it to the head. +UPDATE chat_queued_messages AS target +SET position = COALESCE( + (SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = @chat_id::uuid), + 0 +) - 1 +WHERE target.id = @id::bigint + AND target.chat_id = @chat_id::uuid + AND target.position > COALESCE( + (SELECT MIN(position) FROM chat_queued_messages WHERE chat_id = @chat_id::uuid), + target.position + ); + +-- name: UpsertChatHeartbeat :exec +-- Upserts a heartbeat row for the (chat_id, runner_id) lease. Uses +-- database time so callers do not depend on a local clock. +INSERT INTO chat_heartbeats (chat_id, runner_id, heartbeat_at) +VALUES (@chat_id::uuid, @runner_id::uuid, NOW()) +ON CONFLICT (chat_id, runner_id) DO UPDATE +SET heartbeat_at = EXCLUDED.heartbeat_at; + +-- name: GetChatHeartbeat :one +SELECT * FROM chat_heartbeats +WHERE chat_id = @chat_id::uuid AND runner_id = @runner_id::uuid; + +-- name: IsChatHeartbeatStale :one +-- Returns true when there is no heartbeat row for (chat_id, runner_id) +-- or the existing row is older than @stale_seconds seconds by database +-- time. chatstate calls this in a single query so the staleness check +-- is atomic and does not depend on the caller's local clock. +SELECT NOT EXISTS ( + SELECT 1 FROM chat_heartbeats + WHERE chat_id = @chat_id::uuid + AND runner_id = @runner_id::uuid + AND heartbeat_at > NOW() - (INTERVAL '1 second' * @stale_seconds::int) +) AS stale; + +-- name: BatchDeleteChatHeartbeats :execrows +-- Deletes heartbeat rows for the supplied (chat_id, runner_id) pairs. +DELETE FROM chat_heartbeats +USING unnest(@chat_ids::uuid[]) WITH ORDINALITY AS chat_ids(chat_id, ord) +JOIN unnest(@runner_ids::uuid[]) WITH ORDINALITY AS runner_ids(runner_id, ord) USING (ord) +WHERE chat_heartbeats.chat_id = chat_ids.chat_id + AND chat_heartbeats.runner_id = runner_ids.runner_id; + +-- name: DeleteAllChatHeartbeats :exec +-- Deletes all heartbeat rows for the chat. Used during ownership +-- transitions that abandon a lease. +DELETE FROM chat_heartbeats WHERE chat_id = @chat_id::uuid; + + +-- name: GetChatStreamSyncRows :many +SELECT + id, + snapshot_version, + history_version, + queue_version, + retry_state_version, + generation_attempt, + status, + worker_id +FROM chats +WHERE id = ANY(@ids::uuid[]) +ORDER BY id ASC; + -- name: AutoArchiveInactiveChats :many -- Archives inactive root chats (pinned and already-archived chats skipped), -- cascading to children via root_chat_id. Limits apply to roots, not total diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index fd11ab2e06..347435f66a 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -26,6 +26,7 @@ const ( UniqueChatDiffStatusesPkey UniqueConstraint = "chat_diff_statuses_pkey" // ALTER TABLE ONLY chat_diff_statuses ADD CONSTRAINT chat_diff_statuses_pkey PRIMARY KEY (chat_id); UniqueChatFileLinksChatIDFileIDKey UniqueConstraint = "chat_file_links_chat_id_file_id_key" // ALTER TABLE ONLY chat_file_links ADD CONSTRAINT chat_file_links_chat_id_file_id_key UNIQUE (chat_id, file_id); UniqueChatFilesPkey UniqueConstraint = "chat_files_pkey" // ALTER TABLE ONLY chat_files ADD CONSTRAINT chat_files_pkey PRIMARY KEY (id); + UniqueChatHeartbeatsPkey UniqueConstraint = "chat_heartbeats_pkey" // ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id); UniqueChatMessagesPkey UniqueConstraint = "chat_messages_pkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id); UniqueChatModelConfigsPkey UniqueConstraint = "chat_model_configs_pkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_pkey PRIMARY KEY (id); UniqueChatQueuedMessagesPkey UniqueConstraint = "chat_queued_messages_pkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_pkey PRIMARY KEY (id); diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 3375267575..bca949f531 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -50,6 +50,7 @@ import ( "github.com/coder/coder/v2/coderd/wsbuilder" "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/coderd/x/chatfiles" "github.com/coder/coder/v2/coderd/x/gitsync" @@ -114,30 +115,6 @@ func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error) return false } -func publishChatTitleChange(logger slog.Logger, ps dbpubsub.Pubsub, chat database.Chat) { - if ps == nil { - return - } - event := codersdk.ChatWatchEvent{ - Kind: codersdk.ChatWatchEventKindTitleChange, - Chat: db2sdk.Chat(chat, nil, nil), - } - payload, err := json.Marshal(event) - if err != nil { - logger.Error(context.Background(), "failed to marshal chat title change event", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - return - } - if err := ps.Publish(pubsub.ChatWatchEventChannel(chat.OwnerID), payload); err != nil { - logger.Error(context.Background(), "failed to publish chat title change event", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - } -} - func publishChatConfigEvent(logger slog.Logger, ps dbpubsub.Pubsub, kind pubsub.ChatConfigEventKind, entityID uuid.UUID) { payload, err := json.Marshal(pubsub.ChatConfigEvent{ Kind: kind, @@ -1133,14 +1110,6 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { title := chatTitleFromMessage(titleSource) - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - modelConfigID, modelConfigStatus, modelConfigError := api.resolveCreateChatModelConfigID(ctx, apiKey.UserID, req) if modelConfigError != nil { httpapi.Write(ctx, rw, modelConfigStatus, *modelConfigError) @@ -1327,6 +1296,12 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { } aReq.New = chat + // Kick off best-effort automatic title generation now that the + // chat and its initial user message are persisted. It runs + // detached so it never blocks the create response, and only acts + // on the first user turn. + api.chatDaemon.GenerateChatTitleAsync(ctx, chat) + chatFiles := api.fetchChatFileMetadata(ctx, chat.ID) response := db2sdk.Chat(chat, nil, chatFiles) if len(unlinked) > 0 { @@ -1352,14 +1327,6 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { func (api *API) listChatModels(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() apiKey := httpmw.APIKey(r) - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - availability, err := api.getUserChatProviderAvailability(ctx, apiKey.UserID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ @@ -2626,35 +2593,7 @@ func (api *API) applyChatTitleUpdate( return chat, false } - var ( - updatedChat database.Chat - wrote bool - err error - ) - if api.chatDaemon != nil { - updatedChat, wrote, err = api.chatDaemon.RenameChatTitle(ctx, chat, trimmedTitle) - } else { - err = api.Database.InTx(func(tx database.Store) error { - currentChat, txErr := tx.GetChatByID(ctx, chat.ID) - if txErr != nil { - return txErr - } - if trimmedTitle == currentChat.Title { - updatedChat = currentChat - wrote = false - return nil - } - updatedChat, txErr = tx.UpdateChatTitleByID(ctx, database.UpdateChatTitleByIDParams{ - ID: chat.ID, - Title: trimmedTitle, - }) - if txErr != nil { - return txErr - } - wrote = true - return nil - }, nil) - } + updatedChat, wrote, err := api.chatDaemon.RenameChatTitle(ctx, chat, trimmedTitle) if err != nil { if errors.Is(err, chatd.ErrManualTitleRegenerationInProgress) { httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ @@ -2673,11 +2612,7 @@ func (api *API) applyChatTitleUpdate( return chat, true } if wrote { - if api.chatDaemon != nil { - api.chatDaemon.PublishTitleChange(updatedChat) - } else { - publishChatTitleChange(api.Logger, api.Pubsub, updatedChat) - } + api.chatDaemon.PublishTitleChange(updatedChat) } return updatedChat, false } @@ -2774,6 +2709,21 @@ func (api *API) patchChat(rw http.ResponseWriter, r *http.Request) { if req.Archived != nil { archived := *req.Archived + + // Archive invariant is one-way: parent archived implies + // child archived. Archive state changes target the root + // chat and cascade atomically across the family; child + // chats cannot be archived or unarchived independently. + // This check precedes the no-op check so any child attempt + // surfaces the root-only error regardless of the chat's + // current archived value. + if chat.ParentChatID.Valid { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Chat archive state can only be changed on the root chat.", + }) + return + } + if archived == chat.Archived { state := "archived" if !archived { @@ -2785,37 +2735,30 @@ func (api *API) patchChat(rw http.ResponseWriter, r *http.Request) { return } - // Archive invariant is one-way: parent archived implies - // child archived. Parent archive/unarchive cascade via - // root_chat_id; individual child archive is permitted; - // child unarchive while the parent is archived is rejected - // (enforced atomically in chatd.Server.UnarchiveChat). - if chat.ParentChatID.Valid && !archived { - if done := api.writeChildUnarchiveGuard(ctx, rw, chat); done { - return - } - } var err error - // Use chatDaemon when available so it can interrupt active - // processing before broadcasting archive state. Fall back to - // direct DB when no daemon is running. if archived { - if api.chatDaemon != nil { - err = api.chatDaemon.ArchiveChat(ctx, chat) - } else { - _, err = api.Database.ArchiveChatByID(ctx, chat.ID) - } + err = api.chatDaemon.ArchiveChat(ctx, chat) } else { - if api.chatDaemon != nil { - err = api.chatDaemon.UnarchiveChat(ctx, chat) - } else { - _, err = api.Database.UnarchiveChatByID(ctx, chat.ID) - } + err = api.chatDaemon.UnarchiveChat(ctx, chat) } if err != nil { - if errors.Is(err, chatd.ErrChildUnarchiveParentArchived) { + if errors.Is(err, chatd.ErrArchiveRequiresRootChat) || errors.Is(err, chatstate.ErrChatNotRoot) { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Cannot unarchive a child chat while its parent is archived. Unarchive the parent chat to cascade.", + Message: "Chat archive state can only be changed on the root chat.", + }) + return + } + if writeChatInvalidState(ctx, rw, err) { + return + } + if errors.Is(err, chatstate.ErrTransitionNotAllowed) { + // Archive only succeeds from idle / error execution + // states (W, E0, E1) per the chatd RFC; active + // chats refuse archive instead of being silently + // transitioned to waiting first. + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Cannot archive an active chat. Interrupt or wait for the chat to finish first.", + Detail: err.Error(), }) return } @@ -2965,36 +2908,38 @@ func (api *API) patchChat(rw http.ResponseWriter, r *http.Request) { rw.WriteHeader(http.StatusNoContent) } -// writeChildUnarchiveGuard returns a 400 early when a child unarchive -// request obviously races an archived parent. The durable invariant -// is enforced atomically in chatd.Server.UnarchiveChat; this guard -// just surfaces the error before we take any locks. -// +// writeChatInvalidState writes the shared invalid-state response for +// chatstate.ErrInvalidState across every chat mutation endpoint. // Returns true when a response has been written. -func (api *API) writeChildUnarchiveGuard( - ctx context.Context, - rw http.ResponseWriter, - chat database.Chat, -) bool { - parent, err := api.Database.GetChatByID(ctx, chat.ParentChatID.UUID) - if err != nil { - if errors.Is(err, sql.ErrNoRows) { - httpapi.ResourceNotFound(rw) - return true +func writeChatInvalidState(ctx context.Context, rw http.ResponseWriter, err error) bool { + if !errors.Is(err, chatstate.ErrInvalidState) { + return false + } + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is in an invalid state.", + }) + return true +} + +// writeCommonChatMutationError writes responses shared by chat +// mutation endpoints. Returns true when a response has been written. +func writeCommonChatMutationError(ctx context.Context, rw http.ResponseWriter, err error, archivedMessage string) bool { + switch { + case xerrors.Is(err, chatd.ErrChatArchived): + if archivedMessage == "" { + archivedMessage = "Cannot mutate an archived chat." } - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to load parent chat.", - Detail: err.Error(), - }) - return true - } - if parent.Archived { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Cannot unarchive a child chat while its parent is archived. Unarchive the parent chat to cascade.", + Message: archivedMessage, }) - return true + case writeChatInvalidState(ctx, rw, err): + // response already written + case errors.Is(err, chatstate.ErrChatNotFound), httpapi.Is404Error(err): + httpapi.ResourceNotFound(rw) + default: + return false } - return false + return true } // EXPERIMENTAL: this endpoint is experimental and is subject to change. @@ -3043,14 +2988,6 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { return } - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - var req codersdk.CreateChatMessageRequest if !httpapi.Read(ctx, rw, r, &req) { return @@ -3152,10 +3089,15 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { }) return } - if xerrors.Is(sendErr, chatd.ErrMessageQueueFull) { + if xerrors.Is(sendErr, chatstate.ErrMessageQueueFull) { + var queueFull *chatstate.MessageQueueFullError + detail := "" + if errors.As(sendErr, &queueFull) { + detail = fmt.Sprintf("Maximum %d messages can be queued.", queueFull.Max) + } httpapi.Write(ctx, rw, http.StatusTooManyRequests, codersdk.Response{ Message: "Message queue is full.", - Detail: fmt.Sprintf("Maximum %d messages can be queued.", chatd.MaxQueueSize), + Detail: detail, }) return } @@ -3165,6 +3107,20 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { }) return } + if errors.Is(sendErr, chatstate.ErrChatNotFound) { + httpapi.ResourceNotFound(rw) + return + } + if writeChatInvalidState(ctx, rw, sendErr) { + return + } + if errors.Is(sendErr, chatstate.ErrTransitionNotAllowed) { + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not in a state that accepts new messages.", + Detail: sendErr.Error(), + }) + return + } httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to create chat message.", Detail: sendErr.Error(), @@ -3235,14 +3191,6 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { return } - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - messageIDStr := chi.URLParam(r, "message") messageID, err := strconv.ParseInt(messageIDStr, 10, 64) if err != nil || messageID <= 0 { @@ -3303,6 +3251,15 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Invalid model config ID.", }) + case errors.Is(editErr, chatstate.ErrChatNotFound): + httpapi.ResourceNotFound(rw) + case writeChatInvalidState(ctx, rw, editErr): + // response already written + case errors.Is(editErr, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not in a state that accepts message edits.", + Detail: editErr.Error(), + }) default: httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to edit chat message.", @@ -3349,19 +3306,28 @@ func (api *API) deleteChatQueuedMessage(rw http.ResponseWriter, r *http.Request) return } - if api.chatDaemon != nil { - err = api.chatDaemon.DeleteQueued(ctx, chatID, queuedMessageID) - } else { - err = api.Database.DeleteChatQueuedMessage(ctx, database.DeleteChatQueuedMessageParams{ - ID: queuedMessageID, - ChatID: chatID, - }) - } + err = api.chatDaemon.DeleteQueued(ctx, chatID, queuedMessageID) if err != nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to delete queued message.", - Detail: err.Error(), - }) + switch { + case xerrors.Is(err, chatstate.ErrQueuedMessageNotFound), xerrors.Is(err, sql.ErrNoRows): + httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ + Message: "Queued message not found.", + }) + case errors.Is(err, chatstate.ErrChatNotFound): + httpapi.ResourceNotFound(rw) + case writeChatInvalidState(ctx, rw, err): + // response already written + case errors.Is(err, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat has no queued messages to delete.", + Detail: err.Error(), + }) + default: + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to delete queued message.", + Detail: err.Error(), + }) + } return } @@ -3408,14 +3374,6 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request return } - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - _, txErr := api.chatDaemon.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ ChatID: chatID, CreatedBy: apiKey.UserID, @@ -3426,16 +3384,30 @@ func (api *API) promoteChatQueuedMessage(rw http.ResponseWriter, r *http.Request if maybeWriteLimitErr(ctx, rw, txErr) { return } - if xerrors.Is(txErr, chatd.ErrChatArchived) { + switch { + case xerrors.Is(txErr, chatd.ErrChatArchived): httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Cannot promote queued messages in an archived chat.", }) - return + case xerrors.Is(txErr, chatstate.ErrQueuedMessageNotFound): + httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ + Message: "Queued message not found.", + }) + case errors.Is(txErr, chatstate.ErrChatNotFound): + httpapi.ResourceNotFound(rw) + case writeChatInvalidState(ctx, rw, txErr): + // response already written + case errors.Is(txErr, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat has no queued messages to promote.", + Detail: txErr.Error(), + }) + default: + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to promote queued message.", + Detail: txErr.Error(), + }) } - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to promote queued message.", - Detail: txErr.Error(), - }) return } @@ -3494,14 +3466,6 @@ func (api *API) streamChat(rw http.ResponseWriter, r *http.Request) { chatID := chat.ID logger := api.Logger.Named("chat_streamer").With(slog.F("chat_id", chatID)) - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat streaming is not available.", - Detail: "Chat processor is not configured.", - }) - return - } - var afterMessageID int64 if v := r.URL.Query().Get("after_id"); v != "" { var err error @@ -3518,9 +3482,7 @@ func (api *API) streamChat(rw http.ResponseWriter, r *http.Request) { // Subscribe before accepting the WebSocket so that failures // can still be reported as normal HTTP errors. snapshot, events, cancelSub, ok := api.chatDaemon.SubscribeAuthorized(ctx, chat, r.Header, afterMessageID) - // Subscribe only fails today when the receiver is nil, which - // the chatDaemon == nil guard above already catches. This is - // defensive against future Subscribe failure modes. + // Defensive against future SubscribeAuthorized failure modes. if !ok { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Chat streaming is not available.", @@ -3646,31 +3608,79 @@ func (api *API) interruptChat(rw http.ResponseWriter, r *http.Request) { return } - if api.chatDaemon != nil { - chat = api.chatDaemon.InterruptChat(ctx, chat) - } else { - updatedChat, updateErr := api.Database.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - if updateErr != nil { - logger.Error(ctx, "failed to mark chat as waiting", slog.Error(updateErr)) - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to interrupt chat.", - Detail: updateErr.Error(), - }) + updated, err := api.chatDaemon.InterruptChat(ctx, chat) + if err != nil { + if writeCommonChatMutationError(ctx, rw, err, "Cannot interrupt an archived chat.") { return } - chat = updatedChat + switch { + case errors.Is(err, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not in an interruptible state.", + Detail: err.Error(), + }) + default: + logger.Error(ctx, "failed to interrupt chat", slog.Error(err)) + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to interrupt chat.", + Detail: err.Error(), + }) + } + return } + chat = updated httpapi.Write(ctx, rw, http.StatusOK, db2sdk.Chat(chat, nil, nil)) } +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +// @Summary Reconcile invalid chat state +// @ID reconcile-invalid-chat-state +// @Security CoderSessionToken +// @Tags Chats +// @Produce json +// @Param chat path string true "Chat ID" format(uuid) +// @Success 200 {object} codersdk.Chat +// @Router /api/experimental/chats/{chat}/reconcile-invalid [post] +// @Description Experimental: this endpoint is subject to change. +// +//nolint:revive // HTTP handler writes to ResponseWriter. +func (api *API) reconcileInvalidChatState(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + chat := httpmw.ChatParam(r) + chatID := chat.ID + logger := api.Logger.Named("chat_reconcile_invalid").With(slog.F("chat_id", chatID)) + + if !api.Authorize(r, policy.ActionUpdate, chat.RBACObject()) { + httpapi.ResourceNotFound(rw) + return + } + + updated, err := api.chatDaemon.ReconcileInvalidStateChat(ctx, chat) + if err != nil { + if writeCommonChatMutationError(ctx, rw, err, "") { + return + } + switch { + case errors.Is(err, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not in an invalid state.", + Detail: err.Error(), + }) + default: + logger.Error(ctx, "failed to reconcile invalid chat state", slog.Error(err)) + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to reconcile chat state.", + Detail: err.Error(), + }) + } + return + } + + httpapi.Write(ctx, rw, http.StatusOK, db2sdk.Chat(updated, nil, nil)) +} + // EXPERIMENTAL: this endpoint is experimental and is subject to change. // // @Summary Regenerate chat title @@ -3703,14 +3713,6 @@ func (api *API) regenerateChatTitle(rw http.ResponseWriter, r *http.Request) { return } - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) updatedChat, err := api.chatDaemon.RegenerateChatTitle(ctx, chat) if err != nil { @@ -3757,14 +3759,6 @@ func (api *API) proposeChatTitle(rw http.ResponseWriter, r *http.Request) { return } - if api.chatDaemon == nil { - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Chat processor is unavailable.", - Detail: "Chat processor is not configured.", - }) - return - } - ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) title, err := api.chatDaemon.ProposeChatTitle(ctx, chat) if err != nil { @@ -7857,15 +7851,10 @@ func (api *API) postChatToolResults(rw http.ResponseWriter, r *http.Request) { return } - // Fast-path check outside the transaction. The authoritative - // check happens inside SubmitToolResults under a row lock. - if chat.Status != database.ChatStatusRequiresAction { - httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ - Message: "Chat is not waiting for tool results.", - Detail: fmt.Sprintf("Chat status is %q, expected %q.", chat.Status, database.ChatStatusRequiresAction), - }) - return - } + // The authoritative status check happens inside SubmitToolResults + // under the row lock; that path also surfaces the shared + // invalid-state response for chats that are not in a valid + // execution state at all. var dynamicTools json.RawMessage if chat.DynamicTools.Valid { @@ -7897,6 +7886,15 @@ func (api *API) postChatToolResults(rw http.ResponseWriter, r *http.Request) { Message: validationErr.Message, Detail: validationErr.Detail, }) + case errors.Is(err, chatstate.ErrChatNotFound): + httpapi.ResourceNotFound(rw) + case writeChatInvalidState(ctx, rw, err): + // response already written + case errors.Is(err, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not waiting for tool results.", + Detail: err.Error(), + }) default: api.Logger.Error(ctx, "tool results submission failed", slog.F("chat_id", chat.ID), @@ -8005,3 +8003,23 @@ func (api *API) getChatDebugRun(rw http.ResponseWriter, r *http.Request) { httpapi.Write(ctx, rw, http.StatusOK, db2sdk.ChatDebugRunDetail(run, steps)) } + +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +// @Summary Stream chat parts via WebSockets +// @ID stream-chat-parts-via-websockets +// @Security CoderSessionToken +// @Tags Chats +// @Produce json +// @Param chat path string true "Chat ID" format(uuid) +// @Success 200 {object} codersdk.ChatStreamEvent +// @Router /api/experimental/chats/{chat}/stream/parts [get] +// @x-apidocgen {"skip": true} +// @Description Experimental: this endpoint is subject to change. +func (api *API) streamChatParts(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + chat := httpmw.ChatParam(r) + if err := api.chatDaemon.ServeStreamPartsAuthorized(rw, r, chat); err != nil { + api.Logger.Named("chat_stream_parts").Debug(ctx, "chat stream parts closed", slog.Error(err)) + } +} diff --git a/coderd/exp_chats_chatstate_test.go b/coderd/exp_chats_chatstate_test.go new file mode 100644 index 0000000000..485b5d21b7 --- /dev/null +++ b/coderd/exp_chats_chatstate_test.go @@ -0,0 +1,771 @@ +package coderd_test + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd" + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/util/ptr" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// driveChatToWaiting transitions the chat from `running` (its initial +// state per the RFC) to `waiting` by running chatstate.FinishTurn. +// Tests use this when they need to exercise endpoint behavior that +// only succeeds from idle execution states (W, E0). +func driveChatToWaiting(ctx context.Context, t *testing.T, api *coderd.API, chatID uuid.UUID) { + t.Helper() + chatdCtx := dbauthz.AsChatd(ctx) //nolint:gocritic // Test fixture mirrors chatd background transitions. + machine := chatstate.NewChatMachine(api.Database, api.Pubsub, chatID) + require.NoError(t, machine.Update(chatdCtx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) +} + +// driveChatToRequiresAction commits an assistant message with a single +// dynamic tool_call part and then transitions the chat to +// `requires_action`. The tool_call_id returned lets the caller +// assemble a valid SubmitToolResultsRequest. +func driveChatToRequiresAction( + ctx context.Context, + t *testing.T, + api *coderd.API, + chat codersdk.Chat, + toolName string, +) (toolCallID string) { + t.Helper() + chatdCtx := dbauthz.AsChatd(ctx) //nolint:gocritic // Test fixture mirrors chatd background transitions. + + toolCallID = "call-" + uuid.NewString() + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("dispatching dynamic tool"), + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: toolCallID, + ToolName: toolName, + Args: json.RawMessage(`{}`), + }, + }) + require.NoError(t, err) + + machine := chatstate.NewChatMachine(api.Database, api.Pubsub, chat.ID) + require.NoError(t, machine.Update(chatdCtx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{{ + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + }}, + }) + if err != nil { + return err + } + _, err = tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + return toolCallID +} + +// TestPostChatsStartsRunning verifies the RFC-mandated `running` +// initial status surfaced by the create-chat endpoint. +func TestPostChatsStartsRunning(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, api := newChatClientWithAPI(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "hello", + }}, + }) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusRunning, chat.Status, + "new chats must start in `running` per chatd RFC") + + // Re-reading also reports `running` because the chat row is + // authoritative and no worker has advanced it. + gotChat, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusRunning, gotChat.Status) + require.NotNil(t, api.Pubsub) +} + +// TestArchiveChatStateTransitions covers the two RFC-mandated archive +// behaviors at the endpoint contract level: archiving from an idle +// chat (W) succeeds, and archiving from an active chat (R0) returns +// a state conflict and leaves the chat unarchived. +func TestArchiveChatStateTransitions(t *testing.T) { + t.Parallel() + + t.Run("IdleSucceeds", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, api := newChatClientWithAPI(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "archive me"}}, + }) + require.NoError(t, err) + + driveChatToWaiting(ctx, t, api, chat.ID) + + err = client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) + require.NoError(t, err) + + got, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.True(t, got.Archived) + }) + + t.Run("ActiveChatReturnsConflict", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "no archive"}}, + }) + require.NoError(t, err) + + err = client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) + requireSDKError(t, err, http.StatusConflict) + + got, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.False(t, got.Archived, "active chat must remain unarchived after a conflict") + }) +} + +// TestPostChatMessagesBusyInterrupt verifies that a busy-interrupt +// send returns a queued response and leaves the chat in `interrupting` +// from the endpoint's perspective. +func TestPostChatMessagesBusyInterrupt(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusRunning, chat.Status) + + // CreateChat leaves the chat in `running`; an interrupt-style + // follow-up should land it in `interrupting`. + resp, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "stop"}}, + BusyBehavior: codersdk.ChatBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, resp.Queued, "busy interrupt must return queued=true") + require.NotNil(t, resp.QueuedMessage) + + got, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusInterrupting, got.Status, + "busy interrupt send must land the chat in `interrupting`") +} + +// TestDeleteChatQueuedMessageMissingReturns404 covers the new +// chatstate-driven 404 path for missing queued IDs. The chat must +// have at least one queued message so the request is in a state where +// DeleteQueuedMessage is allowed; the looked-up ID then mismatches +// and the endpoint returns 404 instead of a state-conflict 409. +func TestDeleteChatQueuedMessageMissingReturns404(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + // Seed one queued message via the public endpoint (the chat + // starts in R0, so a queue send lands in R1). + _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "queued"}}, + BusyBehavior: codersdk.ChatBusyBehaviorQueue, + }) + require.NoError(t, err) + + res, err := client.Request( + ctx, + http.MethodDelete, + fmt.Sprintf("/api/experimental/chats/%s/queue/99999999", chat.ID), + nil, + ) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusNotFound, res.StatusCode) +} + +// TestDeleteChatQueuedMessageEmptyQueueReturnsConflict covers the +// state-conflict 409 path when the chat has no queued messages. +func TestDeleteChatQueuedMessageEmptyQueueReturnsConflict(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + res, err := client.Request( + ctx, + http.MethodDelete, + fmt.Sprintf("/api/experimental/chats/%s/queue/99999999", chat.ID), + nil, + ) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusConflict, res.StatusCode) +} + +// TestPromoteChatQueuedMessageMissingReturns404 mirrors the delete +// test for the promote endpoint: with a non-empty queue, an unknown +// queued-message ID returns 404 rather than a 409. +func TestPromoteChatQueuedMessageMissingReturns404(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + // Seed one queued message so the promote transition is allowed. + _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "queued"}}, + BusyBehavior: codersdk.ChatBusyBehaviorQueue, + }) + require.NoError(t, err) + + res, err := client.Request( + ctx, + http.MethodPost, + fmt.Sprintf("/api/experimental/chats/%s/queue/99999999/promote", chat.ID), + nil, + ) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusNotFound, res.StatusCode) +} + +// TestPromoteChatQueuedMessageEmptyQueueReturnsConflict verifies the +// state-conflict 409 path when the chat has no queued messages. +func TestPromoteChatQueuedMessageEmptyQueueReturnsConflict(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + res, err := client.Request( + ctx, + http.MethodPost, + fmt.Sprintf("/api/experimental/chats/%s/queue/99999999/promote", chat.ID), + nil, + ) + require.NoError(t, err) + defer res.Body.Close() + require.Equal(t, http.StatusConflict, res.StatusCode) +} + +// TestInterruptChatIdleReturnsConflict verifies that interrupting an +// idle chat is now rejected. The fixture composes chatstate +// transitions to reach the W state without depending on the +// background worker. +func TestInterruptChatIdleReturnsConflict(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, api := newChatClientWithAPI(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "interrupt me"}}, + }) + require.NoError(t, err) + + driveChatToWaiting(ctx, t, api, chat.ID) + + _, err = client.InterruptChat(ctx, chat.ID) + requireSDKError(t, err, http.StatusConflict) +} + +// TestSubmitToolResultsWrongStateReturnsConflict covers the wrong +// chat-status response when the chat is not in requires_action. +func TestSubmitToolResultsWrongStateReturnsConflict(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusRunning, chat.Status) + + err = client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + Results: []codersdk.ToolResult{{ + ToolCallID: "unknown-call", + Output: json.RawMessage(`{}`), + }}, + }) + requireSDKError(t, err, http.StatusConflict) +} + +// TestSubmitToolResultsRequiresActionSucceeds drives a chat into +// requires_action with a single dynamic tool call and verifies a +// matching SubmitToolResults call returns 204 with the tool result +// persisted. +func TestSubmitToolResultsRequiresActionSucceeds(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, api := newChatClientWithAPI(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + dynamicTools := []codersdk.DynamicTool{{ + Name: "echo", + Description: "test echo tool", + InputSchema: json.RawMessage(`{"type":"object"}`), + }} + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + UnsafeDynamicTools: dynamicTools, + }) + require.NoError(t, err) + + toolCallID := driveChatToRequiresAction(ctx, t, api, chat, "echo") + + err = client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + Results: []codersdk.ToolResult{{ + ToolCallID: toolCallID, + Output: json.RawMessage(`{"ok":true}`), + }}, + }) + require.NoError(t, err) + + // The tool result must be persisted as a visible tool message. + got, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + foundToolResult := false + for _, msg := range got.Messages { + if msg.Role != codersdk.ChatMessageRoleTool { + continue + } + for _, part := range msg.Content { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolCallID == toolCallID { + foundToolResult = true + break + } + } + } + require.True(t, foundToolResult, "tool result message must be visible in chat history") +} + +// TestPatchChatArchiveChildRejected verifies that PATCH /api/experimental/chats/{child} +// with archived=true returns the root-only error regardless of the +// child's current archived value, and does not change archive state on +// any family member. +func TestPatchChatArchiveChildRejected(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + root, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "root"}}, + }) + require.NoError(t, err) + driveChatToWaiting(ctx, t, api, root.ID) + + // Sibling child A and B; both unarchived. + childA := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child-a", + Status: database.ChatStatusWaiting, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + childB := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child-b", + Status: database.ChatStatusWaiting, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + err = client.UpdateChat(ctx, childA.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) + requireSDKError(t, err, http.StatusBadRequest) + + for _, id := range []uuid.UUID{root.ID, childA.ID, childB.ID} { + got, gerr := loadChatRow(ctx, db, id) + require.NoError(t, gerr) + require.False(t, got.Archived, "no family member may flip archive state after a rejected child archive") + } +} + +// TestPatchChatUnarchiveChildRejected verifies that PATCH /api/experimental/chats/{child} +// with archived=false on an archived family is rejected with the +// root-only error and leaves every family member archived. The child +// already matches the requested value? No, the family is archived; +// we are asking to unarchive a child individually. +func TestPatchChatUnarchiveChildRejected(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + root, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "root"}}, + }) + require.NoError(t, err) + driveChatToWaiting(ctx, t, api, root.ID) + + childA := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child-a", + Status: database.ChatStatusWaiting, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + childB := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child-b", + Status: database.ChatStatusWaiting, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + // Archive the whole family via the root. + err = client.UpdateChat(ctx, root.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) + require.NoError(t, err) + for _, id := range []uuid.UUID{root.ID, childA.ID, childB.ID} { + got, gerr := loadChatRow(ctx, db, id) + require.NoError(t, gerr) + require.True(t, got.Archived, "precondition: family archived after root archive") + } + + // Unarchiving a child must be rejected. + err = client.UpdateChat(ctx, childA.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(false)}) + requireSDKError(t, err, http.StatusBadRequest) + + for _, id := range []uuid.UUID{root.ID, childA.ID, childB.ID} { + got, gerr := loadChatRow(ctx, db, id) + require.NoError(t, gerr) + require.True(t, got.Archived, "no family member may flip archive state after a rejected child unarchive") + } +} + +// TestPatchChatArchiveRootRollsBackWhenChildCannotArchive verifies the +// family-archive atomicity guarantee surfaced through the endpoint: +// when a child is in a state that rejects SetArchived (running here), +// the whole cascade rolls back and no family member changes archive +// state. +func TestPatchChatArchiveRootRollsBackWhenChildCannotArchive(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + root, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "root"}}, + }) + require.NoError(t, err) + driveChatToWaiting(ctx, t, api, root.ID) + + // Child is running (R0) which is NOT archive-eligible. + child := dbgen.Chat(t, db, database.Chat{ + OrganizationID: firstUser.OrganizationID, + OwnerID: firstUser.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child", + Status: database.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + err = client.UpdateChat(ctx, root.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) + requireSDKError(t, err, http.StatusConflict) + + for _, id := range []uuid.UUID{root.ID, child.ID} { + got, gerr := loadChatRow(ctx, db, id) + require.NoError(t, gerr) + require.False(t, got.Archived, "rolled-back family archive must not leave any member archived") + } +} + +// TestPostChatMessagesInvalidStateReturnsSharedResponse drives a chat +// into the chatstate-invalid state (waiting with a queued backlog) +// and asserts the shared invalid-state response. This is the +// representative endpoint required by the review. +func TestPostChatMessagesInvalidStateReturnsSharedResponse(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, _, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + // Drive the chat to an invalid combination: status=waiting (W), + // archived=false, and a queued message. ClassifyExecutionState + // returns StateInvalid for (waiting, queue=true). + driveChatToInvalidWaitingWithQueue(ctx, t, api, chat.ID) + + _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "send"}}, + }) + sdkErr := requireSDKError(t, err, http.StatusConflict) + require.Equal(t, "Chat is in an invalid state.", sdkErr.Message, + "invalid-state endpoint response uses the shared message") +} + +// TestPostChatToolResultsInvalidStateReturnsSharedResponse drives a +// chat into the chatstate-invalid state and asserts that the tool +// results endpoint returns the shared invalid-state response instead +// of the old "Chat is not waiting for tool results." status-conflict +// message. This locks the fix that removes the endpoint fast-path +// and routes invalid chats through the chatstate-backed transaction. +func TestPostChatToolResultsInvalidStateReturnsSharedResponse(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, _, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + // Drive the chat to an invalid combination so the tool-results + // endpoint must surface the shared invalid-state response rather + // than the requires_action status conflict. + driveChatToInvalidWaitingWithQueue(ctx, t, api, chat.ID) + + err = client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + Results: []codersdk.ToolResult{{ + ToolCallID: "call-irrelevant", + Output: json.RawMessage(`{}`), + }}, + }) + sdkErr := requireSDKError(t, err, http.StatusConflict) + require.Equal(t, "Chat is in an invalid state.", sdkErr.Message, + "tool-results invalid-state response uses the shared message") +} + +// TestReconcileInvalidChatStateSucceeds drives a chat into the +// chatstate-invalid combination (waiting with a queued backlog) and +// verifies the reconcile endpoint moves it into a valid error state +// while preserving the queued message. +func TestReconcileInvalidChatStateSucceeds(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db, api := newChatClientWithAPIAndDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + + // Drive the chat to an invalid combination: status=waiting (W), + // archived=false, with a queued message. ClassifyExecutionState + // returns StateInvalid for (waiting, queue=true). + driveChatToInvalidWaitingWithQueue(ctx, t, api, chat.ID) + + reconciled, err := client.ReconcileInvalidChatState(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, chat.ID, reconciled.ID) + require.Equal(t, codersdk.ChatStatusError, reconciled.Status) + + // The persisted row must reflect a valid error state with the + // queued message preserved (E1) and a populated last_error. + persisted, err := loadChatRow(ctx, db, chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusError, persisted.Status) + require.False(t, persisted.Archived) + require.True(t, persisted.LastError.Valid) + + queueCount, err := db.CountChatQueuedMessages(dbauthz.AsChatd(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, int64(1), queueCount, "queued message is preserved by reconcile") +} + +// TestReconcileInvalidChatStateNotInvalidReturnsConflict verifies that +// reconciling a chat that is in a valid execution state is rejected +// with a 409 conflict. +func TestReconcileInvalidChatStateNotInvalidReturnsConflict(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + // A freshly created chat starts in the valid running state (R0). + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{Type: codersdk.ChatInputPartTypeText, Text: "hello"}}, + }) + require.NoError(t, err) + require.Equal(t, codersdk.ChatStatusRunning, chat.Status) + + _, err = client.ReconcileInvalidChatState(ctx, chat.ID) + sdkErr := requireSDKError(t, err, http.StatusConflict) + require.Equal(t, "Chat is not in an invalid state.", sdkErr.Message) +} + +// TestReconcileInvalidChatStateNotFound verifies the reconcile +// endpoint returns 404 for a chat that does not exist. +func TestReconcileInvalidChatStateNotFound(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + + _, err := client.ReconcileInvalidChatState(ctx, uuid.New()) + requireSDKError(t, err, http.StatusNotFound) +} + +// loadChatRow reads a chat row directly through dbauthz.AsChatd so +// endpoint tests verify side effects with the daemon's narrower +// permission set. +func loadChatRow(ctx context.Context, db database.Store, id uuid.UUID) (database.Chat, error) { + chatdCtx := dbauthz.AsChatd(ctx) //nolint:gocritic // Test fixture reads rows with chatd permissions. + return db.GetChatByID(chatdCtx, id) +} + +// driveChatToInvalidWaitingWithQueue forces a chat into the +// chatstate-invalid combination (status=waiting, archived=false, +// queue non-empty) by writing directly through the database. This is +// an intentional invalid fixture: chatstate transitions reject +// driving toward this combination, so AsChatd is not used here. +func driveChatToInvalidWaitingWithQueue( + ctx context.Context, + t *testing.T, + api *coderd.API, + chatID uuid.UUID, +) { + t.Helper() + sysCtx := dbauthz.AsSystemRestricted(ctx) //nolint:gocritic // Test fixture writes invalid combination by design. + + // Seed the queue with one row attributed to the chat owner. The + // content is a minimal valid JSON payload; only the row's + // presence matters for ClassifyExecutionState. The owner_id is + // filled from the chat row by the SQL. + rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued"), + }) + require.NoError(t, err) + _, err = api.Database.InsertChatQueuedMessage(sysCtx, database.InsertChatQueuedMessageParams{ + ChatID: chatID, + Content: rawContent.RawMessage, + ModelConfigID: uuid.NullUUID{}, + }) + require.NoError(t, err) + + // Flip the chat's status to waiting via a raw execution-state + // update. This bypasses the transition matrix to produce the + // (waiting, queued) invalid pairing. + _, err = api.Database.UpdateChatExecutionState(sysCtx, database.UpdateChatExecutionStateParams{ + ID: chatID, + Status: database.ChatStatusWaiting, + Archived: false, + }) + require.NoError(t, err) +} diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index dea596aae8..06a2fbea9b 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -61,6 +61,7 @@ func chatDeploymentValues(t testing.TB) *codersdk.DeploymentValues { t.Helper() values := coderdtest.DeploymentValues(t) + require.NoError(t, values.AI.Chat.AIGatewayRoutingEnabled.Set("false")) return values } @@ -131,6 +132,39 @@ func newChatClientWithAPIAndDatabase(t testing.TB, overrides ...func(*coderdtest return codersdk.NewExperimentalClient(client), api.Database, api } +func currentTestAPIKeyID(t testing.TB, client *codersdk.ExperimentalClient) string { + t.Helper() + + apiKeyID, _, ok := strings.Cut(client.SessionToken(), "-") + require.True(t, ok) + require.NotEmpty(t, apiKeyID) + return apiKeyID +} + +func insertTestChatQueuedMessage( + ctx context.Context, + t testing.TB, + db database.Store, + chatID uuid.UUID, + content json.RawMessage, + modelConfigID uuid.UUID, + apiKeyID string, +) database.ChatQueuedMessage { + t.Helper() + + queued, err := db.InsertChatQueuedMessage( + dbauthz.AsSystemRestricted(ctx), + database.InsertChatQueuedMessageParams{ + ChatID: chatID, + Content: content, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + }, + ) + require.NoError(t, err) + return queued +} + // findUserMessage returns the first user-role message from a slice of chat // messages, failing the test if none is found. func findUserMessage(t testing.TB, messages []database.ChatMessage) database.ChatMessage { @@ -1927,7 +1961,9 @@ func TestWatchChats(t *testing.T) { require.Equal(t, createdChat.OwnerID, got.OwnerID) require.Equal(t, modelConfig.ID, got.LastModelConfigID) require.Equal(t, createdChat.Title, got.Title) - require.Equal(t, codersdk.ChatStatusPending, got.Status) + // CreateChat inserts new chats in the running state under the + // chatstate state machine, so the created event carries running. + require.Equal(t, codersdk.ChatStatusRunning, got.Status) require.NotNil(t, got.RootChatID) require.Equal(t, createdChat.ID, *got.RootChatID) require.NotZero(t, got.CreatedAt) @@ -2040,7 +2076,7 @@ func TestWatchChats(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -2055,6 +2091,11 @@ func TestWatchChats(t *testing.T) { }) require.NoError(t, err) + // The parent chat is created via the API, so the chat worker moves + // it to running. Archiving is only allowed from a terminal state, + // so wait for it to settle before archiving below. + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) + childOne := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, OwnerID: user.UserID, @@ -4658,7 +4699,7 @@ func TestGetChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -4673,6 +4714,11 @@ func TestGetChat(t *testing.T) { }) require.NoError(t, err) + // The parent chat is created via the API, so the chat worker moves + // it to running. Archiving is only allowed from a terminal state, + // so wait for it to settle before archiving below. + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) + child := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, OwnerID: user.UserID, @@ -4728,9 +4774,11 @@ func TestGetChatUserPrompts(t *testing.T) { t.Helper() content, err := chatprompt.MarshalParts(parts) require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: userID}) msgs, err := db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ ChatID: chatID, CreatedBy: []uuid.UUID{userID}, + APIKeyID: []string{apiKey.ID}, ModelConfigID: []uuid.UUID{modelConfigID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.CurrentContentVersion}, @@ -4837,9 +4885,11 @@ func TestGetChatUserPrompts(t *testing.T) { // without the guard, jsonb_array_elements would raise // "cannot extract elements from a scalar" and the request // would 500. + legacyAPIKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.UserID}) _, err = db.InsertChatMessages(dbauthz.AsSystemRestricted(ctx), database.InsertChatMessagesParams{ ChatID: chat.ID, CreatedBy: []uuid.UUID{user.UserID}, + APIKeyID: []string{legacyAPIKey.ID}, ModelConfigID: []uuid.UUID{modelConfig.ID}, Role: []database.ChatMessageRole{database.ChatMessageRoleUser}, ContentVersion: []int16{chatprompt.ContentVersionV0}, @@ -5636,7 +5686,7 @@ func TestArchiveChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) mAudit := audit.NewMock() - client := newChatClient(t, func(o *coderdtest.Options) { + client, api := newChatClientWithAPI(t, func(o *coderdtest.Options) { o.Auditor = mAudit }) firstUser := coderdtest.CreateFirstUser(t, client.Client) @@ -5663,6 +5713,8 @@ func TestArchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, chatToArchive.ID) + coderdtest.WaitForChatSettled(ctx, t, api, chatToKeep.ID) chatsBeforeArchive, err := client.ListChats(ctx, nil) require.NoError(t, err) @@ -5718,7 +5770,7 @@ func TestArchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -5733,6 +5785,7 @@ func TestArchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) // Insert child chats directly via the database. child1 := dbgen.Chat(t, db, database.Chat{ @@ -5806,7 +5859,7 @@ func TestArchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -5821,6 +5874,7 @@ func TestArchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) // Insert a child chat directly via the database. child := dbgen.Chat(t, db, database.Chat{ @@ -5832,41 +5886,18 @@ func TestArchiveChat(t *testing.T) { RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, }) - // Individual child archive is permitted and leaves the - // parent active; the invariant is one-way. + // Archive state changes must target the root chat and cascade. + // Child archive attempts are rejected to preserve the family invariant. err = client.UpdateChat(ctx, child.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) - require.NoError(t, err) + requireSDKError(t, err, http.StatusBadRequest) dbChild, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), child.ID) require.NoError(t, err) - require.True(t, dbChild.Archived, "child should be archived") + require.False(t, dbChild.Archived, "child should remain active") dbParent, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), parentChat.ID) require.NoError(t, err) require.False(t, dbParent.Archived, "parent should stay active") - - // Archived child is hidden under an active parent. - activeChats, err := client.ListChats(ctx, &codersdk.ListChatsOptions{Query: "archived:false"}) - require.NoError(t, err) - var activeParent *codersdk.Chat - for i := range activeChats { - if activeChats[i].ID == parentChat.ID { - activeParent = &activeChats[i] - break - } - } - require.NotNil(t, activeParent, "parent should appear in active list") - for _, c := range activeParent.Children { - require.NotEqual(t, child.ID, c.ID, "archived child must not appear under active parent") - } - - // Nor does the child surface in the archived list (only - // roots paginate there). - archivedChats, err := client.ListChats(ctx, &codersdk.ListChatsOptions{Query: "archived:true"}) - require.NoError(t, err) - for _, c := range archivedChats { - require.NotEqual(t, child.ID, c.ID, "archived child should not surface as a root in archived list") - } }) } @@ -5877,7 +5908,7 @@ func TestUnarchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client := newChatClient(t) + client, api := newChatClientWithAPI(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) @@ -5891,6 +5922,7 @@ func TestUnarchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) // Archive the chat first. err = client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(true)}) @@ -5928,7 +5960,7 @@ func TestUnarchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -5942,6 +5974,7 @@ func TestUnarchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) child1 := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, @@ -6020,7 +6053,7 @@ func TestUnarchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client := newChatClient(t) + client, api := newChatClientWithAPI(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) @@ -6034,6 +6067,7 @@ func TestUnarchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) // Trying to unarchive a non-archived chat should fail. err = client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(false)}) @@ -6044,7 +6078,7 @@ func TestUnarchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -6059,6 +6093,7 @@ func TestUnarchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) // Insert a child directly via the database, then archive the // parent so the whole family is archived (cascade). @@ -6091,7 +6126,7 @@ func TestUnarchiveChat(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, db := newChatClientWithDatabase(t) + client, db, api := newChatClientWithAPIAndDatabase(t) user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) @@ -6105,6 +6140,7 @@ func TestUnarchiveChat(t *testing.T) { }, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, parentChat.ID) // Simulate legacy lone-archived child (from before the // child-archive gate existed) by inserting it directly @@ -6121,15 +6157,14 @@ func TestUnarchiveChat(t *testing.T) { _, err = db.ArchiveChatByID(dbauthz.AsSystemRestricted(ctx), child.ID) require.NoError(t, err) - // Unarchiving the child is permitted because the parent is - // already active; this is the recovery path for legacy - // data. + // Archive state changes must target the root chat, even when + // the child is a legacy lone-archived row. err = client.UpdateChat(ctx, child.ID, codersdk.UpdateChatRequest{Archived: ptr.Ref(false)}) - require.NoError(t, err) + requireSDKError(t, err, http.StatusBadRequest) dbChild, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), child.ID) require.NoError(t, err) - require.False(t, dbChild.Archived, "child should be unarchived") + require.True(t, dbChild.Archived, "child should remain archived") }) t.Run("NotFound", func(t *testing.T) { @@ -6222,12 +6257,14 @@ func TestChatPinOrder(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client := newChatClient(t) + client, api := newChatClientWithAPI(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) first := createChat(ctx, t, client, firstUser.OrganizationID, "pinned then archived") second := createChat(ctx, t, client, firstUser.OrganizationID, "stays pinned") + coderdtest.WaitForChatSettled(ctx, t, api, first.ID) + coderdtest.WaitForChatSettled(ctx, t, api, second.ID) // Pin both. err := client.UpdateChat(ctx, first.ID, codersdk.UpdateChatRequest{PinOrder: ptr.Ref(int32(1))}) @@ -6506,7 +6543,7 @@ func TestPostChatMessages(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client := newChatClient(t) + client, api := newChatClientWithAPI(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) @@ -6518,6 +6555,7 @@ func TestPostChatMessages(t *testing.T) { }}, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) err = client.UpdateChat(ctx, chat.ID, codersdk.UpdateChatRequest{ Archived: ptr.Ref(true), @@ -6815,7 +6853,7 @@ func TestWatchChatsStatusChangeCarriesUpdatedLastModelConfigID(t *testing.T) { _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, - Status: database.ChatStatusWaiting, + Status: database.ChatStatusError, WorkerID: uuid.NullUUID{}, StartedAt: sql.NullTime{}, HeartbeatAt: sql.NullTime{}, @@ -8006,7 +8044,7 @@ func TestPatchChatMessage(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client := newChatClient(t) + client, api := newChatClientWithAPI(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) _ = createChatModelConfig(t, client) @@ -8018,6 +8056,7 @@ func TestPatchChatMessage(t *testing.T) { }}, }) require.NoError(t, err) + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) messagesResult, err := client.GetChatMessages(ctx, chat.ID, nil) require.NoError(t, err) @@ -8290,14 +8329,14 @@ func TestInterruptChat(t *testing.T) { interrupted, err := client.InterruptChat(ctx, chat.ID) require.NoError(t, err) require.Equal(t, chat.ID, interrupted.ID) - require.Equal(t, codersdk.ChatStatusWaiting, interrupted.Status) + require.Equal(t, codersdk.ChatStatusInterrupting, interrupted.Status) persisted, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, persisted.Status) - require.False(t, persisted.WorkerID.Valid) - require.False(t, persisted.StartedAt.Valid) - require.False(t, persisted.HeartbeatAt.Valid) + require.Equal(t, database.ChatStatusInterrupting, persisted.Status) + require.True(t, persisted.WorkerID.Valid) + require.True(t, persisted.StartedAt.Valid) + require.True(t, persisted.HeartbeatAt.Valid) }) t.Run("ChatNotFound", func(t *testing.T) { @@ -8726,6 +8765,7 @@ func TestManualTitleEndpointsPassCallerAPIKeyToAIGateway(t *testing.T) { _ = dbgen.ChatMessage(t, db, database.ChatMessage{ ChatID: chat.ID, CreatedBy: uuid.NullUUID{UUID: firstUser.UserID, Valid: true}, + APIKeyID: sql.NullString{String: wantAPIKeyID, Valid: true}, ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, Role: database.ChatMessageRoleUser, Visibility: database.ChatMessageVisibilityBoth, @@ -8738,6 +8778,58 @@ func TestManualTitleEndpointsPassCallerAPIKeyToAIGateway(t *testing.T) { } } +func TestPostChats_AutomaticTitleGeneration(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + + // titleRequested is signaled when the provider receives the structured + // title-generation request. Automatic title generation issues a + // non-streaming request using the "propose_title" schema, which uniquely + // identifies it (the turn status label uses "propose_turn_status_label"). + titleRequested := make(chan struct{}, 1) + baseURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if req.Stream { + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("Hello from test server.")...) + } + if bytes.Contains(req.RawBody, []byte("propose_title")) { + select { + case titleRequested <- struct{}{}: + default: + } + } + return chattest.OpenAINonStreamingResponse(`{"title": "Generated Title"}`) + }) + + client, api := newChatClientWithAPI(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfigWithBaseURL(t, client, baseURL) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: firstUser.OrganizationID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "automatic title generation please", + }}, + }) + require.NoError(t, err) + // The create response carries the synchronous fallback title derived from + // the message, not the asynchronously generated one. + require.Equal(t, "automatic title generation please", chat.Title) + + // The create endpoint kicks off detached title generation; the provider + // should receive the title request without any further client action. + select { + case <-titleRequested: + case <-ctx.Done(): + t.Fatal("timed out waiting for automatic title generation to be triggered") + } + + // Drain background work so the detached goroutine finishes before the test + // (and its fake provider) tears down. + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) +} + func TestGetChatDiffStatus(t *testing.T) { t.Parallel() @@ -8983,20 +9075,14 @@ func TestDeleteChatQueuedMessage(t *testing.T) { OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "delete queued message route test", + Status: database.ChatStatusError, }) deleteContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued message for delete route"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: deleteContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, deleteContent, modelConfig.ID, currentTestAPIKeyID(t, client)) res, err := client.Request( ctx, @@ -9069,6 +9155,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued message route test", + Status: database.ChatStatusError, }) const queuedText = "queued message for promote route" @@ -9076,14 +9163,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) promoteRes, err := client.Request( ctx, @@ -9139,6 +9219,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { OwnerID: user.UserID, LastModelConfigID: modelConfig.ID, Title: "promote queued usage limit", + Status: database.ChatStatusError, }) const queuedText = "queued message for promote route" @@ -9147,27 +9228,10 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) insertAssistantCostMessage(t, db, chat.ID, modelConfig.ID, 100) - _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) - promoteRes, err := client.Request( ctx, http.MethodPost, @@ -9259,14 +9323,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued message no agents access"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) promoteRes, err := memberClient.Request( ctx, @@ -9298,14 +9355,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) // Archive the chat. _, err = db.ArchiveChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) @@ -9395,14 +9445,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText(queuedText), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) promoteRes, err := client.Request( ctx, @@ -9479,9 +9522,9 @@ func TestPromoteChatQueuedMessage(t *testing.T) { // Simulate an active worker by setting status to running. // We do not start a real worker; the running-case behavior - // (reorder + set waiting + clear worker) does not depend on - // one. The deferred auto-promote is exercised by the - // chatd-package tests where a real worker is involved. + // reorders the queue and moves the chat to interrupting. The + // deferred auto-promote is exercised by chatd-package tests + // where a real worker is involved. _, err = db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ ID: chat.ID, Status: database.ChatStatusRunning, @@ -9495,14 +9538,7 @@ func TestPromoteChatQueuedMessage(t *testing.T) { codersdk.ChatMessageText("running-promote"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, client)) promoteRes, err := client.Request( ctx, @@ -9520,10 +9556,10 @@ func TestPromoteChatQueuedMessage(t *testing.T) { after, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, after.Status, - "running-case promote must transition chat to waiting") - require.False(t, after.WorkerID.Valid, - "running-case promote must clear WorkerID") + require.Equal(t, database.ChatStatusInterrupting, after.Status, + "running-case promote must transition chat to interrupting") + require.True(t, after.WorkerID.Valid, + "running-case promote keeps current worker ownership") queuedRemaining, err := db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID) require.NoError(t, err) @@ -14371,6 +14407,8 @@ func TestGetChatMessages_Pagination(t *testing.T) { t *testing.T, db database.Store, chatID uuid.UUID, + modelConfigID uuid.UUID, + apiKeyID string, ) { t.Helper() @@ -14378,14 +14416,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - _, err = db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chatID, - Content: content, - }, - ) - require.NoError(t, err) + _ = insertTestChatQueuedMessage(ctx, t, db, chatID, content, modelConfigID, apiKeyID) } t.Run("NoCursorReturnsAllDESCPlusQueued", func(t *testing.T) { @@ -14397,7 +14428,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) resp, err := client.GetChatMessages(ctx, chat.ID, nil) require.NoError(t, err) @@ -14422,7 +14453,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ BeforeID: ids[2], @@ -14448,7 +14479,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[1], @@ -14476,7 +14507,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { modelConfig := createChatModelConfig(t, client) chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[0], @@ -14505,7 +14536,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 5) // Seed a queued message so the Empty assertion below verifies // the cursor suppresses queued rows, not just that none exist. - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) resp, err := client.GetChatMessages(ctx, chat.ID, &codersdk.ChatMessagesPaginationOptions{ AfterID: ids[0], @@ -14599,7 +14630,7 @@ func TestGetChatMessages_Pagination(t *testing.T) { chat, ids := seedChat(t, db, user.UserID, user.OrganizationID, modelConfig.ID, 3) // Seed a queued message to prove the cursor path suppresses // it even when nothing else comes back. - seedQueuedMessage(ctx, t, db, chat.ID) + seedQueuedMessage(ctx, t, db, chat.ID, modelConfig.ID, currentTestAPIKeyID(t, client)) // The steady-state polling case: the caller already has every // message, so after_id equals the largest seen id. The server @@ -14836,19 +14867,12 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { t.Run("PromoteChatQueuedMessage", func(t *testing.T) { t.Parallel() - ctx, _, sharedClient, chat, db := setup(t) + ctx, ownerClient, sharedClient, chat, db := setup(t) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) res, err := sharedClient.Request( ctx, @@ -14878,19 +14902,12 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { t.Run("DeleteChatQueuedMessage", func(t *testing.T) { t.Parallel() - ctx, _, sharedClient, chat, db := setup(t) + ctx, ownerClient, sharedClient, chat, db := setup(t) queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) res, err := sharedClient.Request( ctx, @@ -14912,6 +14929,15 @@ func TestChatReadOnlySharedWriteHandlers(t *testing.T) { requireSDKError(t, err, http.StatusNotFound) }) + t.Run("ReconcileInvalidChatState", func(t *testing.T) { + t.Parallel() + + ctx, _, sharedClient, chat, _ := setup(t) + _, err := sharedClient.ReconcileInvalidChatState(ctx, chat.ID) + + requireSDKError(t, err, http.StatusNotFound) + }) + t.Run("RegenerateChatTitle", func(t *testing.T) { t.Parallel() @@ -15026,21 +15052,14 @@ func TestChatOwnerOnlyWriteHandlers(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - _, adminClient, chat, db := setupOrgAdminAndOwnerChat(t) + ownerClient, adminClient, chat, db := setupOrgAdminAndOwnerChat(t) // Insert a queued message directly in the DB. queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued"), }) require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage( - dbauthz.AsSystemRestricted(ctx), - database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }, - ) - require.NoError(t, err) + queuedMessage := insertTestChatQueuedMessage(ctx, t, db, chat.ID, queuedContent, chat.LastModelConfigID, currentTestAPIKeyID(t, ownerClient)) // Org admin tries to promote. promoteRes, err := adminClient.Request( diff --git a/coderd/pubsub/chatstateupdate.go b/coderd/pubsub/chatstateupdate.go new file mode 100644 index 0000000000..b83c2d53c6 --- /dev/null +++ b/coderd/pubsub/chatstateupdate.go @@ -0,0 +1,84 @@ +package pubsub + +import ( + "context" + "encoding/json" + "fmt" + + "github.com/google/uuid" + "golang.org/x/xerrors" +) + +// ChatStateUpdateChannel returns the pubsub channel that receives one +// `chat:update:{chat_id}` message every time the chatstate state +// machine commits a transition for the chat. +func ChatStateUpdateChannel(chatID uuid.UUID) string { + return fmt.Sprintf("chat:update:%s", chatID) +} + +// ChatStateOwnershipChannel is the global pubsub channel that +// receives ownership hints when a chat is runnable but currently has +// missing or stale ownership. Workers listen on this channel to know +// when to attempt acquisition. +const ChatStateOwnershipChannel = "chat:ownership" + +// ChatStateUpdateMessage is the JSON payload published on +// [ChatStateUpdateChannel] after every successful CreateChat or +// ChatMachine.Update commit. It carries the committed post-transition +// versions and ownership identifiers so stream loops and workers can +// decide whether to refetch state. +type ChatStateUpdateMessage struct { + SnapshotVersion int64 `json:"snapshot_version"` + WorkerID *uuid.UUID `json:"worker_id"` + RunnerID *uuid.UUID `json:"runner_id"` + HistoryVersion int64 `json:"history_version"` + QueueVersion int64 `json:"queue_version"` + RetryStateVersion int64 `json:"retry_state_version"` + GenerationAttempt int64 `json:"generation_attempt"` + Status string `json:"status"` + Archived bool `json:"archived"` +} + +// ChatStateOwnershipMessage is the JSON payload published on +// [ChatStateOwnershipChannel] when ownership is missing or stale for +// a runnable chat. Subscribers should reload the chat row to confirm +// ownership before acting. +type ChatStateOwnershipMessage struct { + ChatID uuid.UUID `json:"chat_id"` + SnapshotVersion int64 `json:"snapshot_version"` +} + +// HandleChatStateUpdate wraps a typed callback for +// [ChatStateUpdateMessage] consumption, following the same pattern as +// HandleChatWatchEvent. +func HandleChatStateUpdate(cb func(ctx context.Context, payload ChatStateUpdateMessage, err error)) func(ctx context.Context, message []byte, err error) { + return func(ctx context.Context, message []byte, err error) { + if err != nil { + cb(ctx, ChatStateUpdateMessage{}, xerrors.Errorf("chat state update pubsub: %w", err)) + return + } + var payload ChatStateUpdateMessage + if uerr := json.Unmarshal(message, &payload); uerr != nil { + cb(ctx, ChatStateUpdateMessage{}, xerrors.Errorf("unmarshal chat state update: %w", uerr)) + return + } + cb(ctx, payload, err) + } +} + +// HandleChatStateOwnership wraps a typed callback for +// [ChatStateOwnershipMessage] consumption. +func HandleChatStateOwnership(cb func(ctx context.Context, payload ChatStateOwnershipMessage, err error)) func(ctx context.Context, message []byte, err error) { + return func(ctx context.Context, message []byte, err error) { + if err != nil { + cb(ctx, ChatStateOwnershipMessage{}, xerrors.Errorf("chat state ownership pubsub: %w", err)) + return + } + var payload ChatStateOwnershipMessage + if uerr := json.Unmarshal(message, &payload); uerr != nil { + cb(ctx, ChatStateOwnershipMessage{}, xerrors.Errorf("unmarshal chat state ownership: %w", uerr)) + return + } + cb(ctx, payload, err) + } +} diff --git a/coderd/pubsub/chatstreamnotify.go b/coderd/pubsub/chatstreamnotify.go deleted file mode 100644 index d53605d29c..0000000000 --- a/coderd/pubsub/chatstreamnotify.go +++ /dev/null @@ -1,56 +0,0 @@ -package pubsub - -import ( - "fmt" - - "github.com/google/uuid" - - "github.com/coder/coder/v2/codersdk" -) - -// ChatStreamNotifyChannel returns the pubsub channel for per-chat -// stream notifications. Subscribers receive lightweight notifications -// and read actual content from the database. -func ChatStreamNotifyChannel(chatID uuid.UUID) string { - return fmt.Sprintf("chat:stream:%s", chatID) -} - -// ChatStreamNotifyMessage is the payload published on the per-chat -// stream notification channel. Durable message content is still read -// from the database, while transient control events can be carried -// inline for cross-replica delivery. -type ChatStreamNotifyMessage struct { - // AfterMessageID tells subscribers to query messages after this - // ID. Set when a new message is persisted. - AfterMessageID int64 `json:"after_message_id,omitempty"` - - // Status is set when the chat status changes. Subscribers use - // this to update clients and to manage relay lifecycle. - Status string `json:"status,omitempty"` - - // WorkerID identifies which replica is running the chat. Used - // by enterprise relay to know where to connect. - WorkerID string `json:"worker_id,omitempty"` - - // Retry carries a structured retry event for cross-replica live - // delivery. This is transient stream state and is not read back - // from the database. - Retry *codersdk.ChatStreamRetry `json:"retry,omitempty"` - - // ErrorPayload carries a structured error event for cross-replica - // live delivery. Keep Error for backward compatibility with older - // replicas during rolling deploys. - ErrorPayload *codersdk.ChatError `json:"error_payload,omitempty"` - - // Error is the legacy string-only error payload kept for mixed- - // version compatibility during rollout. - Error string `json:"error,omitempty"` - - // QueueUpdate is set when the queued messages change. - QueueUpdate bool `json:"queue_update,omitempty"` - - // FullRefresh signals that subscribers should re-fetch all - // messages from the beginning (e.g. after an edit that - // truncates message history). - FullRefresh bool `json:"full_refresh,omitempty"` -} diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index 9ea2ef5b5a..0dc91010cc 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -44,6 +44,7 @@ import ( "github.com/coder/coder/v2/coderd/wspubsub" "github.com/coder/coder/v2/coderd/x/chatd" "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/gitsync" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" @@ -2551,10 +2552,11 @@ func (api *API) workspaceAgentAddChatContext(rw http.ResponseWriter, r *http.Req return } - err = api.Database.InTx(func(tx database.Store) error { - locked, err := tx.GetChatByIDForUpdate(sysCtx, chat.ID) + machine := chatstate.NewChatMachine(api.Database, api.Pubsub, chat.ID) + err = machine.Update(sysCtx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(sysCtx, chat.ID) if err != nil { - return xerrors.Errorf("lock chat: %w", err) + return xerrors.Errorf("load chat: %w", err) } if !isActiveAgentChat(locked) { return errChatNotActive @@ -2565,31 +2567,68 @@ func (api *API) workspaceAgentAddChatContext(rw http.ResponseWriter, r *http.Req if locked.OwnerID != workspace.OwnerID { return errChatDoesNotBelongToWorkspaceOwner } - if _, err := tx.InsertChatMessages(sysCtx, chatd.BuildSingleUserChatMessageInsertParams( - chat.ID, - "", // Agent-initiated context injection has no caller API key. - content, - database.ChatMessageVisibilityBoth, - locked.LastModelConfigID, - chatprompt.CurrentContentVersion, - uuid.Nil, - )); err != nil { - return xerrors.Errorf("insert context message: %w", err) + apiKeyID, err := resolveAgentChatContextAPIKeyID(sysCtx, store, locked) + if err != nil { + return err } - if err := updateAgentChatLastInjectedContextFromMessages(sysCtx, api.Logger, tx, chat.ID); err != nil { + sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: locked.LastModelConfigID, Valid: locked.LastModelConfigID != uuid.Nil}, + CreatedBy: uuid.NullUUID{UUID: locked.OwnerID, Valid: locked.OwnerID != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + }, + BusyBehavior: chatstate.BusyBehaviorInterrupt, + }) + if err != nil { + return err + } + if len(sendResult.InsertedMessages) == 0 { + return nil + } + if err := updateAgentChatLastInjectedContextFromMessages(sysCtx, api.Logger, store, chat.ID); err != nil { return xerrors.Errorf("rebuild injected context cache: %w", err) } return nil - }, nil) + }) if err != nil { - if errors.Is(err, errChatNotActive) || errors.Is(err, errChatDoesNotBelongToAgent) || errors.Is(err, errChatDoesNotBelongToWorkspaceOwner) { + switch { + case errors.Is(err, errChatNotActive), errors.Is(err, errChatDoesNotBelongToAgent), errors.Is(err, errChatDoesNotBelongToWorkspaceOwner): writeAgentChatError(ctx, rw, err) - return + case errors.Is(err, errChatAPIKeyAttributionUnavailable): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Cannot modify context: chat has no API key attribution.", + }) + case errors.Is(err, chatstate.ErrMessageQueueFull): + var queueFull *chatstate.MessageQueueFullError + detail := "" + if errors.As(err, &queueFull) { + detail = fmt.Sprintf("Maximum %d messages can be queued.", queueFull.Max) + } + httpapi.Write(ctx, rw, http.StatusTooManyRequests, codersdk.Response{ + Message: "Message queue is full.", + Detail: detail, + }) + case errors.Is(err, chatstate.ErrInvalidState): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is in an invalid state.", + }) + case errors.Is(err, chatstate.ErrTransitionNotAllowed): + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Chat is not in a state that accepts new context.", + Detail: err.Error(), + }) + case errors.Is(err, chatstate.ErrChatNotFound): + writeAgentChatError(ctx, rw, errChatNotFound) + default: + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to persist context message.", + Detail: err.Error(), + }) } - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to persist context message.", - Detail: err.Error(), - }) return } @@ -2658,6 +2697,7 @@ var ( errChatNotActive = xerrors.New("chat is not active") errChatDoesNotBelongToAgent = xerrors.New("chat does not belong to this agent") errChatDoesNotBelongToWorkspaceOwner = xerrors.New("chat does not belong to this workspace owner") + errChatAPIKeyAttributionUnavailable = xerrors.New("chat has no API key attribution") ) type multipleActiveChatsError struct { @@ -2756,6 +2796,56 @@ func isActiveAgentChat(chat database.Chat) bool { } } +func resolveAgentChatContextAPIKeyID(ctx context.Context, db database.Store, chat database.Chat) (string, error) { + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + if err != nil { + return "", xerrors.Errorf("load chat messages for API key attribution: %w", err) + } + for i := len(messages) - 1; i >= 0; i-- { + message := messages[i] + if message.Role != database.ChatMessageRoleUser { + continue + } + if !message.APIKeyID.Valid || message.APIKeyID.String == "" { + continue + } + return message.APIKeyID.String, nil + } + + loginTypes := []database.LoginType{ + database.LoginTypePassword, + database.LoginTypeOIDC, + database.LoginTypeGithub, + database.LoginTypeToken, + database.LoginTypeNone, + } + var newest database.APIKey + hasNewest := false + for _, loginType := range loginTypes { + keys, err := db.GetAPIKeysByUserID(ctx, database.GetAPIKeysByUserIDParams{ + LoginType: loginType, + UserID: chat.OwnerID, + IncludeExpired: false, + }) + if err != nil { + return "", xerrors.Errorf("load owner API keys for attribution: %w", err) + } + for _, key := range keys { + if !hasNewest || key.CreatedAt.After(newest.CreatedAt) { + newest = key + hasNewest = true + } + } + } + if !hasNewest { + return "", errChatAPIKeyAttributionUnavailable + } + return newest.ID, nil +} + func clearAgentChatContext( ctx context.Context, db database.Store, diff --git a/coderd/workspaceagents_chat_context_test.go b/coderd/workspaceagents_chat_context_test.go index 2067fe3ff4..131d24e878 100644 --- a/coderd/workspaceagents_chat_context_test.go +++ b/coderd/workspaceagents_chat_context_test.go @@ -4,6 +4,7 @@ import ( "context" "database/sql" "encoding/json" + "fmt" "net/http" "strings" "testing" @@ -19,8 +20,11 @@ import ( "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/coder/v2/testutil" @@ -70,16 +74,6 @@ func TestAgentChatContext(t *testing.T) { ContextFilePath: "/workspace/AGENTS.md", ContextFileContent: "context from the agent", } - fileAPart := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: "/workspace/file-a.md", - ContextFileContent: "file A context", - } - fileBPart := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: "/workspace/file-b.md", - ContextFileContent: "file B context", - } repoHelperSkillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, SkillName: "repo-helper", @@ -96,14 +90,6 @@ func TestAgentChatContext(t *testing.T) { Type: codersdk.ChatMessagePartTypeContextFile, ContextFilePath: agentInstructionsPart.ContextFilePath, } - cachedFileAPart := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: fileAPart.ContextFilePath, - } - cachedFileBPart := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: fileBPart.ContextFilePath, - } cachedRepoHelperSkillPart := codersdk.ChatMessagePart{ Type: codersdk.ChatMessagePartTypeSkill, SkillName: repoHelperSkillPart.SkillName, @@ -123,14 +109,6 @@ func TestAgentChatContext(t *testing.T) { wantCached: []codersdk.ChatMessagePart{cachedAgentInstructionsPart}, cachedOrdered: true, }, - { - name: "AddSuccessIsAdditive", - steps: []addSuccessStep{{req: agentsdk.AddChatContextRequest{Parts: []codersdk.ChatMessagePart{fileAPart}}, wantCount: 1}, {req: agentsdk.AddChatContextRequest{Parts: []codersdk.ChatMessagePart{fileBPart}}, wantCount: 1}}, - wantStored: [][]codersdk.ChatMessagePart{{fileAPart}, {fileBPart}}, - storedOrdered: false, - wantCached: []codersdk.ChatMessagePart{cachedFileAPart, cachedFileBPart}, - cachedOrdered: false, - }, { name: "AddSuccessWithSkillOnlyPartsGetsSentinel", steps: []addSuccessStep{{req: agentsdk.AddChatContextRequest{Parts: []codersdk.ChatMessagePart{repoHelperSkillPart}}, wantCount: 1}}, @@ -249,6 +227,178 @@ func TestAgentChatContext(t *testing.T) { require.Equal(t, updatedModel.ID, persistedChat.LastModelConfigID) }) + t.Run("AddSuccessUpdatesChatStateVersionsAndPublishes", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + baseDB, pubsub := dbtestutil.NewDB(t) + client := coderdtest.New(t, &coderdtest.Options{ + Database: baseDB, + Pubsub: pubsub, + }) + user := coderdtest.CreateFirstUser(t, client) + workspace := dbfake.WorkspaceBuild(t, baseDB, database.WorkspaceTable{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + }).WithAgent().Do() + agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(workspace.AgentToken)) + model := coderd.InsertAgentChatTestModelConfig(t, baseDB, user.UserID) + chat := createAgentChatContextChat(t, baseDB, user.OrganizationID, user.UserID, model.ID, workspace.Agents[0].ID, t.Name()) + + updateCh := make(chan []byte, 1) + cancelSub, err := pubsub.Subscribe(coderdpubsub.ChatStateUpdateChannel(chat.ID), func(_ context.Context, msg []byte) { + updateCh <- msg + }) + require.NoError(t, err) + defer cancelSub() + + resp, err := agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ + ChatID: chat.ID, + Parts: []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFilePath: "/workspace/instructions.md", + ContextFileContent: "remember this file", + }}, + }) + require.NoError(t, err) + require.Equal(t, chat.ID, resp.ChatID) + require.Equal(t, 1, resp.Count) + + persisted, err := baseDB.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, chat.SnapshotVersion+1, persisted.SnapshotVersion) + require.Equal(t, persisted.SnapshotVersion, persisted.HistoryVersion) + + messages := requireAgentChatContextMessages(ctx, t, baseDB, chat.ID) + require.Len(t, messages, 1) + require.Equal(t, persisted.SnapshotVersion, messages[0].Revision) + + cached := requireAgentChatContextCachedParts(ctx, t, baseDB, chat.ID) + require.Len(t, cached, 1) + require.Equal(t, "/workspace/instructions.md", cached[0].ContextFilePath) + + select { + case raw := <-updateCh: + var update coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(raw, &update)) + require.Equal(t, persisted.SnapshotVersion, update.SnapshotVersion) + require.Equal(t, persisted.HistoryVersion, update.HistoryVersion) + case <-ctx.Done(): + t.Fatal("timed out waiting for chat state update") + } + }) + + t.Run("AddInterruptsAndQueuesWhenChatIsRunning", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + setup := newAgentChatContextTestSetup(t) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat = setAgentChatContextChatStatus(ctx, t, setup.db, chat.ID, database.ChatStatusRunning) + chat = acquireAgentChatContextChat(ctx, t, setup.db, chat.ID) + apiKeyID := currentAgentChatContextAPIKeyID(t, setup.client) + + resp, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ + ChatID: chat.ID, + Parts: []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFilePath: "/workspace/queued.md", + ContextFileContent: "queued context", + }}, + }) + require.NoError(t, err) + require.Equal(t, chat.ID, resp.ChatID) + require.Equal(t, 1, resp.Count) + + require.Empty(t, requireAgentChatContextMessages(ctx, t, setup.db, chat.ID)) + + queued, err := setup.db.GetChatQueuedMessages(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Len(t, queued, 1) + require.Equal(t, setup.user.UserID, queued[0].CreatedBy) + require.True(t, queued[0].ModelConfigID.Valid) + require.Equal(t, model.ID, queued[0].ModelConfigID.UUID) + require.True(t, queued[0].APIKeyID.Valid) + require.Equal(t, apiKeyID, queued[0].APIKeyID.String) + + parts := requireAgentChatContextParts(t, queued[0].Content) + require.Len(t, parts, 1) + require.Equal(t, "/workspace/queued.md", parts[0].ContextFilePath) + require.Equal(t, "queued context", parts[0].ContextFileContent) + require.Equal(t, uuid.NullUUID{UUID: setup.workspace.Agents[0].ID, Valid: true}, parts[0].ContextFileAgentID) + + persisted, err := setup.db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.False(t, persisted.LastInjectedContext.Valid) + require.Equal(t, database.ChatStatusInterrupting, persisted.Status) + require.Equal(t, chat.SnapshotVersion+1, persisted.SnapshotVersion) + require.Equal(t, chat.HistoryVersion, persisted.HistoryVersion) + require.Equal(t, persisted.SnapshotVersion, persisted.QueueVersion) + }) + + t.Run("AddFailsWhenQueueIsFull", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + setup := newAgentChatContextTestSetup(t) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + chat = setAgentChatContextChatStatus(ctx, t, setup.db, chat.ID, database.ChatStatusRunning) + chat = acquireAgentChatContextChat(ctx, t, setup.db, chat.ID) + apiKeyID := currentAgentChatContextAPIKeyID(t, setup.client) + for i := range int(chatstate.MaxQueueSize) { + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(fmt.Sprintf("queued %d", i)), + }) + require.NoError(t, err) + _, err = setup.db.InsertChatQueuedMessageWithCreator( + dbauthz.AsSystemRestricted(ctx), + database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: chat.ID, + Content: content.RawMessage, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKeyID, Valid: true}, + CreatedBy: setup.user.UserID, + }, + ) + require.NoError(t, err) + } + + _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ + ChatID: chat.ID, + Parts: []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFilePath: "/workspace/overflow.md", + ContextFileContent: "overflow context", + }}, + }) + sdkErr := requireSDKError(t, err, http.StatusTooManyRequests) + require.Equal(t, "Message queue is full.", sdkErr.Message) + require.Contains(t, sdkErr.Detail, "Maximum") + }) + + t.Run("AddFailsWhenChatStateIsInvalid", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + setup := newAgentChatContextTestSetup(t) + model := coderd.InsertAgentChatTestModelConfig(t, setup.db, setup.user.UserID) + chat := createAgentChatContextChat(t, setup.db, setup.user.OrganizationID, setup.user.UserID, model.ID, setup.workspace.Agents[0].ID, t.Name()) + _ = setAgentChatContextChatStatus(ctx, t, setup.db, chat.ID, database.ChatStatusPending) + + _, err := setup.agentClient.AddChatContext(ctx, agentsdk.AddChatContextRequest{ + ChatID: chat.ID, + Parts: []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFilePath: "/workspace/invalid.md", + ContextFileContent: "invalid state context", + }}, + }) + sdkErr := requireSDKError(t, err, http.StatusConflict) + require.Equal(t, "Chat is in an invalid state.", sdkErr.Message) + }) + t.Run("ClearDeletesSkillMessages", func(t *testing.T) { t.Parallel() @@ -986,6 +1136,45 @@ func newAgentChatContextTestSetup(t *testing.T) agentChatContextTestSetup { } } +func currentAgentChatContextAPIKeyID(t testing.TB, client *codersdk.Client) string { + t.Helper() + + apiKeyID, _, ok := strings.Cut(client.SessionToken(), "-") + require.True(t, ok) + require.NotEmpty(t, apiKeyID) + return apiKeyID +} + +func setAgentChatContextChatStatus( + ctx context.Context, + t testing.TB, + db database.Store, + chatID uuid.UUID, + status database.ChatStatus, +) database.Chat { + t.Helper() + + chat, err := db.UpdateChatStatus(dbauthz.AsSystemRestricted(ctx), database.UpdateChatStatusParams{ + ID: chatID, + Status: status, + }) + require.NoError(t, err) + return chat +} + +func acquireAgentChatContextChat(ctx context.Context, t testing.TB, db database.Store, chatID uuid.UUID) database.Chat { + t.Helper() + + machine := chatstate.NewChatMachine(db, dbpubsub.NewInMemory(), chatID) + require.NoError(t, machine.Update(dbauthz.AsSystemRestricted(ctx), func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: uuid.New(), RunnerID: uuid.New()}) + return err + })) + chat, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chatID) + require.NoError(t, err) + return chat +} + func createAgentChatContextChat( t testing.TB, db database.Store, diff --git a/coderd/x/chatd/active_turn_debug.go b/coderd/x/chatd/active_turn_debug.go new file mode 100644 index 0000000000..2fbd653a04 --- /dev/null +++ b/coderd/x/chatd/active_turn_debug.go @@ -0,0 +1,201 @@ +package chatd + +import ( + "context" + "sync" + + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" +) + +type runnerDebugTurn struct { + runnerCtx context.Context + logger slog.Logger + + mu sync.Mutex + + runContext chatdebug.RunContext + seedSummary map[string]any + service *chatdebug.Service + + created bool + disabled bool + finalized bool + + status chatdebug.Status + statusSet bool + + heartbeatDone chan struct{} +} + +func newRunnerDebugTurn(runnerCtx context.Context, logger slog.Logger) *runnerDebugTurn { + return &runnerDebugTurn{ + runnerCtx: runnerCtx, + logger: logger, + } +} + +func (d *runnerDebugTurn) Ensure( + ctx context.Context, + chat database.Chat, + debug *generationDebug, +) context.Context { + if d == nil { + return ctx + } + + d.mu.Lock() + defer d.mu.Unlock() + + // Check finalized/disabled before created: once the turn is + // finalized, new contexts must not be attributed to the + // finalized run, even if it was created earlier. + if d.disabled || d.finalized { + return ctx + } + if d.created { + return d.contextLocked(ctx) + } + if debug == nil || !debug.Enabled || debug.Service == nil || + chat.ID == uuid.Nil || debug.TriggerMessageID == 0 { + d.disabled = true + return ctx + } + + seedSummary := chatdebug.SeedSummary( + chatdebug.TruncateLabel(debug.TriggerLabel, chatdebug.MaxLabelLength), + ) + rootChatID := uuid.Nil + if chat.RootChatID.Valid { + rootChatID = chat.RootChatID.UUID + } + parentChatID := uuid.Nil + if chat.ParentChatID.Valid { + parentChatID = chat.ParentChatID.UUID + } + + createRunCtx, createRunCancel := context.WithTimeout( + context.WithoutCancel(ctx), debugCreateRunTimeout, + ) + run, createRunErr := debug.Service.CreateRun(createRunCtx, chatdebug.CreateRunParams{ + ChatID: chat.ID, + RootChatID: rootChatID, + ParentChatID: parentChatID, + ModelConfigID: debug.ModelConfig.ID, + TriggerMessageID: debug.TriggerMessageID, + HistoryTipMessageID: debug.HistoryTipMessageID, + Kind: chatdebug.KindChatTurn, + Status: chatdebug.StatusInProgress, + Provider: debug.Provider, + Model: debug.Model, + Summary: seedSummary, + }) + createRunCancel() + if createRunErr != nil { + d.disabled = true + d.logger.Warn(ctx, "failed to create chat debug run", + slog.F("chat_id", chat.ID), + slog.Error(createRunErr), + ) + return ctx + } + + d.service = debug.Service + d.runContext = chatdebugRunContext(run) + d.seedSummary = seedSummary + d.created = true + d.heartbeatDone = make(chan struct{}) + d.service.LaunchRunHeartbeat(d.runnerCtx, d.runContext.RunID, d.runContext.ChatID, d.heartbeatDone) + return d.contextLocked(ctx) +} + +func (d *runnerDebugTurn) Context(ctx context.Context) context.Context { + if d == nil { + return ctx + } + d.mu.Lock() + defer d.mu.Unlock() + return d.contextLocked(ctx) +} + +func (d *runnerDebugTurn) contextLocked(ctx context.Context) context.Context { + if !d.created || d.runContext.RunID == uuid.Nil { + return ctx + } + runContext := d.runContext + return chatdebug.ContextWithRun(ctx, &runContext) +} + +func (d *runnerDebugTurn) RecordOutcome(status chatdebug.Status) { + if d == nil || debugTurnOutcomePriority(status) == 0 { + return + } + d.mu.Lock() + defer d.mu.Unlock() + if d.finalized { + return + } + if !d.statusSet || debugTurnOutcomePriority(status) > debugTurnOutcomePriority(d.status) { + d.status = status + d.statusSet = true + } +} + +func (d *runnerDebugTurn) Finalize(ctx context.Context) { + if d == nil { + return + } + + d.mu.Lock() + if d.finalized { + d.mu.Unlock() + return + } + d.finalized = true + if d.heartbeatDone != nil { + close(d.heartbeatDone) + d.heartbeatDone = nil + } + if !d.created || d.service == nil || d.runContext.RunID == uuid.Nil { + d.mu.Unlock() + return + } + service := d.service + runContext := d.runContext + seedSummary := d.seedSummary + status := chatdebug.StatusInterrupted + if d.statusSet { + status = d.status + } + logger := d.logger + d.mu.Unlock() + + if finalizeErr := service.FinalizeRun(ctx, chatdebug.FinalizeRunParams{ + RunID: runContext.RunID, + ChatID: runContext.ChatID, + Status: status, + SeedSummary: seedSummary, + }); finalizeErr != nil { + logger.Warn(ctx, "failed to finalize chat debug run", + slog.F("chat_id", runContext.ChatID), + slog.F("run_id", runContext.RunID), + slog.Error(finalizeErr), + ) + } +} + +func debugTurnOutcomePriority(status chatdebug.Status) int { + switch status { + case chatdebug.StatusCompleted: + return 1 + case chatdebug.StatusInterrupted: + return 2 + case chatdebug.StatusError: + return 3 + default: + return 0 + } +} diff --git a/coderd/x/chatd/active_turn_debug_internal_test.go b/coderd/x/chatd/active_turn_debug_internal_test.go new file mode 100644 index 0000000000..f599021853 --- /dev/null +++ b/coderd/x/chatd/active_turn_debug_internal_test.go @@ -0,0 +1,155 @@ +package chatd + +import ( + "context" + "database/sql" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" + "github.com/coder/coder/v2/testutil" +) + +func TestRunnerDebugTurnEnsureCreatesOnce(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + runnerCtx, cancel := context.WithCancel(ctx) + defer cancel() + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + chatID := uuid.New() + runID := uuid.New() + modelConfigID := uuid.New() + svc := chatdebug.NewService(db, testutil.Logger(t), nil) + turn := newRunnerDebugTurn(runnerCtx, testutil.Logger(t)) + + db.EXPECT().InsertChatDebugRun(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, params database.InsertChatDebugRunParams) (database.ChatDebugRun, error) { + require.Equal(t, chatID, params.ChatID) + require.Equal(t, string(chatdebug.KindChatTurn), params.Kind) + require.Equal(t, string(chatdebug.StatusInProgress), params.Status) + require.Equal(t, sql.NullInt64{Int64: 123, Valid: true}, params.TriggerMessageID) + return database.ChatDebugRun{ + ID: runID, + ChatID: chatID, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + TriggerMessageID: sql.NullInt64{Int64: 123, Valid: true}, + HistoryTipMessageID: sql.NullInt64{Int64: 456, Valid: true}, + Kind: string(chatdebug.KindChatTurn), + Status: string(chatdebug.StatusInProgress), + Provider: sql.NullString{String: "anthropic", Valid: true}, + Model: sql.NullString{String: "claude", Valid: true}, + }, nil + }).Times(1) + + debug := &generationDebug{ + Enabled: true, + Service: svc, + Provider: "anthropic", + Model: "claude", + TriggerMessageID: 123, + HistoryTipMessageID: 456, + TriggerLabel: "hello", + ModelConfig: database.ChatModelConfig{ID: modelConfigID}, + } + chat := database.Chat{ID: chatID} + + firstCtx := turn.Ensure(ctx, chat, debug) + firstRun, ok := chatdebug.RunFromContext(firstCtx) + require.True(t, ok) + require.Equal(t, runID, firstRun.RunID) + + secondCtx := turn.Ensure(ctx, chat, debug) + secondRun, ok := chatdebug.RunFromContext(secondCtx) + require.True(t, ok) + require.Equal(t, runID, secondRun.RunID) +} + +func TestRunnerDebugTurnEnsureDisabledFirstAttemptStaysDisabled(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + svc := chatdebug.NewService(db, testutil.Logger(t), nil) + turn := newRunnerDebugTurn(ctx, testutil.Logger(t)) + chat := database.Chat{ID: uuid.New()} + + firstCtx := turn.Ensure(ctx, chat, nil) + _, ok := chatdebug.RunFromContext(firstCtx) + require.False(t, ok) + + secondCtx := turn.Ensure(ctx, chat, &generationDebug{ + Enabled: true, + Service: svc, + TriggerMessageID: 1, + ModelConfig: database.ChatModelConfig{ID: uuid.New()}, + }) + _, ok = chatdebug.RunFromContext(secondCtx) + require.False(t, ok) +} + +func TestRunnerDebugTurnRecordOutcomePrecedence(t *testing.T) { + t.Parallel() + + turn := newRunnerDebugTurn(context.Background(), testutil.Logger(t)) + turn.RecordOutcome(chatdebug.StatusCompleted) + require.True(t, turn.statusSet) + require.Equal(t, chatdebug.StatusCompleted, turn.status) + + turn.RecordOutcome(chatdebug.StatusInterrupted) + require.Equal(t, chatdebug.StatusInterrupted, turn.status) + + turn.RecordOutcome(chatdebug.StatusCompleted) + require.Equal(t, chatdebug.StatusInterrupted, turn.status) + + turn.RecordOutcome(chatdebug.StatusError) + require.Equal(t, chatdebug.StatusError, turn.status) +} + +func TestRunnerDebugTurnFinalizeOnce(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + runnerCtx, cancel := context.WithCancel(ctx) + defer cancel() + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + chatID := uuid.New() + runID := uuid.New() + svc := chatdebug.NewService(db, testutil.Logger(t), nil) + turn := newRunnerDebugTurn(runnerCtx, testutil.Logger(t)) + + db.EXPECT().InsertChatDebugRun(gomock.Any(), gomock.Any()). + Return(database.ChatDebugRun{ + ID: runID, + ChatID: chatID, + Kind: string(chatdebug.KindChatTurn), + Status: string(chatdebug.StatusInProgress), + }, nil). + Times(1) + db.EXPECT().GetChatDebugStepsByRunID(gomock.Any(), runID).Return(nil, nil).Times(1) + db.EXPECT().UpdateChatDebugRun(gomock.Any(), gomock.Any()). + DoAndReturn(func(_ context.Context, params database.UpdateChatDebugRunParams) (database.ChatDebugRun, error) { + require.Equal(t, runID, params.ID) + require.Equal(t, chatID, params.ChatID) + require.Equal(t, sql.NullString{String: string(chatdebug.StatusError), Valid: true}, params.Status) + return database.ChatDebugRun{ID: runID, ChatID: chatID}, nil + }).Times(1) + + turn.Ensure(ctx, database.Chat{ID: chatID}, &generationDebug{ + Enabled: true, + Service: svc, + TriggerMessageID: 1, + ModelConfig: database.ChatModelConfig{ID: uuid.New()}, + }) + turn.RecordOutcome(chatdebug.StatusError) + turn.Finalize(ctx) + turn.Finalize(ctx) +} diff --git a/coderd/x/chatd/attempt.go b/coderd/x/chatd/attempt.go new file mode 100644 index 0000000000..0b803e8586 --- /dev/null +++ b/coderd/x/chatd/attempt.go @@ -0,0 +1,64 @@ +package chatd + +import ( + "database/sql" + "time" + + "charm.land/fantasy" + + "github.com/coder/coder/v2/codersdk" +) + +type runnerActionKind string + +type runnerActionMessage struct { + ID int64 + Role codersdk.ChatMessageRole +} + +const ( + runnerActionKindEnterRequiresAction runnerActionKind = "enter_requires_action" + runnerActionKindFinishTurn runnerActionKind = "finish_turn" + runnerActionKindFinishError runnerActionKind = "finish_error" + runnerActionKindFinishInterruption runnerActionKind = "finish_interruption" +) + +// stepData is the durable content produced by one provider attempt. +type stepData struct { + Content []fantasy.Content + Usage fantasy.Usage + ContextLimit sql.NullInt64 + ProviderResponseID string + Runtime time.Duration + + ToolCallCreatedAt map[string]time.Time + ToolResultCreatedAt map[string]time.Time + ReasoningStartedAt []time.Time + ReasoningCompletedAt []time.Time +} + +// pendingDynamicToolCall describes a dynamic tool call parked for a user. +type pendingDynamicToolCall struct { + ToolCallID string + ToolName string + Args string +} + +// compactionOutcome contains a generated context summary. +type compactionOutcome struct { + SystemSummary string + SummaryReport string + ThresholdPercent int32 + UsagePercent float64 + ContextTokens int64 + ContextLimit int64 +} + +type compactionStatus int + +const ( + compactionStatusNotNeeded compactionStatus = iota + compactionStatusNeeded + compactionStatusAfterCompaction + compactionStatusStillOverLimit +) diff --git a/coderd/x/chatd/auto_archive.go b/coderd/x/chatd/auto_archive.go new file mode 100644 index 0000000000..e045447632 --- /dev/null +++ b/coderd/x/chatd/auto_archive.go @@ -0,0 +1,315 @@ +package chatd + +import ( + "cmp" + "context" + "database/sql" + "errors" + "net/http" + "slices" + "strconv" + "time" + + "github.com/dustin/go-humanize" + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/audit" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" +) + +const chatAutoArchiveDigestMaxChats = 25 + +type autoArchivedChat struct { + Chat database.Chat + LastActivityAt time.Time +} + +func (w *chatWorker) archiveLoop(ctx context.Context) { + ticker := w.opts.Clock.NewTicker(w.opts.ArchiveInterval, "chatworker", "auto-archive") + defer ticker.Stop() + w.archiveOnce(ctx, dbtime.Time(w.opts.Clock.Now("chatworker", "auto-archive")).UTC()) + for { + select { + case tick := <-ticker.C: + w.archiveOnce(ctx, dbtime.Time(tick).UTC()) + case <-ctx.Done(): + return + } + } +} + +func (w *chatWorker) archiveOnce(ctx context.Context, start time.Time) { + autoArchiveDays, err := w.opts.Store.GetChatAutoArchiveDays(ctx, codersdk.DefaultChatAutoArchiveDays) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker auto-archive config read failed", slogError(err)) + } + return + } + if autoArchiveDays <= 0 { + return + } + retentionDays, err := w.opts.Store.GetChatRetentionDays(ctx) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker chat retention config read failed", slogError(err)) + } + return + } + + archiveCutoff := dbtime.StartOfDay(start).Add(-time.Duration(autoArchiveDays) * 24 * time.Hour) + rows, err := w.opts.Store.GetAutoArchiveInactiveChatCandidates(ctx, database.GetAutoArchiveInactiveChatCandidatesParams{ + ArchiveCutoff: archiveCutoff, + LimitCount: w.opts.ArchiveBatchSize, + }) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker auto-archive query failed", slogError(err)) + } + return + } + if len(rows) == 0 { + return + } + + archived := make([]autoArchivedChat, 0, len(rows)) + for _, row := range rows { + family, err := w.archiveCandidateSafely(ctx, row) + if err != nil { + if ctx.Err() != nil { + return + } + if isExpectedAutoArchiveError(err) { + w.opts.Logger.Debug(ctx, "chatworker auto-archive skipped chat", + slog.F("chat_id", row.ID), + slog.Error(err), + ) + continue + } + w.opts.Logger.Warn(ctx, "chatworker auto-archive candidate failed", + slog.F("chat_id", row.ID), + slog.Error(err), + ) + continue + } + archived = append(archived, family...) + } + if len(archived) == 0 { + return + } + if w.opts.AutoArchiveRecords != nil { + w.opts.AutoArchiveRecords.Add(float64(len(archived))) + } + w.dispatchChatAutoArchive(context.WithoutCancel(ctx), ctx, start, autoArchiveDays, retentionDays, archived) +} + +func (w *chatWorker) archiveCandidateSafely( + ctx context.Context, + row database.GetAutoArchiveInactiveChatCandidatesRow, +) (family []autoArchivedChat, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = xerrors.Errorf("chatworker auto-archive panic: %v", recovered) + } + }() + return w.archiveCandidate(ctx, row) +} + +func (w *chatWorker) archiveCandidate( + ctx context.Context, + row database.GetAutoArchiveInactiveChatCandidatesRow, +) ([]autoArchivedChat, error) { + familyChats, err := chatstate.SetFamilyArchived(ctx, w.opts.Store, w.opts.Pubsub, chatstate.SetFamilyArchivedInput{ + RootID: row.ID, + Archived: true, + }) + if err != nil { + return nil, err + } + if len(familyChats) == 0 { + return nil, nil + } + w.scheduleArchiveDebugCleanup(ctx, familyChats) + w.publishArchiveWatchEvents(familyChats) + + archived := make([]autoArchivedChat, 0, len(familyChats)) + for _, chat := range familyChats { + lastActivityAt := row.LastActivityAt + if lastActivityAt.IsZero() { + lastActivityAt = chat.CreatedAt + } + archived = append(archived, autoArchivedChat{ + Chat: chat, + LastActivityAt: lastActivityAt, + }) + } + return archived, nil +} + +func isExpectedAutoArchiveError(err error) bool { + return errors.Is(err, sql.ErrNoRows) || + errors.Is(err, chatstate.ErrChatNotFound) || + errors.Is(err, chatstate.ErrChatNotRoot) || + errors.Is(err, chatstate.ErrInvalidState) || + errors.Is(err, chatstate.ErrTransitionNotAllowed) +} + +func (w *chatWorker) publishArchiveWatchEvents(familyChats []database.Chat) { + if w.server != nil { + w.server.publishChatPubsubEvents(familyChats, codersdk.ChatWatchEventKindDeleted) + return + } + for _, chat := range familyChats { + if err := publishChatWatchEvent(w.opts.Pubsub, chat, codersdk.ChatWatchEventKindDeleted); err != nil { + w.opts.Logger.Warn(context.Background(), "chatworker auto-archive watch publish failed", + slog.F("chat_id", chat.ID), + slog.Error(err), + ) + } + } +} + +func (w *chatWorker) scheduleArchiveDebugCleanup(ctx context.Context, familyChats []database.Chat) { + if w.server == nil || len(familyChats) == 0 { + return + } + w.server.scheduleArchiveDebugCleanup(ctx, familyChats) +} + +func (p *Server) scheduleArchiveDebugCleanup(ctx context.Context, familyChats []database.Chat) { + if len(familyChats) == 0 { + return + } + archiveCutoff := familyChats[0].UpdatedAt.Add(-debugCleanupClockSkew) + for _, archivedChat := range familyChats { + p.scheduleDebugCleanup( + ctx, + "failed to delete chat debug rows after archive", + []slog.Field{slog.F("chat_id", archivedChat.ID)}, + func(cleanupCtx context.Context, debugSvc *chatdebug.Service) error { + _, err := debugSvc.DeleteByChatID(cleanupCtx, archivedChat.ID, archiveCutoff) + return err + }, + ) + } +} + +func (w *chatWorker) dispatchChatAutoArchive( + auditCtx context.Context, + enqueueCtx context.Context, + tickStart time.Time, + autoArchiveDays int32, + retentionDays int32, + archived []autoArchivedChat, +) { + roots := make([]autoArchivedChat, 0, len(archived)) + for _, record := range archived { + if !record.Chat.ParentChatID.Valid { + roots = append(roots, record) + } + } + w.auditAutoArchivedChats(auditCtx, roots) + w.enqueueAutoArchiveDigests(enqueueCtx, tickStart, autoArchiveDays, retentionDays, roots) +} + +func (w *chatWorker) auditAutoArchivedChats(ctx context.Context, roots []autoArchivedChat) { + if w.opts.Auditor == nil { + return + } + auditor := w.opts.Auditor.Load() + if auditor == nil { + return + } + for _, record := range roots { + after := record.Chat + before := after + before.Archived = false + audit.BackgroundAudit(ctx, &audit.BackgroundAuditParams[database.Chat]{ + Audit: *auditor, + Log: w.opts.Logger, + UserID: after.OwnerID, + OrganizationID: after.OrganizationID, + Action: database.AuditActionWrite, + Old: before, + New: after, + Status: http.StatusOK, + AdditionalFields: audit.BackgroundTaskFieldsBytes(ctx, w.opts.Logger, audit.BackgroundSubsystemChatAutoArchive), + }) + } +} + +func (w *chatWorker) enqueueAutoArchiveDigests( + ctx context.Context, + tickStart time.Time, + autoArchiveDays int32, + retentionDays int32, + roots []autoArchivedChat, +) { + rootsByOwner := make(map[uuid.UUID][]autoArchivedChat, len(roots)) + for _, record := range roots { + rootsByOwner[record.Chat.OwnerID] = append(rootsByOwner[record.Chat.OwnerID], record) + } + ownerIDs := make([]uuid.UUID, 0, len(rootsByOwner)) + for id := range rootsByOwner { + ownerIDs = append(ownerIDs, id) + } + slices.SortFunc(ownerIDs, func(a, b uuid.UUID) int { + return cmp.Compare(a.String(), b.String()) + }) + for i, ownerID := range ownerIDs { + if err := ctx.Err(); err != nil { + w.opts.Logger.Warn(ctx, "chat auto-archive digest dispatch canceled", + slog.F("remaining_owners", len(ownerIDs)-i), + slog.Error(err), + ) + return + } + data := buildAutoArchiveDigestData(rootsByOwner[ownerID], autoArchiveDays, retentionDays, tickStart) + //nolint:gocritic // Background digest dispatch runs as the notifier subject. + if _, err := w.opts.NotificationsEnqueuer.EnqueueWithData( + dbauthz.AsNotifier(ctx), + ownerID, + notifications.TemplateChatAutoArchiveDigest, + map[string]string{}, + data, + string(audit.BackgroundSubsystemChatAutoArchive), + ); err != nil { + w.opts.Logger.Warn(ctx, "failed to enqueue chat auto-archive digest", + slog.F("owner_id", ownerID), + slog.Error(err), + ) + } + } +} + +func buildAutoArchiveDigestData(rows []autoArchivedChat, autoArchiveDays, retentionDays int32, tickStart time.Time) map[string]any { + overflow := 0 + if len(rows) > chatAutoArchiveDigestMaxChats { + overflow = len(rows) - chatAutoArchiveDigestMaxChats + rows = rows[:chatAutoArchiveDigestMaxChats] + } + chats := make([]map[string]any, 0, len(rows)) + for _, r := range rows { + chats = append(chats, map[string]any{ + "title": r.Chat.Title, + "last_activity_humanized": humanize.RelTime(r.LastActivityAt, tickStart, "ago", "from now"), + }) + } + data := map[string]any{ + "auto_archive_days": strconv.Itoa(int(autoArchiveDays)), + "retention_days": strconv.Itoa(int(retentionDays)), + "archived_chats": chats, + } + if overflow > 0 { + data["additional_archived_count"] = strconv.Itoa(overflow) + } + return data +} diff --git a/coderd/x/chatd/auto_archive_internal_test.go b/coderd/x/chatd/auto_archive_internal_test.go new file mode 100644 index 0000000000..8c2e68b924 --- /dev/null +++ b/coderd/x/chatd/auto_archive_internal_test.go @@ -0,0 +1,819 @@ +package chatd + +import ( + "context" + "database/sql" + "fmt" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + promtestutil "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/audit" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/notifications/notificationstest" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestWorker_AutoArchiveDisabled(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, codersdk.DefaultChatAutoArchiveDays)) + + pubsub := newRecordingPubsub(f.pubsub) + worker := f.newArchiveWorker(t, pubsub, nil, nil) + worker.archiveOnce(ctx, now) + + refreshed, err := f.db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.False(t, refreshed.Archived) + require.Empty(t, pubsub.watchEvents(t)) + require.Empty(t, pubsub.stateUpdateMessages(t, chat.ID)) +} + +func TestWorker_AutoArchivesInactiveRoot(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + insertArchiveMessage(t, f, chat.ID, now.Add(-100*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + require.NoError(t, f.db.UpsertChatRetentionDays(ctx, 30)) + + pubsub := newRecordingPubsub(f.pubsub) + auditor := audit.NewMock() + enqueuer := notificationstest.NewFakeEnqueuer() + worker := f.newArchiveWorker(t, pubsub, mockAuditorPtr(auditor), enqueuer) + worker.archiveOnce(ctx, now) + + refreshed, err := f.db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, refreshed.Archived) + require.Greater(t, refreshed.SnapshotVersion, chat.SnapshotVersion) + + updates := pubsub.stateUpdateMessages(t, chat.ID) + require.NotEmpty(t, updates) + require.True(t, updates[len(updates)-1].Archived) + requireWatchEvent(t, pubsub, chat.ID, codersdk.ChatWatchEventKindDeleted) + + logs := auditor.AuditLogs() + require.Len(t, logs, 1) + require.Equal(t, chat.ID, logs[0].ResourceID) + require.Equal(t, database.ResourceTypeChat, logs[0].ResourceType) + require.Equal(t, database.AuditActionWrite, logs[0].Action) + require.Contains(t, string(logs[0].AdditionalFields), string(audit.BackgroundSubsystemChatAutoArchive)) + + sent := enqueuer.Sent() + require.Len(t, sent, 1) + require.Equal(t, notifications.TemplateChatAutoArchiveDigest, sent[0].TemplateID) + require.Equal(t, f.user.ID, sent[0].UserID) + require.Equal(t, "90", sent[0].Data["auto_archive_days"]) + require.Equal(t, "30", sent[0].Data["retention_days"]) +} + +func TestWorker_AutoArchiveRejectsActiveChild(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + root := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + child := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + f.linkChild(t, root.ID, child.ID) + forceExecutionState(t, f, child.ID, database.ChatStatusRunning, false) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + pubsub := newRecordingPubsub(f.pubsub) + worker := f.newArchiveWorker(t, pubsub, nil, nil) + worker.archiveOnce(ctx, now) + + refreshedRoot, err := f.db.GetChatByID(ctx, root.ID) + require.NoError(t, err) + require.False(t, refreshedRoot.Archived) + refreshedChild, err := f.db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.False(t, refreshedChild.Archived) + require.Empty(t, pubsub.watchEvents(t)) +} + +func TestWorker_AutoArchivePublishesStateUpdatesForFamily(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + root := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + child := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + f.linkChild(t, root.ID, child.ID) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + pubsub := newRecordingPubsub(f.pubsub) + worker := f.newArchiveWorker(t, pubsub, nil, nil) + worker.archiveOnce(ctx, now) + + refreshedRoot, err := f.db.GetChatByID(ctx, root.ID) + require.NoError(t, err) + require.True(t, refreshedRoot.Archived) + refreshedChild, err := f.db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.True(t, refreshedChild.Archived) + require.NotEmpty(t, pubsub.stateUpdateMessages(t, root.ID)) + require.NotEmpty(t, pubsub.stateUpdateMessages(t, child.ID)) + requireWatchEvent(t, pubsub, root.ID, codersdk.ChatWatchEventKindDeleted) + requireWatchEvent(t, pubsub, child.ID, codersdk.ChatWatchEventKindDeleted) +} + +func TestWorker_AutoArchiveExpectedTransitionFailureDoesNotAbortTick(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + blockedRoot := f.createArchiveCandidate(t, now.Add(-130*24*time.Hour)) + blockedChild := f.createArchiveCandidate(t, now.Add(-130*24*time.Hour)) + f.linkChild(t, blockedRoot.ID, blockedChild.ID) + forceExecutionState(t, f, blockedChild.ID, database.ChatStatusRunning, false) + valid := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + pubsub := newRecordingPubsub(f.pubsub) + worker := f.newArchiveWorker(t, pubsub, nil, nil) + worker.archiveOnce(ctx, now) + + blockedAfter, err := f.db.GetChatByID(ctx, blockedRoot.ID) + require.NoError(t, err) + require.False(t, blockedAfter.Archived) + validAfter, err := f.db.GetChatByID(ctx, valid.ID) + require.NoError(t, err) + require.True(t, validAfter.Archived) + requireWatchEvent(t, pubsub, valid.ID, codersdk.ChatWatchEventKindDeleted) +} + +func TestWorker_AutoArchiveDateBoundary(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + onCutoff := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + insertArchiveMessage(t, f, onCutoff.ID, time.Date(2026, 2, 28, 23, 59, 59, 0, time.UTC)) + beforeCutoff := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + insertArchiveMessage(t, f, beforeCutoff.ID, time.Date(2026, 2, 27, 23, 59, 59, 0, time.UTC)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + worker.archiveOnce(ctx, now) + + refreshedOn, err := f.db.GetChatByID(ctx, onCutoff.ID) + require.NoError(t, err) + require.False(t, refreshedOn.Archived) + refreshedBefore, err := f.db.GetChatByID(ctx, beforeCutoff.ID) + require.NoError(t, err) + require.True(t, refreshedBefore.Archived) +} + +func (f *workerTestFixture) createArchiveCandidate(t *testing.T, createdAt time.Time) database.Chat { + t.Helper() + return f.createArchiveCandidateForOwner(t, f.user.ID, createdAt) +} + +func (f *workerTestFixture) createArchiveCandidateForOwner(t *testing.T, ownerID uuid.UUID, createdAt time.Time) database.Chat { + t.Helper() + chat := dbgen.Chat(t, f.db, database.Chat{ + OrganizationID: f.org.ID, + OwnerID: ownerID, + LastModelConfigID: f.model.ID, + Title: testutil.GetRandomName(t), + Status: database.ChatStatusWaiting, + }) + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitShort), "UPDATE chats SET created_at = $1, updated_at = $1 WHERE id = $2", createdAt, chat.ID) + require.NoError(t, err) + chat.CreatedAt = createdAt + chat.UpdatedAt = createdAt + return chat +} + +func (f *workerTestFixture) setPinOrder(t *testing.T, chatID uuid.UUID, order int32) { + t.Helper() + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitShort), "UPDATE chats SET pin_order = $1 WHERE id = $2", order, chatID) + require.NoError(t, err) +} + +func (f *workerTestFixture) softDeleteMessages(t *testing.T, chatID uuid.UUID) { + t.Helper() + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitShort), "UPDATE chat_messages SET deleted = true WHERE chat_id = $1", chatID) + require.NoError(t, err) +} + +func (f *workerTestFixture) archived(t *testing.T, chatID uuid.UUID) bool { + t.Helper() + chat, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chatID) + require.NoError(t, err) + return chat.Archived +} + +func (f *workerTestFixture) linkChild(t *testing.T, rootID uuid.UUID, childID uuid.UUID) { + t.Helper() + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitShort), "UPDATE chats SET parent_chat_id = $1, root_chat_id = $1 WHERE id = $2", rootID, childID) + require.NoError(t, err) +} + +func insertArchiveMessage(t *testing.T, f *workerTestFixture, chatID uuid.UUID, createdAt time.Time) { + t.Helper() + msg := dbgen.ChatMessage(t, f.db, database.ChatMessage{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: f.user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: f.model.ID, Valid: true}, + Role: database.ChatMessageRoleUser, + }) + _, err := f.sqlDB.ExecContext(testutil.Context(t, testutil.WaitShort), "UPDATE chat_messages SET created_at = $1 WHERE id = $2", createdAt, msg.ID) + require.NoError(t, err) +} + +func (f *workerTestFixture) newArchiveWorker( + t *testing.T, + pubsub *recordingPubsub, + auditor *atomic.Pointer[audit.Auditor], + enqueuer *notificationstest.FakeEnqueuer, +) *chatWorker { + t.Helper() + if pubsub == nil { + pubsub = newRecordingPubsub(f.pubsub) + } + if enqueuer == nil { + enqueuer = notificationstest.NewFakeEnqueuer() + } + opts := f.archiveWorkerOptions() + opts.Pubsub = pubsub + opts.NotificationsEnqueuer = enqueuer + opts.Auditor = auditor + return f.newArchiveWorkerWithOptions(t, opts) +} + +// archiveWorkerOptions returns a baseline chatWorkerOptions with the long +// intervals and channel sizes the archive tests rely on. Callers override +// Pubsub, Store, Clock, and the dispatch dependencies as needed. +func (f *workerTestFixture) archiveWorkerOptions() chatWorkerOptions { + return chatWorkerOptions{ + WorkerID: uuid.New(), + Store: f.db, + Logger: slog.Make(), + TaskStarter: newRecordingTaskStarter(), + AcquisitionInterval: time.Hour, + AcquisitionBatchSize: 10, + ArchiveInterval: time.Hour, + ArchiveBatchSize: 10, + RunnerSyncInterval: time.Hour, + HeartbeatInterval: time.Hour, + HeartbeatCleanupInterval: time.Hour, + HeartbeatStaleSeconds: 30, + StateChannelSize: 16, + RunnerManagerChannelSize: 16, + AcquisitionWakeChannelSize: 1, + } +} + +func (f *workerTestFixture) newArchiveWorkerWithOptions(t *testing.T, opts chatWorkerOptions) *chatWorker { + t.Helper() + if opts.Pubsub == nil { + opts.Pubsub = newRecordingPubsub(f.pubsub) + } + if opts.NotificationsEnqueuer == nil { + opts.NotificationsEnqueuer = notificationstest.NewFakeEnqueuer() + } + worker, err := newChatWorker(nil, opts) + require.NoError(t, err) + return worker +} + +func mockAuditorPtr(auditor *audit.MockAuditor) *atomic.Pointer[audit.Auditor] { + var ptr atomic.Pointer[audit.Auditor] + var asInterface audit.Auditor = auditor + ptr.Store(&asInterface) + return &ptr +} + +func requireWatchEvent(t *testing.T, pubsub *recordingPubsub, chatID uuid.UUID, kind codersdk.ChatWatchEventKind) { + t.Helper() + for _, event := range pubsub.watchEvents(t) { + if event.Kind == kind && event.Chat.ID == chatID { + return + } + } + t.Fatalf("missing watch event kind=%s chat_id=%s", kind, chatID) +} + +// Candidate selection (query) semantics. + +func TestWorker_AutoArchiveSkipsPinnedRoot(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + f.setPinOrder(t, chat.ID, 1) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID), "pinned root must not be auto-archived") +} + +func TestWorker_AutoArchiveSkipsActiveStatusRoot(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + forceExecutionState(t, f, chat.ID, database.ChatStatusRunning, false) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID), "running root must not be auto-archived") +} + +func TestWorker_AutoArchiveIgnoresDeletedMessages(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + insertArchiveMessage(t, f, chat.ID, now.Add(-10*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + worker.archiveOnce(ctx, now) + require.False(t, f.archived(t, chat.ID), "recent message must keep the chat active") + + // Once the only recent message is soft-deleted, activity falls back to + // created_at and the chat becomes eligible. + f.softDeleteMessages(t, chat.ID) + worker.archiveOnce(ctx, now) + require.True(t, f.archived(t, chat.ID), "chat with only deleted messages must archive on created_at") +} + +func TestWorker_AutoArchiveChildActivityKeepsRootAlive(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + root := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + child := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + f.linkChild(t, root.ID, child.ID) + insertArchiveMessage(t, f, child.ID, now.Add(-5*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, root.ID), "recent child activity must keep the root alive") + require.False(t, f.archived(t, child.ID)) +} + +func TestWorker_AutoArchiveBatchSizeLimitsAndPaginates(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + oldest := f.createArchiveCandidate(t, now.Add(-122*24*time.Hour)) + middle := f.createArchiveCandidate(t, now.Add(-121*24*time.Hour)) + newest := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + opts := f.archiveWorkerOptions() + opts.Pubsub = newRecordingPubsub(f.pubsub) + opts.ArchiveBatchSize = 2 + worker := f.newArchiveWorkerWithOptions(t, opts) + + // First tick archives the two oldest roots (created_at ASC, limited). + worker.archiveOnce(ctx, now) + require.True(t, f.archived(t, oldest.ID), "oldest root should archive in the first batch") + require.True(t, f.archived(t, middle.ID), "middle root should archive in the first batch") + require.False(t, f.archived(t, newest.ID), "newest root should wait for the next tick") + + // Second tick drains the remaining backlog. + worker.archiveOnce(ctx, now) + require.True(t, f.archived(t, newest.ID), "newest root should archive on the second tick") +} + +func TestWorker_AutoArchiveNoEligibleChats(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + // A recent chat is well within the inactivity window. + chat := f.createArchiveCandidate(t, now.Add(-24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + auditor := audit.NewMock() + enqueuer := notificationstest.NewFakeEnqueuer() + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), mockAuditorPtr(auditor), enqueuer) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID)) + require.Empty(t, auditor.AuditLogs()) + require.Empty(t, enqueuer.Sent()) +} + +// Dispatch (audit + digest) semantics. + +func TestWorker_AutoArchiveMultipleOwnersGetSeparateDigests(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + user2 := dbgen.User(t, f.db, database.User{}) + chat1 := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + chat2 := f.createArchiveCandidateForOwner(t, user2.ID, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + auditor := audit.NewMock() + enqueuer := notificationstest.NewFakeEnqueuer() + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), mockAuditorPtr(auditor), enqueuer) + worker.archiveOnce(ctx, now) + + require.True(t, f.archived(t, chat1.ID)) + require.True(t, f.archived(t, chat2.ID)) + + sent := enqueuer.Sent() + require.Len(t, sent, 2, "each owner should receive its own digest") + require.ElementsMatch(t, []uuid.UUID{f.user.ID, user2.ID}, []uuid.UUID{sent[0].UserID, sent[1].UserID}) + require.Len(t, auditor.AuditLogs(), 2, "each archived root should be audited") +} + +func TestWorker_AutoArchiveAuditsAndDigestsRootOnlyForFamily(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + root := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + child := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + f.linkChild(t, root.ID, child.ID) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + auditor := audit.NewMock() + enqueuer := notificationstest.NewFakeEnqueuer() + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), mockAuditorPtr(auditor), enqueuer) + worker.archiveOnce(ctx, now) + + require.True(t, f.archived(t, root.ID)) + require.True(t, f.archived(t, child.ID)) + + logs := auditor.AuditLogs() + require.Len(t, logs, 1, "only the root should be audited; children inherit the decision") + require.Equal(t, root.ID, logs[0].ResourceID) + require.Len(t, enqueuer.Sent(), 1, "a single-owner family produces one digest") +} + +func TestWorker_AutoArchiveIncrementsRecordsCounter(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + counter := prometheus.NewCounter(prometheus.CounterOpts{Name: "test_chat_auto_archive_records_total"}) + opts := f.archiveWorkerOptions() + opts.Pubsub = newRecordingPubsub(f.pubsub) + opts.AutoArchiveRecords = counter + worker := f.newArchiveWorkerWithOptions(t, opts) + + worker.archiveOnce(ctx, now) + require.True(t, f.archived(t, chat.ID)) + require.InDelta(t, 1.0, promtestutil.ToFloat64(counter), 0.0001, "counter should reflect one archived root") +} + +func TestWorker_AutoArchiveSecondTickIdempotent(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + auditor := audit.NewMock() + enqueuer := notificationstest.NewFakeEnqueuer() + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), mockAuditorPtr(auditor), enqueuer) + + worker.archiveOnce(ctx, now) + require.True(t, f.archived(t, chat.ID)) + require.Len(t, auditor.AuditLogs(), 1) + require.Len(t, enqueuer.Sent(), 1) + + // An already-archived chat is no longer a candidate, so a second tick is a + // no-op for both audit and digest dispatch. + worker.archiveOnce(ctx, now) + require.Len(t, auditor.AuditLogs(), 1, "second tick must not re-audit") + require.Len(t, enqueuer.Sent(), 1, "second tick must not re-notify") +} + +func TestWorker_AutoArchiveCutoffStableAcrossSameDayTicks(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + // created_at is far in the past so the boundary decision is driven purely + // by message activity sitting exactly on the cutoff date. + chat := f.createArchiveCandidate(t, time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)) + // StartOfDay(2026-05-29) - 90d = 2026-02-28; activity on that date is not + // strictly before the cutoff. + insertArchiveMessage(t, f, chat.ID, time.Date(2026, 2, 28, 12, 0, 0, 0, time.UTC)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + worker := f.newArchiveWorker(t, newRecordingPubsub(f.pubsub), nil, nil) + + // Tick early in the UTC day. + worker.archiveOnce(ctx, time.Date(2026, 5, 29, 23, 49, 0, 0, time.UTC)) + require.False(t, f.archived(t, chat.ID), "activity on the cutoff date must survive") + + // Tick later the same UTC day: advancing wall-clock time within a day must + // not change the cutoff ("no trickle"). + worker.archiveOnce(ctx, time.Date(2026, 5, 29, 23, 59, 0, 0, time.UTC)) + require.False(t, f.archived(t, chat.ID), "same-day tick must not change the decision") + + // Tick on the next UTC day: the cutoff advances to 2026-03-01 and the chat + // becomes eligible. + worker.archiveOnce(ctx, time.Date(2026, 5, 30, 0, 9, 0, 0, time.UTC)) + require.True(t, f.archived(t, chat.ID), "cutoff advances on the next UTC day") +} + +func TestWorker_AutoArchiveDigestDispatchContinuesAfterEnqueueError(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + owner1 := uuid.New() + owner2 := uuid.New() + enq := &recordingEnqueuer{failOwner: owner1} + opts := f.archiveWorkerOptions() + opts.NotificationsEnqueuer = enq + worker := f.newArchiveWorkerWithOptions(t, opts) + + roots := []autoArchivedChat{ + {Chat: database.Chat{OwnerID: owner1, OrganizationID: f.org.ID, Title: "a"}, LastActivityAt: time.Now()}, + {Chat: database.Chat{OwnerID: owner2, OrganizationID: f.org.ID, Title: "b"}, LastActivityAt: time.Now()}, + } + worker.enqueueAutoArchiveDigests(context.Background(), time.Now(), 90, 30, roots) + + require.ElementsMatch(t, []uuid.UUID{owner1, owner2}, enq.enqueuedOwners(), + "a transient enqueue failure must not abort the dispatch loop") +} + +func TestWorker_AutoArchiveDigestDispatchStopsWhenCanceled(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + enq := &recordingEnqueuer{} + opts := f.archiveWorkerOptions() + opts.NotificationsEnqueuer = enq + worker := f.newArchiveWorkerWithOptions(t, opts) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + roots := []autoArchivedChat{ + {Chat: database.Chat{OwnerID: uuid.New(), OrganizationID: f.org.ID}, LastActivityAt: time.Now()}, + {Chat: database.Chat{OwnerID: uuid.New(), OrganizationID: f.org.ID}, LastActivityAt: time.Now()}, + } + worker.enqueueAutoArchiveDigests(ctx, time.Now(), 90, 30, roots) + + require.Empty(t, enq.enqueuedOwners(), "canceled dispatch must enqueue nothing") +} + +// Config / query error handling. + +func TestWorker_AutoArchiveDaysConfigReadFailureSkipsTick(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + enqueuer := notificationstest.NewFakeEnqueuer() + opts := f.archiveWorkerOptions() + opts.Store = &archiveErrStore{Store: f.db, autoArchiveDaysErr: xerrors.New("boom")} + opts.NotificationsEnqueuer = enqueuer + worker := f.newArchiveWorkerWithOptions(t, opts) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID), "auto-archive config read failure must skip the tick") + require.Empty(t, enqueuer.Sent()) +} + +func TestWorker_AutoArchiveRetentionConfigReadFailureSkipsTick(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + enqueuer := notificationstest.NewFakeEnqueuer() + opts := f.archiveWorkerOptions() + opts.Store = &archiveErrStore{Store: f.db, retentionDaysErr: xerrors.New("boom")} + opts.NotificationsEnqueuer = enqueuer + worker := f.newArchiveWorkerWithOptions(t, opts) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID), "retention config read failure must skip the tick") + require.Empty(t, enqueuer.Sent()) +} + +func TestWorker_AutoArchiveCandidateQueryFailureSkipsTick(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + now := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + chat := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + enqueuer := notificationstest.NewFakeEnqueuer() + opts := f.archiveWorkerOptions() + opts.Store = &archiveErrStore{Store: f.db, candidatesErr: xerrors.New("boom")} + opts.NotificationsEnqueuer = enqueuer + worker := f.newArchiveWorkerWithOptions(t, opts) + worker.archiveOnce(ctx, now) + + require.False(t, f.archived(t, chat.ID), "candidate query failure must skip the tick") + require.Empty(t, enqueuer.Sent()) +} + +// Loop wiring. + +func TestWorker_AutoArchiveLoopRunsImmediatelyAndOnTick(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + ctx := testutil.Context(t, testutil.WaitLong) + require.NoError(t, f.db.UpsertChatAutoArchiveDays(ctx, 90)) + + mClock := quartz.NewMock(t) + now := mClock.Now().UTC() + first := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + + opts := f.archiveWorkerOptions() + opts.Pubsub = newRecordingPubsub(f.pubsub) + opts.Clock = mClock + opts.ArchiveInterval = time.Minute + worker := f.newArchiveWorkerWithOptions(t, opts) + + trap := mClock.Trap().NewTicker("chatworker", "auto-archive") + defer trap.Close() + + loopCtx, cancel := context.WithCancel(ctx) + done := make(chan struct{}) + go func() { + defer close(done) + worker.archiveLoop(loopCtx) + }() + + // archiveLoop creates the ticker before the immediate startup tick. + trap.MustWait(ctx).MustRelease(ctx) + testutil.Eventually(ctx, t, func(context.Context) bool { + return f.archived(t, first.ID) + }, testutil.IntervalFast, "immediate startup tick should archive the first candidate") + + // A second candidate is only archived once the interval ticker fires. + second := f.createArchiveCandidate(t, now.Add(-120*24*time.Hour)) + mClock.Advance(time.Minute).MustWait(ctx) + testutil.Eventually(ctx, t, func(context.Context) bool { + return f.archived(t, second.ID) + }, testutil.IntervalFast, "interval tick should archive the second candidate") + + cancel() + select { + case <-done: + case <-ctx.Done(): + t.Fatal("archiveLoop did not exit after context cancellation") + } +} + +func TestBuildAutoArchiveDigestData(t *testing.T) { + t.Parallel() + tickStart := time.Date(2026, 5, 29, 12, 0, 0, 0, time.UTC) + + t.Run("UnderCap", func(t *testing.T) { + t.Parallel() + rows := make([]autoArchivedChat, 0, 3) + for i := range 3 { + rows = append(rows, autoArchivedChat{ + Chat: database.Chat{Title: fmt.Sprintf("chat-%d", i)}, + LastActivityAt: tickStart.Add(-time.Duration(i+1) * 24 * time.Hour), + }) + } + data := buildAutoArchiveDigestData(rows, 90, 30, tickStart) + require.Equal(t, "90", data["auto_archive_days"]) + require.Equal(t, "30", data["retention_days"]) + chats, ok := data["archived_chats"].([]map[string]any) + require.True(t, ok) + require.Len(t, chats, 3) + require.Equal(t, "chat-0", chats[0]["title"]) + require.Contains(t, chats[0]["last_activity_humanized"].(string), "ago") + require.NotContains(t, data, "additional_archived_count") + }) + + t.Run("OverflowCap", func(t *testing.T) { + t.Parallel() + total := chatAutoArchiveDigestMaxChats + 5 + rows := make([]autoArchivedChat, 0, total) + for i := range total { + rows = append(rows, autoArchivedChat{ + Chat: database.Chat{Title: fmt.Sprintf("chat-%d", i)}, + LastActivityAt: tickStart.Add(-24 * time.Hour), + }) + } + data := buildAutoArchiveDigestData(rows, 90, 0, tickStart) + chats, ok := data["archived_chats"].([]map[string]any) + require.True(t, ok) + require.Len(t, chats, chatAutoArchiveDigestMaxChats, "titles are capped") + require.Equal(t, "5", data["additional_archived_count"]) + require.Equal(t, "0", data["retention_days"]) + }) +} + +func TestIsExpectedAutoArchiveError(t *testing.T) { + t.Parallel() + expected := []error{ + sql.ErrNoRows, + chatstate.ErrChatNotFound, + chatstate.ErrChatNotRoot, + chatstate.ErrInvalidState, + chatstate.ErrTransitionNotAllowed, + } + for _, err := range expected { + require.True(t, isExpectedAutoArchiveError(err), "%v should be classified as expected", err) + require.True(t, isExpectedAutoArchiveError(xerrors.Errorf("wrapped: %w", err)), + "wrapped %v should still be classified as expected", err) + } + require.False(t, isExpectedAutoArchiveError(xerrors.New("unexpected"))) +} + +// recordingEnqueuer records the owner of every enqueue and can be configured to +// fail for a specific owner (or all owners) to exercise dispatch resilience. +type recordingEnqueuer struct { + mu sync.Mutex + owners []uuid.UUID + failOwner uuid.UUID + failAll bool +} + +func (e *recordingEnqueuer) Enqueue(ctx context.Context, userID, templateID uuid.UUID, labels map[string]string, createdBy string, targets ...uuid.UUID) ([]uuid.UUID, error) { + return e.EnqueueWithData(ctx, userID, templateID, labels, nil, createdBy, targets...) +} + +func (e *recordingEnqueuer) EnqueueWithData(_ context.Context, userID, _ uuid.UUID, _ map[string]string, _ map[string]any, _ string, _ ...uuid.UUID) ([]uuid.UUID, error) { + e.mu.Lock() + e.owners = append(e.owners, userID) + e.mu.Unlock() + if e.failAll || userID == e.failOwner { + return nil, xerrors.New("enqueue failed") + } + return []uuid.UUID{uuid.New()}, nil +} + +func (e *recordingEnqueuer) enqueuedOwners() []uuid.UUID { + e.mu.Lock() + defer e.mu.Unlock() + return append([]uuid.UUID(nil), e.owners...) +} + +// archiveErrStore wraps a real store and injects errors on the reads performed +// at the start of an auto-archive tick. +type archiveErrStore struct { + database.Store + autoArchiveDaysErr error + retentionDaysErr error + candidatesErr error +} + +func (s *archiveErrStore) GetChatAutoArchiveDays(ctx context.Context, defaultAutoArchiveDays int32) (int32, error) { + if s.autoArchiveDaysErr != nil { + return 0, s.autoArchiveDaysErr + } + return s.Store.GetChatAutoArchiveDays(ctx, defaultAutoArchiveDays) +} + +func (s *archiveErrStore) GetChatRetentionDays(ctx context.Context) (int32, error) { + if s.retentionDaysErr != nil { + return 0, s.retentionDaysErr + } + return s.Store.GetChatRetentionDays(ctx) +} + +func (s *archiveErrStore) GetAutoArchiveInactiveChatCandidates(ctx context.Context, arg database.GetAutoArchiveInactiveChatCandidatesParams) ([]database.GetAutoArchiveInactiveChatCandidatesRow, error) { + if s.candidatesErr != nil { + return nil, s.candidatesErr + } + return s.Store.GetAutoArchiveInactiveChatCandidates(ctx, arg) +} diff --git a/coderd/x/chatd/chatadvisor/runner.go b/coderd/x/chatd/chatadvisor/runner.go index d95ef226fb..fe31afbc5c 100644 --- a/coderd/x/chatd/chatadvisor/runner.go +++ b/coderd/x/chatd/chatadvisor/runner.go @@ -50,20 +50,14 @@ func (rt *Runtime) RunAdvisor( nestedProviderOptions := cloneProviderOptions(rt.cfg.ProviderOptions) resetProviderOptionsForNestedCall(nestedProviderOptions) - var persistedStep chatloop.PersistedStep - chatLoopOpts := chatloop.RunOptions{ + assistantOpts := chatloop.GenerateAssistantOptions{ Model: rt.cfg.Model, Messages: BuildAdvisorMessages(question, conversationSnapshot), - MaxSteps: 1, ModelConfig: rt.cfg.ModelConfig, ProviderOptions: nestedProviderOptions, - PersistStep: func(_ context.Context, step chatloop.PersistedStep) error { - persistedStep = step - return nil - }, } if opts != nil && opts.OnAdviceDelta != nil { - chatLoopOpts.PublishMessagePart = func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + assistantOpts.PublishMessagePart = func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { if role != codersdk.ChatMessageRoleAssistant || part.Type != codersdk.ChatMessagePartTypeText || part.Text == "" { @@ -72,13 +66,17 @@ func (rt *Runtime) RunAdvisor( opts.OnAdviceDelta(part.Text) } } - if opts != nil && opts.OnAdviceReset != nil { - chatLoopOpts.OnRetry = func(int, error, chatretry.ClassifiedError, time.Duration) { + + var outcome chatloop.AssistantOutcome + if err := chatretry.Retry(ctx, func(retryCtx context.Context) error { + var err error + outcome, err = chatloop.GenerateAssistant(retryCtx, assistantOpts) + return err + }, func(int, error, chatretry.ClassifiedError, time.Duration) { + if opts != nil && opts.OnAdviceReset != nil { opts.OnAdviceReset() } - } - - if err := chatloop.Run(ctx, chatLoopOpts); err != nil { + }); err != nil { // Refund the use so a transient provider failure does not // permanently exhaust the per-run advisor budget. rt.release() @@ -89,7 +87,7 @@ func (rt *Runtime) RunAdvisor( }, nil } - advice := extractAdvisorText(persistedStep) + advice := extractAdvisorText(outcome.Step) if advice == "" { // Refund: the run did not produce advice, so the contract // "increments on every successful advisor call" treats this diff --git a/coderd/x/chatd/chatadvisor/tool.go b/coderd/x/chatd/chatadvisor/tool.go index ea15becbd9..5878a71c4e 100644 --- a/coderd/x/chatd/chatadvisor/tool.go +++ b/coderd/x/chatd/chatadvisor/tool.go @@ -6,6 +6,9 @@ import ( "strings" "charm.land/fantasy" + + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" + "github.com/coder/coder/v2/codersdk" ) // ToolName is the identifier the advisor tool registers under. The parent @@ -22,8 +25,6 @@ const advisorQuestionMaxRunes = 2000 type ToolOptions struct { Runtime *Runtime GetConversationSnapshot func() []fantasy.Message - PublishAdviceDelta func(toolCallID string, delta string) - PublishAdviceReset func(toolCallID string) } // Tool returns a fantasy.AgentTool that asks a nested model for concise @@ -46,18 +47,37 @@ func Tool(opts ToolOptions) fantasy.AgentTool { return fantasy.NewTextErrorResponse("question is required"), nil } + // The publisher is injected into the execution context by + // chatloop.ExecuteLocalTools; a missing publisher indicates + // a wiring bug rather than a recoverable condition, so fail + // loudly instead of silently dropping the streamed advice. + publish := chatloop.MessagePartPublisherFromContext(ctx) + if publish == nil { + return fantasy.NewTextErrorResponse("advisor tool requires a message-part publisher on the context; this is an internal tool bug"), nil + } + var runOpts *RunAdvisorOptions - if call.ID != "" && (opts.PublishAdviceDelta != nil || opts.PublishAdviceReset != nil) { - runOpts = &RunAdvisorOptions{} - if opts.PublishAdviceDelta != nil { - runOpts.OnAdviceDelta = func(delta string) { - opts.PublishAdviceDelta(call.ID, delta) - } - } - if opts.PublishAdviceReset != nil { - runOpts.OnAdviceReset = func() { - opts.PublishAdviceReset(call.ID) - } + if call.ID != "" { + runOpts = &RunAdvisorOptions{ + OnAdviceDelta: func(delta string) { + if delta == "" { + return + } + publish(codersdk.ChatMessageRoleTool, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: call.ID, + ToolName: ToolName, + ResultDelta: delta, + }) + }, + OnAdviceReset: func() { + publish(codersdk.ChatMessageRoleTool, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: call.ID, + ToolName: ToolName, + ResultReset: true, + }) + }, } } diff --git a/coderd/x/chatd/chatadvisor/tool_test.go b/coderd/x/chatd/chatadvisor/tool_test.go index 28734e9707..0b69ee0abd 100644 --- a/coderd/x/chatd/chatadvisor/tool_test.go +++ b/coderd/x/chatd/chatadvisor/tool_test.go @@ -12,7 +12,9 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/x/chatd/chatadvisor" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/codersdk" ) func TestAdvisorToolSuccess(t *testing.T) { @@ -62,12 +64,6 @@ func TestAdvisorToolSuccess(t *testing.T) { func TestAdvisorToolPublishesAdviceDeltasWithToolCallID(t *testing.T) { t.Parallel() - type publishedDelta struct { - toolCallID string - delta string - } - var published []publishedDelta - runtime, err := chatadvisor.NewRuntime(chatadvisor.RuntimeConfig{ Model: &chattest.FakeModel{ ProviderName: "test-provider", @@ -90,17 +86,23 @@ func TestAdvisorToolPublishesAdviceDeltasWithToolCallID(t *testing.T) { tool := chatadvisor.Tool(chatadvisor.ToolOptions{ Runtime: runtime, GetConversationSnapshot: func() []fantasy.Message { return nil }, - PublishAdviceDelta: func(toolCallID string, delta string) { - published = append(published, publishedDelta{toolCallID: toolCallID, delta: delta}) - }, }) - resp := runAdvisorTool(t, tool, chatadvisor.AdvisorArgs{Question: "What's safest?"}) + var published []codersdk.ChatMessagePart + resp := runAdvisorToolWithPublisher(t, tool, chatadvisor.AdvisorArgs{Question: "What's safest?"}, + func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + require.Equal(t, codersdk.ChatMessageRoleTool, role) + published = append(published, part) + }) require.False(t, resp.IsError) - require.Equal(t, []publishedDelta{ - {toolCallID: "call-1", delta: "Prefer "}, - {toolCallID: "call-1", delta: "the small diff."}, - }, published) + require.Len(t, published, 2) + for _, part := range published { + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, part.Type) + require.Equal(t, "call-1", part.ToolCallID) + require.Equal(t, chatadvisor.ToolName, part.ToolName) + } + require.Equal(t, "Prefer ", published[0].ResultDelta) + require.Equal(t, "the small diff.", published[1].ResultDelta) var result chatadvisor.AdvisorResult require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) @@ -116,10 +118,7 @@ func TestAdvisorToolPublishesAdviceResetWithToolCallID(t *testing.T) { toolCallID string delta string } - var ( - calls int - published []publishedEvent - ) + var calls int runtime, err := chatadvisor.NewRuntime(chatadvisor.RuntimeConfig{ Model: &chattest.FakeModel{ @@ -150,22 +149,24 @@ func TestAdvisorToolPublishesAdviceResetWithToolCallID(t *testing.T) { tool := chatadvisor.Tool(chatadvisor.ToolOptions{ Runtime: runtime, GetConversationSnapshot: func() []fantasy.Message { return nil }, - PublishAdviceDelta: func(toolCallID string, delta string) { - published = append(published, publishedEvent{ - kind: "delta", - toolCallID: toolCallID, - delta: delta, - }) - }, - PublishAdviceReset: func(toolCallID string) { - published = append(published, publishedEvent{ - kind: "reset", - toolCallID: toolCallID, - }) - }, }) - resp := runAdvisorTool(t, tool, chatadvisor.AdvisorArgs{Question: "What's safest?"}) + var published []publishedEvent + resp := runAdvisorToolWithPublisher(t, tool, chatadvisor.AdvisorArgs{Question: "What's safest?"}, + func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + require.Equal(t, codersdk.ChatMessageRoleTool, role) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, part.Type) + require.Equal(t, chatadvisor.ToolName, part.ToolName) + kind := "delta" + if part.ResultReset { + kind = "reset" + } + published = append(published, publishedEvent{ + kind: kind, + toolCallID: part.ToolCallID, + delta: part.ResultDelta, + }) + }) require.False(t, resp.IsError) require.Equal(t, []publishedEvent{ {kind: "delta", toolCallID: "call-1", delta: "stale "}, @@ -194,6 +195,30 @@ func TestAdvisorToolRejectsEmptyQuestion(t *testing.T) { require.Contains(t, resp.Content, "question is required") } +func TestAdvisorToolMissingPublisherReturnsError(t *testing.T) { + t.Parallel() + + tool := chatadvisor.Tool(chatadvisor.ToolOptions{ + Runtime: mustAdvisorRuntime(t), + GetConversationSnapshot: func() []fantasy.Message { + return nil + }, + }) + + data, err := json.Marshal(chatadvisor.AdvisorArgs{Question: "anything?"}) + require.NoError(t, err) + + resp, err := tool.Run(t.Context(), fantasy.ToolCall{ + ID: "call-1", + Name: "advisor", + Input: string(data), + }) + require.NoError(t, err) + require.True(t, resp.IsError) + require.Contains(t, resp.Content, "message-part publisher") + require.Contains(t, resp.Content, "internal tool bug") +} + func TestAdvisorToolPassesNormalQuestion(t *testing.T) { t.Parallel() @@ -444,11 +469,22 @@ func runAdvisorTool( args chatadvisor.AdvisorArgs, ) fantasy.ToolResponse { t.Helper() + return runAdvisorToolWithPublisher(t, tool, args, func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}) +} + +func runAdvisorToolWithPublisher( + t *testing.T, + tool fantasy.AgentTool, + args chatadvisor.AdvisorArgs, + publish func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), +) fantasy.ToolResponse { + t.Helper() data, err := json.Marshal(args) require.NoError(t, err) - resp, err := tool.Run(t.Context(), fantasy.ToolCall{ + ctx := chatloop.WithMessagePartPublisher(t.Context(), publish) + resp, err := tool.Run(ctx, fantasy.ToolCall{ ID: "call-1", Name: "advisor", Input: string(data), diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index f610915ea1..af448f23c3 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -8,8 +8,6 @@ import ( "encoding/json" "errors" "fmt" - "maps" - "math" "net/http" "slices" "strconv" @@ -30,10 +28,12 @@ import ( "cdr.dev/slog/v3" "github.com/coder/coder/v2/coderd/aibridge" + "github.com/coder/coder/v2/coderd/audit" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/db2sdk" "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/pubsub" + "github.com/coder/coder/v2/coderd/notifications" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/ptr" @@ -48,11 +48,11 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chatopenai" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" - "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/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/coderd/x/chatd/internal/agentselect" "github.com/coder/coder/v2/coderd/x/chatd/mcpclient" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" skillspkg "github.com/coder/coder/v2/coderd/x/skills" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" @@ -69,7 +69,6 @@ const ( homeInstructionLookupTimeout = 5 * time.Second planPathLookupTimeout = 5 * time.Second - instructionCacheTTL = 5 * time.Minute workspaceDialValidationDelay = 5 * time.Second // Must exceed agent/x/agentmcp.connectTimeout (30s) so a // cold-start agent's first MCP reload can settle before @@ -98,20 +97,6 @@ const ( // heartbeat updates while a chat is being processed. DefaultChatHeartbeatInterval = 30 * time.Second maxChatSteps = 1200 - // maxStreamBufferSize caps the number of message_part events buffered - // per chat during a single LLM step. When exceeded the oldest event is - // evicted so memory stays bounded. - maxStreamBufferSize = 10000 - // RelaySentinelAfterID is the after_id sentinel used by cross-replica - // relay subscribers. It instructs the peer to skip the durable DB - // snapshot and only deliver buffered message_part events. The - // buffer itself filters committed parts out (see snapshotBufferLocked), - // so the sentinel resolves to "send me any in-progress streaming - // parts you have; I will receive durable messages through pubsub." - RelaySentinelAfterID = math.MaxInt64 - // maxDurableMessageCacheSize caps the number of recent durable message - // events cached per chat for same-replica stream catch-up. - maxDurableMessageCacheSize = 256 // maxConcurrentRecordingUploads caps the number of recording // stop-and-store operations that can run concurrently. Each @@ -120,38 +105,6 @@ const ( // to roughly maxConcurrentRecordingUploads * 110 MB. maxConcurrentRecordingUploads = 25 - // staleRecoveryIntervalDivisor determines how often the stale - // recovery loop runs relative to the stale threshold. A value - // of 5 means recovery runs at 1/5 of the stale-after duration. - staleRecoveryIntervalDivisor = 5 - - // streamDropWarnInterval controls how often WARN-level logs are - // emitted when stream events are dropped. Between intervals the - // drop is logged at DEBUG to avoid log spam. This uses a - // timestamp comparison rather than a quartz.Ticker because the - // state is per-chat — a ticker per chat would require extra - // goroutines and lifecycle management. - streamDropWarnInterval = 10 * time.Second - - // bufferRetainGracePeriod is how long the per-chat stream - // state is kept after processing completes. The retained - // state lets late-connecting cross-replica relay subscribers - // register against the live stream before the next worker - // run starts, preventing a race between cleanupStreamIfIdle - // and subscriber registration. The buffer itself is no - // longer useful at this point: every part has been claimed - // by its durable assistant message and is filtered out of - // the subscriber snapshot. - bufferRetainGracePeriod = 5 * time.Second - // chatStreamControlFetchTimeout bounds subscriber-owned - // control-path DB reads when the caller has no deadline. - chatStreamControlFetchTimeout = 5 * time.Second - - // streamJanitorInterval is how often sweepIdleStreams runs. - // Worst-case retention is bufferRetainGracePeriod + - // streamJanitorInterval. - streamJanitorInterval = 30 * time.Second - // agentDisconnectedRecoveryThreshold is how long the latest // workspace agent must be disconnected before chatd suggests // destructive stop/start recovery. This is intentionally longer @@ -219,7 +172,7 @@ type Server struct { workerID uuid.UUID logger slog.Logger - subscribeFn SubscribeFn + streamPartsDialer StreamPartsDialer agentConnFn AgentConnFunc agentInactiveDisconnectTimeout time.Duration @@ -240,20 +193,18 @@ type Server struct { configCache *chatConfigCache configCacheUnsubscribe func() - // chatStreams stores per-chat stream state. Using sync.Map - // gives each chat independent locking — concurrent chats - // never contend with each other. - chatStreams sync.Map // uuid.UUID -> *chatStreamState - // workspaceMCPToolsCache caches workspace MCP tool definitions // per chat to avoid re-fetching on every turn. The cache is // keyed by chat ID and invalidated when the agent changes. workspaceMCPToolsCache sync.Map // uuid.UUID -> *cachedWorkspaceMCPTools - usageTracker *workspacestats.UsageTracker - clock quartz.Clock - metrics *chatloop.Metrics - recordingSem chan struct{} + usageTracker *workspacestats.UsageTracker + clock quartz.Clock + metrics *chatloop.Metrics + chatWorker *chatWorker + messagePartBuffer *messagepartbuffer.Buffer + streamSyncPoller *streamSyncPoller + recordingSem chan struct{} aibridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] aiGatewayRoutingEnabled bool @@ -263,17 +214,6 @@ type Server struct { maxChatsPerAcquire int32 inFlightChatStaleAfter time.Duration chatHeartbeatInterval time.Duration - - // heartbeatMu guards heartbeatRegistry. - heartbeatMu sync.Mutex - // heartbeatRegistry maps chat IDs to their cancel functions - // and workspace state for the centralized heartbeat loop. - heartbeatRegistry map[uuid.UUID]*heartbeatEntry - - // wakeCh is signaled whenever a chat transitions to - // pending so the run loop calls processOnce immediately - // instead of waiting for the next ticker. - wakeCh chan struct{} } // chatTemplateAllowlist returns the deployment-wide template @@ -541,8 +481,8 @@ func (p *Server) loadCachedWorkspaceContext( // cache. Returns nil (and never an error) on every failure mode so the // caller can continue without MCP tools. // -// This helper is shared between the top-of-turn discovery path and the -// mid-turn PrepareTools path triggered after create_workspace / +// This helper is shared between the initial discovery path and the +// mid-turn workspace binding path triggered after create_workspace or // start_workspace bind a workspace to a chat that started without one. func (p *Server) discoverWorkspaceMCPTools( ctx context.Context, @@ -617,9 +557,9 @@ func (p *Server) discoverWorkspaceMCPTools( // When the agent's MCP server is still racing with agent startup, // ListMCPTools may return an empty list (no error) on the first call; // the primer retries with a short backoff up to -// workspaceMCPPrimeMaxWait so the LLM step that follows the tool call -// sees the workspace MCP tools in the cache and PrepareTools does not -// need to dial again. +// workspaceMCPPrimeMaxWait so the generation action that follows the +// tool call sees the workspace MCP tools in the cache and does not need +// to dial again. // // Returns silently on every failure mode. The chat continues without // workspace MCP tools when the agent does not advertise any within @@ -721,6 +661,17 @@ func (c *turnWorkspaceContext) currentWorkspaceMatches(expected uuid.NullUUID) ( return chatSnapshot, nullUUIDEqual(chatSnapshot.WorkspaceID, expected) } +func (c *turnWorkspaceContext) trackWorkspaceUsage(ctx context.Context, chatSnapshot database.Chat) { + if c.server == nil || !chatSnapshot.WorkspaceID.Valid { + return + } + logger := c.server.logger.With( + slog.F("chat_id", chatSnapshot.ID), + slog.F("owner_id", chatSnapshot.OwnerID), + ) + c.server.trackWorkspaceUsage(ctx, chatSnapshot.ID, chatSnapshot.WorkspaceID, logger) +} + func nullUUIDEqual(left, right uuid.NullUUID) bool { if left.Valid != right.Valid { return false @@ -1078,6 +1029,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces // row so we see the latest heartbeat rather than // a potentially stale cached copy. if currentConn != nil { + chatSnapshot := c.currentChatSnapshot() if agentID != uuid.Nil { freshAgent, err := c.server.db.GetWorkspaceAgentByID(ctx, agentID) if err != nil { @@ -1096,6 +1048,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces continue } } + c.trackWorkspaceUsage(ctx, chatSnapshot) return currentConn, nil } if staleRelease != nil { @@ -1233,6 +1186,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces slog.F("workspace_id", chatSnapshot.WorkspaceID.UUID), slog.F("agent_id", dialResult.AgentID), ) + c.trackWorkspaceUsage(ctx, chatSnapshot) return agentConn, nil } currentConn = c.conn @@ -1241,6 +1195,7 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces if agentRelease != nil { agentRelease() } + c.trackWorkspaceUsage(ctx, chatSnapshot) return currentConn, nil } @@ -1250,178 +1205,9 @@ func (c *turnWorkspaceContext) getWorkspaceConn(ctx context.Context) (workspaces // AgentConnFunc provides access to workspace agent connections. type AgentConnFunc func(ctx context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) -// SubscribeFn replaces the default local-only subscription with a -// multi-replica-aware implementation that merges pubsub notifications, -// remote relay streams, and local parts into a single event channel. -// When set, Subscribe delegates the event-merge goroutine to this -// function instead of using simple local forwarding. -// -// Parameters: -// - ctx: subscription lifetime context (canceled on unsubscribe). -// - params: all state needed to build the merged stream. -// -// Returns the merged event channel. Cleanup is driven by ctx -// cancellation — the merge goroutine tears down all relay state -// in its defer when ctx is done. -// Set by enterprise for HA deployments. Nil in AGPL single-replica. -type SubscribeFn func( - ctx context.Context, - params SubscribeFnParams, -) <-chan codersdk.ChatStreamEvent - -// StatusNotification informs the enterprise relay manager of chat -// status changes so it can open or close relay connections. -type StatusNotification struct { - Status database.ChatStatus - WorkerID uuid.UUID -} - -// SubscribeFnParams carries the state that the enterprise -// SubscribeFn implementation needs from the OSS Subscribe preamble. -type SubscribeFnParams struct { - ChatID uuid.UUID - Chat database.Chat - WorkerID uuid.UUID - StatusNotifications <-chan StatusNotification - RequestHeader http.Header - DB database.Store - Logger slog.Logger -} - -// bufferedStreamPart is a buffered message_part event with its -// committed-message linkage. Parts that have not yet been claimed by -// a durable assistant message carry committedMessageID == 0 and are -// considered "in progress"; when an assistant message is published -// every still-in-progress part is claimed by that durable message -// ID, marking the part as redundant for any subscriber that will -// receive the durable message via REST or pubsub. -type bufferedStreamPart struct { - event codersdk.ChatStreamEvent - // committedMessageID is the durable assistant message ID that - // claimed this part, or 0 while the part belongs to the - // in-progress turn. snapshotBufferLocked drops parts with - // committedMessageID != 0 because the subscriber will receive - // the durable message through a different channel (REST snapshot, - // initial DB query in SubscribeAuthorized, or pubsub). - committedMessageID int64 -} - -type chatStreamState struct { - mu sync.Mutex - buffer []bufferedStreamPart - buffering bool - durableMessages []codersdk.ChatStreamEvent - durableEvictedBefore int64 // highest message ID evicted from durable cache - subscribers map[uuid.UUID]chan codersdk.ChatStreamEvent - bufferDropCount int64 - bufferLastWarnAt time.Time - subscriberDropCount int64 - subscriberLastWarnAt time.Time - // currentRetry records the current retry phase for late-joining - // same-replica subscribers. Nil when the stream is not waiting - // to retry. - currentRetry *codersdk.ChatStreamRetry - // bufferRetainedAt records when processing completed and - // the per-chat stream state entered the post-completion - // grace window. Zero while buffering is active. When - // non-zero, cleanupStreamIfIdle skips GC until the grace - // period expires so cross-replica relay subscribers can - // register without racing state deletion. The buffer - // itself does not deliver content here: every part is - // claimed by a durable assistant message before - // bufferRetainedAt is set, so snapshotBufferLocked - // returns no parts during the grace window. - bufferRetainedAt time.Time -} - -// heartbeatEntry tracks a single chat's cancel function and workspace -// state for the centralized heartbeat loop. Instead of spawning a -// per-chat goroutine, processChat registers an entry here and the -// single heartbeatLoop goroutine handles all chats. -type heartbeatEntry struct { - cancelWithCause context.CancelCauseFunc - chatID uuid.UUID - workspaceID uuid.NullUUID - logger slog.Logger -} - -// resetDropCounters zeroes the rate-limiting state for both buffer -// and subscriber drop warnings. The caller must hold s.mu. -func (s *chatStreamState) resetDropCounters() { - s.bufferDropCount = 0 - s.bufferLastWarnAt = time.Time{} - s.subscriberDropCount = 0 - s.subscriberLastWarnAt = time.Time{} -} - -// streamStateCollector exposes scrape-time gauges derived from -// p.chatStreams. Scrape cost is O(n) with a brief per-state mutex -// held for two len() reads; acceptable at typical scrape cadences. -type streamStateCollector struct { - server *Server -} - -var ( - streamsActiveDesc = prometheus.NewDesc( - "coderd_chatd_streams_active", - "Current number of chat stream state entries (in-flight plus retained).", - nil, nil, - ) - streamBufferSizeMaxDesc = prometheus.NewDesc( - "coderd_chatd_stream_buffer_size_max", - "Maximum current buffer length across all chat streams.", - nil, nil, - ) - streamBufferEventsDesc = prometheus.NewDesc( - "coderd_chatd_stream_buffer_events", - "Sum of current buffer lengths across all chat streams.", - nil, nil, - ) - streamSubscribersDesc = prometheus.NewDesc( - "coderd_chatd_stream_subscribers", - "Current number of chat stream subscribers across all chat streams.", - nil, nil, - ) -) - -func (*streamStateCollector) Describe(ch chan<- *prometheus.Desc) { - ch <- streamsActiveDesc - ch <- streamBufferSizeMaxDesc - ch <- streamBufferEventsDesc - ch <- streamSubscribersDesc -} - -func (c *streamStateCollector) Collect(ch chan<- prometheus.Metric) { - var active, totalEvents, maxBufLen, totalSubs int - c.server.chatStreams.Range(func(_, v any) bool { - state, ok := v.(*chatStreamState) - if !ok { - return true - } - active++ - state.mu.Lock() - bufLen := len(state.buffer) - subs := len(state.subscribers) - state.mu.Unlock() - totalEvents += bufLen - totalSubs += subs - maxBufLen = max(maxBufLen, bufLen) - return true - }) - ch <- prometheus.MustNewConstMetric(streamsActiveDesc, prometheus.GaugeValue, float64(active)) - ch <- prometheus.MustNewConstMetric(streamBufferSizeMaxDesc, prometheus.GaugeValue, float64(maxBufLen)) - ch <- prometheus.MustNewConstMetric(streamBufferEventsDesc, prometheus.GaugeValue, float64(totalEvents)) - ch <- prometheus.MustNewConstMetric(streamSubscribersDesc, prometheus.GaugeValue, float64(totalSubs)) -} - -// MaxQueueSize is the maximum number of queued user messages per chat. -const MaxQueueSize = 20 - var ( // ErrInvalidModelConfigID indicates the requested model config does not exist. ErrInvalidModelConfigID = xerrors.New("invalid model config ID") - // ErrMessageQueueFull indicates the per-chat queue limit was reached. - ErrMessageQueueFull = xerrors.New("chat message queue is full") // ErrEditedMessageNotFound indicates the edited message does not exist // in the target chat. ErrEditedMessageNotFound = xerrors.New("edited message not found") @@ -1431,12 +1217,6 @@ var ( // accept modifications (messages, edits, promotions, or // tool-result submissions). ErrChatArchived = xerrors.New("chat is archived") - - // errChatTakenByOtherWorker is a sentinel used inside the - // processChat cleanup transaction to signal that another - // worker acquired the chat, so all post-TX side effects - // (status publish, pubsub, web push) must be skipped. - errChatTakenByOtherWorker = xerrors.New("chat acquired by another worker") ) // UsageLimitExceededError indicates the user has exceeded their chat spend @@ -1549,8 +1329,16 @@ type PromoteQueuedResult struct { PromotedMessage database.ChatMessage } -// CreateChat creates a chat, inserts optional system prompt and initial user -// message, and moves the chat into pending status. +func validateChatUserMessageAPIKeyID(apiKeyID string) error { + if apiKeyID == "" { + return xerrors.New("api_key_id is required for user chat messages") + } + return nil +} + +// CreateChat creates a chat with its initial history through +// chatstate.CreateChat. The new chat starts in `running` status per +// the chat execution state model. Ownership hints wake chat workers. func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.Chat, error) { if opts.OrganizationID == uuid.Nil { return database.Chat{}, xerrors.New("organization_id is required") @@ -1564,6 +1352,9 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C if len(opts.InitialUserContent) == 0 { return database.Chat{}, xerrors.New("initial user content is required") } + if err := validateChatUserMessageAPIKeyID(opts.APIKeyID); err != nil { + return database.Chat{}, err + } // Ensure MCPServerIDs is non-nil so pq.Array produces '{}' // instead of SQL NULL, which violates the NOT NULL column // constraint. @@ -1573,150 +1364,107 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C if opts.Labels == nil { opts.Labels = database.StringMap{} } + opts.ClientType = cmp.Or(opts.ClientType, database.ChatClientTypeApi) + if !opts.ClientType.Valid() { + return database.Chat{}, xerrors.Errorf("invalid client_type: %q", opts.ClientType) + } // Resolve the deployment prompt before opening the transaction so // chat creation does not hold one DB connection while waiting for // another pool checkout. deploymentPrompt := p.resolveDeploymentSystemPrompt(ctx) - effectivePlanMode := opts.PlanMode - opts.ClientType = cmp.Or(opts.ClientType, database.ChatClientTypeApi) - if !opts.ClientType.Valid() { - return database.Chat{}, xerrors.Errorf("invalid client_type: %q", opts.ClientType) - } - var chat database.Chat - txErr := p.db.InTx(func(tx database.Store) error { - if limitErr := p.checkUsageLimit(ctx, tx, opts.OwnerID, uuid.NullUUID{UUID: opts.OrganizationID, Valid: true}); limitErr != nil { - return limitErr - } - - labelsJSON, err := json.Marshal(opts.Labels) - if err != nil { - return xerrors.Errorf("marshal labels: %w", err) - } - - insertedChat, err := tx.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: opts.OrganizationID, - OwnerID: opts.OwnerID, - WorkspaceID: opts.WorkspaceID, - BuildID: opts.BuildID, - AgentID: opts.AgentID, - ParentChatID: opts.ParentChatID, - RootChatID: opts.RootChatID, - LastModelConfigID: opts.ModelConfigID, - Title: opts.Title, - Mode: opts.ChatMode, - PlanMode: effectivePlanMode, - ClientType: opts.ClientType, - // Chats created with an initial user message start pending. - // Waiting is reserved for idle chats with no pending work. - Status: database.ChatStatusPending, - MCPServerIDs: opts.MCPServerIDs, - Labels: pqtype.NullRawMessage{ - RawMessage: labelsJSON, - Valid: true, - }, - DynamicTools: pqtype.NullRawMessage{ - RawMessage: opts.DynamicTools, - Valid: len(opts.DynamicTools) > 0, - }, - }) - if err != nil { - return xerrors.Errorf("insert chat: %w", err) - } - - userPrompt := SanitizePromptText(opts.SystemPrompt) - workspaceAwareness := workspaceDetachedAwareness - if opts.WorkspaceID.Valid { - workspaceAwareness = workspaceAttachedAwareness - } - workspaceAwarenessContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(workspaceAwareness), - }) - if err != nil { - return xerrors.Errorf("marshal workspace awareness: %w", err) - } - userContent, err := chatprompt.MarshalParts(opts.InitialUserContent) - if err != nil { - return xerrors.Errorf("marshal initial user content: %w", err) - } - - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. - ChatID: insertedChat.ID, - } - - if deploymentPrompt != "" { - deploymentContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(deploymentPrompt), - }) - if err != nil { - return xerrors.Errorf("marshal deployment system prompt: %w", err) - } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleSystem, - deploymentContent, - database.ChatMessageVisibilityModel, - opts.ModelConfigID, - chatprompt.CurrentContentVersion, - )) - } - - if userPrompt != "" { - userPromptContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(userPrompt), - }) - if err != nil { - return xerrors.Errorf("marshal user system prompt: %w", err) - } - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleSystem, - userPromptContent, - database.ChatMessageVisibilityModel, - opts.ModelConfigID, - chatprompt.CurrentContentVersion, - )) - } - - appendChatMessage(&msgParams, newChatMessage( - database.ChatMessageRoleSystem, - workspaceAwarenessContent, - database.ChatMessageVisibilityModel, - opts.ModelConfigID, - chatprompt.CurrentContentVersion, - )) - - userMsg := newUserChatMessage( - opts.APIKeyID, - userContent, - database.ChatMessageVisibilityBoth, - opts.ModelConfigID, - chatprompt.CurrentContentVersion, - ) - userMsg = userMsg.withCreatedBy(opts.OwnerID) - appendUserChatMessage(&msgParams, userMsg) - - _, err = tx.InsertChatMessages(ctx, msgParams) - if err != nil { - return xerrors.Errorf("insert initial chat messages: %w", err) - } - - chat = insertedChat - - if !chat.RootChatID.Valid && !chat.ParentChatID.Valid { - chat.RootChatID = uuid.NullUUID{UUID: chat.ID, Valid: true} - } - return nil - }, nil) - if txErr != nil { - return database.Chat{}, txErr + // Usage limits gate the create before we touch the state machine. + if limitErr := p.checkUsageLimit(ctx, p.db, opts.OwnerID, uuid.NullUUID{UUID: opts.OrganizationID, Valid: true}); limitErr != nil { + return database.Chat{}, limitErr } + labelsJSON, err := json.Marshal(opts.Labels) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal labels: %w", err) + } + + userPrompt := SanitizePromptText(opts.SystemPrompt) + workspaceAwareness := workspaceDetachedAwareness + if opts.WorkspaceID.Valid { + workspaceAwareness = workspaceAttachedAwareness + } + workspaceAwarenessContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(workspaceAwareness), + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal workspace awareness: %w", err) + } + userContent, err := chatprompt.MarshalParts(opts.InitialUserContent) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal initial user content: %w", err) + } + + var initialMessages []chatstate.Message + if deploymentPrompt != "" { + deploymentContent, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(deploymentPrompt), + }) + if marshalErr != nil { + return database.Chat{}, xerrors.Errorf("marshal deployment system prompt: %w", marshalErr) + } + initialMessages = append(initialMessages, systemMessage(deploymentContent, opts.ModelConfigID)) + } + if userPrompt != "" { + userPromptContent, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(userPrompt), + }) + if marshalErr != nil { + return database.Chat{}, xerrors.Errorf("marshal user system prompt: %w", marshalErr) + } + initialMessages = append(initialMessages, systemMessage(userPromptContent, opts.ModelConfigID)) + } + initialMessages = append(initialMessages, systemMessage(workspaceAwarenessContent, opts.ModelConfigID)) + initialMessages = append(initialMessages, userMessageWithAPIKeyID(userContent, opts.ModelConfigID, opts.OwnerID, opts.APIKeyID)) + + result, err := chatstate.CreateChat(ctx, p.db, p.pubsub, chatstate.CreateChatInput{ + OrganizationID: opts.OrganizationID, + OwnerID: opts.OwnerID, + WorkspaceID: opts.WorkspaceID, + BuildID: opts.BuildID, + AgentID: opts.AgentID, + ParentChatID: opts.ParentChatID, + RootChatID: opts.RootChatID, + LastModelConfigID: opts.ModelConfigID, + Title: opts.Title, + Mode: opts.ChatMode, + PlanMode: opts.PlanMode, + MCPServerIDs: opts.MCPServerIDs, + Labels: pqtype.NullRawMessage{ + RawMessage: labelsJSON, + Valid: true, + }, + DynamicTools: pqtype.NullRawMessage{ + RawMessage: opts.DynamicTools, + Valid: len(opts.DynamicTools) > 0, + }, + ClientType: opts.ClientType, + InitialMessages: initialMessages, + }) + if err != nil { + return database.Chat{}, err + } + chat := result.Chat + if !chat.RootChatID.Valid && !chat.ParentChatID.Valid { + chat.RootChatID = uuid.NullUUID{UUID: chat.ID, Valid: true} + } + + // Publish the sidebar watch event explicitly after chatstate has + // committed and emitted its own state-machine notifications. The + // watch endpoint is maintained separately from chatstate notifications. p.publishChatPubsubEvent(chat, codersdk.ChatWatchEventKindCreated, nil) - p.signalWake() return chat, nil } -// SendMessage inserts a user message and optionally queues it while the chat -// is busy, then publishes stream + pubsub updates. +// SendMessage admits a user message through the chatstate.SendMessage +// transition. Pre-transition admission policy (usage limit, plan-mode +// metadata update, MCP server ID update, model-config resolution, queue +// cap) runs inside the same chatstate transaction via the transactional +// store so everything commits or rolls back together. func (p *Server) SendMessage( ctx context.Context, opts SendMessageOptions, @@ -1727,6 +1475,9 @@ func (p *Server) SendMessage( if len(opts.Content) == 0 { return SendMessageResult{}, xerrors.New("content is required") } + if err := validateChatUserMessageAPIKeyID(opts.APIKeyID); err != nil { + return SendMessageResult{}, err + } busyBehavior := opts.BusyBehavior if busyBehavior == "" { @@ -1744,29 +1495,27 @@ func (p *Server) SendMessage( } requestedPlanMode := opts.PlanMode + requestedMCPServerIDs := opts.MCPServerIDs - var ( - result SendMessageResult - queuedMessagesSDK []codersdk.ChatQueuedMessage - ) - - txErr := p.db.InTx(func(tx database.Store) error { - lockedChat, err := tx.GetChatByIDForUpdate(ctx, opts.ChatID) + var result SendMessageResult + machine := p.newChatMachine(opts.ChatID) + updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { - return xerrors.Errorf("lock chat: %w", err) + return xerrors.Errorf("load chat: %w", err) } if lockedChat.Archived { return ErrChatArchived } - // Enforce usage limits before queueing or inserting. - if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { + // Enforce usage limits before any state-machine work. + if limitErr := p.checkUsageLimit(ctx, store, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { return limitErr } if requestedPlanMode != nil { - lockedChat, err = tx.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{ + lockedChat, err = store.UpdateChatPlanModeByID(ctx, database.UpdateChatPlanModeByIDParams{ PlanMode: *requestedPlanMode, ID: opts.ChatID, }) @@ -1777,7 +1526,7 @@ func (p *Server) SendMessage( modelConfigID, err := resolveSendMessageModelConfigID( ctx, - tx, + store, lockedChat, opts.ModelConfigID, ) @@ -1787,16 +1536,16 @@ func (p *Server) SendMessage( // Update MCP server IDs on the chat when explicitly provided. // Explore child chats keep the spawn-time snapshot immutable. - if opts.MCPServerIDs != nil { + if requestedMCPServerIDs != nil { if isExploreSubagentMode(lockedChat.Mode) { p.logger.Warn(ctx, "ignoring explore subagent mcp server ids update, snapshot is immutable after spawn", slog.F("chat_id", opts.ChatID), ) } else { - lockedChat, err = tx.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{ + lockedChat, err = store.UpdateChatMCPServerIDs(ctx, database.UpdateChatMCPServerIDsParams{ ID: opts.ChatID, - MCPServerIDs: *opts.MCPServerIDs, + MCPServerIDs: *requestedMCPServerIDs, }) if err != nil { return xerrors.Errorf("update chat mcp server ids: %w", err) @@ -1804,114 +1553,48 @@ func (p *Server) SendMessage( } } - existingQueued, err := tx.GetChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("get queued messages: %w", err) + messageCreatedBy := opts.CreatedBy + if messageCreatedBy == uuid.Nil { + messageCreatedBy = lockedChat.OwnerID } - // Both queue and interrupt behaviors queue messages - // when the chat is busy. We also keep queueing while a - // backlog exists so waiting chats blocked by spend limits - // preserve FIFO user-message order. Interrupt additionally - // signals the running loop to stop so the queued message - // is promoted sooner. Crucially, this guarantees the - // interrupted assistant response is persisted (with a - // lower id/created_at) before the user message is - // promoted into chat_messages, preserving correct - // conversation order. - if shouldQueueUserMessage(lockedChat.Status) || len(existingQueued) > 0 { - if len(existingQueued) >= MaxQueueSize { - return ErrMessageQueueFull - } - - queued, err := tx.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: opts.ChatID, - Content: content.RawMessage, - ModelConfigID: uuid.NullUUID{ - UUID: modelConfigID, - Valid: modelConfigID != uuid.Nil, - }, - APIKeyID: sql.NullString{ - String: opts.APIKeyID, - Valid: opts.APIKeyID != "", - }, - }) - if err != nil { - return xerrors.Errorf("insert queued message: %w", err) - } - - queuedMessages, err := tx.GetChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("get queued messages: %w", err) - } - - result.Queued = true - result.QueuedMessage = &queued - result.Chat = lockedChat - queuedMessagesSDK = db2sdk.ChatQueuedMessages(queuedMessages) - return nil - } - - message, updatedChat, err := insertUserMessageAndSetPending( - ctx, - tx, - lockedChat, - modelConfigID, - content, - opts.CreatedBy, - opts.APIKeyID, - ) + // Queue capacity is enforced inside tx.SendMessage; this + // wrapper only propagates the typed error. + sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: userMessageWithAPIKeyID(content, modelConfigID, messageCreatedBy, opts.APIKeyID), + BusyBehavior: busyBehaviorToChatState(busyBehavior), + }) if err != nil { return err } - result.Message = message - result.Chat = updatedChat - return nil - }, nil) - if txErr != nil { - return SendMessageResult{}, txErr - } - - if result.Queued { - p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - ChatID: opts.ChatID, - QueuedMessages: queuedMessagesSDK, - }) - p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, - }) - - // For interrupt behavior, signal the running loop to - // stop. setChatWaiting publishes a status notification - // that the worker's control subscriber detects, causing - // it to cancel with ErrInterrupted. The deferred cleanup - // in processChat then auto-promotes the queued message - // after persisting the partial assistant response. - if busyBehavior == SendMessageBusyBehaviorInterrupt { - updatedChat, err := p.setChatWaiting(ctx, opts.ChatID) - if err != nil { - // The message is already queued so the chat is - // not in a broken state — the user can still - // wait for the current run to finish. Log the - // error but don't fail the request. - p.logger.Error(ctx, "failed to interrupt chat for queued message", - slog.F("chat_id", opts.ChatID), - slog.Error(err), - ) - } else { - result.Chat = updatedChat - } + if sendResult.QueuedMessage != nil { + result.Queued = true + result.QueuedMessage = sendResult.QueuedMessage + } else if len(sendResult.InsertedMessages) > 0 { + // The state machine prepends synthetic tool-result + // cancellation messages; the user message is always + // last in the inserted slice. + result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] } - - return result, nil + // Capture the post-transition chat inside the same + // transaction so the returned chat and the watch event + // reflect the snapshot bump and status change produced by + // the transition itself. + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after send: %w", err) + } + result.Chat = refreshed + return nil + }) + if updateErr != nil { + return SendMessageResult{}, updateErr } - p.publishMessage(opts.ChatID, result.Message) - p.publishStatus(opts.ChatID, result.Chat.Status, result.Chat.WorkerID) + // Sidebar watch event keeps the chat list in sync. Stream side + // effects are handled by chat:update consumers. p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) - p.signalWake() return result, nil } @@ -1974,28 +1657,6 @@ func resolveSendMessageModelConfigID( return requested, nil } -func resolveQueuedMessageModelConfigID( - ctx context.Context, - store database.Store, - chat database.Chat, - queuedModelConfigID uuid.NullUUID, -) (uuid.UUID, error) { - chatdCtx := chatdModelConfigLookupContext(ctx) - if queuedModelConfigID.Valid && queuedModelConfigID.UUID != uuid.Nil { - if _, err := store.GetChatModelConfigByID(chatdCtx, queuedModelConfigID.UUID); err == nil { - return queuedModelConfigID.UUID, nil - } else if !errors.Is(err, sql.ErrNoRows) { - return uuid.Nil, xerrors.Errorf( - "get queued model config %s: %w", - queuedModelConfigID.UUID, - err, - ) - } - } - - return resolveFallbackModelConfigID(ctx, store, chat.LastModelConfigID) -} - func resolveFallbackModelConfigID( ctx context.Context, store database.Store, @@ -2024,9 +1685,10 @@ func resolveFallbackModelConfigID( return defaultConfig.ID, nil } -// EditMessage marks the old user message as deleted, soft-deletes all -// following messages, inserts a new message with the updated content, -// clears queued messages, and moves the chat into pending status. +// EditMessage replaces an earlier user message and discards the +// active-history suffix through chatstate.EditMessage. Model-config +// override validation and usage-limit admission run in the same +// transaction as the state-machine transition. func (p *Server) EditMessage( ctx context.Context, opts EditMessageOptions, @@ -2040,6 +1702,9 @@ func (p *Server) EditMessage( if len(opts.Content) == 0 { return EditMessageResult{}, xerrors.New("content is required") } + if err := validateChatUserMessageAPIKeyID(opts.APIKeyID); err != nil { + return EditMessageResult{}, err + } content, err := chatprompt.MarshalParts(opts.Content) if err != nil { @@ -2047,60 +1712,44 @@ func (p *Server) EditMessage( } var ( - result EditMessageResult - editedMsg database.ChatMessage + result EditMessageResult + editedMsg database.ChatMessage + editedCutoffT time.Time ) - txErr := p.db.InTx(func(tx database.Store) error { - lockedChat, err := tx.GetChatByIDForUpdate(ctx, opts.ChatID) + machine := p.newChatMachine(opts.ChatID) + err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { - return xerrors.Errorf("lock chat: %w", err) + return xerrors.Errorf("load chat: %w", err) } - if lockedChat.Archived { return ErrChatArchived } - - if limitErr := p.checkUsageLimit(ctx, tx, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { + if limitErr := p.checkUsageLimit(ctx, store, lockedChat.OwnerID, uuid.NullUUID{UUID: lockedChat.OrganizationID, Valid: true}); limitErr != nil { return limitErr } - editedMsg, err = tx.GetChatMessageByID(ctx, opts.EditedMessageID) + // Capture the target message for the post-commit debug + // cleanup hook below. The transition itself revalidates + // chat ownership and user-message constraints. + target, err := store.GetChatMessageByID(ctx, opts.EditedMessageID) if err != nil { if errors.Is(err, sql.ErrNoRows) { return ErrEditedMessageNotFound } return xerrors.Errorf("get edited message: %w", err) } - if editedMsg.ChatID != opts.ChatID { + if target.ChatID != opts.ChatID { return ErrEditedMessageNotFound } - if editedMsg.Role != database.ChatMessageRoleUser { - return ErrEditedMessageNotUser - } + editedMsg = target - // Soft-delete the original message instead of updating in place - // so that usage/cost data is preserved. - err = tx.SoftDeleteChatMessageByID(ctx, opts.EditedMessageID) - if err != nil { - return xerrors.Errorf("soft-delete edited message: %w", err) - } - - // Soft-delete all messages that came after the edited one. - err = tx.SoftDeleteChatMessagesAfterID(ctx, database.SoftDeleteChatMessagesAfterIDParams{ - ChatID: opts.ChatID, - AfterID: opts.EditedMessageID, - }) - if err != nil { - return xerrors.Errorf("soft-delete later chat messages: %w", err) - } - - // Resolve the model for the replacement message. When the - // caller does not specify a model, preserve the original - // message's model so an edit that only changes text keeps - // behaving as before. - messageModelConfigID := editedMsg.ModelConfigID.UUID + // Validate the optional model-config override up front so + // the user sees ErrInvalidModelConfigID instead of a + // foreign-key error from the message-insert path. + var modelOverride uuid.NullUUID if opts.ModelConfigID != uuid.Nil { - if _, err := tx.GetChatModelConfigByID( + if _, err := store.GetChatModelConfigByID( chatdModelConfigLookupContext(ctx), opts.ModelConfigID, ); err != nil { @@ -2117,76 +1766,51 @@ func (p *Server) EditMessage( err, ) } - messageModelConfigID = opts.ModelConfigID + modelOverride = uuid.NullUUID{UUID: opts.ModelConfigID, Valid: true} } - // Insert a new message with the updated content. The - // InsertChatMessages CTE updates chats.last_model_config_id - // when the new message's model differs, so the assistant turn - // that follows picks up the new selection. - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. - ChatID: opts.ChatID, - } - editUserMsg := newUserChatMessage( - opts.APIKeyID, - content, - editedMsg.Visibility, - messageModelConfigID, - chatprompt.CurrentContentVersion, - ) - editUserMsg = editUserMsg.withCreatedBy(opts.CreatedBy) - appendUserChatMessage(&msgParams, editUserMsg) - newMessages, err := insertChatMessageWithStore(ctx, tx, msgParams) - if err != nil { - return xerrors.Errorf("insert replacement message: %w", err) - } - newMessage := newMessages[0] - - err = tx.DeleteAllChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("delete queued messages: %w", err) - } - updatedChat, err := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: opts.ChatID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, + editResult, err := tx.EditMessage(chatstate.EditMessageInput{ + MessageID: opts.EditedMessageID, + CreatedBy: opts.CreatedBy, + Content: content, + ModelConfigIDOverride: modelOverride, + APIKeyID: sql.NullString{String: opts.APIKeyID, Valid: opts.APIKeyID != ""}, }) if err != nil { - return xerrors.Errorf("set chat pending: %w", err) + if errors.Is(err, chatstate.ErrEditedMessageNotUser) { + return ErrEditedMessageNotUser + } + return err } - - result.Message = newMessage - result.Chat = updatedChat + result.Message = editResult.ReplacementMessage + // Capture the post-edit chat inside the same transaction so + // the returned chat and the debug-cleanup cutoff use the + // snapshot bump and updated_at stamped by the transition. + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after edit: %w", err) + } + result.Chat = refreshed + editedCutoffT = refreshed.UpdatedAt return nil - }, nil) - if txErr != nil { - return EditMessageResult{}, txErr + }) + if err != nil { + return EditMessageResult{}, err } - p.publishEditedMessage(opts.ChatID, result.Message) - p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - QueuedMessages: []codersdk.ChatQueuedMessage{}, - }) - p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, - }) - p.publishStatus(opts.ChatID, result.Chat.Status, result.Chat.WorkerID) + // Sidebar watch event keeps the chat list responsive. Stream + // side effects are handled by chat:update consumers. p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) // Editing can race with an interrupted worker still flushing its // final debug writes. Run a short bounded retry loop so we converge // quickly without relying on the much longer stale-finalization // sweep. Source editCutoff from the DB-stamped updated_at returned - // by UpdateChatStatus so the filter uses the same clock that - // FinalizeStale and other DB timestamps use; subtract + // by the post-edit chat row so the filter uses the same clock that + // stamps replacement-turn debug rows; subtract // debugCleanupClockSkew so replica clock drift cannot let the retry - // delete a replacement turn's debug rows (see the constant for the - // full rationale). - editCutoff := result.Chat.UpdatedAt.Add(-debugCleanupClockSkew) + // delete a replacement turn's debug rows. + editCutoff := editedCutoffT.Add(-debugCleanupClockSkew) p.scheduleDebugCleanup( ctx, "failed to delete chat debug rows after edit", @@ -2199,176 +1823,95 @@ func (p *Server) EditMessage( return err }, ) - p.signalWake() return result, nil } -// ArchiveChat archives a chat family and broadcasts deleted events for each -// affected chat so watching clients converge without a full refetch. If the -// target chat is pending or running, it first transitions the chat back to -// waiting so active processing stops before the archive is broadcast. +// ErrArchiveRequiresRootChat is returned by [Server.ArchiveChat] and +// [Server.UnarchiveChat] when the supplied chat is a child chat. +// Archive state changes must always target the root chat so the +// whole family flips together. +var ErrArchiveRequiresRootChat = xerrors.New( + "chat archive state can only be changed on the root chat", +) + +// ArchiveChat archives a root chat and every child in its family +// through the chatstate state machine. The transition is atomic over +// the whole family: either every member is archived or none is. The +// state machine only permits archive from the idle / error execution +// states (W, E0, E1); active members cause a state conflict that the +// HTTP handler maps to a client error. +// +// Child chats must not be archived independently. ArchiveChat +// rejects them with [ErrArchiveRequiresRootChat] so callers cannot +// silently break the parent-implies-child archive invariant. func (p *Server) ArchiveChat(ctx context.Context, chat database.Chat) error { if chat.ID == uuid.Nil { return xerrors.New("chat_id is required") } - - var ( - archivedChats []database.Chat - interruptedChats []database.Chat - ) - if err := p.db.InTx(func(tx database.Store) error { - if _, err := tx.GetChatByIDForUpdate(ctx, chat.ID); err != nil { - return xerrors.Errorf("lock chat for archive: %w", err) - } - - var err error - archivedChats, err = tx.ArchiveChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("archive chat: %w", err) - } - - for i, archivedChat := range archivedChats { - if archivedChat.Status != database.ChatStatusPending && - archivedChat.Status != database.ChatStatusRunning { - continue - } - - updatedChat, updateErr := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: archivedChat.ID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - if updateErr != nil { - return xerrors.Errorf("set archived chat waiting before cleanup: %w", updateErr) - } - archivedChats[i] = updatedChat - interruptedChats = append(interruptedChats, updatedChat) - } - return nil - }, nil); err != nil { - return err + if chat.ParentChatID.Valid { + return ErrArchiveRequiresRootChat } - - for _, interruptedChat := range interruptedChats { - p.publishStatus(interruptedChat.ID, interruptedChat.Status, interruptedChat.WorkerID) - p.publishChatPubsubEvent(interruptedChat, codersdk.ChatWatchEventKindStatusChange, nil) - } - - // Archiving can race with an interrupted worker still flushing its - // final debug writes. Retry a few times so orphaned rows are - // removed quickly instead of waiting for the stale sweeper. Source - // archiveCutoff from the DB-stamped updated_at returned by - // ArchiveChatByID so the filter uses the same clock that stamps - // replacement-turn debug rows; subtract debugCleanupClockSkew so - // replica clock drift cannot let the retry delete a replacement's - // debug rows if an unarchive races ahead (see the constant for the - // full rationale). All archived chats share the transaction-start - // NOW() so any entry's UpdatedAt is equivalent. - if len(archivedChats) > 0 { - archiveCutoff := archivedChats[0].UpdatedAt.Add(-debugCleanupClockSkew) - for _, archivedChat := range archivedChats { - p.scheduleDebugCleanup( - ctx, - "failed to delete chat debug rows after archive", - []slog.Field{slog.F("chat_id", archivedChat.ID)}, - func(cleanupCtx context.Context, debugSvc *chatdebug.Service) error { - _, err := debugSvc.DeleteByChatID(cleanupCtx, archivedChat.ID, archiveCutoff) - return err - }, - ) - } - } - - p.publishChatPubsubEvents(archivedChats, codersdk.ChatWatchEventKindDeleted) - return nil + return p.setChatFamilyArchived(ctx, chat, true, codersdk.ChatWatchEventKindDeleted) } -// ErrChildUnarchiveParentArchived is returned by UnarchiveChat when a -// child unarchive is rejected because the parent is still archived. -// The patchChat handler maps this to a 400 response. -var ErrChildUnarchiveParentArchived = xerrors.New( - "cannot unarchive child chat while parent is archived", -) - -// UnarchiveChat unarchives a chat family and broadcasts created events. -// Root chats cascade through UnarchiveChatByID. Child chats run under -// a row-level lock on the child (GetChatByIDForUpdate) with an -// in-transaction re-read of the parent, returning -// ErrChildUnarchiveParentArchived when the parent is archived and a -// no-op when the child is already active. -// -// The child is locked before the parent is read to avoid deadlocking -// with a concurrent ArchiveChatByID cascade, which visits child rows -// before the parent. +// UnarchiveChat unarchives a root chat and every child in its family +// through the chatstate state machine. Like ArchiveChat the cascade +// is atomic; ChildChat unarchive attempts are rejected with +// [ErrArchiveRequiresRootChat]. func (p *Server) UnarchiveChat(ctx context.Context, chat database.Chat) error { if chat.ID == uuid.Nil { return xerrors.New("chat_id is required") } + if chat.ParentChatID.Valid { + return ErrArchiveRequiresRootChat + } + return p.setChatFamilyArchived(ctx, chat, false, codersdk.ChatWatchEventKindCreated) +} - if !chat.ParentChatID.Valid { - return p.applyChatLifecycleTransition( - ctx, - chat.ID, - "unarchive", - codersdk.ChatWatchEventKindCreated, - p.db.UnarchiveChatByID, - ) +// setChatFamilyArchived applies SetArchived(archived) to every chat +// in chat's family through chatstate. The transaction-captured +// family rows feed the post-commit debug cleanup and sidebar watch +// events. Callers must only invoke this for root chats. +// +//nolint:revive // Existing API takes the target archive state as a boolean. +func (p *Server) setChatFamilyArchived( + ctx context.Context, + chat database.Chat, + archived bool, + watchKind codersdk.ChatWatchEventKind, +) error { + if chat.ID == uuid.Nil { + return xerrors.New("chat_id is required") + } + if chat.ParentChatID.Valid { + return ErrArchiveRequiresRootChat } - var updated []database.Chat - if err := p.db.InTx(func(tx database.Store) error { - locked, err := tx.GetChatByIDForUpdate(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("lock child for unarchive: %w", err) - } - if !locked.Archived { - // Already unarchived by a concurrent caller; idempotent no-op. - return nil - } - parent, err := tx.GetChatByID(ctx, chat.ParentChatID.UUID) - if err != nil { - return xerrors.Errorf("load parent chat: %w", err) - } - if parent.Archived { - return ErrChildUnarchiveParentArchived - } - updated, err = tx.UnarchiveChatByID(ctx, chat.ID) - if err != nil { - return xerrors.Errorf("unarchive child chat: %w", err) - } - return nil - }, nil); err != nil { - if errors.Is(err, ErrChildUnarchiveParentArchived) { - return ErrChildUnarchiveParentArchived - } + familyChats, err := chatstate.SetFamilyArchived( + ctx, + p.db, + p.pubsub, + chatstate.SetFamilyArchivedInput{ + RootID: chat.ID, + Archived: archived, + }, + ) + if err != nil { return err } - p.publishChatPubsubEvents(updated, codersdk.ChatWatchEventKindCreated) - return nil -} - -func (p *Server) applyChatLifecycleTransition( - ctx context.Context, - chatID uuid.UUID, - action string, - kind codersdk.ChatWatchEventKind, - transition func(context.Context, uuid.UUID) ([]database.Chat, error), -) error { - updatedChats, err := transition(ctx, chatID) - if err != nil { - return xerrors.Errorf("%s chat: %w", action, err) + if archived { + p.scheduleArchiveDebugCleanup(ctx, familyChats) } - p.publishChatPubsubEvents(updatedChats, kind) + p.publishChatPubsubEvents(familyChats, watchKind) return nil } -// DeleteQueued removes a queued user message and publishes the queue update. +// DeleteQueued removes a queued user message through the chatstate +// state machine. Stream side effects are handled by chat:update +// consumers. func (p *Server) DeleteQueued( ctx context.Context, chatID uuid.UUID, @@ -2378,61 +1921,22 @@ func (p *Server) DeleteQueued( return xerrors.New("chat_id is required") } - var queuedMessages []database.ChatQueuedMessage - var queueLoadedOK bool - - txErr := p.db.InTx(func(tx database.Store) error { - // Lock the chat row to prevent processChat from - // auto-promoting a message the user intended to delete. - if _, err := tx.GetChatByIDForUpdate(ctx, chatID); err != nil { - return xerrors.Errorf("lock chat: %w", err) - } - - err := tx.DeleteChatQueuedMessage(ctx, database.DeleteChatQueuedMessageParams{ - ID: queuedMessageID, - ChatID: chatID, + machine := p.newChatMachine(chatID) + err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.DeleteQueuedMessage(chatstate.DeleteQueuedMessageInput{ + QueuedMessageID: queuedMessageID, }) - if err != nil { - return xerrors.Errorf("delete queued message: %w", err) - } - - var err2 error - queuedMessages, err2 = tx.GetChatQueuedMessages(ctx, chatID) - if err2 != nil { - p.logger.Warn(ctx, "failed to load queued messages after delete", - slog.F("chat_id", chatID), - slog.F("queued_message_id", queuedMessageID), - slog.Error(err2), - ) - // Non-fatal: the delete succeeded, so we still commit. - return nil - } - queueLoadedOK = true - - return nil - }, nil) - if txErr != nil { - return txErr - } - - if queueLoadedOK { - p.publishEvent(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - QueuedMessages: db2sdk.ChatQueuedMessages(queuedMessages), - }) - } - // Always notify subscribers so they can re-fetch, even if we - // failed to load the updated queue payload above. - p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, + return err }) - return nil + return err } -// PromoteQueued promotes a queued message into chat history. On a -// running chat with a fresh worker heartbeat the promote is deferred -// to the worker's persist+auto-promote so partial assistant output -// is not lost; otherwise it inserts the user message synchronously. +// PromoteQueued promotes a queued message through the chatstate state +// machine. From running / interrupting states the state machine +// transitions the chat to `interrupting` so the worker can drain the +// in-flight generation before promoting; from idle / error / requires +// action states it inserts the user message into history +// synchronously. func (p *Server) PromoteQueued( ctx context.Context, opts PromoteQueuedOptions, @@ -2442,187 +1946,47 @@ func (p *Server) PromoteQueued( } var ( - result PromoteQueuedResult - promoted database.ChatMessage - updatedChat database.Chat - remainingQueue []database.ChatQueuedMessage - deferred bool - syntheticResults []database.ChatMessage + result PromoteQueuedResult + refreshChat database.Chat + refreshedOK bool ) - - txErr := p.db.InTx(func(tx database.Store) error { - lockedChat, err := tx.GetChatByIDForUpdate(ctx, opts.ChatID) + machine := p.newChatMachine(opts.ChatID) + updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + lockedChat, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { - return xerrors.Errorf("lock chat: %w", err) + return xerrors.Errorf("load chat: %w", err) } - if lockedChat.Archived { return ErrChatArchived } - queuedMessages, err := tx.GetChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("get queued messages: %w", err) - } - - var ( - targetContent json.RawMessage - targetModelConfigID uuid.NullUUID - targetAPIKeyID sql.NullString - found bool - ) - for _, qm := range queuedMessages { - if qm.ID == opts.QueuedMessageID { - targetContent = qm.Content - targetModelConfigID = qm.ModelConfigID - targetAPIKeyID = qm.APIKeyID - found = true - break - } - } - if !found { - return xerrors.Errorf("queued message %d not found in chat %s", opts.QueuedMessageID, opts.ChatID) - } - - // Setting pending would trip persistStep's ownership guard - // and drop the worker's partial output. Set waiting and - // reorder the queued row so the worker's auto-promote picks - // it up after the persist. - heartbeatFresh := lockedChat.HeartbeatAt.Valid && - p.clock.Now().Sub(lockedChat.HeartbeatAt.Time) < p.inFlightChatStaleAfter - if lockedChat.Status == database.ChatStatusRunning && heartbeatFresh { - rowsAffected, err := tx.ReorderChatQueuedMessageToFront(ctx, database.ReorderChatQueuedMessageToFrontParams{ - ChatID: opts.ChatID, - TargetID: opts.QueuedMessageID, - }) - if err != nil { - return xerrors.Errorf("reorder queued message to front: %w", err) - } - // Defensive guard against a future non-chat-locked - // queue mutator. The found check above makes this a - // no-op on the current code path. - if rowsAffected != 1 { - return xerrors.Errorf("reorder queued message to front affected %d rows, want 1", rowsAffected) - } - updatedChat, err = tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: opts.ChatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - if err != nil { - return xerrors.Errorf("set chat to waiting for deferred promote: %w", err) - } - remainingQueue, err = tx.GetChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("get remaining queue after reorder: %w", err) - } - deferred = true - return nil - } - - effectiveModelConfigID, err := resolveQueuedMessageModelConfigID( - ctx, - tx, - lockedChat, - targetModelConfigID, - ) - if err != nil { - return err - } - - // Without synthetic results, the next turn would carry - // unresolved tool_call parts; the LLM API rejects this and the - // chat dead-ends in error. - if lockedChat.Status == database.ChatStatusRequiresAction { - inserted, err := insertSyntheticToolResultsTx( - ctx, tx, lockedChat, - "Tool execution interrupted by queued message promotion", - ) - if err != nil { - return xerrors.Errorf("insert synthetic tool results: %w", err) - } - syntheticResults = inserted - } - - err = tx.DeleteChatQueuedMessage(ctx, database.DeleteChatQueuedMessageParams{ - ID: opts.QueuedMessageID, - ChatID: opts.ChatID, + promoteResult, err := tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: opts.QueuedMessageID, }) - if err != nil { - return xerrors.Errorf("delete queued message: %w", err) - } - - promoted, updatedChat, err = insertUserMessageAndSetPending( - ctx, - tx, - lockedChat, - effectiveModelConfigID, - pqtype.NullRawMessage{ - RawMessage: targetContent, - Valid: len(targetContent) > 0, - }, - opts.CreatedBy, - targetAPIKeyID.String, - ) if err != nil { return err } - - remainingQueue, err = tx.GetChatQueuedMessages(ctx, opts.ChatID) - if err != nil { - return xerrors.Errorf("get remaining queue: %w", err) + if promoteResult.InsertedMessage != nil { + result.PromotedMessage = *promoteResult.InsertedMessage } - result.PromotedMessage = promoted - + // Capture the chat inside the transaction so the watch event + // published below uses the snapshot bump and status change + // produced by the transition itself. + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after promote: %w", err) + } + refreshChat = refreshed + refreshedOK = true return nil - }, nil) - if txErr != nil { - return PromoteQueuedResult{}, txErr - } - - if deferred { - // Skip publishMessage and signalWake: there is no synchronous - // user message yet, and the active worker's interrupt path - // signals its own auto-promote follow-up. - p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - QueuedMessages: db2sdk.ChatQueuedMessages(remainingQueue), - }) - p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, - }) - p.publishStatus(opts.ChatID, updatedChat.Status, updatedChat.WorkerID) - p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) - return result, nil - } - - p.publishEvent(opts.ChatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - QueuedMessages: db2sdk.ChatQueuedMessages(remainingQueue), }) - p.publishChatStreamNotify(opts.ChatID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, - }) - // Publish synth rows before the user message so live viewers - // see the interruption inline. - for _, msg := range syntheticResults { - p.publishMessage(opts.ChatID, msg) + if updateErr != nil { + return PromoteQueuedResult{}, updateErr } - p.publishMessage(opts.ChatID, promoted) - p.publishStatus(opts.ChatID, updatedChat.Status, updatedChat.WorkerID) - p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) - // Marker for ENG-2645: confirms post-TX publishes ran. - p.logger.Debug(ctx, "promote queued completed", - slog.F("chat_id", opts.ChatID), - slog.F("promoted_id", promoted.ID), - slog.F("synthetic_count", len(syntheticResults)), - slog.F("status", updatedChat.Status), - ) - p.signalWake() + if refreshedOK { + p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) + } return result, nil } @@ -2664,251 +2028,198 @@ func (e *ToolResultStatusConflictError) Error() string { } // SubmitToolResults validates and persists client-provided tool -// results, transitions the chat to pending, and wakes the run -// loop. The caller is responsible for the fast-path status check; -// this method performs an authoritative re-check under a row lock. +// results, returning the chat to running through the chatstate state +// machine. Validation runs inside the same transaction as the +// transition so the assistant message and pending tool calls cannot +// drift between reads. func (p *Server) SubmitToolResults( ctx context.Context, opts SubmitToolResultsOptions, ) error { - dynamicToolNames, err := parseDynamicToolNames(pqtype.NullRawMessage{ - RawMessage: opts.DynamicTools, - Valid: len(opts.DynamicTools) > 0, - }) - if err != nil { - return xerrors.Errorf("parse chat dynamic tools: %w", err) - } - - // The GetLastChatMessageByRole lookup and all subsequent - // validation and persistence run inside a single transaction - // so the assistant message cannot change between reads. - var statusConflict *ToolResultStatusConflictError - txErr := p.db.InTx(func(tx database.Store) error { - // Authoritative status check under row lock. - locked, lockErr := tx.GetChatByIDForUpdate(ctx, opts.ChatID) - if lockErr != nil { - return xerrors.Errorf("lock chat for update: %w", lockErr) + var ( + statusConflict *ToolResultStatusConflictError + refreshChat database.Chat + refreshedOK bool + ) + machine := p.newChatMachine(opts.ChatID) + updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) } if locked.Archived { return ErrChatArchived } - if locked.Status != database.ChatStatusRequiresAction { - statusConflict = &ToolResultStatusConflictError{ - ActualStatus: locked.Status, + + toolResults := make([]chatstate.ToolResultInput, 0, len(opts.Results)) + for _, r := range opts.Results { + toolResults = append(toolResults, chatstate.ToolResultInput{ + ToolCallID: r.ToolCallID, + Output: r.Output, + IsError: r.IsError, + }) + } + modelConfigID := opts.ModelConfigID + if modelConfigID == uuid.Nil { + modelConfigID = locked.LastModelConfigID + } + if _, err := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ + CreatedBy: opts.UserID, + ModelConfigID: modelConfigID, + Results: toolResults, + }); err != nil { + if !errors.Is(err, chatstate.ErrInvalidState) && + locked.Status != database.ChatStatusRequiresAction && + errors.Is(err, chatstate.ErrTransitionNotAllowed) { + statusConflict = &ToolResultStatusConflictError{ + ActualStatus: locked.Status, + } + return statusConflict } + return xerrors.Errorf("complete requires action: %w", err) + } + // Capture the chat inside the transaction so the watch event + // uses the snapshot bump and status change produced by the + // transition itself. + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after tool results: %w", err) + } + refreshChat = refreshed + refreshedOK = true + return nil + }) + if updateErr != nil { + if statusConflict != nil { return statusConflict } - - // Get the last assistant message inside the transaction - // for consistency with the row lock above. - lastAssistant, err := tx.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ - ChatID: opts.ChatID, - Role: database.ChatMessageRoleAssistant, - }) - if err != nil { - return xerrors.Errorf("get last assistant message: %w", err) - } - - // Collect tool-call IDs that already have results. - // When a dynamic tool name collides with a built-in, - // the chatloop executes it as a built-in and persists - // the result. Those calls must not count as pending. - afterMsgs, afterErr := tx.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: opts.ChatID, - AfterID: lastAssistant.ID, - }) - if afterErr != nil { - return xerrors.Errorf("get messages after assistant: %w", afterErr) - } - handledCallIDs := make(map[string]bool) - for _, msg := range afterMsgs { - if msg.Role != database.ChatMessageRoleTool { - continue - } - msgParts, msgParseErr := chatprompt.ParseContent(msg) - if msgParseErr != nil { - continue - } - for _, mp := range msgParts { - if mp.Type == codersdk.ChatMessagePartTypeToolResult { - handledCallIDs[mp.ToolCallID] = true - } - } - } - - // Extract pending dynamic tool-call IDs, skipping any - // that were already handled by the chatloop. - pendingCallIDs := make(map[string]bool) - toolCallIDToName := make(map[string]string) - parts, parseErr := chatprompt.ParseContent(lastAssistant) - if parseErr != nil { - return xerrors.Errorf("parse assistant message: %w", parseErr) - } - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolCall && - dynamicToolNames[part.ToolName] && - !handledCallIDs[part.ToolCallID] { - pendingCallIDs[part.ToolCallID] = true - toolCallIDToName[part.ToolCallID] = part.ToolName - } - } - - // Validate submitted results match pending calls exactly. - submittedIDs := make(map[string]bool, len(opts.Results)) - for _, result := range opts.Results { - if submittedIDs[result.ToolCallID] { - return &ToolResultValidationError{ - Message: "Duplicate tool_call_id in results.", - Detail: fmt.Sprintf("Duplicate tool call ID %q.", result.ToolCallID), - } - } - submittedIDs[result.ToolCallID] = true - } - for id := range pendingCallIDs { - if !submittedIDs[id] { - return &ToolResultValidationError{ - Message: "Missing tool result.", - Detail: fmt.Sprintf("Missing result for tool call %q.", id), - } - } - } - for id := range submittedIDs { - if !pendingCallIDs[id] { - return &ToolResultValidationError{ - Message: "Unexpected tool result.", - Detail: fmt.Sprintf("No pending tool call with ID %q.", id), - } - } - } - - // Marshal each tool result into a separate message row. - resultContents := make([]pqtype.NullRawMessage, 0, len(opts.Results)) - for _, result := range opts.Results { - if !json.Valid(result.Output) { - return &ToolResultValidationError{ - Message: "Tool result output must be valid JSON.", - Detail: fmt.Sprintf("Output for tool call %q is not valid JSON.", result.ToolCallID), - } - } - part := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: result.ToolCallID, - ToolName: toolCallIDToName[result.ToolCallID], - Result: result.Output, - IsError: result.IsError, - } - marshaled, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part}) - if marshalErr != nil { - return xerrors.Errorf("marshal tool result: %w", marshalErr) - } - resultContents = append(resultContents, marshaled) - } - - // Insert tool-result messages. - n := len(resultContents) - params := database.InsertChatMessagesParams{ - ChatID: opts.ChatID, - CreatedBy: make([]uuid.UUID, n), - APIKeyID: make([]string, n), - ModelConfigID: make([]uuid.UUID, n), - Role: make([]database.ChatMessageRole, n), - Content: make([]string, n), - ContentVersion: make([]int16, n), - Visibility: make([]database.ChatMessageVisibility, n), - InputTokens: make([]int64, n), - OutputTokens: make([]int64, n), - TotalTokens: make([]int64, n), - ReasoningTokens: make([]int64, n), - CacheCreationTokens: make([]int64, n), - CacheReadTokens: make([]int64, n), - ContextLimit: make([]int64, n), - Compressed: make([]bool, n), - TotalCostMicros: make([]int64, n), - RuntimeMs: make([]int64, n), - ProviderResponseID: make([]string, n), - } - for i, rc := range resultContents { - params.CreatedBy[i] = opts.UserID - params.ModelConfigID[i] = opts.ModelConfigID - params.Role[i] = database.ChatMessageRoleTool - params.Content[i] = string(rc.RawMessage) - params.ContentVersion[i] = chatprompt.CurrentContentVersion - params.Visibility[i] = database.ChatMessageVisibilityBoth - } - if _, insertErr := tx.InsertChatMessages(ctx, params); insertErr != nil { - return xerrors.Errorf("insert tool results: %w", insertErr) - } - - // Transition chat to pending. - if _, updateErr := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: opts.ChatID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }); updateErr != nil { - return xerrors.Errorf("update chat status: %w", updateErr) - } - - return nil - }, nil) - if txErr != nil { - return txErr + return translateToolResultValidationError(updateErr) } - // Wake the chatd run loop so it processes the chat immediately. - p.signalWake() + if refreshedOK { + p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) + } return nil } -// InterruptChat interrupts execution, sets waiting status, and broadcasts status updates. +// translateToolResultValidationError converts a chatstate tool-result +// validation error into the legacy chatd.ToolResultValidationError +// shape so HTTP handlers preserve their existing response detail. If +// err is not a tool-result validation error, it is returned +// unchanged. +func translateToolResultValidationError(err error) error { + var v *chatstate.ToolResultValidationError + if !errors.As(err, &v) { + return err + } + switch { + case xerrors.Is(v, chatstate.ErrToolResultDuplicate): + return &ToolResultValidationError{ + Message: "Duplicate tool_call_id in results.", + Detail: fmt.Sprintf("Duplicate tool call ID %q.", v.ToolCallID), + } + case xerrors.Is(v, chatstate.ErrToolResultMissing): + return &ToolResultValidationError{ + Message: "Missing tool result.", + Detail: fmt.Sprintf("Missing result for tool call %q.", v.ToolCallID), + } + case xerrors.Is(v, chatstate.ErrToolResultUnexpected): + return &ToolResultValidationError{ + Message: "Unexpected tool result.", + Detail: fmt.Sprintf("No pending tool call with ID %q.", v.ToolCallID), + } + case xerrors.Is(v, chatstate.ErrToolResultInvalidJSON): + return &ToolResultValidationError{ + Message: "Tool result output must be valid JSON.", + Detail: fmt.Sprintf("Output for tool call %q is not valid JSON.", v.ToolCallID), + } + default: + return err + } +} + +// InterruptChat interrupts execution through the chatstate.Interrupt +// transition. Active runs land in `interrupting`; requires-action +// chats synthesize cancellation messages and return to running. +// +// Returns the post-transition chat and an error so callers can map +// state conflicts deliberately. Idle chats return a +// chatstate.ErrTransitionNotAllowed wrapper. func (p *Server) InterruptChat( ctx context.Context, chat database.Chat, -) database.Chat { +) (database.Chat, error) { if chat.ID == uuid.Nil { - return chat + return chat, xerrors.New("chat_id is required") } - // If the chat is in requires_action, insert synthetic error - // tool-result messages for each pending dynamic tool call - // before transitioning to waiting. Without this, the LLM - // would see unmatched tool-call parts on the next run. - if chat.Status == database.ChatStatusRequiresAction { - if txErr := p.db.InTx(func(tx database.Store) error { - locked, lockErr := tx.GetChatByIDForUpdate(ctx, chat.ID) - if lockErr != nil { - return xerrors.Errorf("lock chat for interrupt: %w", lockErr) - } - // Another request may have already transitioned - // the chat (e.g. SubmitToolResults committed - // between our snapshot and this lock). - if locked.Status != database.ChatStatusRequiresAction { - return nil - } - _, err := insertSyntheticToolResultsTx(ctx, tx, locked, "Tool execution interrupted by user") + var refreshed database.Chat + machine := p.newChatMachine(chat.ID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := tx.Interrupt(chatstate.InterruptInput{ + Reason: "Tool execution interrupted by user", + }); err != nil { return err - }, nil); txErr != nil { - p.logger.Error(ctx, "failed to insert synthetic tool results during interrupt", - slog.F("chat_id", chat.ID), - slog.Error(txErr), - ) - // Fall through — still try to set waiting status. } + // Capture the post-interrupt chat inside the transaction so + // the returned chat and the watch event reflect the snapshot + // bump and status change produced by the transition itself. + latest, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("reload chat after interrupt: %w", err) + } + refreshed = latest + return nil + }) + if err != nil { + return chat, err } - // Debug runs are finalized in the execution path when the owning - // goroutine observes cancellation, so we do not mutate debug state here. - updatedChat, err := p.setChatWaiting(ctx, chat.ID) - if err != nil { - p.logger.Error(ctx, "failed to mark chat as waiting", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - return chat + p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) + return refreshed, nil +} + +// ReconcileInvalidStateChat recovers a chat stuck in an invalid +// execution-state combination by running the +// chatstate.ReconcileInvalidState transition. The chat lands in an +// error state (E0/E1); queued messages are preserved and pending +// dynamic-tool calls are closed with synthetic cancellations. +// +// Returns the post-transition chat. When the chat is not actually in an +// invalid state the transition returns a wrapped +// chatstate.ErrTransitionNotAllowed; a missing chat returns +// chatstate.ErrChatNotFound. Callers map these to deliberate HTTP +// responses. +func (p *Server) ReconcileInvalidStateChat( + ctx context.Context, + chat database.Chat, +) (database.Chat, error) { + if chat.ID == uuid.Nil { + return chat, xerrors.New("chat_id is required") } - return updatedChat + + var refreshed database.Chat + machine := p.newChatMachine(chat.ID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{}); err != nil { + return err + } + // Capture the post-reconcile chat inside the transaction so + // the returned chat and the watch event reflect the snapshot + // bump and status change produced by the transition itself. + latest, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("reload chat after reconcile: %w", err) + } + refreshed = latest + return nil + }) + if err != nil { + return chat, err + } + + p.publishChatPubsubEvent(refreshed, codersdk.ChatWatchEventKindStatusChange, nil) + return refreshed, nil } const manualTitleMessageWindowLimit = 50 @@ -2932,6 +2243,39 @@ type manualTitleGenerationError struct { activeAPIKeyID string } +// generatedChatTitle carries the title produced by the detached +// automatic title-generation goroutine. maybeGenerateChatTitle stores +// the generated title here so tests can observe it without a database +// read; the title_change pubsub event it publishes remains the source of +// truth for clients. +type generatedChatTitle struct { + mu sync.RWMutex + title string +} + +func (t *generatedChatTitle) Store(title string) { + if t == nil || title == "" { + return + } + + t.mu.Lock() + t.title = title + t.mu.Unlock() +} + +func (t *generatedChatTitle) Load() (string, bool) { + if t == nil { + return "", false + } + + t.mu.RLock() + defer t.mu.RUnlock() + if t.title == "" { + return "", false + } + return t.title, true +} + func (e *manualTitleGenerationError) Error() string { return e.cause.Error() } @@ -3509,107 +2853,6 @@ func deriveChatDebugSeed(messages []database.ChatMessage) ( return triggerMessageID, historyTipMessageID, triggerLabel } -func prepareChatTurnDebugRun( - ctx context.Context, - logger slog.Logger, - chat database.Chat, - modelConfig database.ChatModelConfig, - debugSvc *chatdebug.Service, - debugProvider string, - debugModel string, - triggerMessageID int64, - historyTipMessageID int64, - triggerLabel string, -) (context.Context, func(error, any)) { - finishDebugRun := func(error, any) {} - if debugSvc == nil { - return ctx, finishDebugRun - } - - seedSummary := chatdebug.SeedSummary( - chatdebug.TruncateLabel(triggerLabel, chatdebug.MaxLabelLength), - ) - rootChatID := uuid.Nil - if chat.RootChatID.Valid { - rootChatID = chat.RootChatID.UUID - } - parentChatID := uuid.Nil - if chat.ParentChatID.Valid { - parentChatID = chat.ParentChatID.UUID - } - - // Debug instrumentation must never block the user turn. Detach - // from the chat-processing context and bound the insert so a slow - // or locked DB makes debug logging degrade silently rather than - // stalling chatloop.Run. Matches the pattern used by - // prepareManualTitleDebugRun. - createRunCtx, createRunCancel := context.WithTimeout( - context.WithoutCancel(ctx), debugCreateRunTimeout, - ) - run, createRunErr := debugSvc.CreateRun(createRunCtx, chatdebug.CreateRunParams{ - ChatID: chat.ID, - RootChatID: rootChatID, - ParentChatID: parentChatID, - ModelConfigID: modelConfig.ID, - TriggerMessageID: triggerMessageID, - HistoryTipMessageID: historyTipMessageID, - Kind: chatdebug.KindChatTurn, - Status: chatdebug.StatusInProgress, - Provider: debugProvider, - Model: debugModel, - Summary: seedSummary, - }) - createRunCancel() - if createRunErr != nil { - logger.Warn(ctx, "failed to create chat debug run", - slog.F("chat_id", chat.ID), - slog.Error(createRunErr), - ) - return ctx, finishDebugRun - } - - runCtx := chatdebug.ContextWithRun(ctx, &chatdebug.RunContext{ - RunID: run.ID, - ChatID: chat.ID, - RootChatID: rootChatID, - ParentChatID: parentChatID, - ModelConfigID: modelConfig.ID, - TriggerMessageID: triggerMessageID, - HistoryTipMessageID: historyTipMessageID, - Kind: chatdebug.KindChatTurn, - Provider: debugProvider, - Model: debugModel, - }) - finishDebugRun = func(loopErr error, panicValue any) { - status := chatdebug.ClassifyError(loopErr) - switch { - case panicValue != nil: - status = chatdebug.StatusError - case errors.Is(loopErr, chatloop.ErrInterrupted): - status = chatdebug.StatusInterrupted - case errors.Is(loopErr, chatloop.ErrDynamicToolCall): - // Dynamic tool calls are a successful pause; the run completed - // its model round-trip. - status = chatdebug.StatusCompleted - } - - if finalizeErr := debugSvc.FinalizeRun(runCtx, chatdebug.FinalizeRunParams{ - RunID: run.ID, - ChatID: chat.ID, - Status: status, - SeedSummary: seedSummary, - }); finalizeErr != nil { - logger.Warn(ctx, "failed to finalize chat debug run", - slog.F("chat_id", chat.ID), - slog.F("run_id", run.ID), - slog.Error(finalizeErr), - ) - } - } - - return runCtx, finishDebugRun -} - func (p *Server) resolveManualTitleModel( ctx context.Context, store database.Store, @@ -3853,71 +3096,6 @@ func recordManualTitleUsage( return updatedChat, nil } -// RefreshStatus loads the latest chat status and publishes it to stream subscribers. -func (p *Server) RefreshStatus(ctx context.Context, chatID uuid.UUID) error { - if chatID == uuid.Nil { - return xerrors.New("chat_id is required") - } - - chat, err := p.db.GetChatByID(ctx, chatID) - if err != nil { - return xerrors.Errorf("get chat: %w", err) - } - - p.publishStatus(chat.ID, chat.Status, chat.WorkerID) - return nil -} - -func (p *Server) setChatWaiting(ctx context.Context, chatID uuid.UUID) (database.Chat, error) { - var updatedChat database.Chat - err := p.db.InTx(func(tx database.Store) error { - locked, lockErr := tx.GetChatByIDForUpdate(ctx, chatID) - if lockErr != nil { - return xerrors.Errorf("lock chat for waiting: %w", lockErr) - } - // If the chat has already transitioned to pending (e.g. - // SendMessage with interrupt behavior), don't overwrite - // it — the pending status takes priority so the new - // message gets processed. - if locked.Status == database.ChatStatusPending { - updatedChat = locked - return nil - } - var updateErr error - updatedChat, updateErr = tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - return updateErr - }, nil) - if err != nil { - return database.Chat{}, err - } - p.publishStatus(chatID, updatedChat.Status, updatedChat.WorkerID) - p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) - return updatedChat, nil -} - -func insertChatMessageWithStore( - ctx context.Context, - store database.Store, - params database.InsertChatMessagesParams, -) ([]database.ChatMessage, error) { - messages, err := store.InsertChatMessages(ctx, params) - if err != nil { - return nil, xerrors.Errorf("insert chat message: %w", err) - } - return messages, nil -} - -// chatMessage is the base message type for batch inserts. Use directly -// only for non-user messages; for user messages, use userChatMessage. -// For nullable UUID fields (ModelConfigID, CreatedBy), use uuid.Nil to -// represent NULL. For nullable int64 fields, use 0 to represent NULL. type chatMessage struct { role database.ChatMessageRole content pqtype.NullRawMessage @@ -3938,8 +3116,6 @@ type chatMessage struct { providerResponseID string } -// userChatMessage wraps chatMessage with a required apiKeyID so that -// omitting it for user messages is a compile error, not a silent data bug. type userChatMessage struct { chatMessage apiKeyID string @@ -3950,11 +3126,6 @@ func (m userChatMessage) withCreatedBy(id uuid.UUID) userChatMessage { return m } -func (m userChatMessage) withCompressed() userChatMessage { - m.chatMessage = m.chatMessage.withCompressed() - return m -} - func newChatMessage( role database.ChatMessageRole, content pqtype.NullRawMessage, @@ -3971,8 +3142,6 @@ func newChatMessage( } } -// newUserChatMessage creates a user message. apiKeyID is required so -// that forgetting it is a compile error rather than a silent data bug. func newUserChatMessage( apiKeyID string, content pqtype.NullRawMessage, @@ -3997,47 +3166,6 @@ func (m chatMessage) withCreatedBy(id uuid.UUID) chatMessage { return m } -func (m chatMessage) withCompressed() chatMessage { - m.compressed = true - return m -} - -func (m chatMessage) withUsage( - inputTokens, outputTokens, totalTokens, reasoningTokens, - cacheCreationTokens, cacheReadTokens int64, -) chatMessage { - m.inputTokens = inputTokens - m.outputTokens = outputTokens - m.totalTokens = totalTokens - m.reasoningTokens = reasoningTokens - m.cacheCreationTokens = cacheCreationTokens - m.cacheReadTokens = cacheReadTokens - return m -} - -func (m chatMessage) withContextLimit(limit int64) chatMessage { - m.contextLimit = limit - return m -} - -func (m chatMessage) withTotalCostMicros(cost int64) chatMessage { - m.totalCostMicros = cost - return m -} - -func (m chatMessage) withRuntimeMs(ms int64) chatMessage { - m.runtimeMs = ms - return m -} - -func (m chatMessage) withProviderResponseID(id string) chatMessage { - m.providerResponseID = id - return m -} - -// appendMessageFields writes all chatMessage fields into the batch insert -// params. apiKeyID is explicit so non-user messages always get "" while -// user messages carry the caller's key for AI Gateway routing. func appendMessageFields( params *database.InsertChatMessagesParams, msg chatMessage, @@ -4063,27 +3191,45 @@ func appendMessageFields( params.ProviderResponseID = append(params.ProviderResponseID, msg.providerResponseID) } -// appendChatMessage appends a non-user message to the batch insert params. -func appendChatMessage( - params *database.InsertChatMessagesParams, - msg chatMessage, -) { +func appendChatMessage(params *database.InsertChatMessagesParams, msg chatMessage) { if msg.role == database.ChatMessageRoleUser { panic("developer error: use appendUserChatMessage for user-role messages") } appendMessageFields(params, msg, "") } -// appendUserChatMessage inserts a user message with its apiKeyID preserved. -func appendUserChatMessage( - params *database.InsertChatMessagesParams, - msg userChatMessage, -) { +func appendUserChatMessage(params *database.InsertChatMessagesParams, msg userChatMessage) { appendMessageFields(params, msg.chatMessage, msg.apiKeyID) } // BuildSingleUserChatMessageInsertParams creates batch insert params for // one user message, requiring an apiKeyID for AI Gateway attribution. +// BuildSingleChatMessageInsertParams creates batch insert params for one +// non-user message using the shared chat message builder. +func BuildSingleChatMessageInsertParams( + chatID uuid.UUID, + role database.ChatMessageRole, + content pqtype.NullRawMessage, + visibility database.ChatMessageVisibility, + modelConfigID uuid.UUID, + contentVersion int16, + createdBy uuid.UUID, +) database.InsertChatMessagesParams { + params := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. + ChatID: chatID, + } + msg := newChatMessage(role, content, visibility, modelConfigID, contentVersion) + if createdBy != uuid.Nil { + msg = msg.withCreatedBy(createdBy) + } + if role == database.ChatMessageRoleUser { + appendMessageFields(¶ms, msg, "") + } else { + appendChatMessage(¶ms, msg) + } + return params +} + func BuildSingleUserChatMessageInsertParams( chatID uuid.UUID, apiKeyID string, @@ -4104,79 +3250,14 @@ func BuildSingleUserChatMessageInsertParams( return params } -// insertUserMessageAndSetPending inserts a user message, transitions the -// chat to pending when needed, and returns the refreshed chat row. -func insertUserMessageAndSetPending( - ctx context.Context, - store database.Store, - lockedChat database.Chat, - modelConfigID uuid.UUID, - content pqtype.NullRawMessage, - createdBy uuid.UUID, - apiKeyID string, -) (database.ChatMessage, database.Chat, error) { - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. - ChatID: lockedChat.ID, - } - insertUserMsg := newUserChatMessage( - apiKeyID, - content, - database.ChatMessageVisibilityBoth, - modelConfigID, - chatprompt.CurrentContentVersion, - ) - insertUserMsg = insertUserMsg.withCreatedBy(createdBy) - appendUserChatMessage(&msgParams, insertUserMsg) - messages, err := insertChatMessageWithStore(ctx, store, msgParams) - if err != nil { - return database.ChatMessage{}, database.Chat{}, err - } - message := messages[0] - - if lockedChat.Status == database.ChatStatusPending { - if modelConfigID == uuid.Nil || lockedChat.LastModelConfigID == modelConfigID { - return message, lockedChat, nil - } - // The InsertChatMessages CTE updates chats.last_model_config_id when - // the message's model config differs. Reload to surface that change. - updatedChat, err := store.GetChatByID(ctx, lockedChat.ID) - if err != nil { - return database.ChatMessage{}, database.Chat{}, xerrors.Errorf("get chat after model config update: %w", err) - } - return message, updatedChat, nil - } - - updatedChat, err := store.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: lockedChat.ID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - if err != nil { - return database.ChatMessage{}, database.Chat{}, xerrors.Errorf("set chat pending: %w", err) - } - return message, updatedChat, nil -} - -// shouldQueueUserMessage reports whether a user message should be -// queued while a chat is active. -func shouldQueueUserMessage(status database.ChatStatus) bool { - switch status { - case database.ChatStatusRunning, database.ChatStatusPending, database.ChatStatusRequiresAction: - return true - default: - return false - } -} - // Config configures a chat processor. type Config struct { - Logger slog.Logger - Database database.Store - ReplicaID uuid.UUID - SubscribeFn SubscribeFn + Logger slog.Logger + Database database.Store + ReplicaID uuid.UUID + // StreamPartsDialer dials remote stream parts. Nil uses the local + // in-process channel dialer for every stream. + StreamPartsDialer StreamPartsDialer PendingChatAcquireInterval time.Duration MaxChatsPerAcquire int32 InFlightChatStaleAfter time.Duration @@ -4187,7 +3268,6 @@ type Config struct { CreateWorkspace chattool.CreateWorkspaceFn StartWorkspace chattool.StartWorkspaceFn StopWorkspace chattool.StopWorkspaceFn - Pubsub pubsub.Pubsub ProviderAPIKeys chatprovider.ProviderAPIKeys AllowBYOK bool AllowBYOKSet bool @@ -4205,12 +3285,15 @@ type Config struct { // May be nil if the deployment has no OIDC provider; servers // using user_oidc will then send no Authorization header. OIDCTokenSource mcpclient.UserOIDCTokenSource + + NotificationsEnqueuer notifications.Enqueuer + Auditor *atomic.Pointer[audit.Auditor] } -// New creates a new chat processor. The processor polls for pending -// chats and processes them. It is the caller's responsibility to call Close -// on the returned instance. -func New(cfg Config) *Server { +// New creates a new chat processor with the required pubsub dependency. +// The processor polls for pending chats and processes them. It is the +// caller's responsibility to call Close on the returned instance. +func New(ps pubsub.Pubsub, cfg Config) *Server { ctx, cancel := context.WithCancel(context.Background()) pendingChatAcquireInterval := cfg.PendingChatAcquireInterval @@ -4238,6 +3321,11 @@ func New(cfg Config) *Server { clk = quartz.NewReal() } + notificationsEnqueuer := cfg.NotificationsEnqueuer + if notificationsEnqueuer == nil { + notificationsEnqueuer = notifications.NewNoopEnqueuer() + } + instructionLookupTimeout := cfg.InstructionLookupTimeout if instructionLookupTimeout == 0 { instructionLookupTimeout = homeInstructionLookupTimeout @@ -4252,13 +3340,11 @@ func New(cfg Config) *Server { if cfg.AllowBYOKSet { allowBYOK = cfg.AllowBYOK } - p := &Server{ cancel: cancel, db: cfg.Database, workerID: workerID, logger: cfg.Logger.Named("processor"), - subscribeFn: cfg.SubscribeFn, agentConnFn: cfg.AgentConn, agentInactiveDisconnectTimeout: cfg.AgentInactiveDisconnectTimeout, dialTimeout: defaultDialTimeout, @@ -4266,7 +3352,7 @@ func New(cfg Config) *Server { createWorkspaceFn: cfg.CreateWorkspace, startWorkspaceFn: cfg.StartWorkspace, stopWorkspaceFn: cfg.StopWorkspace, - pubsub: cfg.Pubsub, + pubsub: ps, webpushDispatcher: cfg.WebpushDispatcher, providerAPIKeys: cfg.ProviderAPIKeys, allowBYOK: allowBYOK, @@ -4275,7 +3361,7 @@ func New(cfg Config) *Server { debugSvc := chatdebug.NewService( cfg.Database, cfg.Logger.Named("chatdebug"), - cfg.Pubsub, + ps, chatdebug.WithAlwaysEnable(cfg.AlwaysEnableDebugLogs), ) // Debug runs do not heartbeat during model streams; their @@ -4294,58 +3380,80 @@ func New(cfg Config) *Server { usageTracker: cfg.UsageTracker, clock: clk, recordingSem: make(chan struct{}, maxConcurrentRecordingUploads), - wakeCh: make(chan struct{}, 1), - heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), } + var chatAutoArchiveRecords prometheus.Counter if cfg.PrometheusRegistry != nil { p.metrics = chatloop.NewMetrics(cfg.PrometheusRegistry) - cfg.PrometheusRegistry.MustRegister(&streamStateCollector{server: p}) + chatAutoArchiveRecords = prometheus.NewCounter(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "chat_auto_archive", + Name: "records_archived_total", + Help: "Total number of chats archived by the auto-archive job (counting both roots and cascaded children).", + }) + cfg.PrometheusRegistry.MustRegister(chatAutoArchiveRecords) } else { p.metrics = chatloop.NopMetrics() } + p.messagePartBuffer = messagepartbuffer.New(messagepartbuffer.Options{Clock: clk}) + localStreamPartsDialer := NewLocalStreamPartsDialer(LocalStreamPartsDialerConfig{ + Buffer: p.messagePartBuffer, + Logger: cfg.Logger, + }) + p.streamPartsDialer = streamPartsDialerForServer(workerID, localStreamPartsDialer, cfg.StreamPartsDialer) + p.streamSyncPoller = newStreamSyncPoller(ctx, cfg.Database, clk, cfg.Logger.Named("chatstream")) + p.streamSyncPoller.Start() + chatWorker, err := newChatWorker(p, chatWorkerOptions{ + WorkerID: workerID, + Store: cfg.Database, + Pubsub: ps, + Logger: cfg.Logger.Named("chatworker"), + Clock: clk, + MessagePartBuffer: p.messagePartBuffer, + AcquisitionInterval: pendingChatAcquireInterval, + AcquisitionBatchSize: maxChatsPerAcquire, + HeartbeatInterval: chatHeartbeatInterval, + HeartbeatStaleSeconds: int32(inFlightChatStaleAfter.Seconds()), + NotificationsEnqueuer: notificationsEnqueuer, + Auditor: cfg.Auditor, + AutoArchiveRecords: chatAutoArchiveRecords, + }) + if err != nil { + panic("chatd: create chat worker: " + err.Error()) + } + p.chatWorker = chatWorker + //nolint:gocritic // The chat processor uses a scoped chatd context. ctx = dbauthz.AsChatd(ctx) p.configCache = newChatConfigCache(ctx, cfg.Database, clk) - if p.pubsub != nil { - cancelConfigSub, err := p.pubsub.SubscribeWithErr( - coderdpubsub.ChatConfigEventChannel, - coderdpubsub.HandleChatConfigEvent(func(ctx context.Context, ev coderdpubsub.ChatConfigEvent, err error) { - if err != nil { - p.logger.Warn(ctx, "chat config event error", slog.Error(err)) - return - } - switch ev.Kind { - case coderdpubsub.ChatConfigEventProviders: - p.configCache.InvalidateProviders() - case coderdpubsub.ChatConfigEventModelConfig: - p.configCache.InvalidateModelConfig(ev.EntityID) - case coderdpubsub.ChatConfigEventUserPrompt: - p.configCache.InvalidateUserPrompt(ev.EntityID) - case coderdpubsub.ChatConfigEventAdvisorConfig: - p.configCache.InvalidateAdvisorConfig() - } - }), - ) - if err != nil { - p.logger.Error(ctx, "subscribe to chat config events", slog.Error(err)) - } + cancelConfigSub, err := p.pubsub.SubscribeWithErr( + coderdpubsub.ChatConfigEventChannel, + coderdpubsub.HandleChatConfigEvent(func(ctx context.Context, ev coderdpubsub.ChatConfigEvent, err error) { + if err != nil { + p.logger.Warn(ctx, "chat config event error", slog.Error(err)) + return + } + switch ev.Kind { + case coderdpubsub.ChatConfigEventProviders: + p.configCache.InvalidateProviders() + case coderdpubsub.ChatConfigEventModelConfig: + p.configCache.InvalidateModelConfig(ev.EntityID) + case coderdpubsub.ChatConfigEventUserPrompt: + p.configCache.InvalidateUserPrompt(ev.EntityID) + case coderdpubsub.ChatConfigEventAdvisorConfig: + p.configCache.InvalidateAdvisorConfig() + } + }), + ) + if err != nil { + p.logger.Error(ctx, "subscribe to chat config events", slog.Error(err)) + } else { p.configCacheUnsubscribe = cancelConfigSub } p.ctx = ctx - // Recover stale chats on startup. - p.recoverStaleChats(ctx) - if debugSvc := p.debugService(); debugSvc != nil { - if _, err := debugSvc.FinalizeStale(ctx); err != nil { - p.logger.Warn(ctx, "failed to finalize stale chat debug rows", slog.Error(err)) - } - } - // Spawn background goroutines that all servers need. - p.wg.Go(func() { p.heartbeatLoop(ctx) }) - p.wg.Go(func() { p.streamJanitorLoop(ctx) }) return p } @@ -4355,570 +3463,14 @@ func New(cfg Config) *Server { // server (e.g. tests) can skip this call; heartbeat, stream // janitor, and stale recovery still run. func (p *Server) Start() *Server { - p.wg.Go(func() { p.acquireLoop(p.ctx) }) + if p.chatWorker != nil { + if err := p.chatWorker.Start(p.ctx); err != nil { + p.logger.Error(p.ctx, "failed to start chat worker", slog.Error(err)) + } + } return p } -func (p *Server) acquireLoop(ctx context.Context) { - acquireTicker := p.clock.NewTicker( - p.pendingChatAcquireInterval, - "chatd", - "acquire", - ) - defer acquireTicker.Stop() - - staleRecoveryInterval := p.inFlightChatStaleAfter / staleRecoveryIntervalDivisor - staleTicker := p.clock.NewTicker( - staleRecoveryInterval, - "chatd", - "stale-recovery", - ) - defer staleTicker.Stop() - - for { - select { - case <-ctx.Done(): - return - case <-acquireTicker.C: - p.processOnce(ctx) - case <-p.wakeCh: - p.processOnce(ctx) - case <-staleTicker.C: - p.recoverStaleChats(ctx) - if debugSvc := p.existingDebugService(); debugSvc != nil { - if _, err := debugSvc.FinalizeStale(ctx); err != nil { - p.logger.Warn(ctx, "failed to finalize stale chat debug rows", slog.Error(err)) - } - } - } - } -} - -// signalWake wakes the run loop so it calls processOnce immediately. -// Non-blocking: if a signal is already pending it is a no-op. -func (p *Server) signalWake() { - select { - case p.wakeCh <- struct{}{}: - default: - } -} - -func (p *Server) processOnce(ctx context.Context) { - if ctx.Err() != nil { - return - } - - // We detach from the server lifetime to prevent a - // phantom-acquire race: when the server context is - // canceled, the pq driver's watchCancel goroutine - // races with the actual query on the wire. Using a - // context that cannot be canceled ensures the driver - // sees the query result if Postgres executed it. - acquireCtx, acquireCancel := context.WithTimeout( - context.WithoutCancel(ctx), 10*time.Second, - ) - chats, err := p.db.AcquireChats(acquireCtx, database.AcquireChatsParams{ - StartedAt: time.Now(), - WorkerID: p.workerID, - NumChats: p.maxChatsPerAcquire, - }) - acquireCancel() - if err != nil { - p.logger.Error(ctx, "failed to acquire chats", slog.Error(err)) - return - } - if len(chats) == 0 { - return - } - - // If the server context was canceled while we were - // acquiring, release the chats back to pending. - if ctx.Err() != nil { - releaseCtx, releaseCancel := context.WithTimeout( - context.WithoutCancel(ctx), 10*time.Second, - ) - for _, chat := range chats { - _, updateErr := p.db.UpdateChatStatus(releaseCtx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - if updateErr != nil { - p.logger.Error(ctx, "failed to release chat acquired during shutdown", - slog.F("chat_id", chat.ID), slog.Error(updateErr)) - } - } - releaseCancel() - return - } - - p.inflightMu.Lock() - for _, chat := range chats { - p.inflight.Add(1) - go func() { - defer p.inflight.Done() - p.processChat(ctx, chat) - }() - } - p.inflightMu.Unlock() -} - -func shouldClearRetryPhaseForStatus(status codersdk.ChatStatus) bool { - switch status { - case codersdk.ChatStatusWaiting, - codersdk.ChatStatusPending, - codersdk.ChatStatusPaused, - codersdk.ChatStatusCompleted, - codersdk.ChatStatusError, - codersdk.ChatStatusRequiresAction: - return true - default: - return false - } -} - -func (p *Server) clearProvisionalStreamParts(chatID uuid.UUID) { - val, ok := p.chatStreams.Load(chatID) - if !ok { - return - } - rs, ok := val.(*chatStreamState) - if !ok { - return - } - - // Streamed parts are provisional until a durable message commits - // them. A retry rolls back the failed attempt before replacement - // parts are streamed. - rs.mu.Lock() - rs.buffer = nil - rs.resetDropCounters() - rs.mu.Unlock() -} - -func (p *Server) publishToStream(chatID uuid.UUID, event codersdk.ChatStreamEvent) { - state := p.getOrCreateStreamState(chatID) - state.mu.Lock() - switch event.Type { - case codersdk.ChatStreamEventTypeRetry: - if event.Retry != nil { - retryCopy := *event.Retry - state.currentRetry = &retryCopy - } - case codersdk.ChatStreamEventTypeMessagePart: - // Any streamed part means the provider is making forward - // progress again, so the stream has left the retry backoff - // window regardless of role. - state.currentRetry = nil - case codersdk.ChatStreamEventTypeError: - state.currentRetry = nil - case codersdk.ChatStreamEventTypeStatus: - if event.Status != nil && shouldClearRetryPhaseForStatus(event.Status.Status) { - state.currentRetry = nil - } - } - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - if !state.buffering { - p.cleanupStreamIfIdle(chatID, state) - state.mu.Unlock() - return - } - if len(state.buffer) >= maxStreamBufferSize { - p.metrics.RecordStreamBufferDropped() - state.bufferDropCount++ - now := p.clock.Now() - if now.Sub(state.bufferLastWarnAt) >= streamDropWarnInterval { - p.logger.Warn(context.Background(), "chat stream buffer full, dropping oldest event", - slog.F("chat_id", chatID), - slog.F("buffer_size", len(state.buffer)), - slog.F("dropped_count", state.bufferDropCount), - ) - state.bufferDropCount = 0 - state.bufferLastWarnAt = now - } - // Zero the dropped slot so its *ChatStreamMessagePart is - // GC-eligible; the later append reuses this slot in place - // whenever cap > len. - state.buffer[0] = bufferedStreamPart{} - state.buffer = state.buffer[1:] - } - state.buffer = append(state.buffer, bufferedStreamPart{ - event: event, - // committedMessageID stays 0 here: the part belongs to - // the in-progress turn until publishMessage claims it - // with the committed assistant message ID. - }) - } - subscribers := make([]chan codersdk.ChatStreamEvent, 0, len(state.subscribers)) - for _, ch := range state.subscribers { - subscribers = append(subscribers, ch) - } - state.mu.Unlock() - - var subDropped int64 - for _, ch := range subscribers { - select { - case ch <- event: - default: - subDropped++ - } - } - - // Re-acquire the lock once for both subscriber-drop logging and - // idle cleanup. Merging these avoids an unnecessary unlock/re-lock - // gap between the two sections. - state.mu.Lock() - if subDropped > 0 { - state.subscriberDropCount += subDropped - now := p.clock.Now() - if now.Sub(state.subscriberLastWarnAt) >= streamDropWarnInterval { - p.logger.Warn(context.Background(), "dropping chat stream event", - slog.F("chat_id", chatID), - slog.F("type", event.Type), - slog.F("dropped_count", state.subscriberDropCount), - ) - state.subscriberDropCount = 0 - state.subscriberLastWarnAt = now - } - } - p.cleanupStreamIfIdle(chatID, state) - state.mu.Unlock() -} - -// cacheDurableMessage stores a recently persisted message event in the -// per-chat stream state so that same-replica subscribers can catch up -// from memory instead of the database. The afterMessageID is the -// message ID that precedes this message (i.e. message.ID - 1). -func (p *Server) cacheDurableMessage(chatID uuid.UUID, event codersdk.ChatStreamEvent) { - state := p.getOrCreateStreamState(chatID) - state.mu.Lock() - defer state.mu.Unlock() - - if len(state.durableMessages) >= maxDurableMessageCacheSize { - if evicted := state.durableMessages[0]; evicted.Message != nil { - state.durableEvictedBefore = evicted.Message.ID - } - // Zero the dropped slot so the evicted *ChatMessage is - // GC-eligible; see publishToStream for the same pattern. - state.durableMessages[0] = codersdk.ChatStreamEvent{} - state.durableMessages = state.durableMessages[1:] - } - state.durableMessages = append(state.durableMessages, event) -} - -// getCachedDurableMessages returns cached durable messages with IDs -// greater than afterID. Returns nil when the cache has no relevant -// entries. -func (p *Server) getCachedDurableMessages( - chatID uuid.UUID, - afterID int64, -) []codersdk.ChatStreamEvent { - state := p.getOrCreateStreamState(chatID) - state.mu.Lock() - defer state.mu.Unlock() - - if afterID < state.durableEvictedBefore { - return nil - } - - var result []codersdk.ChatStreamEvent - for _, event := range state.durableMessages { - if event.Message != nil && event.Message.ID > afterID { - result = append(result, event) - } - } - return result -} - -// snapshotBufferLocked returns the buffered message_part events that -// the caller should receive in their initial snapshot. -// -// Parts whose committedMessageID != 0 are dropped: those parts were -// claimed by a durable assistant message that the subscriber will -// receive through a different channel (REST snapshot, the initial DB -// query in SubscribeAuthorized, or pubsub catch-up). Delivering them -// here would render the same content twice on the client, once in the -// streaming UI and once as a durable message. -// -// Every caller receives the same view: in-progress parts are always -// delivered and committed parts are always dropped, regardless of -// cursor or relay sentinel. This keeps the buffer free of duplicate -// work for every subscriber, including cross-replica relay -// subscribers whose user-facing peers receive the durable message -// via pubsub. -// -// The caller must hold the per-chat stream state lock. -func snapshotBufferLocked(buffer []bufferedStreamPart) []codersdk.ChatStreamEvent { - if len(buffer) == 0 { - return nil - } - snapshot := make([]codersdk.ChatStreamEvent, 0, len(buffer)) - for _, part := range buffer { - if part.committedMessageID != 0 { - continue - } - snapshot = append(snapshot, part.event) - } - return snapshot -} - -// subscribeToStream registers a subscriber to the per-chat in-memory -// stream and returns a snapshot of currently in-progress message_part -// events plus the current retry phase, the live subscriber channel, -// and a cancel func. -// -// Parts that were claimed by a committed durable assistant message -// (committedMessageID != 0) are excluded from the snapshot. The -// subscriber will receive those durable messages through the REST -// snapshot, the initial DB query in SubscribeAuthorized, or pubsub, -// so re-delivering their constituent parts here would render the -// same content twice. -func (p *Server) subscribeToStream(chatID uuid.UUID) ( - []codersdk.ChatStreamEvent, - *codersdk.ChatStreamRetry, - <-chan codersdk.ChatStreamEvent, - func(), -) { - state := p.getOrCreateStreamState(chatID) - state.mu.Lock() - snapshot := snapshotBufferLocked(state.buffer) - var currentRetry *codersdk.ChatStreamRetry - if state.currentRetry != nil { - retryCopy := *state.currentRetry - currentRetry = &retryCopy - } - id := uuid.New() - ch := make(chan codersdk.ChatStreamEvent, 128) - state.subscribers[id] = ch - state.mu.Unlock() - - cancel := func() { - state.mu.Lock() - // Remove the subscriber but do not close the channel. - // publishToStream copies subscriber references under - // the per-chat lock then sends outside; closing here - // races with that send and can panic. The channel - // becomes unreachable once removed and will be GC'd. - delete(state.subscribers, id) - p.cleanupStreamIfIdle(chatID, state) - state.mu.Unlock() - } - - return snapshot, currentRetry, ch, cancel -} - -// getOrCreateStreamState returns the per-chat stream state, -// creating one atomically if it doesn't exist. The returned -// state has its own mutex — callers must lock state.mu for -// access. -func (p *Server) getOrCreateStreamState(chatID uuid.UUID) *chatStreamState { - if val, ok := p.chatStreams.Load(chatID); ok { - state, _ := val.(*chatStreamState) - return state - } - val, _ := p.chatStreams.LoadOrStore(chatID, &chatStreamState{ - subscribers: make(map[uuid.UUID]chan codersdk.ChatStreamEvent), - }) - state, _ := val.(*chatStreamState) - return state -} - -// cleanupStreamIfIdle removes the chat entry from the sync.Map when -// there are no subscribers, the stream is not buffering, and any -// grace period for late-connecting relay subscribers has elapsed. If -// the grace window is still open it returns without rescheduling. -// streamJanitorLoop is the backstop that re-checks on a timer. -// -// The caller must hold state.mu. The state pointer may have been -// captured outside this lock (sync.Map.Load or Range); we use -// CompareAndDelete so a stale pointer cannot evict a fresh entry -// installed by a racing getOrCreateStreamState. Returns true -// if the state was deleted, false otherwise. -func (p *Server) cleanupStreamIfIdle(chatID uuid.UUID, state *chatStreamState) bool { - if state.buffering || len(state.subscribers) > 0 { - return false - } - // Keep stream state alive during the grace period so - // late-connecting cross-replica relay subscribers can - // register against this chat before GC. - if !state.bufferRetainedAt.IsZero() && - p.clock.Now().Before(state.bufferRetainedAt.Add(bufferRetainGracePeriod)) { - return false - } - if !p.chatStreams.CompareAndDelete(chatID, state) { - return false - } - p.workspaceMCPToolsCache.Delete(chatID) - return true -} - -// streamJanitorLoop periodically reaps idle chat stream states whose -// grace period has expired. It is the backstop for the grace-window -// early-return in cleanupStreamIfIdle; without it, a subscriber that -// detaches inside grace (the common enterprise relay-drain case, -// relayDrainTimeout = 200ms vs. 5s grace) pins the state forever. -func (p *Server) streamJanitorLoop(ctx context.Context) { - ticker := p.clock.NewTicker(streamJanitorInterval, "chatd", "stream-janitor") - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - p.safeSweepIdleStreams(ctx) - } - } -} - -// safeSweepIdleStreams runs sweepIdleStreams under a panic recovery -// so an unexpected panic in the sweep cannot kill the janitor -// goroutine and silently reintroduce the very leak it exists to -// prevent. The next tick retries. -func (p *Server) safeSweepIdleStreams(ctx context.Context) { - defer func() { - if r := recover(); r != nil { - p.logger.Error(ctx, "stream janitor sweep panicked, will retry next tick", - slog.F("panic", r)) - } - }() - p.sweepIdleStreams() -} - -// sweepIdleStreams iterates chatStreams once and delegates each entry -// to cleanupStreamIfIdle. Range may skip entries that become reapable -// concurrently. Any such entry is reaped on the next tick. -func (p *Server) sweepIdleStreams() { - var reaped atomic.Int64 - defer func() { - if count := reaped.Load(); count > 0 { - p.logger.Info(context.Background(), "reaped idle chat streams", slog.F("count", count)) - } - }() - p.chatStreams.Range(func(key, value any) bool { - chatID, ok := key.(uuid.UUID) - if !ok { - return true - } - state, ok := value.(*chatStreamState) - if !ok { - return true - } - // guard against any panic from cleanupStreamIfIdle locking state.mu for all time - func() { - state.mu.Lock() - defer state.mu.Unlock() - if p.cleanupStreamIfIdle(chatID, state) { - reaped.Add(1) - } - }() - return true - }) -} - -// registerHeartbeat enrolls a chat in the centralized batch -// heartbeat loop. Must be called after chatCtx is created. -func (p *Server) registerHeartbeat(entry *heartbeatEntry) { - p.heartbeatMu.Lock() - defer p.heartbeatMu.Unlock() - if _, exists := p.heartbeatRegistry[entry.chatID]; exists { - p.logger.Warn(context.Background(), - "duplicate heartbeat registration, skipping", - slog.F("chat_id", entry.chatID)) - return - } - p.heartbeatRegistry[entry.chatID] = entry -} - -// unregisterHeartbeat removes a chat from the centralized -// heartbeat loop when chat processing finishes. -func (p *Server) unregisterHeartbeat(chatID uuid.UUID) { - p.heartbeatMu.Lock() - defer p.heartbeatMu.Unlock() - delete(p.heartbeatRegistry, chatID) -} - -// heartbeatLoop runs in a single goroutine, issuing one batch -// heartbeat query per interval for all registered chats. -func (p *Server) heartbeatLoop(ctx context.Context) { - ticker := p.clock.NewTicker(p.chatHeartbeatInterval, "chatd", "batch-heartbeat") - defer ticker.Stop() - for { - select { - case <-ctx.Done(): - return - case <-ticker.C: - p.heartbeatTick(ctx) - } - } -} - -// heartbeatTick issues a single batch UPDATE for all running chats -// owned by this worker. Chats missing from the result set are -// interrupted (stolen by another replica or already completed). -func (p *Server) heartbeatTick(ctx context.Context) { - // Snapshot the registry under the lock. - p.heartbeatMu.Lock() - snapshot := maps.Clone(p.heartbeatRegistry) - p.heartbeatMu.Unlock() - - if len(snapshot) == 0 { - return - } - - // Collect the IDs we believe we own. - ids := slices.Collect(maps.Keys(snapshot)) - - //nolint:gocritic // AsChatd provides narrowly-scoped daemon - // access for batch-updating heartbeats. - chatdCtx := dbauthz.AsChatd(ctx) - updatedIDs, err := p.db.UpdateChatHeartbeats(chatdCtx, database.UpdateChatHeartbeatsParams{ - IDs: ids, - WorkerID: p.workerID, - Now: p.clock.Now(), - }) - if err != nil { - p.logger.Error(ctx, "batch heartbeat failed", slog.Error(err)) - return - } - - // Build a set of IDs that were successfully updated. - updated := make(map[uuid.UUID]struct{}, len(updatedIDs)) - for _, id := range updatedIDs { - updated[id] = struct{}{} - } - - // Interrupt registered chats that were not in the result - // (stolen by another replica or already completed). - for id, entry := range snapshot { - if _, ok := updated[id]; !ok { - entry.logger.Warn(ctx, "chat not in batch heartbeat result, interrupting") - entry.cancelWithCause(chatloop.ErrInterrupted) - continue - } - // Bump workspace usage for surviving chats. - newWsID := p.trackWorkspaceUsage(ctx, entry.chatID, entry.workspaceID, entry.logger) - // Update workspace ID in the registry for next tick. - p.heartbeatMu.Lock() - if current, exists := p.heartbeatRegistry[id]; exists { - current.workspaceID = newWsID - } - p.heartbeatMu.Unlock() - } -} - -// streamSubscriberControlFetchContext keeps a control-path lookup tied to the -// requesting subscriber while applying a fallback timeout when the caller has -// no deadline. -func streamSubscriberControlFetchContext(ctx context.Context) (context.Context, context.CancelFunc) { - if _, ok := ctx.Deadline(); ok { - return ctx, func() {} - } - return context.WithTimeout(ctx, chatStreamControlFetchTimeout) -} - func subscribeWithInitialError(chatID uuid.UUID, message string) ( []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, @@ -4934,589 +3486,6 @@ func subscribeWithInitialError(chatID uuid.UUID, message string) ( }}, events, func() {}, true } -func (p *Server) Subscribe( - ctx context.Context, - chatID uuid.UUID, - requestHeader http.Header, - afterMessageID int64, -) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - bool, -) { - if p == nil { - return nil, nil, nil, false - } - - chat, err := p.db.GetChatByID(ctx, chatID) - if err != nil { - if dbauthz.IsNotAuthorizedError(err) { - return nil, nil, nil, false - } - p.logger.Warn(ctx, "failed to load chat for stream subscription", - slog.F("chat_id", chatID), - slog.Error(err), - ) - return subscribeWithInitialError(chatID, "failed to load initial snapshot") - } - return p.SubscribeAuthorized(ctx, chat, requestHeader, afterMessageID) -} - -// SubscribeAuthorized subscribes an already-authorized chat to merged stream -// updates. The passed chat row proves authorization, but SubscribeAuthorized -// still reloads the chat after the stream subscriptions are armed so the -// initial status and relay setup use fresh state. -func (p *Server) SubscribeAuthorized( - ctx context.Context, - chat database.Chat, - requestHeader http.Header, - afterMessageID int64, -) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - bool, -) { - if p == nil { - return nil, nil, nil, false - } - chatID := chat.ID - - // Subscribe to the local stream for message_parts and same-replica - // persisted messages. Capture the current retry phase under the same - // lock so the transient snapshot and subscriber registration reflect - // a single moment in time. - localSnapshot, localRetry, localParts, localCancel := p.subscribeToStream(chatID) - - // Merge all event sources. - mergedCtx, mergedCancel := context.WithCancel(ctx) - mergedEvents := make(chan codersdk.ChatStreamEvent, 128) - - var allCancels []func() - allCancels = append(allCancels, localCancel) - - // Subscribe to pubsub for durable and structured control - // events (status, messages, queue updates, retry, errors). - // When pubsub is nil (e.g. in-memory - // single-instance) we skip this and deliver all local events. - // - // This MUST happen before the DB queries below so that any - // notification published between the query and the subscription - // is not lost (subscribe-first-then-query pattern). - var notifications <-chan coderdpubsub.ChatStreamNotifyMessage - var errCh <-chan error - if p.pubsub != nil { - notifyCh := make(chan coderdpubsub.ChatStreamNotifyMessage, 10) - errNotifyCh := make(chan error, 1) - notifications = notifyCh - errCh = errNotifyCh - - listener := func(_ context.Context, message []byte, listenErr error) { - if listenErr != nil { - select { - case <-mergedCtx.Done(): - case errNotifyCh <- listenErr: - } - return - } - var notify coderdpubsub.ChatStreamNotifyMessage - if unmarshalErr := json.Unmarshal(message, ¬ify); unmarshalErr != nil { - select { - case <-mergedCtx.Done(): - case errNotifyCh <- xerrors.Errorf("unmarshal chat stream notify: %w", unmarshalErr): - } - return - } - select { - case <-mergedCtx.Done(): - case notifyCh <- notify: - } - } - - if pubsubCancel, pubsubErr := p.pubsub.SubscribeWithErr( - coderdpubsub.ChatStreamNotifyChannel(chatID), - listener, - ); pubsubErr == nil { - allCancels = append(allCancels, pubsubCancel) - } else { - p.logger.Warn(ctx, "failed to subscribe to chat stream notifications", - slog.F("chat_id", chatID), - slog.Error(pubsubErr), - ) - } - } - - cancel := func() { - mergedCancel() - for _, cancelFn := range allCancels { - if cancelFn != nil { - cancelFn() - } - } - } - - // Re-read the chat after the local/pubsub subscriptions are active so - // the initial status event and any enterprise relay setup use fresh - // state instead of the middleware-loaded row. - refreshCtx, refreshCancel := streamSubscriberControlFetchContext(ctx) - snapshotChat, err := func() (database.Chat, error) { - defer refreshCancel() - //nolint:gocritic // SubscribeAuthorized already validated the - // caller; this refresh only loads the latest status/worker for - // the already-authorized stream subscription. - return p.db.GetChatByID(dbauthz.AsChatd(refreshCtx), chatID) - }() - if err != nil { - p.logger.Warn(ctx, "failed to refresh chat for stream subscription; using stale state", - slog.F("chat_id", chatID), - slog.Error(err), - ) - snapshotChat = chat - } - - // Build initial snapshot synchronously. The pubsub subscription - // is already active so no notifications can be lost during this - // window. - initialSnapshot := make([]codersdk.ChatStreamEvent, 0) - delivered := map[int64]struct{}{} - // Add local same-replica message_parts to the snapshot. Retry comes - // from state.currentRetry, not the event buffer, so late joiners see - // only the latest phase rather than a stale buffered retry event. - for _, event := range localSnapshot { - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - initialSnapshot = append(initialSnapshot, event) - } - } - - var retryEvent *codersdk.ChatStreamEvent - if localRetry != nil { - retryEvent = &codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeRetry, - ChatID: chatID, - Retry: localRetry, - } - } - - // Load initial messages from DB. When afterMessageID > 0 the - // caller already has messages up to that ID (e.g. from the REST - // endpoint), so we only fetch newer ones to avoid sending - // duplicate data. - messages, err := p.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: afterMessageID, - }) - if err != nil { - p.logger.Error(ctx, "failed to load initial chat messages", - slog.Error(err), - slog.F("chat_id", chatID), - ) - initialSnapshot = append(initialSnapshot, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: &codersdk.ChatError{Message: "failed to load initial snapshot"}, - }) - } else { - for _, msg := range messages { - sdkMsg := db2sdk.ChatMessage(msg) - initialSnapshot = append(initialSnapshot, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, - ChatID: chatID, - Message: &sdkMsg, - }) - delivered[msg.ID] = struct{}{} - } - } - - // Load initial queue. Queue snapshots are intentionally not - // singleflighted because a chat-scoped key cannot distinguish the - // pre- and post-notification queue state. - queueCtx, queueCancel := streamSubscriberControlFetchContext(ctx) - queued, err := p.db.GetChatQueuedMessages(queueCtx, chatID) - queueCancel() - if err != nil { - p.logger.Error(ctx, "failed to load initial queued messages", - slog.Error(err), - slog.F("chat_id", chatID), - ) - initialSnapshot = append(initialSnapshot, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: &codersdk.ChatError{Message: "failed to load initial snapshot"}, - }) - } else if len(queued) > 0 { - initialSnapshot = append(initialSnapshot, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - ChatID: chatID, - QueuedMessages: db2sdk.ChatQueuedMessages(queued), - }) - } - - // Include the current chat status in the snapshot so the - // frontend can gate message_part processing correctly from - // the very first batch, without waiting for a separate REST - // query. - statusEvent := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeStatus, - ChatID: chatID, - Status: &codersdk.ChatStreamStatus{ - Status: codersdk.ChatStatus(snapshotChat.Status), - }, - } - // Prepend so the frontend sees the current stream phases - // before any message_part events. - prefix := []codersdk.ChatStreamEvent{statusEvent} - if retryEvent != nil { - prefix = append(prefix, *retryEvent) - } - initialSnapshot = append(prefix, initialSnapshot...) - - // Track the highest durable message ID delivered to this subscriber, - // whether it came from the initial DB snapshot, the same-replica local - // stream, or a later DB/cache catch-up. - lastMessageID := afterMessageID - if len(messages) > 0 { - lastMessageID = messages[len(messages)-1].ID - } - - // When an enterprise SubscribeFn is provided, call it to get relay events - // (message_parts from remote replicas). OSS owns pubsub subscription, - // message catch-up, queue updates, and status forwarding; enterprise only - // manages relay dialing. - var relayEvents <-chan codersdk.ChatStreamEvent - var statusNotifications chan StatusNotification - if p.subscribeFn != nil { - statusNotifications = make(chan StatusNotification, 10) - relayEvents = p.subscribeFn(mergedCtx, SubscribeFnParams{ - ChatID: chatID, - Chat: snapshotChat, - WorkerID: p.workerID, - StatusNotifications: statusNotifications, - RequestHeader: requestHeader, - DB: p.db, - Logger: p.logger, - }) - } - hasPubsub := false - if p.pubsub != nil { - // hasPubsub is only true when we actually subscribed - // successfully above (allCancels will contain the pubsub - // cancel func in that case). - hasPubsub = len(allCancels) > 1 - } - - //nolint:nestif - go func() { - defer close(mergedEvents) - if statusNotifications != nil { - defer close(statusNotifications) - } - for { - select { - case <-mergedCtx.Done(): - return - case psErr := <-errCh: - p.logger.Error(mergedCtx, "chat stream pubsub error", - slog.F("chat_id", chatID), - slog.Error(psErr), - ) - select { - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: &codersdk.ChatError{ - Message: psErr.Error(), - }, - }: - case <-mergedCtx.Done(): - } - return - case notify := <-notifications: - // Marker for ENG-2645: subscriber received pubsub notify. - p.logger.Debug(mergedCtx, "stream subscriber received notify", - slog.F("chat_id", chatID), - slog.F("after_message_id", notify.AfterMessageID), - slog.F("status", notify.Status), - slog.F("queue_update", notify.QueueUpdate), - slog.F("last_message_id", lastMessageID), - ) - if notify.AfterMessageID > 0 || notify.FullRefresh { - if notify.FullRefresh { - lastMessageID = 0 - clear(delivered) - } - var ( - deliveredCount int - source string - ) - // Notifies can arrive out of order. Rescan from - // min(AfterMessageID, lastMessageID) to cover the gap, - // floored at afterMessageID to respect the subscription - // boundary. The delivered set deduplicates. - lookupAfter := lastMessageID - if !notify.FullRefresh { - lookupAfter = max(afterMessageID, min(notify.AfterMessageID, lastMessageID)) - } - cached := p.getCachedDurableMessages(chatID, lookupAfter) - if !notify.FullRefresh && len(cached) > 0 { - for _, event := range cached { - if event.Message == nil { - continue - } - if _, ok := delivered[event.Message.ID]; ok { - continue - } - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- event: - } - delivered[event.Message.ID] = struct{}{} - if event.Message.ID > lastMessageID { - lastMessageID = event.Message.ID - } - deliveredCount++ - source = "cache" - } - } - // DB pass picks up cross-replica messages the local cache - // cannot have. Delivered set dedupes against the cache pass. - newMessages, msgErr := p.db.GetChatMessagesByChatID(mergedCtx, database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: lookupAfter, - }) - if msgErr != nil { - p.logger.Warn(mergedCtx, "failed to get chat messages after pubsub notification", - slog.F("chat_id", chatID), - slog.Error(msgErr), - ) - } else { - for _, msg := range newMessages { - if msg.ID <= lookupAfter { - continue - } - if _, ok := delivered[msg.ID]; ok { - continue - } - sdkMsg := db2sdk.ChatMessage(msg) - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, - ChatID: chatID, - Message: &sdkMsg, - }: - } - delivered[msg.ID] = struct{}{} - if msg.ID > lastMessageID { - lastMessageID = msg.ID - } - deliveredCount++ - switch source { - case "": - source = "db" - case "cache": - source = "cache+db" - } - } - } - // Marker for ENG-2645: subscriber delivered durable messages. - p.logger.Debug(mergedCtx, "stream subscriber delivered messages", - slog.F("chat_id", chatID), - slog.F("after_message_id", notify.AfterMessageID), - slog.F("lookup_after", lookupAfter), - slog.F("source", source), - slog.F("delivered_count", deliveredCount), - slog.F("last_message_id", lastMessageID), - ) - } - if notify.Status != "" { - status := database.ChatStatus(notify.Status) - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeStatus, - ChatID: chatID, - Status: &codersdk.ChatStreamStatus{Status: codersdk.ChatStatus(status)}, - }: - } - // Notify enterprise relay manager if present. - if statusNotifications != nil { - workerID := uuid.Nil - if notify.WorkerID != "" { - if parsed, parseErr := uuid.Parse(notify.WorkerID); parseErr == nil { - workerID = parsed - } - } - select { - case statusNotifications <- StatusNotification{Status: status, WorkerID: workerID}: - case <-mergedCtx.Done(): - return - } - } - } - if notify.Retry != nil { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeRetry, - ChatID: chatID, - Retry: notify.Retry, - }: - } - } - if notify.ErrorPayload != nil { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: notify.ErrorPayload, - }: - } - } else if notify.Error != "" { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: &codersdk.ChatError{ - Message: notify.Error, - }, - }: - } - } - if notify.QueueUpdate { - queueCtx, queueCancel := streamSubscriberControlFetchContext(mergedCtx) - queuedMsgs, queueErr := p.db.GetChatQueuedMessages(queueCtx, chatID) - queueCancel() - if queueErr != nil { - p.logger.Warn(mergedCtx, "failed to get queued messages after pubsub notification", - slog.F("chat_id", chatID), - slog.Error(queueErr), - ) - } else { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - ChatID: chatID, - QueuedMessages: db2sdk.ChatQueuedMessages(queuedMsgs), - }: - } - } - } - case event, ok := <-localParts: - if !ok { - localParts = nil - // Local parts channel closed. If pubsub is - // active we continue with pubsub-driven events. - // Otherwise terminate. - if !hasPubsub { - return - } - continue - } - if hasPubsub { - // Forward transient events from local. - // Durable events (messages, queue updates) - // come via pubsub + cache. Status is - // included alongside message_part because - // both travel through the same ordered - // channel: publishStatus is called before - // the first message_part, so FIFO delivery - // guarantees the frontend sees - // status=running before any content. - // Pubsub will deliver a duplicate status - // later; the frontend deduplicates it - // (setChatStatus is idempotent). - // action_required is also transient and - // only published on the local stream, so - // it must be forwarded here. - if event.Type == codersdk.ChatStreamEventTypeMessagePart || - event.Type == codersdk.ChatStreamEventTypeStatus || - event.Type == codersdk.ChatStreamEventTypeActionRequired { - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- event: - } - } - } else { - // No pubsub: forward all event types. - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- event: - } - } - case event, ok := <-relayEvents: - if !ok { - relayEvents = nil - continue - } - select { - case <-mergedCtx.Done(): - return - case mergedEvents <- event: - } - } - } - }() - - return initialSnapshot, mergedEvents, cancel, true -} - -func (p *Server) publishEvent(chatID uuid.UUID, event codersdk.ChatStreamEvent) { - if event.ChatID == uuid.Nil { - event.ChatID = chatID - } - p.publishToStream(chatID, event) -} - -func (p *Server) publishStatus(chatID uuid.UUID, status database.ChatStatus, workerID uuid.NullUUID) { - p.publishEvent(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeStatus, - Status: &codersdk.ChatStreamStatus{Status: codersdk.ChatStatus(status)}, - }) - notify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(status), - } - if workerID.Valid { - notify.WorkerID = workerID.UUID.String() - } - p.publishChatStreamNotify(chatID, notify) -} - -// publishChatStreamNotify broadcasts a per-chat stream notification via -// PostgreSQL pubsub so that all replicas can merge durable database updates -// with transient control events. -func (p *Server) publishChatStreamNotify(chatID uuid.UUID, notify coderdpubsub.ChatStreamNotifyMessage) { - if p.pubsub == nil { - return - } - payload, err := json.Marshal(notify) - if err != nil { - p.logger.Error(context.Background(), "failed to marshal chat stream notify", - slog.F("chat_id", chatID), - slog.Error(err), - ) - return - } - if err := p.pubsub.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), payload); err != nil { - p.logger.Error(context.Background(), "failed to publish chat stream notify", - slog.F("chat_id", chatID), - slog.Error(err), - ) - } -} - // publishChatPubsubEvents broadcasts a lifecycle event for each affected chat. func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.ChatWatchEventKind) { for _, chat := range chats { @@ -5559,60 +3528,11 @@ func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWa } } -// pendingToStreamToolCalls converts a slice of chatloop pending -// tool calls into the SDK streaming representation. -func pendingToStreamToolCalls(pending []chatloop.PendingToolCall) []codersdk.ChatStreamToolCall { - calls := make([]codersdk.ChatStreamToolCall, len(pending)) - for i, tc := range pending { - calls[i] = codersdk.ChatStreamToolCall{ - ToolCallID: tc.ToolCallID, - ToolName: tc.ToolName, - Args: tc.Args, - } - } - return calls -} - -// publishChatActionRequired broadcasts an action_required event via -// PostgreSQL pubsub so that global watchers can react to dynamic -// tool calls without streaming each chat individually. -func (p *Server) publishChatActionRequired(chat database.Chat, pending []chatloop.PendingToolCall) { - if p.pubsub == nil { - return - } - toolCalls := pendingToStreamToolCalls(pending) - sdkChat := db2sdk.Chat(chat, nil, nil) - - event := codersdk.ChatWatchEvent{ - Kind: codersdk.ChatWatchEventKindActionRequired, - Chat: sdkChat, - ToolCalls: toolCalls, - } - payload, err := json.Marshal(event) - if err != nil { - p.logger.Error(context.Background(), "failed to marshal chat action_required pubsub event", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - return - } - if err := p.pubsub.Publish(coderdpubsub.ChatWatchEventChannel(chat.OwnerID), payload); err != nil { - p.logger.Error(context.Background(), "failed to publish chat action_required pubsub event", - slog.F("chat_id", chat.ID), - slog.Error(err), - ) - } -} - // PublishDiffStatusChange broadcasts a diff_status_change event for // the given chat so that watching clients know to re-fetch the diff // status. This is called from the HTTP layer after the diff status // is updated in the database. func (p *Server) PublishDiffStatusChange(ctx context.Context, chatID uuid.UUID) error { - if p.pubsub == nil { - return nil - } - chat, err := p.db.GetChatByID(ctx, chatID) if err != nil { return xerrors.Errorf("get chat: %w", err) @@ -5628,226 +3548,6 @@ func (p *Server) PublishDiffStatusChange(ctx context.Context, chatID uuid.UUID) return nil } -func (p *Server) publishRetry(chatID uuid.UUID, payload *codersdk.ChatStreamRetry) { - if payload == nil { - return - } - p.publishEvent(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeRetry, - Retry: payload, - }) - p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - Retry: payload, - }) -} - -func (p *Server) publishError(chatID uuid.UUID, classified chaterror.ClassifiedError) { - payload := chaterror.TerminalErrorPayload(classified) - if payload == nil { - return - } - p.publishEvent(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - Error: payload, - }) - p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - ErrorPayload: payload, - Error: payload.Message, - }) -} - -func processingFailure(err error) (chaterror.ClassifiedError, bool) { - if err == nil { - return chaterror.ClassifiedError{}, false - } - - classified := chaterror.Classify(err) - if classified.Message == "" { - return chaterror.ClassifiedError{}, false - } - return classified, true -} - -func encodeChatLastErrorPayload(payload *codersdk.ChatError) (pqtype.NullRawMessage, error) { - if payload == nil { - return pqtype.NullRawMessage{}, nil - } - encoded, err := json.Marshal(payload) - if err != nil { - return pqtype.NullRawMessage{}, err - } - return pqtype.NullRawMessage{RawMessage: encoded, Valid: true}, nil -} - -func panicFailureReason(recovered any) string { - var reason string - switch typed := recovered.(type) { - case string: - reason = strings.TrimSpace(typed) - case error: - reason = strings.TrimSpace(typed.Error()) - default: - reason = strings.TrimSpace(fmt.Sprint(typed)) - } - - if reason == "" || reason == "" { - return "chat processing panicked" - } - return "chat processing panicked: " + reason -} - -func (p *Server) publishMessage(chatID uuid.UUID, message database.ChatMessage) { - sdkMessage := db2sdk.ChatMessage(message) - event := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, - ChatID: chatID, - Message: &sdkMessage, - } - p.cacheDurableMessage(chatID, event) - // Claim every still-in-progress buffered message_part for this - // durable assistant message BEFORE publishing it, so any new - // subscriber that races publishEvent below takes a buffer - // snapshot in which the parts for this turn are already - // suppressed. Existing subscribers already received the - // constituent parts on the live channel; the frontend - // dedupes those against the durable message via - // clearStreamState in the same batch. - p.claimCommittedParts(chatID, message) - p.publishEvent(chatID, event) - p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - AfterMessageID: message.ID - 1, - }) -} - -// claimCommittedParts walks the chat's buffered message_part events -// and assigns every in-progress part (committedMessageID == 0) to -// the supplied assistant message ID. Subsequent subscriber snapshots -// drop those parts so a reconnecting client does not re-render the -// content of an assistant turn that has already been delivered as a -// durable message via REST or pubsub. -// -// Tool and user messages do not end an assistant streaming turn, so -// only assistant-role messages claim parts. -func (p *Server) claimCommittedParts(chatID uuid.UUID, message database.ChatMessage) { - if message.Role != database.ChatMessageRoleAssistant { - return - } - val, ok := p.chatStreams.Load(chatID) - if !ok { - return - } - state, ok := val.(*chatStreamState) - if !ok { - return - } - state.mu.Lock() - defer state.mu.Unlock() - for i := range state.buffer { - if state.buffer[i].committedMessageID == 0 { - state.buffer[i].committedMessageID = message.ID - } - } -} - -// publishEditedMessage is like publishMessage but uses FullRefresh -// so remote subscribers re-fetch from the beginning, ensuring the -// edit is never silently dropped. The durable cache is replaced -// with only the edited message. -func (p *Server) publishEditedMessage(chatID uuid.UUID, message database.ChatMessage) { - sdkMessage := db2sdk.ChatMessage(message) - event := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, - ChatID: chatID, - Message: &sdkMessage, - } - state := p.getOrCreateStreamState(chatID) - state.mu.Lock() - state.durableMessages = []codersdk.ChatStreamEvent{event} - state.durableEvictedBefore = 0 - state.mu.Unlock() - p.publishEvent(chatID, event) - p.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - FullRefresh: true, - }) -} - -func (p *Server) publishMessagePart(chatID uuid.UUID, role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { - if part.Type == "" { - return - } - // Strip internal-only fields before client delivery. - // Mirrors db2sdk.chatMessageParts stripping for REST. - part.StripInternal() - p.publishEvent(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: role, - Part: part, - }, - }) -} - -func shouldCancelChatFromControlNotification( - notify coderdpubsub.ChatStreamNotifyMessage, - workerID uuid.UUID, -) bool { - status := database.ChatStatus(strings.TrimSpace(notify.Status)) - switch status { - case database.ChatStatusWaiting, database.ChatStatusPending, database.ChatStatusError: - return true - case database.ChatStatusRunning: - worker := strings.TrimSpace(notify.WorkerID) - if worker == "" { - return false - } - notifyWorkerID, err := uuid.Parse(worker) - if err != nil { - return false - } - return notifyWorkerID != workerID - default: - return false - } -} - -func (p *Server) subscribeChatControl( - ctx context.Context, - chatID uuid.UUID, - cancel context.CancelCauseFunc, - logger slog.Logger, -) func() { - if p.pubsub == nil { - return nil - } - - listener := func(_ context.Context, message []byte, err error) { - if err != nil { - logger.Warn(ctx, "chat control pubsub error", slog.Error(err)) - return - } - - var notify coderdpubsub.ChatStreamNotifyMessage - if unmarshalErr := json.Unmarshal(message, ¬ify); unmarshalErr != nil { - logger.Warn(ctx, "failed to unmarshal chat control notify", slog.Error(unmarshalErr)) - return - } - - if shouldCancelChatFromControlNotification(notify, p.workerID) { - cancel(chatloop.ErrInterrupted) - } - } - - controlCancel, err := p.pubsub.SubscribeWithErr( - coderdpubsub.ChatStreamNotifyChannel(chatID), - listener, - ) - if err != nil { - logger.Warn(ctx, "failed to subscribe to chat control notifications", slog.Error(err)) - return nil - } - return controlCancel -} - // Rejects oversize images on capped providers before any upstream // request is issued. // @@ -5902,74 +3602,6 @@ func (p *Server) chatFileResolver(provider string) chatprompt.FileResolver { } } -// tryAutoPromoteQueuedMessage pops the next queued message and converts it -// into a pending user message inside the caller's transaction. Queued -// messages were already admitted through SendMessage, so this preserves FIFO -// order without re-checking usage limits. -func (p *Server) tryAutoPromoteQueuedMessage( - ctx context.Context, - tx database.Store, - chat database.Chat, -) (*database.ChatMessage, []database.ChatQueuedMessage, bool, error) { - logger := p.logger.With(slog.F("chat_id", chat.ID)) - - queuedMessages, err := tx.GetChatQueuedMessages(ctx, chat.ID) - if err != nil { - return nil, nil, false, xerrors.Errorf("get queued messages: %w", err) - } - if len(queuedMessages) == 0 { - return nil, nil, false, nil - } - nextQueued := queuedMessages[0] - effectiveModelConfigID, err := resolveQueuedMessageModelConfigID( - ctx, - tx, - chat, - nextQueued.ModelConfigID, - ) - if err != nil { - return nil, nil, false, err - } - - poppedQueued, err := tx.PopNextQueuedMessage(ctx, chat.ID) - if err != nil { - return nil, nil, false, xerrors.Errorf("pop next queued message: %w", err) - } - if poppedQueued.ID != nextQueued.ID { - return nil, nil, false, xerrors.New("popped queued message out of order") - } - - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. - ChatID: chat.ID, - } - queuedUserMsg := newUserChatMessage( - nextQueued.APIKeyID.String, - pqtype.NullRawMessage{ - RawMessage: nextQueued.Content, - Valid: len(nextQueued.Content) > 0, - }, - database.ChatMessageVisibilityBoth, - effectiveModelConfigID, - chatprompt.CurrentContentVersion, - ) - queuedUserMsg = queuedUserMsg.withCreatedBy(chat.OwnerID) - appendUserChatMessage(&msgParams, queuedUserMsg) - msgs, err := insertChatMessageWithStore(ctx, tx, msgParams) - if err != nil { - return nil, nil, false, xerrors.Errorf("insert promoted message: %w", err) - } - msg := msgs[0] - - remainingQueuedMessages, err := tx.GetChatQueuedMessages(ctx, chat.ID) - if err != nil { - logger.Error(ctx, "failed to load remaining queued messages after auto-promotion", - slog.F("queued_message_id", nextQueued.ID), slog.Error(err)) - return &msg, nil, false, nil - } - - return &msg, remainingQueuedMessages, true, nil -} - // trackWorkspaceUsage bumps the workspace's last_used_at via the // usage tracker and extends the workspace's autostop deadline. If // wsID is not yet valid, it re-reads the chat from the DB to pick @@ -6004,7 +3636,7 @@ func (p *Server) trackWorkspaceUsage( // so no prebuild guard is needed (unlike reporter.go). // // This fires every heartbeat (~30s) but the SQL only - // writes when 5% of the deadline has elapsed — most calls + // writes when 5% of the deadline has elapsed, most calls // perform a read-only CTE lookup with no UPDATE. // // Scaling note: for 10,000 active chats, this could lead to @@ -6016,467 +3648,16 @@ func (p *Server) trackWorkspaceUsage( return wsID } -type finishActiveChatResult struct { - updatedChat database.Chat - promotedMessage *database.ChatMessage - syntheticToolResults []database.ChatMessage - remainingQueuedMessages []database.ChatQueuedMessage - shouldPublishQueueUpdate bool -} - -func (p *Server) finishActiveChat( - ctx context.Context, - logger slog.Logger, - chat database.Chat, - status database.ChatStatus, - lastError pqtype.NullRawMessage, -) (finishActiveChatResult, error) { - result := finishActiveChatResult{} - - err := p.db.InTx(func(tx database.Store) error { - // Re-read the chat status under lock — another caller - // (e.g. promote) may have already set it to pending. - latestChat, lockErr := tx.GetChatByIDForUpdate(ctx, chat.ID) - if lockErr != nil { - return xerrors.Errorf("lock chat for release: %w", lockErr) - } - - // If another worker has already acquired this chat, - // bail out — we must not overwrite their running - // status or publish spurious events. - if latestChat.Status == database.ChatStatusRunning && - latestChat.WorkerID.Valid && - latestChat.WorkerID.UUID != p.workerID { - return errChatTakenByOtherWorker - } - - // If someone else already set the chat to pending (e.g. - // the promote endpoint), don't overwrite it — just clear - // the worker and let the processor pick it back up. - switch { - case latestChat.Status == database.ChatStatusPending: - status = database.ChatStatusPending - case latestChat.Status == database.ChatStatusWaiting && status != database.ChatStatusWaiting && !latestChat.Archived: - // PromoteQueued's deferred path won the status race. - // Insert synthetic tool results before auto-promoting, - // or a RequiresAction worker outcome reintroduces the - // stops-dead bug this PR exists to fix. - inserted, synthErr := insertSyntheticToolResultsTx( - ctx, tx, latestChat, - "Tool execution interrupted by queued message promotion", - ) - if synthErr != nil { - return xerrors.Errorf("insert synthetic tool results during promote-driven cleanup: %w", synthErr) - } - result.syntheticToolResults = inserted - var promoteErr error - result.promotedMessage, result.remainingQueuedMessages, result.shouldPublishQueueUpdate, promoteErr = p.tryAutoPromoteQueuedMessage(ctx, tx, latestChat) - if promoteErr != nil { - logger.Error(ctx, "auto-promote queued message failed during promote-driven cleanup", slog.Error(promoteErr)) - return xerrors.Errorf("auto-promote queued message: %w", promoteErr) - } - if result.promotedMessage != nil { - status = database.ChatStatusPending - } else { - // Queue drained between snapshot and lock; honor - // the external Waiting. - status = database.ChatStatusWaiting - } - case status == database.ChatStatusWaiting && !latestChat.Archived: - // Queued messages were already admitted through SendMessage, - // so auto-promotion only preserves FIFO order here. Archived - // chats skip promotion so archiving behaves like a hard stop. - var promoteErr error - result.promotedMessage, result.remainingQueuedMessages, result.shouldPublishQueueUpdate, promoteErr = p.tryAutoPromoteQueuedMessage(ctx, tx, latestChat) - if promoteErr != nil { - logger.Error(ctx, "auto-promote queued message failed, rolling back", slog.Error(promoteErr)) - return xerrors.Errorf("auto-promote queued message: %w", promoteErr) - } else if result.promotedMessage != nil { - status = database.ChatStatusPending - } - } - - var updateErr error - result.updatedChat, updateErr = tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: status, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: lastError, - }) - return updateErr - }, nil) - if err != nil { - return finishActiveChatResult{}, err - } - - return result, nil -} - -func (p *Server) shouldPublishFinishedChatState( - ctx context.Context, - logger slog.Logger, - updatedChat database.Chat, -) bool { - latestChat, err := p.db.GetChatByID(ctx, updatedChat.ID) - if err != nil { - logger.Warn(ctx, "failed to re-read chat before publishing finished state", - slog.F("chat_id", updatedChat.ID), - slog.Error(err), - ) - return true - } - - if latestChat.Status != updatedChat.Status || latestChat.WorkerID != updatedChat.WorkerID { - logger.Debug(ctx, "skipping stale finished chat publish", - slog.F("chat_id", updatedChat.ID), - slog.F("expected_status", updatedChat.Status), - slog.F("expected_worker_id", updatedChat.WorkerID), - slog.F("latest_status", latestChat.Status), - slog.F("latest_worker_id", latestChat.WorkerID), - ) - return false - } - - return true -} - -func (p *Server) processChat(ctx context.Context, chat database.Chat) { - logger := p.logger.With(slog.F("chat_id", chat.ID)) - logger.Info(ctx, "processing chat request") - - p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Inc() - defer p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Dec() - - chatCtx, cancel := context.WithCancelCause(ctx) - defer cancel(nil) - - // Gate the control subscriber behind a channel that is closed - // after we publish "running" status. This prevents stale - // pubsub notifications (e.g. the "pending" notification from - // SendMessage that triggered this processing) from - // interrupting us before we start work. Due to async - // PostgreSQL NOTIFY delivery, a notification published before - // subscribeChatControl registers its queue can still arrive - // after registration. - controlArmed := make(chan struct{}) - gatedCancel := func(cause error) { - select { - case <-controlArmed: - cancel(cause) - default: - logger.Debug(ctx, "ignoring control notification before armed") - } - } - - controlCancel := p.subscribeChatControl(chatCtx, chat.ID, gatedCancel, logger) - defer func() { - if controlCancel != nil { - controlCancel() - } - }() - - // Register with the centralized heartbeat loop instead of - // running a per-chat goroutine. The loop issues a single batch - // UPDATE for all chats on this worker and detects stolen chats - // via set-difference. - p.registerHeartbeat(&heartbeatEntry{ - cancelWithCause: cancel, - chatID: chat.ID, - workspaceID: chat.WorkspaceID, - logger: logger, - }) - defer p.unregisterHeartbeat(chat.ID) - - // Start buffering stream events BEFORE publishing the running - // status. This closes a race where a subscriber sees - // status=running but misses message_part events because - // buffering hasn't started yet — the subscriber gets an empty - // snapshot and publishToStream drops message_parts while - // buffering is false. - streamState := p.getOrCreateStreamState(chat.ID) - streamState.mu.Lock() - streamState.buffer = nil - streamState.bufferRetainedAt = time.Time{} - streamState.resetDropCounters() - streamState.buffering = true - streamState.mu.Unlock() - defer func() { - streamState.mu.Lock() - // Fallback cleanup for exit paths that return before a - // terminal stream event is published. - streamState.currentRetry = nil - streamState.resetDropCounters() - streamState.buffering = false - // Retain the per-chat stream state for a grace period - // so cross-replica relay subscribers can register - // against this chat after processing completes, - // without racing cleanupStreamIfIdle. The buffer is - // cleared when the next processChat starts or when - // cleanupStreamIfIdle runs after the grace period; on - // the normal-completion path every part has been - // claimed by its durable assistant message, so the - // snapshot is empty. On error or panic exit some parts - // may still be in-progress; those are likewise - // discarded when the buffer is cleared, and the - // frontend recovers via the next REST snapshot. - streamState.bufferRetainedAt = p.clock.Now() - streamState.mu.Unlock() - }() - - p.publishStatus(chat.ID, database.ChatStatusRunning, uuid.NullUUID{ - UUID: p.workerID, - Valid: true, - }) - - // Arm the control subscriber. Closing the channel is a - // happens-before guarantee in the Go memory model — any - // notification dispatched after this point will correctly - // interrupt processing. - close(controlArmed) - - // Determine the final status and last error payload to set when we're done. - status := database.ChatStatusWaiting - wasInterrupted := false - var lastErrorPayload *codersdk.ChatError - generatedTitle := &generatedChatTitle{} - runResult := runChatResult{} - remainingQueuedMessages := []database.ChatQueuedMessage{} - shouldPublishQueueUpdate := false - var promotedMessage *database.ChatMessage - - defer func() { - // Use a context that is not canceled by Close() so we can - // reliably update the chat status in the database during - // graceful shutdown. - cleanupCtx := context.WithoutCancel(ctx) - - // Handle panics gracefully. - if r := recover(); r != nil { - logger.Error(cleanupCtx, "panic during chat processing", slog.F("panic", r)) - classified := chaterror.ClassifiedError{ - Message: panicFailureReason(r), - Kind: codersdk.ChatErrorKindGeneric, - } - lastErrorPayload = chaterror.TerminalErrorPayload(classified) - p.publishError(chat.ID, classified) - status = database.ChatStatusError - } - - encodedLastError, err := encodeChatLastErrorPayload(lastErrorPayload) - if err != nil { - logger.Warn(cleanupCtx, "failed to marshal chat last error payload", - slog.Error(err), - ) - lastErrorPayload = nil - encodedLastError = pqtype.NullRawMessage{} - } - - // Check for queued messages and auto-promote the next one. - // This must be done atomically with the status update to avoid - // races with the promote endpoint (which also sets status to - // pending). We use a transaction with FOR UPDATE to ensure we - // don't overwrite a status change made by another caller. - finishResult, err := p.finishActiveChat(cleanupCtx, logger, chat, status, encodedLastError) - if errors.Is(err, errChatTakenByOtherWorker) { - // Another worker owns this chat now — skip all - // post-TX side effects (status publish, pubsub, - // web push) to avoid overwriting their state. - return - } - if err != nil { - logger.Error(cleanupCtx, "failed to release chat", slog.Error(err)) - return - } - status = finishResult.updatedChat.Status - promotedMessage = finishResult.promotedMessage - remainingQueuedMessages = finishResult.remainingQueuedMessages - shouldPublishQueueUpdate = finishResult.shouldPublishQueueUpdate - - // Publish synth rows before the promoted user message. - for _, msg := range finishResult.syntheticToolResults { - p.publishMessage(chat.ID, msg) - } - if promotedMessage != nil { - p.publishMessage(chat.ID, *promotedMessage) - } - if shouldPublishQueueUpdate { - p.publishEvent(chat.ID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeQueueUpdate, - QueuedMessages: db2sdk.ChatQueuedMessages(remainingQueuedMessages), - }) - p.publishChatStreamNotify(chat.ID, coderdpubsub.ChatStreamNotifyMessage{ - QueueUpdate: true, - }) - } - if p.shouldPublishFinishedChatState(cleanupCtx, logger, finishResult.updatedChat) { - p.publishStatus(chat.ID, status, uuid.NullUUID{}) - // Best-effort: use any generated title captured during - // processing so push notifications and the status snapshot - // can reflect it without another DB read. The dedicated - // title_change event remains the source of truth. - if title, ok := generatedTitle.Load(); ok { - finishResult.updatedChat.Title = title - } - p.publishChatPubsubEvent(finishResult.updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) - } - - if promotedMessage != nil { - // Wake the processor so it picks up the newly pending - // chat immediately instead of waiting for the next - // acquire-interval tick. - p.signalWake() - } - - // When the chat is parked in requires_action, - // publish the stream event and global pubsub event - // after the DB status has committed. Publishing - // here (not in runChat) prevents a race where a - // fast client reacts before the status is visible. - if status == database.ChatStatusRequiresAction && len(runResult.PendingDynamicToolCalls) > 0 { - toolCalls := pendingToStreamToolCalls(runResult.PendingDynamicToolCalls) - p.publishEvent(chat.ID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeActionRequired, - ActionRequired: &codersdk.ChatStreamActionRequired{ - ToolCalls: toolCalls, - }, - }) - p.publishChatActionRequired(finishResult.updatedChat, runResult.PendingDynamicToolCalls) - } - if wasInterrupted { - p.maybeClearLastTurnSummaryAsync(cleanupCtx, finishResult.updatedChat, logger) - } else { - lastErrorMessage := "" - if lastErrorPayload != nil { - lastErrorMessage = lastErrorPayload.Message - } - p.maybeFinalizeTurnStatusLabelAndPush( - cleanupCtx, - finishResult.updatedChat, - status, - lastErrorMessage, - runResult, - logger, - ) - } - }() - - p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Dec() - p.metrics.Chats.WithLabelValues(chatloop.StateStreaming).Inc() - defer func() { - p.metrics.Chats.WithLabelValues(chatloop.StateStreaming).Dec() - p.metrics.Chats.WithLabelValues(chatloop.StateWaiting).Inc() - }() - runResult, err := p.runChat(chatCtx, chat, generatedTitle, logger) - if err != nil { - if errors.Is(err, chatloop.ErrInterrupted) || errors.Is(context.Cause(chatCtx), chatloop.ErrInterrupted) { - logger.Info(ctx, "chat interrupted") - status = database.ChatStatusWaiting - lastErrorPayload = nil - wasInterrupted = true - return - } - if isShutdownCancellation(ctx, chatCtx, err) { - logger.Info(ctx, "chat canceled during shutdown; returning to pending") - status = database.ChatStatusPending - lastErrorPayload = nil - wasInterrupted = true - return - } - logger.Error(ctx, "failed to process chat", slog.Error(err)) - if classified, ok := processingFailure(err); ok { - lastErrorPayload = chaterror.TerminalErrorPayload(classified) - p.publishError(chat.ID, classified) - } - status = database.ChatStatusError - return - } - - // The LLM invoked a dynamic tool — park the chat in - // requires_action so the client can supply tool results. - if len(runResult.PendingDynamicToolCalls) > 0 { - status = database.ChatStatusRequiresAction - return - } - - // If runChat completed successfully but the server context was - // canceled (e.g. during Close()), the chat should be returned - // to pending so another replica can pick it up. There is a - // race where the LLM stream finishes just as the server is - // shutting down — the HTTP response completes before context - // cancellation propagates, so runChat returns nil instead of - // a context.Canceled error. Without this check the chat would - // be marked "waiting" and never retried. - if ctx.Err() != nil { - logger.Info(ctx, "chat completed during shutdown; returning to pending") - status = database.ChatStatusPending - lastErrorPayload = nil - wasInterrupted = true - return - } -} - -func isShutdownCancellation( - serverCtx context.Context, - chatCtx context.Context, - err error, -) bool { - if err == nil { - return false - } - // During Close(), the server context is canceled. In-flight chats should - // be returned to pending so another replica can retry them. - if serverCtx.Err() == nil { - return false - } - if errors.Is(err, context.Canceled) { - return true - } - return errors.Is(context.Cause(chatCtx), context.Canceled) -} - -// generatedChatTitle shares an asynchronously generated title between the -// detached title-generation goroutine and the deferred cleanup path. -type generatedChatTitle struct { - mu sync.RWMutex - title string -} - -func (t *generatedChatTitle) Store(title string) { - if t == nil || title == "" { - return - } - - t.mu.Lock() - t.title = title - t.mu.Unlock() -} - -func (t *generatedChatTitle) Load() (string, bool) { - if t == nil { - return "", false - } - - t.mu.RLock() - defer t.mu.RUnlock() - if t.title == "" { - return "", false - } - return t.title, true -} - type runChatResult struct { - FinalAssistantText string - StatusLabelModel fantasy.LanguageModel - ProviderKeys chatprovider.ProviderAPIKeys - PendingDynamicToolCalls []chatloop.PendingToolCall - FallbackProvider string - FallbackRoute resolvedModelRoute - FallbackModel string - ModelBuildOptions modelBuildOptions - TriggerMessageID int64 - HistoryTipMessageID int64 + FinalAssistantText string + StatusLabelModel fantasy.LanguageModel + ProviderKeys chatprovider.ProviderAPIKeys + FallbackProvider string + FallbackRoute resolvedModelRoute + FallbackModel string + ModelBuildOptions modelBuildOptions + TriggerMessageID int64 + HistoryTipMessageID int64 } func activeTurnAPIKeyIDFromMessages(messages []database.ChatMessage) (string, bool) { @@ -6804,8 +3985,6 @@ type rootChatToolsOptions struct { modelConfigID uuid.UUID workspaceCtx *turnWorkspaceContext workspaceMu *sync.Mutex - instruction *string - skills *[]chattool.SkillMeta resolvePlanPath func(context.Context) (string, string, error) storeFile chattool.StoreFileFunc isPlanModeTurn bool @@ -6934,31 +4113,14 @@ func (p *Server) appendRootChatTools( // build logs before the tool completes. p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) - // When a workspace is first attached mid-turn (e.g. via - // create_workspace), fetch and persist instruction files - // immediately so the LLM has AGENTS.md context for the remainder - // of this turn. The persisted marker prevents redundant fetches on - // subsequent turns. - if *opts.instruction == "" && updatedChat.WorkspaceID.Valid { - newInstruction, discoveredSkills, persistErr := p.persistInstructionFiles( - ctx, - updatedChat, - opts.modelConfigID, - opts.workspaceCtx.getWorkspaceAgent, - opts.workspaceCtx.getWorkspaceConn, - ) - if persistErr != nil { - p.logger.Warn(ctx, "failed to persist instruction files on workspace attach", - slog.F("chat_id", updatedChat.ID), - slog.Error(persistErr), - ) - } else { - *opts.instruction = newInstruction - if len(discoveredSkills) > 0 { - *opts.skills = discoveredSkills - } - } - } + // Note: we intentionally do not insert AGENTS.md / workspace + // context here. Local tool callbacks must not mutate chat + // history while a local-tool generation task is in flight, + // because that advances history_version before the tool + // result is committed and exits the local-tool commit as + // stale. Workspace context is persisted by the + // persist_workspace_context generation action in a later + // pass. // Prime the workspace MCP tools cache while the create_workspace // or start_workspace tool is still running. The AgentID guard @@ -6967,14 +4129,14 @@ func (p *Server) appendRootChatTools( // empty list on the first try when the agent's MCP Connect is // racing with agent startup; primeWorkspaceMCPCache retries // with a short backoff up to workspaceMCPPrimeMaxWait. Priming - // here lets the next LLM step's PrepareTools hit the cache + // here lets the next assistant-generation action hit the cache // instead of dialing again on a separate timeout budget. // // Run asynchronously: the tool itself must not block on the // primer because the agent may not advertise any MCP tools at // all (e.g. minimal templates), in which case the primer waits - // the full budget before giving up. PrepareTools on the next - // step covers the cache miss path; the primer is purely an + // the full budget before giving up. The next assistant-generation + // action covers the cache miss path; the primer is purely an // optimization that warms the cache while the LLM is thinking. // inflight tracking ensures server shutdown still waits for any // in-progress primer. @@ -6987,21 +4149,11 @@ func (p *Server) appendRootChatTools( // the pre-build and stop-side firings would otherwise spawn a // primer goroutine that dials a missing or dying agent and // burns the full budget for nothing. - // - // Read the snapshot from workspaceCtx rather than the - // updatedChat parameter: persistInstructionFiles above runs - // ensureWorkspaceAgent which calls persistBuildAgentBinding and - // setCurrentChat, so by the time we get here the in-memory - // snapshot has the freshly bound AgentID even when the - // updatedChat parameter (read from the DB before the binding - // was persisted) does not. snapshot := opts.workspaceCtx.currentChatSnapshot() if snapshot.WorkspaceID.Valid && snapshot.AgentID.Valid { - p.inflight.Add(1) - go func() { - defer p.inflight.Done() + p.inflight.Go(func() { p.primeWorkspaceMCPCache(opts.primerCtx, p.logger, snapshot.ID, opts.workspaceCtx) - }() + }) } } @@ -7107,1333 +4259,6 @@ func appendDynamicTools( return append(tools, dynamicToolsFromSDK(logger, filteredDefs)...), dynamicToolNames, nil } -func (p *Server) runChat( - ctx context.Context, - chat database.Chat, - generatedTitle *generatedChatTitle, - logger slog.Logger, -) (runChatResult, error) { - result := runChatResult{} - var ( - model fantasy.LanguageModel - modelConfig database.ChatModelConfig - providerKeys chatprovider.ProviderAPIKeys - callConfig codersdk.ChatModelCallConfig - messages []database.ChatMessage - err error - debugEnabled bool - debugProvider string - modelRoute resolvedModelRoute - debugModel string - ) - - messages, err = p.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) - if err != nil { - return result, xerrors.Errorf("get chat messages: %w", err) - } - modelOpts := modelBuildOptionsFromMessages(messages) - if modelOpts.ActiveAPIKeyID != "" { - ctx = aibridge.WithDelegatedAPIKeyID(ctx, modelOpts.ActiveAPIKeyID) - } - - // Load MCP server configs and user tokens in parallel with model - // resolution. These queries have no dependencies on each other and all - // hit different tables. - var ( - mcpConfigs []database.MCPServerConfig - mcpTokens []database.MCPServerUserToken - ) - var g errgroup.Group - g.Go(func() error { - var err error - model, modelConfig, providerKeys, modelRoute, debugEnabled, debugProvider, debugModel, err = p.resolveChatModel(ctx, chat, modelOpts) - if err != nil { - return err - } - if len(modelConfig.Options) > 0 { - if err := json.Unmarshal(modelConfig.Options, &callConfig); err != nil { - return xerrors.Errorf("parse model call config: %w", err) - } - } - return nil - }) - if len(chat.MCPServerIDs) > 0 { - g.Go(func() error { - var err error - mcpConfigs, err = p.db.GetMCPServerConfigsByIDs( - ctx, chat.MCPServerIDs, - ) - if err != nil { - logger.Warn(ctx, - "failed to load MCP server configs", - slog.Error(err), - ) - } - return nil - }) - g.Go(func() error { - var err error - // If token loading fails, ConnectAll will still - // proceed but oauth2-authenticated servers will - // attempt to connect without credentials. Those - // connections may succeed or fail depending on - // the remote server's auth requirements. - mcpTokens, err = p.db.GetMCPServerUserTokensByUserID( - ctx, chat.OwnerID, - ) - if err != nil { - logger.Warn(ctx, - "failed to load MCP user tokens", - slog.Error(err), - ) - } - return nil - }) - } - if err := g.Wait(); err != nil { - return result, err - } - - // Capture the current turn's mode so prompt and tool behavior can - // be resolved consistently for the rest of the turn. - currentPlanMode := chat.PlanMode - isPlanModeTurn := currentPlanMode.Valid && currentPlanMode.ChatPlanMode == database.ChatPlanModePlan - isExploreSubagent := isExploreSubagentMode(chat.Mode) - isRootChat := !chat.ParentChatID.Valid - var mcpConnectConfigs []database.MCPServerConfig - var approvedPlanMCPConfigIDs map[uuid.UUID]struct{} - // Explore subagents rely on the immutable spawn-time snapshot - // persisted in chat.MCPServerIDs. SendMessage cannot mutate that - // snapshot, so no runtime re-filter against parent state is needed. - // The child's persisted set is authoritative. - mcpConnectConfigs, approvedPlanMCPConfigIDs = filterExternalMCPConfigsForTurn( - mcpConfigs, - currentPlanMode, - chat.ParentChatID, - ) - if isExploreSubagent && isRootChat { - // Root Explore chats stay builtin-only per the accepted plan, so - // strip any persisted external MCP configs at runtime regardless of - // what's on the chat row. Explore children get their snapshot via - // the spawn-time inheritance path and are handled below. - mcpConnectConfigs = nil - approvedPlanMCPConfigIDs = map[uuid.UUID]struct{}{} - } - planModeInstructions := p.loadPlanModeInstructions(ctx, currentPlanMode, logger) - - advisorCfg := p.loadAdvisorConfig(ctx, logger) - - var advisorRuntime *chatadvisor.Runtime - // Plan mode filters the advisor tool out of the turn's tool set via - // filterToolsForTurn, so enabling the runtime there would inject - // guidance and enforce advisor exclusivity for a tool the model - // cannot actually call. Explore chats (root or subagent) run under - // allowedExploreToolNames, whose policy does not include advisor, so - // registering the runtime there would inject guidance for a tool - // that is never exposed to the model. - if advisorCfg.Enabled && isRootChat && !isPlanModeTurn && !isExploreSubagent { - var advisorErr error - advisorRuntime, advisorErr = p.newAdvisorRuntime( - ctx, - chat, - advisorCfg, - model, - callConfig, - providerKeys, - modelOpts, - logger, - ) - if advisorErr != nil { - return result, advisorErr - } - } - - var advisorPromptSnapshot []fantasy.Message - // setAdvisorPromptSnapshot captures the final prompt state the outer - // model sees so the advisor tool can forward it as nested context. - // It is invoked at four lifecycle points (after initial system-prompt - // assembly, inside PrepareMessages before and after instruction - // injection, and after ReloadMessages rebuilds the prompt) because - // the prompt mutates at each of them and the advisor must snapshot - // the post-mutation state. Removing any of those calls would leave - // the advisor with a stale view of the conversation. - // - // The no-op guard keeps the common disabled/filtered paths (advisor - // off, plan mode, explore, child chats) from paying an O(n) prompt - // clone per step for a snapshot that is never consumed. - setAdvisorPromptSnapshot := func(msgs []fantasy.Message) { - if advisorRuntime == nil { - return - } - advisorPromptSnapshot = slices.Clone(msgs) - } - - chainInfo := chatopenai.ResolveChainMode(messages) - result.StatusLabelModel = model - result.ProviderKeys = providerKeys - result.FallbackProvider = modelConfig.Provider - result.FallbackRoute = modelRoute - result.FallbackModel = modelConfig.Model - result.ModelBuildOptions = modelOpts - debugSvc := p.existingDebugService() - // Fire title generation asynchronously so it doesn't block the - // chat response. It uses a detached context so it can finish - // even after the chat processing context is canceled. - // Snapshot values captured by the goroutine because model, providerKeys, - // logger, and ctx are reassigned below. - titleModel := model - titleProviderKeys := providerKeys - titleLogger := logger - titleCtx := context.WithoutCancel(ctx) - p.inflight.Add(1) - go func() { - defer p.inflight.Done() - p.maybeGenerateChatTitle( - titleCtx, - chat, - messages, - modelConfig.Provider, - modelConfig.Model, - titleModel, - modelRoute, - titleProviderKeys, - modelOpts, - generatedTitle, - titleLogger, - debugSvc, - ) - }() - - // Detect computer-use subagent via the mode column. - isComputerUse := chat.Mode.Valid && chat.Mode.ChatMode == database.ChatModeComputerUse - - var ( - computerUseProvider string - computerUseModelProvider string - computerUseModelName string - ) - if isComputerUse { - var err error - computerUseProvider, computerUseModelProvider, computerUseModelName, err = p.computerUseProviderAndModelFromConfig(ctx) - if err != nil { - return result, xerrors.Errorf( - "resolve computer use provider and model: %w", - err, - ) - } - } - - // NOTE: Buffering was already started in processChat before - // the running status was published, so message_part events - // are captured from the moment subscribers can see - // status=running. The deferred cleanup also lives in - // processChat. - - currentChat := chat - loadChatSnapshot := func( - loadCtx context.Context, - chatID uuid.UUID, - ) (database.Chat, error) { - return p.db.GetChatByID(loadCtx, chatID) - } - var ( - chatStateMu sync.Mutex - workspaceMu sync.Mutex - ) - workspaceCtx := turnWorkspaceContext{ - server: p, - chatStateMu: &chatStateMu, - currentChat: ¤tChat, - loadChatSnapshot: loadChatSnapshot, - } - // primerCtx scopes the workspace MCP cache primer goroutines that - // onChatUpdated launches. We cancel it before workspaceCtx.close() - // so an in-flight primer cannot wake from its retry backoff, - // observe a cleared cached conn, dial a fresh one, and leak it - // when no subsequent close() runs. - primerCtx, primerCancel := context.WithCancel(ctx) - defer func() { - primerCancel() - workspaceCtx.close() - }() - - planPathFn := func(ctx context.Context) (string, string, error) { - conn, err := workspaceCtx.getWorkspaceConn(ctx) - if err != nil { - return "", "", err - } - home, err := chattool.ResolveWorkspaceHome(ctx, conn) - if err != nil { - return "", "", err - } - return chattool.PlanPathForChat(home, chat.ID), home, nil - } - resolvePlanPathForTools := func(ctx context.Context) (string, string, error) { - ctx, cancel := context.WithTimeout(ctx, planPathLookupTimeout) - defer cancel() - return planPathFn(ctx) - } - resolvePlanPathBlock := func(resolveCtx context.Context) string { - if chat.ParentChatID.Valid { - return "" - } - - planCtx, cancel := context.WithTimeout(resolveCtx, planPathLookupTimeout) - defer cancel() - - if _, _, err := workspaceCtx.workspaceAgentIDForConn(planCtx); err != nil { - p.logger.Debug(resolveCtx, "plan path instruction: agent not reachable", - slog.Error(err), - slog.F("chat_id", chat.ID), - ) - return "" - } - - planPath, home, err := planPathFn(planCtx) - if err != nil { - p.logger.Debug(resolveCtx, "plan path instruction: failed to resolve plan path", - slog.Error(err), - slog.F("chat_id", chat.ID), - ) - return "" - } - - return formatPlanPathBlock(planPath, home) - } - - // Connect to MCP servers in parallel with instruction - // resolution. ConnectAll only depends on mcpConfigs and - // mcpTokens which are available after g.Wait() above. - var ( - instruction string - resolvedUserPrompt string - mcpTools []fantasy.AgentTool - mcpCleanup func() - workspaceMCPTools []fantasy.AgentTool - workspaceSkills []chattool.SkillMeta - personalSkills []skillspkg.Skill - ) - // Check if instruction files need to be (re-)persisted. - // This happens when no context-file parts exist yet, or when - // the workspace agent has changed (e.g. workspace rebuilt). - needsInstructionPersist := false - hasContextFiles := false - persistedSkills := skillsFromParts(messages) - latestInjectedAgentID, hasLatestInjectedAgent := latestContextAgentID(messages) - currentWorkspaceAgentID := uuid.Nil - hasCurrentWorkspaceAgent := false - if chat.WorkspaceID.Valid { - if agent, agentErr := workspaceCtx.getWorkspaceAgent(ctx); agentErr == nil { - currentWorkspaceAgentID = agent.ID - hasCurrentWorkspaceAgent = true - } - persistedAgentID, found := contextFileAgentID(messages) - hasContextFiles = found - if !hasPersistedInstructionFiles(messages) { - needsInstructionPersist = true - } else if hasCurrentWorkspaceAgent && currentWorkspaceAgentID != persistedAgentID { - // Agent changed. Persist fresh instruction files. - // Old context-file messages remain in the conversation - // to preserve the prompt cache prefix. - needsInstructionPersist = true - } - } - // Convert messages to prompt format in parallel with g2 work. - // ConvertMessagesWithFiles only reads `messages` (available - // after g.Wait()) and resolves file references via the DB. - // No g2 task reads or writes `prompt`, so this is safe. - var prompt []fantasy.Message - var g2 errgroup.Group - g2.Go(func() error { - var err error - prompt, err = chatprompt.ConvertMessagesWithFiles(ctx, messages, p.chatFileResolver(modelConfig.Provider), logger) - if err != nil { - return xerrors.Errorf("build chat prompt: %w", err) - } - return nil - }) - if needsInstructionPersist { - g2.Go(func() error { - var persistErr error - var discoveredSkills []chattool.SkillMeta - instruction, discoveredSkills, persistErr = p.persistInstructionFiles( - ctx, - chat, - modelConfig.ID, - workspaceCtx.getWorkspaceAgent, - func(instructionCtx context.Context) (workspacesdk.AgentConn, error) { - if _, _, err := workspaceCtx.workspaceAgentIDForConn(instructionCtx); err != nil { - return nil, err - } - return workspaceCtx.getWorkspaceConn(instructionCtx) - }, - ) - workspaceSkills = selectSkillMetasForInstructionRefresh( - persistedSkills, - discoveredSkills, - uuid.NullUUID{UUID: currentWorkspaceAgentID, Valid: hasCurrentWorkspaceAgent}, - uuid.NullUUID{UUID: latestInjectedAgentID, Valid: hasLatestInjectedAgent}, - ) - if persistErr != nil { - p.logger.Warn(ctx, "failed to persist instruction files", - slog.F("chat_id", chat.ID), - slog.Error(persistErr), - ) - } - return nil - }) - } else if hasContextFiles { - // On subsequent turns, extract the instruction text and - // skill index from persisted parts so they can be - // re-injected via InsertSystem after compaction drops - // those messages. No workspace dial needed. - instruction = instructionFromContextFiles(messages) - workspaceSkills = persistedSkills - } - g2.Go(func() error { - personalSkills = p.fetchPersonalSkillMetadata(ctx, chat.OwnerID, logger) - return nil - }) - g2.Go(func() error { - resolvedUserPrompt = p.resolveUserPrompt(ctx, chat.OwnerID) - return nil - }) - if len(mcpConnectConfigs) > 0 { - g2.Go(func() error { - // Refresh expired OAuth2 tokens before connecting. - mcpTokens = p.refreshExpiredMCPTokens(ctx, logger, mcpConnectConfigs, mcpTokens) - mcpTools, mcpCleanup = mcpclient.ConnectAll( - ctx, logger, mcpConnectConfigs, mcpTokens, chat.OwnerID, p.oidcTokenSource, - chatprovider.CoderHeaders(chat), - ) - return nil - }) - } - // Workspace MCP discovery stays disabled for all plan-mode turns. - // Root plan mode only gets approved external MCP servers, and - // plan-mode subagents get no MCP tools. When the chat has no - // workspace yet, discovery happens mid-turn via the chatloop - // PrepareTools callback installed below in chatloop.Run options. - if chat.WorkspaceID.Valid && !isPlanModeTurn { - g2.Go(func() error { - workspaceMCPTools = p.discoverWorkspaceMCPTools( - ctx, logger, chat.ID, &workspaceCtx, - ) - return nil - }) - } - if err := g2.Wait(); err != nil { - return result, err - } - prompt, sanitizeStats := chatsanitize.SanitizeAnthropicProviderToolHistory(model.Provider(), prompt) - chatsanitize.LogAnthropicProviderToolSanitization( - ctx, logger, "persisted_history_replay", model.Provider(), model.Model(), sanitizeStats, - ) - subagentInstruction := "" - if !isRootChat { - subagentInstruction = defaultSubagentInstruction - } - resolvedSkillsFor := func(workspaceSkills []chattool.SkillMeta) []skillspkg.ResolvedSkill { - return mergeTurnSkills(personalSkills, workspaceSkills) - } - resolveSkillAlias := func(alias string) (skillspkg.ResolvedSkill, error) { - return skillspkg.Lookup(resolvedSkillsFor(workspaceSkills), alias) - } - initialResolvedSkills := resolvedSkillsFor(workspaceSkills) - injectedSkillIndex := chattool.FormatResolvedSkillIndex(initialResolvedSkills) - prompt = buildSystemPrompt( - prompt, - subagentInstruction, - instruction, - initialResolvedSkills, - resolvedUserPrompt, - systemPromptBehaviorContext{ - planMode: currentPlanMode, - chatMode: chat.Mode, - planModeInstructions: planModeInstructions, - isRootChat: isRootChat, - }, - ) - // Inject advisor guidance when the advisor runtime is available. - if advisorRuntime != nil { - prompt = chatprompt.InsertSystem(prompt, chatadvisor.ParentGuidanceBlock) - } - if mcpCleanup != nil { - defer mcpCleanup() - } - - // Build a lookup from tool name to MCP server config ID - // so we can annotate persisted parts with the originating - // server. - toolNameToConfigID := make(map[string]uuid.UUID) - for _, t := range mcpTools { - if mcpTool, ok := t.(mcpclient.MCPToolIdentifier); ok { - toolNameToConfigID[t.Info().Name] = mcpTool.MCPServerConfigID() - } - } - - instructionInjected := instruction != "" - // workspaceMCPDiscovered tracks whether workspace MCP discovery - // has already been attempted for this turn. The top-of-turn - // discovery path above only fires when chat.WorkspaceID is - // valid at the start of the turn. For chats that bind a - // workspace mid-turn (e.g. via create_workspace) the chatloop - // PrepareTools callback below triggers discovery on the next - // step. After discovery has run once (here or in PrepareTools), - // this flag prevents redundant dials. - workspaceMCPDiscovered := chat.WorkspaceID.Valid || isPlanModeTurn - prompt = renderPlanPathPrompt(prompt, resolvePlanPathBlock(ctx)) - setAdvisorPromptSnapshot(prompt) - // Use the model config's context_limit as a fallback when the LLM - // provider doesn't include context_limit in its response metadata - // (which is the common case). - modelConfigContextLimit := modelConfig.ContextLimit - var finalAssistantText string - var pendingDynamicCalls []chatloop.PendingToolCall - - compactionHistoryTipMessageID := int64(0) - if len(messages) > 0 { - compactionHistoryTipMessageID = messages[len(messages)-1].ID - } - - var compactionOptions *chatloop.CompactionOptions - - persistStep := func(persistCtx context.Context, step chatloop.PersistedStep) error { - // If the chat context has been canceled, bail out before - // inserting any messages. We distinguish the cause so that - // the caller can tell an intentional interruption (e.g. - // EditMessage, user stop) from a server shutdown: - // - ErrInterrupted cause → return ErrInterrupted - // (processChat sets status = waiting). - // - Any other cause (e.g. context.Canceled during - // Close()) → return the original context error so - // isShutdownCancellation can match and set status = - // pending, allowing another replica to retry. - if persistCtx.Err() != nil { - if errors.Is(context.Cause(persistCtx), chatloop.ErrInterrupted) { - return chatloop.ErrInterrupted - } - return persistCtx.Err() - } - - // Capture pending dynamic tool calls so the caller - // can surface them after chatloop.Run returns. - pendingDynamicCalls = step.PendingDynamicToolCalls - - // Split the step content into assistant blocks and tool - // result blocks so they can be stored as separate messages - // with the appropriate roles. Provider-executed tool results - // (e.g. web_search) stay in the assistant content because - // the LLM provider expects them inline in the assistant - // turn, not as separate tool messages. - var assistantBlocks []fantasy.Content - var toolResults []fantasy.ToolResultContent - for _, block := range step.Content { - if tr, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok { - if !tr.ProviderExecuted { - toolResults = append(toolResults, tr) - continue - } - } - if trPtr, ok := fantasy.AsContentType[*fantasy.ToolResultContent](block); ok && trPtr != nil { - if !trPtr.ProviderExecuted { - toolResults = append(toolResults, *trPtr) - continue - } - } - assistantBlocks = append(assistantBlocks, block) - } - - // Pre-marshal all content outside the transaction so the - // FOR UPDATE lock is held only for the INSERT statements. - // Marshaling is pure CPU work with no database dependency. - assistantParts := buildAssistantPartsForPersist( - persistCtx, - p.logger, - assistantBlocks, - toolResults, - step, - toolNameToConfigID, - ) - - var assistantContent pqtype.NullRawMessage - if len(assistantParts) > 0 { - finalAssistantText = strings.TrimSpace(contentBlocksToText(assistantParts)) - var marshalErr error - assistantContent, marshalErr = chatprompt.MarshalParts(assistantParts) - if marshalErr != nil { - return xerrors.Errorf("marshal assistant content: %w", marshalErr) - } - } - - toolResultContents := make([]pqtype.NullRawMessage, len(toolResults)) - for i, tr := range toolResults { - trPart := chatprompt.PartFromContentWithLogger(ctx, logger, tr) - if trPart.ToolName != "" { - if configID, ok := toolNameToConfigID[trPart.ToolName]; ok { - trPart.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} - } - } - // Apply recorded timestamps so persisted - // tool-result parts carry accurate CreatedAt. - if trPart.ToolCallID != "" && step.ToolResultCreatedAt != nil { - if ts, ok := step.ToolResultCreatedAt[trPart.ToolCallID]; ok { - trPart.CreatedAt = &ts - } - } - var marshalErr error - toolResultContents[i], marshalErr = chatprompt.MarshalParts([]codersdk.ChatMessagePart{trPart}) - if marshalErr != nil { - return xerrors.Errorf("marshal tool result %d: %w", i, marshalErr) - } - } - - hasUsage := step.Usage != (fantasy.Usage{}) - usageForCost := fantasyUsageToChatMessageUsage(step.Usage) - totalCostMicros := chatcost.CalculateTotalCostMicros(usageForCost, callConfig.Cost) - - var insertedMessages []database.ChatMessage - if err := p.db.InTx(func(tx database.Store) error { - // Verify this worker still owns the chat before - // inserting messages. This closes the race where - // EditMessage soft-deletes history and clears worker_id - // while persistInterruptedStep (which uses an - // uncancelable context) is still running. - // - // When the chat is in "waiting" status (set by - // InterruptChat / setChatWaiting), the worker_id has - // already been cleared but we still want to persist - // the partial assistant response. We allow the write - // because the history has NOT been truncated — the - // user simply asked to stop. In contrast, EditMessage - // sets the chat to "pending" after truncating, so the - // pending check still correctly blocks stale writes. - lockedChat, lockErr := tx.GetChatByIDForUpdate(persistCtx, chat.ID) - if lockErr != nil { - return xerrors.Errorf("lock chat for persist: %w", lockErr) - } - if !lockedChat.WorkerID.Valid || lockedChat.WorkerID.UUID != p.workerID { - // The worker_id was cleared. Only allow the persist - // if the chat transitioned to "waiting" (interrupt), - // not "pending" (edit) or any other status. - if lockedChat.Status != database.ChatStatusWaiting { - return chatloop.ErrInterrupted - } - } - - stepParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. - ChatID: chat.ID, - } - - var contextLimit int64 - if step.ContextLimit.Valid { - contextLimit = step.ContextLimit.Int64 - } - - var runtimeMs int64 - if step.Runtime > 0 { - runtimeMs = step.Runtime.Milliseconds() - } - - var totalCostVal int64 - if totalCostMicros != nil { - totalCostVal = *totalCostMicros - } - - var inputTokens, outputTokens, totalTokens int64 - var reasoningTokens, cacheCreationTokens, cacheReadTokens int64 - if hasUsage { - inputTokens = step.Usage.InputTokens - outputTokens = step.Usage.OutputTokens - totalTokens = step.Usage.TotalTokens - reasoningTokens = step.Usage.ReasoningTokens - cacheCreationTokens = step.Usage.CacheCreationTokens - cacheReadTokens = step.Usage.CacheReadTokens - } - - if assistantContent.Valid { - appendChatMessage(&stepParams, newChatMessage( - database.ChatMessageRoleAssistant, - assistantContent, - database.ChatMessageVisibilityBoth, - modelConfig.ID, - chatprompt.CurrentContentVersion, - ).withUsage( - inputTokens, outputTokens, totalTokens, - reasoningTokens, cacheCreationTokens, cacheReadTokens, - ).withContextLimit(contextLimit). - withTotalCostMicros(totalCostVal). - withRuntimeMs(runtimeMs). - withProviderResponseID(step.ProviderResponseID)) - } - - for _, resultContent := range toolResultContents { - appendChatMessage(&stepParams, newChatMessage( - database.ChatMessageRoleTool, - resultContent, - database.ChatMessageVisibilityBoth, - modelConfig.ID, - chatprompt.CurrentContentVersion, - )) - } - - if len(stepParams.Role) > 0 { - inserted, insertErr := tx.InsertChatMessages(persistCtx, stepParams) - if insertErr != nil { - return xerrors.Errorf("insert step messages: %w", insertErr) - } - insertedMessages = append(insertedMessages, inserted...) - } - - return nil - }, nil); err != nil { - return xerrors.Errorf("persist step transaction: %w", err) - } - - for _, msg := range insertedMessages { - p.publishMessage(chat.ID, msg) - } - if len(insertedMessages) > 0 { - compactionHistoryTipMessageID = insertedMessages[len(insertedMessages)-1].ID - if compactionOptions != nil { - compactionOptions.HistoryTipMessageID = compactionHistoryTipMessageID - } - } - - // Do NOT clear the stream buffer here. The per-chat - // stream state must remain alive for the post-completion - // grace window so cross-replica relay subscribers can - // register without racing cleanupStreamIfIdle. The buffer - // is bounded by maxStreamBufferSize and is cleared when - // the next processChat starts or when the stream state - // is garbage-collected after the retention grace period. - - return nil - } - // Apply the default MaxOutputTokens if the model config - // does not specify one. - if callConfig.MaxOutputTokens == nil { - maxOutputTokens := int64(32_000) - callConfig.MaxOutputTokens = &maxOutputTokens - } - - // 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() - effectiveThreshold := modelConfig.CompressionThreshold - thresholdSource := "model_default" - if override, ok := p.resolveUserCompactionThreshold(ctx, chat.OwnerID, modelConfig.ID); ok { - effectiveThreshold = override - thresholdSource = "user_override" - } - compactionOptions = &chatloop.CompactionOptions{ - ThresholdPercent: effectiveThreshold, - ContextLimit: modelConfig.ContextLimit, - HistoryTipMessageID: compactionHistoryTipMessageID, - Persist: func( - persistCtx context.Context, - result chatloop.CompactionResult, - ) error { - if err := p.persistChatContextSummary( - persistCtx, - chat.ID, - modelConfig.ID, - modelOpts.ActiveAPIKeyID, - compactionToolCallID, - result, - ); err != nil { - return xerrors.Errorf("persist context summary: %w", err) - } - logger.Info(persistCtx, "chat context summarized", - slog.F("chat_id", chat.ID), - slog.F("threshold_source", thresholdSource), - slog.F("threshold_percent", result.ThresholdPercent), - slog.F("usage_percent", result.UsagePercent), - slog.F("context_tokens", result.ContextTokens), - slog.F("context_limit", result.ContextLimit), - ) - return nil - }, - ToolCallID: compactionToolCallID, - ToolName: "chat_summarized", - PublishMessagePart: func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { - p.publishMessagePart(chat.ID, role, part) - }, - OnError: func(err error) { - logger.Warn(ctx, "failed to compact chat context", slog.Error(err)) - }, - } - - if isComputerUse { - computerUseRoute, keyErr := p.resolveModelRouteForProviderType(ctx, chat.OwnerID, computerUseModelProvider) - if keyErr != nil { - return result, xerrors.Errorf("resolve computer use provider route: %w", keyErr) - } - providerKeys = computerUseRoute.directProviderKeys() - - // Override model for computer use subagent. - cuModel, cuDebugEnabled, resolvedProvider, resolvedModel, cuErr := p.resolveComputerUseModel( - ctx, - chat, - computerUseRoute, - computerUseProvider, - computerUseModelProvider, - computerUseModelName, - modelOpts, - ) - if cuErr != nil { - return result, cuErr - } - model = cuModel - debugEnabled = cuDebugEnabled - debugProvider = resolvedProvider - debugModel = resolvedModel - } - if debugEnabled { - if debugSvc == nil { - return result, xerrors.New("chat debug service missing after enablement check") - } - compactionOptions.DebugSvc = debugSvc - compactionOptions.ChatID = chat.ID - } - - // Enrich the scoped logger with provider/model for this turn. - // Bound once after the cuModel swap; slog.Logger.With appends - // rather than deduping. - logger = logger.With( - slog.F("provider", model.Provider()), - slog.F("model", model.Model()), - ) - - allowAskUserQuestion := isPlanModeTurn && isRootChat - storeChatAttachment := p.newStoreChatAttachmentFunc(&workspaceCtx) - tools := []fantasy.AgentTool{ - chattool.ReadFile(chattool.ReadFileOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - }), - chattool.WriteFile(chattool.WriteFileOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - ResolvePlanPath: resolvePlanPathForTools, - IsPlanTurn: isPlanModeTurn, - }), - chattool.EditFiles(chattool.EditFilesOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - ResolvePlanPath: resolvePlanPathForTools, - IsPlanTurn: isPlanModeTurn, - }), - chattool.AttachFile(chattool.AttachFileOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - StoreFile: storeChatAttachment, - }), - chattool.Execute(chattool.ExecuteOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - }), - chattool.ProcessOutput(chattool.ProcessToolOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - }), - chattool.ProcessList(chattool.ProcessToolOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - }), - chattool.ProcessSignal(chattool.ProcessToolOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - }), - } - if allowAskUserQuestion { - tools = append(tools, chattool.NewAskUserQuestionTool()) - } - // Only root chats (not delegated subagents) get workspace - // provisioning and subagent tools. Child agents must not - // create workspaces or spawn further subagents. They should - // focus on completing their delegated task. - if isRootChat { - tools = p.appendRootChatTools(ctx, tools, rootChatToolsOptions{ - chat: chat, - modelConfigID: modelConfig.ID, - workspaceCtx: &workspaceCtx, - workspaceMu: &workspaceMu, - instruction: &instruction, - skills: &workspaceSkills, - resolvePlanPath: resolvePlanPathForTools, - storeFile: storeChatAttachment, - isPlanModeTurn: isPlanModeTurn, - primerCtx: primerCtx, - }) - } - - skillOpts := chattool.ReadSkillOptions{ - GetWorkspaceConn: workspaceCtx.getWorkspaceConn, - GetSkills: func() []chattool.SkillMeta { - return workspaceSkills - }, - ResolveAlias: resolveSkillAlias, - LoadPersonalSkillBody: func(ctx context.Context, name string) (skillspkg.ParsedSkill, error) { - return p.loadPersonalSkillBody(ctx, chat.OwnerID, name) - }, - } - appendCurrentSkillTools := func(current []fantasy.AgentTool) ([]fantasy.AgentTool, bool) { - if len(personalSkills) == 0 && len(workspaceSkills) == 0 { - return current, false - } - - updated := current - changed := false - appendTool := func(tool fantasy.AgentTool) { - name := tool.Info().Name - if slices.ContainsFunc(current, func(existing fantasy.AgentTool) bool { - return existing.Info().Name == name - }) { - return - } - if !changed { - updated = slices.Clone(current) - changed = true - } - updated = append(updated, tool) - } - appendTool(chattool.ReadSkill(skillOpts)) - if len(workspaceSkills) > 0 { - appendTool(chattool.ReadSkillFile(skillOpts)) - } - return updated, changed - } - tools, _ = appendCurrentSkillTools(tools) - if advisorRuntime != nil { - tools = append(tools, chatadvisor.Tool(chatadvisor.ToolOptions{ - Runtime: advisorRuntime, - GetConversationSnapshot: func() []fantasy.Message { - // The outer prompt contains ParentGuidanceBlock, which - // tells the parent when to call the advisor tool. That - // instruction is meaningless (and slightly confusing) - // when forwarded to the advisor, whose nested run has - // no tools. Strip it before handing the snapshot over. - return stripAdvisorGuidanceBlock(slices.Clone(advisorPromptSnapshot)) - }, - PublishAdviceDelta: func(toolCallID string, delta string) { - if toolCallID == "" || delta == "" { - return - } - p.publishMessagePart(chat.ID, codersdk.ChatMessageRoleTool, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: toolCallID, - ToolName: chatadvisor.ToolName, - ResultDelta: delta, - }) - }, - PublishAdviceReset: func(toolCallID string) { - if toolCallID == "" { - return - } - p.publishMessagePart(chat.ID, codersdk.ChatMessageRoleTool, codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: toolCallID, - ToolName: chatadvisor.ToolName, - ResultReset: true, - }) - }, - })) - } - - var exclusiveToolNames map[string]bool - if advisorRuntime != nil { - exclusiveToolNames = map[string]bool{chatadvisor.ToolName: true} - } - - // Record builtin tool names before appending MCP tools - // so the metrics layer can differentiate between built-in and MCP tools. - builtinToolNames := make(map[string]bool, len(tools)) - for _, t := range tools { - builtinToolNames[t.Info().Name] = true - } - - // Append external MCP tools from the chat's persisted snapshot after the - // built-ins so the LLM sees them as additional capabilities. Explore chats - // trust only the persisted MCPServerIDs snapshot, and workspace-local MCP - // tools stay unavailable to Explore chats. - tools = append(tools, mcpTools...) - if !isExploreSubagent { - tools = append(tools, workspaceMCPTools...) - } - tools = filterToolsForTurn( - tools, - currentPlanMode, - chat.ParentChatID, - approvedPlanMCPConfigIDs, - ) - // Append dynamic tools declared by the client at chat - // creation time. These appear in the LLM's tool list but - // are never executed by the chatloop. The client handles - // execution via POST /tool-results. - var dynamicToolNames map[string]bool - tools, dynamicToolNames, err = appendDynamicTools( - ctx, - logger, - tools, - chat.DynamicTools, - currentPlanMode, - chat.Mode, - ) - if err != nil { - return result, err - } - - // Build provider-native tools (e.g. web search) based on the - // current model configuration. Root Explore chats stay builtin-only per - // the accepted plan, so delegated Explore children are the only Explore - // chats that can inherit web_search. Write-style provider tools stay - // blocked for all Explore chats. - var providerTools []chatloop.ProviderTool - if !isPlanModeTurn && callConfig.ProviderOptions != nil { - providerTools = buildProviderTools(callConfig.ProviderOptions) - if isExploreSubagent { - if !chat.ParentChatID.Valid { - providerTools = nil - } else { - providerTools = slices.DeleteFunc(providerTools, func(tool chatloop.ProviderTool) bool { - return tool.Definition.GetName() != "web_search" - }) - } - } - } - - providerTools, err = appendComputerUseProviderTool( - providerTools, - computerUseProviderToolOptions{ - provider: computerUseProvider, - isPlanModeTurn: isPlanModeTurn, - isComputerUse: isComputerUse, - getWorkspaceConn: workspaceCtx.getWorkspaceConn, - storeFile: storeChatAttachment, - clock: p.clock, - logger: p.logger.Named("computer_use"), - }, - ) - if err != nil { - return result, xerrors.Errorf( - "register computer use provider tool for provider %q: %w", - computerUseProvider, - err, - ) - } - - providerOptions := chatprovider.ProviderOptionsFromChatModelConfig( - model, - callConfig.ProviderOptions, - ) - // When the OpenAI Responses API has store=true, the provider - // retains conversation history server-side. For follow-up turns, - // we set previous_response_id and send only system instructions - // plus the new user input, avoiding redundant replay of prior - // assistant and tool messages that the provider already has. - chainModeActive := chatopenai.ShouldActivateChainMode( - providerOptions, - chainInfo, - modelConfig.ID, - isPlanModeTurn, - ) - if !chainModeActive && chainInfo.PreviousResponseID() != "" { - logger.Debug(ctx, "chain mode disabled", - slog.F("has_unresolved_local_tool_calls", chainInfo.HasUnresolvedLocalToolCalls()), - slog.F("provider_missing_tool_results", chainInfo.ProviderMissingToolResults()), - slog.F("is_plan_mode_turn", isPlanModeTurn), - slog.F("model_config_match", chainInfo.ModelConfigID() == modelConfig.ID), - slog.F("store_enabled", chatopenai.IsResponsesStoreEnabled(providerOptions)), - slog.F("contributing_trailing_user_count", chainInfo.ContributingTrailingUserCount()), - ) - } - if chainModeActive { - providerOptions = chatopenai.WithPreviousResponseID( - providerOptions, - chainInfo.PreviousResponseID(), - ) - prompt = chatopenai.FilterPromptForChainMode(prompt, chainInfo) - } - activeToolNames := activeToolNamesForTurn( - tools, - currentPlanMode, - chat.ParentChatID, - approvedPlanMCPConfigIDs, - ) - if isExploreSubagent { - activeToolNames = allowedExploreToolNames(tools) - } - - var loopErr error - triggerMessageID, historyTipMessageID, triggerLabel := deriveChatDebugSeed(messages) - - // Enrich the logger with correlation fields useful for - // diagnosing tool-call errors inside the chatloop. - loopLogger := logger.With( - slog.F("owner_id", chat.OwnerID), - slog.F("organization_id", chat.OrganizationID), - slog.F("trigger_message_id", triggerMessageID), - ) - if chat.WorkspaceID.Valid { - loopLogger = loopLogger.With(slog.F("workspace_id", chat.WorkspaceID.UUID)) - } - if chat.AgentID.Valid { - loopLogger = loopLogger.With(slog.F("agent_id", chat.AgentID.UUID)) - } - if chat.ParentChatID.Valid { - loopLogger = loopLogger.With(slog.F("parent_chat_id", chat.ParentChatID.UUID)) - } - result.TriggerMessageID = triggerMessageID - result.HistoryTipMessageID = historyTipMessageID - finishDebugRun := func(error, any) {} - if debugEnabled { - ctx, finishDebugRun = prepareChatTurnDebugRun( - ctx, - logger, - chat, - modelConfig, - debugSvc, - debugProvider, - debugModel, - triggerMessageID, - historyTipMessageID, - triggerLabel, - ) - } - defer func() { - panicValue := recover() - finishDebugRun(loopErr, panicValue) - if panicValue != nil { - panic(panicValue) - } - }() - - loopErr = chatloop.Run(ctx, chatloop.RunOptions{ - Model: model, - Messages: prompt, - Tools: tools, - ActiveTools: activeToolNames, - StopAfterTools: stopAfterBehaviorTools(currentPlanMode, chat.Mode, chat.ParentChatID), - MaxSteps: maxChatSteps, - Metrics: p.metrics, - Logger: loopLogger, - BuiltinToolNames: builtinToolNames, - ExclusiveToolNames: exclusiveToolNames, - - ModelConfig: callConfig, - ProviderOptions: providerOptions, - ProviderTools: providerTools, - // dynamicToolNames now contains only names that don't - // collide with built-in/MCP tools. - DynamicToolNames: dynamicToolNames, - - ContextLimitFallback: modelConfigContextLimit, - - PersistStep: persistStep, - PublishMessagePart: func( - role codersdk.ChatMessageRole, - part codersdk.ChatMessagePart, - ) { - if part.ToolName != "" { - if configID, ok := toolNameToConfigID[part.ToolName]; ok { - part.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} - } - } - p.publishMessagePart(chat.ID, role, part) - }, - Compaction: compactionOptions, - ReloadMessages: func(reloadCtx context.Context) ([]fantasy.Message, error) { - reloadedMsgs, err := p.db.GetChatMessagesForPromptByChatID(reloadCtx, chat.ID) - if err != nil { - return nil, xerrors.Errorf("reload chat messages: %w", err) - } - compactionHistoryTipMessageID = 0 - if len(reloadedMsgs) > 0 { - compactionHistoryTipMessageID = reloadedMsgs[len(reloadedMsgs)-1].ID - } - if compactionOptions != nil { - compactionOptions.HistoryTipMessageID = compactionHistoryTipMessageID - } - reloadedPrompt, err := chatprompt.ConvertMessagesWithFiles(reloadCtx, reloadedMsgs, p.chatFileResolver(modelConfig.Provider), logger) - if err != nil { - return nil, xerrors.Errorf("convert reloaded messages: %w", err) - } - reloadedPrompt, sanitizeStats := chatsanitize.SanitizeAnthropicProviderToolHistory(model.Provider(), reloadedPrompt) - chatsanitize.LogAnthropicProviderToolSanitization( - reloadCtx, logger, "reload_messages", model.Provider(), model.Model(), sanitizeStats, - ) - // Re-derive instruction and skills from the reloaded - // messages so that any context added during the - // chatloop (e.g. via persistInstructionFiles when - // the agent changes) is picked up after compaction. - // The captured instruction takes priority; fall - // back to persisted DB content otherwise. - reloadedInstruction := instruction - if reloadedInstruction == "" { - reloadedInstruction = instructionFromContextFiles(reloadedMsgs) - } - if reloadedInstruction != "" { - instructionInjected = true - } - reloadedSkills := skillsFromParts(reloadedMsgs) - if len(reloadedSkills) == 0 { - reloadedSkills = workspaceSkills - } - reloadedResolvedSkills := resolvedSkillsFor(reloadedSkills) - injectedSkillIndex = chattool.FormatResolvedSkillIndex(reloadedResolvedSkills) - reloadUserPrompt := p.resolveUserPrompt(reloadCtx, chat.OwnerID) - reloadedPrompt = buildSystemPrompt( - reloadedPrompt, - subagentInstruction, - reloadedInstruction, - reloadedResolvedSkills, - reloadUserPrompt, - systemPromptBehaviorContext{ - planMode: currentPlanMode, - chatMode: chat.Mode, - planModeInstructions: planModeInstructions, - isRootChat: isRootChat, - }, - ) - // Re-inject advisor guidance after rebuilding system - // blocks so compaction/reload preserves the same - // system-message ordering as the initial prompt path. - if advisorRuntime != nil { - reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, chatadvisor.ParentGuidanceBlock) - } - reloadedPrompt = renderPlanPathPrompt(reloadedPrompt, resolvePlanPathBlock(reloadCtx)) - // Snapshot the full reloaded prompt before chain-mode - // filtering so the advisor runs with complete - // assistant/tool context. The nested advisor call - // clears previous_response_id, so provider-side - // history is unavailable. - setAdvisorPromptSnapshot(reloadedPrompt) - if chainModeActive { - reloadedPrompt = chatopenai.FilterPromptForChainMode( - reloadedPrompt, - chainInfo, - ) - } - return reloadedPrompt, nil - }, - DisableChainMode: func() { - chainModeActive = false - }, - PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { - updatedTools, toolsChanged := appendCurrentSkillTools(currentTools) - - // Mid-turn workspace MCP discovery for chats that bind a - // workspace via create_workspace or start_workspace after the - // turn has already started. The top-of-turn discovery path is - // gated on chat.WorkspaceID.Valid; this callback bridges the - // gap so the LLM sees workspace MCP tools on the very next - // step instead of the turn after. - // - // create_workspace and start_workspace prime - // workspaceMCPToolsCache via onChatUpdated after - // waitForAgentReady returns, so the call below is almost - // always a cache hit. The primer's bounded wait means the - // dial fallback here only runs when priming itself failed. - if workspaceMCPDiscovered || isExploreSubagent { - if toolsChanged { - return updatedTools - } - return nil - } - snapshot := workspaceCtx.currentChatSnapshot() - if !snapshot.WorkspaceID.Valid { - if toolsChanged { - return updatedTools - } - return nil - } - discovered := p.discoverWorkspaceMCPTools( - ctx, loopLogger, chat.ID, &workspaceCtx, - ) - if len(discovered) == 0 { - // Leave workspaceMCPDiscovered false so a subsequent - // step retries discovery. PrepareTools fires once per - // LLM step, so retries are unbounded for the rest of - // the turn. Per-step cost is one - // GetWorkspaceAgentsInLatestBuildByWorkspaceID query - // plus one ListMCPTools RPC, both fast against a live - // conn. The primer's 30s budget applies to its own - // loop only. - if toolsChanged { - return updatedTools - } - return nil - } - workspaceMCPDiscovered = true - return append(slices.Clone(updatedTools), discovered...) - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - // Skip the snapshot update when chain mode is active; - // the chatloop passes in the chain-filtered prompt - // (system plus trailing user messages) and the advisor - // needs the full pre-chain history captured at the - // initial-prompt and ReloadMessages sites. - if !chainModeActive { - setAdvisorPromptSnapshot(msgs) - } - result := msgs - changed := false - if !instructionInjected && instruction != "" { - instructionInjected = true - result = chatprompt.InsertSystem(result, instruction) - changed = true - } - if skillIndex := chattool.FormatResolvedSkillIndex(resolvedSkillsFor(workspaceSkills)); skillIndex != "" && skillIndex != injectedSkillIndex { - result = removeSkillIndexMessages(result) - result = chatprompt.InsertSystem(result, skillIndex) - injectedSkillIndex = skillIndex - changed = true - } - if !changed { - return nil - } - if !chainModeActive { - setAdvisorPromptSnapshot(result) - } - return result - }, - OnRetry: func( - attempt int, - retryErr error, - classified chatretry.ClassifiedError, - delay time.Duration, - ) { - p.clearProvisionalStreamParts(chat.ID) - logger.Warn(ctx, "retrying LLM stream", - slog.F("attempt", attempt), - slog.F("delay", delay.String()), - slog.F("kind", classified.Kind), - slog.Error(retryErr), - ) - payload := chaterror.StreamRetryPayload(attempt, delay, classified) - p.publishRetry(chat.ID, payload) - }, - - OnInterruptedPersistError: func(err error) { - p.logger.Warn(ctx, "failed to persist interrupted chat step", slog.Error(err)) - }, - }) - if errors.Is(loopErr, chatloop.ErrStopAfterTool) { - loopErr = nil - } - if errors.Is(loopErr, chatloop.ErrDynamicToolCall) { - // The stream event is published in processChat's - // defer after the DB status transitions to - // requires_action, preventing a race where a fast - // client reacts before the status is committed. - result.FinalAssistantText = finalAssistantText - result.PendingDynamicToolCalls = pendingDynamicCalls - return result, nil - } - if loopErr != nil { - classified := chaterror.Classify(loopErr).WithProvider(model.Provider()) - return result, chaterror.WithClassification(loopErr, classified) - } - result.FinalAssistantText = finalAssistantText - return result, nil -} - // buildProviderTools creates provider-native tool definitions // (like web search) based on the model configuration. These // tools are executed server-side by the LLM provider. @@ -8471,126 +4296,6 @@ func buildProviderTools(options *codersdk.ChatModelProviderOptions) []chatloop.P return tools } -// persistChatContextSummary is called from the chat loop's compaction -// callback. activeAPIKeyID is stamped onto the summary user message. When -// empty, it falls back to the delegated key in ctx. -func (p *Server) persistChatContextSummary( - ctx context.Context, - chatID uuid.UUID, - modelConfigID uuid.UUID, - activeAPIKeyID string, - toolCallID string, - result chatloop.CompactionResult, -) error { - if strings.TrimSpace(result.SystemSummary) == "" || - strings.TrimSpace(result.SummaryReport) == "" { - return nil - } - - systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(result.SystemSummary), - }) - if err != nil { - return xerrors.Errorf("encode system summary: %w", err) - } - - args, err := json.Marshal(map[string]any{ - "source": "automatic", - "threshold_percent": result.ThresholdPercent, - }) - if err != nil { - return xerrors.Errorf("encode summary tool args: %w", err) - } - - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageToolCall(toolCallID, "chat_summarized", args), - }) - if err != nil { - return xerrors.Errorf("encode summary tool call: %w", err) - } - - summaryResult, err := json.Marshal(map[string]any{ - "summary": result.SummaryReport, - "source": "automatic", - "threshold_percent": result.ThresholdPercent, - "usage_percent": result.UsagePercent, - "context_tokens": result.ContextTokens, - "context_limit_tokens": result.ContextLimit, - }) - if err != nil { - return xerrors.Errorf("encode summary result payload: %w", err) - } - toolResult, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageToolResult(toolCallID, "chat_summarized", summaryResult, false, false), - }) - if err != nil { - return xerrors.Errorf("encode summary tool result: %w", err) - } - - summaryAPIKeyID := activeAPIKeyID - if summaryAPIKeyID == "" { - summaryAPIKeyID, _ = aibridge.DelegatedAPIKeyIDFromContext(ctx) - } - - var insertedMessages []database.ChatMessage - - txErr := p.db.InTx(func(tx database.Store) error { - summaryParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. - ChatID: chatID, - } - - // Hidden summary user message (not published to subscribers). - summaryUserMsg := newUserChatMessage( - summaryAPIKeyID, - systemContent, - database.ChatMessageVisibilityModel, - modelConfigID, - chatprompt.CurrentContentVersion, - ) - summaryUserMsg = summaryUserMsg.withCompressed() - appendUserChatMessage(&summaryParams, summaryUserMsg) - - // Assistant tool-call message. - appendChatMessage(&summaryParams, newChatMessage( - database.ChatMessageRoleAssistant, - assistantContent, - database.ChatMessageVisibilityUser, - modelConfigID, - chatprompt.CurrentContentVersion, - ).withCompressed()) - - // Tool result message. - appendChatMessage(&summaryParams, newChatMessage( - database.ChatMessageRoleTool, - toolResult, - database.ChatMessageVisibilityBoth, - modelConfigID, - chatprompt.CurrentContentVersion, - ).withCompressed()) - - allInserted, txErr := tx.InsertChatMessages(ctx, summaryParams) - if txErr != nil { - return xerrors.Errorf("insert summary messages: %w", txErr) - } - // Skip the first message (hidden summary user msg) when - // publishing — only the assistant and tool messages are - // visible to subscribers. - insertedMessages = allInserted[1:] - - return nil - }, nil) - if txErr != nil { - return txErr - } - - // Publish after transaction commits to avoid notifying - // subscribers about messages that could be rolled back. - for _, msg := range insertedMessages { - p.publishMessage(chatID, msg) - } - return nil -} - func (p *Server) resolveChatModel( ctx context.Context, chat database.Chat, @@ -9000,7 +4705,7 @@ func (p *Server) fetchWorkspaceContext( // Stamp server-side fields and sanitize content. The // agent cannot know its own UUID, OS metadata, or - // directory — those are added here at the trust boundary. + // directory, those are added here at the trust boundary. agentID := uuid.NullUUID{UUID: loadedAgent.ID, Valid: true} for i := range agentParts { @@ -9023,6 +4728,16 @@ func (p *Server) fetchWorkspaceContext( return &loadedAgent, agentParts, discoveredSkills, workspaceConnOK } +func filterSkillParts(parts []codersdk.ChatMessagePart) []codersdk.ChatMessagePart { + var filtered []codersdk.ChatMessagePart + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeSkill { + filtered = append(filtered, part) + } + } + return filtered +} + // persistInstructionFiles fetches AGENTS.md instruction files and // skills from the workspace agent, persisting both as message // parts. This is called once when a workspace is first attached @@ -9039,10 +4754,6 @@ func (p *Server) persistInstructionFiles( agent, agentParts, discoveredSkills, workspaceConnOK := p.fetchWorkspaceContext( ctx, chat, getWorkspaceAgent, getWorkspaceConn, ) - // Defensive guard: fetchWorkspaceContext returns nil when the - // chat has no valid workspace or the agent lookup fails. It's - // cheaper to guard here than push the precondition up to all - // callers. if agent == nil { return "", nil, nil } @@ -9063,12 +4774,11 @@ func (p *Server) persistInstructionFiles( directory = agent.Directory } + contextAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) if !hasContent { if !workspaceConnOK { return "", nil, nil } - // Persist a blank context-file marker (plus any skill-only - // parts) so subsequent turns skip the workspace agent dial. if !hasContextFilePart { agentParts = append([]codersdk.ChatMessagePart{{ Type: codersdk.ChatMessagePartTypeContextFile, @@ -9079,7 +4789,6 @@ func (p *Server) persistInstructionFiles( if err != nil { return "", nil, nil } - contextAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chat.ID, } @@ -9091,9 +4800,6 @@ func (p *Server) persistInstructionFiles( chatprompt.CurrentContentVersion, )) _, _ = p.db.InsertChatMessages(ctx, msgParams) - // Update the cache column: persist skills if any - // exist, or clear to NULL so stale data from a - // previous agent doesn't linger. skillParts := filterSkillParts(agentParts) p.updateLastInjectedContext(ctx, chat.ID, skillParts) return "", discoveredSkills, nil @@ -9103,7 +4809,6 @@ func (p *Server) persistInstructionFiles( return "", nil, xerrors.Errorf("marshal context-file parts: %w", err) } - contextAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. ChatID: chat.ID, } @@ -9117,9 +4822,6 @@ func (p *Server) persistInstructionFiles( if _, err := p.db.InsertChatMessages(ctx, msgParams); err != nil { return "", nil, xerrors.Errorf("persist instruction files: %w", err) } - // Build stripped copies for the cache column so internal - // fields (full file content, OS, directory, skill paths) - // are never persisted or returned to API clients. stripped := make([]codersdk.ChatMessagePart, len(agentParts)) copy(stripped, agentParts) for i := range stripped { @@ -9127,16 +4829,13 @@ func (p *Server) persistInstructionFiles( } p.updateLastInjectedContext(ctx, chat.ID, stripped) - // Return the formatted instruction text and discovered skills - // so the caller can inject them into this turn's prompt (since - // the prompt was built before we persisted). return formatSystemInstructions(agent.OperatingSystem, directory, agentParts), discoveredSkills, nil } // updateLastInjectedContext persists the injected context // parts (AGENTS.md files and skills) on the chat row so they // are directly queryable without scanning messages. This is -// best-effort — a failure here is logged but does not block +// best-effort, a failure here is logged but does not block // the turn. func (p *Server) updateLastInjectedContext(ctx context.Context, chatID uuid.UUID, parts []codersdk.ChatMessagePart) { param := pqtype.NullRawMessage{Valid: false} @@ -9320,281 +5019,6 @@ func formatPlanPathBlock(chatPath, home string) string { return b.String() } -func (p *Server) recoverStaleChats(ctx context.Context) { - staleAfter := p.clock.Now().Add(-p.inFlightChatStaleAfter) - staleChats, err := p.db.GetStaleChats(ctx, staleAfter) - if err != nil { - p.logger.Error(ctx, "failed to get stale chats", slog.Error(err)) - return - } - - recovered := 0 - for _, chat := range staleChats { - p.logger.Info(ctx, "recovering stale chat", - slog.F("chat_id", chat.ID), - slog.F("status", chat.Status)) - - // Use a transaction with FOR UPDATE to avoid a TOCTOU race: - // between GetStaleChats (a bare SELECT) and here, the chat's - // heartbeat may have been refreshed. We re-check freshness - // under the row lock before resetting. - err := p.db.InTx(func(tx database.Store) error { - locked, lockErr := tx.GetChatByIDForUpdate(ctx, chat.ID) - if lockErr != nil { - return xerrors.Errorf("lock chat for recovery: %w", lockErr) - } - - switch locked.Status { - case database.ChatStatusRunning: - // Re-check: only recover if the chat is still stale. - // A valid heartbeat at or after the threshold means - // the chat was refreshed after our snapshot. - if locked.HeartbeatAt.Valid && !locked.HeartbeatAt.Time.Before(staleAfter) { - p.logger.Debug(ctx, "chat heartbeat refreshed since snapshot, skipping recovery", - slog.F("chat_id", chat.ID)) - return nil - } - case database.ChatStatusRequiresAction: - // Re-check: the chat may have been updated after - // our snapshot, similar to the heartbeat check for - // running chats. - if !locked.UpdatedAt.Before(staleAfter) { - p.logger.Debug(ctx, "chat updated since snapshot, skipping recovery", - slog.F("chat_id", chat.ID)) - return nil - } - case database.ChatStatusWaiting: - // Deferred-promote stranding: worker died before its - // post-cancel cleanup ran. Re-check freshness. - if !locked.UpdatedAt.Before(staleAfter) { - p.logger.Debug(ctx, "chat updated since snapshot, skipping recovery", - slog.F("chat_id", chat.ID)) - return nil - } - default: - // Status changed since our snapshot; skip. - p.logger.Debug(ctx, "chat status changed since snapshot, skipping recovery", - slog.F("chat_id", chat.ID), - slog.F("status", locked.Status)) - return nil - } - - lastError := pqtype.NullRawMessage{} - if locked.Status == database.ChatStatusRequiresAction { - lastErrorPayload, marshalErr := encodeChatLastErrorPayload( - chaterror.TerminalErrorPayload(chaterror.ClassifiedError{ - Message: "Dynamic tool execution timed out", - Kind: codersdk.ChatErrorKindGeneric, - }), - ) - if marshalErr != nil { - p.logger.Warn(ctx, "failed to marshal stale recovery last error payload", - slog.F("chat_id", chat.ID), - slog.Error(marshalErr), - ) - } else { - lastError = lastErrorPayload - } - } - - recoverStatus := database.ChatStatusPending - if locked.Status == database.ChatStatusRequiresAction { - // Timed-out requires_action chats have dangling - // tool calls with no matching results. Setting - // them back to pending would replay incomplete - // tool calls to the LLM, so mark them as errors. - recoverStatus = database.ChatStatusError - } - - // Insert synthetic error tool-result messages - // so the LLM history remains valid if the user - // retries the chat later. - if locked.Status == database.ChatStatusRequiresAction { - if _, synthErr := insertSyntheticToolResultsTx(ctx, tx, locked, "Dynamic tool execution timed out"); synthErr != nil { - p.logger.Warn(ctx, "failed to insert synthetic tool results during stale recovery", - slog.F("chat_id", chat.ID), - slog.Error(synthErr), - ) - // Continue with error status even if - // synthetic results fail to insert. - } - } - - if locked.Status == database.ChatStatusWaiting { - // Close pending dynamic tool calls; otherwise the - // promoted user message would feed the LLM a turn it - // rejects. Propagate errors so the next recovery - // tick retries instead of promoting incomplete - // history. - if _, synthErr := insertSyntheticToolResultsTx(ctx, tx, locked, "Tool execution interrupted by queued message promotion"); synthErr != nil { - return xerrors.Errorf("insert synthetic tool results during stale recovery: %w", synthErr) - } - promoted, _, _, promoteErr := p.tryAutoPromoteQueuedMessage(ctx, tx, locked) - if promoteErr != nil { - return xerrors.Errorf("auto-promote during stale recovery: %w", promoteErr) - } - if promoted == nil { - // Empty queue means nothing to recover. - return nil - } - } - - // Reset so any replica can pick it up (pending) or - // the client sees the failure (error). - _, updateErr := tx.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: recoverStatus, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: lastError, - }) - if updateErr != nil { - return updateErr - } - recovered++ - return nil - }, nil) - if err != nil { - p.logger.Error(ctx, "failed to recover stale chat", - slog.F("chat_id", chat.ID), slog.Error(err)) - } - } - - if recovered > 0 { - p.logger.Info(ctx, "recovered stale chats", slog.F("count", recovered)) - } -} - -// insertSyntheticToolResultsTx inserts IsError tool-result messages -// for unresolved dynamic tool calls in the last assistant message, -// skipping calls already handled (e.g. by chatloop dispatching a -// name-colliding dynamic tool as a built-in). It operates on the -// provided store, which may be a transaction handle. -func insertSyntheticToolResultsTx( - ctx context.Context, - store database.Store, - chat database.Chat, - reason string, -) ([]database.ChatMessage, error) { - dynamicToolNames, err := parseDynamicToolNames(chat.DynamicTools) - if err != nil { - return nil, xerrors.Errorf("parse dynamic tools: %w", err) - } - if len(dynamicToolNames) == 0 { - return nil, nil - } - - // No assistant means nothing to close: a deferred promote can - // race a worker that fails before any persist, and the cleanup - // TX must still advance. - lastAssistant, err := store.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ - ChatID: chat.ID, - Role: database.ChatMessageRoleAssistant, - }) - if errors.Is(err, sql.ErrNoRows) { - return nil, nil - } - if err != nil { - return nil, xerrors.Errorf("get last assistant message: %w", err) - } - - parts, err := chatprompt.ParseContent(lastAssistant) - if err != nil { - return nil, xerrors.Errorf("parse assistant message: %w", err) - } - - // Mirrors SubmitToolResults. - afterMsgs, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: lastAssistant.ID, - }) - if err != nil { - return nil, xerrors.Errorf("get messages after assistant: %w", err) - } - handledCallIDs := make(map[string]bool) - for _, msg := range afterMsgs { - if msg.Role != database.ChatMessageRoleTool { - continue - } - msgParts, err := chatprompt.ParseContent(msg) - if err != nil { - continue - } - for _, mp := range msgParts { - if mp.Type == codersdk.ChatMessagePartTypeToolResult { - handledCallIDs[mp.ToolCallID] = true - } - } - } - - // Collect dynamic tool calls that need synthetic results. - var resultContents []pqtype.NullRawMessage - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeToolCall || !dynamicToolNames[part.ToolName] { - continue - } - if handledCallIDs[part.ToolCallID] { - continue - } - resultPart := codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: part.ToolCallID, - ToolName: part.ToolName, - Result: json.RawMessage(fmt.Sprintf("%q", reason)), - IsError: true, - } - marshaled, marshalErr := chatprompt.MarshalParts([]codersdk.ChatMessagePart{resultPart}) - if marshalErr != nil { - return nil, xerrors.Errorf("marshal synthetic tool result: %w", marshalErr) - } - resultContents = append(resultContents, marshaled) - } - - if len(resultContents) == 0 { - return nil, nil - } - - // Insert tool-result messages using the same pattern as - // SubmitToolResults. - n := len(resultContents) - params := database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: make([]uuid.UUID, n), - APIKeyID: make([]string, n), - ModelConfigID: make([]uuid.UUID, n), - Role: make([]database.ChatMessageRole, n), - Content: make([]string, n), - ContentVersion: make([]int16, n), - Visibility: make([]database.ChatMessageVisibility, n), - InputTokens: make([]int64, n), - OutputTokens: make([]int64, n), - TotalTokens: make([]int64, n), - ReasoningTokens: make([]int64, n), - CacheCreationTokens: make([]int64, n), - CacheReadTokens: make([]int64, n), - ContextLimit: make([]int64, n), - Compressed: make([]bool, n), - TotalCostMicros: make([]int64, n), - RuntimeMs: make([]int64, n), - ProviderResponseID: make([]string, n), - } - for i, rc := range resultContents { - params.CreatedBy[i] = uuid.Nil - params.ModelConfigID[i] = chat.LastModelConfigID - params.Role[i] = database.ChatMessageRoleTool - params.Content[i] = string(rc.RawMessage) - params.ContentVersion[i] = chatprompt.CurrentContentVersion - params.Visibility[i] = database.ChatMessageVisibilityBoth - } - inserted, err := store.InsertChatMessages(ctx, params) - if err != nil { - return nil, xerrors.Errorf("insert synthetic tool results: %w", err) - } - - return inserted, nil -} - // parseDynamicToolNames unmarshals the dynamic tools JSON column // and returns a map of tool names. This centralizes the repeated // pattern of deserializing DynamicTools into a name set. @@ -9686,7 +5110,7 @@ func (p *Server) finalizeSuccessfulTurnStatusLabelWithAfterFunc( slog.F("label_length", len(statusLabel)), ) - p.updateLastTurnSummary(finalizeCtx, chat, chat.UpdatedAt, statusLabel, logger) + p.updateLastTurnSummary(finalizeCtx, chat, chat.HistoryVersion, statusLabel, logger) afterFinalize(finalizeCtx, statusLabel) }) @@ -9775,7 +5199,7 @@ func (p *Server) setLastTurnSummaryAsync( // still counted in p.inflight. Do not take inflightMu here because // drainInflight holds it while waiting. p.inflight.Go(func() { - p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.UpdatedAt, summary, logger) + p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, summary, logger) }) } @@ -9784,14 +5208,11 @@ func (p *Server) clearLastTurnSummaryAsync( chat database.Chat, logger slog.Logger, ) { - if !chat.LastTurnSummary.Valid { - return - } // This helper runs during processChat cleanup, while processChat is // still counted in p.inflight. Do not take inflightMu here because // drainInflight holds it while waiting. p.inflight.Go(func() { - p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.UpdatedAt, "", logger) + p.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, "", logger) }) } @@ -9801,7 +5222,7 @@ func (p *Server) clearLastTurnSummaryAsync( func (p *Server) updateLastTurnSummary( ctx context.Context, chat database.Chat, - expectedUpdatedAt time.Time, + expectedHistoryVersion int64, summary string, logger slog.Logger, ) { @@ -9814,9 +5235,9 @@ func (p *Server) updateLastTurnSummary( defer cancel() affected, err := p.db.UpdateChatLastTurnSummary(updateCtx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: expectedUpdatedAt, - LastTurnSummary: lastTurnSummary, + ID: chat.ID, + ExpectedHistoryVersion: expectedHistoryVersion, + LastTurnSummary: lastTurnSummary, }) if err != nil { logger.Warn(updateCtx, "failed to update chat turn summary", @@ -9830,13 +5251,13 @@ func (p *Server) updateLastTurnSummary( logger.Info(updateCtx, "skipped stale chat turn summary update with non-empty summary", slog.F("chat_id", chat.ID), slog.F("summary_length", len(summary)), - slog.F("expected_updated_at", expectedUpdatedAt), + slog.F("expected_history_version", expectedHistoryVersion), ) return } logger.Debug(updateCtx, "skipped stale chat turn summary update", slog.F("chat_id", chat.ID), - slog.F("expected_updated_at", expectedUpdatedAt), + slog.F("expected_history_version", expectedHistoryVersion), ) return } @@ -9844,10 +5265,6 @@ func (p *Server) updateLastTurnSummary( updatedChat := chat updatedChat.LastTurnSummary = lastTurnSummary p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindSummaryChange, nil) - - // AcquireChats uses SKIP LOCKED; re-wake so a wake racing this - // UPDATE's row lock does not strand a freshly-pending chat. - p.signalWake() } func (p *Server) webpushConfigured() bool { @@ -9882,6 +5299,17 @@ func (p *Server) Close() error { p.configCacheUnsubscribe = nil unsub() } + if p.chatWorker != nil { + if err := p.chatWorker.Close(); err != nil { + p.logger.Warn(context.Background(), "failed to close chat worker", slog.Error(err)) + } + } + if p.streamSyncPoller != nil { + p.streamSyncPoller.Close() + } + if p.messagePartBuffer != nil { + p.messagePartBuffer.Close() + } p.cancel() p.wg.Wait() p.drainInflight() diff --git a/coderd/x/chatd/chatd_chainmode_test.go b/coderd/x/chatd/chatd_chainmode_test.go new file mode 100644 index 0000000000..b81354d1fa --- /dev/null +++ b/coderd/x/chatd/chatd_chainmode_test.go @@ -0,0 +1,573 @@ +package chatd_test + +import ( + "context" + "encoding/json" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + + "charm.land/fantasy" + fantasyanthropic "charm.land/fantasy/providers/anthropic" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +func TestActiveServer_ChainBrokenRecovery(t *testing.T) { + t.Parallel() + + const ( + previousResponseID = "resp_poisoned" + recoveredAnswer = "recovered answer" + ) + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newOpenAIRequestRecorder() + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + requests.record(req) + if req.PreviousResponseID != nil { + return chattest.OpenAIErrorResponse(http.StatusNotFound, "invalid_request_error", chainBrokenProviderErrorMessage) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks(recoveredAnswer)...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "first user") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderResponseID(ctx, t, db, chat.ID, "first assistant", model.ID, previousResponseID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("follow up")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + got := requests.all() + require.GreaterOrEqual(t, len(got), 3) + generationRequests := filterStreamingRequests(got) + require.Len(t, generationRequests, 3) + require.Nil(t, generationRequests[0].PreviousResponseID) + require.Equal(t, previousResponseID, requirePreviousResponseID(t, generationRequests[1])) + require.Nil(t, generationRequests[2].PreviousResponseID) + requireRawPromptContains(t, generationRequests[2], "first user") + requireRawPromptContains(t, generationRequests[2], "first assistant") + requireRawPromptContains(t, generationRequests[2], "follow up") + + messages := chatMessages(ctx, t, db, chat.ID) + requireTextPart(t, messages[len(messages)-1], recoveredAnswer) +} + +func TestActiveServer_ChainBrokenRecoveryAppliesProviderPromptPrep(t *testing.T) { + t.Parallel() + + const previousResponseID = "resp_anthropic_chain" + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var streamCalls atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + if streamCalls.Add(1) == 2 { + return chattest.AnthropicErrorResponse(http.StatusInternalServerError, "server_error", chainBrokenProviderErrorMessage) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("anthropic answer")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertSystemTextMessage(ctx, t, db, chat.ID, "sys-1", model.ID) + insertProviderResponseID(ctx, t, db, chat.ID, "hi", model.ID, previousResponseID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("follow up")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + recovered := generationRequests[1] + require.Len(t, recovered.Messages, 4) + require.True(t, anthropicSystemHasEphemeralCacheControl(t, recovered)) + require.False(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[0])) + require.False(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[1])) + require.True(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[2])) + require.True(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[3])) +} + +func TestActiveServer_NonChainBrokenRetryPreservesChainMode(t *testing.T) { + t.Parallel() + + const previousResponseID = "resp_still_valid" + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newOpenAIRequestRecorder() + var streamCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + requests.record(req) + if req.Stream && streamCalls.Add(1) == 2 { + return chattest.OpenAIServerErrorResponse() + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("answer")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "first user") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderResponseID(ctx, t, db, chat.ID, "first assistant", model.ID, previousResponseID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("follow up")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterStreamingRequests(requests.all()) + require.Len(t, generationRequests, 3) + require.Equal(t, previousResponseID, requirePreviousResponseID(t, generationRequests[1])) + require.Equal(t, previousResponseID, requirePreviousResponseID(t, generationRequests[2])) + requireRawPromptNotContains(t, generationRequests[2], "first user") + requireRawPromptContains(t, generationRequests[2], "follow up") +} + +func TestActiveServer_ChainBrokenRecoveryPersistsAcrossGenerationActions(t *testing.T) { + t.Parallel() + + const previousResponseID = "resp_tool_poisoned" + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newOpenAIRequestRecorder() + var streamCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + requests.record(req) + if !req.Stream { + return chattest.OpenAINonStreamingResponse(`{"title":"test"}`) + } + switch streamCalls.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("first answer")...) + case 2: + return chattest.OpenAIErrorResponse(http.StatusNotFound, "invalid_request_error", chainBrokenProviderErrorMessage) + case 3: + return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk("read_skill", `{"name":"x"}`)) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("final answer")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "first user") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderResponseID(ctx, t, db, chat.ID, "first assistant", model.ID, previousResponseID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("follow up")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterStreamingRequests(requests.all()) + require.Len(t, generationRequests, 4) + require.Equal(t, previousResponseID, requirePreviousResponseID(t, generationRequests[1])) + require.Nil(t, generationRequests[2].PreviousResponseID) + require.Nil(t, generationRequests[3].PreviousResponseID) +} + +func TestActiveServer_ChainBrokenWithoutChainModeIsSafe(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newOpenAIRequestRecorder() + var streamCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + requests.record(req) + if req.Stream && streamCalls.Add(1) == 1 { + return chattest.OpenAIErrorResponse(http.StatusNotFound, "invalid_request_error", chainBrokenProviderErrorMessage) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("recovered")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "only user") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + require.Nil(t, generationRequests[0].PreviousResponseID) + require.Nil(t, generationRequests[1].PreviousResponseID) +} + +func TestActiveServer_ChainBrokenRecoveryDropsOrphanProviderToolCall(t *testing.T) { + t.Parallel() + + const previousResponseID = "resp_orphan_provider_tool" + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var streamCalls atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + if streamCalls.Add(1) == 2 { + return chattest.AnthropicErrorResponse(http.StatusInternalServerError, "server_error", chainBrokenProviderErrorMessage) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("cleaned")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateModelForChainMode(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "first user") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderResponseID(ctx, t, db, chat.ID, "first assistant", model.ID, previousResponseID) + insertOrphanProviderToolCall(ctx, t, db, chat.ID, model.ID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + recoveredBody := anthropicRequestBody(t, generationRequests[1]) + require.NotContains(t, recoveredBody, "web_search") + require.Contains(t, recoveredBody, "partial") + require.Contains(t, recoveredBody, "continue") + requireAnthropicRequestRedactedReasoning(t, generationRequests[1], "redacted-payload") +} + +type anthropicRequestRecorder struct { + mu sync.Mutex + requests []chattest.AnthropicRequest +} + +func newAnthropicRequestRecorder() *anthropicRequestRecorder { + return &anthropicRequestRecorder{} +} + +func (r *anthropicRequestRecorder) record(req *chattest.AnthropicRequest) { + r.mu.Lock() + defer r.mu.Unlock() + r.requests = append(r.requests, *req) +} + +func (r *anthropicRequestRecorder) all() []chattest.AnthropicRequest { + r.mu.Lock() + defer r.mu.Unlock() + return append([]chattest.AnthropicRequest(nil), r.requests...) +} + +func filterAnthropicStreamingRequests(requests []chattest.AnthropicRequest) []chattest.AnthropicRequest { + out := make([]chattest.AnthropicRequest, 0, len(requests)) + for _, req := range requests { + if req.Stream { + out = append(out, req) + } + } + return out +} + +func seedAnthropicChatDependencies(t *testing.T, db database.Store, baseURL string) (database.User, database.Organization, database.ChatModelConfig) { + t.Helper() + user := dbgen.User(t, db, database.User{}) + _ = testAPIKeyID(t, db, user.ID) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + provider := dbgen.AIProvider(t, db, database.AIProvider{Type: database.AiProviderTypeAnthropic}, func(params *database.InsertAIProviderParams) { + params.BaseUrl = baseURL + }) + dbgen.AIProviderKey(t, db, database.AIProviderKey{ProviderID: provider.ID}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + }) + return user, org, model +} + +func anthropicSystemHasEphemeralCacheControl(t *testing.T, req chattest.AnthropicRequest) bool { + t.Helper() + return strings.Contains(string(req.System), `"cache_control":{"type":"ephemeral"}`) +} + +func anthropicMessageHasEphemeralCacheControl(t *testing.T, message chattest.AnthropicRequestMessage) bool { + t.Helper() + return strings.Contains(string(message.Content), `"cache_control":{"type":"ephemeral"}`) +} + +func anthropicRequestBody(t *testing.T, req chattest.AnthropicRequest) string { + t.Helper() + data, err := json.Marshal(req.Messages) + require.NoError(t, err) + return string(data) +} + +func insertSystemTextMessage( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + text string, + modelID uuid.UUID, +) { + t.Helper() + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + params := chatd.BuildSingleChatMessageInsertParams( + chatID, + database.ChatMessageRoleSystem, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + uuid.Nil, + ) + _, err = db.InsertChatMessages(ctx, params) + require.NoError(t, err) +} + +func requireAnthropicRequestRedactedReasoning(t *testing.T, req chattest.AnthropicRequest, redactedData string) { + t.Helper() + body := anthropicRequestBody(t, req) + require.Contains(t, body, "redacted-payload") + require.Contains(t, body, redactedData) +} + +func insertOrphanProviderToolCall(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, modelID uuid.UUID) { + t.Helper() + reasoningMetadata, err := json.Marshal(fantasy.ProviderMetadata{ + fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{RedactedData: "redacted-payload"}, + }) + require.NoError(t, err) + parts := []codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeReasoning, + ProviderMetadata: reasoningMetadata, + }, + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "ws-orphan", + ToolName: "web_search", + Args: json.RawMessage(`{"query":"coder"}`), + ProviderExecuted: true, + }, + codersdk.ChatMessageText("partial"), + } + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + params := chatd.BuildSingleChatMessageInsertParams( + chatID, + database.ChatMessageRoleAssistant, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + uuid.Nil, + ) + _, err = db.InsertChatMessages(ctx, params) + require.NoError(t, err) +} + +const chainBrokenProviderErrorMessage = "Previous response with id 'resp_abc' not found." + +type openAIRequestRecorder struct { + mu sync.Mutex + requests []chattest.OpenAIRequest +} + +func newOpenAIRequestRecorder() *openAIRequestRecorder { + return &openAIRequestRecorder{} +} + +func (r *openAIRequestRecorder) record(req *chattest.OpenAIRequest) { + r.mu.Lock() + defer r.mu.Unlock() + r.requests = append(r.requests, *req) +} + +func (r *openAIRequestRecorder) all() []chattest.OpenAIRequest { + r.mu.Lock() + defer r.mu.Unlock() + return append([]chattest.OpenAIRequest(nil), r.requests...) +} + +func updateModelForChainMode(t *testing.T, db database.Store, model database.ChatModelConfig) database.ChatModelConfig { + t.Helper() + store := true + options, err := json.Marshal(codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + OpenAI: &codersdk.ChatModelOpenAIProviderOptions{Store: &store}, + }, + }) + require.NoError(t, err) + updated, err := db.UpdateChatModelConfig(context.Background(), database.UpdateChatModelConfigParams{ + ID: model.ID, + DisplayName: model.DisplayName, + Model: model.Model, + Provider: model.Provider, + Enabled: model.Enabled, + ContextLimit: model.ContextLimit, + CompressionThreshold: model.CompressionThreshold, + Options: options, + AIProviderID: model.AIProviderID, + }) + require.NoError(t, err) + return updated +} + +func createChatThroughServer( + ctx context.Context, + t *testing.T, + db database.Store, + server *chatd.Server, + orgID uuid.UUID, + userID uuid.UUID, + modelID uuid.UUID, + text string, +) database.Chat { + t.Helper() + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: orgID, + OwnerID: userID, + APIKeyID: testAPIKeyID(t, db, userID), + Title: "chain mode test", + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}, + ModelConfigID: modelID, + }) + require.NoError(t, err) + return chat +} + +func waitForChatStatus(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID, status database.ChatStatus) database.Chat { + t.Helper() + var chat database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + latest, err := db.GetChatByID(ctx, chatID) + if err != nil { + return false + } + chat = latest + return latest.Status == status && !latest.WorkerID.Valid && !latest.RunnerID.Valid + }, testutil.IntervalFast) + return chat +} + +func insertProviderResponseID( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + text string, + modelID uuid.UUID, + providerResponseID string, +) { + t.Helper() + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + params := chatd.BuildSingleChatMessageInsertParams( + chatID, + database.ChatMessageRoleAssistant, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + uuid.Nil, + ) + params.ProviderResponseID[0] = providerResponseID + _, err = db.InsertChatMessages(ctx, params) + require.NoError(t, err) +} + +func chatMessages(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID) []database.ChatMessage { + t.Helper() + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chatID}) + require.NoError(t, err) + return messages +} + +func filterStreamingRequests(requests []chattest.OpenAIRequest) []chattest.OpenAIRequest { + out := make([]chattest.OpenAIRequest, 0, len(requests)) + for _, req := range requests { + if req.Stream { + out = append(out, req) + } + } + return out +} + +func requirePreviousResponseID(t *testing.T, req chattest.OpenAIRequest) string { + t.Helper() + require.NotNil(t, req.PreviousResponseID) + return *req.PreviousResponseID +} + +func requireRawPromptContains(t *testing.T, req chattest.OpenAIRequest, text string) { + t.Helper() + require.Contains(t, string(req.RawBody), text) +} + +func requireRawPromptNotContains(t *testing.T, req chattest.OpenAIRequest, text string) { + t.Helper() + require.NotContains(t, string(req.RawBody), text) +} + +func requireTextPart(t *testing.T, msg database.ChatMessage, text string) { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText && part.Text == text { + return + } + } + t.Fatalf("missing text part %q in message %d", text, msg.ID) +} diff --git a/coderd/x/chatd/chatd_debug.go b/coderd/x/chatd/chatd_debug.go index 79dc419418..bbb66cb82a 100644 --- a/coderd/x/chatd/chatd_debug.go +++ b/coderd/x/chatd/chatd_debug.go @@ -20,6 +20,9 @@ const ( // best-effort, so the turn proceeds without debug rows if the // DB is slow or locked. Matches the manual-title budget. debugCreateRunTimeout = 5 * time.Second + // debugFinalizeTimeout caps best-effort debug run finalization + // outside the runner's canceled context. + debugFinalizeTimeout = 5 * time.Second // debugCleanupClockSkew gives cleanup cutoffs tolerance for cross- // replica clock drift. The cutoff is sampled from the DB // (updated_at returned by the status transition), and diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index b8d9766b9b..61864c2894 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -25,10 +25,11 @@ import ( "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" - "github.com/coder/coder/v2/coderd/x/chatd/chaterror" + "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/coderd/x/chatd/chatloop" openaicomputeruse "github.com/coder/coder/v2/coderd/x/chatd/chatopenai/computeruse" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" @@ -242,6 +243,10 @@ func TestAppendComputerUseProviderTool(t *testing.T) { fantasy.NewImageResponse([]byte("png"), "image/png"), ) require.NotNil(t, metadata) + + errorResponse := fantasy.NewTextErrorResponse("failed") + require.Nil(t, providerTools[0].ResultProviderMetadata(errorResponse)) + require.Nil(t, providerTools[0].ResultProviderMetadata(fantasy.NewTextResponse("not media"))) } func TestAppendComputerUseProviderTool_Gates(t *testing.T) { @@ -2114,615 +2119,6 @@ func TestTurnWorkspaceContext_EnsureWorkspaceAgentIgnoresCachedAgentForDifferent require.Equal(t, updatedChat, currentChat) } -func TestSubscribeDedupesLocallyDeliveredMessageOnNotifyCatchup(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - initialMessage := database.ChatMessage{ - ID: 1, - ChatID: chatID, - Role: database.ChatMessageRoleUser, - } - localMessage := database.ChatMessage{ - ID: 2, - ChatID: chatID, - Role: database.ChatMessageRoleAssistant, - } - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{initialMessage}, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - // DB catchup runs unconditionally on every notify; the delivered - // set dedupes against locally-delivered messages. - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 1, - }).Return(nil, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - server.publishMessage(chatID, localMessage) - - event := requireStreamMessageEvent(t, events) - require.Equal(t, int64(2), event.Message.ID) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeUsesDurableCacheWhenLocalMessageWasNotDelivered(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - initialMessage := database.ChatMessage{ - ID: 1, - ChatID: chatID, - Role: database.ChatMessageRoleUser, - } - cachedMessage := codersdk.ChatMessage{ - ID: 2, - ChatID: chatID, - Role: codersdk.ChatMessageRoleAssistant, - } - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{initialMessage}, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - // DB catchup runs unconditionally; cached id=2 is deduped via - // the delivered set so this query returning nil is sufficient. - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 1, - }).Return(nil, nil), - ) - - server := newSubscribeTestServer(t, db) - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, - ChatID: chatID, - Message: &cachedMessage, - }) - - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - AfterMessageID: 1, - }) - - event := requireStreamMessageEvent(t, events) - require.Equal(t, int64(2), event.Message.ID) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeQueriesDatabaseWhenDurableCacheMisses(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - initialMessage := database.ChatMessage{ - ID: 1, - ChatID: chatID, - Role: database.ChatMessageRoleUser, - } - catchupMessage := database.ChatMessage{ - ID: 2, - ChatID: chatID, - Role: database.ChatMessageRoleAssistant, - } - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{initialMessage}, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 1, - }).Return([]database.ChatMessage{catchupMessage}, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - AfterMessageID: 1, - }) - - event := requireStreamMessageEvent(t, events) - require.Equal(t, int64(2), event.Message.ID) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeFullRefreshStillUsesDatabaseCatchup(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - initialMessage := database.ChatMessage{ - ID: 1, - ChatID: chatID, - Role: database.ChatMessageRoleUser, - } - editedMessage := database.ChatMessage{ - ID: 1, - ChatID: chatID, - Role: database.ChatMessageRoleUser, - } - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{initialMessage}, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{editedMessage}, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - server.publishEditedMessage(chatID, editedMessage) - - event := requireStreamMessageEvent(t, events) - require.Equal(t, int64(1), event.Message.ID) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeDeliversRetryEventViaPubsubOnce(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - expected := newTestRetryPayload() - - server.publishRetry(chatID, expected) - - event := requireStreamRetryEvent(t, events) - require.Equal(t, expected, event.Retry) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeReplaysCurrentRetryPhaseInSnapshot(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newBufferedSubscribeTestServer(t, db, chatID) - - expected := newTestRetryPayload() - server.publishRetry(chatID, expected) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - require.Len(t, snapshot, 2) - require.Equal(t, codersdk.ChatStreamEventTypeStatus, snapshot[0].Type) - require.Equal(t, codersdk.ChatStreamEventTypeRetry, snapshot[1].Type) - event := requireSnapshotRetryEvent(t, snapshot) - require.Equal(t, expected, event.Retry) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeCapturesRetryPhaseAtSubscriptionBoundary(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - expected := newTestRetryPayload() - - server := newSubscribeTestServer(t, db) - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).DoAndReturn(func(context.Context, database.GetChatMessagesByChatIDParams) ([]database.ChatMessage, error) { - server.publishRetry(chatID, expected) - return nil, nil - }), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - requireNoSnapshotRetryEvent(t, snapshot) - event := requireStreamRetryEvent(t, events) - require.Equal(t, expected, event.Retry) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeDoesNotReplayRetryAfterStreamResumes(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newBufferedSubscribeTestServer(t, db, chatID) - - server.publishRetry(chatID, newTestRetryPayload()) - server.publishMessagePart(chatID, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("retry recovered")) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - requireNoSnapshotRetryEvent(t, snapshot) - requireSnapshotMessagePartEvent(t, snapshot) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeDoesNotReplayFailedAttemptPartsAfterRetry(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newBufferedSubscribeTestServer(t, db, chatID) - - server.publishMessagePart(chatID, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("failed partial")) - server.clearProvisionalStreamParts(chatID) - server.publishRetry(chatID, newTestRetryPayload()) - server.publishMessagePart(chatID, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("retry recovered")) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - requireNoSnapshotRetryEvent(t, snapshot) - partEvent := requireSnapshotMessagePartEvent(t, snapshot) - require.Equal(t, "retry recovered", partEvent.MessagePart.Part.Text) - for _, event := range snapshot { - if event.Type != codersdk.ChatStreamEventTypeMessagePart { - continue - } - require.NotEqual(t, "failed partial", event.MessagePart.Part.Text) - } - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeDoesNotReplayRetryAfterTerminalError(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newBufferedSubscribeTestServer(t, db, chatID) - - server.publishRetry(chatID, newTestRetryPayload()) - server.publishError(chatID, chaterror.ClassifiedError{ - Message: "OpenAI is rate limiting requests.", - Kind: codersdk.ChatErrorKindRateLimit, - Provider: "openai", - Retryable: true, - StatusCode: 429, - }) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - requireNoSnapshotRetryEvent(t, snapshot) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeDoesNotReplayRetryAfterTerminalStatus(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusCompleted} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newBufferedSubscribeTestServer(t, db, chatID) - - server.publishRetry(chatID, newTestRetryPayload()) - server.publishStatus(chatID, database.ChatStatusCompleted, uuid.NullUUID{}) - - snapshot, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - requireNoSnapshotRetryEvent(t, snapshot) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribePrefersStructuredErrorPayloadViaPubsub(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - classified := chaterror.ClassifiedError{ - Message: "OpenAI is rate limiting requests.", - Kind: codersdk.ChatErrorKindRateLimit, - Provider: "openai", - Retryable: true, - StatusCode: 429, - } - server.publishError(chatID, classified) - - event := requireStreamErrorEvent(t, events) - require.Equal(t, chaterror.TerminalErrorPayload(classified), event.Error) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func TestSubscribeFallsBackToLegacyErrorStringViaPubsub(t *testing.T) { - t.Parallel() - - ctx, cancelCtx := context.WithCancel(context.Background()) - defer cancelCtx() - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - server := newSubscribeTestServer(t, db) - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{ - Error: "legacy error only", - }) - - event := requireStreamErrorEvent(t, events) - require.Equal(t, &codersdk.ChatError{Message: "legacy error only"}, event.Error) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -func newTestRetryPayload() *codersdk.ChatStreamRetry { - payload := chaterror.StreamRetryPayload(1, 1500*time.Millisecond, chaterror.ClassifiedError{ - Message: "OpenAI is rate limiting requests.", - Kind: codersdk.ChatErrorKindRateLimit, - Provider: "openai", - Retryable: true, - StatusCode: 429, - }) - if payload == nil { - panic("expected retry payload") - } - payload.RetryingAt = time.Unix(1_700_000_000, 0).UTC() - return payload -} - -func TestSubscribeAuthorizedFallsBackToStaleRowWhenRefreshFails(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - server := newSubscribeTestServer(t, db) - - chatID := uuid.New() - staleChat := database.Chat{ID: chatID, Status: database.ChatStatusPending} - - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffer = []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - ChatID: chatID, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("thinking"), - }, - }, - }} - state.mu.Unlock() - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(database.Chat{}, xerrors.New("refresh failed")), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - - initialSnapshot, events, cancel, ok := server.SubscribeAuthorized(ctx, staleChat, nil, 0) - require.True(t, ok) - defer cancel() - - require.Len(t, initialSnapshot, 2) - require.Equal(t, codersdk.ChatStreamEventTypeStatus, initialSnapshot[0].Type) - require.NotNil(t, initialSnapshot[0].Status) - require.Equal(t, codersdk.ChatStatusPending, initialSnapshot[0].Status.Status) - require.Equal(t, codersdk.ChatStreamEventTypeMessagePart, initialSnapshot[1].Type) - require.NotNil(t, initialSnapshot[1].MessagePart) - require.Equal(t, "thinking", initialSnapshot[1].MessagePart.Part.Text) - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - func TestSubscribeRejectsUnauthorizedCallerBeforeSharedFetches(t *testing.T) { t.Parallel() @@ -2740,9 +2136,6 @@ func TestSubscribeRejectsUnauthorizedCallerBeforeSharedFetches(t *testing.T) { require.Nil(t, snapshot) require.Nil(t, events) require.Nil(t, cancel) - - _, exists := server.chatStreams.Load(chatID) - require.False(t, exists) } func TestSubscribeSurfacesTransientLookupFailureAsInitialError(t *testing.T) { @@ -2767,220 +2160,22 @@ func TestSubscribeSurfacesTransientLookupFailureAsInitialError(t *testing.T) { _, open := <-events require.False(t, open) - - _, exists := server.chatStreams.Load(chatID) - require.False(t, exists) } func newSubscribeTestServer(t *testing.T, db database.Store) *Server { t.Helper() + poller := newStreamSyncPoller(context.Background(), db, nil, slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})) + t.Cleanup(poller.Close) return &Server{ - db: db, - logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - pubsub: dbpubsub.NewInMemory(), + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + pubsub: dbpubsub.NewInMemory(), + clock: quartz.NewReal(), + streamSyncPoller: poller, } } -func newBufferedSubscribeTestServer(t *testing.T, db database.Store, chatID uuid.UUID) *Server { - t.Helper() - - server := newSubscribeTestServer(t, db) - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffering = true - state.mu.Unlock() - return server -} - -func requireStreamMessageEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { - t.Helper() - - select { - case event, ok := <-events: - require.True(t, ok, "chat stream closed before delivering an event") - require.Equal(t, codersdk.ChatStreamEventTypeMessage, event.Type) - require.NotNil(t, event.Message) - return event - case <-time.After(time.Second): - t.Fatal("timed out waiting for chat stream message event") - return codersdk.ChatStreamEvent{} - } -} - -func requireStreamRetryEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { - t.Helper() - - select { - case event, ok := <-events: - require.True(t, ok, "chat stream closed before delivering an event") - require.Equal(t, codersdk.ChatStreamEventTypeRetry, event.Type) - require.NotNil(t, event.Retry) - return event - case <-time.After(time.Second): - t.Fatal("timed out waiting for chat stream retry event") - return codersdk.ChatStreamEvent{} - } -} - -func requireSnapshotRetryEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { - t.Helper() - - var retryEvents []codersdk.ChatStreamEvent - for _, event := range snapshot { - if event.Type == codersdk.ChatStreamEventTypeRetry { - retryEvents = append(retryEvents, event) - } - } - - require.Len(t, retryEvents, 1, "expected exactly one retry event in snapshot") - require.NotNil(t, retryEvents[0].Retry) - return retryEvents[0] -} - -func requireNoSnapshotRetryEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) { - t.Helper() - - for _, event := range snapshot { - require.NotEqual(t, codersdk.ChatStreamEventTypeRetry, event.Type, - "unexpected retry event in snapshot: %+v", event) - } -} - -func requireSnapshotMessagePartEvent(t *testing.T, snapshot []codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { - t.Helper() - - for _, event := range snapshot { - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - require.NotNil(t, event.MessagePart) - return event - } - } - - t.Fatal("expected message_part event in snapshot") - return codersdk.ChatStreamEvent{} -} - -func requireStreamErrorEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent) codersdk.ChatStreamEvent { - t.Helper() - - select { - case event, ok := <-events: - require.True(t, ok, "chat stream closed before delivering an event") - require.Equal(t, codersdk.ChatStreamEventTypeError, event.Type) - require.NotNil(t, event.Error) - return event - case <-time.After(time.Second): - t.Fatal("timed out waiting for chat stream error event") - return codersdk.ChatStreamEvent{} - } -} - -func requireNoStreamEvent(t *testing.T, events <-chan codersdk.ChatStreamEvent, wait time.Duration) { - t.Helper() - - select { - case event, ok := <-events: - if !ok { - t.Fatal("chat stream closed unexpectedly") - } - t.Fatalf("unexpected chat stream event: %+v", event) - case <-time.After(wait): - } -} - -// TestPublishToStream_DropWarnRateLimiting walks through a -// realistic lifecycle: buffer fills up, subscriber channel fills -// up, counters get reset between steps. It verifies that WARN -// logs are rate-limited to at most once per streamDropWarnInterval -// and that counter resets re-enable an immediate WARN. -func TestPublishToStream_DropWarnRateLimiting(t *testing.T) { - t.Parallel() - - sink := testutil.NewFakeSink(t) - mClock := quartz.NewMock(t) - - server := &Server{ - logger: sink.Logger(), - clock: mClock, - } - - chatID := uuid.New() - subCh := make(chan codersdk.ChatStreamEvent, 1) - subCh <- codersdk.ChatStreamEvent{} // pre-fill so sends always drop - - // Set up state that mirrors a running chat: buffer at capacity, - // buffering enabled, one saturated subscriber. - state := &chatStreamState{ - buffering: true, - buffer: make([]bufferedStreamPart, maxStreamBufferSize), - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{ - uuid.New(): subCh, - }, - } - server.chatStreams.Store(chatID, state) - - bufferMsg := "chat stream buffer full, dropping oldest event" - subMsg := "dropping chat stream event" - - filter := func(level slog.Level, msg string) func(slog.SinkEntry) bool { - return func(e slog.SinkEntry) bool { - return e.Level == level && e.Message == msg - } - } - - // --- Phase 1: buffer-full rate limiting --- - // message_part events hit both the buffer-full and subscriber-full - // paths. The first publish triggers a WARN for each; the rest - // within the window are DEBUG. - partEvent := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - } - for i := 0; i < 50; i++ { - server.publishToStream(chatID, partEvent) - } - - require.Len(t, sink.Entries(filter(slog.LevelWarn, bufferMsg)), 1) - require.Empty(t, sink.Entries(filter(slog.LevelDebug, bufferMsg))) - requireFieldValue(t, sink.Entries(filter(slog.LevelWarn, bufferMsg))[0], "dropped_count", int64(1)) - - // Subscriber also saw 50 drops (one per publish). - require.Len(t, sink.Entries(filter(slog.LevelWarn, subMsg)), 1) - require.Empty(t, sink.Entries(filter(slog.LevelDebug, subMsg))) - requireFieldValue(t, sink.Entries(filter(slog.LevelWarn, subMsg))[0], "dropped_count", int64(1)) - - // --- Phase 2: clock advance triggers second WARN with count --- - mClock.Advance(streamDropWarnInterval + time.Second) - server.publishToStream(chatID, partEvent) - - bufWarn := sink.Entries(filter(slog.LevelWarn, bufferMsg)) - require.Len(t, bufWarn, 2) - requireFieldValue(t, bufWarn[1], "dropped_count", int64(50)) - - subWarn := sink.Entries(filter(slog.LevelWarn, subMsg)) - require.Len(t, subWarn, 2) - requireFieldValue(t, subWarn[1], "dropped_count", int64(50)) - - // --- Phase 3: counter reset (simulates step persist) --- - state.mu.Lock() - state.buffer = make([]bufferedStreamPart, maxStreamBufferSize) - state.resetDropCounters() - state.mu.Unlock() - - // The very next drop should WARN immediately — the reset zeroed - // lastWarnAt so the interval check passes. - server.publishToStream(chatID, partEvent) - - bufWarn = sink.Entries(filter(slog.LevelWarn, bufferMsg)) - require.Len(t, bufWarn, 3, "expected WARN immediately after counter reset") - requireFieldValue(t, bufWarn[2], "dropped_count", int64(1)) - - subWarn = sink.Entries(filter(slog.LevelWarn, subMsg)) - require.Len(t, subWarn, 3, "expected subscriber WARN immediately after counter reset") - requireFieldValue(t, subWarn[2], "dropped_count", int64(1)) -} - func TestResolveUserCompactionThreshold(t *testing.T) { t.Parallel() @@ -3194,8 +2389,7 @@ func TestSkillsFromParts(t *testing.T) { t.Run("RoundTrip", func(t *testing.T) { // Simulate persist -> reconstruct cycle: marshal skill - // parts the same way persistInstructionFiles does, then - // verify skillsFromParts recovers the metadata. + // parts, then verify skillsFromParts recovers the metadata. t.Parallel() want := []chattool.SkillMeta{ {Name: "deep-review", Description: "Multi-reviewer review", Dir: "/skills/deep-review"}, @@ -3657,40 +2851,6 @@ func TestContextFileAgentID(t *testing.T) { }) } -func TestHasPersistedInstructionFiles(t *testing.T) { - t.Parallel() - - t.Run("IgnoresAgentChatContextSentinel", func(t *testing.T) { - t.Parallel() - agentID := uuid.New() - msgs := []database.ChatMessage{ - chattest.ChatMessageWithParts([]codersdk.ChatMessagePart{{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: AgentChatContextSentinelPath, - ContextFileAgentID: uuid.NullUUID{ - UUID: agentID, - Valid: true, - }, - }}), - } - require.False(t, hasPersistedInstructionFiles(msgs)) - }) - - t.Run("AcceptsPersistedInstructionFile", func(t *testing.T) { - t.Parallel() - agentID := uuid.New() - msgs := []database.ChatMessage{ - chattest.ChatMessageWithParts([]codersdk.ChatMessagePart{{ - Type: codersdk.ChatMessagePartTypeContextFile, - ContextFilePath: "/workspace/AGENTS.md", - ContextFileContent: "repo instructions", - ContextFileAgentID: uuid.NullUUID{UUID: agentID, Valid: true}, - }}), - } - require.True(t, hasPersistedInstructionFiles(msgs)) - }) -} - func TestInstructionFromContextFilesUsesLatestContextAgent(t *testing.T) { t.Parallel() @@ -3854,709 +3014,6 @@ func TestSkillsFromPartsUsesLatestContextAgent(t *testing.T) { }}, got) } -func TestMergeSkillMetas(t *testing.T) { - t.Parallel() - - persisted := []chattool.SkillMeta{{ - Name: "repo-helper", - Description: "Persisted skill", - Dir: "/skills/repo-helper-old", - }} - discovered := []chattool.SkillMeta{ - { - Name: "repo-helper", - Description: "Discovered replacement", - Dir: "/skills/repo-helper-new", - MetaFile: "SKILL.md", - }, - { - Name: "deep-review", - Description: "Discovered skill", - Dir: "/skills/deep-review", - }, - } - - got := mergeSkillMetas(persisted, discovered) - require.Equal(t, []chattool.SkillMeta{ - discovered[0], - discovered[1], - }, got) -} - -func TestSelectSkillMetasForInstructionRefresh(t *testing.T) { - t.Parallel() - - persisted := []chattool.SkillMeta{{Name: "persisted", Dir: "/skills/persisted"}} - discovered := []chattool.SkillMeta{{Name: "discovered", Dir: "/skills/discovered"}} - currentAgentID := uuid.New() - otherAgentID := uuid.New() - - t.Run("MergesCurrentAgentSkills", func(t *testing.T) { - t.Parallel() - got := selectSkillMetasForInstructionRefresh( - persisted, - discovered, - uuid.NullUUID{UUID: currentAgentID, Valid: true}, - uuid.NullUUID{UUID: currentAgentID, Valid: true}, - ) - require.Equal(t, []chattool.SkillMeta{discovered[0], persisted[0]}, got) - }) - - t.Run("DropsStalePersistedSkillsWhenAgentChanged", func(t *testing.T) { - t.Parallel() - got := selectSkillMetasForInstructionRefresh( - persisted, - discovered, - uuid.NullUUID{UUID: currentAgentID, Valid: true}, - uuid.NullUUID{UUID: otherAgentID, Valid: true}, - ) - require.Equal(t, discovered, got) - }) - - t.Run("PreservesPersistedSkillsWhenAgentLookupFails", func(t *testing.T) { - t.Parallel() - got := selectSkillMetasForInstructionRefresh( - persisted, - nil, - uuid.NullUUID{}, - uuid.NullUUID{UUID: otherAgentID, Valid: true}, - ) - require.Equal(t, persisted, got) - }) -} - -// TestProcessChat_IgnoresStaleControlNotification verifies that -// processChat is not interrupted by a "pending" notification -// published before processing begins. This is the race that caused -// TestOpenAIReasoningWithWebSearchRoundTripStoreFalse to flake: -// SendMessage publishes "pending" via PostgreSQL NOTIFY, and due -// to async delivery the notification can arrive at the control -// subscriber after it registers but before the processor publishes -// "running". -func TestProcessChat_IgnoresStaleControlNotification(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - ps := dbpubsub.NewInMemory() - clock := quartz.NewMock(t) - - chatID := uuid.New() - workerID := uuid.New() - - server := &Server{ - db: db, - logger: logger, - pubsub: ps, - clock: clock, - workerID: workerID, - chatHeartbeatInterval: time.Minute, - metrics: chatloop.NopMetrics(), - configCache: newChatConfigCache(ctx, db, clock), - heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), - } - - // Publish a stale "pending" notification on the control channel - // BEFORE processChat subscribes. In production this is the - // notification from SendMessage that triggered the processing. - staleNotify, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusPending), - }) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), staleNotify) - require.NoError(t, err) - - // Track which status processChat writes during cleanup. - var finalStatus database.ChatStatus - - // The deferred cleanup in processChat runs a transaction. - db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn( - func(fn func(database.Store) error, _ *database.TxOptions) error { - return fn(db) - }, - ) - db.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return( - database.Chat{ID: chatID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}}, nil, - ) - db.EXPECT().UpdateChatStatus(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, params database.UpdateChatStatusParams) (database.Chat, error) { - finalStatus = params.Status - return database.Chat{ - ID: chatID, - Status: params.Status, - LastTurnSummary: sql.NullString{String: "previous summary", Valid: true}, - }, nil - }, - ) - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return( - database.Chat{ID: chatID, Status: database.ChatStatusError}, - nil, - ) - - db.EXPECT().UpdateChatLastTurnSummary(gomock.Any(), gomock.Any()).Return(int64(1), nil) - - // resolveChatModel fails immediately — that's fine, we only - // need processChat to get past initialization without being - // interrupted by the stale notification. - db.EXPECT().GetChatModelConfigByID(gomock.Any(), gomock.Any()).Return( - database.ChatModelConfig{}, xerrors.New("no model configured"), - ).AnyTimes() - db.EXPECT().GetAIProviders(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil).AnyTimes() - db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return( - database.ChatUsageLimitConfig{}, sql.ErrNoRows, - ).AnyTimes() - db.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chatID).Return(nil, nil).AnyTimes() - - chat := database.Chat{ID: chatID, LastModelConfigID: uuid.New()} - done := make(chan struct{}) - go func() { - defer close(done) - server.processChat(ctx, chat) - }() - - // Wait for processChat to finish entirely. It re-reads chat state and - // runs more cleanup after UpdateChatStatus, so signaling completion from - // the status update itself races test teardown. - testutil.TryReceive(ctx, t, done) - - WaitUntilIdleForTest(server) - - // If the stale notification interrupted us, status would be - // "waiting" (the ErrInterrupted path). Since the gate blocked - // it, processChat reached runChat, which failed on model - // resolution → status is "error". - require.Equal(t, database.ChatStatusError, finalStatus, - "processChat should have reached runChat (error), not been interrupted (waiting)") -} - -func TestShouldPublishFinishedChatState(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - chatID := uuid.New() - workerID := uuid.New() - - server := &Server{db: db} - updatedChat := database.Chat{ - ID: chatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - } - - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(database.Chat{ - ID: chatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - }, nil) - - require.True(t, server.shouldPublishFinishedChatState(ctx, logger, updatedChat)) - - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(database.Chat{ - ID: chatID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, - }, nil) - - require.False(t, server.shouldPublishFinishedChatState(ctx, logger, updatedChat)) -} - -// TestShouldPublishFinishedChatState_DBErrorPublishes pins the -// deliberate fail-open behavior when the re-read query errors: we -// surface the finished state anyway so watchers don't get stuck -// waiting for a status update that never arrives. The error path is -// easy to regress into a fail-closed default otherwise. -func TestShouldPublishFinishedChatState_DBErrorPublishes(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - chatID := uuid.New() - - server := &Server{db: db} - updatedChat := database.Chat{ - ID: chatID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - } - - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return( - database.Chat{}, xerrors.New("boom"), - ) - - require.True(t, server.shouldPublishFinishedChatState(ctx, logger, updatedChat), - "fail-open: a re-read error must not swallow the status change") -} - -// TestHeartbeatTick_StolenChatIsInterrupted verifies that when the -// batch heartbeat UPDATE does not return a registered chat's ID -// (because another replica stole it or it was completed), the -// heartbeat tick cancels that chat's context with ErrInterrupted -// while leaving surviving chats untouched. -func TestHeartbeatTick_StolenChatIsInterrupted(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - clock := quartz.NewMock(t) - - workerID := uuid.New() - - server := &Server{ - db: db, - logger: logger, - clock: clock, - workerID: workerID, - chatHeartbeatInterval: time.Minute, - metrics: chatloop.NopMetrics(), - heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), - } - - // Create three chats with independent cancel functions. - chat1 := uuid.New() - chat2 := uuid.New() - chat3 := uuid.New() - - _, cancel1 := context.WithCancelCause(ctx) - _, cancel2 := context.WithCancelCause(ctx) - ctx3, cancel3 := context.WithCancelCause(ctx) - - server.registerHeartbeat(&heartbeatEntry{ - cancelWithCause: cancel1, - chatID: chat1, - logger: logger, - }) - server.registerHeartbeat(&heartbeatEntry{ - cancelWithCause: cancel2, - chatID: chat2, - logger: logger, - }) - server.registerHeartbeat(&heartbeatEntry{ - cancelWithCause: cancel3, - chatID: chat3, - logger: logger, - }) - - // The batch UPDATE returns only chat1 and chat2 — - // chat3 was "stolen" by another replica. - db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, params database.UpdateChatHeartbeatsParams) ([]uuid.UUID, error) { - require.Equal(t, workerID, params.WorkerID) - require.Len(t, params.IDs, 3) - // Return only chat1 and chat2 as surviving. - return []uuid.UUID{chat1, chat2}, nil - }, - ) - - server.heartbeatTick(ctx) - - // chat3's context should be canceled with ErrInterrupted. - require.ErrorIs(t, context.Cause(ctx3), chatloop.ErrInterrupted, - "stolen chat should be interrupted") - - // chat3 should have been removed from the registry by - // unregister (in production this happens via defer in - // processChat). The heartbeat tick itself does not - // unregister — it only cancels. Verify the entry is - // still present (processChat's defer would clean it up). - server.heartbeatMu.Lock() - _, chat1Exists := server.heartbeatRegistry[chat1] - _, chat2Exists := server.heartbeatRegistry[chat2] - _, chat3Exists := server.heartbeatRegistry[chat3] - server.heartbeatMu.Unlock() - - require.True(t, chat1Exists, "surviving chat1 should remain registered") - require.True(t, chat2Exists, "surviving chat2 should remain registered") - require.True(t, chat3Exists, - "stolen chat3 should still be in registry (processChat defer removes it)") -} - -// TestHeartbeatTick_DBErrorDoesNotInterruptChats verifies that a -// transient database failure causes the tick to log and return -// without canceling any registered chats. -func TestHeartbeatTick_DBErrorDoesNotInterruptChats(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - clock := quartz.NewMock(t) - - server := &Server{ - db: db, - logger: logger, - clock: clock, - workerID: uuid.New(), - chatHeartbeatInterval: time.Minute, - metrics: chatloop.NopMetrics(), - heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), - } - - chatID := uuid.New() - chatCtx, cancel := context.WithCancelCause(ctx) - - server.registerHeartbeat(&heartbeatEntry{ - cancelWithCause: cancel, - chatID: chatID, - logger: logger, - }) - - // Simulate a transient DB error. - db.EXPECT().UpdateChatHeartbeats(gomock.Any(), gomock.Any()).Return( - nil, xerrors.New("connection reset"), - ) - - server.heartbeatTick(ctx) - - // Chat should NOT be interrupted — the tick logged and - // returned early. - require.NoError(t, chatCtx.Err(), - "chat context should not be canceled on transient DB error") -} - -// TestSubscribeCancelDuringGrace_ReapedBySweep verifies that a -// subscriber detach inside bufferRetainGracePeriod (the OSS trigger -// for the retained-buffer leak) leaves the state mapped, and the -// next sweep past the grace window reaps it. -func TestSubscribeCancelDuringGrace_ReapedBySweep(t *testing.T) { - t.Parallel() - - logger := slogtest.Make(t, nil) - mClock := quartz.NewMock(t) - - server := &Server{ - logger: logger, - clock: mClock, - } - - chatID := uuid.New() - start := mClock.Now() - - // Just-finished chat: processing done, buffer retained for - // late-connecting relay subscribers. - state := &chatStreamState{ - buffering: false, - bufferRetainedAt: start, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - buffer: []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - }, - }, - }}, - } - server.chatStreams.Store(chatID, state) - - // Real subscribeToStream cancel path: the WS subscriber detach - // that leaks in prod. - snapshot, currentRetry, events, cancelSub := server.subscribeToStream(chatID) - require.Len(t, snapshot, 1) - require.Nil(t, currentRetry) - require.NotNil(t, events) - - mClock.Advance(bufferRetainGracePeriod / 2) - cancelSub() - - _, ok := server.chatStreams.Load(chatID) - require.True(t, ok, - "entry should remain during grace window after subscriber detach") - - mClock.Advance(bufferRetainGracePeriod) - server.sweepIdleStreams() - - _, ok = server.chatStreams.Load(chatID) - require.False(t, ok, - "entry should be reaped after grace period expires and sweep runs") -} - -// TestSweepIdleStreams_ReapsStaleRetainedBuffer: grace expired, no -// subscribers, not buffering -> reaped. -func TestSweepIdleStreams_ReapsStaleRetainedBuffer(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - state := &chatStreamState{ - buffering: false, - bufferRetainedAt: mClock.Now(), - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - buffer: []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - }, - }}, - } - server.chatStreams.Store(chatID, state) - - mClock.Advance(bufferRetainGracePeriod + time.Second) - server.sweepIdleStreams() - - _, ok := server.chatStreams.Load(chatID) - require.False(t, ok, "stale retained state should be reaped") -} - -// TestSweepIdleStreams_DoesNotReapActiveBuffering: buffering=true -// blocks reap even long after any grace would have expired. -func TestSweepIdleStreams_DoesNotReapActiveBuffering(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - state := &chatStreamState{ - buffering: true, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - buffer: []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - }, - }}, - } - server.chatStreams.Store(chatID, state) - - mClock.Advance(time.Hour) - server.sweepIdleStreams() - - _, ok := server.chatStreams.Load(chatID) - require.True(t, ok, "actively-buffering state must not be reaped") -} - -// TestSweepIdleStreams_DoesNotReapWithSubscribers: attached -// subscribers block reap even when grace has expired. -func TestSweepIdleStreams_DoesNotReapWithSubscribers(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - state := &chatStreamState{ - buffering: false, - bufferRetainedAt: mClock.Now(), - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{ - uuid.New(): make(chan codersdk.ChatStreamEvent, 1), - }, - buffer: []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - }, - }}, - } - server.chatStreams.Store(chatID, state) - - mClock.Advance(bufferRetainGracePeriod + time.Second) - server.sweepIdleStreams() - - _, ok := server.chatStreams.Load(chatID) - require.True(t, ok, "state with subscribers must not be reaped") -} - -// TestSweepIdleStreams_DefersDuringGracePeriod: sweep inside grace -// is a no-op; the next sweep past grace reaps. -func TestSweepIdleStreams_DefersDuringGracePeriod(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - start := mClock.Now() - state := &chatStreamState{ - buffering: false, - bufferRetainedAt: start, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - buffer: []bufferedStreamPart{{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - }, - }}, - } - server.chatStreams.Store(chatID, state) - - mClock.Advance(bufferRetainGracePeriod / 2) - server.sweepIdleStreams() - - _, ok := server.chatStreams.Load(chatID) - require.True(t, ok, "sweep inside grace window must not reap") - - mClock.Advance(bufferRetainGracePeriod) - server.sweepIdleStreams() - - _, ok = server.chatStreams.Load(chatID) - require.False(t, ok, "sweep after grace window must reap") -} - -// TestPublishToStream_DropZeroesBackingSlot verifies that evicting -// the oldest buffered event at capacity zeroes the dropped slot so -// its *ChatStreamMessagePart becomes GC-eligible immediately. -func TestPublishToStream_DropZeroesBackingSlot(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - - // Over-allocate by one so the post-drop append fits in place and - // exercises the backing-array reuse this test is checking. - buf := make([]bufferedStreamPart, maxStreamBufferSize, maxStreamBufferSize+1) - for i := range buf { - buf[i] = bufferedStreamPart{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - }, - } - } - // Sentinel in slot 0 distinguishes "slot was zeroed" from "slot - // was overwritten by a later append". - sentinel := &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - } - buf[0] = bufferedStreamPart{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: sentinel, - }, - } - // Alias over the full backing array so we can still observe slot - // 0 after publishToStream reslices state.buffer forward. - origBacking := buf[:cap(buf)] - - state := &chatStreamState{ - buffering: true, - buffer: buf, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - } - server.chatStreams.Store(chatID, state) - - newPart := &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - } - server.publishToStream(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: newPart, - }) - - require.Equal(t, bufferedStreamPart{}, origBacking[0], - "dropped slot must be zero-valued so its *ChatStreamMessagePart "+ - "is eligible for GC; got %+v", origBacking[0]) - - // Sanity-check the in-place append path the fix targets: if Go's - // growth policy ever makes this append reallocate, this fails - // loudly so the test author revisits the setup. - require.Same(t, newPart, origBacking[len(origBacking)-1].event.MessagePart, - "append must have landed in the original backing array; the "+ - "zero-out invariant only matters when cap > len") -} - -// TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry covers -// the race where a caller holds a pointer to a no-longer-mapped -// state (e.g. a janitor Range callback racing a fresh -// getOrCreateStreamState) and would otherwise evict the fresh entry. -// With CompareAndDelete in cleanupStreamIfIdle the stale delete is -// a no-op. -func TestCleanupStreamIfIdle_StalePointerDoesNotDeleteFreshEntry(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - - // Stale pointer: reapable (not buffering, no subscribers, grace - // expired) but no longer the map's live entry. - stale := &chatStreamState{ - buffering: false, - bufferRetainedAt: mClock.Now(), - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - } - - // Fresh entry: the state getOrCreateStreamState would install - // after a racing processChat run. Actively buffering, so not - // reapable. Only this state is in the map. - fresh := &chatStreamState{ - buffering: true, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - } - server.chatStreams.Store(chatID, fresh) - - mClock.Advance(bufferRetainGracePeriod + time.Second) - - // Stale caller mirrors the janitor Range callback after the map - // entry has already been replaced. - stale.mu.Lock() - server.cleanupStreamIfIdle(chatID, stale) - stale.mu.Unlock() - - got, ok := server.chatStreams.Load(chatID) - require.True(t, ok, - "fresh entry must remain mapped when cleanup is called with a stale pointer") - require.Same(t, fresh, got, - "cleanup must not replace the fresh entry with the stale one") -} - -// TestSafeSweepIdleStreams_RecoversFromPanic verifies that an -// unexpected panic inside sweepIdleStreams is recovered rather than -// killing the janitor goroutine. Without this guard, a panic would -// silently reintroduce the very leak the janitor exists to prevent. -func TestSafeSweepIdleStreams_RecoversFromPanic(t *testing.T) { - t.Parallel() - - server := &Server{ - logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - clock: quartz.NewMock(t), - } - - chatID := uuid.New() - // A nil *chatStreamState passes the type assertion in sweepIdleStreams - // but panics on state.mu.Lock with a nil-pointer deref. Any future - // panic source in the sweep would trigger the same recovery path. - var nilState *chatStreamState - server.chatStreams.Store(chatID, nilState) - - require.NotPanics(t, func() { - server.safeSweepIdleStreams(context.Background()) - }, "safeSweepIdleStreams must recover panics so the janitor loop keeps running") -} - func TestGetWorkspaceConn_StaleAgentRecovery(t *testing.T) { // Regression test: when a workspace is rebuilt, the chat's stored // agent ID points to a disconnected agent from the old build. The @@ -5671,571 +4128,6 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { require.ErrorContains(t, err, "authentication failed") } -// TestAutoPromote_InsertFailureRollsBackTransaction verifies that when -// tryAutoPromoteQueuedMessage pops a queued message but the subsequent -// insert fails, the error propagates to the InTx callback, causing the -// transaction to roll back and preserving the queued message. -func TestAutoPromote_InsertFailureRollsBackTransaction(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitShort) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - tx := dbmock.NewMockStore(ctrl) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ps := dbpubsub.NewInMemory() - clock := quartz.NewReal() - - chatID := uuid.New() - workerID := uuid.New() - ownerID := uuid.New() - modelConfigID := uuid.New() - - waitingChat := database.Chat{ - ID: chatID, - OwnerID: ownerID, - LastModelConfigID: modelConfigID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, - } - queuedMsg := database.ChatQueuedMessage{ - ID: 1, - ChatID: chatID, - Content: []byte(`[{"type":"text","text":"queued"}]`), - } - insertErr := xerrors.New("insert failed") - - server := &Server{ - db: db, - logger: logger, - pubsub: ps, - configCache: newChatConfigCache(ctx, db, clock), - } - - // The caller runs tryAutoPromoteQueuedMessage inside InTx. - // Wire the mock to execute the callback against the TX mock. - var txErr error - db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn( - func(fn func(database.Store) error, _ *database.TxOptions) error { - txErr = fn(tx) - return txErr - }, - ) - - // Inside the TX: lock chat, get queued messages, resolve model - // config, pop queued message, insert fails. - tx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(waitingChat, nil) - tx.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return([]database.ChatQueuedMessage{queuedMsg}, nil) - tx.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(database.ChatModelConfig{ID: modelConfigID}, nil) - tx.EXPECT().PopNextQueuedMessage(gomock.Any(), chatID).Return(queuedMsg, nil) - tx.EXPECT().InsertChatMessages(gomock.Any(), gomock.Any()).Return(nil, insertErr) - - // Invoke tryAutoPromoteQueuedMessage through the same InTx - // pattern the processChat defer uses. The test directly calls - // the production path to verify error propagation. - _ = db.InTx(func(txStore database.Store) error { - latestChat, err := txStore.GetChatByIDForUpdate(ctx, chatID) - if err != nil { - return err - } - - _, _, _, promoteErr := server.tryAutoPromoteQueuedMessage(ctx, txStore, latestChat) - if promoteErr != nil { - return promoteErr - } - - // This code path should not be reached when the insert - // fails, because promoteErr should be non-nil. - return nil - }, nil) - - // The InTx callback must return a non-nil error so the - // transaction rolls back, preserving the queued message. - require.Error(t, txErr, "InTx callback should return error when insert fails") -} - -// TestAutoPromote_WakesRunLoopAfterPromotion verifies that after the -func TestAutoPromote_InsertFailureSkipsStatusUpdate(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitLong) - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - tx := dbmock.NewMockStore(ctrl) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - ps := dbpubsub.NewInMemory() - clock := quartz.NewReal() - - chatID := uuid.New() - workerID := uuid.New() - ownerID := uuid.New() - modelConfigID := uuid.New() - - waitingChat := database.Chat{ - ID: chatID, - OwnerID: ownerID, - LastModelConfigID: modelConfigID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, - } - queuedMsg := database.ChatQueuedMessage{ - ID: 1, - ChatID: chatID, - Content: []byte(`[{"type":"text","text":"queued"}]`), - } - - wakeCh := make(chan struct{}, 1) - server := &Server{ - db: db, - logger: logger, - pubsub: ps, - clock: clock, - workerID: workerID, - wakeCh: wakeCh, - chatHeartbeatInterval: time.Minute, - metrics: chatloop.NopMetrics(), - configCache: newChatConfigCache(ctx, db, clock), - heartbeatRegistry: make(map[uuid.UUID]*heartbeatEntry), - } - - // Hold model resolution until the interrupt has canceled the chat - // context. Returning ErrInterrupted keeps processChat on the - // interrupted path regardless of whether the cache singleflight sees - // the caller cancellation or the DB fetch result first. - modelBlocked := make(chan struct{}) - modelRelease := make(chan struct{}) - var modelBlockedOnce sync.Once - db.EXPECT().GetChatModelConfigByID(gomock.Any(), gomock.Any()).DoAndReturn( - func(_ context.Context, _ uuid.UUID) (database.ChatModelConfig, error) { - modelBlockedOnce.Do(func() { close(modelBlocked) }) - <-modelRelease - return database.ChatModelConfig{}, chatloop.ErrInterrupted - }, - ).AnyTimes() - db.EXPECT().GetAIProviders(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil).AnyTimes() - db.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return( - database.ChatUsageLimitConfig{}, sql.ErrNoRows, - ).AnyTimes() - db.EXPECT().GetChatMessagesForPromptByChatID(gomock.Any(), chatID).Return(nil, nil).AnyTimes() - - // The deferred cleanup transaction: InsertChatMessages fails, - // so UpdateChatStatus must NOT be called. - db.EXPECT().InTx(gomock.Any(), gomock.Any()).DoAndReturn( - func(fn func(database.Store) error, _ *database.TxOptions) error { - return fn(tx) - }, - ) - tx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(waitingChat, nil) - tx.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return([]database.ChatQueuedMessage{queuedMsg}, nil) - tx.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(database.ChatModelConfig{ID: modelConfigID}, nil) - tx.EXPECT().PopNextQueuedMessage(gomock.Any(), chatID).Return(queuedMsg, nil) - tx.EXPECT().InsertChatMessages(gomock.Any(), gomock.Any()).Return( - nil, xerrors.New("insert failed"), - ) - tx.EXPECT().UpdateChatStatus(gomock.Any(), gomock.Any()).Times(0) - - // Subscribe BEFORE launching the goroutine. - runningCh := make(chan struct{}, 1) - unsubRunning, err := ps.SubscribeWithErr( - coderdpubsub.ChatStreamNotifyChannel(chatID), - func(_ context.Context, msg []byte, err error) { - if err != nil { - return - } - var notify coderdpubsub.ChatStreamNotifyMessage - if json.Unmarshal(msg, ¬ify) != nil { - return - } - if notify.Status == string(database.ChatStatusRunning) { - select { - case runningCh <- struct{}{}: - default: - } - } - }, - ) - require.NoError(t, err) - defer unsubRunning() - - chat := database.Chat{ID: chatID, OwnerID: ownerID, LastModelConfigID: modelConfigID} - processDone := make(chan struct{}) - go func() { - defer close(processDone) - server.processChat(ctx, chat) - }() - - select { - case <-runningCh: - case <-ctx.Done(): - t.Fatal("timed out waiting for running status") - } - - select { - case <-modelBlocked: - case <-ctx.Done(): - t.Fatal("timed out waiting for model resolution") - } - - // Publish an interrupt so processChat exits runChat. - interruptMsg, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusWaiting), - }) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), interruptMsg) - require.NoError(t, err) - close(modelRelease) - - select { - case <-processDone: - case <-ctx.Done(): - t.Fatal("processChat did not complete") - } - - // The wake channel should NOT have a signal because the - // transaction failed before reaching UpdateChatStatus. - select { - case <-wakeCh: - t.Fatal("wake channel should not have a signal after insert failure") - default: - // No signal, as expected. - } -} - -// makeInProgressPart is a small constructor for buffered message_part -// fixtures used by snapshotBufferLocked / subscribeToStream tests. It -// builds an in-progress part (committedMessageID == 0) with a -// recognizable text body so failing assertions can identify which -// part survived the filter. -func makeInProgressPart(text string) bufferedStreamPart { - return bufferedStreamPart{ - event: codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - Part: codersdk.ChatMessageText(text), - }, - }, - } -} - -// makeCommittedPart builds a part already claimed by the given -// durable assistant message ID. -func makeCommittedPart(committedID int64, text string) bufferedStreamPart { - p := makeInProgressPart(text) - p.committedMessageID = committedID - return p -} - -func partText(event codersdk.ChatStreamEvent) string { - if event.MessagePart == nil { - return "" - } - return event.MessagePart.Part.Text -} - -// TestSnapshotBufferLocked_DropsCommittedParts asserts the core -// dedup contract: parts that were claimed by a durable assistant -// message (committedMessageID != 0) are dropped from the snapshot -// because the subscriber will receive that durable message through -// the REST snapshot, the initial DB query, or pubsub. -func TestSnapshotBufferLocked_DropsCommittedParts(t *testing.T) { - t.Parallel() - - buffer := []bufferedStreamPart{ - makeCommittedPart(100, "turnA-1"), - makeCommittedPart(100, "turnA-2"), - makeCommittedPart(200, "turnB-1"), - makeInProgressPart("in-progress-1"), - makeInProgressPart("in-progress-2"), - } - - snapshot := snapshotBufferLocked(buffer) - - require.Len(t, snapshot, 2, - "only in-progress (committedMessageID == 0) parts should be kept") - require.Equal(t, "in-progress-1", partText(snapshot[0])) - require.Equal(t, "in-progress-2", partText(snapshot[1])) -} - -// TestSnapshotBufferLocked_AllInProgressReturnsAll covers the -// fresh-load convention: when no assistant message has committed -// yet, every buffered part is in-progress and must be delivered. -func TestSnapshotBufferLocked_AllInProgressReturnsAll(t *testing.T) { - t.Parallel() - - buffer := []bufferedStreamPart{ - makeInProgressPart("a"), - makeInProgressPart("b"), - makeInProgressPart("c"), - } - - snapshot := snapshotBufferLocked(buffer) - - require.Len(t, snapshot, 3, - "all in-progress parts must be delivered to the subscriber") - require.Equal(t, "a", partText(snapshot[0])) - require.Equal(t, "b", partText(snapshot[1])) - require.Equal(t, "c", partText(snapshot[2])) -} - -// TestSnapshotBufferLocked_EmptyBufferReturnsNil documents that -// snapshotBufferLocked returns nil (not an empty slice) for an -// empty buffer, matching the prior append-from-nil behavior. -func TestSnapshotBufferLocked_EmptyBufferReturnsNil(t *testing.T) { - t.Parallel() - - require.Nil(t, snapshotBufferLocked(nil)) - require.Nil(t, snapshotBufferLocked([]bufferedStreamPart{})) -} - -// TestSnapshotBufferLocked_AllCommittedReturnsEmpty covers the -// natural resting point after an assistant turn commits and before -// the next turn starts streaming: every buffered part has been -// claimed and must be filtered out. The snapshot must be empty so -// reconnecting subscribers do not re-render content that is already -// available as a durable message. -func TestSnapshotBufferLocked_AllCommittedReturnsEmpty(t *testing.T) { - t.Parallel() - - buffer := []bufferedStreamPart{ - makeCommittedPart(100, "a"), - makeCommittedPart(100, "b"), - makeCommittedPart(200, "c"), - } - - require.Empty(t, snapshotBufferLocked(buffer)) -} - -// TestPublishToStream_AppendsAsInProgress verifies that parts -// buffered while the chat is streaming are tagged as in-progress -// (committedMessageID == 0) until publishMessage claims them via a -// committed assistant message. -func TestPublishToStream_AppendsAsInProgress(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - - chatID := uuid.New() - state := &chatStreamState{ - buffering: true, - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - } - server.chatStreams.Store(chatID, state) - - server.publishToStream(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - Part: codersdk.ChatMessageText("hello"), - }, - }) - - state.mu.Lock() - defer state.mu.Unlock() - require.Len(t, state.buffer, 1) - require.Equal(t, int64(0), state.buffer[0].committedMessageID, - "newly buffered parts must be in-progress until publishMessage claims them") - require.Equal(t, "hello", partText(state.buffer[0].event)) -} - -// TestClaimCommittedParts covers the per-role behavior of -// claimCommittedParts: -// - assistant messages claim every in-progress part with the -// committed message ID. -// - tool / user messages do not claim parts. -// - parts already claimed by an earlier assistant message are not -// re-claimed. -// - a chat with no live state is a no-op (does not panic). -func TestClaimCommittedParts(t *testing.T) { - t.Parallel() - - t.Run("AssistantClaimsAllInProgressParts", func(t *testing.T) { - t.Parallel() - - server := &Server{ - logger: slogtest.Make(t, nil), - clock: quartz.NewMock(t), - } - chatID := uuid.New() - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffer = []bufferedStreamPart{ - makeCommittedPart(100, "old-1"), - makeInProgressPart("new-1"), - makeInProgressPart("new-2"), - } - state.mu.Unlock() - - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 200, - Role: database.ChatMessageRoleAssistant, - }) - - state.mu.Lock() - defer state.mu.Unlock() - require.Equal(t, int64(100), state.buffer[0].committedMessageID, - "already-claimed parts must keep their original message ID") - require.Equal(t, int64(200), state.buffer[1].committedMessageID, - "in-progress parts must be claimed by the new message ID") - require.Equal(t, int64(200), state.buffer[2].committedMessageID, - "in-progress parts must be claimed by the new message ID") - }) - - t.Run("ToolMessageIsNoOp", func(t *testing.T) { - t.Parallel() - - server := &Server{ - logger: slogtest.Make(t, nil), - clock: quartz.NewMock(t), - } - chatID := uuid.New() - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffer = []bufferedStreamPart{ - makeInProgressPart("a"), - makeInProgressPart("b"), - } - state.mu.Unlock() - - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 300, - Role: database.ChatMessageRoleTool, - }) - - state.mu.Lock() - defer state.mu.Unlock() - require.Equal(t, int64(0), state.buffer[0].committedMessageID, - "tool messages must not claim buffered parts") - require.Equal(t, int64(0), state.buffer[1].committedMessageID, - "tool messages must not claim buffered parts") - }) - - t.Run("UserMessageIsNoOp", func(t *testing.T) { - t.Parallel() - - server := &Server{ - logger: slogtest.Make(t, nil), - clock: quartz.NewMock(t), - } - chatID := uuid.New() - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffer = []bufferedStreamPart{ - makeInProgressPart("a"), - } - state.mu.Unlock() - - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 400, - Role: database.ChatMessageRoleUser, - }) - - state.mu.Lock() - defer state.mu.Unlock() - require.Equal(t, int64(0), state.buffer[0].committedMessageID, - "user messages must not claim buffered parts") - }) - - t.Run("NoLiveStateIsNoOp", func(t *testing.T) { - t.Parallel() - - server := &Server{ - logger: slogtest.Make(t, nil), - clock: quartz.NewMock(t), - } - chatID := uuid.New() - - // No state stored: claimCommittedParts must not panic and - // must not allocate a new state for an unknown chat. - require.NotPanics(t, func() { - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 500, - Role: database.ChatMessageRoleAssistant, - }) - }) - _, ok := server.chatStreams.Load(chatID) - require.False(t, ok, - "claimCommittedParts must not create stream state for a chat that has none") - }) -} - -// TestSubscribeToStream_FiltersBufferedParts_Integration wires -// publishToStream, claimCommittedParts (via publishMessage), and -// subscribeToStream together to confirm the end-to-end contract: a -// reconnecting subscriber only receives parts that belong to the -// current in-progress turn, not parts that were already committed -// to durable assistant messages. -func TestSubscribeToStream_FiltersBufferedParts_Integration(t *testing.T) { - t.Parallel() - - mClock := quartz.NewMock(t) - server := &Server{ - logger: slogtest.Make(t, nil), - clock: mClock, - } - chatID := uuid.New() - - // Simulate the lifecycle: - // 1. Stream parts of turn A (still in-progress, no commit yet). - // 2. Commit turn A; its parts are claimed by message 100. - // 3. Stream parts of turn B (in-progress). - // 4. Commit turn B; its parts are claimed by message 200. - // 5. Stream parts of turn C (in-progress, never committed). - state := server.getOrCreateStreamState(chatID) - state.mu.Lock() - state.buffering = true - state.mu.Unlock() - - publishPart := func(text string) { - server.publishToStream(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: codersdk.ChatMessageRoleAssistant, - Part: codersdk.ChatMessageText(text), - }, - }) - } - - publishPart("A-1") - publishPart("A-2") - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 100, - Role: database.ChatMessageRoleAssistant, - }) - publishPart("B-1") - publishPart("B-2") - server.claimCommittedParts(chatID, database.ChatMessage{ - ID: 200, - Role: database.ChatMessageRoleAssistant, - }) - publishPart("C-1") - - // Reconnecting subscriber: only the currently in-progress turn - // (turn C) survives the filter, no matter what cursor the - // client passes through SubscribeAuthorized (the filter no - // longer depends on the cursor). - snapshot, _, _, cancel := server.subscribeToStream(chatID) - defer cancel() - - texts := make([]string, 0, len(snapshot)) - for _, ev := range snapshot { - texts = append(texts, partText(ev)) - } - require.Equal(t, []string{"C-1"}, texts, - "only in-progress (un-claimed) buffered parts must survive the filter") -} - -// TestPrimeWorkspaceMCPCache_SuccessOnFirstAttempt verifies the -// onChatUpdated cache primer path: when create_workspace / -// start_workspace finish waitForAgentReady and the agent's MCP -// server is already advertising tools, a single ListMCPTools call -// populates the cache so the next PrepareTools step is a cache hit -// and does not need to dial. func TestPrimeWorkspaceMCPCache_SuccessOnFirstAttempt(t *testing.T) { t.Parallel() @@ -6668,81 +4560,134 @@ func TestPrimeWorkspaceMCPCache_ExitsOnContextCancel(t *testing.T) { require.False(t, ok, "primer must not cache anything when canceled") } -func TestPersistChatContextSummarySetsAPIKeyID(t *testing.T) { +// TestGetWorkspaceConnBumpsWorkspaceUsage verifies that acquiring a +// workspace agent connection bumps the workspace's last_used_at via +// the usage tracker and extends the build's autostop deadline. +func TestGetWorkspaceConnBumpsWorkspaceUsage(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) - ctx := context.Background() + ctx := testutil.Context(t, testutil.WaitLong) user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + + // Create a workspace with a full build chain so we can verify + // both last_used_at (dormancy) and deadline (autostop) bumps. + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + tmpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + CreatedBy: user.ID, + }) + require.NoError(t, db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{ + ID: tmpl.ID, + UpdatedAt: dbtime.Now(), + AllowUserAutostop: true, + ActivityBump: int64(time.Hour), + })) + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: org.ID, + TemplateID: tmpl.ID, + Ttl: sql.NullInt64{Valid: true, Int64: int64(8 * time.Hour)}, + }) + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + CompletedAt: sql.NullTime{ + Valid: true, + Time: dbtime.Now().Add(-30 * time.Minute), + }, + }) + // Build deadline is 30 minutes in the past, close enough to + // be bumped by the 1-hour activity bump. + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + TemplateVersionID: tv.ID, + JobID: pj.ID, + Transition: database.WorkspaceTransitionStart, + Deadline: dbtime.Now().Add(-30 * time.Minute), + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: pj.ID, + }) + dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: res.ID, + }) + originalDeadline := build.Deadline + chat := dbgen.Chat(t, db, database.Chat{ OwnerID: user.ID, OrganizationID: org.ID, LastModelConfigID: modelConfig.ID, - }) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{ - UserID: user.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, }) - server := &Server{db: db} - persistAndAssertSummaryKey := func( - summaryCtx context.Context, - chatID uuid.UUID, - activeAPIKeyID string, - wantAPIKeyID string, - toolCallID string, - ) { - t.Helper() + // Usage tracker with manual tick/flush so the test controls + // when last_used_at is written to the DB. + flushTick := make(chan time.Time) + flushDone := make(chan int, 1) + tracker := workspacestats.NewTracker(db, + workspacestats.TrackerWithTickFlush(flushTick, flushDone), + workspacestats.TrackerWithLogger(slogtest.Make(t, nil)), + ) + t.Cleanup(func() { tracker.Close() }) - err := server.persistChatContextSummary( - summaryCtx, - chatID, - modelConfig.ID, - activeAPIKeyID, - toolCallID, - chatloop.CompactionResult{ - SystemSummary: "summarized context", - SummaryReport: "context was summarized", - ThresholdPercent: 70, - UsagePercent: 85.0, - ContextTokens: 8500, - ContextLimit: 10000, - }, - ) - require.NoError(t, err) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + mockConn.EXPECT().AwaitReachable(gomock.Any()).Return(true).AnyTimes() - msgs, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) - require.NoError(t, err) - - // GetChatMessagesForPromptByChatID uses a compaction boundary CTE - // that selects compressed=true, visibility='model'. Only the user - // summary qualifies; the assistant (visibility=user) and tool - // result (visibility=both) are excluded by the CTE filter. - require.NotEmpty(t, msgs) - - var foundUserSummary bool - for _, msg := range msgs { - if msg.Role == database.ChatMessageRoleUser { - foundUserSummary = true - require.True(t, msg.APIKeyID.Valid, "summary user message must have APIKeyID set") - require.Equal(t, wantAPIKeyID, msg.APIKeyID.String, "summary user message APIKeyID must match") - } - } - require.True(t, foundUserSummary, "expected to find compressed user summary message") + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: quartz.NewReal(), + agentInactiveDisconnectTimeout: 30 * time.Second, + dialTimeout: testutil.WaitLong, + usageTracker: tracker, + agentConnFn: func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + }, } - persistAndAssertSummaryKey(ctx, chat.ID, apiKey.ID, apiKey.ID, "tool-call-id-1") + currentChat := chat + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: &sync.Mutex{}, + currentChat: ¤tChat, + loadChatSnapshot: db.GetChatByID, + } + t.Cleanup(workspaceCtx.close) - fallbackChat := dbgen.Chat(t, db, database.Chat{ - OwnerID: user.ID, - OrganizationID: org.ID, - LastModelConfigID: modelConfig.ID, - }) - fallbackKey, _ := dbgen.APIKey(t, db, database.APIKey{ - UserID: user.ID, - }) - fallbackCtx := aibridge.WithDelegatedAPIKeyID(ctx, fallbackKey.ID) - persistAndAssertSummaryKey(fallbackCtx, fallbackChat.ID, "", fallbackKey.ID, "tool-call-id-2") + _, err := workspaceCtx.getWorkspaceConn(ctx) + require.NoError(t, err) + + // getWorkspaceConn tracks usage synchronously; flushing the + // tracker must write last_used_at for the linked workspace. + testutil.RequireSend(ctx, t, flushTick, time.Now()) + count := testutil.RequireReceive(ctx, t, flushDone) + require.Greater(t, count, 0, + "expected the usage tracker to flush the chat workspace") + + updatedWs, err := db.GetWorkspaceByID(ctx, ws.ID) + require.NoError(t, err) + require.True(t, updatedWs.LastUsedAt.After(ws.LastUsedAt), + "workspace last_used_at should have been bumped") + + // The activity bump runs synchronously inside + // getWorkspaceConn, so the deadline is already extended. + // ±2 minute tolerance mirrors activitybump_test.go. + updatedBuild, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) + require.NoError(t, err) + require.True(t, updatedBuild.Deadline.After(originalDeadline), + "workspace build deadline should have been bumped") + now := dbtime.Now() + require.True(t, updatedBuild.Deadline.After(now.Add(time.Hour-2*time.Minute))) + require.True(t, updatedBuild.Deadline.Before(now.Add(time.Hour+2*time.Minute))) } diff --git a/coderd/x/chatd/chatd_retry_test.go b/coderd/x/chatd/chatd_retry_test.go new file mode 100644 index 0000000000..07f8cd5934 --- /dev/null +++ b/coderd/x/chatd/chatd_retry_test.go @@ -0,0 +1,231 @@ +package chatd_test + +import ( + "context" + "encoding/json" + "sync/atomic" + "testing" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd" + "github.com/coder/coder/v2/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestActiveServer_RetryStatePersistedDuringBackoff(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) + var calls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if calls.Add(1) == 1 { + return chattest.OpenAIRateLimitResponse() + } + return chattest.OpenAIStreamingResponse(openAITextChunksWithStop("recovered")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.Clock = clock + }) + + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + withRetry := waitForChatRetryState(ctx, t, db, chat.ID) + require.Equal(t, database.ChatStatusRunning, withRetry.Status) + require.True(t, withRetry.RetryState.Valid) + require.Equal(t, withRetry.SnapshotVersion, withRetry.RetryStateVersion) + require.Equal(t, int64(1), withRetry.GenerationAttempt) + + var retryPayload codersdk.ChatStreamRetry + require.NoError(t, json.Unmarshal(withRetry.RetryState.RawMessage, &retryPayload)) + require.Equal(t, 1, retryPayload.Attempt) + require.Equal(t, int64(1000), retryPayload.DelayMs) + require.Equal(t, "OpenAI is rate limiting requests.", retryPayload.Error) + require.Equal(t, codersdk.ChatErrorKindRateLimit, retryPayload.Kind) + require.Equal(t, "openai", retryPayload.Provider) + require.Equal(t, 429, retryPayload.StatusCode) + require.False(t, retryPayload.RetryingAt.IsZero()) + + advanceToNextTimer(ctx, clock) + advanceUntilProviderCall(ctx, clock, &calls, 2) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(2), calls.Load()) + latest, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.False(t, latest.RetryState.Valid) + require.Greater(t, latest.RetryStateVersion, withRetry.RetryStateVersion) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + requireTextPart(t, messages[len(messages)-1], "recovered") +} + +func TestActiveServer_RetryStreamSilenceTimeoutAndClassification(t *testing.T) { + t.Parallel() + + t.Run("rate limit retry recovers and records metric", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + reg := prometheus.NewRegistry() + var calls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if calls.Add(1) == 1 { + return chattest.OpenAIRateLimitResponse() + } + return chattest.OpenAIStreamingResponse(openAITextChunksWithStop("recovered")...) + }) + user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "gpt-4o", + Enabled: true, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.PrometheusRegistry = reg + }) + + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(2), calls.Load()) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + requireTextPart(t, messages[len(messages)-1], "recovered") + requireRetryCounter(t, reg, "coderd_chatd_stream_retries_total", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o", + "kind": string(codersdk.ChatErrorKindRateLimit), + "chain_broken": "false", + }) + }) + + t.Run("stream silence timeout retry recovers", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + clock := quartz.NewMock(t).WithLogger(quartz.NoOpLogger) + reg := prometheus.NewRegistry() + var calls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if calls.Add(1) == 1 { + <-req.Request.Context().Done() + return chattest.OpenAIStreamingResponse(openAITextChunksWithStop("timed out")...) + } + return chattest.OpenAIStreamingResponse(openAITextChunksWithStop("recovered")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.Clock = clock + cfg.PrometheusRegistry = reg + }) + + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + advanceUntilProviderCall(ctx, clock, &calls, 1) + advanceToNextTimer(ctx, clock) + advanceUntilProviderCall(ctx, clock, &calls, 2) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(2), calls.Load()) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + requireTextPart(t, messages[len(messages)-1], "recovered") + requireRetryCounter(t, reg, "coderd_chatd_stream_retries_total", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o-mini", + "kind": string(codersdk.ChatErrorKindStreamSilenceTimeout), + "chain_broken": "false", + }) + }) +} + +func requireRetryCounter(t *testing.T, reg *prometheus.Registry, name string, wantValue float64, wantLabels map[string]string) { + t.Helper() + require.True(t, hasRetryCounter(t, reg, name, wantValue, wantLabels), "metric %s not found", name) +} + +func hasRetryCounter(t *testing.T, reg *prometheus.Registry, name string, wantValue float64, wantLabels map[string]string) bool { + t.Helper() + + families, err := reg.Gather() + require.NoError(t, err) + for _, family := range families { + if family.GetName() != name { + continue + } + for _, metric := range family.GetMetric() { + if metric.GetCounter().GetValue() != wantValue { + continue + } + labels := map[string]string{} + for _, label := range metric.GetLabel() { + labels[label.GetName()] = label.GetValue() + } + matches := true + for key, want := range wantLabels { + if labels[key] != want { + matches = false + break + } + } + if matches { + return true + } + } + return false + } + return false +} + +func waitForChatRetryState(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID) database.Chat { + t.Helper() + var chat database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + latest, err := db.GetChatByID(ctx, chatID) + if err != nil { + return false + } + chat = latest + return latest.RetryState.Valid + }, testutil.IntervalFast) + return chat +} + +func advanceUntilProviderCall(ctx context.Context, clock *quartz.Mock, calls *atomic.Int32, want int32) { + for calls.Load() < want { + advanceToNextTimer(ctx, clock) + } +} + +func advanceToNextTimer(ctx context.Context, clock *quartz.Mock) { + _, waiter := clock.AdvanceNext() + waiter.MustWait(ctx) +} + +func openAITextChunksWithStop(deltas ...string) []chattest.OpenAIChunk { + chunks := chattest.OpenAITextChunks(deltas...) + if len(chunks) == 0 { + return nil + } + chunks[len(chunks)-1].Choices[0].FinishReason = "stop" + return chunks +} diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 7b0a66eb43..b25a25fc24 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -21,16 +21,19 @@ import ( "testing" "time" + "charm.land/fantasy" + fantasyanthropic "charm.land/fantasy/providers/anthropic" "github.com/google/uuid" mcpgo "github.com/mark3labs/mcp-go/mcp" mcpserver "github.com/mark3labs/mcp-go/server" "github.com/prometheus/client_golang/prometheus" + io_prometheus_client "github.com/prometheus/client_model/go" "github.com/sqlc-dev/pqtype" - "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "golang.org/x/xerrors" + "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/agent/agentcontextconfig" "github.com/coder/coder/v2/agent/agenttest" @@ -44,13 +47,13 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" - coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" - "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/slice" "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/coderd/x/chatd/chatadvisor" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatsanitize" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/codersdk" @@ -70,6 +73,12 @@ type recordedOpenAIRequest struct { ContentLength int64 } +func testAPIKeyID(t testing.TB, db database.Store, userID uuid.UUID) string { + t.Helper() + key, _ := dbgen.APIKey(t, db, database.APIKey{ID: uuid.NewString(), UserID: userID}) + return key.ID +} + type chatAIGatewayRecordedRequest struct { ProviderName string Source aibridge.Source @@ -81,10 +90,11 @@ type chatAIGatewayRecordedRequest struct { } type chatAIGatewayTestFactory struct { - target *url.URL - transport http.RoundTripper - mu sync.Mutex - requests []chatAIGatewayRecordedRequest + target *url.URL + transport http.RoundTripper + preservePath bool + mu sync.Mutex + requests []chatAIGatewayRecordedRequest } func newChatAIGatewayTestFactory(t testing.TB, targetBaseURL string) *chatAIGatewayTestFactory { @@ -95,6 +105,14 @@ func newChatAIGatewayTestFactory(t testing.TB, targetBaseURL string) *chatAIGate return &chatAIGatewayTestFactory{target: target, transport: http.DefaultTransport} } +func newChatAIGatewayPreservePathTestFactory(t testing.TB, targetBaseURL string) *chatAIGatewayTestFactory { + t.Helper() + + target, err := url.Parse(targetBaseURL) + require.NoError(t, err) + return &chatAIGatewayTestFactory{target: target, transport: http.DefaultTransport, preservePath: true} +} + func (f *chatAIGatewayTestFactory) TransportFor(providerName string, source aibridge.Source) (http.RoundTripper, error) { return chatAIGatewayRoundTripper{factory: f, providerName: providerName, source: source}, nil } @@ -126,9 +144,13 @@ func (t chatAIGatewayRoundTripper) RoundTrip(req *http.Request) (*http.Response, t.factory.mu.Unlock() targetURL := *t.factory.target - targetURL.Path = strings.TrimPrefix(req.URL.Path, "/v1") - if targetURL.Path == "" { - targetURL.Path = "/" + if t.factory.preservePath { + targetURL.Path = req.URL.Path + } else { + targetURL.Path = strings.TrimPrefix(req.URL.Path, "/v1") + if targetURL.Path == "" { + targetURL.Path = "/" + } } targetURL.RawQuery = req.URL.RawQuery @@ -244,10 +266,10 @@ func newWorkspaceToolTestServer( mockConn.EXPECT().ListMCPTools(gomock.Any()). Return(workspacesdk.ListMCPToolsResponse{}, nil).AnyTimes() mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()). - Return(workspacesdk.LSResponse{}, nil).AnyTimes() + Return(workspacesdk.LSResponse{AbsolutePathString: "/home/coder"}, nil).AnyTimes() mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). DoAndReturn(func(_ context.Context, path string, _, _ int64) (io.ReadCloser, string, error) { - if path == "/home/coder/PLAN.md" { + if strings.HasPrefix(path, "/home/coder/.coder/plans/PLAN-") || path == "/home/coder/PLAN.md" { return io.NopCloser(strings.NewReader(planContent)), "", nil } return io.NopCloser(strings.NewReader("")), "", nil @@ -261,57 +283,6 @@ func newWorkspaceToolTestServer( }) } -func TestInterruptChatBroadcastsStatusAcrossInstances(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replicaA := newTestServer(t, db, ps, uuid.New()) - replicaB := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replicaA.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "interrupt-me", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - runningWorker := uuid.New() - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: runningWorker, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - _, events, cancel, ok := replicaB.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - updated := replicaA.InterruptChat(ctx, chat) - require.Equal(t, database.ChatStatusWaiting, updated.Status) - require.False(t, updated.WorkerID.Valid) - - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeStatus && event.Status != nil { - return event.Status.Status == codersdk.ChatStatusWaiting - } - t.Logf("skipping unexpected event: type=%s", event.Type) - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) -} - func TestSubagentChatExcludesWorkspaceProvisioningTools(t *testing.T) { t.Parallel() @@ -887,7 +858,12 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { ClientType: database.ChatClientTypeApi, }) - exploreChat := dbgen.Chat(t, db, database.Chat{ + userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("inspect the codebase"), + }) + require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + createdExplore, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -900,21 +876,22 @@ func TestExploreChatUsesPersistedMCPSnapshot(t *testing.T) { ChatMode: database.ChatModeExplore, Valid: true, }, - Status: database.ChatStatusPending, MCPServerIDs: []uuid.UUID{mcpConfig.ID}, ClientType: database.ChatClientTypeApi, - }) - - dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: exploreChat.ID, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - ModelConfigID: uuid.NullUUID{UUID: webSearchModel.ID, Valid: true}, - Role: database.ChatMessageRoleUser, - Content: pqtype.NullRawMessage{ - RawMessage: json.RawMessage(`[{"type":"text","text":"inspect the codebase"}]`), - Valid: true, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: userContent, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: webSearchModel.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, }, }) + require.NoError(t, err) + exploreChat := createdExplore.Chat ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) @@ -1024,6 +1001,7 @@ func TestRootExploreChatStaysBuiltinOnlyAtRuntime(t *testing.T) { exploreChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root-explore-builtin-only", ModelConfigID: model.ID, ChatMode: database.NullChatMode{ @@ -1109,6 +1087,7 @@ func TestRootExploreChatExcludesWebSearchProviderToolAtRuntime(t *testing.T) { exploreChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root-explore-no-provider-web-search", ModelConfigID: webSearchModel.ID, ChatMode: database.NullChatMode{ @@ -1238,6 +1217,7 @@ func TestExploreChatSendMessageCannotMutateMCPSnapshot(t *testing.T) { rootChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "runtime-parent", ModelConfigID: model.ID, MCPServerIDs: []uuid.UUID{parentConfig.ID}, @@ -1280,6 +1260,7 @@ func TestExploreChatSendMessageCannotMutateMCPSnapshot(t *testing.T) { _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: exploreChat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("inspect the codebase again")}, MCPServerIDs: &updatedMCPServerIDs, }) @@ -1442,6 +1423,7 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { planChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-mode-root-mcp-visibility", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -1461,6 +1443,7 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { askChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "ask-mode-root-mcp-visibility", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -1507,91 +1490,13 @@ func TestPlanModeRootChatAllowsApprovedExternalMCPTools(t *testing.T) { "ask mode should continue exposing workspace MCP tools") } -func TestInterruptChatClearsWorkerInDatabase(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "db-transition", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - updated := replica.InterruptChat(ctx, chat) - require.Equal(t, database.ChatStatusWaiting, updated.Status) - require.False(t, updated.WorkerID.Valid) - - fromDB, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, fromDB.Status) - require.False(t, fromDB.WorkerID.Valid) -} - -func TestArchiveChatMovesPendingChatToWaiting(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "archive-pending", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - fromDB, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, fromDB.Status) - require.False(t, fromDB.WorkerID.Valid) - require.False(t, fromDB.StartedAt.Valid) - require.False(t, fromDB.HeartbeatAt.Valid) - require.True(t, fromDB.Archived) - require.Zero(t, fromDB.PinOrder) -} - // TestUnarchiveChildChat covers the deterministic branches of the -// Server.UnarchiveChat child path: happy path, archived-parent reject, -// and already-active no-op. +// Server.UnarchiveChat child path: every child unarchive attempt is +// rejected with chatd.ErrArchiveRequiresRootChat. func TestUnarchiveChildChat(t *testing.T) { t.Parallel() - t.Run("ChildWithActiveParentUnarchives", func(t *testing.T) { + t.Run("ChildWithActiveParentRejected", func(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -1601,11 +1506,12 @@ func TestUnarchiveChildChat(t *testing.T) { parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) - require.NoError(t, replica.UnarchiveChat(ctx, child)) + err := replica.UnarchiveChat(ctx, child) + require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) dbChild, err := db.GetChatByID(ctx, child.ID) require.NoError(t, err) - require.False(t, dbChild.Archived, "child should be unarchived") + require.True(t, dbChild.Archived, "child should remain archived") dbParent, err := db.GetChatByID(ctx, parent.ID) require.NoError(t, err) @@ -1625,14 +1531,14 @@ func TestUnarchiveChildChat(t *testing.T) { require.NoError(t, err) err = replica.UnarchiveChat(ctx, child) - require.ErrorIs(t, err, chatd.ErrChildUnarchiveParentArchived) + require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) dbChild, err := db.GetChatByID(ctx, child.ID) require.NoError(t, err) require.True(t, dbChild.Archived, "child should remain archived") }) - t.Run("AlreadyActiveChildNoOp", func(t *testing.T) { + t.Run("ActiveChildRejected", func(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -1642,7 +1548,8 @@ func TestUnarchiveChildChat(t *testing.T) { _, child := insertParentWithActiveChild(t, db, user, org, model) - require.NoError(t, replica.UnarchiveChat(ctx, child)) + err := replica.UnarchiveChat(ctx, child) + require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) dbChild, err := db.GetChatByID(ctx, child.ID) require.NoError(t, err) @@ -1650,6 +1557,60 @@ func TestUnarchiveChildChat(t *testing.T) { }) } +// TestArchiveChat_RejectsChildChat verifies that Server.ArchiveChat +// refuses every child chat with chatd.ErrArchiveRequiresRootChat +// regardless of the family's current archive state. Archive state +// changes must always be issued against the root chat so the whole +// family flips together. +func TestArchiveChat_RejectsChildChat(t *testing.T) { + t.Parallel() + + t.Run("ActiveChildRejected", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + replica := newTestServer(t, db, ps, uuid.New()) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + + parent, child := insertParentWithActiveChild(t, db, user, org, model) + + err := replica.ArchiveChat(ctx, child) + require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat) + + dbChild, err := db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.False(t, dbChild.Archived, "child should stay active after rejected archive") + + dbParent, err := db.GetChatByID(ctx, parent.ID) + require.NoError(t, err) + require.False(t, dbParent.Archived, "parent should stay active after rejected child archive") + }) + + t.Run("AlreadyArchivedChildRejected", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + replica := newTestServer(t, db, ps, uuid.New()) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + + parent, child := insertParentWithArchivedChild(ctx, t, db, user, org, model) + + err := replica.ArchiveChat(ctx, child) + require.ErrorIs(t, err, chatd.ErrArchiveRequiresRootChat, + "child archive must be rejected even when the child is already archived") + + dbChild, err := db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.True(t, dbChild.Archived, "child archived flag should not change") + + dbParent, err := db.GetChatByID(ctx, parent.ID) + require.NoError(t, err) + require.False(t, dbParent.Archived, "parent should stay active") + }) +} + // insertParentWithActiveChild creates a parent chat and an active // child chat linked to it. Both are returned in their initial // (active) state. @@ -1698,141 +1659,6 @@ func insertParentWithArchivedChild( return parent, child } -func TestArchiveChatInterruptsActiveProcessing(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - streamStarted := make(chan struct{}) - streamCanceled := make(chan struct{}) - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("title") - } - chunks := make(chan chattest.OpenAIChunk, 1) - go func() { - defer close(chunks) - chunks <- chattest.OpenAITextChunks("partial")[0] - select { - case <-streamStarted: - default: - close(streamStarted) - } - <-req.Context().Done() - select { - case <-streamCanceled: - default: - close(streamCanceled) - } - }() - return chattest.OpenAIResponse{StreamingChunks: chunks} - }) - - server := newActiveTestServer(t, db, ps) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "archive-interrupt", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusRunning && fromDB.WorkerID.Valid - }, testutil.IntervalFast) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case <-streamStarted: - return true - default: - return false - } - }, testutil.IntervalFast) - - _, events, cancel, ok := server.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer cancel() - - queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - require.NotNil(t, queuedResult.QueuedMessage) - - err = server.ArchiveChat(ctx, chat) - require.NoError(t, err) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case <-streamCanceled: - return true - default: - return false - } - }, testutil.IntervalFast) - - gotWaitingStatus := false - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - for { - select { - case ev := <-events: - if ev.Type == codersdk.ChatStreamEventTypeStatus && - ev.Status != nil && - ev.Status.Status == codersdk.ChatStatusWaiting { - gotWaitingStatus = true - return true - } - default: - return gotWaitingStatus - } - } - }, testutil.IntervalFast) - require.True(t, gotWaitingStatus, "expected a waiting status event after archive") - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Archived && - fromDB.Status == database.ChatStatusWaiting && - !fromDB.WorkerID.Valid && - !fromDB.StartedAt.Valid && - !fromDB.HeartbeatAt.Valid - }, testutil.IntervalFast) - - queuedMessages, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, queuedMessages, 1) - require.Equal(t, queuedResult.QueuedMessage.ID, queuedMessages[0].ID) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - userMessages := 0 - for _, msg := range messages { - if msg.Role == database.ChatMessageRoleUser { - userMessages++ - } - } - require.Equal(t, 1, userMessages, "expected queued message to stay queued after archive") -} - func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { t.Parallel() @@ -1845,6 +1671,7 @@ func TestUpdateChatHeartbeatsRequiresOwnership(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "heartbeat-ownership", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -1960,6 +1787,7 @@ func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "queue-when-busy", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -1978,6 +1806,7 @@ func TestSendMessageQueueBehaviorQueuesWhenBusy(t *testing.T) { result, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, BusyBehavior: chatd.SendMessageBusyBehaviorQueue, }) @@ -2033,6 +1862,7 @@ func TestPlanTurnPromptContract(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), OrganizationID: org.ID, Title: "plan-turn-prompt-contract", ModelConfigID: model.ID, @@ -2073,73 +1903,6 @@ func TestPlanTurnPromptContract(t *testing.T) { } } -func TestSendMessageQueuesWhenWaitingWithQueuedBacklog(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "queue-when-waiting-with-backlog", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("older queued"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }) - require.NoError(t, err) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) - - result, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("newer queued")}, - }) - require.NoError(t, err) - require.True(t, result.Queued) - require.NotNil(t, result.QueuedMessage) - require.Equal(t, database.ChatStatusWaiting, result.Chat.Status) - - queued, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, queued, 2) - - olderSDK := db2sdk.ChatQueuedMessage(queued[0]) - require.Len(t, olderSDK.Content, 1) - require.Equal(t, "older queued", olderSDK.Content[0].Text) - - newerSDK := db2sdk.ChatQueuedMessage(queued[1]) - require.Len(t, newerSDK.Content, 1) - require.Equal(t, "newer queued", newerSDK.Content[0].Text) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) -} - func TestSendMessageRejectsInvalidQueuedModelConfigID(t *testing.T) { t.Parallel() @@ -2160,6 +1923,7 @@ func TestSendMessageRejectsInvalidQueuedModelConfigID(t *testing.T) { invalidModelConfigID := uuid.New() _, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, ModelConfigID: invalidModelConfigID, }) @@ -2170,185 +1934,6 @@ func TestSendMessageRejectsInvalidQueuedModelConfigID(t *testing.T) { require.Empty(t, queued) } -func TestSendMessageInterruptBehaviorQueuesAndInterruptsWhenBusy(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newStartedTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "interrupt-when-busy", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // CreateChat calls signalWake which triggers processOnce in - // the background. Wait for that processing to finish so it - // doesn't race with the manual status update below. - waitForChatProcessed(ctx, t, db, chat.ID, replica) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - result, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("interrupt")}, - BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, - }) - require.NoError(t, err) - - // The message should be queued, not inserted directly. - require.True(t, result.Queued) - require.NotNil(t, result.QueuedMessage) - - // The chat should transition to waiting (interrupt signal), - // not pending. - require.Equal(t, database.ChatStatusWaiting, result.Chat.Status) - - fromDB, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, fromDB.Status) - - // The message should be in the queue, not in chat_messages. - queued, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, queued, 1) - - // Only messages from the initial processing round should be in - // chat_messages (user + assistant). The "interrupt" message must - // be in the queue, not inserted directly. - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 2) -} - -func TestEditMessageUpdatesAndTruncatesAndClearsQueue(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "edit-message", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, - }) - require.NoError(t, err) - - initialMessages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, initialMessages, 1) - editedMessageID := initialMessages[0].ID - - _, err = replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("follow-up")}, - BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, - }) - require.NoError(t, err) - _, err = replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("another")}, - BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, - }) - require.NoError(t, err) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("queued"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }) - require.NoError(t, err) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) - apiKeyID := apiKey.ID - editResult, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ - ChatID: chat.ID, - EditedMessageID: editedMessageID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, - APIKeyID: apiKeyID, - }) - require.NoError(t, err) - // The edited message is soft-deleted and a new message is inserted, - // so the returned message ID will differ from the original. - require.NotEqual(t, editedMessageID, editResult.Message.ID) - require.True(t, editResult.Message.APIKeyID.Valid) - require.Equal(t, apiKeyID, editResult.Message.APIKeyID.String) - require.Equal(t, database.ChatStatusPending, editResult.Chat.Status) - require.False(t, editResult.Chat.WorkerID.Valid) - - editedSDK := db2sdk.ChatMessage(editResult.Message) - require.Len(t, editedSDK.Content, 1) - require.Equal(t, "edited", editedSDK.Content[0].Text) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) - require.Equal(t, editResult.Message.ID, messages[0].ID) - require.True(t, messages[0].APIKeyID.Valid) - require.Equal(t, apiKeyID, messages[0].APIKeyID.String) - onlyMessage := db2sdk.ChatMessage(messages[0]) - require.Len(t, onlyMessage.Content, 1) - require.Equal(t, "edited", onlyMessage.Content[0].Text) - - queued, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, queued, 0) - - // WaitUntilIdleForTest drains the debug-cleanup goroutine - // from EditMessage. Must be called from the test goroutine - // (not inside require.Eventually) to avoid Add/Wait race. - chatd.WaitUntilIdleForTest(replica) - var chatFromDB database.Chat - require.Eventually(t, func() bool { - c, e := db.GetChatByID(ctx, chat.ID) - if e != nil { - return false - } - chatFromDB = c - return chatFromDB.Status != database.ChatStatusRunning - }, testutil.WaitShort, testutil.IntervalFast) - require.False(t, chatFromDB.WorkerID.Valid) -} - func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { t.Parallel() @@ -2379,6 +1964,7 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, Title: "test-with-workspace", ModelConfigID: model.ID, @@ -2416,6 +2002,7 @@ func TestCreateChatInsertsWorkspaceAwarenessMessage(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "test-without-workspace", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -2495,6 +2082,7 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { _, err = replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "over-limit", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -2517,244 +2105,6 @@ func TestCreateChatRejectsWhenUsageLimitReached(t *testing.T) { require.Len(t, afterChats, len(beforeChats)) } -func TestPromoteQueuedAllowsAlreadyQueuedMessageWhenUsageLimitReached(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newStartedTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) - - _, err := db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ - Enabled: true, - DefaultLimitMicros: 100, - Period: string(codersdk.ChatUsageLimitPeriodDay), - }) - require.NoError(t, err) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "queued-limit-reached", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // CreateChat calls signalWake which triggers processOnce in - // the background. Wait for that processing to finish so it - // doesn't race with the manual status update below. - waitForChatProcessed(ctx, t, db, chat.ID, replica) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - queuedResult, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, - APIKeyID: apiKey.ID, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - require.NotNil(t, queuedResult.QueuedMessage) - require.True(t, queuedResult.QueuedMessage.APIKeyID.Valid) - require.Equal(t, apiKey.ID, queuedResult.QueuedMessage.APIKeyID.String) - - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("assistant"), - }) - require.NoError(t, err) - - _ = dbgen.ChatMessage(t, db, database.ChatMessage{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Role: database.ChatMessageRoleAssistant, - ContentVersion: chatprompt.CurrentContentVersion, - Content: assistantContent, - TotalCostMicros: sql.NullInt64{Int64: 100, Valid: true}, - }) - - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) - - result, err := replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedResult.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.Equal(t, database.ChatMessageRoleUser, result.PromotedMessage.Role) - require.True(t, result.PromotedMessage.APIKeyID.Valid) - require.Equal(t, apiKey.ID, result.PromotedMessage.APIKeyID.String) - - queued, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Empty(t, queued) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 4) - require.Equal(t, database.ChatMessageRoleUser, messages[3].Role) -} - -func TestPromoteQueuedMessageUsesQueuedModelConfigID(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(t, db) - modelConfigB := insertChatModelConfigWithCallConfig( - t, - db, - user.ID, - "openai", - "gpt-4o-mini-promote-"+uuid.NewString(), - codersdk.ChatModelCallConfig{}, - ) - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: modelConfigA.ID, - Title: "promote queued uses stored model", - }) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued with model b")}) - require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - ModelConfigID: uuid.NullUUID{ - UUID: modelConfigB.ID, - Valid: true, - }, - }) - require.NoError(t, err) - - result, err := replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.True(t, result.PromotedMessage.ModelConfigID.Valid) - require.Equal(t, modelConfigB.ID, result.PromotedMessage.ModelConfigID.UUID) - - storedChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, modelConfigB.ID, storedChat.LastModelConfigID) - // The processor can pick up the pending chat immediately after - // promotion, so this test only requires that promotion moved it out of - // waiting and preserved the queued model configuration. - require.Contains(t, []database.ChatStatus{ - database.ChatStatusPending, - database.ChatStatusRunning, - }, storedChat.Status) -} - -func TestPromoteQueuedMessageReloadsChatWhenModelConfigChangesDuringPending(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(t, db) - modelConfigB := insertChatModelConfigWithCallConfig( - t, - db, - user.ID, - "openai", - "gpt-4o-mini-promote-pending-"+uuid.NewString(), - codersdk.ChatModelCallConfig{}, - ) - - watchEvents := make(chan struct { - payload codersdk.ChatWatchEvent - err error - }, 1) - cancelWatch, err := ps.SubscribeWithErr( - coderdpubsub.ChatWatchEventChannel(user.ID), - coderdpubsub.HandleChatWatchEvent(func(_ context.Context, payload codersdk.ChatWatchEvent, err error) { - select { - case watchEvents <- struct { - payload codersdk.ChatWatchEvent - err error - }{payload: payload, err: err}: - default: - } - }), - ) - require.NoError(t, err) - defer cancelWatch() - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - Status: database.ChatStatusPending, - OwnerID: user.ID, - LastModelConfigID: modelConfigA.ID, - Title: "promote queued reloads pending chat", - }) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued with new model")}) - require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - ModelConfigID: uuid.NullUUID{ - UUID: modelConfigB.ID, - Valid: true, - }, - }) - require.NoError(t, err) - - result, err := replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.True(t, result.PromotedMessage.ModelConfigID.Valid) - require.Equal(t, modelConfigB.ID, result.PromotedMessage.ModelConfigID.UUID) - - storedChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusPending, storedChat.Status) - require.Equal(t, modelConfigB.ID, storedChat.LastModelConfigID) - - select { - case event := <-watchEvents: - require.NoError(t, event.err) - require.Equal(t, codersdk.ChatWatchEventKindStatusChange, event.payload.Kind) - require.Equal(t, chat.ID, event.payload.Chat.ID) - require.Equal(t, codersdk.ChatStatusPending, event.payload.Chat.Status) - require.Equal(t, modelConfigB.ID, event.payload.Chat.LastModelConfigID) - case <-ctx.Done(): - t.Fatal("timed out waiting for status change watch event") - } -} - func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { t.Parallel() // TODO(CODAGT-353): Re-enable this test after the chatd notification flow @@ -2834,6 +2184,7 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "auto-promote per-turn model order", ModelConfigID: modelConfigA.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -2844,6 +2195,7 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { queuedB, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued b")}, ModelConfigID: modelConfigB.ID, BusyBehavior: chatd.SendMessageBusyBehaviorQueue, @@ -2853,6 +2205,7 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { queuedC, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued c")}, ModelConfigID: modelConfigC.ID, BusyBehavior: chatd.SendMessageBusyBehaviorQueue, @@ -2898,200 +2251,6 @@ func TestAutoPromoteQueuedMessagesPreservesPerTurnModelOrder(t *testing.T) { require.Equal(t, []uuid.UUID{modelConfigA.ID, modelConfigB.ID, modelConfigC.ID}, userModelConfigIDs) } -func TestAutoPromoteQueuedMessageFallsBackForLegacyQueuedRows(t *testing.T) { - t.Parallel() - - testAutoPromoteQueuedMessageFallback(t, uuid.NullUUID{}) -} - -func TestAutoPromoteQueuedMessageFallsBackForInvalidQueuedModelConfigID(t *testing.T) { - t.Parallel() - - testAutoPromoteQueuedMessageFallback(t, uuid.NullUUID{ - UUID: uuid.New(), - Valid: true, - }) -} - -func testAutoPromoteQueuedMessageFallback(t *testing.T, queuedModelConfigID uuid.NullUUID) { - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitSuperLong) - - firstRunStarted := make(chan struct{}) - secondRunStarted := make(chan struct{}, 1) - allowFirstRunFinish := make(chan struct{}) - var requestCount atomic.Int32 - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("title") - } - - switch requestCount.Add(1) { - case 1: - chunks := make(chan chattest.OpenAIChunk, 1) - go func() { - defer close(chunks) - chunks <- chattest.OpenAITextChunks("first run partial")[0] - select { - case <-firstRunStarted: - default: - close(firstRunStarted) - } - <-allowFirstRunFinish - }() - return chattest.OpenAIResponse{StreamingChunks: chunks} - default: - select { - case secondRunStarted <- struct{}{}: - default: - } - return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("fallback run done")...) - } - }) - - server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { - // Disable periodic polling so only signalWake can - // trigger the next processing run. - cfg.PendingChatAcquireInterval = time.Hour - }) - user, org, modelConfig := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "auto-promote queued fallback", - ModelConfigID: modelConfig.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - testutil.TryReceive(ctx, t, firstRunStarted) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("legacy queued row")}) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - ModelConfigID: queuedModelConfigID, - }) - require.NoError(t, err) - - close(allowFirstRunFinish) - - testutil.TryReceive(ctx, t, secondRunStarted) - require.GreaterOrEqual(t, requestCount.Load(), int32(2)) - chatd.WaitUntilIdleForTest(server) - - queuedMessages, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Empty(t, queuedMessages) - - storedChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, storedChat.Status) - require.Equal(t, modelConfig.ID, storedChat.LastModelConfigID) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var found bool - for _, message := range messages { - if message.Role != database.ChatMessageRoleUser { - continue - } - sdkMessage := db2sdk.ChatMessage(message) - require.Len(t, sdkMessage.Content, 1) - if sdkMessage.Content[0].Text != "legacy queued row" { - continue - } - require.True(t, message.ModelConfigID.Valid) - require.Equal(t, modelConfig.ID, message.ModelConfigID.UUID) - found = true - } - require.True(t, found) -} - -func TestPromoteQueuedMessageFallsBackForLegacyQueuedRows(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfigA := seedChatDependencies(t, db) - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: modelConfigA.ID, - Title: "promote queued legacy fallback", - }) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("legacy queued row")}) - require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - }) - require.NoError(t, err) - - result, err := replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.True(t, result.PromotedMessage.ModelConfigID.Valid) - require.Equal(t, modelConfigA.ID, result.PromotedMessage.ModelConfigID.UUID) - - storedChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, modelConfigA.ID, storedChat.LastModelConfigID) -} - -func TestPromoteQueuedMessageFallsBackForInvalidQueuedModelConfigID(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, modelConfig := seedChatDependencies(t, db) - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - LastModelConfigID: modelConfig.ID, - Title: "promote queued invalid fallback", - }) - - queuedContent, err := json.Marshal([]codersdk.ChatMessagePart{codersdk.ChatMessageText("invalid queued model")}) - require.NoError(t, err) - queuedMessage, err := db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent, - ModelConfigID: uuid.NullUUID{ - UUID: uuid.New(), - Valid: true, - }, - }) - require.NoError(t, err) - - result, err := replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.True(t, result.PromotedMessage.ModelConfigID.Valid) - require.Equal(t, modelConfig.ID, result.PromotedMessage.ModelConfigID.UUID) - - storedChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, modelConfig.ID, storedChat.LastModelConfigID) -} - func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { t.Parallel() @@ -3191,6 +2350,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "interrupt-autopromote-limit", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -3201,6 +2361,7 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, }) @@ -3214,8 +2375,9 @@ func TestInterruptAutoPromotionIgnoresLaterUsageLimitIncrease(t *testing.T) { testutil.TryReceive(ctx, t, secondRequestStarted) laterQueuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("later queued")}, + ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("later queued")}, }) require.NoError(t, err) require.True(t, laterQueuedResult.Queued) @@ -3297,6 +2459,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "edit-limit-reached", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, @@ -3327,6 +2490,7 @@ func TestEditMessageRejectsWhenUsageLimitReached(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: editedMessageID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -3360,6 +2524,7 @@ func TestEditMessageRejectsMissingMessage(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "missing-edited-message", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -3368,6 +2533,7 @@ func TestEditMessageRejectsMissingMessage(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: 999999, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -3387,6 +2553,7 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "non-user-edited-message", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -3408,6 +2575,7 @@ func TestEditMessageRejectsNonUserMessage(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: assistantMessage.ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -3434,6 +2602,7 @@ func TestEditMessageDebugCleanupDeletesPreEditRuns(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "debug-edit-cleanup", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("first")}, @@ -3483,6 +2652,7 @@ func TestEditMessageDebugCleanupDeletesPreEditRuns(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: editedMsgID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -3540,6 +2710,7 @@ func TestEditMessageDebugCleanupPreservesRecentRuns(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "debug-edit-buffer", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("first")}, @@ -3572,6 +2743,7 @@ func TestEditMessageDebugCleanupPreservesRecentRuns(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: editedMsgID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -3602,180 +2774,6 @@ func TestEditMessageDebugCleanupPreservesRecentRuns(t *testing.T) { "the buffered run must survive the fast retry") } -// TestArchiveChatDebugCleanupDeletesPreArchiveRuns verifies that -// ArchiveChat schedules cleanup that deletes pre-archive debug runs -// for the archived chat. Covers the archiveCutoff sampled from -// ArchiveChatByID's DB-stamped updated_at and the DeleteByChatID -// delete path. -func TestArchiveChatDebugCleanupDeletesPreArchiveRuns(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newDebugEnabledTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "debug-archive-cleanup", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - staleStart := time.Now().Add(-time.Hour).UTC().Truncate(time.Microsecond) - staleRun, err := db.InsertChatDebugRun(ctx, database.InsertChatDebugRunParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Kind: "chat_turn", - Status: "in_progress", - Provider: sql.NullString{String: "openai", Valid: true}, - Model: sql.NullString{String: model.Model, Valid: true}, - StartedAt: sql.NullTime{Time: staleStart, Valid: true}, - UpdatedAt: sql.NullTime{Time: staleStart, Valid: true}, - }) - require.NoError(t, err) - - // Freshly-inserted run inside the skew buffer must survive the - // fast retry for the same reason as the edit-cleanup buffer test. - recentStart := time.Now().Add(-time.Second).UTC().Truncate(time.Microsecond) - recentRun, err := db.InsertChatDebugRun(ctx, database.InsertChatDebugRunParams{ - ChatID: chat.ID, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - Kind: "chat_turn", - Status: "in_progress", - Provider: sql.NullString{String: "openai", Valid: true}, - Model: sql.NullString{String: model.Model, Valid: true}, - StartedAt: sql.NullTime{Time: recentStart, Valid: true}, - UpdatedAt: sql.NullTime{Time: recentStart, Valid: true}, - }) - require.NoError(t, err) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - chatd.WaitUntilIdleForTest(replica) - - // ErrNoRows proves the fast-retry path DELETED the row: - // FinalizeStale only UPDATEs in place, never deletes. - _, err = db.GetChatDebugRunByID(ctx, staleRun.ID) - require.ErrorIs(t, err, sql.ErrNoRows, - "pre-archive run outside the buffer should be deleted") - - remaining, err := db.GetChatDebugRunByID(ctx, recentRun.ID) - require.NoError(t, err, - "runs inside the clock-skew buffer must survive the fast retry") - require.Equal(t, recentRun.ID, remaining.ID) - - // Count the seeded survivors directly so the delete is verified - // not just by absence of a specific row. Scoped to seeded IDs - // because the archive transition may still race with other - // background debug writes. - remainingRuns, err := db.GetChatDebugRunsByChatID(ctx, database.GetChatDebugRunsByChatIDParams{ - ChatID: chat.ID, LimitVal: 100, - }) - require.NoError(t, err) - seeded := map[uuid.UUID]bool{staleRun.ID: true, recentRun.ID: true} - survivors := 0 - for _, r := range remainingRuns { - if seeded[r.ID] { - survivors++ - } - } - require.Equal(t, 1, survivors, - "only the recent (buffered) seeded run should survive") -} - -func TestRecoverStaleChatsPeriodically(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Use a very short stale threshold so the periodic recovery - // kicks in quickly during the test. - staleAfter := 500 * time.Millisecond - - // Create a chat and simulate a dead worker by setting the chat - // to running with a heartbeat in the past. - deadWorkerID := uuid.New() - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "stale-recovery-periodic", - LastModelConfigID: model.ID, - }) - - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: deadWorkerID, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, - }) - require.NoError(t, err) - - // Start a new replica. Its startup recovery will reset the - // chat (since the heartbeat is old), but the key point is that - // the periodic loop also recovers newly-stale chats. - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) - server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) - }) - - // The startup recovery should have already reset our stale - // chat. - require.Eventually(t, func() bool { - fromDB, err := db.GetChatByID(ctx, chat.ID) - if err != nil { - return false - } - return fromDB.Status == database.ChatStatusPending - }, testutil.WaitMedium, testutil.IntervalFast) - - // Now simulate a second stale chat appearing AFTER startup. - // This tests the periodic recovery, not just the startup one. - deadWorkerID2 := uuid.New() - chat2 := dbgen.Chat(t, db, database.Chat{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "stale-recovery-periodic-2", - LastModelConfigID: model.ID, - }) - - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat2.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: deadWorkerID2, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, - }) - require.NoError(t, err) - - // The periodic stale recovery loop (running at staleAfter/5 = - // 100ms intervals) should pick this up without a restart. - require.Eventually(t, func() bool { - fromDB, err := db.GetChatByID(ctx, chat2.ID) - if err != nil { - return false - } - return fromDB.Status == database.ChatStatusPending - }, testutil.WaitMedium, testutil.IntervalFast) -} - func TestRecoverStaleRequiresActionChat(t *testing.T) { t.Parallel() @@ -3784,107 +2782,168 @@ func TestRecoverStaleRequiresActionChat(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) user, org, model := seedChatDependencies(t, db) - // Use a very short stale threshold so the periodic recovery - // kicks in quickly during the test. - staleAfter := 500 * time.Millisecond + toolName := "my_dynamic_tool" + dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ + Name: toolName, + Description: "A test dynamic tool.", + InputSchema: mcpgo.ToolInputSchema{ + Type: "object", + Properties: map[string]any{}, + }, + }}) + require.NoError(t, err) - // Create a chat and set it to requires_action to simulate a - // client that disappeared while the chat was waiting for - // dynamic tool results. - chat := dbgen.Chat(t, db, database.Chat{ + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("hello"), + }) + require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, - Title: "stale-requires-action", LastModelConfigID: model.ID, - }) - - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRequiresAction, + Title: "stale-requires-action", + DynamicTools: nullRawMessage(dynamicToolsJSON), + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, }) require.NoError(t, err) - // Backdate updated_at so the chat appears stale to the - // recovery loop without needing time.Sleep. + toolCallID := "call_" + uuid.NewString() + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: toolCallID, + ToolName: toolName, + Args: json.RawMessage(`{}`), + }, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(db, ps, created.Chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }, + }) + return err + })) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) _, err = rawDB.ExecContext(ctx, - "UPDATE chats SET updated_at = $1 WHERE id = $2", - time.Now().Add(-time.Hour), chat.ID) + "UPDATE chats SET requires_action_deadline_at = $1 WHERE id = $2", + time.Now().Add(-time.Hour), created.Chat.ID) require.NoError(t, err) - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) + server := newTestServer(t, db, ps, uuid.New()) server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) + + chatResult := waitForTerminalChat(ctx, t, db, created.Chat.ID) + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + require.False(t, chatResult.RequiresActionDeadlineAt.Valid) + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: created.Chat.ID, }) - - // The stale recovery should transition the requires_action - // chat to error with the timeout message. - var chatResult database.Chat - require.Eventually(t, func() bool { - chatResult, err = db.GetChatByID(ctx, chat.ID) - if err != nil { - return false - } - return chatResult.Status == database.ChatStatusError - }, testutil.WaitMedium, testutil.IntervalFast) - - persistedError := requireChatLastErrorPayload(t, chatResult.LastError) - require.Equal(t, codersdk.ChatError{ - Message: "Dynamic tool execution timed out", - Kind: codersdk.ChatErrorKindGeneric, - }, persistedError) - require.False(t, chatResult.WorkerID.Valid) + require.NoError(t, err) + require.Len(t, messages, 4) + parts, err := chatprompt.ParseContent(messages[2]) + require.NoError(t, err) + require.Len(t, parts, 1) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, parts[0].Type) + require.Equal(t, toolCallID, parts[0].ToolCallID) + require.Equal(t, toolName, parts[0].ToolName) + require.True(t, parts[0].IsError) + require.JSONEq(t, `"Tool execution timed out"`, string(parts[0].Result)) } func TestNewReplicaRecoversStaleChatFromDeadReplica(t *testing.T) { t.Parallel() - db, ps := dbtestutil.NewDB(t) + db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) ctx := testutil.Context(t, testutil.WaitLong) user, org, model := seedChatDependencies(t, db) - // Simulate a chat left running by a dead replica with a stale - // heartbeat (well beyond the stale threshold). - deadReplicaID := uuid.New() - chat := dbgen.Chat(t, db, database.Chat{ + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("hello"), + }) + require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, OwnerID: user.ID, - Title: "orphaned-chat", LastModelConfigID: model.ID, - }) - - // Set the heartbeat far in the past so it's definitely stale. - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: deadReplicaID, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now().Add(-time.Hour), Valid: true}, + Title: "orphaned-chat", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, }) require.NoError(t, err) + deadWorkerID := uuid.New() + deadRunnerID := uuid.New() + machine := chatstate.NewChatMachine(db, ps, created.Chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: deadWorkerID, RunnerID: deadRunnerID}) + return err + })) + // Simulate a chat left running by a dead replica with a stale + // heartbeat (well beyond the stale threshold). + _, err = rawDB.ExecContext(ctx, + "UPDATE chat_heartbeats SET heartbeat_at = $1 WHERE chat_id = $2 AND runner_id = $3", + time.Now().Add(-time.Hour), created.Chat.ID, deadRunnerID) + require.NoError(t, err) + + newWorkerID := uuid.New() + server := newTestServer(t, db, ps, newWorkerID) // Start a new replica. It should recover the stale chat on // startup. - newReplica := newTestServer(t, db, ps, uuid.New()) - _ = newReplica + server.Start() + var recovered database.Chat require.Eventually(t, func() bool { - fromDB, err := db.GetChatByID(ctx, chat.ID) + recovered, err = db.GetChatByID(ctx, created.Chat.ID) if err != nil { return false } - return fromDB.Status == database.ChatStatusPending && - !fromDB.WorkerID.Valid + return recovered.Status == database.ChatStatusRunning && + recovered.WorkerID.Valid && recovered.WorkerID.UUID == newWorkerID && + recovered.RunnerID.Valid && recovered.RunnerID.UUID != deadRunnerID }, testutil.WaitMedium, testutil.IntervalFast) + + _, err = db.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: recovered.RunnerID.UUID, + }) + require.NoError(t, err) } func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) { @@ -3906,11 +2965,10 @@ func TestWaitingChatsAreNotRecoveredAsStale(t *testing.T) { // Start a replica with a short stale threshold. logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: testutil.WaitLong, InFlightChatStaleAfter: 500 * time.Millisecond, }) @@ -4003,6 +3061,7 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "status-snapshot", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -4015,8 +3074,15 @@ func TestSubscribeSnapshotIncludesStatusEvent(t *testing.T) { // Passive server: status is always Pending. require.NotEmpty(t, snapshot) - require.Equal(t, codersdk.ChatStreamEventTypeStatus, snapshot[0].Type) - require.NotNil(t, snapshot[0].Status) + statusIdx := -1 + for i, event := range snapshot { + if event.Type == codersdk.ChatStreamEventTypeStatus { + statusIdx = i + break + } + } + require.NotEqual(t, -1, statusIdx) + require.NotNil(t, snapshot[statusIdx].Status) } func TestPersistToolResultWithBinaryData(t *testing.T) { @@ -4117,6 +3183,7 @@ func TestPersistToolResultWithBinaryData(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "binary-tool-result", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -4217,11 +3284,10 @@ func TestRequiresActionChatPersistsWaitingStatusLabel(t *testing.T) { mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -4248,6 +3314,7 @@ func TestRequiresActionChatPersistsWaitingStatusLabel(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "requires-action-status-label", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -4285,6 +3352,684 @@ func TestRequiresActionChatPersistsWaitingStatusLabel(t *testing.T) { "expected no web push dispatch for a requires_action chat") } +func TestActiveServer_InterruptionBehavior(t *testing.T) { + t.Parallel() + + t.Run("partial stream commits synthetic tool result and promotes queued message", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + streamStarted := make(chan struct{}) + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + + if requestCount.Add(1) != 1 { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("queued response")...) + } + chunks := make(chan chattest.AnthropicChunk, 5) + go func() { + defer close(chunks) + chunks <- chattest.AnthropicChunk{ + Type: "message_start", + Message: chattest.AnthropicChunkMessage{ + ID: "msg-partial-interrupt", + Type: "message", + Role: "assistant", + Model: "claude-3-opus-20240229", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 0, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "text", + Text: "", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_delta", + Index: 0, + Delta: chattest.AnthropicDeltaBlock{Type: "text_delta", Text: "partial assistant output"}, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 1, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "tool_use", + ID: "interrupt-tool-1", + Name: "read_file", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_delta", + Index: 1, + Delta: chattest.AnthropicDeltaBlock{Type: "input_json_delta", PartialJSON: `{"path":"main.go"}`}, + } + select { + case <-streamStarted: + default: + close(streamStarted) + } + <-req.Context().Done() + }() + return chattest.AnthropicResponse{StreamingChunks: chunks} + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(0) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "interrupt-partial-tool", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("start and call a tool"), + }, + }) + require.NoError(t, err) + + testutil.TryReceive(ctx, t, streamStarted) + queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued after interrupt")}, + BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, queued.Queued) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.GreaterOrEqual(t, requestCount.Load(), int32(2)) + + messages := chatMessages(ctx, t, db, chat.ID) + var userTexts []string + var foundPartial bool + for _, msg := range messages { + parts, parseErr := chatprompt.ParseContent(msg) + require.NoError(t, parseErr) + switch msg.Role { + case database.ChatMessageRoleUser: + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText { + userTexts = append(userTexts, part.Text) + } + } + case database.ChatMessageRoleAssistant: + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText && strings.Contains(part.Text, "partial assistant output") { + foundPartial = true + } + } + } + } + require.Equal(t, []string{"start and call a tool", "queued after interrupt"}, userTexts) + require.True(t, foundPartial) + + parts := chatToolParts(ctx, t, db, chat.ID) + call := requireToolCallPart(t, parts, "read_file") + require.Equal(t, "interrupt-tool-1", call.ToolCallID) + require.Empty(t, call.Args) + require.Nil(t, call.CreatedAt, "incomplete streamed call should not have a durable call timestamp") + result := requireToolResultPart(t, parts, "read_file") + require.Equal(t, "interrupt-tool-1", result.ToolCallID) + require.True(t, result.IsError) + require.JSONEq(t, `{"error":"tool call was interrupted before it produced a result"}`, string(result.Result)) + require.NotNil(t, result.CreatedAt) + }) + + t.Run("tool execution cancellation commits interrupted result", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var requestCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + if requestCount.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/slow.txt"}`) + chunk.Choices[0].ToolCalls[0].ID = "tc-slow" + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("calling tool")[0], + chunk, + ) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("after interrupt")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + toolStarted := make(chan struct{}) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/slow.txt", int64(1), int64(0), gomock.Any()). + DoAndReturn(func(ctx context.Context, _ string, _, _ int64, _ workspacesdk.ReadFileLinesLimits) (workspacesdk.ReadFileLinesResponse, error) { + close(toolStarted) + <-ctx.Done() + return workspacesdk.ReadFileLinesResponse{}, ctx.Err() + }).Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "interrupt-tool-execution", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("run the slow tool"), + }, + }) + require.NoError(t, err) + + testutil.TryReceive(ctx, t, toolStarted) + queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue after interrupt")}, + BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, queued.Queued) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.GreaterOrEqual(t, requestCount.Load(), int32(2)) + + messages := chatMessages(ctx, t, db, chat.ID) + var foundText bool + for _, msg := range messages { + if msg.Role != database.ChatMessageRoleAssistant { + continue + } + parts, parseErr := chatprompt.ParseContent(msg) + require.NoError(t, parseErr) + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText && strings.Contains(part.Text, "calling tool") { + foundText = true + } + } + } + require.True(t, foundText) + + parts := chatToolParts(ctx, t, db, chat.ID) + call := requireToolCallPart(t, parts, "read_file") + require.Equal(t, "tc-slow", call.ToolCallID) + require.NotNil(t, call.CreatedAt) + result := requireToolResultPart(t, parts, "read_file") + require.Equal(t, "tc-slow", result.ToolCallID) + require.True(t, result.IsError) + require.JSONEq(t, `{"error":"tool call was interrupted before it produced a result"}`, string(result.Result)) + require.NotNil(t, result.CreatedAt) + require.False(t, result.CreatedAt.Before(*call.CreatedAt)) + }) + + t.Run("anthropic provider-only interruption commits no synthetic result", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + webSearchEnabled := true + providerToolStarted := make(chan struct{}) + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + + if requestCount.Add(1) != 1 { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("after interrupt")...) + } + chunks := make(chan chattest.AnthropicChunk, 2) + go func() { + defer close(chunks) + chunks <- chattest.AnthropicChunk{ + Type: "message_start", + Message: chattest.AnthropicChunkMessage{ + ID: "msg-provider-interrupt", + Type: "message", + Role: "assistant", + Model: "claude-3-opus-20240229", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 0, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "server_tool_use", + ID: "ws-interrupt", + Name: "web_search", + Input: json.RawMessage(`{"query":"coder"}`), + }, + } + select { + case <-providerToolStarted: + default: + close(providerToolStarted) + } + <-req.Context().Done() + }() + return chattest.AnthropicResponse{StreamingChunks: chunks} + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{WebSearchEnabled: &webSearchEnabled}, + }, + }) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") + testutil.TryReceive(ctx, t, providerToolStarted) + queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue after provider interrupt")}, + BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, queued.Queued) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + parts := chatToolParts(ctx, t, db, chat.ID) + require.False(t, toolResultPartExists(parts, "web_search"), + "provider-executed web_search should not get a synthetic local result") + }) + + t.Run("anthropic mixed provider and local interruption keeps local synthetic result", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + webSearchEnabled := true + streamStarted := make(chan struct{}) + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + + if requestCount.Add(1) != 1 { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("after interrupt")...) + } + chunks := make(chan chattest.AnthropicChunk, 3) + go func() { + defer close(chunks) + chunks <- chattest.AnthropicChunk{ + Type: "message_start", + Message: chattest.AnthropicChunkMessage{ + ID: "msg-mixed-interrupt", + Type: "message", + Role: "assistant", + Model: "claude-3-opus-20240229", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 0, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "server_tool_use", + ID: "ws-interrupt", + Name: "web_search", + Input: json.RawMessage(`{"query":"coder"}`), + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 1, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "tool_use", + ID: "tc-local", + Name: "read_file", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_delta", + Index: 1, + Delta: chattest.AnthropicDeltaBlock{Type: "input_json_delta", PartialJSON: `{"path":"main.go"}`}, + } + select { + case <-streamStarted: + default: + close(streamStarted) + } + <-req.Context().Done() + }() + return chattest.AnthropicResponse{StreamingChunks: chunks} + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{WebSearchEnabled: &webSearchEnabled}, + }, + }) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(0) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "anthropic-mixed-interrupt", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search and read"), + }, + }) + require.NoError(t, err) + testutil.TryReceive(ctx, t, streamStarted) + queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue after mixed interrupt")}, + BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, queued.Queued) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + parts := chatToolParts(ctx, t, db, chat.ID) + require.False(t, toolResultPartExists(parts, "web_search")) + call := requireToolCallPart(t, parts, "read_file") + require.Equal(t, "tc-local", call.ToolCallID) + require.False(t, call.ProviderExecuted) + result := requireToolResultPart(t, parts, "read_file") + require.Equal(t, "tc-local", result.ToolCallID) + require.False(t, result.ProviderExecuted) + require.True(t, result.IsError) + require.JSONEq(t, `{"error":"tool call was interrupted before it produced a result"}`, string(result.Result)) + }) + + t.Run("interrupted reasoning persists timestamps", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + sendReasoning := true + thinkingBudget := int64(1024) + reasoningStarted := make(chan struct{}) + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + + if requestCount.Add(1) != 1 { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("after interrupt")...) + } + chunks := make(chan chattest.AnthropicChunk, 3) + go func() { + defer close(chunks) + chunks <- chattest.AnthropicChunk{ + Type: "message_start", + Message: chattest.AnthropicChunkMessage{ + ID: "msg-reasoning-interrupt", + Type: "message", + Role: "assistant", + Model: "claude-3-opus-20240229", + }, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 0, + ContentBlock: chattest.AnthropicContentBlock{Type: "thinking"}, + } + chunks <- chattest.AnthropicChunk{ + Type: "content_block_delta", + Index: 0, + Delta: chattest.AnthropicDeltaBlock{Type: "thinking_delta", Thinking: "interrupted thought"}, + } + select { + case <-reasoningStarted: + default: + close(reasoningStarted) + } + <-req.Context().Done() + }() + return chattest.AnthropicResponse{StreamingChunks: chunks} + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{ + SendReasoning: &sendReasoning, + Thinking: &codersdk.ChatModelAnthropicThinkingOptions{BudgetTokens: &thinkingBudget}, + }, + }, + }) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "think") + testutil.TryReceive(ctx, t, reasoningStarted) + queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue after reasoning")}, + BusyBehavior: chatd.SendMessageBusyBehaviorInterrupt, + }) + require.NoError(t, err) + require.True(t, queued.Queued) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + messages := chatMessages(ctx, t, db, chat.ID) + var reasoningParts []codersdk.ChatMessagePart + for _, msg := range messages { + if msg.Role != database.ChatMessageRoleAssistant { + continue + } + reasoningParts = append(reasoningParts, reasoningPartsFromMessage(t, msg)...) + } + require.Len(t, reasoningParts, 1) + require.Equal(t, "interrupted thought", strings.TrimSpace(reasoningParts[0].Text)) + require.NotNil(t, reasoningParts[0].CreatedAt) + require.NotNil(t, reasoningParts[0].CompletedAt) + require.False(t, reasoningParts[0].CreatedAt.IsZero()) + require.False(t, reasoningParts[0].CompletedAt.IsZero()) + require.False(t, reasoningParts[0].CompletedAt.Before(*reasoningParts[0].CreatedAt)) + }) +} + +func TestActiveServer_DynamicToolsAndStopAfterToolBehavior(t *testing.T) { + t.Parallel() + + t.Run("dynamic tool enters requires action", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + streamedCallCount.Add(1) + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"query":"test"}`), + ) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + dynamicToolsJSON := dynamicToolJSON(t, "my_dynamic_tool") + + server := newActiveTestServer(t, db, ps) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "dynamic-tool-requires-action", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolsJSON, + }) + require.NoError(t, err) + + var chatResult database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + got, getErr := db.GetChatByID(ctx, chat.ID) + if getErr != nil { + return false + } + chatResult = got + return got.Status == database.ChatStatusRequiresAction || got.Status == database.ChatStatusError + }, testutil.IntervalFast) + require.Equal(t, database.ChatStatusRequiresAction, chatResult.Status, + "expected requires_action, got %s (last_error=%q)", + chatResult.Status, chatLastErrorMessage(chatResult.LastError)) + require.True(t, chatResult.RequiresActionDeadlineAt.Valid) + require.Equal(t, int32(1), streamedCallCount.Load()) + + parts := chatToolParts(ctx, t, db, chat.ID) + call := requireToolCallPart(t, parts, "my_dynamic_tool") + require.JSONEq(t, `{"query":"test"}`, string(call.Args)) + require.False(t, toolResultPartExists(parts, "my_dynamic_tool"), + "dynamic tool should wait for submitted results") + }) + + t.Run("successful stop after tool finishes turn", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamedCallCount.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("propose_plan", `{}`), + ) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("should not continue")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + server := newWorkspaceToolTestServer(t, db, ps, dbAgent.ID, "# Plan\n") + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "stop-after-success", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("propose a plan"), + }, + }) + require.NoError(t, err) + chatResult := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.False(t, chatResult.WorkerID.Valid) + require.False(t, chatResult.RunnerID.Valid) + require.Equal(t, int32(1), streamedCallCount.Load(), + "stop after tool should finish without another assistant call") + + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), "propose_plan") + require.False(t, result.IsError, + "stop after tool should be based on a successful tool result") + }) + + t.Run("error stop after tool continues generation", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamedCallCount.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("propose_plan", `{"path":"/tmp/not-plan.txt"}`), + ) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("tool failed, continue")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + server := newWorkspaceToolTestServer(t, db, ps, dbAgent.ID, "# Plan\n") + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "stop-after-error", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("propose a plan with a bad path"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(2), streamedCallCount.Load(), + "error stop after tool result should not finish the turn by itself") + + parts := chatToolParts(ctx, t, db, chat.ID) + result := requireToolResultPart(t, parts, "propose_plan") + require.True(t, result.IsError) + messages := chatMessages(ctx, t, db, chat.ID) + requireTextPart(t, messages[len(messages)-1], "tool failed, continue") + }) +} + func TestDynamicToolCallPausesAndResumes(t *testing.T) { t.Parallel() @@ -4353,6 +4098,7 @@ func TestDynamicToolCallPausesAndResumes(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "dynamic-tool-pause-resume", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -4520,6 +4266,7 @@ func TestDynamicToolNamedProposePlanRemainsAvailableOutsidePlanMode(t *testing.T chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "dynamic-propose-plan-collision", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -4634,6 +4381,7 @@ func TestDynamicToolCallMixedWithBuiltIn(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "mixed-builtin-dynamic", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -4773,6 +4521,7 @@ func TestSubmitToolResultsConcurrency(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "concurrency-tool-results", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -4889,12 +4638,11 @@ func ptrRef[T any](v T) *T { return &v } -func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { +func TestSubscribeNoDuplicateMessageParts(t *testing.T) { t.Parallel() - // Use nil pubsub to force the no-pubsub path. - db, _ := dbtestutil.NewDB(t) - replica := newStartedTestServer(t, db, nil, uuid.New()) + db, ps := dbtestutil.NewDB(t) + replica := newTestServer(t, db, ps, uuid.New()) ctx := testutil.Context(t, testutil.WaitLong) user, org, model := seedChatDependencies(t, db) @@ -4902,20 +4650,13 @@ func TestSubscribeNoPubsubNoDuplicateMessageParts(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "no-dup-parts", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, }) require.NoError(t, err) - // Wait for any wake-triggered processing to settle before - // subscribing, so the snapshot captures the final state. - // The wake signal may trigger processOnce which will fail - // (no LLM configured) and set the chat to error status. - // Poll until the chat reaches a terminal state (not pending - // and not running), then wait for the goroutine to finish. - waitForChatProcessed(ctx, t, db, chat.ID, replica) - snapshot, events, cancel, ok := replica.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) t.Cleanup(cancel) @@ -5378,6 +5119,7 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) chat, err := inactive.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "stopped-workspace-regression", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -5387,19 +5129,10 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) }) require.NoError(t, err) - // Close the inactive server so its wake-triggered processing - // stops and releases the chat. Then reset to pending so the - // active server (created below) can acquire it cleanly. + // Close the inactive server. The chat remains in the valid + // state-machine `running` state created by CreateChat, and the + // active server created below can acquire it because it is unowned. require.NoError(t, inactive.Close()) - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusPending, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - LastError: pqtype.NullRawMessage{}, - }) - require.NoError(t, err) build, err := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) require.NoError(t, err) @@ -5506,181 +5239,6 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) require.Contains(t, string(parts[0].Result), "workspace has no running agent") } -func TestHeartbeatBumpsWorkspaceUsage(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("ok") - } - // Block until the request context is canceled so the chat - // stays in a processing state long enough for heartbeats - // to fire. - chunks := make(chan chattest.OpenAIChunk) - go func() { - defer close(chunks) - <-req.Context().Done() - }() - return chattest.OpenAIResponse{StreamingChunks: chunks} - })) - - // Create a workspace with a full build chain so we can verify - // both last_used_at (dormancy) and deadline (autostop) bumps. - tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - tmpl := dbgen.Template(t, db, database.Template{ - OrganizationID: org.ID, - ActiveVersionID: tv.ID, - CreatedBy: user.ID, - }) - require.NoError(t, db.UpdateTemplateScheduleByID(ctx, database.UpdateTemplateScheduleByIDParams{ - ID: tmpl.ID, - UpdatedAt: dbtime.Now(), - AllowUserAutostop: true, - ActivityBump: int64(time.Hour), - })) - ws := dbgen.Workspace(t, db, database.WorkspaceTable{ - OwnerID: user.ID, - OrganizationID: org.ID, - TemplateID: tmpl.ID, - Ttl: sql.NullInt64{Valid: true, Int64: int64(8 * time.Hour)}, - }) - pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ - OrganizationID: org.ID, - CompletedAt: sql.NullTime{ - Valid: true, - Time: dbtime.Now().Add(-30 * time.Minute), - }, - }) - // Build deadline is 30 minutes in the past, close enough to - // be bumped by the default 1-hour activity bump. - build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ - WorkspaceID: ws.ID, - TemplateVersionID: tv.ID, - JobID: pj.ID, - Transition: database.WorkspaceTransitionStart, - Deadline: dbtime.Now().Add(-30 * time.Minute), - }) - originalDeadline := build.Deadline - - // Set up a short heartbeat interval and a UsageTracker that - // flushes frequently so last_used_at gets updated in the DB. - flushTick := make(chan time.Time) - flushDone := make(chan int, 1) - tracker := workspacestats.NewTracker(db, - workspacestats.TrackerWithTickFlush(flushTick, flushDone), - workspacestats.TrackerWithLogger(slogtest.Make(t, nil)), - ) - t.Cleanup(func() { tracker.Close() }) - - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - // Wrap the database with dbauthz so the chatd server's - // AsChatd context is enforced on every query, matching - // production behavior. - authzDB := dbauthz.New(db, rbac.NewStrictCachingAuthorizer(prometheus.NewRegistry()), slogtest.Make(t, nil), coderdtest.AccessControlStorePointer()) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: authzDB, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: 10 * time.Millisecond, - InFlightChatStaleAfter: testutil.WaitLong, - ChatHeartbeatInterval: 100 * time.Millisecond, - UsageTracker: tracker, - }) - server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) - }) - - // Create a chat WITHOUT a workspace, the normal starting state. - // In production, CreateChat is called from the HTTP handler with - // the authenticated user's context. Here we use AsChatd since - // the chatd server processes everything under that role. - chatCtx := dbauthz.AsChatd(ctx) - chat, err := server.CreateChat(chatCtx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "usage-tracking-test", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // Wait for the chat to start processing and at least one - // heartbeat to fire. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - fromDB, listErr := db.GetChatByID(ctx, chat.ID) - if listErr != nil { - return false - } - return fromDB.Status == database.ChatStatusRunning && - fromDB.HeartbeatAt.Valid && - fromDB.HeartbeatAt.Time.After(fromDB.CreatedAt) - }, testutil.IntervalFast, - "chat should be running with at least one heartbeat") - - // Flush the tracker and verify nothing was tracked yet - // (no workspace linked). - testutil.RequireSend(ctx, t, flushTick, time.Now()) - count := testutil.RequireReceive(ctx, t, flushDone) - require.Equal(t, 0, count, - "expected no workspaces to be flushed before association") - - // Link the workspace to the chat in the DB, simulating what - // the create_workspace tool does mid-conversation. - _, err = db.UpdateChatWorkspaceBinding(ctx, database.UpdateChatWorkspaceBindingParams{ - WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, - ID: chat.ID, - }) - require.NoError(t, err) - - // The heartbeat re-reads the workspace association from the DB - // on each tick. Wait for the tracker to pick it up. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case flushTick <- time.Now(): - case <-ctx.Done(): - return false - } - select { - case c := <-flushDone: - return c > 0 - case <-ctx.Done(): - return false - } - }, testutil.IntervalMedium, - "expected usage tracker to flush the late-associated workspace") - - // Verify the workspace's last_used_at was actually updated. - updatedWs, err := db.GetWorkspaceByID(ctx, ws.ID) - require.NoError(t, err) - require.True(t, updatedWs.LastUsedAt.After(ws.LastUsedAt), - "workspace last_used_at should have been bumped") - - // Verify the workspace build deadline was also extended. - // The SQL only writes when 5% of the deadline has elapsed, - // most calls perform a read-only CTE lookup. Wider ±2 - // minute tolerance than activitybump_test.go because the bump - // happens asynchronously via the heartbeat goroutine. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - updatedBuild, buildErr := db.GetLatestWorkspaceBuildByWorkspaceID(ctx, ws.ID) - if buildErr != nil || !updatedBuild.Deadline.After(originalDeadline) { - return false - } - now := dbtime.Now() - return updatedBuild.Deadline.After(now.Add(time.Hour-2*time.Minute)) && - updatedBuild.Deadline.Before(now.Add(time.Hour+2*time.Minute)) - }, testutil.IntervalFast, - "workspace build deadline should have been bumped to ~now+1h") -} - func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { t.Parallel() @@ -5710,11 +5268,10 @@ func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { t.Cleanup(func() { tracker.Close() }) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitLong, ChatHeartbeatInterval: 100 * time.Millisecond, @@ -5728,22 +5285,28 @@ func TestHeartbeatNoWorkspaceNoBump(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "no-workspace-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, }) require.NoError(t, err) - // Wait for the chat to be acquired and at least one heartbeat - // to fire. + // Wait for the chat to be acquired and at least one runner + // heartbeat to be written. testutil.Eventually(ctx, t, func(ctx context.Context) bool { fromDB, listErr := db.GetChatByID(ctx, chat.ID) - if listErr != nil { + if listErr != nil || fromDB.Status != database.ChatStatusRunning || !fromDB.RunnerID.Valid { return false } - return fromDB.Status == database.ChatStatusRunning && - fromDB.HeartbeatAt.Valid && - fromDB.HeartbeatAt.Time.After(fromDB.CreatedAt) + heartbeat, heartbeatErr := db.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: chat.ID, + RunnerID: fromDB.RunnerID.UUID, + }) + if heartbeatErr != nil { + return false + } + return heartbeat.HeartbeatAt.After(fromDB.CreatedAt) }, testutil.IntervalFast, "chat should be running with at least one heartbeat") @@ -5798,11 +5361,10 @@ func newTestServer( t.Helper() logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: replicaID, - Pubsub: ps, PendingChatAcquireInterval: testutil.WaitLong, }) t.Cleanup(func() { @@ -5811,6 +5373,2786 @@ func newTestServer( return server } +func highUsageTextResponse(text string) chattest.AnthropicResponse { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 80, + OutputTokens: 5, + }, text)...) +} + +func anthropicCompactionResponse(text string) chattest.AnthropicResponse { + return chattest.AnthropicResponse{Response: &chattest.AnthropicMessage{ + ID: "msg-compaction", + Type: "message", + Role: "assistant", + Content: text, + Model: "claude-3-opus-20240229", + StopReason: "end_turn", + }} +} + +func highUsageReadFileResponse(path string) chattest.AnthropicResponse { + chunks := chattest.AnthropicToolCallChunks("read_file", fmt.Sprintf(`{"path":%q}`, path)) + for i := range chunks { + if chunks[i].Type == "message_start" { + chunks[i].Message.Usage = map[string]int{"input_tokens": 80} + } + if chunks[i].Type == "message_delta" { + chunks[i].UsageMap = map[string]int{"output_tokens": 5} + } + } + return chattest.AnthropicStreamingResponse(chunks...) +} + +func TestActiveServer_AIGatewayRoutingPreservesAPIKeyAfterCompaction(t *testing.T) { + t.Parallel() + + const ( + compactionSummary = "summary text for AI Gateway compaction" + contextLimit = int64(100) + thresholdPercent = int32(70) + ) + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var streamCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + body := anthropicRequestBody(t, *req) + if !req.Stream { + if strings.Contains(body, "You are performing a context compaction") { + return chattest.AnthropicNonStreamingResponse(compactionSummary) + } + return chattest.AnthropicNonStreamingResponse("AI Gateway Compaction") + } + + switch streamCount.Add(1) { + case 1: + return highUsageReadFileResponse("/tmp/a.txt") + case 2: + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 20, + OutputTokens: 5, + }, "continued after compaction")...) + default: + t.Fatalf("unexpected streamed model call %d", streamCount.Load()) + return chattest.AnthropicStreamingResponse() + } + }) + factory := newChatAIGatewayPreservePathTestFactory(t, anthropicURL) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCompressionThreshold(t, db, model, contextLimit, thresholdPercent) + provider, err := db.GetAIProviderByID(ctx, model.AIProviderID.UUID) + require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + _, err = db.UpsertUserAIProviderKey(ctx, database.UpsertUserAIProviderKeyParams{ + ID: uuid.New(), + UserID: user.ID, + AIProviderID: provider.ID, + APIKey: "sk-user-aibridge", + }) + require.NoError(t, err) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1\tpackage main"}, nil). + Times(1) + + creator := newTestServer(t, db, ps, uuid.New()) + chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "aigateway-compaction", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + APIKeyID: apiKey.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("trigger compaction"), + }, + }) + require.NoError(t, err) + contextContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFileAgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + ContextFilePath: "/home/coder/project/AGENTS.md", + ContextFileContent: "# Project instructions", + ContextFileOS: "linux", + ContextFileDirectory: "/home/coder/project", + }}) + require.NoError(t, err) + _, err = db.InsertChatMessages(ctx, chatd.BuildSingleUserChatMessageInsertParams( + chat.ID, + apiKey.ID, + contextContent, + database.ChatMessageVisibilityBoth, + model.ID, + chatprompt.CurrentContentVersion, + user.ID, + )) + require.NoError(t, err) + + _ = newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) + cfg.AIGatewayRoutingEnabled = true + cfg.AllowBYOK = true + cfg.AllowBYOKSet = true + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + + chatResult := waitForTerminalChat(ctx, t, db, chat.ID) + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + require.False(t, chatResult.LastError.Valid) + + messages := chatMessages(ctx, t, db, chat.ID) + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + compressed := compressedChatSummarizedMessages(t, append(promptMessages, messages...)) + require.Len(t, compressed.summaries, 1) + require.True(t, compressed.summaries[0].APIKeyID.Valid) + require.Equal(t, apiKey.ID, compressed.summaries[0].APIKeyID.String) + + requests := factory.requestsSnapshot() + require.NotEmpty(t, requests) + for _, req := range requests { + require.Equal(t, provider.Name, req.ProviderName) + require.Equal(t, aibridge.SourceAgents, req.Source) + require.Equal(t, apiKey.ID, req.APIKeyID) + require.Equal(t, "sk-user-aibridge", req.XAPIKey) + require.Equal(t, "delegated", req.CoderToken) + } +} + +func TestActiveServer_CompactionRecordsMetric(t *testing.T) { + t.Parallel() + + const ( + compactionSummary = "summary text for compaction" + contextLimit = int64(100) + thresholdPercent = int32(70) + ) + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + reg := prometheus.NewRegistry() + var streamCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + body := anthropicRequestBody(t, *req) + if !req.Stream { + if strings.Contains(body, "You are performing a context compaction") { + return anthropicCompactionResponse(compactionSummary) + } + return chattest.AnthropicNonStreamingResponse("title") + } + switch streamCount.Add(1) { + case 1: + return highUsageReadFileResponse("/tmp/a.txt") + case 2: + require.Contains(t, body, compactionSummary) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 20, + OutputTokens: 5, + }, "continued after compaction")...) + default: + t.Fatalf("unexpected generation request: %s", body) + return chattest.AnthropicStreamingResponse() + } + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCompressionThreshold(t, db, model, contextLimit, thresholdPercent) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1\tpackage main"}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.PrometheusRegistry = reg + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "compaction-metric", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file and continue"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requireChatdMetricCounter(t, reg, "coderd_chatd_compaction_total", 1, map[string]string{ + "provider": "anthropic", + "model": "claude-sonnet-4-20250514", + "result": "success", + }) +} + +func TestActiveServer_Compaction(t *testing.T) { + t.Parallel() + + const ( + compactionSummary = "summary text for compaction" + contextLimit = int64(100) + thresholdPercent = int32(70) + ) + + newHighUsageReadFileResponse := func(path string) chattest.AnthropicResponse { + chunks := chattest.AnthropicToolCallChunks("read_file", fmt.Sprintf(`{"path":%q}`, path)) + for i := range chunks { + if chunks[i].Type == "message_start" { + chunks[i].Message.Usage = map[string]int{"input_tokens": 80} + } + if chunks[i].Type == "message_delta" { + chunks[i].UsageMap = map[string]int{"output_tokens": 5} + } + } + return chattest.AnthropicStreamingResponse(chunks...) + } + + t.Run("commits summary when threshold reached and continues from committed summary", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var streamCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + body := anthropicRequestBody(t, *req) + if !req.Stream { + if strings.Contains(body, "You are performing a context compaction") { + require.Contains(t, body, "read_file") + require.Contains(t, body, "package main") + return anthropicCompactionResponse(compactionSummary) + } + return chattest.AnthropicNonStreamingResponse("title") + } + switch streamCount.Add(1) { + case 1: + return newHighUsageReadFileResponse("/tmp/a.txt") + default: + require.Contains(t, body, compactionSummary) + require.Contains(t, body, "The following is a summary of the earlier conversation") + require.Contains(t, body, `"role":"user"`) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 20, + OutputTokens: 5, + }, "continued after compaction")...) + } + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCompressionThreshold(t, db, model, contextLimit, thresholdPercent) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1 package main"}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "compaction-continues", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file and continue"), + }, + }) + require.NoError(t, err) + chat = waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.False(t, chat.WorkerID.Valid) + require.False(t, chat.RunnerID.Valid) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.GreaterOrEqual(t, len(generationRequests), 2) + require.Equal(t, int32(2), streamCount.Load()) + + messages := chatMessages(ctx, t, db, chat.ID) + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + compressed := compressedChatSummarizedMessages(t, append(promptMessages, messages...)) + require.Len(t, compressed.summaries, 1) + require.Len(t, compressed.calls, 1) + require.Len(t, compressed.results, 1) + + require.Equal(t, database.ChatMessageRoleUser, compressed.summaries[0].Role) + require.Equal(t, database.ChatMessageVisibilityModel, compressed.summaries[0].Visibility) + summaryText := messageText(t, compressed.summaries[0]) + require.Contains(t, summaryText, "The following is a summary of the earlier conversation") + require.Contains(t, summaryText, compactionSummary) + + callPart := singlePartOfType(t, compressed.calls[0], codersdk.ChatMessagePartTypeToolCall) + resultPart := singlePartOfType(t, compressed.results[0], codersdk.ChatMessagePartTypeToolResult) + require.Equal(t, callPart.ToolCallID, resultPart.ToolCallID) + require.Equal(t, "chat_summarized", resultPart.ToolName) + require.JSONEq(t, `{"summary":"summary text for compaction","source":"automatic","threshold_percent":70,"usage_percent":80,"context_tokens":80,"context_limit_tokens":100}`, string(resultPart.Result)) + requireTextPart(t, messages[len(messages)-1], "continued after compaction") + }) + + t.Run("does not compact when high usage finishes the turn", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamCount atomic.Int32 + var compactionRequests atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + body := anthropicRequestBody(t, *req) + if strings.Contains(body, "You are performing a context compaction") { + compactionRequests.Add(1) + return anthropicCompactionResponse(compactionSummary) + } + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + streamCount.Add(1) + return highUsageTextResponse("done without compaction") + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCompressionThreshold(t, db, model, contextLimit, thresholdPercent) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "finish with high usage") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + require.Equal(t, int32(1), streamCount.Load()) + require.Equal(t, int32(0), compactionRequests.Load()) + messages := chatMessages(ctx, t, db, chat.ID) + compressed := compressedChatSummarizedMessages(t, messages) + require.Empty(t, compressed.summaries) + require.Empty(t, compressed.calls) + require.Empty(t, compressed.results) + for _, msg := range messages { + require.False(t, msg.Compressed, "message %d should not be compressed", msg.ID) + } + requireTextPart(t, messages[len(messages)-1], "done without compaction") + }) + + t.Run("fails when compaction leaves chat over limit", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + body := anthropicRequestBody(t, *req) + if !req.Stream { + if strings.Contains(body, "You are performing a context compaction") { + return anthropicCompactionResponse(compactionSummary) + } + return chattest.AnthropicNonStreamingResponse("title") + } + switch streamCount.Add(1) { + case 1: + return newHighUsageReadFileResponse("/tmp/a.txt") + default: + require.Contains(t, body, compactionSummary) + return highUsageTextResponse("still too large") + } + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCompressionThreshold(t, db, model, contextLimit, thresholdPercent) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1 package main"}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "compaction-still-over-limit", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file and stay too large"), + }, + }) + require.NoError(t, err) + chat = waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + require.Contains(t, chatLastErrorMessage(chat.LastError), "The chat request failed unexpectedly.") + }) +} + +type compressedCompactionMessages struct { + summaries []database.ChatMessage + calls []database.ChatMessage + results []database.ChatMessage +} + +func compressedChatSummarizedMessages(t *testing.T, messages []database.ChatMessage) compressedCompactionMessages { + t.Helper() + seen := map[int64]bool{} + var out compressedCompactionMessages + for _, msg := range messages { + if !msg.Compressed || seen[msg.ID] { + continue + } + seen[msg.ID] = true + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + for _, part := range parts { + switch part.Type { + case codersdk.ChatMessagePartTypeText: + if msg.Role == database.ChatMessageRoleUser { + out.summaries = append(out.summaries, msg) + } + case codersdk.ChatMessagePartTypeToolCall: + if part.ToolName == "chat_summarized" { + out.calls = append(out.calls, msg) + } + case codersdk.ChatMessagePartTypeToolResult: + if part.ToolName == "chat_summarized" { + out.results = append(out.results, msg) + } + } + } + } + return out +} + +func messageText(t *testing.T, msg database.ChatMessage) string { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + var builder strings.Builder + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText { + _, _ = builder.WriteString(part.Text) + } + } + return builder.String() +} + +func singlePartOfType(t *testing.T, msg database.ChatMessage, typ codersdk.ChatMessagePartType) codersdk.ChatMessagePart { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + var matches []codersdk.ChatMessagePart + for _, part := range parts { + if part.Type == typ { + matches = append(matches, part) + } + } + require.Len(t, matches, 1) + return matches[0] +} + +func TestActiveServer_BasicAssistantGenerationAndPromptPreparation(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model.ContextLimit = 4096 + model = updateChatModelContextLimit(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertSystemTextMessage(ctx, t, db, chat.ID, "sys-2", model.ID) + insertAssistantTextMessage(ctx, t, db, chat.ID, "working", model.ID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + recovered := generationRequests[1] + require.True(t, anthropicSystemHasEphemeralCacheControl(t, recovered)) + require.Len(t, recovered.Messages, 4) + require.False(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[0])) + require.False(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[1])) + require.True(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[2])) + require.True(t, anthropicMessageHasEphemeralCacheControl(t, recovered.Messages[3])) + require.False(t, anthropicRequestContainsPromptSentinel(t, recovered)) + toolNames := anthropicRequestToolNames(recovered) + require.Contains(t, toolNames, "read_file") + require.Contains(t, toolNames, "write_file") + + messages := chatMessages(ctx, t, db, chat.ID) + last := messages[len(messages)-1] + require.Equal(t, database.ChatMessageRoleAssistant, last.Role) + require.True(t, last.ContextLimit.Valid) + require.Equal(t, int64(4096), last.ContextLimit.Int64) + require.GreaterOrEqual(t, last.RuntimeMs.Int64, int64(0)) + requireTextPart(t, last, "done") + + requests = newAnthropicRequestRecorder() + server = newActiveTestServer(t, db, ps) + planChat := createPlanSubagentChatWithHistory(ctx, t, db, org.ID, user.ID, model.ID) + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: planChat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, planChat.ID, database.ChatStatusWaiting) + + planRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, planRequests, 1) + toolNames = anthropicRequestToolNames(planRequests[0]) + require.Contains(t, toolNames, "read_file") + require.NotContains(t, toolNames, "write_file") +} + +func TestActiveServer_ToolExecutionAndPolicy(t *testing.T) { + t.Parallel() + + t.Run("rejects disallowed active tool", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("write_file", `{"path":"/tmp/nope","content":"blocked"}`), + ) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().WriteFile(gomock.Any(), gomock.Any(), gomock.Any()).Times(0) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "active-tool-reject", + ModelConfigID: model.ID, + ChatMode: database.NullChatMode{ChatMode: database.ChatModeExplore, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("try to write a file"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + result := requireToolResultPart(t, parts, "write_file") + require.True(t, result.IsError) + require.JSONEq(t, `{"error":"Tool not active in this turn: write_file"}`, string(result.Result)) + }) + + t.Run("provider runner executes and preserves metadata", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + const computerResultMetadata = `{"openai":{"type":"openai.responses.computer_call_output_options","data":{"detail":"original"}}}` + var streamedCallCount atomic.Int32 + var secondRawBody []byte + var callsMu sync.Mutex + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + callsMu.Lock() + secondRawBody = append([]byte(nil), req.RawBody...) + callsMu.Unlock() + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, _, model := seedChatDependenciesWithProviderPolicy(t, db, "openai", openAIURL, "test-key", true, false, true) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + model.Model = "gpt-5.5" + model = updateChatModelContextLimit(t, db, model) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { cfg.AllowBYOKSet = true; cfg.AllowBYOK = false }) + result := codersdk.ChatMessageToolResult( + "computer-call", + "computer", + json.RawMessage(`{"data":"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4n539HwAHFwLVF8kc1wAAAABJRU5ErkJggg==","mime_type":"image/png"}`), + false, + true, + ) + result.ProviderMetadata = json.RawMessage(computerResultMetadata) + computerCall := codersdk.ChatMessageToolCall( + "computer-call", + "computer", + json.RawMessage(`{"type":"screenshot"}`), + ) + computerCall.ProviderExecuted = true + created, err := chatstate.CreateChat(dbauthz.AsSystemRestricted(ctx), db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "provider-runner-replay-active", + MCPServerIDs: []uuid.UUID{}, + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + userMessageForTest(t, "use provider runner", model.ID, user.ID, apiKey.ID), + assistantMessageForTest(t, []codersdk.ChatMessagePart{computerCall}, model.ID), + toolMessageForTest(t, []codersdk.ChatMessagePart{result}, model.ID), + }, + }) + require.NoError(t, err) + chat := created.Chat + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + waitForTerminalChat(ctx, t, db, chat.ID) + gotChat, gotErr := db.GetChatByID(ctx, chat.ID) + require.NoError(t, gotErr) + require.Equal(t, database.ChatStatusWaiting, gotChat.Status) + require.Eventually(t, func() bool { return streamedCallCount.Load() >= 1 }, testutil.WaitShort, testutil.IntervalFast) + + callsMu.Lock() + body := string(secondRawBody) + callsMu.Unlock() + require.Contains(t, body, "computer_call_output") + require.Contains(t, body, `"detail":"original"`) + }) + + t.Run("multi step local tool execution", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + var secondCallMessages []chattest.OpenAIMessage + var callsMu sync.Mutex + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/a.txt"}`), + ) + } + callsMu.Lock() + secondCallMessages = append([]chattest.OpenAIMessage(nil), req.Messages...) + callsMu.Unlock() + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("all done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1\tpackage main", + }, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "multi-step-tool", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2)) + parts := chatToolParts(ctx, t, db, chat.ID) + call := requireToolCallPart(t, parts, "read_file") + result := requireToolResultPart(t, parts, "read_file") + require.False(t, result.IsError) + require.NotNil(t, call.CreatedAt) + require.NotNil(t, result.CreatedAt) + require.False(t, result.CreatedAt.Before(*call.CreatedAt)) + messages := chatMessages(ctx, t, db, chat.ID) + requireTextPart(t, messages[len(messages)-1], "all done") + + callsMu.Lock() + secondMessages := append([]chattest.OpenAIMessage(nil), secondCallMessages...) + callsMu.Unlock() + require.NotEmpty(t, secondMessages) + require.True(t, openAIMessagesContain(secondMessages, "1\\tpackage main")) + }) + + t.Run("parallel local and provider executed timestamps", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + webSearchEnabled := true + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + readA := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/a.txt"}`) + readB := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/b.txt"}`) + second := readB.Choices[0].ToolCalls[0] + second.Index = 1 + readA.Choices[0].ToolCalls = append(readA.Choices[0].ToolCalls, second) + return chattest.OpenAIResponse{ + StreamingChunks: chattest.OpenAIStreamingResponse(readA).StreamingChunks, + WebSearch: &chattest.OpenAIWebSearchCall{ID: "ws-timestamps", Query: "coder"}, + } + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + OpenAI: &codersdk.ChatModelOpenAIProviderOptions{WebSearchEnabled: &webSearchEnabled}, + }, + }) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, Content: "a", FileSize: 1, TotalLines: 1, LinesRead: 1}, nil). + Times(1) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/b.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, Content: "b", FileSize: 1, TotalLines: 1, LinesRead: 1}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "parallel-timestamps", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search and read files"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + for _, toolName := range []string{"read_file", "web_search"} { + call := requireToolCallPart(t, parts, toolName) + result := requireToolResultPart(t, parts, toolName) + require.NotNil(t, call.CreatedAt, toolName) + require.NotNil(t, result.CreatedAt, toolName) + require.False(t, result.CreatedAt.Before(*call.CreatedAt), toolName) + if toolName == "web_search" { + require.True(t, call.ProviderExecuted) + require.True(t, result.ProviderExecuted) + } else { + require.False(t, call.ProviderExecuted) + require.False(t, result.ProviderExecuted) + } + } + }) +} + +func TestActiveServer_RecordsGenerationMetrics(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + reg := prometheus.NewRegistry() + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(openAITextChunksWithStop("hello")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.PrometheusRegistry = reg + }) + + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requireChatdMetricCounter(t, reg, "coderd_chatd_steps_total", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o-mini", + }) + requireChatdMetricHistogram(t, reg, "coderd_chatd_message_count", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o-mini", + }, chatdMetricHistogramRequirement{}) + requireChatdMetricHistogram(t, reg, "coderd_chatd_prompt_size_bytes", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o-mini", + }, chatdMetricHistogramRequirement{PositiveSum: true}) + requireChatdMetricHistogram(t, reg, "coderd_chatd_ttft_seconds", 1, map[string]string{ + "provider": "openai", + "model": "gpt-4o-mini", + }, chatdMetricHistogramRequirement{}) +} + +func TestActiveServer_ToolErrorRecordsMetric(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + toolName string + toolArgs string + chatMode database.NullChatMode + setupAgent func(*agentconnmock.MockAgentConn) + }{ + { + name: "builtin tool IsError", + toolName: "read_file", + toolArgs: `{"path":"/tmp/missing.txt"}`, + setupAgent: func(mockConn *agentconnmock.MockAgentConn) { + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/missing.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: false, Error: "file not found"}, nil). + Times(1) + }, + }, + { + name: "non builtin MCP style tool IsError", + toolName: "dynamic_error_tool", + toolArgs: `{"input":"hello"}`, + setupAgent: func(mockConn *agentconnmock.MockAgentConn) { + mockConn.EXPECT().CallMCPTool(gomock.Any(), gomock.Any()). + Return(workspacesdk.CallMCPToolResponse{ + IsError: true, + Content: []workspacesdk.MCPToolContent{{ + Type: "text", + Text: "dynamic failed", + }}, + }, nil). + Times(1) + }, + }, + { + name: "tool Run returns error", + toolName: "read_file", + toolArgs: `{"path":"/tmp/error.txt"}`, + setupAgent: func(mockConn *agentconnmock.MockAgentConn) { + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/error.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{}, xerrors.New("connection refused")). + Times(1) + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + reg := prometheus.NewRegistry() + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk(tt.toolName, tt.toolArgs), + ) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + model.Model = "test-model" + model = updateChatModelContextLimit(t, db, model) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn, workspacesdk.MCPToolInfo{ + Name: "dynamic_error_tool", + Description: "dynamic error tool", + Schema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "input": map[string]any{"type": "string"}, + }, + }, + }) + tt.setupAgent(mockConn) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.PrometheusRegistry = reg + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chatOpts := chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "tool-error-metric", + ModelConfigID: model.ID, + ChatMode: tt.chatMode, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("run an erroring tool"), + }, + } + chat, err := server.CreateChat(ctx, chatOpts) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), tt.toolName) + require.True(t, result.IsError) + requireChatdMetricCounter(t, reg, "coderd_chatd_tool_errors_total", 1, map[string]string{ + "provider": "openai-compat", + "model": "test-model", + "tool_name": tt.toolName, + }) + }) + } +} + +func userMessageForTest( + t *testing.T, + text string, + modelID uuid.UUID, + createdBy uuid.UUID, + apiKeyID string, +) chatstate.Message { + t.Helper() + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + } +} + +func assistantMessageForTest( + t *testing.T, + parts []codersdk.ChatMessagePart, + modelID uuid.UUID, +) chatstate.Message { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + } +} + +func toolMessageForTest( + t *testing.T, + parts []codersdk.ChatMessagePart, + modelID uuid.UUID, +) chatstate.Message { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleTool, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + } +} + +func setupToolExecutionAgentConn( + t *testing.T, + mockConn *agentconnmock.MockAgentConn, + mcpTools ...workspacesdk.MCPToolInfo, +) { + t.Helper() + mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + mockConn.EXPECT().ContextConfig(gomock.Any()). + Return(workspacesdk.ContextConfigResponse{}, xerrors.New("not supported")).AnyTimes() + mockConn.EXPECT().ListMCPTools(gomock.Any()). + Return(workspacesdk.ListMCPToolsResponse{Tools: mcpTools}, nil).AnyTimes() + mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.LSResponse{AbsolutePathString: "/home/coder"}, nil).AnyTimes() + mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() +} + +func mustParseChatParts(t *testing.T, msg database.ChatMessage) []codersdk.ChatMessagePart { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + return parts +} + +func dynamicToolJSON(t *testing.T, name string) []byte { + t.Helper() + encoded, err := json.Marshal([]mcpgo.Tool{{ + Name: name, + Description: "A test dynamic tool.", + InputSchema: mcpgo.ToolInputSchema{ + Type: "object", + Properties: map[string]any{ + "query": map[string]any{"type": "string"}, + }, + }, + }}) + require.NoError(t, err) + return encoded +} + +func toolResultPartExists(parts []codersdk.ChatMessagePart, toolName string) bool { + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == toolName { + return true + } + } + return false +} + +func chatToolParts( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, +) []codersdk.ChatMessagePart { + t.Helper() + var parts []codersdk.ChatMessagePart + for _, msg := range chatMessages(ctx, t, db, chatID) { + parsed, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + for _, part := range parsed { + if part.Type == codersdk.ChatMessagePartTypeToolCall || + part.Type == codersdk.ChatMessagePartTypeToolResult { + parts = append(parts, part) + } + } + } + return parts +} + +func requireToolCallPart( + t *testing.T, + parts []codersdk.ChatMessagePart, + toolName string, +) codersdk.ChatMessagePart { + t.Helper() + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == toolName { + return part + } + } + t.Fatalf("missing tool-call part for %q", toolName) + return codersdk.ChatMessagePart{} +} + +func requireToolResultPart( + t *testing.T, + parts []codersdk.ChatMessagePart, + toolName string, +) codersdk.ChatMessagePart { + t.Helper() + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == toolName { + return part + } + } + t.Fatalf("missing tool-result part for %q", toolName) + return codersdk.ChatMessagePart{} +} + +func openAIMessagesContain(messages []chattest.OpenAIMessage, text string) bool { + for _, msg := range messages { + if strings.Contains(msg.Content, text) { + return true + } + } + return false +} + +func requireChatdMetricCounter( + t *testing.T, + reg *prometheus.Registry, + name string, + wantValue float64, + wantLabels map[string]string, +) { + t.Helper() + families, err := reg.Gather() + require.NoError(t, err) + for _, family := range families { + if family.GetName() != name { + continue + } + for _, metric := range family.GetMetric() { + labels := metricLabels(metric) + if !metricLabelsMatch(labels, wantLabels) { + continue + } + require.Equal(t, wantValue, metric.GetCounter().GetValue()) + return + } + t.Fatalf("metric %s with labels %v not found", name, wantLabels) + } + t.Fatalf("metric %s not found", name) +} + +type chatdMetricHistogramRequirement struct { + PositiveSum bool +} + +func requireChatdMetricHistogram( + t *testing.T, + reg *prometheus.Registry, + name string, + wantSampleCount uint64, + wantLabels map[string]string, + requirement chatdMetricHistogramRequirement, +) { + t.Helper() + families, err := reg.Gather() + require.NoError(t, err) + for _, family := range families { + if family.GetName() != name { + continue + } + for _, metric := range family.GetMetric() { + labels := metricLabels(metric) + if !metricLabelsMatch(labels, wantLabels) { + continue + } + histogram := metric.GetHistogram() + require.Equal(t, wantSampleCount, histogram.GetSampleCount()) + if requirement.PositiveSum { + require.Positive(t, histogram.GetSampleSum()) + } + return + } + t.Fatalf("metric %s with labels %v not found", name, wantLabels) + } + t.Fatalf("metric %s not found", name) +} + +func metricLabels(metric interface { + GetLabel() []*io_prometheus_client.LabelPair +}, +) map[string]string { + labels := map[string]string{} + for _, label := range metric.GetLabel() { + labels[label.GetName()] = label.GetValue() + } + return labels +} + +func metricLabelsMatch(labels, wantLabels map[string]string) bool { + for key, value := range wantLabels { + if labels[key] != value { + return false + } + } + return true +} + +func TestActiveServer_AnthropicUsageMatchesFinalDelta(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + anthropicURL := chattest.NewAnthropic(t, func(_ *chattest.AnthropicRequest) chattest.AnthropicResponse { + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 200, + OutputTokens: 75, + CacheCreationInputTokens: 30, + CacheReadInputTokens: 150, + }, "cached response")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + messages := chatMessages(ctx, t, db, chat.ID) + last := messages[len(messages)-1] + require.Equal(t, database.ChatMessageRoleAssistant, last.Role) + require.Equal(t, sql.NullInt64{Int64: 200, Valid: true}, last.InputTokens) + require.Equal(t, sql.NullInt64{Int64: 75, Valid: true}, last.OutputTokens) + require.Equal(t, sql.NullInt64{Int64: 275, Valid: true}, last.TotalTokens) + require.Equal(t, sql.NullInt64{Int64: 30, Valid: true}, last.CacheCreationTokens) + require.Equal(t, sql.NullInt64{Int64: 150, Valid: true}, last.CacheReadTokens) +} + +func TestActiveServer_ChatTurnDebugRunRecordsStreamStep(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse(`{"label":"Debug response"}`) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 200, + OutputTokens: 75, + CacheCreationInputTokens: 30, + CacheReadInputTokens: 150, + }, "debug response")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AlwaysEnableDebugLogs = true + }) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "hello debug") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.NoError(t, server.Close()) + debugCtx := testutil.Context(t, testutil.WaitLong) + + var chatTurnRuns []database.ChatDebugRun + testutil.Eventually(debugCtx, t, func(ctx context.Context) bool { + runs, err := db.GetChatDebugRunsByChatID(ctx, database.GetChatDebugRunsByChatIDParams{ + ChatID: chat.ID, + LimitVal: 100, + }) + if err != nil { + return false + } + chatTurnRuns = chatTurnRuns[:0] + for _, run := range runs { + if run.Kind == string(codersdk.ChatDebugRunKindChatTurn) { + chatTurnRuns = append(chatTurnRuns, run) + } + } + return len(chatTurnRuns) == 1 && chatTurnRuns[0].FinishedAt.Valid + }, testutil.IntervalFast) + + require.Len(t, chatTurnRuns, 1) + run := chatTurnRuns[0] + require.Equal(t, string(codersdk.ChatDebugStatusCompleted), run.Status) + + steps, err := db.GetChatDebugStepsByRunID(debugCtx, run.ID) + require.NoError(t, err) + require.Len(t, steps, 1) + step := steps[0] + require.Equal(t, string(codersdk.ChatDebugStepOperationStream), step.Operation) + require.Equal(t, string(codersdk.ChatDebugStatusCompleted), step.Status) + require.NotEmpty(t, step.NormalizedRequest) + require.True(t, step.NormalizedResponse.Valid) + require.True(t, step.Usage.Valid) + require.NotEmpty(t, step.Attempts) + require.True(t, step.FinishedAt.Valid) + + var normalizedRequest map[string]any + require.NoError(t, json.Unmarshal(step.NormalizedRequest, &normalizedRequest)) + require.NotEmpty(t, normalizedRequest["messages"]) + + var normalizedResponse map[string]any + require.NoError(t, json.Unmarshal(step.NormalizedResponse.RawMessage, &normalizedResponse)) + require.NotEmpty(t, normalizedResponse["content"]) + require.NotEmpty(t, normalizedResponse["usage"]) + + var usage map[string]any + require.NoError(t, json.Unmarshal(step.Usage.RawMessage, &usage)) + require.EqualValues(t, 200, usage["input_tokens"]) + require.EqualValues(t, 75, usage["output_tokens"]) + require.EqualValues(t, 30, usage["cache_creation_tokens"]) + require.EqualValues(t, 150, usage["cache_read_tokens"]) + + var attempts []map[string]any + require.NoError(t, json.Unmarshal(step.Attempts, &attempts)) + require.Len(t, attempts, 1) + require.NotEmpty(t, attempts[0]["request_body"]) + require.NotEmpty(t, attempts[0]["response_body"]) + + var summary map[string]any + require.NoError(t, json.Unmarshal(run.Summary, &summary)) + require.Equal(t, "POST /v1/messages", summary["endpoint_label"]) + require.Equal(t, "hello debug", summary["first_message"]) + require.EqualValues(t, 1, summary["step_count"]) + require.EqualValues(t, 200, summary["total_input_tokens"]) + require.EqualValues(t, 75, summary["total_output_tokens"]) + require.EqualValues(t, 30, summary["total_cache_creation_tokens"]) + require.EqualValues(t, 150, summary["total_cache_read_tokens"]) +} + +func TestActiveServer_ChatTurnDebugRunRecordsMultipleStreamSteps(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + if !req.Stream { + return chattest.AnthropicNonStreamingResponse(`{"label":"Read file"}`) + } + switch streamCount.Add(1) { + case 1: + return chattest.AnthropicStreamingResponse( + chattest.AnthropicToolCallChunks("read_file", `{"path":"/tmp/a.txt"}`)..., + ) + case 2: + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunksWithCacheUsage(chattest.AnthropicUsage{ + InputTokens: 20, + OutputTokens: 7, + }, "final debug response")...) + default: + t.Fatalf("unexpected stream request %d", streamCount.Load()) + return chattest.AnthropicStreamingResponse() + } + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/a.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, FileSize: 12, TotalLines: 1, LinesRead: 1, Content: "1\tpackage main"}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AlwaysEnableDebugLogs = true + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "multi-step-debug", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file and continue"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.NoError(t, server.Close()) + debugCtx := testutil.Context(t, testutil.WaitLong) + + var chatTurnRuns []database.ChatDebugRun + testutil.Eventually(debugCtx, t, func(ctx context.Context) bool { + runs, err := db.GetChatDebugRunsByChatID(ctx, database.GetChatDebugRunsByChatIDParams{ + ChatID: chat.ID, + LimitVal: 100, + }) + if err != nil { + return false + } + chatTurnRuns = chatTurnRuns[:0] + for _, run := range runs { + if run.Kind == string(codersdk.ChatDebugRunKindChatTurn) { + chatTurnRuns = append(chatTurnRuns, run) + } + } + if len(chatTurnRuns) != 1 || !chatTurnRuns[0].FinishedAt.Valid { + return false + } + steps, err := db.GetChatDebugStepsByRunID(ctx, chatTurnRuns[0].ID) + return err == nil && len(steps) == 2 + }, testutil.IntervalFast) + + require.Len(t, chatTurnRuns, 1) + run := chatTurnRuns[0] + require.Equal(t, string(codersdk.ChatDebugStatusCompleted), run.Status) + + steps, err := db.GetChatDebugStepsByRunID(debugCtx, run.ID) + require.NoError(t, err) + require.Len(t, steps, 2) + for i, step := range steps { + require.EqualValues(t, i+1, step.StepNumber) + require.Equal(t, string(codersdk.ChatDebugStepOperationStream), step.Operation) + require.Equal(t, string(codersdk.ChatDebugStatusCompleted), step.Status) + require.NotEmpty(t, step.Attempts) + require.True(t, step.FinishedAt.Valid) + } + + var firstResponse map[string]any + require.NoError(t, json.Unmarshal(steps[0].NormalizedResponse.RawMessage, &firstResponse)) + require.NotEmpty(t, firstResponse["content"]) + + var secondResponse map[string]any + require.NoError(t, json.Unmarshal(steps[1].NormalizedResponse.RawMessage, &secondResponse)) + require.NotEmpty(t, secondResponse["content"]) + + var summary map[string]any + require.NoError(t, json.Unmarshal(run.Summary, &summary)) + require.Equal(t, "POST /v1/messages", summary["endpoint_label"]) + require.EqualValues(t, 2, summary["step_count"]) + require.EqualValues(t, 30, summary["total_input_tokens"]) + require.EqualValues(t, 12, summary["total_output_tokens"]) +} + +func TestActiveServer_AnthropicSanitizesProviderToolBeforeRequest(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertOrphanProviderToolCall(ctx, t, db, chat.ID, model.ID) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + body := anthropicRequestBody(t, generationRequests[1]) + require.NotContains(t, body, "web_search") + require.Contains(t, body, "partial") + require.Contains(t, body, "continue") + requireAnthropicRequestRedactedReasoning(t, generationRequests[1], "redacted-payload") +} + +func TestActiveServer_AnthropicProviderToolPreRequestGuard(t *testing.T) { + t.Parallel() + + webSearchEnabled := true + callConfig := codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{ + WebSearchEnabled: &webSearchEnabled, + }, + }, + } + + t.Run("allowed web search survives when provider tool is enabled", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCallConfig(t, db, model, callConfig) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderToolPairMessageWithLocalTool(ctx, t, db, chat.ID, model.ID, "ws-allowed") + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + body := anthropicRequestBody(t, generationRequests[1]) + require.Contains(t, body, "ws-allowed") + require.Contains(t, body, "web_search") + }) + + t.Run("web search history survives when provider tool is disabled", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search and read") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + insertProviderToolPairMessageWithLocalTool(ctx, t, db, chat.ID, model.ID, "ws-disabled") + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + ModelConfigID: model.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("continue")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + body := anthropicRequestBody(t, generationRequests[1]) + require.Contains(t, body, "ws-disabled") + require.Contains(t, body, "web_search") + require.Contains(t, body, "tc-1") + require.Contains(t, body, "file") + }) +} + +func TestActiveServer_AnthropicDropsUnpairedProviderToolBeforePersist(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + toolName string + toolInput json.RawMessage + }{ + { + name: "web_search", + toolName: "web_search", + toolInput: json.RawMessage(`{"query":"coder"}`), + }, + { + name: "code_execution", + toolName: "code_execution", + toolInput: json.RawMessage(`{"code":"print(1)"}`), + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + if requestCount.Add(1) == 1 { + return chattest.AnthropicStreamingResponse( + anthropicServerToolUseChunks("pt-1", tt.toolName, tt.toolInput, "tool_use")..., + ) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("after sanitized step")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = enableAnthropicWebSearchForTest(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "run provider tool") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 1) + messages := chatMessages(ctx, t, db, chat.ID) + last := messages[len(messages)-1] + require.Equal(t, database.ChatMessageRoleUser, last.Role) + requireTextPart(t, last, "run provider tool") + require.False(t, toolPartExists(chatToolParts(ctx, t, db, chat.ID), tt.toolName), + "unpaired provider tool content should not be committed") + }) + } +} + +func TestActiveServer_AnthropicKeepsPairedWebSearchBeforePersist(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + return chattest.AnthropicStreamingResponse( + anthropicWebSearchPairChunks("ws-1", `{"query":"coder"}`, "search done", "end_turn")..., + ) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = enableAnthropicWebSearchForTest(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 1) + parts := chatToolParts(ctx, t, db, chat.ID) + toolCall := requireToolCallPart(t, parts, "web_search") + require.Equal(t, "ws-1", toolCall.ToolCallID) + require.True(t, toolCall.ProviderExecuted) + toolResult := requireToolResultPart(t, parts, "web_search") + require.Equal(t, "ws-1", toolResult.ToolCallID) + require.True(t, toolResult.ProviderExecuted) + require.NotEmpty(t, toolResult.ProviderMetadata) + messages := chatMessages(ctx, t, db, chat.ID) + requireTextPart(t, messages[len(messages)-1], "search done") +} + +// TestActiveServer_AnthropicWebSearchFollowUpHasNoSyntheticCancellation +// reproduces a bug where sending a follow-up user message after a +// completed provider-executed web_search turn inserted a synthetic +// cancellation tool-result ("Tool execution interrupted by new user +// message") for the server tool call. The provider-executed result +// lives inside the assistant message, so the cancellation synthesizer +// saw the call as outstanding and emitted a client-style tool-role +// result for a srvtoolu_ ID. On the next request that result replays +// as a plain tool_result block, which Anthropic rejects: +// +// unexpected `tool_use_id` found in `tool_result` blocks: +// srvtoolu_... Each `tool_result` block must have a +// corresponding `tool_use` block in the previous message. +func TestActiveServer_AnthropicWebSearchFollowUpHasNoSyntheticCancellation(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var streamingRequestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + if streamingRequestCount.Add(1) == 1 { + return chattest.AnthropicStreamingResponse( + anthropicWebSearchPairChunks("srvtoolu_ws1", `{"query":"coder"}`, "search done", "end_turn")..., + ) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("follow-up done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = enableAnthropicWebSearchForTest(t, db, model) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + // Simulate a web search turn followed by a user follow-up. + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("thanks, tell me more")}, + }) + require.NoError(t, err) + + // Wait for the follow-up turn to run and the chat to settle. + testutil.Eventually(ctx, t, func(context.Context) bool { + return streamingRequestCount.Load() >= 2 + }, testutil.IntervalFast) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + // The provider-executed web_search call is answered by the + // provider-executed result inside the assistant message. No + // tool-role message may carry a synthetic result for it. + for _, msg := range chatMessages(ctx, t, db, chat.ID) { + if msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + for _, part := range parts { + if part.Type != codersdk.ChatMessagePartTypeToolResult { + continue + } + require.NotEqual(t, "srvtoolu_ws1", part.ToolCallID, + "provider-executed web_search call received a synthetic tool-role result: %s", string(part.Result)) + } + } +} + +func TestActiveServer_AnthropicSanitizesWebSearchBeforeContinuation(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + requests := newAnthropicRequestRecorder() + var requestCount atomic.Int32 + anthropicURL := chattest.NewAnthropic(t, func(req *chattest.AnthropicRequest) chattest.AnthropicResponse { + requests.record(req) + if !req.Stream { + return chattest.AnthropicNonStreamingResponse("title") + } + if requestCount.Add(1) == 1 { + chunks := anthropicServerToolUseChunks("ws-1", "web_search", json.RawMessage(`{"query":"coder"}`), "tool_use") + chunks = append(chunks[:len(chunks)-2], anthropicToolUseChunksWithoutMessageEnvelope(1, "tc-1", "read_file", `{"path":"main.go"}`)...) + chunks = append(chunks, + chattest.AnthropicChunk{ + Type: "message_delta", + StopReason: "tool_use", + Usage: chattest.AnthropicUsage{InputTokens: 10, OutputTokens: 5}, + }, + chattest.AnthropicChunk{Type: "message_stop"}, + ) + return chattest.AnthropicStreamingResponse(chunks...) + } + return chattest.AnthropicStreamingResponse(chattest.AnthropicTextChunks("done")...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = enableAnthropicWebSearchForTest(t, db, model) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "main.go", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true, Content: "package main", FileSize: 12, TotalLines: 1, LinesRead: 1}, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "anthropic-web-search-continuation", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("search and read"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + generationRequests := filterAnthropicStreamingRequests(requests.all()) + require.Len(t, generationRequests, 2) + continuationBody := anthropicRequestBody(t, generationRequests[1]) + require.NotContains(t, continuationBody, "server_tool_use") + require.NotContains(t, continuationBody, "web_search_tool_result") + require.NotContains(t, continuationBody, "ws-1") + require.Contains(t, continuationBody, "tc-1") + require.Contains(t, continuationBody, "package main") + + parts := chatToolParts(ctx, t, db, chat.ID) + require.False(t, toolPartExists(parts, "web_search")) + toolCall := requireToolCallPart(t, parts, "read_file") + require.Equal(t, "tc-1", toolCall.ToolCallID) + require.False(t, toolCall.ProviderExecuted) + toolResult := requireToolResultPart(t, parts, "read_file") + require.Equal(t, "tc-1", toolResult.ToolCallID) + require.False(t, toolResult.ProviderExecuted) +} + +func TestActiveServer_ExclusiveToolPolicy(t *testing.T) { + t.Parallel() + + t.Run("mixed exclusive and local tools commit policy errors", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + advisorChunk := chattest.OpenAIToolCallChunk("advisor", `{"question":"help"}`) + readChunk := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/a.txt"}`) + readCall := readChunk.Choices[0].ToolCalls[0] + readCall.Index = 1 + advisorChunk.Choices[0].ToolCalls = append(advisorChunk.Choices[0].ToolCalls, readCall) + return chattest.OpenAIStreamingResponse(advisorChunk) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{Enabled: true, MaxUsesPerRun: 3, MaxOutputTokens: 1024}) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Times(0) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "exclusive-local-policy", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("advise and read"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + advisorResult := requireToolResultPart(t, parts, "advisor") + readResult := requireToolResultPart(t, parts, "read_file") + require.True(t, advisorResult.IsError) + require.True(t, readResult.IsError) + require.Contains(t, string(advisorResult.Result), "advisor must be called alone, without other tools in the same batch") + require.Contains(t, string(readResult.Result), "this tool was skipped because advisor must run alone in its batch") + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2)) + }) + + t.Run("mixed exclusive and dynamic tools commit policy errors", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if streamedCallCount.Add(1) == 1 { + advisorChunk := chattest.OpenAIToolCallChunk("advisor", `{"question":"help"}`) + dynamicChunk := chattest.OpenAIToolCallChunk("mcp_tool", `{"q":"docs"}`) + dynamicCall := dynamicChunk.Choices[0].ToolCalls[0] + dynamicCall.Index = 1 + advisorChunk.Choices[0].ToolCalls = append(advisorChunk.Choices[0].ToolCalls, dynamicCall) + return chattest.OpenAIStreamingResponse(advisorChunk) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{Enabled: true, MaxUsesPerRun: 3, MaxOutputTokens: 1024}) + dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ + Name: "mcp_tool", + Description: "dynamic test tool", + InputSchema: mcpgo.ToolInputSchema{Type: "object", Properties: map[string]any{"q": map[string]any{"type": "string"}}}, + }}) + require.NoError(t, err) + + server := newActiveTestServer(t, db, ps) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "exclusive-dynamic-policy", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("advise and call dynamic"), + }, + DynamicTools: dynamicToolsJSON, + }) + require.NoError(t, err) + chatResult := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.NotEqual(t, database.ChatStatusRequiresAction, chatResult.Status) + + parts := chatToolParts(ctx, t, db, chat.ID) + advisorResult := requireToolResultPart(t, parts, "advisor") + dynamicResult := requireToolResultPart(t, parts, "mcp_tool") + require.True(t, advisorResult.IsError) + require.True(t, dynamicResult.IsError) + require.Contains(t, string(advisorResult.Result), "advisor must be called alone, without other tools in the same batch") + require.Contains(t, string(dynamicResult.Result), "this tool was skipped because advisor must run alone in its batch") + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2)) + }) + + t.Run("solo exclusive tool executes", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamedCallCount.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("advisor", `{"question":"help me decide"}`), + ) + case 2: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("nested advice")...) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{Enabled: true, MaxUsesPerRun: 3, MaxOutputTokens: 1024}) + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "advise only") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + result := requireToolResultPart(t, parts, "advisor") + require.False(t, result.IsError) + require.Contains(t, string(result.Result), "nested advice") + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(3)) + }) + + t.Run("exclusive tool with provider executed tool executes", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + webSearchEnabled := true + var streamedCallCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch streamedCallCount.Add(1) { + case 1: + return chattest.OpenAIResponse{ + StreamingChunks: chattest.OpenAIStreamingResponse( + chattest.OpenAIToolCallChunk("advisor", `{"question":"search informed advice"}`), + ).StreamingChunks, + WebSearch: &chattest.OpenAIWebSearchCall{ID: "ws-advisor", Query: "coder"}, + } + case 2: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("nested advice")...) + default: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + OpenAI: &codersdk.ChatModelOpenAIProviderOptions{WebSearchEnabled: &webSearchEnabled}, + }, + }) + seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{Enabled: true, MaxUsesPerRun: 3, MaxOutputTokens: 1024}) + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search then advise") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + advisorResult := requireToolResultPart(t, parts, "advisor") + webResult := requireToolResultPart(t, parts, "web_search") + require.False(t, advisorResult.IsError) + require.True(t, webResult.ProviderExecuted) + require.GreaterOrEqual(t, streamedCallCount.Load(), int32(3)) + }) +} + +func TestActiveServer_ReasoningTimestamps(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + sendReasoning := true + thinkingBudget := int64(1024) + anthropicURL := chattest.NewAnthropic(t, func(_ *chattest.AnthropicRequest) chattest.AnthropicResponse { + return chattest.AnthropicStreamingResponse(chattest.AnthropicReasoningTextChunks( + []chattest.AnthropicReasoningBlock{ + {Text: "first thought", Signature: "sig_1"}, + {Text: "second thought", Signature: "sig_2"}, + }, + "answer", + )...) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{ + SendReasoning: &sendReasoning, + Thinking: &codersdk.ChatModelAnthropicThinkingOptions{ + BudgetTokens: &thinkingBudget, + }, + }, + }, + }) + + server := newActiveTestServer(t, db, ps) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "think") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + messages := chatMessages(ctx, t, db, chat.ID) + assistant := messages[len(messages)-1] + reasoningParts := reasoningPartsFromMessage(t, assistant) + require.Len(t, reasoningParts, 2) + require.Equal(t, []string{"first thought", "second thought"}, []string{ + strings.TrimSpace(reasoningParts[0].Text), + strings.TrimSpace(reasoningParts[1].Text), + }) + for i := range reasoningParts { + require.NotNil(t, reasoningParts[i].CreatedAt) + require.NotNil(t, reasoningParts[i].CompletedAt) + require.False(t, reasoningParts[i].CreatedAt.IsZero()) + require.False(t, reasoningParts[i].CompletedAt.IsZero()) + require.False(t, reasoningParts[i].CompletedAt.Before(*reasoningParts[i].CreatedAt)) + } + require.False(t, reasoningParts[1].CreatedAt.Before(*reasoningParts[0].CompletedAt)) +} + +func TestAnthropicProviderToolPreRequestGuard(t *testing.T) { + t.Parallel() + + providerPair := func(id string) []fantasy.MessagePart { + return []fantasy.MessagePart{ + fantasy.ToolCallPart{ + ToolCallID: id, + ToolName: "web_search", + Input: `{"query":"coder"}`, + ProviderExecuted: true, + }, + fantasy.ToolResultPart{ + ToolCallID: id, + Output: fantasy.ToolResultOutputContentText{Text: "ok"}, + ProviderExecuted: true, + ProviderOptions: fantasy.ProviderOptions(validWebSearchProviderMetadataForTest()), + }, + } + } + + t.Run("orphan provider result is textified", func(t *testing.T) { + t.Parallel() + + guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( + context.Background(), + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + fantasyanthropic.Name, + "claude-test", + []fantasy.Message{ + { + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{ + fantasy.TextPart{Text: "keep"}, + fantasy.ToolResultPart{ + ToolCallID: "ws-orphan", + Output: fantasy.ToolResultOutputContentText{Text: "search result"}, + ProviderExecuted: true, + }, + }, + }, + }, + ) + require.NoError(t, err) + + requireNoProviderExecutedToolResultPrompt(t, guarded) + requireAnthropicProviderToolPromptSafe(t, guarded) + require.Len(t, guarded, 1) + require.Len(t, guarded[0].Content, 2) + textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](guarded[0].Content[0]) + require.True(t, ok) + require.Equal(t, "keep", textPart.Text) + textPart, ok = fantasy.AsMessagePart[fantasy.TextPart](guarded[0].Content[1]) + require.True(t, ok) + require.Equal(t, "search result", textPart.Text) + }) + + t.Run("valid provider history is unchanged", func(t *testing.T) { + t.Parallel() + + content := []fantasy.MessagePart{fantasy.TextPart{Text: "keep"}} + content = append(content, providerPair("ws-one")...) + content = append(content, providerPair("ws-two")...) + guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( + context.Background(), + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + fantasyanthropic.Name, + "claude-test", + []fantasy.Message{{Role: fantasy.MessageRoleAssistant, Content: content}}, + ) + require.NoError(t, err) + + requireAnthropicProviderToolPromptSafe(t, guarded) + require.Len(t, guarded, 1) + require.Len(t, guarded[0].Content, len(content)) + requireProviderExecutedToolCallPrompt(t, guarded, "ws-one") + requireProviderExecutedToolResultPrompt(t, guarded, "ws-one") + requireProviderExecutedToolCallPrompt(t, guarded, "ws-two") + requireProviderExecutedToolResultPrompt(t, guarded, "ws-two") + }) + + t.Run("non Anthropic providers are unchanged", func(t *testing.T) { + t.Parallel() + + prompt := []fantasy.Message{ + { + Role: fantasy.MessageRoleAssistant, + Content: providerPair("ws-other-provider"), + }, + } + guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( + context.Background(), + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + "fake", + "fake-model", + prompt, + ) + require.NoError(t, err) + require.Equal(t, prompt, guarded) + }) + + t.Run("logs removals", func(t *testing.T) { + t.Parallel() + + logSink := testutil.NewFakeSink(t) + logger := logSink.Logger() + logPair := providerPair("ws-log") + guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( + context.Background(), + logger, + fantasyanthropic.Name, + "claude-test", + []fantasy.Message{ + { + Role: fantasy.MessageRoleAssistant, + Content: []fantasy.MessagePart{ + logPair[1], + logPair[0], + }, + }, + }, + ) + require.NoError(t, err) + + requireNoProviderExecutedToolCallPrompt(t, guarded) + requireNoProviderExecutedToolResultPrompt(t, guarded) + requireTextPrompt(t, guarded, "ok") + entries := logSink.Entries(func(e slog.SinkEntry) bool { + return e.Level == slog.LevelWarn && + e.Message == "removed provider-executed tool history" + }) + require.Len(t, entries, 1) + require.Equal(t, "pre_request_guard", requireLogField(t, entries[0], "phase")) + require.Equal(t, 1, requireLogField(t, entries[0], "removed_tool_calls")) + require.Equal(t, 1, requireLogField(t, entries[0], "removed_tool_results")) + }) +} + +func enableAnthropicWebSearchForTest( + t *testing.T, + db database.Store, + model database.ChatModelConfig, +) database.ChatModelConfig { + t.Helper() + webSearchEnabled := true + return updateChatModelCallConfig(t, db, model, codersdk.ChatModelCallConfig{ + ProviderOptions: &codersdk.ChatModelProviderOptions{ + Anthropic: &codersdk.ChatModelAnthropicProviderOptions{ + WebSearchEnabled: &webSearchEnabled, + }, + }, + }) +} + +func anthropicMessageStartChunk(messageID string) chattest.AnthropicChunk { + return chattest.AnthropicChunk{ + Type: "message_start", + Message: chattest.AnthropicChunkMessage{ + ID: messageID, + Type: "message", + Role: "assistant", + Model: "claude-3-opus-20240229", + }, + } +} + +func anthropicServerToolUseChunks( + toolCallID string, + toolName string, + input json.RawMessage, + stopReason string, +) []chattest.AnthropicChunk { + chunks := []chattest.AnthropicChunk{ + anthropicMessageStartChunk("msg-" + toolCallID), + } + chunks = append(chunks, anthropicServerToolUseChunksWithoutMessageEnvelope(0, toolCallID, toolName, input)...) + chunks = append(chunks, + chattest.AnthropicChunk{ + Type: "message_delta", + StopReason: stopReason, + Usage: chattest.AnthropicUsage{InputTokens: 10, OutputTokens: 5}, + }, + chattest.AnthropicChunk{Type: "message_stop"}, + ) + return chunks +} + +func anthropicServerToolUseChunksWithoutMessageEnvelope( + index int, + toolCallID string, + toolName string, + input json.RawMessage, +) []chattest.AnthropicChunk { + return []chattest.AnthropicChunk{ + { + Type: "content_block_start", + Index: index, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "server_tool_use", + ID: toolCallID, + Name: toolName, + Input: input, + }, + }, + { + Type: "content_block_stop", + Index: index, + }, + } +} + +func anthropicToolUseChunksWithoutMessageEnvelope( + index int, + toolCallID string, + toolName string, + input string, +) []chattest.AnthropicChunk { + return []chattest.AnthropicChunk{ + { + Type: "content_block_start", + Index: index, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "tool_use", + ID: toolCallID, + Name: toolName, + Input: json.RawMessage(`{}`), + }, + }, + { + Type: "content_block_delta", + Index: index, + Delta: chattest.AnthropicDeltaBlock{ + Type: "input_json_delta", + PartialJSON: input, + }, + }, + { + Type: "content_block_stop", + Index: index, + }, + } +} + +func anthropicWebSearchPairChunks( + toolCallID string, + queryInput string, + text string, + stopReason string, +) []chattest.AnthropicChunk { + resultContent := []map[string]any{{ + "type": "web_search_result", + "url": "https://example.com/coder", + "title": "Coder", + "encrypted_content": "encrypted-coder", + }} + chunks := []chattest.AnthropicChunk{ + anthropicMessageStartChunk("msg-" + toolCallID), + } + chunks = append(chunks, anthropicServerToolUseChunksWithoutMessageEnvelope(0, toolCallID, "web_search", json.RawMessage(queryInput))...) + chunks = append(chunks, + chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 1, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "web_search_tool_result", + ToolUseID: toolCallID, + Content: resultContent, + }, + }, + chattest.AnthropicChunk{Type: "content_block_stop", Index: 1}, + chattest.AnthropicChunk{ + Type: "content_block_start", + Index: 2, + ContentBlock: chattest.AnthropicContentBlock{ + Type: "text", + }, + }, + chattest.AnthropicChunk{ + Type: "content_block_delta", + Index: 2, + Delta: chattest.AnthropicDeltaBlock{ + Type: "text_delta", + Text: text, + }, + }, + chattest.AnthropicChunk{Type: "content_block_stop", Index: 2}, + chattest.AnthropicChunk{ + Type: "message_delta", + StopReason: stopReason, + Usage: chattest.AnthropicUsage{InputTokens: 10, OutputTokens: 5}, + }, + chattest.AnthropicChunk{Type: "message_stop"}, + ) + return chunks +} + +func toolPartExists(parts []codersdk.ChatMessagePart, toolName string) bool { + for _, part := range parts { + if (part.Type == codersdk.ChatMessagePartTypeToolCall || part.Type == codersdk.ChatMessagePartTypeToolResult) && + part.ToolName == toolName { + return true + } + } + return false +} + +func updateChatModelCompressionThreshold(t *testing.T, db database.Store, model database.ChatModelConfig, contextLimit int64, threshold int32) database.ChatModelConfig { + t.Helper() + model.ContextLimit = contextLimit + model.CompressionThreshold = threshold + updated, err := db.UpdateChatModelConfig(context.Background(), database.UpdateChatModelConfigParams{ + ID: model.ID, + DisplayName: model.DisplayName, + Model: model.Model, + Provider: model.Provider, + Enabled: model.Enabled, + ContextLimit: model.ContextLimit, + CompressionThreshold: model.CompressionThreshold, + Options: model.Options, + AIProviderID: model.AIProviderID, + }) + require.NoError(t, err) + return updated +} + +func updateChatModelContextLimit(t *testing.T, db database.Store, model database.ChatModelConfig) database.ChatModelConfig { + t.Helper() + updated, err := db.UpdateChatModelConfig(context.Background(), database.UpdateChatModelConfigParams{ + ID: model.ID, + DisplayName: model.DisplayName, + Model: model.Model, + Provider: model.Provider, + Enabled: model.Enabled, + ContextLimit: model.ContextLimit, + CompressionThreshold: model.CompressionThreshold, + Options: model.Options, + AIProviderID: model.AIProviderID, + }) + require.NoError(t, err) + return updated +} + +func updateChatModelCallConfig(t *testing.T, db database.Store, model database.ChatModelConfig, callConfig codersdk.ChatModelCallConfig) database.ChatModelConfig { + t.Helper() + options, err := json.Marshal(callConfig) + require.NoError(t, err) + updated, err := db.UpdateChatModelConfig(context.Background(), database.UpdateChatModelConfigParams{ + ID: model.ID, + DisplayName: model.DisplayName, + Model: model.Model, + Provider: model.Provider, + Enabled: model.Enabled, + ContextLimit: model.ContextLimit, + CompressionThreshold: model.CompressionThreshold, + Options: options, + AIProviderID: model.AIProviderID, + }) + require.NoError(t, err) + return updated +} + +func insertAssistantTextMessage( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + text string, + modelID uuid.UUID, +) { + t.Helper() + insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, []codersdk.ChatMessagePart{ + codersdk.ChatMessageText(text), + }) +} + +func insertProviderToolPairMessageWithLocalTool( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + modelID uuid.UUID, + toolCallID string, +) { + t.Helper() + metadata, err := json.Marshal(fantasy.ProviderMetadata{ + fantasyanthropic.Name: &fantasyanthropic.WebSearchResultMetadata{ + Results: []fantasyanthropic.WebSearchResultItem{{ + URL: "https://example.com", + Title: "Example", + EncryptedContent: "encrypted", + }}, + }, + }) + require.NoError(t, err) + parts := []codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: toolCallID, + ToolName: "web_search", + Args: json.RawMessage(`{"query":"coder"}`), + ProviderExecuted: true, + }, + { + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: toolCallID, + ToolName: "web_search", + Result: json.RawMessage(`"ok"`), + ProviderExecuted: true, + ProviderMetadata: metadata, + }, + } + parts = append(parts, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "tc-1", + ToolName: "read_file", + Args: json.RawMessage(`{"path":"main.go"}`), + }) + insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleAssistant, modelID, uuid.Nil, parts) + insertChatMessageParts(ctx, t, db, chatID, database.ChatMessageRoleTool, modelID, uuid.Nil, []codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: "tc-1", + ToolName: "read_file", + Result: json.RawMessage(`"file"`), + }, + }) +} + +func insertChatMessageParts( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + role database.ChatMessageRole, + modelID uuid.UUID, + createdBy uuid.UUID, + parts []codersdk.ChatMessagePart, +) database.ChatMessage { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + var params database.InsertChatMessagesParams + if role == database.ChatMessageRoleUser { + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: createdBy}) + params = chatd.BuildSingleUserChatMessageInsertParams( + chatID, + apiKey.ID, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + createdBy, + ) + } else { + params = chatd.BuildSingleChatMessageInsertParams( + chatID, + role, + content, + database.ChatMessageVisibilityBoth, + modelID, + chatprompt.CurrentContentVersion, + createdBy, + ) + } + messages, err := db.InsertChatMessages(ctx, params) + require.NoError(t, err) + require.Len(t, messages, 1) + return messages[0] +} + +func createPlanSubagentChatWithHistory( + ctx context.Context, + t *testing.T, + db database.Store, + orgID uuid.UUID, + userID uuid.UUID, + modelID uuid.UUID, +) database.Chat { + t.Helper() + rootChat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: orgID, + OwnerID: userID, + LastModelConfigID: modelID, + Title: "plan subagent active tools root", + Status: database.ChatStatusWaiting, + PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, + MCPServerIDs: []uuid.UUID{}, + ClientType: database.ChatClientTypeApi, + }) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: orgID, + OwnerID: userID, + LastModelConfigID: modelID, + Title: "plan subagent active tools", + Status: database.ChatStatusWaiting, + PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, + ParentChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: rootChat.ID, Valid: true}, + MCPServerIDs: []uuid.UUID{}, + ClientType: database.ChatClientTypeApi, + }) + insertSystemTextMessage(ctx, t, db, chat.ID, "You are not currently connected to a workspace.", modelID) + insertChatMessageParts(ctx, t, db, chat.ID, database.ChatMessageRoleUser, modelID, userID, []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("hello"), + }) + return chat +} + +func anthropicRequestToolNames(req chattest.AnthropicRequest) []string { + names := make([]string, 0, len(req.Tools)) + for _, tool := range req.Tools { + names = append(names, tool.Name) + } + return names +} + +func anthropicRequestContainsPromptSentinel(t *testing.T, req chattest.AnthropicRequest) bool { + t.Helper() + body := anthropicRequestBody(t, req) + return strings.Contains(body, "__chatd_agent_prompt_sentinel_") +} + +func reasoningPartsFromMessage(t *testing.T, msg database.ChatMessage) []codersdk.ChatMessagePart { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + var reasoning []codersdk.ChatMessagePart + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeReasoning { + reasoning = append(reasoning, part) + } + } + return reasoning +} + +func validWebSearchProviderMetadataForTest() fantasy.ProviderMetadata { + return fantasy.ProviderMetadata{ + fantasyanthropic.Name: &fantasyanthropic.WebSearchResultMetadata{ + Results: []fantasyanthropic.WebSearchResultItem{ + { + URL: "https://example.com", + Title: "Example", + EncryptedContent: "encrypted", + }, + }, + }, + } +} + +func safeToolCallPart(part fantasy.MessagePart) (fantasy.ToolCallPart, bool) { + var zero fantasy.ToolCallPart + if part == nil { + return zero, false + } + if value, ok := part.(*fantasy.ToolCallPart); ok && value == nil { + return zero, false + } + type toolCallPart = fantasy.ToolCallPart + return fantasy.AsMessagePart[toolCallPart](part) +} + +func safeToolResultPart(part fantasy.MessagePart) (fantasy.ToolResultPart, bool) { + var zero fantasy.ToolResultPart + if part == nil { + return zero, false + } + if value, ok := part.(*fantasy.ToolResultPart); ok && value == nil { + return zero, false + } + type toolResultPart = fantasy.ToolResultPart + return fantasy.AsMessagePart[toolResultPart](part) +} + +func requireProviderExecutedToolCallPrompt( + t *testing.T, + prompt []fantasy.Message, + id string, +) fantasy.ToolCallPart { + t.Helper() + for _, message := range prompt { + for _, part := range message.Content { + toolCall, ok := safeToolCallPart(part) + if ok && toolCall.ProviderExecuted && toolCall.ToolCallID == id { + return toolCall + } + } + } + t.Fatalf("missing provider-executed prompt tool call %q", id) + return fantasy.ToolCallPart{} +} + +func requireProviderExecutedToolResultPrompt( + t *testing.T, + prompt []fantasy.Message, + id string, +) fantasy.ToolResultPart { + t.Helper() + for _, message := range prompt { + for _, part := range message.Content { + toolResult, ok := safeToolResultPart(part) + if ok && toolResult.ProviderExecuted && toolResult.ToolCallID == id { + return toolResult + } + } + } + t.Fatalf("missing provider-executed prompt tool result %q", id) + return fantasy.ToolResultPart{} +} + +func requireNoProviderExecutedToolCallPrompt(t *testing.T, prompt []fantasy.Message) { + t.Helper() + for i, message := range prompt { + for j, part := range message.Content { + toolCall, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](part) + if ok && toolCall.ProviderExecuted { + t.Fatalf("prompt[%d].content[%d]: unexpected provider-executed call", i, j) + } + } + } +} + +func requireNoProviderExecutedToolResultPrompt(t *testing.T, prompt []fantasy.Message) { + t.Helper() + for i, message := range prompt { + for j, part := range message.Content { + toolResult, ok := safeToolResultPart(part) + if ok && toolResult.ProviderExecuted { + t.Fatalf("prompt[%d].content[%d]: unexpected provider-executed result", i, j) + } + } + } +} + +func requireTextPrompt(t *testing.T, prompt []fantasy.Message, text string) fantasy.TextPart { + t.Helper() + for _, message := range prompt { + for _, part := range message.Content { + textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](part) + if ok && textPart.Text == text { + return textPart + } + } + } + t.Fatalf("missing prompt text %q", text) + return fantasy.TextPart{} +} + +func requireAnthropicProviderToolPromptSafe(t *testing.T, prompt []fantasy.Message) { + t.Helper() + require.Empty(t, chatsanitize.ValidateAnthropicProviderToolHistory(prompt)) +} + +func requireLogField(t *testing.T, entry slog.SinkEntry, name string) any { + t.Helper() + for _, field := range entry.Fields { + if field.Name == name { + return field.Value + } + } + t.Fatalf("missing log field %q", name) + return nil +} + func TestPassiveServerDoesNotProcess(t *testing.T) { t.Parallel() @@ -5822,6 +8164,7 @@ func TestPassiveServerDoesNotProcess(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "should-stay-pending", InitialUserContent: []codersdk.ChatMessagePart{{Type: codersdk.ChatMessagePartTypeText, Text: "hello"}}, ModelConfigID: model.ID, @@ -5830,36 +8173,12 @@ func TestPassiveServerDoesNotProcess(t *testing.T) { chatd.WaitUntilIdleForTest(server) - // Re-read from DB to catch any unexpected state transition. + // Re-read from DB to catch any unexpected processing. stored, err := db.GetChatByID(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, database.ChatStatusPending, stored.Status) -} - -// newStartedTestServer creates a server with Start() called. -// Uses a long acquire interval so processing is triggered by -// wake signals, not polling. -func newStartedTestServer( - t *testing.T, - db database.Store, - ps dbpubsub.Pubsub, - replicaID uuid.UUID, -) *chatd.Server { - t.Helper() - - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: replicaID, - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - }) - server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) - }) - return server + require.Equal(t, database.ChatStatusRunning, stored.Status) + require.False(t, stored.WorkerID.Valid) + require.False(t, stored.RunnerID.Valid) } // newDebugEnabledTestServer creates a passive test server with @@ -5876,11 +8195,10 @@ func newDebugEnabledTestServer( t.Helper() logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: replicaID, - Pubsub: ps, PendingChatAcquireInterval: testutil.WaitLong, AlwaysEnableDebugLogs: true, }) @@ -5907,14 +8225,13 @@ func newActiveTestServer( Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, } for _, o := range overrides { o(&cfg) } - server := chatd.New(cfg) + server := chatd.New(ps, cfg) server.Start() t.Cleanup(func() { require.NoError(t, server.Close()) @@ -5985,11 +8302,10 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { "openai", openAIURL, ) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: testutil.WaitLong, AlwaysEnableDebugLogs: tt.alwaysEnableDebugLogs, }) @@ -6072,6 +8388,7 @@ func seedChatDependenciesWithProvider( t.Helper() user := dbgen.User(t, db, database.User{}) + _ = testAPIKeyID(t, db, user.ID) org := dbgen.Organization(t, db, database.Organization{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -6102,6 +8419,7 @@ func seedChatDependenciesWithProviderPolicy( t.Helper() user := dbgen.User(t, db, database.User{}) + _ = testAPIKeyID(t, db, user.ID) org := dbgen.Organization(t, db, database.Organization{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -6138,45 +8456,14 @@ func seedLastTurnSummary( t.Helper() affected, err := db.UpdateChatLastTurnSummary(ctx, database.UpdateChatLastTurnSummaryParams{ - ID: chat.ID, - ExpectedUpdatedAt: chat.UpdatedAt, - LastTurnSummary: sql.NullString{String: summary, Valid: true}, + ID: chat.ID, + ExpectedHistoryVersion: chat.HistoryVersion, + LastTurnSummary: sql.NullString{String: summary, Valid: true}, }) require.NoError(t, err) require.Equal(t, int64(1), affected) } -func waitForTerminalChatStatusEvent( - ctx context.Context, - t *testing.T, - events <-chan codersdk.ChatStreamEvent, -) codersdk.ChatStatus { - t.Helper() - - var terminalStatus codersdk.ChatStatus - testutil.Eventually(ctx, t, func(context.Context) bool { - for { - select { - case event, ok := <-events: - if !ok { - return false - } - if event.Type != codersdk.ChatStreamEventTypeStatus || event.Status == nil { - continue - } - if event.Status.Status == codersdk.ChatStatusWaiting || event.Status.Status == codersdk.ChatStatusError { - terminalStatus = event.Status.Status - return true - } - default: - return false - } - } - }, testutil.IntervalFast) - - return terminalStatus -} - func waitForTerminalChat( ctx context.Context, t *testing.T, @@ -6289,8 +8576,17 @@ func seedWorkspaceWithAgent( JobID: pj.ID, }) dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ - ResourceID: res.ID, + ResourceID: res.ID, + Directory: "/home/coder/project", + OperatingSystem: "linux", }) + require.NoError(t, db.UpdateWorkspaceAgentStartupByID(context.Background(), database.UpdateWorkspaceAgentStartupByIDParams{ + ID: dbAgent.ID, + Version: "v1.0.0", + ExpandedDirectory: "/home/coder/project", + })) + dbAgent, err := db.GetWorkspaceAgentByID(context.Background(), dbAgent.ID) + require.NoError(t, err) return ws, dbAgent } @@ -6355,11 +8651,10 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -6374,6 +8669,7 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "interrupt-no-push", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -6401,9 +8697,9 @@ func TestInterruptChatDoesNotSendWebPushNotification(t *testing.T) { } }, testutil.IntervalFast) - // Interrupt the chat. - updated := server.InterruptChat(ctx, chat) - require.Equal(t, database.ChatStatusWaiting, updated.Status) + // Interrupt the chat. The worker finalizes the interruption asynchronously. + updated, _ := server.InterruptChat(ctx, chat) + require.Equal(t, database.ChatStatusInterrupting, updated.Status) // Wait for the chat to finish processing and return to waiting. testutil.Eventually(ctx, t, func(ctx context.Context) bool { @@ -6476,11 +8772,10 @@ func TestSuccessfulChatSendsWebPushWithNavigationData(t *testing.T) { mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -6496,6 +8791,7 @@ func TestSuccessfulChatSendsWebPushWithNavigationData(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "push-nav-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -6563,11 +8859,10 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) }) loggerA := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - serverA := chatd.New(chatd.Config{ + serverA := chatd.New(ps, chatd.Config{ Logger: loggerA, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitLong, }) @@ -6582,6 +8877,7 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) chat, err := serverA.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "shutdown-retry", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -6607,22 +8903,11 @@ func TestCloseDuringShutdownContextCanceledShouldRetryOnNewReplica(t *testing.T) require.NoError(t, serverA.Close()) - require.Eventually(t, func() bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusPending && - !fromDB.WorkerID.Valid && - !fromDB.LastError.Valid - }, testutil.WaitMedium, testutil.IntervalFast) - loggerB := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - serverB := chatd.New(chatd.Config{ + serverB := chatd.New(ps, chatd.Config{ Logger: loggerB, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitLong, }) @@ -6672,11 +8957,10 @@ func TestSuccessfulChatSendsWebPushWithSummary(t *testing.T) { mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -6692,6 +8976,7 @@ func TestSuccessfulChatSendsWebPushWithSummary(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "summary-push-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do the thing")}, @@ -6748,6 +9033,7 @@ func TestSuccessfulChatPersistsTurnSummaryWithoutWebPush(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "summary-no-webpush-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do the thing")}, @@ -6790,11 +9076,10 @@ func TestSuccessfulChatSendsWebPushFallbackWithoutSummaryForEmptyAssistantText(t mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -6809,6 +9094,7 @@ func TestSuccessfulChatSendsWebPushFallbackWithoutSummaryForEmptyAssistantText(t chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "empty-summary-push-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do the thing")}, @@ -6850,11 +9136,10 @@ func TestErroredChatClearsLastTurnSummaryAndSendsWebPush(t *testing.T) { mockPush := &mockWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, WebpushDispatcher: mockPush, @@ -6869,6 +9154,7 @@ func TestErroredChatClearsLastTurnSummaryAndSendsWebPush(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "error-summary-clear-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do the thing")}, @@ -7097,6 +9383,7 @@ func TestComputerUseSubagentToolsAndModel(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "computer-use-detection", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -7243,11 +9530,10 @@ func TestInterruptChatPersistsPartialResponse(t *testing.T) { }) logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ + server := chatd.New(ps, chatd.Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, PendingChatAcquireInterval: 10 * time.Millisecond, InFlightChatStaleAfter: testutil.WaitSuperLong, }) @@ -7262,19 +9548,13 @@ func TestInterruptChatPersistsPartialResponse(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "interrupt-persist-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, }) require.NoError(t, err) - // Subscribe to the chat's event stream so we can observe - // message_part events. This proves the chatloop has actually - // processed the streamed chunks. - _, events, subCancel, ok := server.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer subCancel() - // Wait for the mock to finish sending chunks. testutil.Eventually(ctx, t, func(ctx context.Context) bool { select { @@ -7285,27 +9565,9 @@ func TestInterruptChatPersistsPartialResponse(t *testing.T) { } }, testutil.IntervalFast) - // Drain the event channel until we see a message_part event, - // which means the chatloop has consumed and published the chunk. - gotMessagePart := false - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - for { - select { - case ev := <-events: - if ev.Type == codersdk.ChatStreamEventTypeMessagePart { - gotMessagePart = true - return true - } - default: - return gotMessagePart - } - } - }, testutil.IntervalFast) - require.True(t, gotMessagePart, "should have received at least one message_part event") - - // Now interrupt the chat. The chatloop has processed content. - updated := server.InterruptChat(ctx, chat) - require.Equal(t, database.ChatStatusWaiting, updated.Status) + // Now interrupt the chat. The provider has sent partial content. + updated, _ := server.InterruptChat(ctx, chat) + require.Equal(t, database.ChatStatusInterrupting, updated.Status) // Wait for the partial assistant message to be persisted. // After the interrupt, the chatloop runs persistInterruptedStep @@ -7391,6 +9653,7 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "user-provider-key-success", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -7399,15 +9662,8 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { }) require.NoError(t, err) - _, events, cancel, ok := creator.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - _ = newActiveTestServer(t, db, ps) - terminalStatus := waitForTerminalChatStatusEvent(ctx, t, events) - require.Equal(t, codersdk.ChatStatusWaiting, terminalStatus) - chatResult := waitForTerminalChat(ctx, t, db, chat.ID) require.Equal(t, database.ChatStatusWaiting, chatResult.Status) require.False(t, chatResult.LastError.Valid) @@ -7418,21 +9674,12 @@ func TestProcessChat_UserProviderKey_Success(t *testing.T) { require.Contains(t, recordedAuthHeaders, "Bearer "+userAPIKey) } -func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if req.Stream { - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("hello through AI Gateway")..., - ) - } - return chattest.OpenAINonStreamingResponse(`{"title":"AI Gateway Chat"}`) - }) - factory := newChatAIGatewayTestFactory(t, openAIURL) +func seedAIGatewayOpenAITestDependencies( + t *testing.T, + db database.Store, + openAIURL string, +) (database.User, database.Organization, database.AIProvider, database.ChatModelConfig, database.APIKey) { + t.Helper() user := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) @@ -7452,7 +9699,7 @@ func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, }) apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) - _, err := db.UpsertUserAIProviderKey(ctx, database.UpsertUserAIProviderKeyParams{ + _, err := db.UpsertUserAIProviderKey(context.Background(), database.UpsertUserAIProviderKeyParams{ ID: uuid.New(), UserID: user.ID, AIProviderID: provider.ID, @@ -7460,6 +9707,27 @@ func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { }) require.NoError(t, err) + return user, org, provider, model, apiKey +} + +func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if req.Stream { + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("hello through AI Gateway")..., + ) + } + return chattest.OpenAINonStreamingResponse(`{"title":"AI Gateway Chat"}`) + }) + factory := newChatAIGatewayTestFactory(t, openAIURL) + + user, org, provider, model, apiKey := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) + creator := newTestServer(t, db, ps, uuid.New()) chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, @@ -7484,8 +9752,7 @@ func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { cfg.AllowBYOKSet = true }) - terminalStatus := waitForTerminalChatStatusEvent(ctx, t, events) - require.Equal(t, codersdk.ChatStatusWaiting, terminalStatus) + _ = events chatResult := waitForTerminalChat(ctx, t, db, chat.ID) require.Equal(t, database.ChatStatusWaiting, chatResult.Status) @@ -7512,6 +9779,87 @@ func TestProcessChat_AIGatewayRoutingUsesDelegatedAPIKey(t *testing.T) { } } +func TestProcessChat_AIGatewayRoutingPreservesAPIKeyAfterWorkspaceContext(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if req.Stream { + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("hello after workspace context")..., + ) + } + return chattest.OpenAINonStreamingResponse(`{"title":"AI Gateway Workspace"}`) + }) + factory := newChatAIGatewayTestFactory(t, openAIURL) + user, org, provider, model, apiKey := seedAIGatewayOpenAITestDependencies(t, db, openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + creator := newTestServer(t, db, ps, uuid.New()) + chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "aigateway-workspace-context", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + APIKeyID: apiKey.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("use the workspace context"), + }, + }) + require.NoError(t, err) + + const contextText = "# Project instructions\nAlways keep routing metadata." + _ = newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(factory) + cfg.AIGatewayRoutingEnabled = true + cfg.AllowBYOK = true + cfg.AllowBYOKSet = true + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupWorkspaceContextAgentConn(t, mockConn, dbAgent, contextText, nil) + return mockConn, func() {}, nil + } + }) + + chatResult := waitForTerminalChat(ctx, t, db, chat.ID) + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + require.False(t, chatResult.LastError.Valid) + + messages := persistedChatMessages(ctx, t, db, chat.ID) + var contextMessages []database.ChatMessage + for _, msg := range messages { + if msg.Role != database.ChatMessageRoleUser || + msg.Visibility != database.ChatMessageVisibilityBoth { + continue + } + for _, part := range mustParseChatParts(t, msg) { + if part.Type == codersdk.ChatMessagePartTypeContextFile && + part.ContextFileAgentID.Valid && + part.ContextFileAgentID.UUID == dbAgent.ID { + contextMessages = append(contextMessages, msg) + } + } + } + require.Len(t, contextMessages, 1) + require.True(t, contextMessages[0].APIKeyID.Valid) + require.Equal(t, apiKey.ID, contextMessages[0].APIKeyID.String) + + requests := factory.requestsSnapshot() + require.NotEmpty(t, requests) + for _, req := range requests { + require.Equal(t, provider.Name, req.ProviderName) + require.Equal(t, aibridge.SourceAgents, req.Source) + require.Equal(t, apiKey.ID, req.APIKeyID) + require.Equal(t, "Bearer sk-user-aibridge", req.Authorization) + require.Equal(t, "delegated", req.CoderToken) + } +} + func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { t.Parallel() @@ -7544,6 +9892,7 @@ func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { chat, err := creator.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "user-provider-key-missing", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -7552,15 +9901,8 @@ func TestProcessChat_UserProviderKey_MissingKeyError(t *testing.T) { }) require.NoError(t, err) - _, events, cancel, ok := creator.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - _ = newActiveTestServer(t, db, ps) - terminalStatus := waitForTerminalChatStatusEvent(ctx, t, events) - require.Equal(t, codersdk.ChatStatusError, terminalStatus) - chatResult := waitForTerminalChat(ctx, t, db, chat.ID) require.Equal(t, database.ChatStatusError, chatResult.Status) persistedError := requireChatLastErrorPayload(t, chatResult.LastError) @@ -7584,10 +9926,19 @@ func TestProcessChatPanicRecovery(t *testing.T) { // the processChat goroutine. panicWrapper := &panicOnInTxDB{Store: db} + firstOpenAICallStarted := make(chan struct{}) + continueFirstOpenAICall := make(chan struct{}) + var openAICallCount atomic.Int32 openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("Panic recovery test") } + + if openAICallCount.Add(1) == 1 { + close(firstOpenAICallStarted) + <-continueFirstOpenAICall + } + return chattest.OpenAIStreamingResponse( chattest.OpenAITextChunks("hello")..., ) @@ -7603,6 +9954,7 @@ func TestProcessChatPanicRecovery(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "panic-recovery", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -7611,13 +9963,14 @@ func TestProcessChatPanicRecovery(t *testing.T) { }) require.NoError(t, err) - // Enable the panic now that CreateChat's InTx has completed. - // The next InTx call is PersistStep inside the chatloop, - // running synchronously on the processChat goroutine. + testutil.TryReceive(ctx, t, firstOpenAICallStarted) + + // Enable the panic while the first provider call is blocked. The next InTx + // call is PersistStep inside the chatloop, running synchronously on the + // processChat goroutine after the provider returns. panicWrapper.enablePanic() + close(continueFirstOpenAICall) - // Wait for the panic to be recovered and the chat to - // transition to error status. var chatResult database.Chat require.Eventually(t, func() bool { got, getErr := db.GetChatByID(ctx, chat.ID) @@ -7625,13 +9978,31 @@ func TestProcessChatPanicRecovery(t *testing.T) { return false } chatResult = got - return got.Status == database.ChatStatusError + return got.Status == database.ChatStatusWaiting }, testutil.WaitLong, testutil.IntervalFast) + require.Equal(t, int32(2), openAICallCount.Load()) - persistedError := requireChatLastErrorPayload(t, chatResult.LastError) - require.Contains(t, persistedError.Message, "chat processing panicked") - require.Contains(t, persistedError.Message, "intentional test panic") - require.Equal(t, codersdk.ChatErrorKindGeneric, persistedError.Kind) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: 0, + }) + require.NoError(t, err) + + var assistantText string + for _, message := range messages { + if message.Role != database.ChatMessageRoleAssistant { + continue + } + parts, parseErr := chatprompt.ParseContent(message) + require.NoError(t, parseErr) + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText { + assistantText += part.Text + } + } + } + require.Equal(t, "hello", assistantText) + require.False(t, chatResult.LastError.Valid) } // panicOnInTxDB wraps a database.Store and panics on the first InTx @@ -7768,6 +10139,7 @@ func TestMCPServerToolInvocation(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "mcp-tool-test", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -7924,6 +10296,7 @@ func TestPlanModeRootChatApprovedExternalMCPToolInvocation(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-mode-mcp-invocation", ModelConfigID: model.ID, PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, @@ -8053,6 +10426,7 @@ func TestPlanModeRootChatApprovedExternalMCPWorkflowCanReachProposePlan(t *testi chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-mode-mcp-propose-plan", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -8273,6 +10647,7 @@ func TestMCPServerOAuth2TokenRefresh(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "oauth2-refresh-test", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -8382,6 +10757,7 @@ func TestMCPServerOAuth2TokenRefreshFailureGraceful(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "graceful-degradation-test", ModelConfigID: model.ID, MCPServerIDs: []uuid.UUID{mcpConfig.ID}, @@ -8506,6 +10882,7 @@ func TestChatTemplateAllowlistEnforcement(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "allowlist-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -8620,6 +10997,7 @@ func TestSignalWakeImmediateAcquisition(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "wake-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -8686,6 +11064,7 @@ func TestSignalWakeSendMessage(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "wake-send-test", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("first")}, @@ -8701,8 +11080,9 @@ func TestSignalWakeSendMessage(t *testing.T) { // Now send a follow-up message, which should also be // processed immediately via signalWake. _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("second")}, + ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("second")}, }) require.NoError(t, err) @@ -8873,71 +11253,6 @@ func TestAgentContextFilesAndSkillsLoadedIntoChat(t *testing.T) { "plan-file-path block should be part of the main system prompt, not a standalone message") } -func TestSendMessageRejectsArchivedChat(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "send-archived", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - _, err = replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("should fail")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.ErrorIs(t, err, chatd.ErrChatArchived) -} - -func TestEditMessageRejectsArchivedChat(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "edit-archived", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}, - }) - require.NoError(t, err) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - require.Len(t, messages, 1) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ - ChatID: chat.ID, - EditedMessageID: messages[0].ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, - }) - require.ErrorIs(t, err, chatd.ErrChatArchived) -} - // TestEditMessageWithModelConfigOverride verifies that callers can // change the model when editing a previous user message. The // replacement message must persist with the new model and the chat's @@ -8962,6 +11277,7 @@ func TestEditMessageWithModelConfigOverride(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), OrganizationID: org.ID, Title: "edit-with-model-override", ModelConfigID: modelA.ID, @@ -8979,6 +11295,7 @@ func TestEditMessageWithModelConfigOverride(t *testing.T) { result, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: initial[0].ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, ModelConfigID: modelB.ID, @@ -9007,6 +11324,7 @@ func TestEditMessagePreservesModelConfigByDefault(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), OrganizationID: org.ID, Title: "edit-preserves-model", ModelConfigID: modelA.ID, @@ -9023,6 +11341,7 @@ func TestEditMessagePreservesModelConfigByDefault(t *testing.T) { result, err := replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: initial[0].ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, }) @@ -9050,6 +11369,7 @@ func TestEditMessageRejectsUnknownModelConfig(t *testing.T) { chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), OrganizationID: org.ID, Title: "edit-unknown-model", ModelConfigID: modelA.ID, @@ -9066,6 +11386,7 @@ func TestEditMessageRejectsUnknownModelConfig(t *testing.T) { _, err = replica.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), EditedMessageID: initial[0].ID, Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, ModelConfigID: uuid.New(), @@ -9087,521 +11408,7 @@ func TestEditMessageRejectsUnknownModelConfig(t *testing.T) { require.Equal(t, modelA.ID, storedChat.LastModelConfigID) } -func TestPromoteQueuedRejectsArchivedChat(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "promote-archived", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // Queue a message by setting the chat to running first. - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - queuedResult, err := replica.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - - // Move back to waiting, then archive. - chat, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - WorkerID: uuid.NullUUID{}, - StartedAt: sql.NullTime{}, - HeartbeatAt: sql.NullTime{}, - }) - require.NoError(t, err) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - _, err = replica.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedResult.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.ErrorIs(t, err, chatd.ErrChatArchived) -} - -// TestPromoteQueuedWhileRequiresAction guards against the -// stops-dead failure mode: promoting on requires_action without -// closing pending dynamic tool calls leaves the assistant turn -// with unresolved tool_call parts that the LLM API rejects. It -// also asserts the synthetic tool-result row is published to live -// SSE subscribers before the promoted user message. -func TestPromoteQueuedWhileRequiresAction(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - var streamedCallCount atomic.Int32 - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("requires-action-promote") - } - if streamedCallCount.Add(1) == 1 { - return chattest.OpenAIStreamingResponse( - chattest.OpenAIToolCallChunk( - "my_dynamic_tool", - `{"input":"hello"}`, - ), - ) - } - // Second call: the resumed run after promote completes. - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("Resumed after promotion.")..., - ) - }) - - user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - server := newActiveTestServer(t, db, ps) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{ - Type: "object", - Properties: map[string]any{ - "input": map[string]any{"type": "string"}, - }, - Required: []string{"input"}, - }, - }}) - require.NoError(t, err) - - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "promote-while-requires-action", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{ - codersdk.ChatMessageText("Please call the dynamic tool."), - }, - DynamicTools: dynamicToolsJSON, - }) - require.NoError(t, err) - - var chatBeforePromote database.Chat - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - got, getErr := db.GetChatByID(ctx, chat.ID) - if getErr != nil { - return false - } - chatBeforePromote = got - return got.Status == database.ChatStatusRequiresAction || - got.Status == database.ChatStatusError - }, testutil.IntervalFast) - require.Equal(t, database.ChatStatusRequiresAction, chatBeforePromote.Status, - "expected requires_action, got %s (last_error=%q)", - chatBeforePromote.Status, chatLastErrorMessage(chatBeforePromote.LastError)) - - var pendingToolCallID string - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if dbErr != nil { - return false - } - for _, msg := range messages { - if msg.Role != database.ChatMessageRoleAssistant { - continue - } - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - continue - } - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == "my_dynamic_tool" { - pendingToolCallID = part.ToolCallID - return true - } - } - } - return false - }, testutil.IntervalFast) - require.NotEmpty(t, pendingToolCallID, "expected pending dynamic tool call") - - queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("promote me")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - require.NotNil(t, queuedResult.QueuedMessage) - - // Subscribe before promoting to capture published events. - _, events, subCancel, ok := server.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer subCancel() - promoteResult, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedResult.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.Equal(t, database.ChatMessageRoleUser, promoteResult.PromotedMessage.Role) - - // Synthetic row must publish before the promoted user message. - var ( - syntheticPublishedAt int - userPublishedAt int - messagesSeen int - ) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - for { - select { - case ev := <-events: - if ev.Type != codersdk.ChatStreamEventTypeMessage || ev.Message == nil { - continue - } - messagesSeen++ - switch ev.Message.Role { - case codersdk.ChatMessageRoleTool: - if syntheticPublishedAt == 0 { - syntheticPublishedAt = messagesSeen - } - case codersdk.ChatMessageRoleUser: - if ev.Message.ID == promoteResult.PromotedMessage.ID { - userPublishedAt = messagesSeen - } - } - if syntheticPublishedAt > 0 && userPublishedAt > 0 { - return true - } - default: - return false - } - } - }, testutil.IntervalFast) - - require.Less(t, syntheticPublishedAt, userPublishedAt, - "synthetic tool-result must be published before the promoted user message") - - queuedAfter, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Empty(t, queuedAfter, "queued message should be removed after sync promotion") - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var ( - syntheticToolResult *database.ChatMessage - promotedUserMessage *database.ChatMessage - ) - for i := range messages { - msg := messages[i] - if msg.Role == database.ChatMessageRoleTool { - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeToolResult { - continue - } - if part.ToolCallID != pendingToolCallID { - continue - } - require.True(t, part.IsError, - "synthetic tool result should have IsError=true") - syntheticToolResult = &messages[i] - } - } - if msg.ID == promoteResult.PromotedMessage.ID { - promotedUserMessage = &messages[i] - } - } - require.NotNil(t, syntheticToolResult, - "expected a synthetic error tool result for the pending tool call") - require.NotNil(t, promotedUserMessage) - require.Less(t, syntheticToolResult.ID, promotedUserMessage.ID, - "synthetic tool result must precede the promoted user message") - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - got, getErr := db.GetChatByID(ctx, chat.ID) - if getErr != nil { - return false - } - return got.Status == database.ChatStatusWaiting || got.Status == database.ChatStatusError - }, testutil.IntervalFast) - final, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, final.Status, - "chat should resume to waiting after promotion (last_error=%q)", - chatLastErrorMessage(final.LastError)) -} - // TestPromoteQueuedWhileRequiresActionMixedTools guards against -// duplicating already-resolved built-in tool results: synthetic -// results must be scoped to dynamic tool names only. -func TestPromoteQueuedWhileRequiresActionMixedTools(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - var streamedCallCount atomic.Int32 - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("mixed-tools-promote") - } - if streamedCallCount.Add(1) == 1 { - builtinChunk := chattest.OpenAIToolCallChunk( - "read_file", - `{"path":"/tmp/test.txt"}`, - ) - dynamicChunk := chattest.OpenAIToolCallChunk( - "my_dynamic_tool", - `{"input":"hello world"}`, - ) - mergedChunk := builtinChunk - dynCall := dynamicChunk.Choices[0].ToolCalls[0] - dynCall.Index = 1 - mergedChunk.Choices[0].ToolCalls = append( - mergedChunk.Choices[0].ToolCalls, - dynCall, - ) - return chattest.OpenAIStreamingResponse(mergedChunk) - } - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("Resumed after mixed-tool promotion.")..., - ) - }) - - user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - server := newActiveTestServer(t, db, ps) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{ - Type: "object", - Properties: map[string]any{ - "input": map[string]any{"type": "string"}, - }, - Required: []string{"input"}, - }, - }}) - require.NoError(t, err) - - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "promote-while-requires-action-mixed", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{ - codersdk.ChatMessageText("Call both tools."), - }, - DynamicTools: dynamicToolsJSON, - }) - require.NoError(t, err) - - var chatBeforePromote database.Chat - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - got, getErr := db.GetChatByID(ctx, chat.ID) - if getErr != nil { - return false - } - chatBeforePromote = got - return got.Status == database.ChatStatusRequiresAction || - got.Status == database.ChatStatusError - }, testutil.IntervalFast) - require.Equal(t, database.ChatStatusRequiresAction, chatBeforePromote.Status, - "expected requires_action, got %s (last_error=%q)", - chatBeforePromote.Status, chatLastErrorMessage(chatBeforePromote.LastError)) - - // The built-in tool resolves before requires_action; capture - // its row ID to assert the dynamic synthetic comes after. - var ( - dynamicToolCallID string - builtinToolResultID int64 - builtinToolResultSeen bool - ) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - messages, dbErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if dbErr != nil { - return false - } - for _, msg := range messages { - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - continue - } - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolName == "read_file" { - builtinToolResultID = msg.ID - builtinToolResultSeen = true - } - if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolName == "my_dynamic_tool" { - dynamicToolCallID = part.ToolCallID - } - } - } - return builtinToolResultSeen && dynamicToolCallID != "" - }, testutil.IntervalFast) - require.NotEmpty(t, dynamicToolCallID) - require.NotZero(t, builtinToolResultID) - - queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("promote me")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - require.NotNil(t, queuedResult.QueuedMessage) - - _, events, subCancel, ok := server.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer subCancel() - promoteResult, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedResult.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.NotZero(t, promoteResult.PromotedMessage.ID, - "requires_action promotion is synchronous and returns the inserted message") - - // Only the dynamic tool's synth row publishes; the built-in's - // pre-existing result is not republished. - var ( - syntheticPublishCount int - userPublished bool - ) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - for { - select { - case ev := <-events: - if ev.Type != codersdk.ChatStreamEventTypeMessage || ev.Message == nil { - t.Logf("subscriber consumed non-message event type=%s", ev.Type) - continue - } - t.Logf("subscriber consumed message id=%d role=%s match_promoted=%t", ev.Message.ID, ev.Message.Role, ev.Message.ID == promoteResult.PromotedMessage.ID) - switch ev.Message.Role { - case codersdk.ChatMessageRoleTool: - syntheticPublishCount++ - case codersdk.ChatMessageRoleUser: - if ev.Message.ID == promoteResult.PromotedMessage.ID { - userPublished = true - } - } - if userPublished { - return true - } - default: - return false - } - } - }, testutil.IntervalFast) - - require.Equal(t, 1, syntheticPublishCount, - "only the dynamic tool's synthetic result must be published; the built-in's pre-existing result must not be republished") - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var ( - dynamicSyntheticCount int - builtinResultsForReadFile int - ) - for _, msg := range messages { - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeToolResult { - continue - } - switch part.ToolName { - case "read_file": - builtinResultsForReadFile++ - case "my_dynamic_tool": - if part.IsError && part.ToolCallID == dynamicToolCallID && msg.ID > builtinToolResultID { - dynamicSyntheticCount++ - } - } - } - } - require.Equal(t, 1, dynamicSyntheticCount, - "expected exactly one synthetic error tool result for the dynamic tool call") - require.Equal(t, 1, builtinResultsForReadFile, - "built-in tool result should not be duplicated by promotion") - - require.Greater(t, promoteResult.PromotedMessage.ID, builtinToolResultID) -} - -func TestSubmitToolResultsRejectsArchivedChat(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - replica := newTestServer(t, db, ps, uuid.New()) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat, err := replica.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "submit-tool-archived", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - err = replica.ArchiveChat(ctx, chat) - require.NoError(t, err) - - // Set requires_action so the test exercises a realistic - // scenario where SubmitToolResults would be called. - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRequiresAction, - }) - require.NoError(t, err) - - err = replica.SubmitToolResults(ctx, chatd.SubmitToolResultsOptions{ - ChatID: chat.ID, - UserID: user.ID, - ModelConfigID: model.ID, - Results: []codersdk.ToolResult{{ - ToolCallID: "fake-tool-call-id", - Output: json.RawMessage(`{"result":"ignored"}`), - }}, - }) - require.ErrorIs(t, err, chatd.ErrChatArchived) -} - func TestAcquireChatsSkipsArchivedPendingChat(t *testing.T) { t.Parallel() @@ -9689,6 +11496,7 @@ func TestAdvisorGating_Disabled(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-disabled", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -9786,6 +11594,7 @@ func TestAdvisorGating_RootChat(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-root", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -9873,6 +11682,19 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { finalCallMessages []chattest.OpenAIMessage ) + // Declared before the OpenAI handler so the nested advisor stream can + // gate its completion on the live collector below having observed the + // streamed deltas. + var ( + livePartsMu sync.Mutex + liveAdvisorDeltas []string + ) + liveDeltasCaptured := func() bool { + livePartsMu.Lock() + defer livePartsMu.Unlock() + return slices.Equal(advisorDeltas, liveAdvisorDeltas) + } + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("title") @@ -9881,10 +11703,12 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { switch streamedCallCount.Add(1) { case 1: // Parent turn 1: call advisor solo. - return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk( + chunk := chattest.OpenAIToolCallChunk( "advisor", `{"question":"how should I approach this refactor?"}`, - )) + ) + chunk.Choices[0].ToolCalls[0].ID = "advisor-happy-path-call" + return chattest.OpenAIStreamingResponse(chunk) case 2: // Nested advisor turn. The nested call has no tools because // chatadvisor.RunAdvisor runs with MaxSteps=1 and no tool @@ -9895,9 +11719,33 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { advisorMessages = append([]chattest.OpenAIMessage(nil), req.Messages...) streamedCallsMu.Unlock() advisorCallSeen.Store(true) - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks(advisorDeltas...)..., - ) + // Stream the deltas, then hold the nested response open until + // the live subscriber has captured them. Advisor deltas are + // stream-only: they live in the generation attempt's message + // part episode, and once the tool result is committed the + // subscriber's stream loop targets the next episode and never + // replays this one. Without the hold, a slow pubsub sync makes + // the subscriber skip the episode entirely and the deltas are + // lost, flaking the streaming assertions below. + chunks := make(chan chattest.OpenAIChunk) + go func() { + defer close(chunks) + for _, chunk := range chattest.OpenAITextChunks(advisorDeltas...) { + chunks <- chunk + } + deadline := time.NewTimer(testutil.WaitLong) + defer deadline.Stop() + for !liveDeltasCaptured() { + select { + case <-deadline.C: + // Give up and let the assertions below report the + // failure instead of hanging the stream forever. + return + case <-time.After(testutil.IntervalFast): + } + } + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} default: // Parent turn 2: observe the advisor tool result and close // out with a final text reply. @@ -9921,6 +11769,7 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-happy-path", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -9932,11 +11781,7 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { // Advisor deltas are transient; a late subscriber misses them. _, liveEvents, cancelLive, ok := server.Subscribe(ctx, chat.ID, nil, 0) require.True(t, ok) - var ( - livePartsMu sync.Mutex - liveAdvisorDeltas []string - liveCollectorDone = make(chan struct{}) - ) + liveCollectorDone := make(chan struct{}) go func() { defer close(liveCollectorDone) for { @@ -9955,6 +11800,7 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { if event.MessagePart.Role != codersdk.ChatMessageRoleTool || part.Type != codersdk.ChatMessagePartTypeToolResult || part.ToolName != chatadvisor.ToolName || + part.ToolCallID != "advisor-happy-path-call" || part.ResultDelta == "" { continue } @@ -10021,15 +11867,19 @@ func TestAdvisorHappyPath_RootChat(t *testing.T) { require.True(t, parentSawAdvisorResult, "parent must see the advisor reply in its continuation call") - require.EventuallyWithT(t, func(c *assert.CollectT) { - livePartsMu.Lock() - defer livePartsMu.Unlock() - assert.Equal(c, advisorDeltas, liveAdvisorDeltas, - "advisor nested text deltas must stream into the parent tool card") - }, testutil.WaitLong, testutil.IntervalFast) - + // Stop the live collector and assert it captured the streaming + // advisor deltas during processing. Late subscribers no longer + // see committed parts because publishMessage claims them out of + // new snapshots, so the assertion must use the live collector. + require.Eventually(t, liveDeltasCaptured, testutil.WaitLong, testutil.IntervalFast, + "advisor nested text deltas must stream into the parent tool card") cancelLive() <-liveCollectorDone + livePartsMu.Lock() + collectedAdvisorDeltas := append([]string(nil), liveAdvisorDeltas...) + livePartsMu.Unlock() + require.Equal(t, advisorDeltas, collectedAdvisorDeltas, + "advisor nested text deltas must stream into the parent tool card") persisted, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: chat.ID, @@ -10109,6 +11959,7 @@ func TestAdvisorGating_ChildChat(t *testing.T) { childChat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-child", ModelConfigID: model.ID, ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, @@ -10187,6 +12038,7 @@ func TestAdvisorGating_PlanMode(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-plan-mode", ModelConfigID: model.ID, PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true}, @@ -10266,6 +12118,7 @@ func TestAdvisorGating_ExploreSubagent(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-explore", ModelConfigID: model.ID, ChatMode: database.NullChatMode{ @@ -10305,9 +12158,9 @@ func TestAdvisorGating_ExploreSubagent(t *testing.T) { // runtime together with chain mode and asserts the snapshot captured for // the nested advisor call retains the full pre-chain prompt. Chain mode // otherwise strips assistant and tool turns from the prompt the outer -// loop sees, so a regression that moves setAdvisorPromptSnapshot behind -// filterPromptForChainMode, or drops the !chainModeActive guards in -// PrepareMessages, would leak the filtered view into the advisor's +// loop sees, so a regression that captures the advisor snapshot after +// filterPromptForChainMode, or removes the chain-mode guard around +// advisor snapshotting, would leak the filtered view into the advisor's // nested call. The advisor would then only see the trailing user // message, losing the context the outer model had been building on. func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) { @@ -10420,6 +12273,7 @@ func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "advisor-chain-mode", ModelConfigID: responsesModel.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -10439,6 +12293,7 @@ func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) { _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Content: []codersdk.ChatMessagePart{ codersdk.ChatMessageText(turn2User), }, @@ -10510,920 +12365,6 @@ func seedAdvisorConfig( require.NoError(t, err) } -// TestPromoteQueuedWhileRunning guards against the data-loss -// failure mode: promoting on a streaming chat must preserve -// partial assistant output by deferring the user-message insert -// to the worker's auto-promote. -func TestPromoteQueuedWhileRunning(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - streamStarted := make(chan struct{}) - streamCanceled := make(chan struct{}) - var streamCallCount atomic.Int32 - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("running-promote") - } - if streamCallCount.Add(1) > 1 { - // Subsequent calls are the resumed run; let it settle. - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("resumed after promotion")..., - ) - } - chunks := make(chan chattest.OpenAIChunk, 1) - go func() { - defer close(chunks) - chunks <- chattest.OpenAITextChunks("partial-running-output")[0] - select { - case <-streamStarted: - default: - close(streamStarted) - } - <-req.Context().Done() - select { - case <-streamCanceled: - default: - close(streamCanceled) - } - }() - return chattest.OpenAIResponse{StreamingChunks: chunks} - }) - - server := newActiveTestServer(t, db, ps) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "promote-while-running", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusRunning && fromDB.WorkerID.Valid - }, testutil.IntervalFast) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case <-streamStarted: - return true - default: - return false - } - }, testutil.IntervalFast) - - queuedResult, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("promote me")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queuedResult.Queued) - require.NotNil(t, queuedResult.QueuedMessage) - - promoteResult, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queuedResult.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - // Deferred promotion: no synchronous user message. - require.Zero(t, promoteResult.PromotedMessage.ID) - - // Worker observes waiting and cancels. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case <-streamCanceled: - return true - default: - return false - } - }, testutil.IntervalFast) - - // Partial assistant output is preserved (not lost as it was - // pre-fix) and precedes the promoted user message. Poll on the - // messages themselves: the status passes through Waiting - // transiently before finishActiveChat's external-Waiting case - // promotes the queued message and flips the chat to Pending. - // Both messages being persisted implies cleanup completed. - var ( - partialAssistantID int64 - promotedUserID int64 - ) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if err != nil { - return false - } - var ( - assistantID int64 - userID int64 - ) - for _, msg := range messages { - switch msg.Role { - case database.ChatMessageRoleAssistant: - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - continue - } - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeText && strings.Contains(part.Text, "partial-running-output") { - assistantID = msg.ID - } - } - case database.ChatMessageRoleUser: - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - continue - } - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeText && strings.Contains(part.Text, "promote me") { - userID = msg.ID - } - } - } - } - if assistantID == 0 || userID == 0 { - return false - } - partialAssistantID = assistantID - promotedUserID = userID - return true - }, testutil.IntervalFast) - require.Less(t, partialAssistantID, promotedUserID, - "promoted user message must follow the persisted partial output") -} - -// TestPromoteQueuedWhileRunningRespectsMessageOrder guards -// against losing or reshuffling sibling queued messages when one -// is promoted out-of-order. -func TestPromoteQueuedWhileRunningRespectsMessageOrder(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - streamStarted := make(chan struct{}) - var streamCallCount atomic.Int32 - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("running-promote-order") - } - if streamCallCount.Add(1) > 1 { - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("resumed")..., - ) - } - chunks := make(chan chattest.OpenAIChunk, 1) - go func() { - defer close(chunks) - chunks <- chattest.OpenAITextChunks("partial")[0] - select { - case <-streamStarted: - default: - close(streamStarted) - } - <-req.Context().Done() - }() - return chattest.OpenAIResponse{StreamingChunks: chunks} - }) - - server := newActiveTestServer(t, db, ps) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - chat, err := server.CreateChat(ctx, chatd.CreateOptions{ - OwnerID: user.ID, - OrganizationID: org.ID, - Title: "promote-while-running-order", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusRunning && fromDB.WorkerID.Valid - }, testutil.IntervalFast) - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - select { - case <-streamStarted: - return true - default: - return false - } - }, testutil.IntervalFast) - - queueA, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("A")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.NotNil(t, queueA.QueuedMessage) - queueB, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("B")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.NotNil(t, queueB.QueuedMessage) - queueC, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("C")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.NotNil(t, queueC.QueuedMessage) - - promoteResult, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queueB.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.Zero(t, promoteResult.PromotedMessage.ID, - "running-case promotion is deferred to auto-promote") - - // Wait for the worker to drain all three queued messages into - // chat history, then verify ordering. Reading queue state right - // after PromoteQueued races the worker's auto-promote pipeline - // (TOCTOU), so we wait for the settled outcome instead. - var posB, posA, posC int - var foundA, foundB, foundC bool - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - messages, getErr := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - if getErr != nil { - return false - } - foundA, foundB, foundC = false, false, false - for i, msg := range messages { - if msg.Role != database.ChatMessageRoleUser { - continue - } - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - return false - } - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeText { - continue - } - // Only A, B, C are tracked; other user messages are ignored. - switch part.Text { - case "A": - posA = i - foundA = true - case "B": - posB = i - foundB = true - case "C": - posC = i - foundC = true - } - } - } - return foundA && foundB && foundC - }, testutil.IntervalFast, - "queued messages not found in chat history: foundA=%v, foundB=%v, foundC=%v", foundA, foundB, foundC) - - // PromoteQueued reorders the queue to [B, A, C], so the worker - // processes B first, then A, then C. Verify that ordering. - require.Less(t, posB, posA, - "promoted message B must appear before A in history") - require.Less(t, posA, posC, - "non-promoted messages must preserve relative order (A before C)") -} - -// TestFinishActiveChatExternalWaitingInsertsSyntheticResults -// asserts the cleanup TX inserts synthetic tool-result rows when -// PromoteQueued's deferred path set Status=Waiting while the -// worker concluded with RequiresAction. Without it, the next -// chatloop run would feed the LLM an assistant turn with -// unresolved tool_call parts and the API would reject it. -func TestFinishActiveChatExternalWaitingInsertsSyntheticResults(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - server := newActiveTestServer(t, db, ps) - user, org, model := seedChatDependencies(t, db) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{ - Type: "object", - Properties: map[string]any{}, - }, - }}) - require.NoError(t, err) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "external-waiting-stops-dead-guard", - LastModelConfigID: model.ID, - DynamicTools: nullRawMessage(dynamicToolsJSON), - }) - require.NoError(t, err) - - // Seed a user message and an assistant message with an - // unresolved dynamic tool call. This mirrors what the worker - // would have persisted before the deferred promote arrived. - insertUserTextMessage(t, db, chat.ID, user.ID, model.ID, "user input") - - pendingCallID := "call_pending_dynamic" - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - { - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: pendingCallID, - ToolName: "my_dynamic_tool", - Args: json.RawMessage(`{}`), - }, - }) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, - }) - require.NoError(t, err) - - // Queue a message and put the chat in the post-promote - // Waiting state (no worker, queue at front). - queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("queued-after-promote"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent.RawMessage, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - }) - require.NoError(t, err) - - // Refresh chat with current status (Waiting, no worker). - latestChat, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - - // Drive the cleanup path with the local-RequiresAction outcome. - updated, promoted, syntheticToolResults, finishErr := chatd.FinishActiveChatForTest( - ctx, server, latestChat, database.ChatStatusRequiresAction, "", - ) - require.NoError(t, finishErr) - require.NotNil(t, promoted, "queued message must be auto-promoted into history") - require.Equal(t, database.ChatStatusPending, updated.Status, - "chat must end Pending so the run loop picks it up") - require.Len(t, syntheticToolResults, 1, - "cleanup TX must return the inserted synthetic tool-result row so the post-TX caller can publish it") - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var ( - assistantIdx = -1 - synthToolIdx = -1 - promotedUserIdx = -1 - ) - for i, msg := range messages { - switch msg.Role { - case database.ChatMessageRoleAssistant: - assistantIdx = i - case database.ChatMessageRoleTool: - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolResult && - part.ToolCallID == pendingCallID && part.IsError { - synthToolIdx = i - } - } - case database.ChatMessageRoleUser: - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeText && - part.Text == "queued-after-promote" { - promotedUserIdx = i - } - } - } - } - require.NotEqual(t, -1, assistantIdx, "assistant tool-call message present") - require.NotEqual(t, -1, synthToolIdx, - "synthetic tool result for the unresolved dynamic tool call must be inserted") - require.NotEqual(t, -1, promotedUserIdx, - "promoted queued message must be inserted as a user message") - require.Less(t, assistantIdx, synthToolIdx, - "synthetic tool result must follow the assistant message") - require.Less(t, synthToolIdx, promotedUserIdx, - "promoted user message must follow the synthetic tool result") -} - -// TestPromoteQueuedFallsThroughOnStaleHeartbeat asserts a stale -// heartbeat takes the synchronous path so the chat does not strand -// in Waiting waiting on a worker that will not return. -func TestPromoteQueuedFallsThroughOnStaleHeartbeat(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - staleAfter := 100 * time.Millisecond - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) - t.Cleanup(func() { require.NoError(t, server.Close()) }) - - user, org, model := seedChatDependencies(t, db) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "stale-heartbeat-promote-fallthrough", - LastModelConfigID: model.ID, - }) - require.NoError(t, err) - - // Place the chat in Running with a stale heartbeat. We do not - // start the server's run loop, so no worker will ever pick this - // chat up; the test isolates the fall-through decision in - // PromoteQueued. - deadWorker := uuid.New() - staleTime := time.Now().Add(-2 * staleAfter) - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: deadWorker, Valid: true}, - StartedAt: sql.NullTime{Time: staleTime, Valid: true}, - HeartbeatAt: sql.NullTime{Time: staleTime, Valid: true}, - }) - require.NoError(t, err) - - queued, err := server.SendMessage(ctx, chatd.SendMessageOptions{ - ChatID: chat.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("promote me")}, - BusyBehavior: chatd.SendMessageBusyBehaviorQueue, - }) - require.NoError(t, err) - require.True(t, queued.Queued) - require.NotNil(t, queued.QueuedMessage) - - result, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ - ChatID: chat.ID, - QueuedMessageID: queued.QueuedMessage.ID, - CreatedBy: user.ID, - }) - require.NoError(t, err) - require.NotZero(t, result.PromotedMessage.ID, - "stale heartbeat must take the synchronous path and insert a user message inline") - - got, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusPending, got.Status, - "synchronous promote ends Pending") - require.False(t, got.WorkerID.Valid, - "worker_id is cleared by the synchronous promote") -} - -// TestRecoverStaleChatsRecoversWaitingWithQueue asserts a Waiting -// chat with a non-empty queue and stale updated_at gets recovered -// to Pending, closing the post-promote-stranding hole. -func TestRecoverStaleChatsRecoversWaitingWithQueue(t *testing.T) { - t.Parallel() - - db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - staleAfter := 100 * time.Millisecond - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) - t.Cleanup(func() { require.NoError(t, server.Close()) }) - user, org, model := seedChatDependencies(t, db) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "stale-waiting-with-queue", - LastModelConfigID: model.ID, - }) - require.NoError(t, err) - - queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("queued-stranded"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent.RawMessage, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - }) - require.NoError(t, err) - // Backdate updated_at directly so the chat is past the stale - // threshold without sleeping. - _, err = rawDB.ExecContext(ctx, - "UPDATE chats SET updated_at = $1 WHERE id = $2", - time.Now().Add(-time.Hour), chat.ID) - require.NoError(t, err) - - chatd.RecoverStaleChatsForTest(ctx, server) - - got, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusPending, got.Status, - "stale-recovery must promote the front-of-queue and set Pending") - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - var foundPromoted bool - for _, msg := range messages { - if msg.Role != database.ChatMessageRoleUser { - continue - } - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeText && - part.Text == "queued-stranded" { - foundPromoted = true - } - } - } - require.True(t, foundPromoted, - "the front-of-queue message must be promoted into history") - - remaining, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Empty(t, remaining, - "the queue is drained after the recovery promotes its only entry") -} - -// TestRecoverStaleChatsWaitingWithUnresolvedToolCallInsertsSyntheticResults -// asserts stale recovery closes pending dynamic tool calls before -// promoting, so the recovery path does not stop the chat dead by -// feeding the LLM unresolved tool_call parts. -func TestRecoverStaleChatsWaitingWithUnresolvedToolCallInsertsSyntheticResults(t *testing.T) { - t.Parallel() - - db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - staleAfter := 100 * time.Millisecond - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) - t.Cleanup(func() { require.NoError(t, server.Close()) }) - - user, org, model := seedChatDependencies(t, db) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{ - Type: "object", - Properties: map[string]any{}, - }, - }}) - require.NoError(t, err) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "stale-waiting-with-unresolved-tool-call", - LastModelConfigID: model.ID, - DynamicTools: nullRawMessage(dynamicToolsJSON), - }) - require.NoError(t, err) - - insertUserTextMessage(t, db, chat.ID, user.ID, model.ID, "please call the tool") - - pendingCallID := "call_unresolved_dynamic" - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - { - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: pendingCallID, - ToolName: "my_dynamic_tool", - Args: json.RawMessage(`{}`), - }, - }) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, - }) - require.NoError(t, err) - - queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("queued-after-crash"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent.RawMessage, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - }) - require.NoError(t, err) - - _, err = rawDB.ExecContext(ctx, - "UPDATE chats SET updated_at = $1 WHERE id = $2", - time.Now().Add(-time.Hour), chat.ID) - require.NoError(t, err) - - chatd.RecoverStaleChatsForTest(ctx, server) - - got, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusPending, got.Status) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var ( - assistantIdx = -1 - synthIdx = -1 - promotedUserIdx = -1 - ) - for i, msg := range messages { - switch msg.Role { - case database.ChatMessageRoleAssistant: - assistantIdx = i - case database.ChatMessageRoleTool: - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeToolResult && - part.ToolCallID == pendingCallID && part.IsError { - synthIdx = i - } - } - case database.ChatMessageRoleUser: - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type == codersdk.ChatMessagePartTypeText && - part.Text == "queued-after-crash" { - promotedUserIdx = i - } - } - } - } - require.NotEqual(t, -1, assistantIdx, "assistant tool-call message present") - require.NotEqual(t, -1, synthIdx, - "stale recovery must insert synthetic tool result for the unresolved dynamic tool call") - require.NotEqual(t, -1, promotedUserIdx, - "queued message must be promoted into history") - require.Less(t, assistantIdx, synthIdx) - require.Less(t, synthIdx, promotedUserIdx) -} - -// TestInsertSyntheticToolResultsTxSkipsAlreadyHandledCalls asserts -// the helper skips tool calls already handled (e.g. when a dynamic -// tool name collides with a built-in the chatloop dispatched). -// Without dedup the LLM would see two results for the same call ID. -func TestInsertSyntheticToolResultsTxSkipsAlreadyHandledCalls(t *testing.T) { - t.Parallel() - - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) - - user, org, model := seedChatDependencies(t, db) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{ - { - Name: "duplicate_call_tool", - Description: "Tool whose call already has a result.", - InputSchema: mcpgo.ToolInputSchema{Type: "object", Properties: map[string]any{}}, - }, - { - Name: "still_pending_tool", - Description: "Tool whose call has no result yet.", - InputSchema: mcpgo.ToolInputSchema{Type: "object", Properties: map[string]any{}}, - }, - }) - require.NoError(t, err) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusRequiresAction, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "synth-results-dedup", - LastModelConfigID: model.ID, - DynamicTools: nullRawMessage(dynamicToolsJSON), - }) - require.NoError(t, err) - - insertUserTextMessage(t, db, chat.ID, user.ID, model.ID, "please call both tools") - - handledCallID := "call_already_handled" - pendingCallID := "call_still_pending" - assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - { - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: handledCallID, - ToolName: "duplicate_call_tool", - Args: json.RawMessage(`{}`), - }, - { - Type: codersdk.ChatMessagePartTypeToolCall, - ToolCallID: pendingCallID, - ToolName: "still_pending_tool", - Args: json.RawMessage(`{}`), - }, - }) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(assistantContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, - }) - require.NoError(t, err) - - // Pre-insert a tool-result for the handled call ID. This - // simulates the chatloop having dispatched the colliding - // dynamic tool name as a built-in. - handledResultContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - { - Type: codersdk.ChatMessagePartTypeToolResult, - ToolCallID: handledCallID, - ToolName: "duplicate_call_tool", - Result: json.RawMessage(`"already done"`), - }, - }) - require.NoError(t, err) - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleTool}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Content: []string{string(handledResultContent.RawMessage)}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, - }) - require.NoError(t, err) - - chatRow, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - - _, err = chatd.InsertSyntheticToolResultsTxForTest( - ctx, db, chatRow, "synth reason", - ) - require.NoError(t, err) - - messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, - AfterID: 0, - }) - require.NoError(t, err) - - var ( - handledCount int - pendingCount int - syntheticForPending bool - ) - for _, msg := range messages { - if msg.Role != database.ChatMessageRoleTool { - continue - } - parts, parseErr := chatprompt.ParseContent(msg) - require.NoError(t, parseErr) - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeToolResult { - continue - } - switch part.ToolCallID { - case handledCallID: - handledCount++ - case pendingCallID: - pendingCount++ - if part.IsError { - syntheticForPending = true - } - } - } - } - require.Equal(t, 1, handledCount, - "handled call must keep exactly one tool result") - require.Equal(t, 1, pendingCount, - "pending call must get exactly one synthetic tool result") - require.True(t, syntheticForPending, - "the new tool result for the pending call must be marked IsError") -} - // nullRawMessage wraps raw JSON in a NullRawMessage. An empty input // becomes the zero value (Valid=false). func nullRawMessage(raw []byte) pqtype.NullRawMessage { @@ -11433,167 +12374,491 @@ func nullRawMessage(raw []byte) pqtype.NullRawMessage { return pqtype.NullRawMessage{RawMessage: raw, Valid: true} } -// TestInsertSyntheticToolResultsTxReturnsNilWhenNoAssistantMessage -// asserts the helper short-circuits cleanly when no assistant -// message exists yet, so a deferred promote racing a worker that -// fails before any persist does not roll back the cleanup TX. -func TestInsertSyntheticToolResultsTxReturnsNilWhenNoAssistantMessage(t *testing.T) { +// Regression for the cold-start race: chatd must wait long enough +// for ListMCPTools to return after the agent's MCP reload settles. +func TestActiveServer_WorkspaceContextAndDynamicToolInjection(t *testing.T) { t.Parallel() - db, _ := dbtestutil.NewDB(t) - ctx := testutil.Context(t, testutil.WaitLong) + t.Run("persists workspace context before provider request", func(t *testing.T) { + t.Parallel() - user, org, model := seedChatDependencies(t, db) + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{Type: "object", Properties: map[string]any{}}, - }}) - require.NoError(t, err) + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "no-assistant-message", - LastModelConfigID: model.ID, - DynamicTools: nullRawMessage(dynamicToolsJSON), + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + const contextText = "# Project instructions\nAlways write tests." + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupWorkspaceContextAgentConn(t, mockConn, dbAgent, contextText, nil) + return mockConn, func() {}, nil + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "workspace-context-before-provider", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("What are the workspace rules?"), + }, + }) + require.NoError(t, err) + + chatResult := waitForTerminalChat(ctx, t, db, chat.ID) + if chatResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(chatResult.LastError)) + } + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + + parts := persistedChatParts(ctx, t, db, chat.ID) + require.Len(t, contextFilePartsForAgent(parts, dbAgent.ID), 1) + contextPart := contextFilePartsForAgent(parts, dbAgent.ID)[0] + require.Equal(t, "/home/coder/project/AGENTS.md", contextPart.ContextFilePath) + require.Equal(t, contextText, contextPart.ContextFileContent) + require.Equal(t, "linux", contextPart.ContextFileOS) + require.Equal(t, "/home/coder/project", contextPart.ContextFileDirectory) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.Len(t, recorded, 1, "expected exactly one streamed model call") + require.True(t, requestHasSystemSubstring(recorded[0], "")) + require.True(t, requestHasSystemSubstring(recorded[0], contextText)) + require.True(t, requestHasSystemSubstring(recorded[0], "AGENTS.md")) }) - require.NoError(t, err) - // No assistant message persisted. The helper must return nil so - // the caller's transaction can still advance. - _, err = chatd.InsertSyntheticToolResultsTxForTest( - ctx, db, chat, "no assistant", - ) - require.NoError(t, err) + t.Run("persists workspace context once for the same agent", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + const contextText = "# Project instructions\nKeep it simple." + var contextConfigCalls atomic.Int32 + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupWorkspaceContextAgentConn(t, mockConn, dbAgent, contextText, &contextConfigCalls) + return mockConn, func() {}, nil + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "workspace-context-once", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("First turn."), + }, + }) + require.NoError(t, err) + firstResult := waitForTerminalChat(ctx, t, db, chat.ID) + if firstResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(firstResult.LastError)) + } + + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Content: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Second turn."), + }, + }) + require.NoError(t, err) + + secondResult := waitForTerminalChat(ctx, t, db, chat.ID) + if secondResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(secondResult.LastError)) + } + require.Equal(t, database.ChatStatusWaiting, secondResult.Status) + + parts := persistedChatParts(ctx, t, db, chat.ID) + require.Len(t, contextFilePartsForAgent(parts, dbAgent.ID), 1) + require.Equal(t, int32(1), contextConfigCalls.Load()) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.GreaterOrEqual(t, len(recorded), 2) + require.True(t, requestHasSystemSubstring(recorded[0], contextText)) + require.True(t, requestHasSystemSubstring(recorded[len(recorded)-1], contextText)) + }) + + t.Run("repersists workspace context after agent changes", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + requestsMu.Unlock() + + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, firstAgent := seedWorkspaceWithAgent(t, db, user.ID) + + oldContext := "# Old instructions\nUse the old agent." + newContext := "# New instructions\nUse the new agent." + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + switch agentID { + case firstAgent.ID: + setupWorkspaceContextAgentConn(t, mockConn, firstAgent, oldContext, nil) + default: + setupWorkspaceContextAgentConn(t, mockConn, database.WorkspaceAgent{ + ID: agentID, + OperatingSystem: "linux", + Directory: "/home/coder/project-new", + ExpandedDirectory: "/home/coder/project-new", + }, newContext, nil) + } + return mockConn, func() {}, nil + } + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "workspace-context-agent-change", + ModelConfigID: model.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("First turn."), + }, + }) + require.NoError(t, err) + firstResult := waitForTerminalChat(ctx, t, db, chat.ID) + if firstResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(firstResult.LastError)) + } + + secondTV := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + secondBuild, secondAgent := seedNewWorkspaceAgentBuild(t, db, user.ID, org.ID, ws.ID, secondTV.ID) + _, err = db.UpdateChatBuildAgentBinding(ctx, database.UpdateChatBuildAgentBindingParams{ + ID: chat.ID, + BuildID: uuid.NullUUID{UUID: secondBuild.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: secondAgent.ID, Valid: true}, + }) + require.NoError(t, err) + + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Content: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Second turn."), + }, + }) + require.NoError(t, err) + + secondResult := waitForTerminalChat(ctx, t, db, chat.ID) + if secondResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(secondResult.LastError)) + } + require.Equal(t, database.ChatStatusWaiting, secondResult.Status) + + parts := persistedChatParts(ctx, t, db, chat.ID) + require.Len(t, contextFilePartsForAgent(parts, firstAgent.ID), 1) + require.Len(t, contextFilePartsForAgent(parts, secondAgent.ID), 1) + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.GreaterOrEqual(t, len(recorded), 2) + latest := recorded[len(recorded)-1] + require.True(t, requestHasSystemSubstring(latest, newContext)) + require.False(t, requestHasSystemSubstring(latest, oldContext)) + }) } -// TestRecoverStaleChatsWaitingPropagatesSynthError asserts stale -// recovery rolls back when synth-result insertion fails, leaving -// the chat Waiting for the next tick instead of promoting on top -// of incomplete history. -func TestRecoverStaleChatsWaitingPropagatesSynthError(t *testing.T) { - t.Parallel() +func setupWorkspaceContextAgentConn( + t *testing.T, + mockConn *agentconnmock.MockAgentConn, + agent database.WorkspaceAgent, + contextText string, + contextConfigCalls *atomic.Int32, +) { + t.Helper() + directory := agent.ExpandedDirectory + if directory == "" { + directory = agent.Directory + } + if directory == "" { + directory = "/home/coder/project" + } + operatingSystem := agent.OperatingSystem + if operatingSystem == "" { + operatingSystem = "linux" + } + mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + mockConn.EXPECT().ContextConfig(gomock.Any()).DoAndReturn( + func(context.Context) (workspacesdk.ContextConfigResponse, error) { + if contextConfigCalls != nil { + contextConfigCalls.Add(1) + } + return workspacesdk.ContextConfigResponse{ + Parts: []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFilePath: directory + "/AGENTS.md", + ContextFileContent: contextText, + ContextFileOS: operatingSystem, + ContextFileDirectory: directory, + }}, + }, nil + }, + ).AnyTimes() + mockConn.EXPECT().ListMCPTools(gomock.Any()). + Return(workspacesdk.ListMCPToolsResponse{}, nil).AnyTimes() + mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.LSResponse{AbsolutePathString: "/home/coder"}, nil).AnyTimes() + mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() +} - db, ps, rawDB := dbtestutil.NewDBWithSQLDB(t) - ctx := testutil.Context(t, testutil.WaitLong) +func persistedChatParts( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, +) []codersdk.ChatMessagePart { + t.Helper() + messages := persistedChatMessages(ctx, t, db, chatID) + var parts []codersdk.ChatMessagePart + for _, msg := range messages { + parts = append(parts, mustParseChatParts(t, msg)...) + } + return parts +} - staleAfter := 100 * time.Millisecond - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := chatd.New(chatd.Config{ - Logger: logger, - Database: db, - ReplicaID: uuid.New(), - Pubsub: ps, - PendingChatAcquireInterval: testutil.WaitLong, - InFlightChatStaleAfter: staleAfter, - }) - t.Cleanup(func() { require.NoError(t, server.Close()) }) - - user, org, model := seedChatDependencies(t, db) - - dynamicToolsJSON, err := json.Marshal([]mcpgo.Tool{{ - Name: "my_dynamic_tool", - Description: "A test dynamic tool.", - InputSchema: mcpgo.ToolInputSchema{Type: "object", Properties: map[string]any{}}, - }}) - require.NoError(t, err) - - chat, err := db.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, - OwnerID: user.ID, - Title: "stale-waiting-synth-error", - LastModelConfigID: model.ID, - DynamicTools: nullRawMessage(dynamicToolsJSON), - }) - require.NoError(t, err) - - insertUserTextMessage(t, db, chat.ID, user.ID, model.ID, "user input") - - // Inject a synth-results error via an unsupported - // ContentVersion: the row is valid JSON so the insert - // succeeds, but chatprompt.ParseContent rejects it inside the - // helper. Brittle if a future migration adds a content_version - // CHECK constraint; switch to a mock store at that point. - _, err = db.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{uuid.Nil}, - ModelConfigID: []uuid.UUID{model.ID}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - ContentVersion: []int16{99}, - Content: []string{`{}`}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityBoth}, - InputTokens: []int64{0}, - OutputTokens: []int64{0}, - TotalTokens: []int64{0}, - ReasoningTokens: []int64{0}, - CacheCreationTokens: []int64{0}, - CacheReadTokens: []int64{0}, - ContextLimit: []int64{0}, - Compressed: []bool{false}, - TotalCostMicros: []int64{0}, - RuntimeMs: []int64{0}, - ProviderResponseID: []string{""}, - }) - require.NoError(t, err) - - queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText("queued-not-promoted-on-synth-error"), - }) - require.NoError(t, err) - _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ - ChatID: chat.ID, - Content: queuedContent.RawMessage, - ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, - }) - require.NoError(t, err) - - _, err = rawDB.ExecContext(ctx, - "UPDATE chats SET updated_at = $1 WHERE id = $2", - time.Now().Add(-time.Hour), chat.ID) - require.NoError(t, err) - - chatd.RecoverStaleChatsForTest(ctx, server) - - got, err := db.GetChatByID(ctx, chat.ID) - require.NoError(t, err) - require.Equal(t, database.ChatStatusWaiting, got.Status, - "recovery must leave the chat in Waiting when synth-results fails so the next tick retries") - - // The queued message must still be in the queue, not promoted. - remaining, err := db.GetChatQueuedMessages(ctx, chat.ID) - require.NoError(t, err) - require.Len(t, remaining, 1, - "queued message must not be promoted when synth-results fails") - - // No promoted user message should appear in history. +func persistedChatMessages( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, +) []database.ChatMessage { + t.Helper() messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ - ChatID: chat.ID, + ChatID: chatID, AfterID: 0, }) require.NoError(t, err) + return messages +} + +func contextFilePartsForAgent( + parts []codersdk.ChatMessagePart, + agentID uuid.UUID, +) []codersdk.ChatMessagePart { + var matched []codersdk.ChatMessagePart + for _, part := range parts { + if part.Type != codersdk.ChatMessagePartTypeContextFile || + !part.ContextFileAgentID.Valid || + part.ContextFileAgentID.UUID != agentID || + part.ContextFileContent == "" { + continue + } + matched = append(matched, part) + } + return matched +} + +func requireChatToolPart( + t *testing.T, + messages []database.ChatMessage, + partType codersdk.ChatMessagePartType, + toolName string, +) codersdk.ChatMessagePart { + t.Helper() for _, msg := range messages { - if msg.Role != database.ChatMessageRoleUser { - continue + for _, part := range mustParseChatParts(t, msg) { + if part.Type == partType && part.ToolName == toolName { + return part + } } - parts, parseErr := chatprompt.ParseContent(msg) - if parseErr != nil { - continue + } + require.FailNowf(t, "missing chat tool part", "type=%q tool=%q", partType, toolName) + return codersdk.ChatMessagePart{} +} + +func openAIRequestContainsToolResult(req recordedOpenAIRequest, toolResultText string) bool { + for _, msg := range req.Messages { + if msg.Role == "tool" && strings.Contains(msg.Content, toolResultText) { + return true } - for _, part := range parts { - require.NotEqual(t, "queued-not-promoted-on-synth-error", part.Text, - "queued message must not be promoted when synth-results fails") + } + return false +} + +func nextWorkspaceBuildNumber(t *testing.T, db database.Store, workspaceID uuid.UUID) int32 { + t.Helper() + builds, err := db.GetWorkspaceBuildsByWorkspaceID(context.Background(), database.GetWorkspaceBuildsByWorkspaceIDParams{ + WorkspaceID: workspaceID, + OffsetOpt: 0, + LimitOpt: 100, + }) + require.NoError(t, err) + var maxBuild int32 + for _, build := range builds { + if build.BuildNumber > maxBuild { + maxBuild = build.BuildNumber } } + return maxBuild + 1 +} + +func seedNewWorkspaceAgentBuild( + t *testing.T, + db database.Store, + userID uuid.UUID, + orgID uuid.UUID, + workspaceID uuid.UUID, + templateVersionID uuid.UUID, +) (database.WorkspaceBuild, database.WorkspaceAgent) { + t.Helper() + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + InitiatorID: userID, + OrganizationID: orgID, + StartedAt: sql.NullTime{Time: dbtime.Now().Add(-time.Minute), Valid: true}, + CompletedAt: sql.NullTime{Time: dbtime.Now(), Valid: true}, + }) + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: workspaceID, + TemplateVersionID: templateVersionID, + JobID: pj.ID, + BuildNumber: nextWorkspaceBuildNumber(t, db, workspaceID), + InitiatorID: userID, + Transition: database.WorkspaceTransitionStart, + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: pj.ID, + }) + now := dbtime.Now() + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ + ResourceID: res.ID, + LifecycleState: database.WorkspaceAgentLifecycleStateReady, + StartedAt: sql.NullTime{Time: now, Valid: true}, + ReadyAt: sql.NullTime{Time: now, Valid: true}, + FirstConnectedAt: sql.NullTime{Time: now, Valid: true}, + LastConnectedAt: sql.NullTime{Time: now, Valid: true}, + Directory: "/home/coder/project-new", + OperatingSystem: "linux", + }) + require.NoError(t, db.UpdateWorkspaceAgentStartupByID(context.Background(), database.UpdateWorkspaceAgentStartupByIDParams{ + ID: agent.ID, + Version: "v1.0.0", + ExpandedDirectory: "/home/coder/project-new", + })) + loadedAgent, err := db.GetWorkspaceAgentByID(context.Background(), agent.ID) + require.NoError(t, err) + return build, loadedAgent +} + +func seedWorkspaceForCreateTool( + t *testing.T, + db database.Store, + user database.User, + org database.Organization, +) (database.Template, database.WorkspaceTable, database.WorkspaceBuild, database.WorkspaceAgent) { + t.Helper() + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + tpl := dbgen.Template(t, db, database.Template{ + CreatedBy: user.ID, + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + }) + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + TemplateID: tpl.ID, + OwnerID: user.ID, + OrganizationID: org.ID, + }) + build, agent := seedNewWorkspaceAgentBuild(t, db, user.ID, org.ID, ws.ID, tv.ID) + return tpl, ws, build, agent } -// Regression for the cold-start race: chatd must wait long enough -// for ListMCPTools to return after the agent's MCP reload settles. func TestRunChat_WorkspaceMCPDiscoveryWaitsForSlowAgent(t *testing.T) { t.Parallel() @@ -11666,6 +12931,7 @@ func TestRunChat_WorkspaceMCPDiscoveryWaitsForSlowAgent(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "workspace-mcp-slow-agent", ModelConfigID: model.ID, WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, @@ -11691,13 +12957,149 @@ func TestRunChat_WorkspaceMCPDiscoveryWaitsForSlowAgent(t *testing.T) { "timeout exceeds the agent's MCP reload time") } -// TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace guards the -// regression where chats that bound their workspace mid-turn (via -// create_workspace) never saw workspace MCP tools on the same turn. The -// chatloop tool list was frozen at the top of the turn, so the first -// post-create_workspace step had no workspace MCP tools and the model -// fell back to bash. See PrepareTools wiring in runChat. -func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) { +// TestActiveServer_WorkspaceMCPToolDiscoveredMidTurnExecutes guards that +// a workspace MCP tool discovered after mid-turn workspace binding is +// active and executable in later generation actions for the same turn. +func TestActiveServer_WorkspaceMCPToolDiscoveredMidTurnExecutes(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + + var ( + requestsMu sync.Mutex + requests []recordedOpenAIRequest + ) + + workspaceToolName := "workspace-exec-mcp__echo" + workspaceCreateToolArgsJSON := "" + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + + requestsMu.Lock() + requests = append(requests, recordOpenAIRequest(req)) + callIdx := len(requests) + requestsMu.Unlock() + + switch callIdx { + case 1: + return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON)) + case 2: + return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk(workspaceToolName, `{"input":"hello"}`)) + default: + return chattest.OpenAIStreamingResponse( + chattest.OpenAITextChunks("done")..., + ) + } + }) + + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + + // Seed a workspace and agent for create_workspace to bind to. + tpl, ws, build, dbAgent := seedWorkspaceForCreateTool(t, db, user, org) + workspaceCreateToolArgsJSON = fmt.Sprintf(`{"template_id":%q}`, tpl.ID.String()) + + workspaceToolsResp := workspacesdk.ListMCPToolsResponse{ + Tools: []workspacesdk.MCPToolInfo{{ + ServerName: "workspace-exec-mcp", + Name: workspaceToolName, + Description: "workspace echo tool", + Schema: map[string]any{ + "type": "object", + "properties": map[string]any{ + "input": map[string]any{"type": "string"}, + }, + }, + Required: []string{"input"}, + }}, + } + + var callMCPToolCount atomic.Int32 + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + mockConn.EXPECT().SetExtraHeaders(gomock.Any()).AnyTimes() + mockConn.EXPECT().ContextConfig(gomock.Any()). + Return(workspacesdk.ContextConfigResponse{}, xerrors.New("not supported")).AnyTimes() + mockConn.EXPECT().ListMCPTools(gomock.Any()). + Return(workspaceToolsResp, nil).AnyTimes() + mockConn.EXPECT().CallMCPTool(gomock.Any(), gomock.Cond(func(req workspacesdk.CallMCPToolRequest) bool { + return req.ToolName == workspaceToolName && req.Arguments["input"] == "hello" + })).DoAndReturn(func(_ context.Context, _ workspacesdk.CallMCPToolRequest) (workspacesdk.CallMCPToolResponse, error) { + callMCPToolCount.Add(1) + return workspacesdk.CallMCPToolResponse{ + Content: []workspacesdk.MCPToolContent{{ + Type: "text", + Text: "echo: hello", + }}, + }, nil + }).Times(1) + mockConn.EXPECT().LS(gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.LSResponse{}, nil).AnyTimes() + mockConn.EXPECT().ReadFile(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(io.NopCloser(strings.NewReader("")), "", nil).AnyTimes() + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true}, nil).AnyTimes() + mockConn.EXPECT().AwaitReachable(gomock.Any()).Return(true).AnyTimes() + + createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) { + return codersdk.Workspace{ + ID: ws.ID, + Name: req.Name, + OwnerName: user.Username, + OrganizationID: org.ID, + TemplateID: tpl.ID, + LatestBuild: codersdk.WorkspaceBuild{ + ID: build.ID, + Status: codersdk.WorkspaceStatusRunning, + }, + }, nil + } + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + cfg.CreateWorkspace = createFn + }) + + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), + Title: "workspace-mcp-midturn-executes", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("Create a workspace and call the workspace MCP tool."), + }, + }) + require.NoError(t, err) + + chatResult := waitForTerminalChat(ctx, t, db, chat.ID) + if chatResult.Status == database.ChatStatusError { + require.FailNowf(t, "chat failed", "last_error=%q", + chatLastErrorMessage(chatResult.LastError)) + } + require.Equal(t, database.ChatStatusWaiting, chatResult.Status) + require.Equal(t, int32(1), callMCPToolCount.Load()) + + messages := persistedChatMessages(ctx, t, db, chat.ID) + toolCall := requireChatToolPart(t, messages, codersdk.ChatMessagePartTypeToolCall, workspaceToolName) + require.NotEmpty(t, toolCall.ToolCallID) + toolResult := requireChatToolPart(t, messages, codersdk.ChatMessagePartTypeToolResult, workspaceToolName) + require.Contains(t, string(toolResult.Result), "echo: hello") + + requestsMu.Lock() + recorded := append([]recordedOpenAIRequest(nil), requests...) + requestsMu.Unlock() + require.GreaterOrEqual(t, len(recorded), 3) + require.Contains(t, recorded[1].Tools, workspaceToolName) + require.True(t, openAIRequestContainsToolResult(recorded[len(recorded)-1], "echo: hello")) +} + +func TestActiveServer_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -11722,9 +13124,7 @@ func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) requestsMu.Unlock() if callIdx == 1 { - return chattest.OpenAIStreamingResponse( - chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON), - ) + return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON)) } return chattest.OpenAIStreamingResponse( chattest.OpenAITextChunks("done")..., @@ -11733,47 +13133,10 @@ func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - // Seed a workspace+agent for create_workspace to bind to. - tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - tpl := dbgen.Template(t, db, database.Template{ - CreatedBy: user.ID, - OrganizationID: org.ID, - ActiveVersionID: tv.ID, - }) + // Seed a workspace and agent for create_workspace to bind to. + tpl, ws, build, dbAgent := seedWorkspaceForCreateTool(t, db, user, org) workspaceCreateToolArgsJSON = fmt.Sprintf(`{"template_id":%q}`, tpl.ID.String()) - ws := dbgen.Workspace(t, db, database.WorkspaceTable{ - TemplateID: tpl.ID, - OwnerID: user.ID, - OrganizationID: org.ID, - }) - pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ - InitiatorID: user.ID, - OrganizationID: org.ID, - CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, - }) - build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ - TemplateVersionID: tv.ID, - WorkspaceID: ws.ID, - JobID: pj.ID, - }) - res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ - Transition: database.WorkspaceTransitionStart, - JobID: pj.ID, - }) - now := dbtime.Now() - dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ - ResourceID: res.ID, - LifecycleState: database.WorkspaceAgentLifecycleStateReady, - StartedAt: sql.NullTime{Time: now, Valid: true}, - ReadyAt: sql.NullTime{Time: now, Valid: true}, - FirstConnectedAt: sql.NullTime{Time: now, Valid: true}, - LastConnectedAt: sql.NullTime{Time: now, Valid: true}, - }) - workspaceToolsResp := workspacesdk.ListMCPToolsResponse{ Tools: []workspacesdk.MCPToolInfo{{ ServerName: "workspace-midturn-mcp", @@ -11824,6 +13187,7 @@ func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "workspace-mcp-midturn", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -11851,21 +13215,11 @@ func TestRunChat_WorkspaceMCPDiscoveryAfterMidTurnCreateWorkspace(t *testing.T) "this is the fix for mid-turn workspace MCP discovery") } -// TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery guards the -// regression on the workspaceMCPDiscovered flag flip: the prior -// implementation set the flag to true before calling -// discoverWorkspaceMCPTools, so a single empty result permanently -// blocked retries within the turn. The fix sets the flag to true -// only after a non-empty discovery, so subsequent PrepareTools -// invocations keep retrying until tools appear. -// -// Scenario: create_workspace binds a workspace mid-turn. The first -// few ListMCPTools calls return empty (simulating the agent's MCP -// Connect still racing with agent startup); a later call returns -// the workspace MCP tool. The chat takes multiple steps before -// finishing, and we assert that one of the post-create_workspace -// streamed model calls advertises the workspace tool. -func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { +// TestActiveServer_WorkspaceMCPDiscoveryRetriesAfterEmptyResult guards +// the regression where an empty workspace MCP discovery result +// permanently blocked retries within the turn. The active worker should +// retry discovery in later generation actions until tools appear. +func TestActiveServer_WorkspaceMCPDiscoveryRetriesAfterEmptyResult(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) @@ -11891,16 +13245,10 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { // Step 1: trigger create_workspace. if callIdx == 1 { - return chattest.OpenAIStreamingResponse( - chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON), - ) + return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk("create_workspace", workspaceCreateToolArgsJSON)) } - // Step 2..N-1: emit empty text to keep the chatloop running so - // PrepareTools fires on each step. The chatloop ends a turn - // when the model returns a non-empty assistant message with no - // tool calls; an empty text chunk would terminate the turn, so - // we attach a dummy tool call to force another step. Use the - // LS tool because it exists for all workspaces and is cheap. + // Step 2..N-1 calls a cheap workspace tool so the active worker + // runs several generation actions before the final assistant text. if callIdx < 6 { return chattest.OpenAIStreamingResponse( chattest.OpenAIToolCallChunk("ls", `{"path":"/tmp"}`), @@ -11914,47 +13262,10 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) - // Seed a workspace+agent for create_workspace to bind to. - tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ - OrganizationID: org.ID, - CreatedBy: user.ID, - }) - tpl := dbgen.Template(t, db, database.Template{ - CreatedBy: user.ID, - OrganizationID: org.ID, - ActiveVersionID: tv.ID, - }) + // Seed a workspace and agent for create_workspace to bind to. + tpl, ws, build, dbAgent := seedWorkspaceForCreateTool(t, db, user, org) workspaceCreateToolArgsJSON = fmt.Sprintf(`{"template_id":%q}`, tpl.ID.String()) - ws := dbgen.Workspace(t, db, database.WorkspaceTable{ - TemplateID: tpl.ID, - OwnerID: user.ID, - OrganizationID: org.ID, - }) - pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ - InitiatorID: user.ID, - OrganizationID: org.ID, - CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, - }) - build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ - TemplateVersionID: tv.ID, - WorkspaceID: ws.ID, - JobID: pj.ID, - }) - res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ - Transition: database.WorkspaceTransitionStart, - JobID: pj.ID, - }) - now := dbtime.Now() - dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ - ResourceID: res.ID, - LifecycleState: database.WorkspaceAgentLifecycleStateReady, - StartedAt: sql.NullTime{Time: now, Valid: true}, - ReadyAt: sql.NullTime{Time: now, Valid: true}, - FirstConnectedAt: sql.NullTime{Time: now, Valid: true}, - LastConnectedAt: sql.NullTime{Time: now, Valid: true}, - }) - workspaceToolsResp := workspacesdk.ListMCPToolsResponse{ Tools: []workspacesdk.MCPToolInfo{{ ServerName: "workspace-empty-retry-mcp", @@ -11967,13 +13278,10 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { }}, } - // First two ListMCPTools calls return empty (no error). One is the - // primer goroutine's only attempt before its retry timer fires; - // the other is PrepareTools on the first post-create_workspace - // step. The third and later calls return the workspace tool. The - // assertion below requires that a post-create_workspace step - // eventually advertises the tool, which can only happen if the - // PrepareTools callback retries discovery on subsequent steps. + // First two ListMCPTools calls return empty (no error). One may + // come from the cache primer and one from the first generation + // action after create_workspace. Later calls return the workspace + // tool, proving discovery retries after empty results. var listCalls atomic.Int32 ctrl := gomock.NewController(t) mockConn := agentconnmock.NewMockAgentConn(ctrl) @@ -12020,6 +13328,7 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "workspace-mcp-empty-retry", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -12041,14 +13350,9 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { require.GreaterOrEqual(t, len(recorded), 3, "expected at least three streamed model calls; chat must run past the empty discovery") - // The first call has no workspace yet; the second call is the - // first post-create_workspace step which sees an empty - // ListMCPTools result. By the third (or later) call PrepareTools - // must have retried discovery, so at least one post-step request - // must advertise the workspace tool. Without the - // workspaceMCPDiscovered flag-flip fix the flag would have been - // set true on the failed first attempt and no subsequent step - // would have re-attempted discovery. + // The first call has no workspace yet. By a later post-binding + // call, workspace MCP discovery must have retried after the empty + // results and advertised the workspace tool. sawWorkspaceTool := false for i := 2; i < len(recorded); i++ { if slices.Contains(recorded[i].Tools, workspaceToolName) { @@ -12057,7 +13361,7 @@ func TestRunChat_PrepareToolsRetriesAfterEmptyDiscovery(t *testing.T) { } } require.True(t, sawWorkspaceTool, - "PrepareTools must retry workspace MCP discovery on subsequent "+ - "steps; without the fix the first empty result would "+ - "permanently block retries within the turn") + "workspace MCP discovery must retry on subsequent steps; "+ + "without the fix the first empty result would permanently "+ + "block retries within the turn") } diff --git a/coderd/x/chatd/chatdebug/service.go b/coderd/x/chatd/chatdebug/service.go index 091d8ece26..89a5667cda 100644 --- a/coderd/x/chatd/chatdebug/service.go +++ b/coderd/x/chatd/chatdebug/service.go @@ -522,6 +522,60 @@ func (s *Service) TouchStep( }) } +// TouchRun bumps the run's updated_at timestamp without changing any +// other fields. Runner-owned debug turns use this while no model step is +// active, such as requires-action waits. +func (s *Service) TouchRun(ctx context.Context, runID uuid.UUID, chatID uuid.UUID) error { + if s == nil || runID == uuid.Nil || chatID == uuid.Nil { + return nil + } + return s.db.TouchChatDebugRunUpdatedAt(chatdContext(ctx), + database.TouchChatDebugRunUpdatedAtParams{ + Now: s.clock.Now(), + ID: runID, + ChatID: chatID, + }) +} + +// LaunchRunHeartbeat starts a goroutine that periodically touches an +// open run until done is closed or ctx is canceled. +func (s *Service) LaunchRunHeartbeat(ctx context.Context, runID uuid.UUID, chatID uuid.UUID, done <-chan struct{}) { + if s == nil || runID == uuid.Nil || chatID == uuid.Nil || done == nil { + return + } + go func() { + thresholdCh := s.thresholdChan() + interval := s.heartbeatInterval() + ticker := s.clock.NewTicker(interval, "chatdebug", "run-heartbeat") + defer ticker.Stop() + resetTicker := func() { + if newInterval := s.heartbeatInterval(); newInterval != interval { + interval = newInterval + ticker.Reset(interval, "chatdebug", "run-heartbeat") + } + } + for { + select { + case <-ctx.Done(): + return + case <-done: + return + case <-thresholdCh: + thresholdCh = s.thresholdChan() + resetTicker() + case <-ticker.C: + if err := s.TouchRun(ctx, runID, chatID); err != nil { + s.log.Debug(ctx, "run heartbeat touch failed", + slog.Error(err), + slog.F("run_id", runID), + ) + } + resetTicker() + } + } + }() +} + // DeleteByChatID deletes debug data for a chat and emits a delete event. // The startedBefore bound scopes deletion to runs created before that // instant so that retried cleanup does not remove runs created by a diff --git a/coderd/x/chatd/chatdebug/summary.go b/coderd/x/chatd/chatdebug/summary.go index 7b69a6b8c3..75c9da8e84 100644 --- a/coderd/x/chatd/chatdebug/summary.go +++ b/coderd/x/chatd/chatdebug/summary.go @@ -41,28 +41,6 @@ func SeedSummary(label string) map[string]any { return map[string]any{"first_message": label} } -// ExtractFirstUserText extracts the plain text content from a -// fantasy.Prompt for the first user message. Used to derive -// first_message labels at run creation time. -func ExtractFirstUserText(prompt fantasy.Prompt) string { - for _, msg := range prompt { - if msg.Role != fantasy.MessageRoleUser { - continue - } - - var sb strings.Builder - for _, part := range msg.Content { - tp, ok := fantasy.AsMessagePart[fantasy.TextPart](part) - if !ok { - continue - } - _, _ = sb.WriteString(tp.Text) - } - return sb.String() - } - return "" -} - // AggregateRunSummary reads all steps for the given run, computes token // totals, and merges them with the run's existing summary (preserving any // seeded first_message label). The baseSummary parameter should be the diff --git a/coderd/x/chatd/chatdebug/summary_test.go b/coderd/x/chatd/chatdebug/summary_test.go index 3c41877cd2..fc329ae0a9 100644 --- a/coderd/x/chatd/chatdebug/summary_test.go +++ b/coderd/x/chatd/chatdebug/summary_test.go @@ -6,7 +6,6 @@ import ( "time" "unicode/utf8" - "charm.land/fantasy" "github.com/google/uuid" "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" @@ -62,68 +61,6 @@ func TestSeedSummary(t *testing.T) { }) } -func TestExtractFirstUserText(t *testing.T) { - t.Parallel() - - t.Run("EmptyPrompt", func(t *testing.T) { - t.Parallel() - got := chatdebug.ExtractFirstUserText(fantasy.Prompt{}) - require.Equal(t, "", got) - }) - - t.Run("NoUserMessages", func(t *testing.T) { - t.Parallel() - prompt := fantasy.Prompt{ - { - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "system"}}, - }, - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "assistant"}}, - }, - } - got := chatdebug.ExtractFirstUserText(prompt) - require.Equal(t, "", got) - }) - - t.Run("FirstUserMessageMixedParts", func(t *testing.T) { - t.Parallel() - prompt := fantasy.Prompt{ - { - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: "hello "}, - fantasy.FilePart{Filename: "test.png"}, - fantasy.TextPart{Text: "world"}, - }, - }, - } - got := chatdebug.ExtractFirstUserText(prompt) - require.Equal(t, "hello world", got) - }) - - t.Run("MultipleUserMessagesReturnsFirst", func(t *testing.T) { - t.Parallel() - prompt := fantasy.Prompt{ - { - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "system"}}, - }, - { - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "first"}}, - }, - { - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "second"}}, - }, - } - got := chatdebug.ExtractFirstUserText(prompt) - require.Equal(t, "first", got) - }) -} - func TestService_AggregateRunSummary(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index efe67083e2..f2aa7dab19 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -6,7 +6,6 @@ import ( "encoding/base64" "encoding/json" "errors" - "maps" "slices" "strconv" "strings" @@ -17,6 +16,7 @@ import ( "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" "charm.land/fantasy/schema" + "github.com/google/uuid" "golang.org/x/xerrors" "cdr.dev/slog/v3" @@ -33,12 +33,6 @@ import ( ) 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 // defaultStreamSilenceTimeout bounds how long an individual // model attempt may go without receiving a stream part before // the attempt is canceled and retried. @@ -214,6 +208,79 @@ type RunOptions struct { BuiltinToolNames map[string]bool } +// GenerateAssistantOptions configures one assistant model call. +type GenerateAssistantOptions struct { + Model fantasy.LanguageModel + Messages []fantasy.Message + Tools []fantasy.AgentTool + ActiveTools []string + ProviderTools []ProviderTool + StreamSilenceTimeout time.Duration + Clock quartz.Clock + + ContextLimitFallback int64 + ModelConfig codersdk.ChatModelCallConfig + ProviderOptions fantasy.ProviderOptions + + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + Logger slog.Logger + Metrics *Metrics +} + +// AssistantOutcome is the durable assistant-side result from one model call. +type AssistantOutcome struct { + Step PersistedStep + ToolCalls []fantasy.ToolCallContent + FinishReason fantasy.FinishReason + ModelStopped bool +} + +// ExecuteLocalToolsOptions configures one local tool execution batch. +type ExecuteLocalToolsOptions struct { + Tools []fantasy.AgentTool + ActiveTools []string + ProviderTools []ProviderTool + ToolCalls []fantasy.ToolCallContent + + ExclusiveToolNames map[string]bool + BuiltinToolNames map[string]bool + ModelProvider string + ModelName string + + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) + Logger slog.Logger + Metrics *Metrics + Clock quartz.Clock +} + +// ToolExecutionOutcome is the durable tool-result content from one batch. +type ToolExecutionOutcome struct { + Step PersistedStep +} + +// GenerateCompactionOptions configures one context compaction call. +type GenerateCompactionOptions struct { + Model fantasy.LanguageModel + Messages []fantasy.Message + + ThresholdPercent int32 + ContextLimit int64 + ContextLimitFallback int64 + SummaryPrompt string + SystemSummaryPrefix string + Timeout time.Duration + StepUsage fantasy.Usage + StepMetadata fantasy.ProviderMetadata + + DebugSvc *chatdebug.Service + ChatID uuid.UUID + HistoryTipMessageID int64 + ToolCallID string + ToolName string + + PublishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) +} + // 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 @@ -245,106 +312,6 @@ type stepResult struct { reasoningCompletedAt []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 { - continue - } - opts := fantasy.ProviderOptions(reasoning.ProviderMetadata) - if strings.TrimSpace(reasoning.Text) == "" && !chatsanitize.HasAnthropicSignedReasoningOptions(opts) { - continue - } - assistantParts = append(assistantParts, fantasy.ReasoningPart{ - Text: reasoning.Text, - ProviderOptions: opts, - }) - 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 { @@ -353,17 +320,11 @@ type reasoningState struct { startedAt time.Time } -// Run executes the chat step-stream loop and delegates -// persistence/publishing to callbacks. -func Run(ctx context.Context, opts RunOptions) error { +// GenerateAssistant performs one assistant model stream and returns the +// durable assistant-side content. It does not execute tools, retry, or persist. +func GenerateAssistant(ctx context.Context, opts GenerateAssistantOptions) (AssistantOutcome, 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 + return AssistantOutcome{}, xerrors.New("chat model is required") } if opts.StreamSilenceTimeout <= 0 { opts.StreamSilenceTimeout = defaultStreamSilenceTimeout @@ -376,360 +337,207 @@ func Run(ctx context.Context, opts RunOptions) error { } publishMessagePart := func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { - if opts.PublishMessagePart == nil { - return + if opts.PublishMessagePart != nil { + opts.PublishMessagePart(role, part) } - opts.PublishMessagePart(role, part) } - tools := buildToolDefinitions(opts.Tools, opts.ActiveTools, opts.ProviderTools) + provider := opts.Model.Provider() + modelName := opts.Model.Model() + runOpts := RunOptions{ + Model: opts.Model, + Logger: opts.Logger, + } + _, prepared, err := prepareMessagesForRequest(ctx, runOpts, opts.Messages, provider, modelName, 0, 1) + if err != nil { + return AssistantOutcome{}, xerrors.Errorf("prepare prompt: %w", err) + } + opts.Metrics.MessageCount.WithLabelValues(provider, modelName).Observe(float64(len(prepared))) + opts.Metrics.PromptSizeBytes.WithLabelValues(provider, modelName).Observe(float64(EstimatePromptSize(prepared))) + opts.Metrics.StepsTotal.WithLabelValues(provider, modelName).Inc() - messages := opts.Messages - var lastUsage fantasy.Usage - var lastProviderMetadata fantasy.ProviderMetadata - needsFullHistoryReload := false - reloadFullHistory := func(stage string) error { - if opts.ReloadMessages == nil { - return nil + call := fantasy.Call{ + Prompt: prepared, + Tools: buildToolDefinitions(opts.Tools, opts.ActiveTools, opts.ProviderTools), + 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, + } + + stepStart := opts.Clock.Now() + stepCtx := chatdebug.ReuseStep(ctx) + attempt, streamErr := guardedStream( + stepCtx, + provider, + modelName, + opts.Clock, + opts.StreamSilenceTimeout, + func(attemptCtx context.Context) (fantasy.StreamResponse, error) { + return opts.Model.Stream(attemptCtx, call) + }, + opts.Metrics, + ) + if streamErr != nil { + wrappedErr := wrapProviderStreamError(provider, streamErr) + classified := chaterror.Classify(wrappedErr).WithProvider(provider) + if classified.Retryable { + opts.Metrics.RecordStreamRetry(provider, modelName, classified) } - reloaded, err := opts.ReloadMessages(ctx) - if err != nil { - return xerrors.Errorf("reload messages %s: %w", stage, err) + return AssistantOutcome{}, wrappedErr + } + defer attempt.release() + + result, processErr := processStepStream(attempt.ctx, attempt.stream, opts.Clock, publishMessagePart) + if err := attempt.finish(processErr); err != nil { + if errors.Is(err, ErrInterrupted) { + return AssistantOutcome{}, ErrInterrupted } - messages = reloaded + wrappedErr := wrapProviderStreamError(provider, err) + classified := chaterror.Classify(wrappedErr).WithProvider(provider) + if classified.Retryable { + opts.Metrics.RecordStreamRetry(provider, modelName, classified) + } + return AssistantOutcome{}, wrappedErr + } + + contextLimit := extractContextLimitWithFallback(result.providerMetadata, opts.ContextLimitFallback) + result.content = chatsanitize.SanitizeAnthropicProviderToolStepContent( + ctx, opts.Logger, provider, modelName, + "assistant_helper", 0, result.finishReason, result.content, + ) + step := PersistedStep{ + Content: result.content, + Usage: result.usage, + ContextLimit: contextLimit, + ProviderResponseID: chatopenai.ExtractResponseIDIfStored(opts.ProviderOptions, result.providerMetadata), + Runtime: opts.Clock.Since(stepStart), + ToolCallCreatedAt: result.toolCallCreatedAt, + ToolResultCreatedAt: result.toolResultCreatedAt, + ReasoningStartedAt: result.reasoningStartedAt, + ReasoningCompletedAt: result.reasoningCompletedAt, + } + return AssistantOutcome{ + Step: step, + ToolCalls: append([]fantasy.ToolCallContent(nil), result.toolCalls...), + FinishReason: result.finishReason, + ModelStopped: len(result.content) == 0, + }, nil +} + +func wrapProviderStreamError(provider string, err error) error { + if err == nil { 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() - if opts.PrepareTools != nil { - if updated := opts.PrepareTools(opts.Tools); updated != nil { - opts.ActiveTools = mergeNewToolNames( - opts.ActiveTools, opts.Tools, updated, - ) - opts.Tools = updated - tools = buildToolDefinitions( - opts.Tools, opts.ActiveTools, opts.ProviderTools, - ) - } - } - var prepared []fantasy.Message - var prepareErr error - messages, prepared, prepareErr = prepareMessagesForRequest( - ctx, opts, messages, provider, modelName, step, totalSteps, - ) - if prepareErr != nil { - return xerrors.Errorf("prepare prompt: %w", prepareErr) - } - 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 - var retryPrepareErr error - stepCtx := chatdebug.ReuseStep(ctx) - err := chatretry.Retry(stepCtx, func(retryCtx context.Context) error { - if retryPrepareErr != nil { - return retryPrepareErr - } - attempt, streamErr := guardedStream( - retryCtx, - provider, - modelName, - opts.Clock, - opts.StreamSilenceTimeout, - 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 classified.ChainBroken { - if chatopenai.HasPreviousResponseID(opts.ProviderOptions) { - opts.ProviderOptions = chatopenai.ClearPreviousResponseID(opts.ProviderOptions) - } - if chatopenai.HasPreviousResponseID(call.ProviderOptions) { - call.ProviderOptions = chatopenai.ClearPreviousResponseID(call.ProviderOptions) - } - if opts.DisableChainMode != nil { - opts.DisableChainMode() - } - if opts.ReloadMessages != nil { - reloaded, err := opts.ReloadMessages(ctx) - if err != nil { - opts.Logger.Warn(ctx, - "chain-broken recovery: reload messages failed", - slog.Error(err), - ) - } else { - // Reloaded history replaces the prompt prepared before - // the failed attempt, so run the same preparation - // pipeline used by normal provider requests. - var ( - reloadedCanonical []fantasy.Message - retryPrompt []fantasy.Message - prepareErr error - ) - call.Prompt = nil - reloadedCanonical, retryPrompt, prepareErr = prepareMessagesForRequest( - ctx, opts, reloaded, provider, modelName, step, totalSteps, - ) - if prepareErr != nil { - retryPrepareErr = prepareErr - } else { - messages = reloadedCanonical - call.Prompt = retryPrompt - } - } - } - } - if opts.OnRetry != nil { - opts.OnRetry(attempt, retryErr, classified, delay) - } - }) - if err != nil { - if errors.Is(err, ErrInterrupted) { - persistInterruptedStep(ctx, opts, &result) - return ErrInterrupted - } - if retryPrepareErr != nil && errors.Is(err, retryPrepareErr) { - return xerrors.Errorf("prepare prompt: %w", err) - } - 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 { - var err error - toolResults, err = executeToolsForStep(ctx, opts, &result, provider, modelName, step, stepStart, publishMessagePart) - if err != nil { - return 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: chatopenai.ExtractResponseIDIfStored(opts.ProviderOptions, result.providerMetadata), - Runtime: time.Since(stepStart), - ToolCallCreatedAt: result.toolCallCreatedAt, - ToolResultCreatedAt: result.toolResultCreatedAt, - ReasoningStartedAt: result.reasoningStartedAt, - ReasoningCompletedAt: result.reasoningCompletedAt, - }); 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 chatopenai.HasPreviousResponseID(opts.ProviderOptions) { - opts.ProviderOptions = chatopenai.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 + classified := chaterror.Classify(err).WithProvider(provider) + if !classified.Retryable && classified.StatusCode == 0 && errors.Is(err, context.Canceled) { + wrapped := errors.Join(chaterror.ErrProviderTransportReset, err) + reclassified := chaterror.Classify(wrapped).WithProvider(provider) + if reclassified.Retryable { + classified = reclassified + err = wrapped } + } + return xerrors.Errorf("stream response: %w", chaterror.WithClassification(err, classified)) +} - if needsFullHistoryReload && stoppedByModel && opts.ReloadMessages != nil { - if err := reloadFullHistory("before post-run compaction after chain mode exit"); err != nil { - return err - } - needsFullHistoryReload = false +// ExecuteLocalTools runs local tool calls and returns durable tool results. It +// does not retry or persist. +func ExecuteLocalTools(ctx context.Context, opts ExecuteLocalToolsOptions) (ToolExecutionOutcome, error) { + if opts.Metrics == nil { + opts.Metrics = NopMetrics() + } + provider := opts.ModelProvider + if provider == "" { + provider = "unknown" + } + modelName := opts.ModelName + if modelName == "" { + modelName = "unknown" + } + publishMessagePart := func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + if opts.PublishMessagePart != nil { + opts.PublishMessagePart(role, part) } - - // 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 + } + // Expose the publisher on the execution context so tools that stream + // intermediate output (e.g. the advisor tool) can publish parts + // without capturing the publisher at construction time. + ctx = WithMessagePartPublisher(ctx, opts.PublishMessagePart) + if ctx.Err() != nil { + return ToolExecutionOutcome{}, ctx.Err() } - return nil + localCalls := make([]fantasy.ToolCallContent, 0, len(opts.ToolCalls)) + for _, tc := range opts.ToolCalls { + if !tc.ProviderExecuted { + localCalls = append(localCalls, tc) + } + } + if len(localCalls) == 0 { + return ToolExecutionOutcome{}, nil + } + + var result stepResult + policyResults, exclusiveViolation := applyExclusiveToolPolicy( + localCalls, + opts.ExclusiveToolNames, + opts.Metrics, + provider, + modelName, + ) + if exclusiveViolation { + now := clockNow(opts.Clock) + for _, tr := range policyResults { + recordToolResultTimestamp(&result, tr.ToolCallID, now) + publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) + ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) + ssePart.CreatedAt = &now + publishMessagePart(codersdk.ChatMessageRoleTool, ssePart) + result.content = append(result.content, tr) + } + if ctx.Err() != nil { + return ToolExecutionOutcome{}, ctx.Err() + } + return ToolExecutionOutcome{Step: PersistedStep{ + Content: result.content, + ToolResultCreatedAt: result.toolResultCreatedAt, + }}, nil + } + + toolResults := executeTools( + ctx, + opts.Clock, + opts.Tools, + opts.ActiveTools, + opts.ProviderTools, + localCalls, + 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) + }, + ) + if ctx.Err() != nil { + return ToolExecutionOutcome{}, ctx.Err() + } + for _, tr := range toolResults { + result.content = append(result.content, tr) + } + return ToolExecutionOutcome{Step: PersistedStep{ + Content: result.content, + ToolResultCreatedAt: result.toolResultCreatedAt, + }}, nil } // prepareMessagesForRequest applies the prompt preparation pipeline used @@ -922,12 +730,19 @@ func guardedStream( }, nil } +// clockNow returns the clock's current time normalized the same +// way as dbtime.Now so persisted timestamps are Postgres-safe. +func clockNow(clock quartz.Clock) time.Time { + return dbtime.Time(clock.Now().UTC()) +} + // 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, + clock quartz.Clock, publishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), ) (stepResult, error) { var result stepResult @@ -962,18 +777,23 @@ func processStepStream( activeReasoningContent[part.ID] = reasoningState{ text: part.Delta, options: part.ProviderMetadata, - startedAt: dbtime.Now(), + startedAt: clockNow(clock), } case fantasy.StreamPartTypeReasoningDelta: + reasoningPart := codersdk.ChatMessageReasoning(part.Delta) if active, exists := activeReasoningContent[part.ID]; exists { active.text += part.Delta if len(part.ProviderMetadata) > 0 { active.options = part.ProviderMetadata } activeReasoningContent[part.ID] = active + if !active.startedAt.IsZero() { + startedAt := active.startedAt + reasoningPart.CreatedAt = &startedAt + } } - publishMessagePart(codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageReasoning(part.Delta)) + publishMessagePart(codersdk.ChatMessageRoleAssistant, reasoningPart) case fantasy.StreamPartTypeReasoningEnd: if active, exists := activeReasoningContent[part.ID]; exists { @@ -986,7 +806,7 @@ func processStepStream( } result.content = append(result.content, content) result.reasoningStartedAt = append(result.reasoningStartedAt, active.startedAt) - result.reasoningCompletedAt = append(result.reasoningCompletedAt, dbtime.Now()) + result.reasoningCompletedAt = append(result.reasoningCompletedAt, clockNow(clock)) delete(activeReasoningContent, part.ID) } case fantasy.StreamPartTypeToolInputStart: @@ -1037,7 +857,7 @@ func processStepStream( // Record when the model emitted this tool call // so the persisted part carries an accurate // timestamp for duration computation. - now := dbtime.Now() + now := clockNow(clock) if result.toolCallCreatedAt == nil { result.toolCallCreatedAt = make(map[string]time.Time) } @@ -1078,7 +898,7 @@ func processStepStream( } result.content = append(result.content, tr) - now := dbtime.Now() + now := clockNow(clock) if result.toolResultCreatedAt == nil { result.toolResultCreatedAt = make(map[string]time.Time) } @@ -1109,6 +929,7 @@ func processStepStream( // still streaming when the interrupt arrived. flushActiveState( &result, + clock, activeTextContent, activeReasoningContent, activeToolCalls, @@ -1129,6 +950,7 @@ func processStepStream( errors.Is(context.Cause(ctx), ErrInterrupted) { flushActiveState( &result, + clock, activeTextContent, activeReasoningContent, activeToolCalls, @@ -1154,6 +976,7 @@ func processStepStream( // event ordering for SSE subscribers. func executeTools( ctx context.Context, + clock quartz.Clock, allTools []fantasy.AgentTool, activeTools []string, providerTools []ProviderTool, @@ -1226,7 +1049,7 @@ func executeTools( // Record when this tool completed (or panicked). // Captured per-goroutine so parallel tools get // accurate individual completion times. - completedAt[i] = dbtime.Now() + completedAt[i] = clockNow(clock) }() results[i] = executeSingleTool( ctx, @@ -1255,187 +1078,6 @@ func executeTools( return results } -// executeToolsForStep runs the tool-execution phase of a single -// chatloop step. It enforces the exclusive-tool policy, partitions -// built-in versus dynamic tool calls, dispatches built-in tools, and -// when dynamic tool calls are present persists the step and returns -// ErrDynamicToolCall so the caller can execute them externally. -// Returns the tool results to append to the step, or an error that the -// caller must propagate (ErrInterrupted, ErrDynamicToolCall, ctx.Err(), -// or a persistence failure). -func executeToolsForStep( - ctx context.Context, - opts RunOptions, - result *stepResult, - provider, modelName string, - step int, - stepStart time.Time, - publishMessagePart func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), -) ([]fantasy.ToolResultContent, error) { - // 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 nil, ErrInterrupted - } - return nil, ctx.Err() - } - - // Enforce exclusivity across ALL locally-executable tool - // calls (both built-in and dynamic) before partitioning. - // Checking only the built-in partition would let the model - // bypass the policy by mixing an exclusive tool with a - // dynamic tool: the exclusive tool would still run and the - // dynamic call would still be handed to the caller for - // external execution, breaking the planning-only contract. - localCandidates := make([]fantasy.ToolCallContent, 0, len(result.toolCalls)) - for _, tc := range result.toolCalls { - if !tc.ProviderExecuted { - localCandidates = append(localCandidates, tc) - } - } - policyResults, exclusiveViolation := applyExclusiveToolPolicy( - localCandidates, - opts.ExclusiveToolNames, - opts.Metrics, - provider, - modelName, - ) - if exclusiveViolation { - now := dbtime.Now() - for _, tr := range policyResults { - recordToolResultTimestamp(result, tr.ToolCallID, now) - publishToolAttachments(ctx, opts.Logger, tr, now, publishMessagePart) - ssePart := chatprompt.PartFromContentWithLogger(ctx, opts.Logger, tr) - ssePart.CreatedAt = &now - publishMessagePart(codersdk.ChatMessageRoleTool, ssePart) - } - for _, tr := range policyResults { - result.content = append(result.content, tr) - } - // Mirror the post-execution interruption check used by the - // non-policy path: if the chat was interrupted while we - // synthesized policy errors, route through - // persistInterruptedStep so the synthesized results are not - // dropped when the regular PersistStep path fails on a - // canceled context. - if ctx.Err() != nil { - if errors.Is(context.Cause(ctx), ErrInterrupted) { - persistInterruptedStep(ctx, opts, result) - return nil, ErrInterrupted - } - return nil, ctx.Err() - } - // Fall through to the normal persistence path so the loop - // continues with error results that the model can observe - // and retry. Skip partitioning, execution, and - // pending-dynamic persistence. - return policyResults, nil - } - - // 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 { - // Strip Anthropic provider-executed tool calls without - // matching results before persisting so the action-required - // step does not carry a malformed tool-call history into - // downstream provider requests. - result.content = chatsanitize.SanitizeAnthropicProviderToolStepContent( - ctx, opts.Logger, provider, modelName, - "dynamic_tool_persist", step, result.finishReason, result.content, - ) - if err := persistPendingDynamicStep(ctx, opts, result, stepStart, dynamicCalls); err != nil { - return nil, err - } - tryCompactOnExit(ctx, opts, result.usage, result.providerMetadata) - return nil, 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 nil, ErrInterrupted - } - return nil, ctx.Err() - } - - return toolResults, nil -} - -// persistPendingDynamicStep persists a step that has pending dynamic -// tool calls awaiting external execution. Returns ErrInterrupted when -// persistence fails because the chat was interrupted. -func persistPendingDynamicStep( - ctx context.Context, - opts RunOptions, - result *stepResult, - stepStart time.Time, - dynamicCalls []fantasy.ToolCallContent, -) error { - 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) - - if err := opts.PersistStep(ctx, PersistedStep{ - Content: result.content, - Usage: result.usage, - ContextLimit: contextLimit, - ProviderResponseID: chatopenai.ExtractResponseIDIfStored(opts.ProviderOptions, result.providerMetadata), - Runtime: time.Since(stepStart), - PendingDynamicToolCalls: pending, - ReasoningStartedAt: result.reasoningStartedAt, - ReasoningCompletedAt: result.reasoningCompletedAt, - }); err != nil { - if errors.Is(err, ErrInterrupted) { - persistInterruptedStep(ctx, opts, result) - return ErrInterrupted - } - return xerrors.Errorf("persist step: %w", err) - } - return nil -} - // applyExclusiveToolPolicy checks whether toolCalls violate the // exclusive-tool policy declared by exclusiveToolNames. When a // violation is detected it synthesizes deterministic policy-error @@ -1643,6 +1285,7 @@ func executeSingleTool( // persistence. func flushActiveState( result *stepResult, + clock quartz.Clock, activeText map[string]string, activeReasoning map[string]reasoningState, activeToolCalls map[string]*fantasy.ToolCallContent, @@ -1659,7 +1302,7 @@ func flushActiveState( // completedAt is filled in here with the interruption // time so partial reasoning shows the time spent before // the interruption. - flushedAt := dbtime.Now() + flushedAt := clockNow(clock) for _, rs := range activeReasoning { if rs.text == "" && !chatsanitize.HasAnthropicSignedReasoningOptions(fantasy.ProviderOptions(rs.options)) { continue @@ -1698,173 +1341,10 @@ func flushActiveState( } } -// persistInterruptedStep saves durable content from a partial stream. -// Provider-executed calls without results are removed because their result -// metadata cannot be synthesized safely, except when removal would mutate -// signed Anthropic replay state. -func persistInterruptedStep( - ctx context.Context, - opts RunOptions, - result *stepResult, -) { - if result == nil || (len(result.content) == 0 && len(result.toolCalls) == 0) { - return - } - - provider := "" - modelName := "" - if opts.Model != nil { - provider = opts.Model.Provider() - modelName = opts.Model.Model() - } - var sanitizeStats chatsanitize.AnthropicProviderToolSanitizationStats - result.content, sanitizeStats = chatsanitize.SanitizeAnthropicProviderToolContent(provider, result.content) - chatsanitize.LogAnthropicProviderToolSanitization( - ctx, opts.Logger, "interrupted_persist", provider, modelName, sanitizeStats, - ) - - // 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{}{} - } - } - - // Copy existing timestamps and add result timestamps for - // interrupted tool calls so the frontend can show partial - // duration. - toolCallCreatedAt := maps.Clone(result.toolCallCreatedAt) - if toolCallCreatedAt == nil { - toolCallCreatedAt = make(map[string]time.Time) - } - toolResultCreatedAt := maps.Clone(result.toolResultCreatedAt) - if toolResultCreatedAt == nil { - toolResultCreatedAt = make(map[string]time.Time) - } - - // 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...) - - interruptedAt := dbtime.Now() - for _, tc := range result.toolCalls { - if tc.ToolCallID == "" { - continue - } - if _, exists := answeredToolCalls[tc.ToolCallID]; exists { - continue - } - if chatsanitize.IsAnthropicProviderExecutedToolCall(provider, tc) { - continue - } - content = append(content, fantasy.ToolResultContent{ - ToolCallID: tc.ToolCallID, - ToolName: tc.ToolName, - ProviderExecuted: tc.ProviderExecuted, - Result: fantasy.ToolResultOutputContentError{ - Error: xerrors.New(interruptedToolResultErrorMessage), - }, - }) - // Only stamp synthetic results; don't clobber - // timestamps from tools that completed before - // the interruption arrived. - if _, exists := toolResultCreatedAt[tc.ToolCallID]; !exists { - toolResultCreatedAt[tc.ToolCallID] = interruptedAt - } - answeredToolCalls[tc.ToolCallID] = struct{}{} - } - - if len(content) == 0 { - return - } - - persistCtx := context.WithoutCancel(ctx) - if err := opts.PersistStep(persistCtx, PersistedStep{ - Content: content, - ToolCallCreatedAt: toolCallCreatedAt, - ToolResultCreatedAt: toolResultCreatedAt, - ReasoningStartedAt: result.reasoningStartedAt, - ReasoningCompletedAt: result.reasoningCompletedAt, - }); err != nil { - if opts.OnInterruptedPersistError != nil { - opts.OnInterruptedPersistError(err) - } - } -} - -// tryCompactOnExit runs compaction when the chatloop is about -// to exit early (e.g. via ErrDynamicToolCall). The normal -// inline and post-run compaction paths are unreachable in -// early-exit scenarios, so this ensures the context window -// doesn't grow unbounded. -func tryCompactOnExit( - ctx context.Context, - opts RunOptions, - usage fantasy.Usage, - metadata fantasy.ProviderMetadata, -) { - if opts.Compaction == nil || opts.ReloadMessages == nil { - return - } - reloaded, err := opts.ReloadMessages(ctx) - if err != nil { - return - } - did, compactErr := tryCompact( - ctx, - opts.Model, - opts.Compaction, - opts.ContextLimitFallback, - usage, - metadata, - reloaded, - ) - opts.Metrics.RecordCompaction(opts.Model.Provider(), opts.Model.Model(), did, compactErr) - if compactErr != nil && opts.Compaction.OnError != nil { - opts.Compaction.OnError(compactErr) - } -} - func isToolActive(name string, activeTools []string) bool { return len(activeTools) == 0 || slices.Contains(activeTools, name) } -// mergeNewToolNames returns activeTools augmented with any tool names -// from newTools that are not present in oldTools and not already in -// activeTools. This keeps newly injected tools (e.g. via PrepareTools) -// callable even when activeTools is non-empty. -// -// When activeTools is empty, all tools are already active and the slice -// is returned unchanged. -func mergeNewToolNames(activeTools []string, oldTools, newTools []fantasy.AgentTool) []string { - if len(activeTools) == 0 { - return activeTools - } - old := make(map[string]struct{}, len(oldTools)) - for _, t := range oldTools { - old[t.Info().Name] = struct{}{} - } - active := make(map[string]struct{}, len(activeTools)) - for _, name := range activeTools { - active[name] = struct{}{} - } - for _, t := range newTools { - name := t.Info().Name - if _, alreadyActive := active[name]; alreadyActive { - continue - } - if _, existedBefore := old[name]; existedBefore { - continue - } - activeTools = append(activeTools, name) - active[name] = struct{}{} - } - return activeTools -} - // buildToolDefinitions converts AgentTool definitions into the // fantasy.Tool slice expected by fantasy.Call. When activeTools // is non-empty, only function tools whose name appears in the @@ -1901,24 +1381,6 @@ func buildToolDefinitions(tools []fantasy.AgentTool, activeTools []string, provi return prepared } -// shouldStopAfterTools returns true if any tool result in the -// slice matches a name in stopTools and produced a successful -// (non-error) result. -func shouldStopAfterTools(stopTools map[string]struct{}, results []fantasy.ToolResultContent) bool { - if len(stopTools) == 0 { - return false - } - for _, tr := range results { - if _, ok := stopTools[tr.ToolName]; !ok { - continue - } - if _, isErr := tr.Result.(fantasy.ToolResultOutputContentError); !isErr { - return true - } - } - return false -} - func shouldApplyAnthropicPromptCaching(model fantasy.LanguageModel) bool { if model == nil { return false diff --git a/coderd/x/chatd/chatloop/chatloop_internal_test.go b/coderd/x/chatd/chatloop/chatloop_internal_test.go index 1d6ff07560..96825127c9 100644 --- a/coderd/x/chatd/chatloop/chatloop_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_internal_test.go @@ -3,587 +3,16 @@ package chatloop import ( "context" "iter" - "sync" "testing" "charm.land/fantasy" fantasyanthropic "charm.land/fantasy/providers/anthropic" - fantasyopenai "charm.land/fantasy/providers/openai" "github.com/stretchr/testify/require" - "golang.org/x/xerrors" - "cdr.dev/slog/v3/sloggers/slogtest" - "github.com/coder/coder/v2/coderd/x/chatd/chatopenai" - "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/codersdk" + "github.com/coder/quartz" ) -func TestRun_ChainBrokenRecovers(t *testing.T) { - t.Parallel() - - // Given: a chain-mode run whose previous provider_response_id is present in - // our database but no longer recognized by the provider for some reason - var ( - streamCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - disableCalls := 0 - reloadCalls := 0 - reloadedHistory := []fantasy.Message{ - {Role: "system", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "sys"}}}, - {Role: "user", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hello"}}}, - {Role: "assistant", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hi"}}}, - {Role: "user", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "follow up"}}}, - } - - chainFiltered := []fantasy.Message{ - {Role: "system", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "sys"}}}, - {Role: "user", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "follow up"}}}, - } - - // When: the first attempt fails with the chain-broken error - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: chainFiltered, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() { - disableCalls++ - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - reloadCalls++ - return reloadedHistory, nil - }, - }) - - // Then: DisableChainMode and ReloadMessages each run once and the - // retry attempt sends the full reloaded history without - // previous_response_id. - require.NoError(t, err) - require.Equal(t, 2, streamCalls, "exactly two stream attempts (one failure, one success)") - require.Equal(t, 1, disableCalls, "DisableChainMode called once on chain-broken recovery") - require.Equal(t, 1, reloadCalls, "ReloadMessages called once on chain-broken recovery") - - require.False(t, - chatopenai.HasPreviousResponseID(secondCallOpt), - "second attempt must not carry previous_response_id; it was poisoned", - ) - require.Equal(t, reloadedHistory, secondPrompt, - "second attempt must use full reloaded history, not chain-filtered prompt", - ) -} - -func TestRun_ChainBrokenRecoveryPreparesReloadedMessages(t *testing.T) { - t.Parallel() - - var ( - streamCalls int - prepareCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - reloadedHistory := []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "full history"), - } - - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "chain-filtered"), - }, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() {}, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return reloadedHistory, nil - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - prepareCalls++ - return append(msgs, textMessage(fantasy.MessageRoleSystem, "prepared")) - }, - }) - - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - require.Equal(t, 2, prepareCalls, - "reloaded history must be prepared before the retry") - require.False(t, chatopenai.HasPreviousResponseID(secondCallOpt)) - requireTextPrompt(t, secondPrompt, "full history") - requireTextPrompt(t, secondPrompt, "prepared") -} - -func TestRun_ChainBrokenRecoveryAppliesProviderPromptPrep(t *testing.T) { - t.Parallel() - - var ( - streamCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - reloadedHistory := []fantasy.Message{ - textMessage(fantasy.MessageRoleSystem, "sys-1"), - textMessage(fantasy.MessageRoleSystem, "sys-2"), - textMessage(fantasy.MessageRoleUser, "hello"), - textMessage(fantasy.MessageRoleAssistant, "hi"), - textMessage(fantasy.MessageRoleUser, "follow up"), - } - - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleSystem, "sys-2"), - textMessage(fantasy.MessageRoleUser, "follow up"), - }, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() {}, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return reloadedHistory, nil - }, - }) - - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - require.False(t, chatopenai.HasPreviousResponseID(secondCallOpt)) - require.Len(t, secondPrompt, 5) - require.False(t, hasAnthropicEphemeralCacheControl(secondPrompt[0])) - require.True(t, hasAnthropicEphemeralCacheControl(secondPrompt[1])) - require.False(t, hasAnthropicEphemeralCacheControl(secondPrompt[2])) - require.True(t, hasAnthropicEphemeralCacheControl(secondPrompt[3])) - require.True(t, hasAnthropicEphemeralCacheControl(secondPrompt[4])) -} - -func TestRun_ChainBrokenReloadWithoutDisableChainModeIsExplicit(t *testing.T) { - t.Parallel() - - var ( - streamCalls int - prepareCalls int - reloadCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "chain-filtered"), - }, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - reloadCalls++ - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "full history"), - }, nil - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - prepareCalls++ - return append(msgs, textMessage(fantasy.MessageRoleSystem, "prepared")) - }, - // DisableChainMode is intentionally nil. This covers callers - // whose ReloadMessages does not depend on chain-mode state. - }) - - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - require.Equal(t, 1, reloadCalls) - require.Equal(t, 2, prepareCalls) - require.False(t, chatopenai.HasPreviousResponseID(secondCallOpt)) - requireTextPrompt(t, secondPrompt, "full history") - requireTextPrompt(t, secondPrompt, "prepared") -} - -func TestRun_ChainBrokenComposesWithPostStepChainExit(t *testing.T) { - t.Parallel() - - // Given a chain-mode run whose recovery succeeds and yields a - // tool call so the step loop continues - var ( - mu sync.Mutex - streamCalls int - capturedOpts []fantasy.ProviderOptions - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - streamCalls++ - attempt := streamCalls - capturedOpts = append(capturedOpts, call.ProviderOptions) - mu.Unlock() - - switch attempt { - case 1: - // Initial chained attempt: 404 from provider. - return nil, xerrors.New(chainBrokenErrorMessage) - case 2: - // Recovery succeeded; emit a tool call so the - // step loop continues to a second step. - 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: end the run. - return finishingStream(), nil - } - }, - } - - // When the second step builds its call from opts.ProviderOptions - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 3, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hi"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool("read_file"), - }, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() {}, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hi"), - }, nil - }, - }) - - // Then it must not re-send the poisoned previous_response_id - // because chain-broken recovery cleared both the current call and - // subsequent step options. - require.NoError(t, err) - require.Equal(t, 3, streamCalls, - "expected three stream calls: chain-broken failure, recovered tool-call step, follow-up step") - for i, providerOpts := range capturedOpts[1:] { - require.False(t, - chatopenai.HasPreviousResponseID(providerOpts), - "every stream call after recovery (index %d) must have cleared previous_response_id", - i+1, - ) - } -} - -func TestRun_ChainBrokenReloadFailureStillClearsChain(t *testing.T) { - t.Parallel() - - // Given: a chain-mode run whose ReloadMessages callback errors - var ( - streamCalls int - prepareCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - disableCalls := 0 - chainFiltered := []fantasy.Message{ - {Role: "system", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "sys"}}}, - {Role: "user", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "follow up"}}}, - } - - // When: the chain-broken error fires - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: chainFiltered, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() { - disableCalls++ - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return nil, xerrors.New("reload exploded") - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - prepareCalls++ - return append(msgs, textMessage(fantasy.MessageRoleSystem, "prepared")) - }, - }) - - // Then: the poisoned previous_response_id is still cleared and - // DisableChainMode still runs, so the retry has any chance of - // succeeding against the chain-filtered prompt. - require.NoError(t, err) - require.Equal(t, 1, disableCalls) - require.Equal(t, 1, prepareCalls) - require.False(t, - chatopenai.HasPreviousResponseID(secondCallOpt), - "chain options must still be cleared even when reload fails", - ) - requireTextPrompt(t, secondPrompt, "follow up") - requireTextPrompt(t, secondPrompt, "prepared") -} - -func TestRun_ChainBrokenRecoveryDropsOrphanProviderToolCall(t *testing.T) { - t.Parallel() - - var ( - streamCalls int - secondCallOpt fantasy.ProviderOptions - secondPrompt []fantasy.Message - ) - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - ModelName: "claude-test", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - secondCallOpt = call.ProviderOptions - secondPrompt = call.Prompt - return finishingStream(), nil - } - }, - } - - reloadCalls := 0 - err := Run(context.Background(), RunOptions{ - Model: model, - Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "chain-filtered"), - }, - ProviderOptions: chainModeProviderOptions("resp_poisoned"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() {}, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - reloadCalls++ - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search"), - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{ - fantasy.ReasoningPart{ProviderOptions: fantasy.ProviderOptions{fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{RedactedData: "redacted-payload"}}}, - fantasy.ToolCallPart{ToolCallID: "ws-orphan", ToolName: "web_search", Input: `{"query":"coder"}`, ProviderExecuted: true}, - fantasy.TextPart{Text: "partial"}, - }, - }, - textMessage(fantasy.MessageRoleUser, "continue"), - }, nil - }, - }) - - require.NoError(t, err) - require.Equal(t, 1, reloadCalls) - require.Equal(t, 2, streamCalls) - require.False(t, chatopenai.HasPreviousResponseID(secondCallOpt)) - requireNoProviderExecutedToolCallPrompt(t, secondPrompt) - requireAnthropicProviderToolPromptSafe(t, secondPrompt) - requireTextPrompt(t, secondPrompt, "search") - requireTextPrompt(t, secondPrompt, "partial") - requireTextPrompt(t, secondPrompt, "continue") - reasoningPart := requireReasoningPrompt(t, secondPrompt) - reasoningMetadata := fantasyanthropic.GetReasoningMetadata(reasoningPart.ProviderOptions) - require.NotNil(t, reasoningMetadata) - require.Equal(t, "redacted-payload", reasoningMetadata.RedactedData) -} - -func TestRun_ChainBrokenWithoutChainModeIsSafe(t *testing.T) { - t.Parallel() - - // Given: a run with no chain-mode options or callbacks - var streamCalls int - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New(chainBrokenErrorMessage) - default: - return finishingStream(), nil - } - }, - } - - // When: a future provider returns a chain-broken signal, - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - // No ProviderOptions, no DisableChainMode, no ReloadMessages. - }) - - // Then: the recovery branch must no-op (no panic, no missing - // callbacks) and the retry runs normally. - require.NoError(t, err) - require.Equal(t, 2, streamCalls) -} - -func TestRun_NonChainBrokenRetryDoesNotTouchChainState(t *testing.T) { - t.Parallel() - - // Given: a chain-mode run with a still-valid previous_response_id - var ( - streamCalls int - secondCallOpt fantasy.ProviderOptions - ) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - switch streamCalls { - case 1: - return nil, xerrors.New("received status 503 from upstream") - default: - secondCallOpt = call.ProviderOptions - return finishingStream(), nil - } - }, - } - - disableCalls := 0 - reloadCalls := 0 - - // When: a non-chain-broken retryable error fires (503) - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - Messages: []fantasy.Message{ - {Role: "user", Content: []fantasy.MessagePart{fantasy.TextPart{Text: "hi"}}}, - }, - ProviderOptions: chainModeProviderOptions("resp_still_valid"), - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - DisableChainMode: func() { - disableCalls++ - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - reloadCalls++ - return nil, nil - }, - }) - - // Then: chain mode stays engaged, ReloadMessages is not called, - // and the retry preserves previous_response_id. - require.NoError(t, err) - require.Equal(t, 0, disableCalls, - "non-chain-broken retry must not exit chain mode") - require.Equal(t, 0, reloadCalls, - "non-chain-broken retry must not reload history") - require.True(t, - chatopenai.HasPreviousResponseID(secondCallOpt), - "non-chain-broken retry must preserve previous_response_id", - ) -} - func TestProcessStepStreamPreservesReasoningMetadataAcrossNilDelta(t *testing.T) { t.Parallel() @@ -605,7 +34,7 @@ func TestProcessStepStreamPreservesReasoningMetadataAcrossNilDelta(t *testing.T) yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) }) - result, err := processStepStream(context.Background(), stream, func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}) + result, err := processStepStream(context.Background(), stream, quartz.NewMock(t), func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}) require.NoError(t, err) require.Len(t, result.content, 1) reasoning, ok := fantasy.AsContentType[fantasy.ReasoningContent](result.content[0]) @@ -641,7 +70,7 @@ func TestProcessStepStreamPersistsRedactedThinkingOnEnd(t *testing.T) { yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) }) - result, err := processStepStream(context.Background(), stream, func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}) + result, err := processStepStream(context.Background(), stream, quartz.NewMock(t), func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}) require.NoError(t, err) require.Len(t, result.content, 2) reasoning, ok := fantasy.AsContentType[fantasy.ReasoningContent](result.content[0]) @@ -652,40 +81,13 @@ func TestProcessStepStreamPersistsRedactedThinkingOnEnd(t *testing.T) { require.Equal(t, "redacted-payload", metadata.RedactedData) } -func TestStepResultToResponseMessagesPreservesEmptySignedReasoning(t *testing.T) { - t.Parallel() - - result := stepResult{ - content: []fantasy.Content{ - fantasy.ReasoningContent{ - ProviderMetadata: fantasy.ProviderMetadata{ - fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{ - RedactedData: "redacted-payload", - }, - }, - }, - fantasy.TextContent{Text: "done"}, - }, - } - - messages := result.toResponseMessages() - - require.Len(t, messages, 1) - require.Len(t, messages[0].Content, 2) - reasoning, ok := fantasy.AsMessagePart[fantasy.ReasoningPart](messages[0].Content[0]) - require.True(t, ok) - require.Empty(t, reasoning.Text) - metadata := fantasyanthropic.GetReasoningMetadata(reasoning.ProviderOptions) - require.NotNil(t, metadata) - require.Equal(t, "redacted-payload", metadata.RedactedData) -} - func TestFlushActiveStatePreservesEmptySignedReasoning(t *testing.T) { t.Parallel() result := &stepResult{} flushActiveState( result, + quartz.NewMock(t), map[string]string{}, map[string]reasoningState{ "signed": { @@ -709,32 +111,3 @@ func TestFlushActiveStatePreservesEmptySignedReasoning(t *testing.T) { require.NotNil(t, metadata) require.Equal(t, "redacted-payload", metadata.RedactedData) } - -// chainBrokenError is what OpenAI returns when previous_response_id -// points at a response it does not have stored. -const chainBrokenErrorMessage = "Previous response with id 'resp_abc' not found." - -// finishingStream returns a stream that emits a single Finish part. -// The chatloop treats a finishReason of Stop as "stoppedByModel" and -// exits the per-step loop after persisting. -func finishingStream() fantasy.StreamResponse { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }) - }) -} - -// chainModeProviderOptions builds a fantasy.ProviderOptions carrying -// the OpenAI Responses options with previous_response_id set, the same -// shape chatd builds when chain mode is active. -func chainModeProviderOptions(previousResponseID string) fantasy.ProviderOptions { - store := true - return fantasy.ProviderOptions{ - fantasyopenai.Name: &fantasyopenai.ResponsesProviderOptions{ - Store: &store, - PreviousResponseID: &previousResponseID, - }, - } -} diff --git a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go index 9769f10d01..4e0c6bf55a 100644 --- a/coderd/x/chatd/chatloop/chatloop_run_internal_test.go +++ b/coderd/x/chatd/chatloop/chatloop_run_internal_test.go @@ -5,7 +5,6 @@ import ( "encoding/base64" "errors" "iter" - "strings" "sync" "sync/atomic" "testing" @@ -21,9 +20,7 @@ import ( "golang.org/x/xerrors" "cdr.dev/slog/v3" - "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" - "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/chattest" "github.com/coder/coder/v2/codersdk" @@ -31,8 +28,6 @@ import ( "github.com/coder/quartz" ) -const activeToolName = "read_file" - func validWebSearchProviderMetadataForTest() fantasy.ProviderMetadata { return fantasy.ProviderMetadata{ fantasyanthropic.Name: &fantasyanthropic.WebSearchResultMetadata{ @@ -77,30 +72,6 @@ func safeToolResultContent(block fantasy.Content) (fantasy.ToolResultContent, bo } } -func safeToolCallPart(part fantasy.MessagePart) (fantasy.ToolCallPart, bool) { - var zero fantasy.ToolCallPart - if part == nil { - return zero, false - } - if value, ok := part.(*fantasy.ToolCallPart); ok && value == nil { - return zero, false - } - type toolCallPart = fantasy.ToolCallPart - return fantasy.AsMessagePart[toolCallPart](part) -} - -func safeToolResultPart(part fantasy.MessagePart) (fantasy.ToolResultPart, bool) { - var zero fantasy.ToolResultPart - if part == nil { - return zero, false - } - if value, ok := part.(*fantasy.ToolResultPart); ok && value == nil { - return zero, false - } - type toolResultPart = fantasy.ToolResultPart - return fantasy.AsMessagePart[toolResultPart](part) -} - func toolCallContentToPart(toolCall fantasy.ToolCallContent) fantasy.ToolCallPart { return fantasy.ToolCallPart{ ToolCallID: toolCall.ToolCallID, @@ -120,467 +91,6 @@ func toolResultContentToPart(toolResult fantasy.ToolResultContent) fantasy.ToolR } } -func awaitRunResult(ctx context.Context, t *testing.T, done <-chan error) error { - t.Helper() - - select { - case err := <-done: - return err - case <-ctx.Done(): - t.Fatal("timed out waiting for Run to complete") - return nil - } -} - -func TestRun_ActiveToolsPrepareBehavior(t *testing.T) { - t.Parallel() - - var capturedCall fantasy.Call - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - capturedCall = call - 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}, - }), nil - }, - } - - persistStepCalls := 0 - var persistedStep PersistedStep - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleSystem, "sys-1"), - textMessage(fantasy.MessageRoleSystem, "sys-2"), - textMessage(fantasy.MessageRoleUser, "hello"), - textMessage(fantasy.MessageRoleAssistant, "working"), - textMessage(fantasy.MessageRoleUser, "continue"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool(activeToolName), - newNoopTool("write_file"), - }, - MaxSteps: 3, - ActiveTools: []string{activeToolName}, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistStepCalls++ - persistedStep = step - return nil - }, - }) - require.NoError(t, err) - - require.Equal(t, 1, persistStepCalls) - require.True(t, persistedStep.ContextLimit.Valid) - require.Equal(t, int64(4096), persistedStep.ContextLimit.Int64) - require.GreaterOrEqual(t, persistedStep.Runtime, time.Duration(0), - "step runtime should be non-negative") - - require.NotEmpty(t, capturedCall.Prompt) - require.False(t, containsPromptSentinel(capturedCall.Prompt)) - require.Len(t, capturedCall.Tools, 1) - require.Equal(t, activeToolName, capturedCall.Tools[0].GetName()) - - require.Len(t, capturedCall.Prompt, 5) - require.False(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[0])) - require.True(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[1])) - require.False(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[2])) - require.True(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[3])) - require.True(t, hasAnthropicEphemeralCacheControl(capturedCall.Prompt[4])) -} - -func TestRun_ActiveToolsRejectsDisallowedExecution(t *testing.T) { - t.Parallel() - - var blockedCalls atomic.Int32 - blockedToolName := "write_file" - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-blocked", ToolCallName: blockedToolName}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-blocked", Delta: `{"path":"/tmp/nope"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-blocked"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-blocked", - ToolCallName: blockedToolName, - ToolCallInput: `{"path":"/tmp/nope"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - blockedTool := fantasy.NewAgentTool( - blockedToolName, - "blocked tool", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - blockedCalls.Add(1) - return fantasy.NewTextResponse("should not run"), nil - }, - ) - - var persistedStep PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "try the blocked tool"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool(activeToolName), - blockedTool, - }, - ActiveTools: []string{activeToolName}, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedStep = step - return nil - }, - }) - require.NoError(t, err) - require.Zero(t, blockedCalls.Load(), "disallowed tool must not execute") - - var foundToolError bool - for _, block := range persistedStep.Content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok || toolResult.ToolName != blockedToolName { - continue - } - errResult, ok := toolResult.Result.(fantasy.ToolResultOutputContentError) - require.True(t, ok) - assert.EqualError(t, errResult.Error, "Tool not active in this turn: "+blockedToolName) - foundToolError = true - } - require.True(t, foundToolError, "persisted step should include the rejected tool result") -} - -func TestRun_ActiveToolsAllowsProviderRunnerExecution(t *testing.T) { - t.Parallel() - - providerRunnerName := "computer" - var runnerCalls atomic.Int32 - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-provider-runner", ToolCallName: providerRunnerName}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-provider-runner", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-provider-runner"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-provider-runner", - ToolCallName: providerRunnerName, - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - runnerTool := fantasy.NewAgentTool( - providerRunnerName, - "provider runner", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - runnerCalls.Add(1) - return fantasy.NewTextResponse("ran provider runner"), nil - }, - ) - - var persistedStep PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "use the computer"), - }, - Tools: []fantasy.AgentTool{newNoopTool(activeToolName)}, - ActiveTools: []string{activeToolName}, - ProviderTools: []ProviderTool{ - { - Definition: fantasy.FunctionTool{ - Name: providerRunnerName, - Description: "provider runner", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{}, - }, - }, - Runner: runnerTool, - }, - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedStep = step - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, int32(1), runnerCalls.Load(), - "provider runner should execute even when omitted from active tools") - - var foundToolResult bool - for _, block := range persistedStep.Content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok || toolResult.ToolName != providerRunnerName { - continue - } - textResult, ok := toolResult.Result.(fantasy.ToolResultOutputContentText) - require.True(t, ok) - assert.Equal(t, "ran provider runner", textResult.Text) - foundToolResult = true - } - require.True(t, foundToolResult, - "persisted step should include the provider runner result") -} - -func TestRun_ProviderToolResultProviderMetadata(t *testing.T) { - t.Parallel() - - expectedMetadata := fantasy.ProviderMetadata{ - "openai": &testProviderData{data: map[string]any{ - "detail": "original", - }}, - } - - tests := []struct { - name string - callback func(fantasy.ToolResponse) fantasy.ProviderMetadata - want fantasy.ProviderMetadata - }{ - { - name: "callback returns metadata", - callback: func(fantasy.ToolResponse) fantasy.ProviderMetadata { - return expectedMetadata - }, - want: expectedMetadata, - }, - { - name: "callback nil", - want: nil, - }, - { - name: "callback returns nil", - callback: func(fantasy.ToolResponse) fantasy.ProviderMetadata { - return nil - }, - want: nil, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - providerRunnerName := "computer" - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-provider-runner", ToolCallName: providerRunnerName}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-provider-runner", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-provider-runner"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-provider-runner", - ToolCallName: providerRunnerName, - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - runnerTool := fantasy.NewAgentTool( - providerRunnerName, - "provider runner", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.ToolResponse{ - Type: "image", - Data: []byte("image bytes"), - MediaType: "image/png", - Content: "screenshot", - }, nil - }, - ) - - var persistedStep PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "use the computer"), - }, - ProviderTools: []ProviderTool{ - { - Definition: fantasy.FunctionTool{ - Name: providerRunnerName, - Description: "provider runner", - InputSchema: map[string]any{ - "type": "object", - "properties": map[string]any{}, - }, - }, - Runner: runnerTool, - ResultProviderMetadata: tt.callback, - }, - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedStep = step - return nil - }, - }) - require.NoError(t, err) - - var foundResult fantasy.ToolResultContent - for _, block := range persistedStep.Content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok || toolResult.ToolName != providerRunnerName { - continue - } - foundResult = toolResult - break - } - require.NotEmpty(t, foundResult.ToolCallID, - "persisted step should include the provider runner result") - - mediaResult, ok := foundResult.Result.(fantasy.ToolResultOutputContentMedia) - require.True(t, ok, "expected media result") - assert.Equal(t, "image/png", mediaResult.MediaType) - assert.Equal(t, tt.want, foundResult.ProviderMetadata) - - if tt.want == nil { - return - } - - messages := stepResult{content: persistedStep.Content}.toResponseMessages() - require.Len(t, messages, 2) - require.Equal(t, fantasy.MessageRoleTool, messages[1].Role) - require.Len(t, messages[1].Content, 1) - - resultPart, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](messages[1].Content[0]) - require.True(t, ok, "expected outbound tool result part") - assert.Equal(t, fantasy.ProviderOptions(tt.want), resultPart.ProviderOptions) - }) - } -} - -func TestProcessStepStream_AnthropicUsageMatchesFinalDelta(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - 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: "cached response"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - { - Type: fantasy.StreamPartTypeFinish, - Usage: fantasy.Usage{ - InputTokens: 200, - OutputTokens: 75, - TotalTokens: 275, - CacheCreationTokens: 30, - CacheReadTokens: 150, - ReasoningTokens: 0, - }, - FinishReason: fantasy.FinishReasonStop, - }, - }), nil - }, - } - - var persistedStep PersistedStep - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedStep = step - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, int64(200), persistedStep.Usage.InputTokens) - require.Equal(t, int64(75), persistedStep.Usage.OutputTokens) - require.Equal(t, int64(275), persistedStep.Usage.TotalTokens) - require.Equal(t, int64(30), persistedStep.Usage.CacheCreationTokens) - require.Equal(t, int64(150), persistedStep.Usage.CacheReadTokens) -} - -func TestRun_OnRetryEnrichesProvider(t *testing.T) { - t.Parallel() - - type retryRecord struct { - attempt int - errMsg string - classified chatretry.ClassifiedError - delay time.Duration - } - - var records []retryRecord - calls := 0 - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - calls++ - if calls == 1 { - return nil, xerrors.New("received status 429 from upstream") - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - OnRetry: func( - attempt int, - retryErr error, - classified chatretry.ClassifiedError, - delay time.Duration, - ) { - records = append(records, retryRecord{ - attempt: attempt, - errMsg: retryErr.Error(), - classified: classified, - delay: delay, - }) - }, - }) - require.NoError(t, err) - require.Len(t, records, 1) - require.Equal(t, 1, records[0].attempt) - require.Equal(t, "received status 429 from upstream", records[0].errMsg) - require.Equal(t, chatretry.Delay(0), records[0].delay) - require.Equal(t, "openai", records[0].classified.Provider) - require.Equal(t, codersdk.ChatErrorKindRateLimit, records[0].classified.Kind) - require.True(t, records[0].classified.Retryable) - require.Equal(t, 429, records[0].classified.StatusCode) - require.Equal( - t, - "OpenAI is rate limiting requests.", - records[0].classified.Message, - ) -} - func TestStreamSilenceGuard_DisarmAndFireRace(t *testing.T) { t.Parallel() @@ -638,546 +148,236 @@ func TestStreamSilenceGuard_DisarmPreservesPermanentError(t *testing.T) { require.Nil(t, context.Cause(attemptCtx)) } -func TestRun_RetriesSilenceTimeoutWhileOpeningStream(t *testing.T) { +func TestGenerateAssistant_ProviderContextSurvivesStreamError(t *testing.T) { t.Parallel() - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitShort, - ) - defer cancel() - - mClock := quartz.NewMock(t) - trap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer trap.Close() - - attempts := 0 - attemptCause := make(chan error, 1) - var retries []chatretry.ClassifiedError model := &chattest.FakeModel{ ProviderName: "openai", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - <-ctx.Done() - attemptCause <- context.Cause(ctx) - return nil, ctx.Err() - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil + ModelName: "test-model", + StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + return nil, xerrors.New("upstream returned status 400") }, } - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - OnRetry: func( - _ int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, classified) - }, - }) - }() - - trap.MustWait(ctx).MustRelease(ctx) - mClock.Advance(silenceTimeout).MustWait(ctx) - trap.MustWait(ctx).MustRelease(ctx) - - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 2, attempts) - require.Len(t, retries, 1) - require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, retries[0].Kind) - require.True(t, retries[0].Retryable) - require.Equal(t, "openai", retries[0].Provider) - require.Equal( - t, - "OpenAI did not send response data in time.", - retries[0].Message, - ) - select { - case cause := <-attemptCause: - require.ErrorIs(t, cause, errStreamSilenceTimeout) - case <-ctx.Done(): - t.Fatal("timed out waiting for silence timeout cause") - } + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + Messages: []fantasy.Message{ + textMessage(fantasy.MessageRoleUser, "hello"), + }, + }) + require.Error(t, err) + classified := chaterror.Classify(err) + require.Equal(t, "openai", classified.Provider) + require.Equal(t, "OpenAI returned an unexpected error.", classified.Message) } -// TestRun_HTTP2TransportErrorClassifiedAsRetryableTimeout proves the -// provider comes from Model.Provider() (not from sniffing the error -// text) by using an error string with no provider hint and running -// the same assertion across two providers. -func TestRun_HTTP2TransportErrorClassifiedAsRetryableTimeout(t *testing.T) { +func TestGenerateAssistant_HTTP2TransportErrorClassifiedAsRetryableTimeout(t *testing.T) { t.Parallel() - providers := []string{"anthropic", "openai"} - for _, provider := range providers { + for _, provider := range []string{"anthropic", "openai"} { + provider := provider t.Run(provider, func(t *testing.T) { t.Parallel() - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitShort, - ) - defer cancel() - - mClock := quartz.NewMock(t) - trap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer trap.Close() - - attempts := 0 - var retries []chatretry.ClassifiedError model := &chattest.FakeModel{ ProviderName: provider, + ModelName: "test-model", StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - // Bare transport error; Provider must - // come from Model.Provider(). - return nil, xerrors.New( - "http2: client connection force closed via ClientConn.Close", - ) - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil + return nil, xerrors.New("http2: client connection force closed via ClientConn.Close") }, } - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - OnRetry: func( - _ int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, classified) - }, - }) - }() - - // One guard per attempt. - trap.MustWait(ctx).MustRelease(ctx) - trap.MustWait(ctx).MustRelease(ctx) - - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 2, attempts) - require.Len(t, retries, 1) - require.Equal(t, codersdk.ChatErrorKindTimeout, retries[0].Kind, "Kind") - require.True(t, retries[0].Retryable, "Retryable") - require.Equal(t, provider, retries[0].Provider, "Provider") + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + }) + require.Error(t, err) + classified := chaterror.Classify(err) + require.Equal(t, codersdk.ChatErrorKindTimeout, classified.Kind) + require.Equal(t, provider, classified.Provider) + require.True(t, classified.Retryable) }) } } -func TestRun_RetriesProviderContextCanceledStreamError(t *testing.T) { +func TestGenerateAssistant_StreamSilenceTimeoutRetryClassification(t *testing.T) { t.Parallel() - attempts := 0 - retryErrs := make(chan error, chatretry.MaxAttempts) - retries := make(chan chatretry.ClassifiedError, chatretry.MaxAttempts) - var persisted []fantasy.Content - ctx := testutil.Context(t, testutil.WaitShort) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "partial"}, - {Type: fantasy.StreamPartTypeError, Error: context.Canceled}, - }), nil - } - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-2"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-2", Delta: "done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-2"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } + t.Run("timeout while opening stream", func(t *testing.T) { + t.Parallel() - err := Run(ctx, RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, step PersistedStep) error { - persisted = append([]fantasy.Content(nil), step.Content...) - return nil - }, - OnRetry: func( - _ int, - retryErr error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retryErrs <- retryErr - retries <- classified - }, - }) - require.NoError(t, err) - require.Equal(t, 2, attempts) - require.Len(t, retryErrs, 1) - require.Len(t, retries, 1) - retryErr := testutil.RequireReceive(ctx, t, retryErrs) - classified := testutil.RequireReceive(ctx, t, retries) - require.ErrorIs(t, retryErr, chaterror.ErrProviderTransportReset) - require.ErrorIs(t, retryErr, context.Canceled) - require.Equal(t, codersdk.ChatErrorKindTimeout, classified.Kind) - require.True(t, classified.Retryable) - require.Equal(t, "openai", classified.Provider) - require.Equal(t, "OpenAI is temporarily unavailable.", classified.Message) - - text := requireTextContent(t, persisted, "done") - require.Equal(t, "done", text.Text) - for _, block := range persisted { - if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { - require.NotContains(t, text.Text, "partial") - } - } -} - -func TestRun_RetriesSilenceTimeoutBeforeFirstPart(t *testing.T) { - t.Parallel() - - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitShort, - ) - defer cancel() - - mClock := quartz.NewMock(t) - trap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer trap.Close() - - attempts := 0 - attemptCause := make(chan error, 1) - var retries []chatretry.ClassifiedError - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { + ctx := testutil.Context(t, testutil.WaitLong) + const silenceTimeout = 5 * time.Millisecond + clock := quartz.NewMock(t) + trap := clock.Trap().AfterFunc(streamSilenceGuardTimerTag) + defer trap.Close() + var calls atomic.Int32 + model := &chattest.FakeModel{ + ProviderName: "openai", + ModelName: "test-model", + StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + if calls.Add(1) == 1 { <-ctx.Done() - attemptCause <- context.Cause(ctx) - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil - }, - } - - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil + return nil, ctx.Err() + } + return streamFromParts([]fantasy.StreamPart{{ + Type: fantasy.StreamPartTypeFinish, + FinishReason: fantasy.FinishReasonStop, + }}), nil }, - OnRetry: func( - _ int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, classified) + } + done := make(chan error, 1) + go func() { + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + Clock: clock, + StreamSilenceTimeout: silenceTimeout, + }) + done <- err + }() + + trap.MustWait(ctx).MustRelease(ctx) + _, waiter := clock.AdvanceNext() + waiter.MustWait(ctx) + require.Error(t, <-done) + require.Equal(t, int32(1), calls.Load()) + }) + + t.Run("timeout before first part", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + const silenceTimeout = 5 * time.Millisecond + clock := quartz.NewMock(t) + trap := clock.Trap().AfterFunc(streamSilenceGuardTimerTag) + defer trap.Close() + var calls atomic.Int32 + model := &chattest.FakeModel{ + ProviderName: "openai", + ModelName: "test-model", + StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + calls.Add(1) + return func(yield func(fantasy.StreamPart) bool) { + <-ctx.Done() + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeError, Error: ctx.Err()}) + }, nil }, - }) - }() + } + done := make(chan error, 1) + go func() { + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + Clock: clock, + StreamSilenceTimeout: silenceTimeout, + }) + done <- err + }() - trap.MustWait(ctx).MustRelease(ctx) - mClock.Advance(silenceTimeout).MustWait(ctx) - trap.MustWait(ctx).MustRelease(ctx) + trap.MustWait(ctx).MustRelease(ctx) + _, waiter := clock.AdvanceNext() + waiter.MustWait(ctx) + err := <-done + require.Error(t, err) + classified := chaterror.Classify(err) + require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, classified.Kind) + require.Equal(t, "openai", classified.Provider) + require.True(t, classified.Retryable) + require.Equal(t, int32(1), calls.Load()) + }) - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 2, attempts) - require.Len(t, retries, 1) - require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, retries[0].Kind) - require.True(t, retries[0].Retryable) - require.Equal(t, "openai", retries[0].Provider) - require.Equal( - t, - "OpenAI did not send response data in time.", - retries[0].Message, - ) - select { - case cause := <-attemptCause: - require.ErrorIs(t, cause, errStreamSilenceTimeout) - case <-ctx.Done(): - t.Fatal("timed out waiting for silence timeout cause") - } -} + t.Run("first part disarms timeout", func(t *testing.T) { + t.Parallel() -func TestRun_StreamPartsResetSilenceTimeout(t *testing.T) { - t.Parallel() - - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitShort, - ) - defer cancel() - - mClock := quartz.NewMock(t) - armTrap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer armTrap.Close() - resetTrap := mClock.Trap().TimerReset(streamSilenceGuardTimerTag) - defer resetTrap.Close() - - attempts := 0 - retried := false - firstPartYielded := make(chan struct{}, 1) - secondPartYielded := make(chan struct{}, 1) - continueToSecond := make(chan struct{}) - continueToFinish := make(chan struct{}) - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}) { - return - } - select { - case firstPartYielded <- struct{}{}: - default: - } - - select { - case <-continueToSecond: - case <-ctx.Done(): - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - return - } - - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeTextDelta, - ID: "text-1", - Delta: "done", - }) { - return - } - select { - case secondPartYielded <- struct{}{}: - default: - } - - select { - case <-continueToFinish: - case <-ctx.Done(): - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - return - } - - parts := []fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - } - for _, part := range parts { - if !yield(part) { - return - } - } - }), nil - }, - } - - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - OnRetry: func( - _ int, - _ error, - _ chatretry.ClassifiedError, - _ time.Duration, - ) { - retried = true - }, - }) - }() - - armTrap.MustWait(ctx).MustRelease(ctx) - resetTrap.MustWait(ctx).MustRelease(ctx) - select { - case <-firstPartYielded: - case <-ctx.Done(): - t.Fatal("timed out waiting for first stream part") - } - - mClock.Advance(silenceTimeout / 2).MustWait(ctx) - close(continueToSecond) - resetTrap.MustWait(ctx).MustRelease(ctx) - select { - case <-secondPartYielded: - case <-ctx.Done(): - t.Fatal("timed out waiting for second stream part") - } - - mClock.Advance(silenceTimeout / 2).MustWait(ctx) - close(continueToFinish) - resetTrap.MustWait(ctx).MustRelease(ctx) - resetTrap.MustWait(ctx).MustRelease(ctx) - - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 1, attempts) - require.False(t, retried) -} - -func TestRun_RetriesSilenceTimeoutBetweenParts(t *testing.T) { - t.Parallel() - - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitLong, - ) - defer cancel() - - mClock := quartz.NewMock(t) - armTrap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer armTrap.Close() - resetTrap := mClock.Trap().TimerReset(streamSilenceGuardTimerTag) - defer resetTrap.Close() - - attempts := 0 - firstPartYielded := make(chan struct{}, 1) - attemptCause := make(chan error, 1) - var retries []chatretry.ClassifiedError - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { + ctx := testutil.Context(t, testutil.WaitLong) + const silenceTimeout = 5 * time.Millisecond + clock := quartz.NewMock(t) + trap := clock.Trap().AfterFunc(streamSilenceGuardTimerTag) + defer trap.Close() + var calls atomic.Int32 + continueStream := make(chan struct{}) + model := &chattest.FakeModel{ + ProviderName: "openai", + ModelName: "test-model", + StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + calls.Add(1) + return func(yield func(fantasy.StreamPart) bool) { if !yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}) { return } select { - case firstPartYielded <- struct{}{}: - default: + case <-continueStream: + case <-ctx.Done(): + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeError, Error: ctx.Err()}) + return } + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "done"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}) + yield(fantasy.StreamPart{Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}) + }, nil + }, + } + done := make(chan error, 1) + go func() { + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + Clock: clock, + StreamSilenceTimeout: silenceTimeout, + }) + done <- err + }() + trap.MustWait(ctx).MustRelease(ctx) + close(continueStream) + require.NoError(t, <-done) + require.Equal(t, int32(1), calls.Load()) + }) + + t.Run("silent stream close after timeout", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + const silenceTimeout = 5 * time.Millisecond + clock := quartz.NewMock(t) + trap := clock.Trap().AfterFunc(streamSilenceGuardTimerTag) + defer trap.Close() + var calls atomic.Int32 + model := &chattest.FakeModel{ + ProviderName: "openai", + ModelName: "test-model", + StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { + calls.Add(1) + return func(func(fantasy.StreamPart) bool) { <-ctx.Done() - attemptCause <- context.Cause(ctx) - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil - }, - } - - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil + }, nil }, - OnRetry: func( - _ int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, classified) - }, - }) - }() + } + done := make(chan error, 1) + go func() { + _, err := GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, + Clock: clock, + StreamSilenceTimeout: silenceTimeout, + }) + done <- err + }() - armTrap.MustWait(ctx).MustRelease(ctx) - resetTrap.MustWait(ctx).MustRelease(ctx) - select { - case <-firstPartYielded: - case <-ctx.Done(): - t.Fatal("timed out waiting for first stream part") - } - - mClock.Advance(silenceTimeout).MustWait(ctx) - armTrap.MustWait(ctx).MustRelease(ctx) - resetTrap.MustWait(ctx).MustRelease(ctx) - - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 2, attempts) - require.Len(t, retries, 1) - require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, retries[0].Kind) - require.True(t, retries[0].Retryable) - require.Equal(t, "openai", retries[0].Provider) - select { - case cause := <-attemptCause: - require.ErrorIs(t, cause, errStreamSilenceTimeout) - case <-ctx.Done(): - t.Fatal("timed out waiting for silence timeout cause") - } + trap.MustWait(ctx).MustRelease(ctx) + _, waiter := clock.AdvanceNext() + waiter.MustWait(ctx) + err := <-done + require.Error(t, err) + classified := chaterror.Classify(err) + require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, classified.Kind) + require.Equal(t, int32(1), calls.Load()) + }) } -func TestRun_PanicInPublishMessagePartReleasesAttempt(t *testing.T) { +func TestGenerateAssistant_PanicInPublishMessagePartReleasesAttempt(t *testing.T) { t.Parallel() attemptReleased := make(chan struct{}) model := &chattest.FakeModel{ ProviderName: "openai", + ModelName: "test-model", StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { go func() { <-ctx.Done() @@ -1200,216 +400,14 @@ func TestRun_PanicInPublishMessagePartReleasesAttempt(t *testing.T) { } }() - _ = Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, + _, _ = GenerateAssistant(context.Background(), GenerateAssistantOptions{ + Model: model, PublishMessagePart: func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) { panic("publish panic") }, }) - t.Fatal("expected Run to panic") -} - -func TestRun_RetriesSilenceTimeoutWhenStreamStaysSilent(t *testing.T) { - t.Parallel() - - const silenceTimeout = 5 * time.Millisecond - - ctx, cancel := context.WithTimeout( - context.Background(), - testutil.WaitShort, - ) - defer cancel() - - mClock := quartz.NewMock(t) - trap := mClock.Trap().AfterFunc(streamSilenceGuardTimerTag) - defer trap.Close() - - attempts := 0 - attemptCause := make(chan error, 1) - var retries []chatretry.ClassifiedError - model := &chattest.FakeModel{ - ProviderName: "openai", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - attempts++ - if attempts == 1 { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - <-ctx.Done() - attemptCause <- context.Cause(ctx) - }), nil - } - return streamFromParts([]fantasy.StreamPart{{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }}), nil - }, - } - - done := make(chan error, 1) - go func() { - done <- Run(context.Background(), RunOptions{ - Model: model, - MaxSteps: 1, - StreamSilenceTimeout: silenceTimeout, - Clock: mClock, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - OnRetry: func( - _ int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, classified) - }, - }) - }() - - trap.MustWait(ctx).MustRelease(ctx) - mClock.Advance(silenceTimeout).MustWait(ctx) - trap.MustWait(ctx).MustRelease(ctx) - - require.NoError(t, awaitRunResult(ctx, t, done)) - require.Equal(t, 2, attempts) - require.Len(t, retries, 1) - require.Equal(t, codersdk.ChatErrorKindStreamSilenceTimeout, retries[0].Kind) - require.True(t, retries[0].Retryable) - require.Equal(t, "openai", retries[0].Provider) - require.Equal( - t, - "OpenAI did not send response data in time.", - retries[0].Message, - ) - select { - case cause := <-attemptCause: - require.ErrorIs(t, cause, errStreamSilenceTimeout) - case <-ctx.Done(): - t.Fatal("timed out waiting for silence timeout cause") - } -} - -func TestRun_InterruptedStepPersistsSyntheticToolResult(t *testing.T) { - t.Parallel() - - started := make(chan struct{}) - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - parts := []fantasy.StreamPart{ - { - Type: fantasy.StreamPartTypeToolInputStart, - ID: "interrupt-tool-1", - ToolCallName: "read_file", - }, - { - Type: fantasy.StreamPartTypeToolInputDelta, - ID: "interrupt-tool-1", - ToolCallName: "read_file", - Delta: `{"path":"main.go"`, - }, - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "partial assistant output"}, - } - for _, part := range parts { - if !yield(part) { - return - } - } - - select { - case <-started: - default: - close(started) - } - - <-ctx.Done() - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - }, - } - - ctx, cancel := context.WithCancelCause(context.Background()) - defer cancel(nil) - - go func() { - <-started - cancel(ErrInterrupted) - }() - - persistedAssistantCtxErr := xerrors.New("unset") - var persistedContent []fantasy.Content - var persistedStep PersistedStep - - err := Run(ctx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool("read_file"), - }, - MaxSteps: 3, - PersistStep: func(persistCtx context.Context, step PersistedStep) error { - persistedAssistantCtxErr = persistCtx.Err() - persistedContent = append([]fantasy.Content(nil), step.Content...) - persistedStep = step - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - require.NoError(t, persistedAssistantCtxErr) - - require.NotEmpty(t, persistedContent) - var ( - foundText bool - foundToolCall bool - foundToolResult bool - ) - for _, block := range persistedContent { - if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { - if strings.Contains(text.Text, "partial assistant output") { - foundText = true - } - continue - } - if toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block); ok { - if toolCall.ToolCallID == "interrupt-tool-1" && - toolCall.ToolName == "read_file" && - strings.Contains(toolCall.Input, `"path":"main.go"`) { - foundToolCall = true - } - continue - } - if toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok { - if toolResult.ToolCallID == "interrupt-tool-1" && - toolResult.ToolName == "read_file" { - _, isErr := toolResult.Result.(fantasy.ToolResultOutputContentError) - require.True(t, isErr, "interrupted tool result should be an error") - foundToolResult = true - } - } - } - require.True(t, foundText) - require.True(t, foundToolCall) - require.True(t, foundToolResult) - - // The interrupted tool was flushed mid-stream (never reached - // StreamPartTypeToolCall), so it has no call timestamp. - // But the synthetic error result must have a result timestamp. - require.Contains(t, persistedStep.ToolResultCreatedAt, "interrupt-tool-1", - "interrupted tool result must have a result timestamp") - require.NotContains(t, persistedStep.ToolCallCreatedAt, "interrupt-tool-1", - "interrupted tool should have no call timestamp (never reached StreamPartTypeToolCall)") + t.Fatal("expected GenerateAssistant to panic") } func requireToolResultErrorMessage( @@ -1435,16 +433,6 @@ func streamFromParts(parts []fantasy.StreamPart) fantasy.StreamResponse { }) } -func newNoopTool(name string) fantasy.AgentTool { - return fantasy.NewAgentTool( - name, - "test noop tool", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.ToolResponse{}, nil - }, - ) -} - func textMessage(role fantasy.MessageRole, text string) fantasy.Message { return fantasy.Message{ Role: role, @@ -1454,71 +442,6 @@ func textMessage(role fantasy.MessageRole, text string) fantasy.Message { } } -func requireNoProviderExecutedToolCallContent(t *testing.T, content []fantasy.Content) { - t.Helper() - - for i, block := range content { - toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block) - if ok && toolCall.ProviderExecuted { - t.Fatalf("content[%d]: unexpected provider-executed call", i) - } - } -} - -func requireNoProviderExecutedToolResultContent(t *testing.T, content []fantasy.Content) { - t.Helper() - - for i, block := range content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if ok && toolResult.ProviderExecuted { - t.Fatalf("content[%d]: unexpected provider-executed result", i) - } - } -} - -func requireReasoningPrompt(t *testing.T, prompt []fantasy.Message) fantasy.ReasoningPart { - t.Helper() - - for _, message := range prompt { - for _, part := range message.Content { - reasoningPart, ok := fantasy.AsMessagePart[fantasy.ReasoningPart](part) - if ok { - return reasoningPart - } - } - } - t.Fatal("missing prompt reasoning") - return fantasy.ReasoningPart{} -} - -func requireTextPrompt(t *testing.T, prompt []fantasy.Message, text string) fantasy.TextPart { - t.Helper() - - for _, message := range prompt { - for _, part := range message.Content { - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](part) - if ok && textPart.Text == text { - return textPart - } - } - } - t.Fatalf("missing prompt text %q", text) - return fantasy.TextPart{} -} - -func requireNoProviderExecutedToolCallPrompt(t *testing.T, prompt []fantasy.Message) { - t.Helper() - - for i, message := range prompt { - for j, part := range message.Content { - toolCall, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](part) - if ok && toolCall.ProviderExecuted { - t.Fatalf("prompt[%d].content[%d]: unexpected provider-executed call", i, j) - } - } - } -} - func requireTextContent(t *testing.T, content []fantasy.Content, text string) fantasy.TextContent { t.Helper() @@ -1532,634 +455,6 @@ func requireTextContent(t *testing.T, content []fantasy.Content, text string) fa return fantasy.TextContent{} } -func requireToolCallContent(t *testing.T, content []fantasy.Content, id, name string) fantasy.ToolCallContent { - t.Helper() - - for _, block := range content { - toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block) - if ok && toolCall.ToolCallID == id && toolCall.ToolName == name { - return toolCall - } - } - t.Fatalf("missing tool call %q", id) - return fantasy.ToolCallContent{} -} - -func requireToolResultContent(t *testing.T, content []fantasy.Content, id, name string) fantasy.ToolResultContent { - t.Helper() - - for _, block := range content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if ok && toolResult.ToolCallID == id && toolResult.ToolName == name { - return toolResult - } - } - t.Fatalf("missing tool result %q", id) - return fantasy.ToolResultContent{} -} - -func requireToolResultPrompt(t *testing.T, prompt []fantasy.Message, id string) fantasy.ToolResultPart { - t.Helper() - - for _, message := range prompt { - for _, part := range message.Content { - toolResult, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](part) - if ok && toolResult.ToolCallID == id { - return toolResult - } - } - } - t.Fatalf("missing prompt tool result %q", id) - return fantasy.ToolResultPart{} -} - -func requireNoProviderExecutedToolResultPrompt(t *testing.T, prompt []fantasy.Message) { - t.Helper() - - for i, message := range prompt { - for j, part := range message.Content { - toolResult, ok := safeToolResultPart(part) - if ok && toolResult.ProviderExecuted { - t.Fatalf("prompt[%d].content[%d]: unexpected provider-executed result", i, j) - } - } - } -} - -func requireProviderExecutedToolCallPrompt( - t *testing.T, - prompt []fantasy.Message, - id string, -) fantasy.ToolCallPart { - t.Helper() - - for _, message := range prompt { - for _, part := range message.Content { - toolCall, ok := safeToolCallPart(part) - if ok && toolCall.ProviderExecuted && toolCall.ToolCallID == id { - return toolCall - } - } - } - t.Fatalf("missing provider-executed prompt tool call %q", id) - return fantasy.ToolCallPart{} -} - -func requireProviderExecutedToolResultPrompt( - t *testing.T, - prompt []fantasy.Message, - id string, -) fantasy.ToolResultPart { - t.Helper() - - for _, message := range prompt { - for _, part := range message.Content { - toolResult, ok := safeToolResultPart(part) - if ok && toolResult.ProviderExecuted && toolResult.ToolCallID == id { - return toolResult - } - } - } - t.Fatalf("missing provider-executed prompt tool result %q", id) - return fantasy.ToolResultPart{} -} - -func requireAnthropicProviderToolPromptSafe(t *testing.T, prompt []fantasy.Message) { - t.Helper() - - require.Empty(t, chatsanitize.ValidateAnthropicProviderToolHistory(prompt)) -} - -func requireLogField(t *testing.T, entry slog.SinkEntry, name string) any { - t.Helper() - - for _, field := range entry.Fields { - if field.Name == name { - return field.Value - } - } - t.Fatalf("missing log field %q", name) - return nil -} - -func containsPromptSentinel(prompt []fantasy.Message) bool { - for _, message := range prompt { - if message.Role != fantasy.MessageRoleUser || len(message.Content) != 1 { - continue - } - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](message.Content[0]) - if !ok { - continue - } - if strings.HasPrefix(textPart.Text, "__chatd_agent_prompt_sentinel_") { - return true - } - } - return false -} - -func TestRun_MultiStepToolExecution(t *testing.T) { - t.Parallel() - - var mu sync.Mutex - var streamCalls int - var secondCallPrompt []fantasy.Message - - model := &chattest.FakeModel{ - ProviderName: "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 - var persistedSteps []PersistedStep - 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, step PersistedStep) error { - persistStepCalls++ - persistedSteps = append(persistedSteps, step) - 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") - - // The first persisted step (tool-call step) must carry - // accurate timestamps for duration computation. - require.Len(t, persistedSteps, 2) - toolStep := persistedSteps[0] - require.Contains(t, toolStep.ToolCallCreatedAt, "tc-1", - "tool-call step must record when the model emitted the call") - require.Contains(t, toolStep.ToolResultCreatedAt, "tc-1", - "tool-call step must record when the tool result was produced") - require.False(t, toolStep.ToolResultCreatedAt["tc-1"].Before(toolStep.ToolCallCreatedAt["tc-1"]), - "tool-result timestamp must be >= tool-call timestamp") -} - -func TestStopAfterTool_Success(t *testing.T) { - t.Parallel() - - streamCalls := 0 - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-plan", ToolCallName: "propose_plan"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-plan", Delta: `{"path":"/tmp/plan.md"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-plan"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-plan", - ToolCallName: "propose_plan", - ToolCallInput: `{"path":"/tmp/plan.md"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - proposePlanTool := fantasy.NewAgentTool( - "propose_plan", - "writes a plan", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.NewTextResponse("plan saved"), nil - }, - ) - - var persistedSteps []PersistedStep - persistStepCalls := 0 - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "propose a plan"), - }, - Tools: []fantasy.AgentTool{proposePlanTool}, - MaxSteps: 5, - StopAfterTools: map[string]struct{}{ - "propose_plan": {}, - }, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistStepCalls++ - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.ErrorIs(t, err, ErrStopAfterTool) - require.Equal(t, 1, streamCalls) - require.Equal(t, 1, persistStepCalls) - require.Len(t, persistedSteps, 1) - - var foundToolResult bool - for _, block := range persistedSteps[0].Content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok || toolResult.ToolName != "propose_plan" { - continue - } - foundToolResult = true - _, isErr := toolResult.Result.(fantasy.ToolResultOutputContentError) - require.False(t, isErr, "stop-after-tool should only trigger on successful tool results") - } - require.True(t, foundToolResult, "persisted step should include the successful tool result before stopping") -} - -func TestStopAfterTool_IgnoresErrorResults(t *testing.T) { - t.Parallel() - - streamCalls := 0 - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - if streamCalls == 1 { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-plan", ToolCallName: "propose_plan"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-plan", Delta: `{"path":"/tmp/plan.md"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-plan"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-plan", - ToolCallName: "propose_plan", - ToolCallInput: `{"path":"/tmp/plan.md"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "tool failed, continue"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - proposePlanTool := fantasy.NewAgentTool( - "propose_plan", - "writes a plan", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.NewTextErrorResponse("plan failed"), nil - }, - ) - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "propose a plan"), - }, - Tools: []fantasy.AgentTool{proposePlanTool}, - MaxSteps: 5, - StopAfterTools: map[string]struct{}{ - "propose_plan": {}, - }, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - require.Len(t, persistedSteps, 2) - - var foundToolError bool - for _, block := range persistedSteps[0].Content { - toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok || toolResult.ToolName != "propose_plan" { - continue - } - _, foundToolError = toolResult.Result.(fantasy.ToolResultOutputContentError) - } - require.True(t, foundToolError, "first step should persist the failed tool result") -} - -func TestRun_ParallelToolExecutionTimestamps(t *testing.T) { - t.Parallel() - - var mu sync.Mutex - var streamCalls int - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - _ = call - - switch step { - case 0: - // Step 0: produce two tool calls in one stream. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "read_file"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{"path":"a.go"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "read_file", - ToolCallInput: `{"path":"a.go"}`, - }, - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-2", ToolCallName: "write_file"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-2", Delta: `{"path":"b.go"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-2"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-2", - ToolCallName: "write_file", - ToolCallInput: `{"path":"b.go"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - default: - // Step 1: return plain text. - 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 persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "do both"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool("read_file"), - newNoopTool("write_file"), - }, - MaxSteps: 5, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - - // Two steps: tool-call step + text step. - require.Equal(t, 2, streamCalls) - require.Len(t, persistedSteps, 2) - - toolStep := persistedSteps[0] - - // Both tool-call IDs must appear in ToolCallCreatedAt. - require.Contains(t, toolStep.ToolCallCreatedAt, "tc-1", - "tool-call step must record when tc-1 was emitted") - require.Contains(t, toolStep.ToolCallCreatedAt, "tc-2", - "tool-call step must record when tc-2 was emitted") - - // Both tool-call IDs must appear in ToolResultCreatedAt. - require.Contains(t, toolStep.ToolResultCreatedAt, "tc-1", - "tool-call step must record when tc-1 result was produced") - require.Contains(t, toolStep.ToolResultCreatedAt, "tc-2", - "tool-call step must record when tc-2 result was produced") - - // Result timestamps must be >= call timestamps for both. - require.False(t, toolStep.ToolResultCreatedAt["tc-1"].Before(toolStep.ToolCallCreatedAt["tc-1"]), - "tc-1 tool-result timestamp must be >= tool-call timestamp") - require.False(t, toolStep.ToolResultCreatedAt["tc-2"].Before(toolStep.ToolCallCreatedAt["tc-2"]), - "tc-2 tool-result timestamp must be >= tool-call timestamp") -} - -// TestRun_ExclusiveToolPolicyViolation exercises the full Run() -> -// executeToolsForStep() -> applyExclusiveToolPolicy() wiring. When an -// exclusive tool is called alongside other locally-executable tools, -// neither runner must fire and every call in the batch must receive a -// synthesized policy error that is both persisted and published via -// SSE. This guards against a regression where -// executeToolsForStep's policy call is accidentally removed: the -// pure-unit tests cover the policy function in isolation, but only -// this test catches a broken wiring path. -func TestRun_ExclusiveToolPolicyViolation(t *testing.T) { - t.Parallel() - - var advisorRuns atomic.Int32 - advisorTool := fantasy.NewAgentTool( - "advisor", - "returns strategic guidance", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - advisorRuns.Add(1) - return fantasy.NewTextResponse(`{"status":"ok"}`), nil - }, - ) - var readRuns atomic.Int32 - readTool := fantasy.NewAgentTool( - "read_file", - "reads a file", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - readRuns.Add(1) - return fantasy.NewTextResponse(`{"contents":"main"}`), nil - }, - ) - - var mu sync.Mutex - var streamCalls int - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - if step == 0 { - // Step 0: model emits an illegal mixed batch. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "advisor-1", ToolCallName: "advisor"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "advisor-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "advisor-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "advisor-1", - ToolCallName: "advisor", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeToolInputStart, ID: "read-1", ToolCallName: "read_file"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "read-1", Delta: `{"path":"main.go"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "read-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "read-1", - ToolCallName: "read_file", - ToolCallInput: `{"path":"main.go"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - // Step 1: the loop re-streams after tool results; end the run. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "ok, retrying"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var persistedSteps []PersistedStep - var publishedToolParts []codersdk.ChatMessagePart - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "please advise and read"), - }, - Tools: []fantasy.AgentTool{advisorTool, readTool}, - ExclusiveToolNames: map[string]bool{"advisor": true}, - MaxSteps: 5, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - PublishMessagePart: func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { - if role != codersdk.ChatMessageRoleTool { - return - } - publishedToolParts = append(publishedToolParts, part) - }, - }) - require.NoError(t, err) - - // Neither runner must have fired: the policy short-circuits - // before partitioning and execution. - require.Equal(t, int32(0), advisorRuns.Load(), - "advisor runner must not fire on mixed batches") - require.Equal(t, int32(0), readRuns.Load(), - "read_file runner must not fire on mixed batches") - - // Two steps: the mixed-batch step plus the follow-up stream. - require.Len(t, persistedSteps, 2) - firstStep := persistedSteps[0] - - advisorErr, ok := findToolResultByID(firstStep.Content, "advisor-1") - require.True(t, ok, "persisted step must contain the advisor policy result") - requireToolResultErrorMessage(t, advisorErr, - "advisor must be called alone, without other tools in the same batch. Retry with only the advisor call.") - - readErr, ok := findToolResultByID(firstStep.Content, "read-1") - require.True(t, ok, "persisted step must contain the read_file policy result") - requireToolResultErrorMessage(t, readErr, - "this tool was skipped because advisor must run alone in its batch. Retry your tool calls without advisor, or call advisor separately first.") - - // Policy-error results must be SSE-published so the client - // can render them immediately. Confirm both tool-result parts - // reached PublishMessagePart with a non-nil CreatedAt, which - // is the dbtime.Now() stamp the policy branch sets. - var sawAdvisorPart, sawReadPart bool - for _, part := range publishedToolParts { - switch part.ToolCallID { - case "advisor-1": - sawAdvisorPart = true - require.NotNil(t, part.CreatedAt, - "policy result SSE part must carry the dbtime.Now() timestamp") - case "read-1": - sawReadPart = true - require.NotNil(t, part.CreatedAt, - "policy result SSE part must carry the dbtime.Now() timestamp") - } - } - require.True(t, sawAdvisorPart, "advisor policy result must be SSE-published") - require.True(t, sawReadPart, "read_file policy result must be SSE-published") -} - -func findToolResultByID( - content []fantasy.Content, - toolCallID string, -) (fantasy.ToolResultContent, bool) { - for _, block := range content { - tr, ok := fantasy.AsContentType[fantasy.ToolResultContent](block) - if !ok { - continue - } - if tr.ToolCallID == toolCallID { - return tr, true - } - } - return fantasy.ToolResultContent{}, false -} - func TestExclusiveToolPolicy_MixedBatchErrors(t *testing.T) { t.Parallel() @@ -2244,1191 +539,6 @@ func TestExclusiveToolPolicy_MultipleExclusive(t *testing.T) { ) } -// TestRun_ExclusiveToolPolicyBlocksMixedWithDynamicTool guards the -// exclusive-over-dynamic bypass: the policy must run before the -// built-in vs dynamic partition. If a future refactor moves the -// policy check beneath the partition (so only built-in calls are -// inspected), an exclusive builtin mixed with a dynamic tool would -// still execute locally while the dynamic call is handed off via -// ErrDynamicToolCall, breaking the planning-only contract. -// -// This test has the model emit an exclusive builtin (advisor) -// alongside a dynamic tool (mcp_tool) in the same batch and asserts -// that Run does NOT exit with ErrDynamicToolCall, the advisor -// runner never fires, and both calls receive a synthesized policy -// error. -func TestRun_ExclusiveToolPolicyBlocksMixedWithDynamicTool(t *testing.T) { - t.Parallel() - - var advisorRuns atomic.Int32 - advisorTool := fantasy.NewAgentTool( - "advisor", - "returns strategic guidance", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - advisorRuns.Add(1) - return fantasy.NewTextResponse(`{"status":"ok"}`), nil - }, - ) - - var mu sync.Mutex - var streamCalls int - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - if step == 0 { - // Step 0: model emits an illegal mixed batch - // combining an exclusive builtin with a - // dynamic tool. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "advisor-1", ToolCallName: "advisor"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "advisor-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "advisor-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "advisor-1", - ToolCallName: "advisor", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeToolInputStart, ID: "mcp-1", ToolCallName: "mcp_tool"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "mcp-1", Delta: `{"q":"docs"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "mcp-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "mcp-1", - ToolCallName: "mcp_tool", - ToolCallInput: `{"q":"docs"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - // Step 1: after the policy error is fed back, - // terminate the run so the test assertions have a - // deterministic exit. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "retrying"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "please advise and fetch"), - }, - Tools: []fantasy.AgentTool{advisorTool}, - DynamicToolNames: map[string]bool{"mcp_tool": true}, - ExclusiveToolNames: map[string]bool{"advisor": true}, - MaxSteps: 5, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - // Run must NOT exit with ErrDynamicToolCall: the policy - // short-circuits before the dynamic partition so the dynamic - // call is never handed off for external execution. - require.NoError(t, err) - - // The advisor runner must not fire on mixed batches; the - // policy blocks the whole batch including the exclusive tool - // itself. - require.Equal(t, int32(0), advisorRuns.Load(), - "advisor runner must not fire on mixed batches") - - // Two steps: the mixed-batch step with synthesized policy - // errors plus the follow-up stream that ends the run. - require.Len(t, persistedSteps, 2) - firstStep := persistedSteps[0] - - // The persisted step must not record the dynamic tool as - // pending: the policy-error path returns before - // persistPendingDynamicStep runs. - require.Empty(t, firstStep.PendingDynamicToolCalls, - "policy-rejected batches must not leak dynamic tool calls to the caller") - - advisorErr, ok := findToolResultByID(firstStep.Content, "advisor-1") - require.True(t, ok, "persisted step must contain the advisor policy result") - requireToolResultErrorMessage(t, advisorErr, - "advisor must be called alone, without other tools in the same batch. Retry with only the advisor call.") - - mcpErr, ok := findToolResultByID(firstStep.Content, "mcp-1") - require.True(t, ok, "persisted step must contain the mcp_tool policy result") - requireToolResultErrorMessage(t, mcpErr, - "this tool was skipped because advisor must run alone in its batch. Retry your tool calls without advisor, or call advisor separately first.") -} - -// TestRun_ExclusiveToolAloneSucceeds is the happy-path counterpart -// to TestRun_ExclusiveToolPolicyViolation: a single exclusive tool -// emitted alone must actually execute. The `len(toolCalls) <= 1` -// guard in firstExclusiveToolName is the sole mechanism that lets -// solo exclusive-tool calls proceed. If that guard regresses to -// `< 1`, every solo exclusive-tool call would enter an infinite -// policy-error/retry loop, and every unit test on the policy -// function in isolation would still pass. Only this Run()-level -// test catches that regression. -func TestRun_ExclusiveToolAloneSucceeds(t *testing.T) { - t.Parallel() - - var advisorRuns atomic.Int32 - advisorTool := fantasy.NewAgentTool( - "advisor", - "returns strategic guidance", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - advisorRuns.Add(1) - return fantasy.NewTextResponse(`{"status":"ok"}`), nil - }, - ) - - var mu sync.Mutex - var streamCalls int - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - if step == 0 { - // Step 0: model emits exactly one - // exclusive-tool call in isolation. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "advisor-1", ToolCallName: "advisor"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "advisor-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "advisor-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "advisor-1", - ToolCallName: "advisor", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - // Step 1: the loop re-streams after the tool - // result; end the run. - 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}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "please advise"), - }, - Tools: []fantasy.AgentTool{advisorTool}, - ExclusiveToolNames: map[string]bool{"advisor": true}, - MaxSteps: 5, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - - // The solo exclusive tool must actually execute exactly once. - require.Equal(t, int32(1), advisorRuns.Load(), - "solo exclusive-tool call must execute") - - // The first persisted step must contain a non-error tool - // result for the advisor call, proving the policy did not - // synthesize an error and the real runner fired. - require.GreaterOrEqual(t, len(persistedSteps), 1) - result, ok := findToolResultByID(persistedSteps[0].Content, "advisor-1") - require.True(t, ok, "persisted step must contain the advisor tool result") - _, isErr := result.Result.(fantasy.ToolResultOutputContentError) - require.Falsef(t, isErr, - "solo exclusive-tool call must produce a real tool result, not a policy error: %+v", result.Result) -} - -// TestRun_ExclusiveToolWithProviderExecutedSucceeds guards the -// interaction between the ProviderExecuted filter and the -// exclusive-tool policy. executeToolsForStep builds localCandidates -// by dropping ProviderExecuted calls before passing them to -// applyExclusiveToolPolicy. That filter is the sole mechanism -// preventing a false policy violation when a solo exclusive tool -// appears in a batch where the provider also server-executed a tool -// (for example Anthropic web_search). -// -// If the filter is removed, localCandidates would contain both the -// provider-executed call and the exclusive call. firstExclusiveToolName -// would then see len > 1, find advisor, and return a violation. The -// advisor would never run and the retry loop would burn steps until -// MaxSteps. -// -// This test emits an advisor call alongside a provider-executed -// web_search call (with its provider-emitted result) and asserts the -// advisor runner actually fires. -func TestRun_ExclusiveToolWithProviderExecutedSucceeds(t *testing.T) { - t.Parallel() - - var advisorRuns atomic.Int32 - advisorTool := fantasy.NewAgentTool( - "advisor", - "returns strategic guidance", - func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) { - advisorRuns.Add(1) - return fantasy.NewTextResponse(`{"status":"ok"}`), nil - }, - ) - - var mu sync.Mutex - var streamCalls int - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - if step == 0 { - // Step 0: provider server-executed web_search and - // returned its result inline, plus the model - // emitted an exclusive advisor call for local - // execution. The ProviderExecuted filter must - // drop web_search from the policy check so the - // advisor is treated as a solo exclusive call. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "ws-1", ToolCallName: "web_search", ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "ws-1", Delta: `{"query":"coder"}`, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "ws-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "ws-1", - ToolCallName: "web_search", - ToolCallInput: `{"query":"coder"}`, - ProviderExecuted: true, - }, - { - Type: fantasy.StreamPartTypeToolResult, - ID: "ws-1", - ToolCallName: "web_search", - ProviderExecuted: true, - }, - {Type: fantasy.StreamPartTypeToolInputStart, ID: "advisor-1", ToolCallName: "advisor"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "advisor-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "advisor-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "advisor-1", - ToolCallName: "advisor", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - // Step 1: end the run after the advisor result is - // fed back. - 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}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search and then advise"), - }, - Tools: []fantasy.AgentTool{advisorTool}, - ExclusiveToolNames: map[string]bool{"advisor": true}, - MaxSteps: 5, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - - // The advisor must execute exactly once: the ProviderExecuted - // filter removes web_search from the exclusivity check, so the - // advisor is treated as a solo exclusive call. - require.Equal(t, int32(1), advisorRuns.Load(), - "advisor must execute when the only other call in the batch was provider-executed") - - // The advisor result must be a real tool result, not a - // synthesized policy error. - require.GreaterOrEqual(t, len(persistedSteps), 1) - advisorResult, ok := findToolResultByID(persistedSteps[0].Content, "advisor-1") - require.True(t, ok, "persisted step must contain the advisor tool result") - _, isErr := advisorResult.Result.(fantasy.ToolResultOutputContentError) - require.Falsef(t, isErr, - "advisor must produce a real tool result, not a policy error: %+v", advisorResult.Result) -} - -func TestRun_PersistStepErrorPropagates(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: "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") -} - -// TestRun_ShutdownDuringToolExecutionReturnsContextCanceled verifies that -// when the parent context is canceled (simulating server shutdown) while -// a tool is blocked, Run returns context.Canceled, not ErrInterrupted. -// This matters because the caller uses the error type to decide whether -// to set chat status to "pending" (retryable on another worker) vs -// "waiting" (stuck forever). -func TestRun_ShutdownDuringToolExecutionReturnsContextCanceled(t *testing.T) { - t.Parallel() - - toolStarted := make(chan struct{}) - - // Model returns a single tool call, then finishes. - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-block", ToolCallName: "blocking_tool"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-block", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-block"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-block", - ToolCallName: "blocking_tool", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - // Tool that blocks until its context is canceled, simulating - // a long-running operation like wait_agent. - blockingTool := fantasy.NewAgentTool( - "blocking_tool", - "blocks until context canceled", - func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - close(toolStarted) - <-ctx.Done() - return fantasy.ToolResponse{}, ctx.Err() - }, - ) - - // Simulate the server context (parent) and chat context - // (child). Canceling the parent simulates graceful shutdown. - serverCtx, serverCancel := context.WithCancel(context.Background()) - defer serverCancel() - - serverCancelDone := make(chan struct{}) - go func() { - defer close(serverCancelDone) - <-toolStarted - t.Logf("tool started, canceling server context to simulate shutdown") - serverCancel() - }() - - // persistStep mirrors the FIXED chatd.go code: it only returns - // ErrInterrupted when the context was actually canceled due to - // an interruption (cause is ErrInterrupted). For shutdown - // (plain context.Canceled), it returns the original error so - // callers can distinguish the two. - persistStep := func(persistCtx context.Context, _ PersistedStep) error { - if persistCtx.Err() != nil { - if errors.Is(context.Cause(persistCtx), ErrInterrupted) { - return ErrInterrupted - } - return persistCtx.Err() - } - return nil - } - - err := Run(serverCtx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "run the blocking tool"), - }, - Tools: []fantasy.AgentTool{blockingTool}, - MaxSteps: 3, - PersistStep: persistStep, - }) - // Wait for the cancel goroutine to finish to aid flake - // diagnosis if the test ever hangs. - <-serverCancelDone - - require.Error(t, err) - // The error must NOT be ErrInterrupted, it should propagate - // as context.Canceled so the caller can distinguish shutdown - // from user interruption. Use assert (not require) so both - // checks are evaluated even if the first fails. - assert.NotErrorIs(t, err, ErrInterrupted, "shutdown cancellation must not be converted to ErrInterrupted") - assert.ErrorIs(t, err, context.Canceled, "shutdown should propagate as context.Canceled") -} - -func TestToResponseMessages_ProviderExecutedToolResultInAssistantMessage(t *testing.T) { - t.Parallel() - - sr := stepResult{ - content: []fantasy.Content{ - // Provider-executed tool call (e.g. web_search). - fantasy.ToolCallContent{ - ToolCallID: "provider-tc-1", - ToolName: "web_search", - Input: `{"query":"coder"}`, - ProviderExecuted: true, - }, - // Provider-executed tool result, must stay in - // assistant message. - fantasy.ToolResultContent{ - ToolCallID: "provider-tc-1", - ToolName: "web_search", - ProviderExecuted: true, - ProviderMetadata: fantasy.ProviderMetadata{"anthropic": nil}, - }, - // Local tool call (e.g. read_file). - fantasy.ToolCallContent{ - ToolCallID: "local-tc-1", - ToolName: "read_file", - Input: `{"path":"main.go"}`, - ProviderExecuted: false, - }, - // Local tool result, should go into tool message. - fantasy.ToolResultContent{ - ToolCallID: "local-tc-1", - ToolName: "read_file", - Result: fantasy.ToolResultOutputContentText{Text: "some result"}, - ProviderExecuted: false, - }, - }, - } - - msgs := sr.toResponseMessages() - require.Len(t, msgs, 2, "expected assistant + tool messages") - - // First message: assistant role. - assistantMsg := msgs[0] - assert.Equal(t, fantasy.MessageRoleAssistant, assistantMsg.Role) - require.Len(t, assistantMsg.Content, 3, - "assistant message should have provider ToolCallPart, provider ToolResultPart, and local ToolCallPart") - - // Part 0: provider tool call. - providerTC, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](assistantMsg.Content[0]) - require.True(t, ok, "part 0 should be ToolCallPart") - assert.Equal(t, "provider-tc-1", providerTC.ToolCallID) - assert.True(t, providerTC.ProviderExecuted) - - // Part 1: provider tool result (inline in assistant turn). - providerTR, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](assistantMsg.Content[1]) - require.True(t, ok, "part 1 should be ToolResultPart") - assert.Equal(t, "provider-tc-1", providerTR.ToolCallID) - assert.True(t, providerTR.ProviderExecuted) - - // Part 2: local tool call. - localTC, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](assistantMsg.Content[2]) - require.True(t, ok, "part 2 should be ToolCallPart") - assert.Equal(t, "local-tc-1", localTC.ToolCallID) - assert.False(t, localTC.ProviderExecuted) - - // Second message: tool role. - toolMsg := msgs[1] - assert.Equal(t, fantasy.MessageRoleTool, toolMsg.Role) - require.Len(t, toolMsg.Content, 1, - "tool message should have only the local ToolResultPart") - - localTR, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](toolMsg.Content[0]) - require.True(t, ok, "tool part should be ToolResultPart") - assert.Equal(t, "local-tc-1", localTR.ToolCallID) - assert.False(t, localTR.ProviderExecuted) -} - -func TestToResponseMessages_FiltersEmptyTextAndReasoningParts(t *testing.T) { - t.Parallel() - - sr := stepResult{ - content: []fantasy.Content{ - // Empty text, should be filtered. - fantasy.TextContent{Text: ""}, - // Whitespace-only text, should be filtered. - fantasy.TextContent{Text: " \t\n"}, - // Empty reasoning, should be filtered. - fantasy.ReasoningContent{Text: ""}, - // Whitespace-only reasoning, should be filtered. - fantasy.ReasoningContent{Text: " \n"}, - // Non-empty text, should pass through. - fantasy.TextContent{Text: "hello world"}, - // Leading/trailing whitespace with content, kept - // with the original value (not trimmed). - fantasy.TextContent{Text: " hello "}, - // Non-empty reasoning, should pass through. - fantasy.ReasoningContent{Text: "let me think"}, - // Tool call, should be unaffected by filtering. - fantasy.ToolCallContent{ - ToolCallID: "tc-1", - ToolName: "read_file", - Input: `{"path":"main.go"}`, - }, - // Local tool result, should be unaffected by filtering. - fantasy.ToolResultContent{ - ToolCallID: "tc-1", - ToolName: "read_file", - Result: fantasy.ToolResultOutputContentText{Text: "file contents"}, - }, - }, - } - - msgs := sr.toResponseMessages() - require.Len(t, msgs, 2, "expected assistant + tool messages") - - // First message: assistant role with non-empty text, reasoning, - // and the tool call. The four empty/whitespace-only parts must - // have been dropped. - assistantMsg := msgs[0] - assert.Equal(t, fantasy.MessageRoleAssistant, assistantMsg.Role) - require.Len(t, assistantMsg.Content, 4, - "assistant message should have 2x TextPart, ReasoningPart, and ToolCallPart") - - // Part 0: non-empty text. - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](assistantMsg.Content[0]) - require.True(t, ok, "part 0 should be TextPart") - assert.Equal(t, "hello world", textPart.Text) - - // Part 1: padded text, original whitespace preserved. - paddedPart, ok := fantasy.AsMessagePart[fantasy.TextPart](assistantMsg.Content[1]) - require.True(t, ok, "part 1 should be TextPart") - assert.Equal(t, " hello ", paddedPart.Text) - - // Part 2: non-empty reasoning. - reasoningPart, ok := fantasy.AsMessagePart[fantasy.ReasoningPart](assistantMsg.Content[2]) - require.True(t, ok, "part 2 should be ReasoningPart") - assert.Equal(t, "let me think", reasoningPart.Text) - - // Part 3: tool call (unaffected by text/reasoning filtering). - toolCallPart, ok := fantasy.AsMessagePart[fantasy.ToolCallPart](assistantMsg.Content[3]) - require.True(t, ok, "part 3 should be ToolCallPart") - assert.Equal(t, "tc-1", toolCallPart.ToolCallID) - assert.Equal(t, "read_file", toolCallPart.ToolName) - - // Second message: tool role with the local tool result. - toolMsg := msgs[1] - assert.Equal(t, fantasy.MessageRoleTool, toolMsg.Role) - require.Len(t, toolMsg.Content, 1, - "tool message should have only the local ToolResultPart") - - toolResultPart, ok := fantasy.AsMessagePart[fantasy.ToolResultPart](toolMsg.Content[0]) - require.True(t, ok, "tool part should be ToolResultPart") - assert.Equal(t, "tc-1", toolResultPart.ToolCallID) -} - -func hasAnthropicEphemeralCacheControl(message fantasy.Message) bool { - if len(message.ProviderOptions) == 0 { - return false - } - - options, ok := message.ProviderOptions[fantasyanthropic.Name] - if !ok { - return false - } - - cacheOptions, ok := options.(*fantasyanthropic.ProviderCacheControlOptions) - return ok && cacheOptions.CacheControl.Type == "ephemeral" -} - -// TestRun_InterruptedDuringToolExecutionPersistsStep verifies that when -// tools are executing and the chat is interrupted, the accumulated step -// content (assistant blocks + tool results) is persisted via the -// interrupt-safe path rather than being lost. -func TestRun_InterruptedDuringToolExecutionPersistsStep(t *testing.T) { - t.Parallel() - - toolStarted := make(chan struct{}) - - // Model returns a completed tool call in the stream. - model := &chattest.FakeModel{ - ProviderName: "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: "calling tool"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeReasoningStart, ID: "reason-1"}, - {Type: fantasy.StreamPartTypeReasoningDelta, ID: "reason-1", Delta: "let me think"}, - {Type: fantasy.StreamPartTypeReasoningEnd, ID: "reason-1"}, - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "slow_tool"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{"key":"value"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "slow_tool", - ToolCallInput: `{"key":"value"}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - }, - } - - // Tool that blocks until context is canceled, simulating - // a long-running operation interrupted by the user. - slowTool := fantasy.NewAgentTool( - "slow_tool", - "blocks until canceled", - func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - close(toolStarted) - <-ctx.Done() - return fantasy.ToolResponse{}, ctx.Err() - }, - ) - - ctx, cancel := context.WithCancelCause(context.Background()) - defer cancel(nil) - - go func() { - <-toolStarted - cancel(ErrInterrupted) - }() - - var persistedContent []fantasy.Content - persistedCtxErr := xerrors.New("unset") - - err := Run(ctx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "run the slow tool"), - }, - Tools: []fantasy.AgentTool{slowTool}, - MaxSteps: 3, - PersistStep: func(persistCtx context.Context, step PersistedStep) error { - persistedCtxErr = persistCtx.Err() - persistedContent = append([]fantasy.Content(nil), step.Content...) - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - // persistInterruptedStep uses context.WithoutCancel, so the - // persist callback should see a non-canceled context. - require.NoError(t, persistedCtxErr) - require.NotEmpty(t, persistedContent) - - var ( - foundText bool - foundReasoning bool - foundToolCall bool - foundToolResult bool - ) - for _, block := range persistedContent { - if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { - if strings.Contains(text.Text, "calling tool") { - foundText = true - } - continue - } - if reasoning, ok := fantasy.AsContentType[fantasy.ReasoningContent](block); ok { - if strings.Contains(reasoning.Text, "let me think") { - foundReasoning = true - } - continue - } - if toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block); ok { - if toolCall.ToolCallID == "tc-1" && toolCall.ToolName == "slow_tool" { - foundToolCall = true - } - continue - } - if toolResult, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok { - if toolResult.ToolCallID == "tc-1" { - foundToolResult = true - } - } - } - require.True(t, foundText, "persisted content should include text from the stream") - require.True(t, foundReasoning, "persisted content should include reasoning from the stream") - require.True(t, foundToolCall, "persisted content should include the tool call") - require.True(t, foundToolResult, "persisted content should include the tool result (error from cancellation)") -} - -// TestRun_ProviderExecutedToolResultTimestamps verifies that -// provider-executed tool results (e.g. web search) have their -// timestamps recorded in PersistedStep.ToolResultCreatedAt so -// the persistence layer can stamp CreatedAt on the parts. -func TestRun_ProviderExecutedToolResultTimestamps(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - // Simulate a provider-executed tool call and result - // (e.g. Anthropic web search) followed by a text - // response, all in a single stream. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "ws-1", ToolCallName: "web_search", ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "ws-1", Delta: `{"query":"coder"}`, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "ws-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "ws-1", - ToolCallName: "web_search", - ToolCallInput: `{"query":"coder"}`, - ProviderExecuted: true, - }, - // Provider-executed tool result, emitted by - // the provider, not our tool runner. - { - Type: fantasy.StreamPartTypeToolResult, - ID: "ws-1", - ToolCallName: "web_search", - ProviderExecuted: true, - }, - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "search done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search for coder"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - require.Len(t, persistedSteps, 1) - - step := persistedSteps[0] - - // Provider-executed tool call should have a call timestamp. - require.Contains(t, step.ToolCallCreatedAt, "ws-1", - "provider-executed tool call must record its timestamp") - - // Provider-executed tool result should have a result - // timestamp so the frontend can compute duration. - require.Contains(t, step.ToolResultCreatedAt, "ws-1", - "provider-executed tool result must record its timestamp") - - require.False(t, - step.ToolResultCreatedAt["ws-1"].Before(step.ToolCallCreatedAt["ws-1"]), - "tool-result timestamp must be >= tool-call timestamp") -} - -func TestRun_AnthropicDropsUnpairedProviderToolBeforePersist(t *testing.T) { - t.Parallel() - - testCases := []struct { - name string - toolName string - toolInput string - }{ - { - name: "web_search", - toolName: "web_search", - toolInput: `{"query":"coder"}`, - }, - { - name: "code_execution", - toolName: "code_execution", - toolInput: `{"code":"print(1)"}`, - }, - } - - for _, tc := range testCases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "pt-1", ToolCallName: tc.toolName, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "pt-1", Delta: tc.toolInput, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "pt-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "pt-1", - ToolCallName: tc.toolName, - ToolCallInput: tc.toolInput, - ProviderExecuted: true, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - persistCalls := 0 - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "run provider tool"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - persistCalls++ - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, 0, persistCalls) - }) - } -} - -func TestRun_AnthropicKeepsPairedWebSearchBeforePersist(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "ws-1", ToolCallName: "web_search", ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "ws-1", Delta: `{"query":"coder"}`, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "ws-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "ws-1", - ToolCallName: "web_search", - ToolCallInput: `{"query":"coder"}`, - ProviderExecuted: true, - }, - { - Type: fantasy.StreamPartTypeToolResult, - ID: "ws-1", - ToolCallName: "web_search", - ProviderExecuted: true, - ProviderMetadata: validWebSearchProviderMetadataForTest(), - }, - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "search done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search for coder"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - require.Len(t, persistedSteps, 1) - - toolCall := requireToolCallContent(t, persistedSteps[0].Content, "ws-1", "web_search") - require.True(t, toolCall.ProviderExecuted) - toolResult := requireToolResultContent(t, persistedSteps[0].Content, "ws-1", "web_search") - require.True(t, toolResult.ProviderExecuted) - requireTextContent(t, persistedSteps[0].Content, "search done") -} - -func TestRun_AnthropicInterruptedWebSearchDoesNotPersistSyntheticResult(t *testing.T) { - t.Parallel() - - started := make(chan struct{}) - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputStart, - ID: "ws-1", - ToolCallName: "web_search", - ProviderExecuted: true, - }) { - return - } - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputDelta, - ID: "ws-1", - Delta: `{"query":"coder"}`, - ProviderExecuted: true, - }) { - return - } - close(started) - <-ctx.Done() - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - }, - } - - ctx, cancel := context.WithCancelCause(context.Background()) - defer cancel(nil) - go func() { - <-started - cancel(ErrInterrupted) - }() - - persistCalls := 0 - err := Run(ctx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search for coder"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - persistCalls++ - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - require.Equal(t, 0, persistCalls) -} - -func TestRun_AnthropicInterruptedProviderToolKeepsLocalSyntheticResult(t *testing.T) { - t.Parallel() - - started := make(chan struct{}) - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputStart, - ID: "ws-1", - ToolCallName: "web_search", - ProviderExecuted: true, - }) { - return - } - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputDelta, - ID: "ws-1", - Delta: `{"query":"coder"}`, - ProviderExecuted: true, - }) { - return - } - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputStart, - ID: "tc-1", - ToolCallName: "read_file", - }) { - return - } - if !yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeToolInputDelta, - ID: "tc-1", - Delta: `{"path":"main.go"}`, - }) { - return - } - close(started) - <-ctx.Done() - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - }, - } - - ctx, cancel := context.WithCancelCause(context.Background()) - defer cancel(nil) - go func() { - <-started - cancel(ErrInterrupted) - }() - - var persistedSteps []PersistedStep - err := Run(ctx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search and read"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - require.Len(t, persistedSteps, 1) - requireNoProviderExecutedToolCallContent(t, persistedSteps[0].Content) - requireNoProviderExecutedToolResultContent(t, persistedSteps[0].Content) - - toolCall := requireToolCallContent(t, persistedSteps[0].Content, "tc-1", "read_file") - require.False(t, toolCall.ProviderExecuted) - toolResult := requireToolResultContent(t, persistedSteps[0].Content, "tc-1", "read_file") - require.False(t, toolResult.ProviderExecuted) - _, isErr := toolResult.Result.(fantasy.ToolResultOutputContentError) - require.True(t, isErr) -} - -func TestRun_AnthropicSanitizesProviderToolBeforeRequest(t *testing.T) { - t.Parallel() - - var capturedPrompt []fantasy.Message - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - capturedPrompt = append([]fantasy.Message(nil), call.Prompt...) - 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}, - }), nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search for coder"), - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{ - fantasy.ToolCallPart{ - ToolCallID: "ws-1", - ToolName: "web_search", - Input: `{"query":"coder"}`, - ProviderExecuted: true, - }, - }, - }, - textMessage(fantasy.MessageRoleUser, "continue"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - }) - require.NoError(t, err) - require.Len(t, capturedPrompt, 1) - require.Equal(t, fantasy.MessageRoleUser, capturedPrompt[0].Role) - require.Len(t, capturedPrompt[0].Content, 2) - requireNoProviderExecutedToolCallPrompt(t, capturedPrompt) -} - -func TestRun_AnthropicSanitizesWebSearchBeforeContinuation(t *testing.T) { - t.Parallel() - - var mu sync.Mutex - var streamCalls int - var secondCallPrompt []fantasy.Message - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - switch step { - case 0: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "ws-1", ToolCallName: "web_search", ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "ws-1", Delta: `{"query":"coder"}`, ProviderExecuted: true}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "ws-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "ws-1", - ToolCallName: "web_search", - ToolCallInput: `{"query":"coder"}`, - ProviderExecuted: true, - }, - {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: - 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: "done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - } - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search and read"), - }, - Tools: []fantasy.AgentTool{ - newNoopTool("read_file"), - }, - MaxSteps: 2, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - require.Len(t, persistedSteps, 2) - requireNoProviderExecutedToolCallContent(t, persistedSteps[0].Content) - requireNoProviderExecutedToolCallPrompt(t, secondCallPrompt) - - toolCall := requireToolCallContent(t, persistedSteps[0].Content, "tc-1", "read_file") - require.False(t, toolCall.ProviderExecuted) - toolResult := requireToolResultContent(t, persistedSteps[0].Content, "tc-1", "read_file") - require.False(t, toolResult.ProviderExecuted) - promptResult := requireToolResultPrompt(t, secondCallPrompt, "tc-1") - require.False(t, promptResult.ProviderExecuted) -} - func TestSanitizeAnthropicProviderToolContent(t *testing.T) { t.Parallel() @@ -3735,764 +845,6 @@ func TestSanitizeAnthropicProviderToolContent(t *testing.T) { } } -func TestRun_AnthropicProviderToolPreRequestGuard(t *testing.T) { - t.Parallel() - - webSearchTool := ProviderTool{ - Definition: fantasy.ProviderDefinedTool{ - ID: "anthropic.web_search", - Name: "web_search", - }, - } - providerPair := func(id string) []fantasy.MessagePart { - return []fantasy.MessagePart{ - fantasy.ToolCallPart{ - ToolCallID: id, - ToolName: "web_search", - Input: `{"query":"coder"}`, - ProviderExecuted: true, - }, - fantasy.ToolResultPart{ - ToolCallID: id, - Output: fantasy.ToolResultOutputContentText{Text: "ok"}, - ProviderExecuted: true, - ProviderOptions: fantasy.ProviderOptions(validWebSearchProviderMetadataForTest()), - }, - } - } - completionModel := func(capturedPrompt *[]fantasy.Message) *chattest.FakeModel { - return &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - ModelName: "claude-test", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - *capturedPrompt = append([]fantasy.Message(nil), call.Prompt...) - 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}, - }), nil - }, - } - } - - t.Run("allowed web search survives when provider tool is enabled", func(t *testing.T) { - t.Parallel() - - var capturedPrompt []fantasy.Message - err := Run(context.Background(), RunOptions{ - Model: completionModel(&capturedPrompt), - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search"), - { - Role: fantasy.MessageRoleAssistant, - Content: providerPair("ws-allowed"), - }, - textMessage(fantasy.MessageRoleUser, "continue"), - }, - ProviderTools: []ProviderTool{webSearchTool}, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - }) - require.NoError(t, err) - - toolCall := requireProviderExecutedToolCallPrompt(t, capturedPrompt, "ws-allowed") - require.Equal(t, "web_search", toolCall.ToolName) - requireProviderExecutedToolResultPrompt(t, capturedPrompt, "ws-allowed") - requireAnthropicProviderToolPromptSafe(t, capturedPrompt) - }) - - t.Run("web search history survives when provider tool is disabled", func(t *testing.T) { - t.Parallel() - - var capturedPrompt []fantasy.Message - err := Run(context.Background(), RunOptions{ - Model: completionModel(&capturedPrompt), - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search and read"), - { - Role: fantasy.MessageRoleAssistant, - Content: append(providerPair("ws-disabled"), fantasy.ToolCallPart{ - ToolCallID: "tc-1", - ToolName: "read_file", - Input: `{"path":"main.go"}`, - }), - }, - { - Role: fantasy.MessageRoleTool, - Content: []fantasy.MessagePart{ - fantasy.ToolResultPart{ - ToolCallID: "tc-1", - Output: fantasy.ToolResultOutputContentText{Text: "file"}, - }, - }, - }, - textMessage(fantasy.MessageRoleUser, "continue"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - }) - require.NoError(t, err) - - requireProviderExecutedToolCallPrompt(t, capturedPrompt, "ws-disabled") - requireProviderExecutedToolResultPrompt(t, capturedPrompt, "ws-disabled") - promptResult := requireToolResultPrompt(t, capturedPrompt, "tc-1") - require.False(t, promptResult.ProviderExecuted) - requireAnthropicProviderToolPromptSafe(t, capturedPrompt) - }) - - t.Run("direct guard textifies orphaned provider result", func(t *testing.T) { - t.Parallel() - - guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( - context.Background(), - slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - fantasyanthropic.Name, - "claude-test", - []fantasy.Message{ - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: "keep"}, - fantasy.ToolResultPart{ - ToolCallID: "ws-orphan", - Output: fantasy.ToolResultOutputContentText{Text: "search result"}, - ProviderExecuted: true, - }, - }, - }, - }, - ) - require.NoError(t, err) - - requireNoProviderExecutedToolResultPrompt(t, guarded) - requireAnthropicProviderToolPromptSafe(t, guarded) - require.Len(t, guarded, 1) - require.Len(t, guarded[0].Content, 2) - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](guarded[0].Content[0]) - require.True(t, ok) - require.Equal(t, "keep", textPart.Text) - textPart, ok = fantasy.AsMessagePart[fantasy.TextPart](guarded[0].Content[1]) - require.True(t, ok) - require.Equal(t, "search result", textPart.Text) - }) - - t.Run("direct guard leaves valid provider history unchanged", func(t *testing.T) { - t.Parallel() - - content := []fantasy.MessagePart{fantasy.TextPart{Text: "keep"}} - content = append(content, providerPair("ws-one")...) - content = append(content, providerPair("ws-two")...) - guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( - context.Background(), - slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - fantasyanthropic.Name, - "claude-test", - []fantasy.Message{{Role: fantasy.MessageRoleAssistant, Content: content}}, - ) - require.NoError(t, err) - - requireAnthropicProviderToolPromptSafe(t, guarded) - require.Len(t, guarded, 1) - require.Len(t, guarded[0].Content, len(content)) - requireProviderExecutedToolCallPrompt(t, guarded, "ws-one") - requireProviderExecutedToolResultPrompt(t, guarded, "ws-one") - requireProviderExecutedToolCallPrompt(t, guarded, "ws-two") - requireProviderExecutedToolResultPrompt(t, guarded, "ws-two") - }) - - t.Run("direct guard leaves non Anthropic providers unchanged", func(t *testing.T) { - t.Parallel() - - prompt := []fantasy.Message{ - { - Role: fantasy.MessageRoleAssistant, - Content: providerPair("ws-other-provider"), - }, - } - guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( - context.Background(), - slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - "fake", - "fake-model", - prompt, - ) - require.NoError(t, err) - require.Equal(t, prompt, guarded) - }) - - t.Run("guard logs removals", func(t *testing.T) { - t.Parallel() - - logSink := testutil.NewFakeSink(t) - logger := logSink.Logger() - logPair := providerPair("ws-log") - guarded, err := chatsanitize.ApplyAnthropicProviderToolGuard( - context.Background(), - logger, - fantasyanthropic.Name, - "claude-test", - []fantasy.Message{ - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{ - logPair[1], - logPair[0], - }, - }, - }, - ) - require.NoError(t, err) - - requireNoProviderExecutedToolCallPrompt(t, guarded) - requireNoProviderExecutedToolResultPrompt(t, guarded) - requireTextPrompt(t, guarded, "ok") - entries := logSink.Entries(func(e slog.SinkEntry) bool { - return e.Level == slog.LevelWarn && - e.Message == "removed provider-executed tool history" - }) - require.Len(t, entries, 1) - require.Equal(t, "pre_request_guard", requireLogField(t, entries[0], "phase")) - require.Equal(t, 1, requireLogField(t, entries[0], "removed_tool_calls")) - require.Equal(t, 1, requireLogField(t, entries[0], "removed_tool_results")) - }) - t.Run("run drops orphan provider call before provider request", func(t *testing.T) { - t.Parallel() - - streamCalls := 0 - var capturedPrompt fantasy.Prompt - model := &chattest.FakeModel{ - ProviderName: fantasyanthropic.Name, - ModelName: "claude-test", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - streamCalls++ - capturedPrompt = call.Prompt - return finishingStream(), nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "search"), - { - Role: fantasy.MessageRoleAssistant, - Content: []fantasy.MessagePart{ - fantasy.ReasoningPart{ - ProviderOptions: fantasy.ProviderOptions{ - fantasyanthropic.Name: &fantasyanthropic.ReasoningOptionMetadata{ - RedactedData: "redacted-payload", - }, - }, - }, - fantasy.ToolCallPart{ - ToolCallID: "ws-orphan", - ToolName: "web_search", - Input: `{"query":"coder"}`, - ProviderExecuted: true, - }, - fantasy.TextPart{Text: "partial"}, - }, - }, - textMessage(fantasy.MessageRoleUser, "continue"), - }, - Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, 1, streamCalls) - requireNoProviderExecutedToolCallPrompt(t, capturedPrompt) - requireAnthropicProviderToolPromptSafe(t, capturedPrompt) - requireTextPrompt(t, capturedPrompt, "partial") - reasoningPart := requireReasoningPrompt(t, capturedPrompt) - reasoningMetadata := fantasyanthropic.GetReasoningMetadata(reasoningPart.ProviderOptions) - require.NotNil(t, reasoningMetadata) - require.Equal(t, "redacted-payload", reasoningMetadata.RedactedData) - }) -} - -// TestRun_PersistStepInterruptedFallback verifies that when the normal -// PersistStep call returns ErrInterrupted (e.g., context canceled in a -// race), the step is retried via the interrupt-safe path. -func TestRun_PersistStepInterruptedFallback(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: "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 world"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var ( - mu sync.Mutex - persistCalls int - savedContent []fantasy.Content - ) - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - mu.Lock() - defer mu.Unlock() - persistCalls++ - if persistCalls == 1 { - // First call: simulate an interrupt race by - // returning ErrInterrupted without persisting. - return ErrInterrupted - } - // Second call (from persistInterruptedStep fallback): - // accept the content. - savedContent = append([]fantasy.Content(nil), step.Content...) - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - - mu.Lock() - defer mu.Unlock() - require.Equal(t, 2, persistCalls, "PersistStep should be called twice: once normally (failing), once via fallback") - require.NotEmpty(t, savedContent) - - var foundText bool - for _, block := range savedContent { - if text, ok := fantasy.AsContentType[fantasy.TextContent](block); ok { - if strings.Contains(text.Text, "hello world") { - foundText = true - } - } - } - require.True(t, foundText, "fallback should persist the text content") -} - -func TestRun_PrepareMessagesInjectsSystemContextMidLoop(t *testing.T) { - t.Parallel() - - const injectedInstruction = "You are working in /home/coder/project. Follow AGENTS.md guidelines." - - var mu sync.Mutex - var streamCalls int - var secondCallPrompt []fantasy.Message - - // Step 0 calls a tool. Step 1 sees the injected system message. - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - switch step { - case 0: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "create_workspace"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "create_workspace", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - default: - 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: "done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - } - }, - } - - // Simulate: after the tool executes (step 0), instruction - // becomes available. PrepareMessages injects it before step 1. - instructionInjected := make(chan struct{}) - var instructionAvailable atomic.Value - // The tool sets instruction after execution. - tool := fantasy.NewAgentTool( - "create_workspace", - "create a workspace", - func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - instructionAvailable.Store(injectedInstruction) - return fantasy.ToolResponse{}, nil - }, - ) - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "create a workspace and open a PR"), - }, - Tools: []fantasy.AgentTool{tool}, - MaxSteps: 5, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - select { - case <-instructionInjected: - return nil - default: - } - instr, ok := instructionAvailable.Load().(string) - if !ok || instr == "" { - return nil - } - close(instructionInjected) - // Insert a system message after existing system messages. - result := make([]fantasy.Message, 0, len(msgs)+1) - inserted := false - for i, msg := range msgs { - result = append(result, msg) - if !inserted && msg.Role == fantasy.MessageRoleSystem { - // Insert after the last system message. - if i+1 >= len(msgs) || msgs[i+1].Role != fantasy.MessageRoleSystem { - result = append(result, fantasy.Message{ - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: instr}, - }, - }) - inserted = true - } - } - } - if !inserted { - // No system messages, prepend. - result = append([]fantasy.Message{{ - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: instr}, - }, - }}, result...) - } - return result - }, - }) - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - - // The second LLM call should contain the injected instruction. - require.NotEmpty(t, secondCallPrompt) - var foundInstruction bool - for _, msg := range secondCallPrompt { - if msg.Role != fantasy.MessageRoleSystem { - continue - } - for _, part := range msg.Content { - if tp, ok := fantasy.AsMessagePart[fantasy.TextPart](part); ok { - if strings.Contains(tp.Text, "AGENTS.md") { - foundInstruction = true - } - } - } - } - require.True(t, foundInstruction, - "step 1 prompt should contain the injected system instruction") -} - -func TestRun_PrepareMessagesOnlyFiresOnce(t *testing.T) { - t.Parallel() - - var mu sync.Mutex - var streamCalls int - - // Three steps: tool call, tool call, text. PrepareMessages - // should inject on step 1 and return nil on step 2. - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - if step < 2 { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-" + strings.Repeat("x", step+1), ToolCallName: "noop"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-" + strings.Repeat("x", step+1), Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-" + strings.Repeat("x", step+1)}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-" + strings.Repeat("x", step+1), - ToolCallName: "noop", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - } - 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}, - }), nil - }, - } - - var prepareCalls atomic.Int32 - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "do something"), - }, - Tools: []fantasy.AgentTool{newNoopTool("noop")}, - MaxSteps: 5, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message { - call := prepareCalls.Add(1) - if call == 1 { - // First call: inject a message. - return append(msgs, fantasy.Message{ - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{fantasy.TextPart{Text: "injected"}}, - }) - } - // Subsequent calls: no changes. - return nil - }, - }) - require.NoError(t, err) - require.Equal(t, 3, streamCalls) - // PrepareMessages is called before each of the 3 steps. - require.Equal(t, 3, int(prepareCalls.Load())) -} - -// TestRun_PrepareToolsInjectsToolMidLoop guards the regression where a -// chat creating its workspace mid-turn (via create_workspace) saw the -// workspace MCP tools only on the next turn. Before the fix, the tool -// list was frozen at the top of the turn and the model could not call -// any workspace MCP tools until turn 2. With the fix, PrepareTools is -// invoked before every step and can inject tools that become available -// mid-loop. -func TestRun_PrepareToolsInjectsToolMidLoop(t *testing.T) { - t.Parallel() - - const injectedToolName = "workspace_mcp__echo" - - var mu sync.Mutex - var streamCalls int - var secondCallTools []fantasy.Tool - - // Step 0 calls create_workspace. Step 1 should see the - // injected workspace MCP tool. - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - switch step { - case 0: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "create_workspace"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "create_workspace", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - default: - mu.Lock() - secondCallTools = append([]fantasy.Tool(nil), call.Tools...) - mu.Unlock() - 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}, - }), nil - } - }, - } - - var workspaceReady atomic.Bool - createWorkspaceTool := fantasy.NewAgentTool( - "create_workspace", - "create a workspace", - func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - workspaceReady.Store(true) - return fantasy.ToolResponse{}, nil - }, - ) - - var prepareCalls atomic.Int32 - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "create a workspace and use MCP"), - }, - Tools: []fantasy.AgentTool{createWorkspaceTool}, - ActiveTools: []string{"create_workspace"}, - MaxSteps: 5, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { - prepareCalls.Add(1) - if !workspaceReady.Load() { - return nil - } - return append(currentTools, newNoopTool(injectedToolName)) - }, - }) - require.NoError(t, err) - require.Equal(t, 2, streamCalls) - // PrepareTools is called before each of the 2 steps. - require.Equal(t, int32(2), prepareCalls.Load()) - - require.NotEmpty(t, secondCallTools) - var foundInjectedTool bool - for _, tool := range secondCallTools { - if tool.GetName() == injectedToolName { - foundInjectedTool = true - break - } - } - require.True(t, foundInjectedTool, - "step 1 prompt should advertise the workspace MCP tool injected by PrepareTools") -} - -// TestRun_PrepareToolsAddsNewToolToActiveSet guards the contract that -// when PrepareTools injects a tool, that tool is callable on the -// next step even when opts.ActiveTools was non-empty (and would -// otherwise filter the new tool out). -func TestRun_PrepareToolsAddsNewToolToActiveSet(t *testing.T) { - t.Parallel() - - const injectedToolName = "workspace_mcp__echo" - - var mu sync.Mutex - var streamCalls int - var injectedToolRan atomic.Bool - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCalls - streamCalls++ - mu.Unlock() - - switch step { - case 0: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "create_workspace"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "create_workspace", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - case 1: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-2", ToolCallName: injectedToolName}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-2", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-2"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-2", - ToolCallName: injectedToolName, - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - }), nil - default: - 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}, - }), nil - } - }, - } - - var workspaceReady atomic.Bool - createWorkspaceTool := fantasy.NewAgentTool( - "create_workspace", - "create a workspace", - func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - workspaceReady.Store(true) - return fantasy.ToolResponse{}, nil - }, - ) - - injectedTool := fantasy.NewAgentTool( - injectedToolName, - "injected workspace MCP tool", - func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - injectedToolRan.Store(true) - return fantasy.ToolResponse{}, nil - }, - ) - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "create a workspace and use MCP"), - }, - Tools: []fantasy.AgentTool{createWorkspaceTool}, - // Active list deliberately excludes the injected tool name; - // PrepareTools must add it so the tool is callable. - ActiveTools: []string{"create_workspace"}, - MaxSteps: 5, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - PrepareTools: func(currentTools []fantasy.AgentTool) []fantasy.AgentTool { - if !workspaceReady.Load() { - return nil - } - for _, t := range currentTools { - if t.Info().Name == injectedToolName { - return nil - } - } - return append(currentTools, injectedTool) - }, - }) - require.NoError(t, err) - require.GreaterOrEqual(t, streamCalls, 2) - require.True(t, injectedToolRan.Load(), - "injected tool must be callable on the step after PrepareTools adds it") -} - func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { t.Parallel() @@ -4633,170 +985,3 @@ func TestExecuteSingleTool_MediaBase64Encoding(t *testing.T) { require.Contains(t, textOutput.Text, "world") }) } - -// TestRun_ReasoningTimestamps verifies that StreamPartTypeReasoningStart -// and StreamPartTypeReasoningEnd produce parallel ReasoningStartedAt / -// ReasoningCompletedAt slices on PersistedStep, in the same occurrence -// order as the reasoning content blocks. The frontend computes -// reasoning duration as completed_at - started_at. -func TestRun_ReasoningTimestamps(t *testing.T) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeReasoningStart, ID: "reason-1"}, - {Type: fantasy.StreamPartTypeReasoningDelta, ID: "reason-1", Delta: "first thought"}, - {Type: fantasy.StreamPartTypeReasoningEnd, ID: "reason-1"}, - {Type: fantasy.StreamPartTypeReasoningStart, ID: "reason-2"}, - {Type: fantasy.StreamPartTypeReasoningDelta, ID: "reason-2", Delta: "second thought"}, - {Type: fantasy.StreamPartTypeReasoningEnd, ID: "reason-2"}, - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "answer"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - }), nil - }, - } - - var persistedSteps []PersistedStep - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "think"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedSteps = append(persistedSteps, step) - return nil - }, - }) - require.NoError(t, err) - require.Len(t, persistedSteps, 1) - - step := persistedSteps[0] - - // Both reasoning blocks must produce parallel timestamp entries. - require.Len(t, step.ReasoningStartedAt, 2, - "each StreamPartTypeReasoningEnd must record a started_at") - require.Len(t, step.ReasoningCompletedAt, 2, - "each StreamPartTypeReasoningEnd must record a completed_at") - - // Timestamps must be monotonic per block (completed_at >= started_at), - // and both timestamps must be populated. Asserting only monotonicity - // is not enough: time.Time{} is year 0001, so completed_at.Before(zero) - // is trivially false and a regression that drops the started_at stamp - // would slip past the comparison. - for i := range step.ReasoningStartedAt { - require.False(t, step.ReasoningStartedAt[i].IsZero(), - "started_at[%d] must be non-zero", i) - require.False(t, step.ReasoningCompletedAt[i].IsZero(), - "completed_at[%d] must be non-zero", i) - require.False(t, - step.ReasoningCompletedAt[i].Before(step.ReasoningStartedAt[i]), - "completed_at[%d] must be >= started_at[%d]", i, i) - } - - // Successive blocks must be ordered: reasoning-2 cannot start - // before reasoning-1 completes. - require.False(t, - step.ReasoningStartedAt[1].Before(step.ReasoningCompletedAt[0]), - "reasoning-2 started_at must be >= reasoning-1 completed_at") - - // The reasoning content blocks must appear in the same order - // in step.Content so the persistence layer can correlate by - // occurrence order. - var reasoningOrder []string - for _, c := range step.Content { - if r, ok := fantasy.AsContentType[fantasy.ReasoningContent](c); ok { - reasoningOrder = append(reasoningOrder, r.Text) - } - } - require.Equal(t, []string{"first thought", "second thought"}, reasoningOrder) -} - -func TestRun_InterruptedReasoningFlushesTimestamps(t *testing.T) { - t.Parallel() - - started := make(chan struct{}) - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(ctx context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return iter.Seq[fantasy.StreamPart](func(yield func(fantasy.StreamPart) bool) { - parts := []fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeReasoningStart, ID: "reason-1"}, - {Type: fantasy.StreamPartTypeReasoningDelta, ID: "reason-1", Delta: "interrupted thought"}, - } - for _, part := range parts { - if !yield(part) { - return - } - } - - select { - case <-started: - default: - close(started) - } - - <-ctx.Done() - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeError, - Error: ctx.Err(), - }) - }), nil - }, - } - - ctx, cancel := context.WithCancelCause(context.Background()) - defer cancel(nil) - - go func() { - <-started - cancel(ErrInterrupted) - }() - - var persistedStep PersistedStep - err := Run(ctx, RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "think"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, step PersistedStep) error { - persistedStep = step - return nil - }, - }) - require.ErrorIs(t, err, ErrInterrupted) - - // flushActiveState must have appended exactly one entry to each - // parallel slice, matching the single in-progress reasoning block. - require.Len(t, persistedStep.ReasoningStartedAt, 1, - "interrupted reasoning must flush its started_at") - require.Len(t, persistedStep.ReasoningCompletedAt, 1, - "interrupted reasoning must flush a completed_at stamp") - - // Both timestamps must be populated and the completed stamp - // must be at or after the started stamp. - require.False(t, persistedStep.ReasoningStartedAt[0].IsZero(), - "flushed reasoning started_at must be non-zero") - require.False(t, persistedStep.ReasoningCompletedAt[0].IsZero(), - "flushed reasoning completed_at must be non-zero") - require.False(t, - persistedStep.ReasoningCompletedAt[0].Before(persistedStep.ReasoningStartedAt[0]), - "flushed completed_at must be >= started_at") - - // The flushed reasoning content must appear in step.Content so - // the persistence layer's occurrence-order correlation lines up - // with the timestamp slices. - var reasoningBlocks []fantasy.ReasoningContent - for _, c := range persistedStep.Content { - if r, ok := fantasy.AsContentType[fantasy.ReasoningContent](c); ok { - reasoningBlocks = append(reasoningBlocks, r) - } - } - require.Len(t, reasoningBlocks, 1) - require.Equal(t, "interrupted thought", reasoningBlocks[0].Text) -} diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index b267f17e2a..330def364f 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -90,45 +90,36 @@ type CompactionResult struct { ContextLimit int64 } -// 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, - model fantasy.LanguageModel, - compaction *CompactionOptions, - contextLimitFallback int64, - stepUsage fantasy.Usage, - stepMetadata fantasy.ProviderMetadata, - allMessages []fantasy.Message, -) (bool, error) { - config, ok := normalizedCompactionConfig(compaction) +// GenerateCompaction generates one context summary and returns it without +// persisting. It publishes compaction progress parts when configured. +func GenerateCompaction(ctx context.Context, opts GenerateCompactionOptions) (CompactionResult, error) { + if opts.Model == nil { + return CompactionResult{}, xerrors.New("chat model is required") + } + config, ok := normalizedCompactionGenerateConfig(opts) if !ok { - return false, nil + return CompactionResult{}, nil } - contextTokens := contextTokensFromUsage(stepUsage) + contextTokens := contextTokensFromUsage(opts.StepUsage) if contextTokens <= 0 { - return false, nil + return CompactionResult{}, nil } - - metadataLimit := extractContextLimit(stepMetadata) + metadataLimit := extractContextLimit(opts.StepMetadata) contextLimit := resolveContextLimit( metadataLimit.Int64, config.ContextLimit, - contextLimitFallback, + opts.ContextLimitFallback, ) - usagePercent, compact := shouldCompact( - contextTokens, contextLimit, config.ThresholdPercent, + contextTokens, + contextLimit, + config.ThresholdPercent, ) if !compact { - return false, nil + return CompactionResult{}, nil } - // Publish the "Summarizing..." tool-call indicator so - // connected clients see activity during summary generation. if config.PublishMessagePart != nil && config.ToolCallID != "" { config.PublishMessagePart( codersdk.ChatMessageRoleAssistant, @@ -136,40 +127,26 @@ func tryCompact( ) } - summary, err := generateCompactionSummary( - ctx, model, allMessages, config, - ) + summary, err := generateCompactionSummary(ctx, opts.Model, opts.Messages, config) if err != nil { - return false, err + publishCompactionError(config, "failed to generate compaction summary") + return CompactionResult{}, 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") + return CompactionResult{}, xerrors.New("compaction produced an empty summary") } - systemSummary := strings.TrimSpace( - config.SystemSummaryPrefix + "\n\n" + summary, - ) - - persistCtx := context.WithoutCancel(ctx) - err = config.Persist(persistCtx, CompactionResult{ - SystemSummary: systemSummary, + result := CompactionResult{ + SystemSummary: strings.TrimSpace( + config.SystemSummaryPrefix + "\n\n" + summary, + ), 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, @@ -184,37 +161,22 @@ func tryCompact( codersdk.ChatMessageToolResult(config.ToolCallID, config.ToolName, resultJSON, false, false), ) } - - return true, nil + return result, 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( - codersdk.ChatMessageRoleTool, - codersdk.ChatMessageToolResult(config.ToolCallID, config.ToolName, errJSON, true, false), - ) -} - -// 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 CompactionOptions{}, false +func normalizedCompactionGenerateConfig(opts GenerateCompactionOptions) (CompactionOptions, bool) { + config := CompactionOptions{ + ThresholdPercent: opts.ThresholdPercent, + ContextLimit: opts.ContextLimit, + SummaryPrompt: opts.SummaryPrompt, + SystemSummaryPrefix: opts.SystemSummaryPrefix, + Timeout: opts.Timeout, + DebugSvc: opts.DebugSvc, + ChatID: opts.ChatID, + HistoryTipMessageID: opts.HistoryTipMessageID, + ToolCallID: opts.ToolCallID, + ToolName: opts.ToolName, + PublishMessagePart: opts.PublishMessagePart, } if strings.TrimSpace(config.SummaryPrompt) == "" { config.SummaryPrompt = defaultCompactionSummaryPrompt @@ -232,10 +194,24 @@ func normalizedCompactionConfig(opts *CompactionOptions) (CompactionOptions, boo if config.ThresholdPercent == maxCompactionThresholdPercent { return CompactionOptions{}, false } - return config, true } +// 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( + codersdk.ChatMessageRoleTool, + codersdk.ChatMessageToolResult(config.ToolCallID, config.ToolName, errJSON, true, false), + ) +} + // 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 diff --git a/coderd/x/chatd/chatloop/compaction_internal_test.go b/coderd/x/chatd/chatloop/compaction_internal_test.go index ae26ed8cf0..ba6580465d 100644 --- a/coderd/x/chatd/chatloop/compaction_internal_test.go +++ b/coderd/x/chatd/chatloop/compaction_internal_test.go @@ -3,7 +3,6 @@ package chatloop import ( "context" "encoding/json" - "sync" "testing" "time" @@ -18,7 +17,6 @@ import ( "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/coderd/x/chatd/chattest" - "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -235,778 +233,3 @@ func TestGenerateCompactionSummary_PanicFinalizesAsError(t *testing.T) { t.Fatal("FinalizeRun never reached UpdateChatDebugRun on panic") } } - -func TestRun_Compaction(t *testing.T) { - t.Parallel() - - t.Run("PersistsWhenThresholdReached", func(t *testing.T) { - t.Parallel() - - persistCompactionCalls := 0 - var persistedCompaction CompactionResult - const summaryText = "summary text for compaction" - - model := &chattest.FakeModel{ - ProviderName: "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: "done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 80, - TotalTokens: 85, - }, - }, - }), nil - }, - GenerateFn: func(_ context.Context, call fantasy.Call) (*fantasy.Response, error) { - require.NotEmpty(t, call.Prompt) - lastPrompt := call.Prompt[len(call.Prompt)-1] - require.Equal(t, fantasy.MessageRoleUser, lastPrompt.Role) - require.Len(t, lastPrompt.Content, 1) - - instruction, ok := fantasy.AsMessagePart[fantasy.TextPart](lastPrompt.Content[0]) - require.True(t, ok) - require.Equal(t, "summarize now", instruction.Text) - - return &fantasy.Response{ - Content: []fantasy.Content{ - fantasy.TextContent{Text: summaryText}, - }, - }, nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ContextLimitFallback: 100, - Compaction: &CompactionOptions{ - ThresholdPercent: 70, - SummaryPrompt: "summarize now", - Persist: func(_ context.Context, result CompactionResult) error { - persistCompactionCalls++ - persistedCompaction = result - return nil - }, - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, nil - }, - }) - require.NoError(t, err) - // Compaction fires twice: once inline when the threshold is - // reached on step 0 (the only step, since MaxSteps=1), and - // once from the post-run safety net during the re-entry - // iteration (where totalSteps already equals MaxSteps so the - // inner loop doesn't execute, but lastUsage still exceeds - // the threshold). - require.Equal(t, 2, persistCompactionCalls) - require.Contains(t, persistedCompaction.SystemSummary, summaryText) - require.Equal(t, summaryText, persistedCompaction.SummaryReport) - require.Equal(t, int64(80), persistedCompaction.ContextTokens) - require.Equal(t, int64(100), persistedCompaction.ContextLimit) - require.InDelta(t, 80.0, persistedCompaction.UsagePercent, 0.0001) - }) - - t.Run("PublishesPartsBeforeAndAfterPersist", func(t *testing.T) { - t.Parallel() - - const summaryText = "compaction summary for ordering test" - - // 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 := &chattest.FakeModel{ - ProviderName: "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: "done"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 80, - TotalTokens: 85, - }, - }, - }), nil - }, - GenerateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { - callOrder = append(callOrder, "generate") - return &fantasy.Response{ - Content: []fantasy.Content{ - fantasy.TextContent{Text: summaryText}, - }, - }, nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ContextLimitFallback: 100, - Compaction: &CompactionOptions{ - ThresholdPercent: 70, - SummaryPrompt: "summarize now", - ToolCallID: "test-tool-call-id", - ToolName: "chat_summarized", - PublishMessagePart: func(role codersdk.ChatMessageRole, 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") - return nil - }, - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, nil - }, - }) - require.NoError(t, err) - // Compaction fires twice (see PersistsWhenThresholdReached - // for the full explanation). Each cycle follows the order: - // publish_tool_call → generate → persist → publish_tool_result. - require.Equal(t, []string{ - "publish_tool_call", - "generate", - "persist", - "publish_tool_result", - "publish_tool_call", - "generate", - "persist", - "publish_tool_result", - }, callOrder) - }) - - t.Run("PublishNotCalledBelowThreshold", func(t *testing.T) { - t.Parallel() - - publishCalled := false - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 10, - }, - }, - }), nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ContextLimitFallback: 100, - Compaction: &CompactionOptions{ - ThresholdPercent: 70, - ToolCallID: "test-tool-call-id", - ToolName: "chat_summarized", - PublishMessagePart: func(_ codersdk.ChatMessageRole, _ codersdk.ChatMessagePart) { - publishCalled = true - }, - Persist: func(_ context.Context, _ CompactionResult) error { - return nil - }, - }, - }) - require.NoError(t, err) - 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 := &chattest.FakeModel{ - ProviderName: "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 := &chattest.FakeModel{ - ProviderName: "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) { - t.Parallel() - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 80, - }, - }, - }), nil - }, - GenerateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { - return nil, xerrors.New("generate failed") - }, - } - - compactionErr := xerrors.New("unset") - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ContextLimitFallback: 100, - Compaction: &CompactionOptions{ - ThresholdPercent: 70, - Persist: func(_ context.Context, _ CompactionResult) error { - return nil - }, - OnError: func(err error) { - compactionErr = err - }, - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, nil - }, - }) - require.NoError(t, err) - require.Error(t, compactionErr) - require.ErrorContains(t, compactionErr, "generate summary text") - }) - - t.Run("PostRunCompactionReEntersStepLoop", func(t *testing.T) { - t.Parallel() - - // When post-run compaction fires (no mid-loop compaction) - // and ReloadMessages is provided, Run should re-enter the - // step loop with the reloaded messages so the agent - // continues working. - - var mu sync.Mutex - var streamCallCount int - persistCompactionCalls := 0 - reloadCalls := 0 - - const summaryText = "post-run compacted summary" - - compactedMessages := []fantasy.Message{ - textMessage(fantasy.MessageRoleSystem, "compacted system"), - textMessage(fantasy.MessageRoleUser, "compacted user"), - } - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCallCount - streamCallCount++ - mu.Unlock() - - switch step { - case 0: - // First turn: text-only response with high usage. - // No tool calls, so shouldContinue = false and - // the inner step loop breaks. Compaction should - // fire, then the outer loop re-enters. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "initial response"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 80, - TotalTokens: 85, - }, - }, - }), nil - default: - // Second turn (after compaction re-entry): - // text-only with low usage — should finish. - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-2"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-2", Delta: "continued after compaction"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-2"}, - { - 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 - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - 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 on the final step of the first pass. - // The inline path fires (ReloadMessages is set) and then - // the outer loop re-enters. On the second pass the usage - // is below threshold so no further compaction occurs. - require.GreaterOrEqual(t, persistCompactionCalls, 1) - // ReloadMessages was called (inline + re-entry). - require.GreaterOrEqual(t, reloadCalls, 1) - // Two stream calls: one before compaction, one after re-entry. - require.Equal(t, 2, streamCallCount) - }) - - t.Run("PostRunCompactionReEntryIncludesUserSummary", func(t *testing.T) { - t.Parallel() - - // After compaction the summary is stored as a user-role - // message. When the loop re-enters, the reloaded prompt - // must contain this user message so the LLM provider - // receives a valid prompt (providers like Anthropic - // require at least one non-system message). - - var mu sync.Mutex - var streamCallCount int - var reEntryPrompt []fantasy.Message - persistCompactionCalls := 0 - - const summaryText = "post-run compacted summary" - - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - mu.Lock() - step := streamCallCount - streamCallCount++ - mu.Unlock() - - switch step { - case 0: - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-1", Delta: "initial response"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-1"}, - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - Usage: fantasy.Usage{ - InputTokens: 80, - TotalTokens: 85, - }, - }, - }), nil - default: - mu.Lock() - reEntryPrompt = append([]fantasy.Message(nil), call.Prompt...) - mu.Unlock() - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "text-2"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "text-2", Delta: "continued"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "text-2"}, - { - 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 - }, - } - - // Simulate real post-compaction DB state: the summary is - // a user-role message (the only non-system content). - compactedMessages := []fantasy.Message{ - textMessage(fantasy.MessageRoleSystem, "system prompt"), - textMessage(fantasy.MessageRoleUser, "Summary of earlier chat context:\n\ncompacted summary"), - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - 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) - - require.GreaterOrEqual(t, persistCompactionCalls, 1) - // Re-entry happened: stream was called at least twice. - require.Equal(t, 2, streamCallCount) - // The re-entry prompt must contain the user summary. - require.NotEmpty(t, reEntryPrompt) - hasUser := false - for _, msg := range reEntryPrompt { - if msg.Role == fantasy.MessageRoleUser { - hasUser = true - break - } - } - require.True(t, hasUser, "re-entry prompt must contain a user message (the compaction summary)") - }) - - t.Run("TriggersOnDynamicToolExit", func(t *testing.T) { - t.Parallel() - - var persistCompactionCalls int - const summaryText = "compaction summary for dynamic tool exit" - - // The LLM calls a dynamic tool. Usage is above the - // compaction threshold so compaction should fire even - // though the chatloop exits via ErrDynamicToolCall. - model := &chattest.FakeModel{ - ProviderName: "fake", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return streamFromParts([]fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc-1", ToolCallName: "my_dynamic_tool"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc-1", Delta: `{"query": "test"}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc-1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc-1", - ToolCallName: "my_dynamic_tool", - ToolCallInput: `{"query": "test"}`, - }, - { - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonToolCalls, - Usage: fantasy.Usage{ - InputTokens: 80, - TotalTokens: 85, - }, - }, - }), nil - }, - GenerateFn: func(_ context.Context, _ fantasy.Call) (*fantasy.Response, error) { - return &fantasy.Response{ - Content: []fantasy.Content{ - fantasy.TextContent{Text: summaryText}, - }, - }, nil - }, - } - - err := Run(context.Background(), RunOptions{ - Model: model, - Messages: []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, - MaxSteps: 5, - DynamicToolNames: map[string]bool{"my_dynamic_tool": true}, - PersistStep: func(_ context.Context, _ PersistedStep) error { - return nil - }, - ContextLimitFallback: 100, - Compaction: &CompactionOptions{ - ThresholdPercent: 70, - SummaryPrompt: "summarize now", - Persist: func(_ context.Context, result CompactionResult) error { - persistCompactionCalls++ - require.Contains(t, result.SystemSummary, summaryText) - return nil - }, - }, - ReloadMessages: func(_ context.Context) ([]fantasy.Message, error) { - return []fantasy.Message{ - textMessage(fantasy.MessageRoleUser, "hello"), - }, nil - }, - }) - require.ErrorIs(t, err, ErrDynamicToolCall) - require.Equal(t, 1, persistCompactionCalls, - "compaction must fire before dynamic tool exit") - }) -} diff --git a/coderd/x/chatd/chatloop/metrics_test.go b/coderd/x/chatd/chatloop/metrics_test.go index 40eabf99ca..e414e91fab 100644 --- a/coderd/x/chatd/chatloop/metrics_test.go +++ b/coderd/x/chatd/chatloop/metrics_test.go @@ -4,7 +4,6 @@ import ( "context" "strconv" "testing" - "time" "charm.land/fantasy" "github.com/prometheus/client_golang/prometheus" @@ -15,7 +14,6 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chaterror" "github.com/coder/coder/v2/coderd/x/chatd/chatloop" - "github.com/coder/coder/v2/coderd/x/chatd/chatretry" "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/codersdk" ) @@ -411,117 +409,12 @@ func TestRecordToolError(t *testing.T) { }) } -func TestRun_RecordsMetrics(t *testing.T) { +func TestGenerateAssistant_StreamRetryRecordsMetric(t *testing.T) { t.Parallel() reg := prometheus.NewRegistry() metrics := chatloop.NewMetrics(reg) - model := &chattest.FakeModel{ - ProviderName: "test-provider", - ModelName: "test-model", - StreamFn: func(_ context.Context, call fantasy.Call) (fantasy.StreamResponse, error) { - return func(yield func(fantasy.StreamPart) bool) { - parts := []fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeTextStart, ID: "t1"}, - {Type: fantasy.StreamPartTypeTextDelta, ID: "t1", Delta: "hello"}, - {Type: fantasy.StreamPartTypeTextEnd, ID: "t1"}, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonStop}, - } - for _, p := range parts { - if !yield(p) { - return - } - } - }, nil - }, - } - - err := chatloop.Run(context.Background(), chatloop.RunOptions{ - Model: model, - Messages: []fantasy.Message{ - { - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: "hello"}, - }, - }, - }, - MaxSteps: 1, - PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error { - return nil - }, - Metrics: metrics, - }) - require.NoError(t, err) - - families, err := reg.Gather() - require.NoError(t, err) - - assertProviderModelLabels := func(t *testing.T, metric *dto.Metric) { - t.Helper() - labels := map[string]string{} - for _, lp := range metric.GetLabel() { - labels[lp.GetName()] = lp.GetValue() - } - assert.Equal(t, "test-provider", labels["provider"]) - assert.Equal(t, "test-model", labels["model"]) - } - - found := make(map[string]bool) - for _, f := range families { - found[f.GetName()] = true - - switch f.GetName() { - case "coderd_chatd_steps_total": - require.Len(t, f.GetMetric(), 1) - assert.Equal(t, float64(1), f.GetMetric()[0].GetCounter().GetValue(), - "steps_total should be 1 after one step") - assertProviderModelLabels(t, f.GetMetric()[0]) - case "coderd_chatd_message_count": - require.Len(t, f.GetMetric(), 1) - assert.Equal(t, uint64(1), f.GetMetric()[0].GetHistogram().GetSampleCount(), - "message_count should have 1 observation") - assertProviderModelLabels(t, f.GetMetric()[0]) - case "coderd_chatd_prompt_size_bytes": - require.Len(t, f.GetMetric(), 1) - assert.Equal(t, uint64(1), f.GetMetric()[0].GetHistogram().GetSampleCount(), - "prompt_size_bytes should have 1 observation") - assertProviderModelLabels(t, f.GetMetric()[0]) - case "coderd_chatd_ttft_seconds": - require.Len(t, f.GetMetric(), 1) - assert.Equal(t, uint64(1), f.GetMetric()[0].GetHistogram().GetSampleCount(), - "ttft_seconds should have 1 observation") - assertProviderModelLabels(t, f.GetMetric()[0]) - } - } - - assert.True(t, found["coderd_chatd_steps_total"], "steps_total not recorded") - assert.True(t, found["coderd_chatd_message_count"], "message_count not recorded") - assert.True(t, found["coderd_chatd_prompt_size_bytes"], "prompt_size_bytes not recorded") - assert.True(t, found["coderd_chatd_ttft_seconds"], "ttft_seconds not recorded") -} - -// TestRun_StreamRetry_RecordsMetric exercises the end-to-end retry -// path: a retryable error on the first Stream call, success on the -// second. Asserts both the metric and the back-compat OnRetry -// callback fire. -// -// Note: chatretry.Retry uses time.NewTimer (not quartz.Clock), so -// this test pays chatretry.InitialDelay (1s) of real wall-clock -// time per retry. Keep it to one retry. -func TestRun_StreamRetry_RecordsMetric(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - metrics := chatloop.NewMetrics(reg) - - type retryCall struct { - attempt int - classified chatretry.ClassifiedError - } - var retries []retryCall - calls := 0 model := &chattest.FakeModel{ ProviderName: "test-provider", @@ -529,7 +422,14 @@ func TestRun_StreamRetry_RecordsMetric(t *testing.T) { StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { calls++ if calls == 1 { - return nil, xerrors.New("received status 429 from upstream") + return nil, chaterror.WithClassification( + xerrors.New("received status 429 from upstream"), + chaterror.ClassifiedError{ + Kind: codersdk.ChatErrorKindRateLimit, + Provider: "test-provider", + Retryable: true, + }, + ) } return func(yield func(fantasy.StreamPart) bool) { yield(fantasy.StreamPart{ @@ -540,35 +440,12 @@ func TestRun_StreamRetry_RecordsMetric(t *testing.T) { }, } - err := chatloop.Run(context.Background(), chatloop.RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error { - return nil - }, + _, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, Metrics: metrics, - OnRetry: func( - attempt int, - _ error, - classified chatretry.ClassifiedError, - _ time.Duration, - ) { - retries = append(retries, retryCall{ - attempt: attempt, - classified: classified, - }) - }, }) - require.NoError(t, err) - - // Back-compat: OnRetry still fires with classified error. - require.Len(t, retries, 1) - assert.Equal(t, 1, retries[0].attempt) - assert.Equal(t, codersdk.ChatErrorKindRateLimit, retries[0].classified.Kind) - assert.Equal(t, "test-provider", retries[0].classified.Provider) - - // Metric assertion. + require.Error(t, err) + require.Equal(t, 1, calls) requireCounter(t, reg, "coderd_chatd_stream_retries_total", 1, map[string]string{ "provider": "test-provider", "model": "test-model", @@ -577,10 +454,10 @@ func TestRun_StreamRetry_RecordsMetric(t *testing.T) { }) } -// TestRun_StreamRetry_ContextCanceledTransportResetIncrements pins the -// invariant that provider-originated context cancellation is counted as -// a retryable transport reset when the chat context is still alive. -func TestRun_StreamRetry_ContextCanceledTransportResetIncrements(t *testing.T) { +// TestGenerateAssistant_StreamRetry_ContextCanceledTransportResetIncrements pins the +// invariant that provider-originated context cancellation is counted as a +// retryable transport reset when the chat context is still alive. +func TestGenerateAssistant_StreamRetry_ContextCanceledTransportResetIncrements(t *testing.T) { t.Parallel() reg := prometheus.NewRegistry() @@ -592,29 +469,16 @@ func TestRun_StreamRetry_ContextCanceledTransportResetIncrements(t *testing.T) { ModelName: "test-model", StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { attempts++ - if attempts == 1 { - return nil, context.Canceled - } - return func(yield func(fantasy.StreamPart) bool) { - _ = yield(fantasy.StreamPart{ - Type: fantasy.StreamPartTypeFinish, - FinishReason: fantasy.FinishReasonStop, - }) - }, nil + return nil, context.Canceled }, } - err := chatloop.Run(context.Background(), chatloop.RunOptions{ - Model: model, - MaxSteps: 1, - ContextLimitFallback: 4096, - PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error { - return nil - }, + _, err := chatloop.GenerateAssistant(context.Background(), chatloop.GenerateAssistantOptions{ + Model: model, Metrics: metrics, }) - require.NoError(t, err) - require.Equal(t, 2, attempts) + require.Error(t, err) + require.Equal(t, 1, attempts) requireCounter(t, reg, "coderd_chatd_stream_retries_total", 1, map[string]string{ "provider": "test-provider", @@ -623,105 +487,3 @@ func TestRun_StreamRetry_ContextCanceledTransportResetIncrements(t *testing.T) { "chain_broken": "false", }) } - -func TestRun_ToolError_RecordsMetric(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - toolFn func(context.Context, struct{}, fantasy.ToolCall) (fantasy.ToolResponse, error) - builtinToolNames map[string]bool - wantLabel string - }{ - { - name: "builtin_tool_IsError", - toolFn: func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.ToolResponse{ - Content: "something went wrong", - IsError: true, - }, nil - }, - builtinToolNames: map[string]bool{"failing_tool": true}, - wantLabel: "failing_tool", - }, - { - name: "mcp_tool_IsError", - toolFn: func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.ToolResponse{ - Content: "something went wrong", - IsError: true, - }, nil - }, - builtinToolNames: map[string]bool{}, - wantLabel: "failing_tool", - }, - { - name: "tool_Run_returns_error", - toolFn: func(_ context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { - return fantasy.ToolResponse{}, xerrors.New("connection refused") - }, - builtinToolNames: map[string]bool{"failing_tool": true}, - wantLabel: "failing_tool", - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - metrics := chatloop.NewMetrics(reg) - - failingTool := fantasy.NewAgentTool( - "failing_tool", - "a tool that always fails", - tt.toolFn, - ) - - model := &chattest.FakeModel{ - ProviderName: "test-provider", - ModelName: "test-model", - StreamFn: func(_ context.Context, _ fantasy.Call) (fantasy.StreamResponse, error) { - return func(yield func(fantasy.StreamPart) bool) { - parts := []fantasy.StreamPart{ - {Type: fantasy.StreamPartTypeToolInputStart, ID: "tc1", ToolCallName: "failing_tool"}, - {Type: fantasy.StreamPartTypeToolInputDelta, ID: "tc1", Delta: `{}`}, - {Type: fantasy.StreamPartTypeToolInputEnd, ID: "tc1"}, - { - Type: fantasy.StreamPartTypeToolCall, - ID: "tc1", - ToolCallName: "failing_tool", - ToolCallInput: `{}`, - }, - {Type: fantasy.StreamPartTypeFinish, FinishReason: fantasy.FinishReasonToolCalls}, - } - for _, p := range parts { - if !yield(p) { - return - } - } - }, nil - }, - } - - err := chatloop.Run(context.Background(), chatloop.RunOptions{ - Model: model, - MaxSteps: 1, - Tools: []fantasy.AgentTool{failingTool}, - ActiveTools: []string{"failing_tool"}, - BuiltinToolNames: tt.builtinToolNames, - PersistStep: func(_ context.Context, _ chatloop.PersistedStep) error { - return nil - }, - Metrics: metrics, - }) - require.NoError(t, err) - - requireCounter(t, reg, "coderd_chatd_tool_errors_total", 1, map[string]string{ - "provider": "test-provider", - "model": "test-model", - "tool_name": tt.wantLabel, - }) - }) - } -} diff --git a/coderd/x/chatd/chatloop/publish_context.go b/coderd/x/chatd/chatloop/publish_context.go new file mode 100644 index 0000000000..25348827a1 --- /dev/null +++ b/coderd/x/chatd/chatloop/publish_context.go @@ -0,0 +1,32 @@ +package chatloop + +import ( + "context" + + "github.com/coder/coder/v2/codersdk" +) + +type messagePartPublisherKey struct{} + +// WithMessagePartPublisher returns a context carrying the streaming +// message-part publisher so tools can stream intermediate output (e.g. +// advisor advice deltas) while they execute. ExecuteLocalTools injects +// the publisher before running tools. +func WithMessagePartPublisher( + ctx context.Context, + publish func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), +) context.Context { + if publish == nil { + return ctx + } + return context.WithValue(ctx, messagePartPublisherKey{}, publish) +} + +// MessagePartPublisherFromContext returns the publisher injected by +// ExecuteLocalTools, or nil when absent. +func MessagePartPublisherFromContext( + ctx context.Context, +) func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) { + publish, _ := ctx.Value(messagePartPublisherKey{}).(func(codersdk.ChatMessageRole, codersdk.ChatMessagePart)) + return publish +} diff --git a/coderd/x/chatd/chatloop/publish_context_internal_test.go b/coderd/x/chatd/chatloop/publish_context_internal_test.go new file mode 100644 index 0000000000..3792df0439 --- /dev/null +++ b/coderd/x/chatd/chatloop/publish_context_internal_test.go @@ -0,0 +1,60 @@ +package chatloop + +import ( + "context" + "testing" + + "charm.land/fantasy" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/codersdk" + "github.com/coder/quartz" +) + +func TestMessagePartPublisherContextRoundTrip(t *testing.T) { + t.Parallel() + + require.Nil(t, MessagePartPublisherFromContext(context.Background())) + + var published []codersdk.ChatMessagePart + publish := func(_ codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + published = append(published, part) + } + ctx := WithMessagePartPublisher(context.Background(), publish) + got := MessagePartPublisherFromContext(ctx) + require.NotNil(t, got) + got(codersdk.ChatMessageRoleTool, codersdk.ChatMessagePart{ToolCallID: "call-1"}) + require.Len(t, published, 1) + require.Equal(t, "call-1", published[0].ToolCallID) + + // A nil publisher must not be stored. + require.Nil(t, MessagePartPublisherFromContext(WithMessagePartPublisher(context.Background(), nil))) +} + +func TestExecuteLocalToolsInjectsMessagePartPublisher(t *testing.T) { + t.Parallel() + + var toolSawPublisher bool + tool := fantasy.NewAgentTool( + "probe", + "reports whether the execution context carries a publisher", + func(ctx context.Context, _ struct{}, _ fantasy.ToolCall) (fantasy.ToolResponse, error) { + toolSawPublisher = MessagePartPublisherFromContext(ctx) != nil + return fantasy.NewTextResponse("ok"), nil + }, + ) + + _, err := ExecuteLocalTools(context.Background(), ExecuteLocalToolsOptions{ + Tools: []fantasy.AgentTool{tool}, + ActiveTools: []string{"probe"}, + ToolCalls: []fantasy.ToolCallContent{{ + ToolCallID: "call-1", + ToolName: "probe", + Input: "{}", + }}, + PublishMessagePart: func(codersdk.ChatMessageRole, codersdk.ChatMessagePart) {}, + Clock: quartz.NewReal(), + }) + require.NoError(t, err) + require.True(t, toolSawPublisher) +} diff --git a/coderd/x/chatd/chatopenai/responses.go b/coderd/x/chatd/chatopenai/responses.go index 2c3cad1b09..134ce31590 100644 --- a/coderd/x/chatd/chatopenai/responses.go +++ b/coderd/x/chatd/chatopenai/responses.go @@ -109,45 +109,6 @@ func WithPreviousResponseID( return cloned } -// HasPreviousResponseID checks whether the provider options contain an OpenAI -// Responses entry with a non-empty PreviousResponseID. -func HasPreviousResponseID(providerOptions fantasy.ProviderOptions) bool { - if len(providerOptions) == 0 { - return false - } - - entry, ok := providerOptions[fantasyopenai.Name] - if !ok { - return false - } - options, ok := entry.(*fantasyopenai.ResponsesProviderOptions) - return ok && options != nil && options.PreviousResponseID != nil && - *options.PreviousResponseID != "" -} - -// ClearPreviousResponseID returns a clone of providerOptions with -// PreviousResponseID cleared on the OpenAI Responses options. The original -// providerOptions is not modified. -func ClearPreviousResponseID(providerOptions fantasy.ProviderOptions) fantasy.ProviderOptions { - cloned := maps.Clone(providerOptions) - if cloned == nil { - return fantasy.ProviderOptions{} - } - - entry, ok := cloned[fantasyopenai.Name] - if !ok { - return cloned - } - options, ok := entry.(*fantasyopenai.ResponsesProviderOptions) - if !ok || options == nil { - return cloned - } - optionsClone := *options - optionsClone.PreviousResponseID = nil - cloned[fantasyopenai.Name] = &optionsClone - return cloned -} - // extractResponseID extracts the OpenAI Responses API response ID from provider // metadata. Returns an empty string if no OpenAI Responses metadata is present. func extractResponseID(metadata fantasy.ProviderMetadata) string { diff --git a/coderd/x/chatd/chatopenai/responses_test.go b/coderd/x/chatd/chatopenai/responses_test.go index 5a6e3b9596..59c5cdb44f 100644 --- a/coderd/x/chatd/chatopenai/responses_test.go +++ b/coderd/x/chatd/chatopenai/responses_test.go @@ -254,86 +254,6 @@ func TestWithPreviousResponseIDNilInput(t *testing.T) { require.Empty(t, got) } -func TestHasPreviousResponseID(t *testing.T) { - t.Parallel() - - emptyID := "" - responseID := "resp-123" - - tests := []struct { - name string - opts fantasy.ProviderOptions - want bool - }{ - { - name: "NilOptions", - }, - { - name: "EmptyID", - opts: fantasy.ProviderOptions{ - fantasyopenai.Name: &fantasyopenai.ResponsesProviderOptions{ - PreviousResponseID: &emptyID, - }, - }, - }, - { - name: "NonEmptyID", - opts: fantasy.ProviderOptions{ - fantasyopenai.Name: &fantasyopenai.ResponsesProviderOptions{ - PreviousResponseID: &responseID, - }, - }, - want: true, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - got := chatopenai.HasPreviousResponseID(tt.opts) - require.Equal(t, tt.want, got) - }) - } -} - -func TestClearPreviousResponseID(t *testing.T) { - t.Parallel() - - responseID := "resp-123" - options := &fantasyopenai.ResponsesProviderOptions{ - PreviousResponseID: &responseID, - } - otherOptions := &fantasyopenai.ProviderOptions{} - opts := fantasy.ProviderOptions{ - fantasyopenai.Name: options, - "other": otherOptions, - } - - got := chatopenai.ClearPreviousResponseID(opts) - - got["new"] = otherOptions - require.NotContains(t, opts, "new") - require.NotNil(t, options.PreviousResponseID) - require.Equal(t, "resp-123", *options.PreviousResponseID) - - gotOtherOptions, ok := got["other"].(*fantasyopenai.ProviderOptions) - require.True(t, ok) - require.True(t, otherOptions == gotOtherOptions) - clonedOptions, ok := got[fantasyopenai.Name].(*fantasyopenai.ResponsesProviderOptions) - require.True(t, ok) - require.NotSame(t, options, clonedOptions) - require.Nil(t, clonedOptions.PreviousResponseID) - - require.NotPanics(t, func() { - got := chatopenai.ClearPreviousResponseID(nil) - require.NotNil(t, got) - chatopenai.ClearPreviousResponseID(fantasy.ProviderOptions{ - fantasyopenai.Name: &fantasyopenai.ProviderOptions{}, - }) - }) -} - func TestExtractResponseIDIfStoredMetadata(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatprompt/chatprompt.go b/coderd/x/chatd/chatprompt/chatprompt.go index ac86b11e96..656d8d100d 100644 --- a/coderd/x/chatd/chatprompt/chatprompt.go +++ b/coderd/x/chatd/chatprompt/chatprompt.go @@ -235,39 +235,6 @@ func ConvertMessagesWithFiles( return prompt, nil } -// PrependSystem prepends a system message unless an existing system -// message already mentions create_workspace guidance. -func PrependSystem(prompt []fantasy.Message, instruction string) []fantasy.Message { - instruction = strings.TrimSpace(instruction) - if instruction == "" { - return prompt - } - for _, message := range prompt { - if message.Role != fantasy.MessageRoleSystem { - continue - } - for _, part := range message.Content { - textPart, ok := fantasy.AsMessagePart[fantasy.TextPart](part) - if !ok { - continue - } - if strings.Contains(strings.ToLower(textPart.Text), "create_workspace") { - return prompt - } - } - } - - out := make([]fantasy.Message, 0, len(prompt)+1) - out = append(out, fantasy.Message{ - Role: fantasy.MessageRoleSystem, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: instruction}, - }, - }) - out = append(out, prompt...) - return out -} - // InsertSystem inserts a system message after the existing system // block and before the first non-system message. func InsertSystem(prompt []fantasy.Message, instruction string) []fantasy.Message { @@ -298,24 +265,6 @@ func InsertSystem(prompt []fantasy.Message, instruction string) []fantasy.Messag return out } -// AppendUser appends an instruction as a user message at the end of -// the prompt. -func AppendUser(prompt []fantasy.Message, instruction string) []fantasy.Message { - instruction = strings.TrimSpace(instruction) - if instruction == "" { - return prompt - } - out := make([]fantasy.Message, 0, len(prompt)+1) - out = append(out, prompt...) - out = append(out, fantasy.Message{ - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: instruction}, - }, - }) - return out -} - const ( // ContentVersionV0 is the legacy content format. Parsing uses // role-aware heuristics to distinguish fantasy envelope format diff --git a/coderd/x/chatd/chatprovider/chatprovider.go b/coderd/x/chatd/chatprovider/chatprovider.go index 545fb71a2e..5862138b0c 100644 --- a/coderd/x/chatd/chatprovider/chatprovider.go +++ b/coderd/x/chatd/chatprovider/chatprovider.go @@ -36,11 +36,6 @@ var supportedProviderNames = []string{ fantasyvercel.Name, } -var envPresetProviderNames = []string{ - fantasyopenai.Name, - fantasyanthropic.Name, -} - var providerDisplayNameByName = map[string]string{ fantasyanthropic.Name: "Anthropic", fantasyazure.Name: "Azure OpenAI", @@ -52,22 +47,6 @@ var providerDisplayNameByName = map[string]string{ fantasyvercel.Name: "Vercel AI Gateway", } -// SupportedProviders returns all chat providers supported by Fantasy. -func SupportedProviders() []string { - return append([]string(nil), supportedProviderNames...) -} - -// IsEnvPresetProvider reports whether provider supports env presets. -func IsEnvPresetProvider(provider string) bool { - normalized := NormalizeProvider(provider) - for _, candidate := range envPresetProviderNames { - if candidate == normalized { - return true - } - } - return false -} - // ProviderDisplayName returns a default display name for a provider. func ProviderDisplayName(provider string) string { normalized := NormalizeProvider(provider) @@ -795,343 +774,6 @@ func AnthropicThinkingDisplayFromChat(value *string) *fantasyanthropic.ThinkingD return &valueCopy } -// MergeMissingModelCostConfig fills unset pricing metadata from defaults. -func MergeMissingModelCostConfig( - dst **codersdk.ModelCostConfig, - defaults *codersdk.ModelCostConfig, -) { - if defaults == nil { - return - } - if *dst == nil { - copied := *defaults - *dst = &copied - return - } - - current := *dst - if current.InputPricePerMillionTokens == nil { - current.InputPricePerMillionTokens = defaults.InputPricePerMillionTokens - } - if current.OutputPricePerMillionTokens == nil { - current.OutputPricePerMillionTokens = defaults.OutputPricePerMillionTokens - } - if current.CacheReadPricePerMillionTokens == nil { - current.CacheReadPricePerMillionTokens = defaults.CacheReadPricePerMillionTokens - } - if current.CacheWritePricePerMillionTokens == nil { - current.CacheWritePricePerMillionTokens = defaults.CacheWritePricePerMillionTokens - } -} - -// MergeMissingProviderOptions fills unset provider option fields from defaults. -func MergeMissingProviderOptions( - dst **codersdk.ChatModelProviderOptions, - defaults *codersdk.ChatModelProviderOptions, -) { - if defaults == nil { - return - } - if *dst == nil { - copied := *defaults - *dst = &copied - return - } - - current := *dst - for _, provider := range []string{ - fantasyopenai.Name, - fantasyanthropic.Name, - fantasygoogle.Name, - fantasyopenaicompat.Name, - fantasyopenrouter.Name, - fantasyvercel.Name, - } { - switch provider { - case fantasyopenai.Name: - if defaults.OpenAI == nil { - continue - } - if current.OpenAI == nil { - copied := *defaults.OpenAI - current.OpenAI = &copied - continue - } - dstOpenAI := current.OpenAI - defaultOpenAI := defaults.OpenAI - if dstOpenAI.Include == nil { - dstOpenAI.Include = defaultOpenAI.Include - } - if dstOpenAI.Instructions == nil { - dstOpenAI.Instructions = defaultOpenAI.Instructions - } - if dstOpenAI.LogitBias == nil { - dstOpenAI.LogitBias = defaultOpenAI.LogitBias - } - if dstOpenAI.LogProbs == nil { - dstOpenAI.LogProbs = defaultOpenAI.LogProbs - } - if dstOpenAI.TopLogProbs == nil { - dstOpenAI.TopLogProbs = defaultOpenAI.TopLogProbs - } - if dstOpenAI.MaxToolCalls == nil { - dstOpenAI.MaxToolCalls = defaultOpenAI.MaxToolCalls - } - if dstOpenAI.ParallelToolCalls == nil { - dstOpenAI.ParallelToolCalls = defaultOpenAI.ParallelToolCalls - } - if dstOpenAI.User == nil { - dstOpenAI.User = defaultOpenAI.User - } - if dstOpenAI.ReasoningEffort == nil { - dstOpenAI.ReasoningEffort = defaultOpenAI.ReasoningEffort - } - if dstOpenAI.ReasoningSummary == nil { - dstOpenAI.ReasoningSummary = defaultOpenAI.ReasoningSummary - } - if dstOpenAI.MaxCompletionTokens == nil { - dstOpenAI.MaxCompletionTokens = defaultOpenAI.MaxCompletionTokens - } - if dstOpenAI.TextVerbosity == nil { - dstOpenAI.TextVerbosity = defaultOpenAI.TextVerbosity - } - if dstOpenAI.Prediction == nil { - dstOpenAI.Prediction = defaultOpenAI.Prediction - } - if dstOpenAI.Store == nil { - dstOpenAI.Store = defaultOpenAI.Store - } - if dstOpenAI.Metadata == nil { - dstOpenAI.Metadata = defaultOpenAI.Metadata - } - if dstOpenAI.PromptCacheKey == nil { - dstOpenAI.PromptCacheKey = defaultOpenAI.PromptCacheKey - } - if dstOpenAI.SafetyIdentifier == nil { - dstOpenAI.SafetyIdentifier = defaultOpenAI.SafetyIdentifier - } - if dstOpenAI.ServiceTier == nil { - dstOpenAI.ServiceTier = defaultOpenAI.ServiceTier - } - if dstOpenAI.StructuredOutputs == nil { - dstOpenAI.StructuredOutputs = defaultOpenAI.StructuredOutputs - } - if dstOpenAI.StrictJSONSchema == nil { - dstOpenAI.StrictJSONSchema = defaultOpenAI.StrictJSONSchema - } - - case fantasyanthropic.Name: - if defaults.Anthropic == nil { - continue - } - if current.Anthropic == nil { - copied := *defaults.Anthropic - current.Anthropic = &copied - continue - } - dstAnthropic := current.Anthropic - defaultAnthropic := defaults.Anthropic - if dstAnthropic.SendReasoning == nil { - dstAnthropic.SendReasoning = defaultAnthropic.SendReasoning - } - if dstAnthropic.Thinking == nil { - dstAnthropic.Thinking = defaultAnthropic.Thinking - } else if defaultAnthropic.Thinking != nil && - dstAnthropic.Thinking.BudgetTokens == nil { - dstAnthropic.Thinking.BudgetTokens = defaultAnthropic.Thinking.BudgetTokens - } - if dstAnthropic.Effort == nil { - dstAnthropic.Effort = defaultAnthropic.Effort - } - if dstAnthropic.ThinkingDisplay == nil { - dstAnthropic.ThinkingDisplay = defaultAnthropic.ThinkingDisplay - } - if dstAnthropic.DisableParallelToolUse == nil { - dstAnthropic.DisableParallelToolUse = defaultAnthropic.DisableParallelToolUse - } - - case fantasygoogle.Name: - if defaults.Google == nil { - continue - } - if current.Google == nil { - copied := *defaults.Google - current.Google = &copied - continue - } - dstGoogle := current.Google - defaultGoogle := defaults.Google - if dstGoogle.ThinkingConfig == nil { - dstGoogle.ThinkingConfig = defaultGoogle.ThinkingConfig - } else if defaultGoogle.ThinkingConfig != nil { - if dstGoogle.ThinkingConfig.ThinkingBudget == nil { - dstGoogle.ThinkingConfig.ThinkingBudget = defaultGoogle.ThinkingConfig.ThinkingBudget - } - if dstGoogle.ThinkingConfig.IncludeThoughts == nil { - dstGoogle.ThinkingConfig.IncludeThoughts = defaultGoogle.ThinkingConfig.IncludeThoughts - } - } - if strings.TrimSpace(dstGoogle.CachedContent) == "" { - dstGoogle.CachedContent = defaultGoogle.CachedContent - } - if dstGoogle.SafetySettings == nil { - dstGoogle.SafetySettings = defaultGoogle.SafetySettings - } - if strings.TrimSpace(dstGoogle.Threshold) == "" { - dstGoogle.Threshold = defaultGoogle.Threshold - } - - case fantasyopenaicompat.Name: - if defaults.OpenAICompat == nil { - continue - } - if current.OpenAICompat == nil { - copied := *defaults.OpenAICompat - current.OpenAICompat = &copied - continue - } - dstCompat := current.OpenAICompat - defaultCompat := defaults.OpenAICompat - if dstCompat.User == nil { - dstCompat.User = defaultCompat.User - } - if dstCompat.ReasoningEffort == nil { - dstCompat.ReasoningEffort = defaultCompat.ReasoningEffort - } - - case fantasyopenrouter.Name: - if defaults.OpenRouter == nil { - continue - } - if current.OpenRouter == nil { - copied := *defaults.OpenRouter - current.OpenRouter = &copied - continue - } - dstRouter := current.OpenRouter - defaultRouter := defaults.OpenRouter - if dstRouter.Reasoning == nil { - dstRouter.Reasoning = defaultRouter.Reasoning - } else if defaultRouter.Reasoning != nil { - if dstRouter.Reasoning.Enabled == nil { - dstRouter.Reasoning.Enabled = defaultRouter.Reasoning.Enabled - } - if dstRouter.Reasoning.Exclude == nil { - dstRouter.Reasoning.Exclude = defaultRouter.Reasoning.Exclude - } - if dstRouter.Reasoning.MaxTokens == nil { - dstRouter.Reasoning.MaxTokens = defaultRouter.Reasoning.MaxTokens - } - if dstRouter.Reasoning.Effort == nil { - dstRouter.Reasoning.Effort = defaultRouter.Reasoning.Effort - } - } - if dstRouter.ExtraBody == nil { - dstRouter.ExtraBody = defaultRouter.ExtraBody - } - if dstRouter.IncludeUsage == nil { - dstRouter.IncludeUsage = defaultRouter.IncludeUsage - } - if dstRouter.LogitBias == nil { - dstRouter.LogitBias = defaultRouter.LogitBias - } - if dstRouter.LogProbs == nil { - dstRouter.LogProbs = defaultRouter.LogProbs - } - if dstRouter.ParallelToolCalls == nil { - dstRouter.ParallelToolCalls = defaultRouter.ParallelToolCalls - } - if dstRouter.User == nil { - dstRouter.User = defaultRouter.User - } - if dstRouter.Provider == nil { - dstRouter.Provider = defaultRouter.Provider - } else if defaultRouter.Provider != nil { - if dstRouter.Provider.Order == nil { - dstRouter.Provider.Order = defaultRouter.Provider.Order - } - if dstRouter.Provider.AllowFallbacks == nil { - dstRouter.Provider.AllowFallbacks = defaultRouter.Provider.AllowFallbacks - } - if dstRouter.Provider.RequireParameters == nil { - dstRouter.Provider.RequireParameters = defaultRouter.Provider.RequireParameters - } - if dstRouter.Provider.DataCollection == nil { - dstRouter.Provider.DataCollection = defaultRouter.Provider.DataCollection - } - if dstRouter.Provider.Only == nil { - dstRouter.Provider.Only = defaultRouter.Provider.Only - } - if dstRouter.Provider.Ignore == nil { - dstRouter.Provider.Ignore = defaultRouter.Provider.Ignore - } - if dstRouter.Provider.Quantizations == nil { - dstRouter.Provider.Quantizations = defaultRouter.Provider.Quantizations - } - if dstRouter.Provider.Sort == nil { - dstRouter.Provider.Sort = defaultRouter.Provider.Sort - } - } - - case fantasyvercel.Name: - if defaults.Vercel == nil { - continue - } - if current.Vercel == nil { - copied := *defaults.Vercel - current.Vercel = &copied - continue - } - dstVercel := current.Vercel - defaultVercel := defaults.Vercel - if dstVercel.Reasoning == nil { - dstVercel.Reasoning = defaultVercel.Reasoning - } else if defaultVercel.Reasoning != nil { - if dstVercel.Reasoning.Enabled == nil { - dstVercel.Reasoning.Enabled = defaultVercel.Reasoning.Enabled - } - if dstVercel.Reasoning.MaxTokens == nil { - dstVercel.Reasoning.MaxTokens = defaultVercel.Reasoning.MaxTokens - } - if dstVercel.Reasoning.Effort == nil { - dstVercel.Reasoning.Effort = defaultVercel.Reasoning.Effort - } - if dstVercel.Reasoning.Exclude == nil { - dstVercel.Reasoning.Exclude = defaultVercel.Reasoning.Exclude - } - } - if dstVercel.ProviderOptions == nil { - dstVercel.ProviderOptions = defaultVercel.ProviderOptions - } else if defaultVercel.ProviderOptions != nil { - if dstVercel.ProviderOptions.Order == nil { - dstVercel.ProviderOptions.Order = defaultVercel.ProviderOptions.Order - } - if dstVercel.ProviderOptions.Models == nil { - dstVercel.ProviderOptions.Models = defaultVercel.ProviderOptions.Models - } - } - if dstVercel.User == nil { - dstVercel.User = defaultVercel.User - } - if dstVercel.LogitBias == nil { - dstVercel.LogitBias = defaultVercel.LogitBias - } - if dstVercel.LogProbs == nil { - dstVercel.LogProbs = defaultVercel.LogProbs - } - if dstVercel.TopLogProbs == nil { - dstVercel.TopLogProbs = defaultVercel.TopLogProbs - } - if dstVercel.ParallelToolCalls == nil { - dstVercel.ParallelToolCalls = defaultVercel.ParallelToolCalls - } - if dstVercel.ExtraBody == nil { - dstVercel.ExtraBody = defaultVercel.ExtraBody - } - } - } -} - // Header constants sent on upstream LLM API requests so that // intermediaries (e.g. aibridged) can correlate traffic back to // Coder entities. @@ -1170,22 +812,6 @@ func CoderHeaders(chat database.Chat) map[string]string { return h } -// CoderHeadersFromIDs is a convenience form of CoderHeaders for call -// sites that do not have a full database.Chat in scope. -func CoderHeadersFromIDs( - ownerID uuid.UUID, - chatID uuid.UUID, - parentChatID uuid.NullUUID, - workspaceID uuid.NullUUID, -) map[string]string { - return CoderHeaders(database.Chat{ - ID: chatID, - OwnerID: ownerID, - ParentChatID: parentChatID, - WorkspaceID: workspaceID, - }) -} - // ModelFromConfig resolves a provider/model pair and constructs a fantasy // language model client using the provided provider credentials. The // userAgent is sent as the User-Agent header on every outgoing LLM diff --git a/coderd/x/chatd/chatprovider/chatprovider_test.go b/coderd/x/chatd/chatprovider/chatprovider_test.go index 261c6e9528..fea015db73 100644 --- a/coderd/x/chatd/chatprovider/chatprovider_test.go +++ b/coderd/x/chatd/chatprovider/chatprovider_test.go @@ -424,24 +424,6 @@ func TestProviderOptionsFromChatModelConfig_AnthropicThinkingDisplay(t *testing. require.Equal(t, fantasyanthropic.ThinkingDisplaySummarized, *anthropicOptions.ThinkingDisplay) } -func TestMergeMissingProviderOptions_AnthropicThinkingDisplay(t *testing.T) { - t.Parallel() - - options := &codersdk.ChatModelProviderOptions{ - Anthropic: &codersdk.ChatModelAnthropicProviderOptions{}, - } - defaults := &codersdk.ChatModelProviderOptions{ - Anthropic: &codersdk.ChatModelAnthropicProviderOptions{ - ThinkingDisplay: ptr.Ref("summarized"), - }, - } - - chatprovider.MergeMissingProviderOptions(&options, defaults) - - require.NotNil(t, options.Anthropic.ThinkingDisplay) - require.Equal(t, "summarized", *options.Anthropic.ThinkingDisplay) -} - func TestResolveUserProviderKeys_UnavailableReason(t *testing.T) { t.Parallel() @@ -1559,66 +1541,6 @@ func TestModelFromConfig_HTTPClient(t *testing.T) { _ = testutil.TryReceive(ctx, t, called) } -func TestMergeMissingProviderOptions_OpenRouterNested(t *testing.T) { - t.Parallel() - - options := &codersdk.ChatModelProviderOptions{ - OpenRouter: &codersdk.ChatModelOpenRouterProviderOptions{ - Reasoning: &codersdk.ChatModelReasoningOptions{ - Enabled: ptr.Ref(true), - }, - Provider: &codersdk.ChatModelOpenRouterProvider{ - Order: []string{"openai"}, - }, - }, - } - defaults := &codersdk.ChatModelProviderOptions{ - OpenRouter: &codersdk.ChatModelOpenRouterProviderOptions{ - Reasoning: &codersdk.ChatModelReasoningOptions{ - Enabled: ptr.Ref(false), - Exclude: ptr.Ref(true), - MaxTokens: ptr.Ref[int64](123), - Effort: ptr.Ref("high"), - }, - IncludeUsage: ptr.Ref(true), - Provider: &codersdk.ChatModelOpenRouterProvider{ - Order: []string{"anthropic"}, - AllowFallbacks: ptr.Ref(true), - RequireParameters: ptr.Ref(false), - DataCollection: ptr.Ref("allow"), - Only: []string{"openai"}, - Ignore: []string{"foo"}, - Quantizations: []string{"int8"}, - Sort: ptr.Ref("latency"), - }, - }, - } - - chatprovider.MergeMissingProviderOptions(&options, defaults) - - require.NotNil(t, options) - require.NotNil(t, options.OpenRouter) - require.NotNil(t, options.OpenRouter.Reasoning) - require.True(t, *options.OpenRouter.Reasoning.Enabled) - require.Equal(t, true, *options.OpenRouter.Reasoning.Exclude) - require.EqualValues(t, 123, *options.OpenRouter.Reasoning.MaxTokens) - require.Equal(t, "high", *options.OpenRouter.Reasoning.Effort) - require.NotNil(t, options.OpenRouter.IncludeUsage) - require.True(t, *options.OpenRouter.IncludeUsage) - - require.NotNil(t, options.OpenRouter.Provider) - require.Equal(t, []string{"openai"}, options.OpenRouter.Provider.Order) - require.NotNil(t, options.OpenRouter.Provider.AllowFallbacks) - require.True(t, *options.OpenRouter.Provider.AllowFallbacks) - require.NotNil(t, options.OpenRouter.Provider.RequireParameters) - require.False(t, *options.OpenRouter.Provider.RequireParameters) - require.Equal(t, "allow", *options.OpenRouter.Provider.DataCollection) - require.Equal(t, []string{"openai"}, options.OpenRouter.Provider.Only) - require.Equal(t, []string{"foo"}, options.OpenRouter.Provider.Ignore) - require.Equal(t, []string{"int8"}, options.OpenRouter.Provider.Quantizations) - require.Equal(t, "latency", *options.OpenRouter.Provider.Sort) -} - func TestResolveModelWithProviderHint(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatretry/chatretry_test.go b/coderd/x/chatd/chatretry/chatretry_test.go index 61fdb047bb..b750548d0a 100644 --- a/coderd/x/chatd/chatretry/chatretry_test.go +++ b/coderd/x/chatd/chatretry/chatretry_test.go @@ -4,7 +4,6 @@ import ( "context" "errors" "fmt" - "io" "sync/atomic" "testing" "time" @@ -18,80 +17,6 @@ import ( "github.com/coder/coder/v2/codersdk" ) -func TestIsRetryableDelegatesToClassification(t *testing.T) { - t.Parallel() - - tests := []struct { - name string - err error - retryable bool - }{ - {name: "Nil", err: nil, retryable: false}, - {name: "RetryableExplicitStatus429", err: xerrors.New("received status 429 from upstream"), retryable: true}, - {name: "RetryableTimeout", err: xerrors.New("service unavailable"), retryable: true}, - { - name: "RetryableAnthropicMissingMessageStop", - err: xerrors.Errorf( - "anthropic stream closed before message_stop: %w", - io.EOF, - ), - retryable: true, - }, - { - name: "RetryableOpenAIResponsesMissingTerminalEvent", - err: xerrors.Errorf( - "openai responses stream closed before terminal event: %w", - io.EOF, - ), - retryable: true, - }, - {name: "NonRetryableAuth", err: xerrors.New("invalid api key"), retryable: false}, - {name: "NonRetryableGeneric", err: xerrors.New("boom"), retryable: false}, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - t.Parallel() - - require.Equal(t, tt.retryable, chatretry.IsRetryable(tt.err)) - require.Equal(t, chaterror.Classify(tt.err).Retryable, chatretry.IsRetryable(tt.err)) - }) - } -} - -func TestRetryabilityFromClassifyStatusCodes(t *testing.T) { - t.Parallel() - - tests := []struct { - code int - retryable bool - }{ - {408, true}, - {429, true}, - {500, true}, - {502, true}, - {503, true}, - {504, true}, - {529, true}, - {200, false}, - {400, false}, - {401, false}, - {403, false}, - {404, false}, - } - - for _, tt := range tests { - t.Run(fmt.Sprintf("Status%d", tt.code), func(t *testing.T) { - t.Parallel() - - err := xerrors.Errorf("status %d from upstream", tt.code) - classified := chaterror.Classify(err) - require.Equal(t, tt.retryable, classified.Retryable) - require.Equal(t, classified.Retryable, chatretry.IsRetryable(err)) - }) - } -} - func TestDelay(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatsanitize/anthropic.go b/coderd/x/chatd/chatsanitize/anthropic.go index f3605ed091..b11be26098 100644 --- a/coderd/x/chatd/chatsanitize/anthropic.go +++ b/coderd/x/chatd/chatsanitize/anthropic.go @@ -466,15 +466,6 @@ func contentHasAnthropicSignedReasoning(content []fantasy.Content) bool { return false } -// IsAnthropicProviderExecutedToolCall reports whether toolCall is an -// Anthropic provider-executed tool call. -func IsAnthropicProviderExecutedToolCall( - provider string, - toolCall fantasy.ToolCallContent, -) bool { - return provider == fantasyanthropic.Name && toolCall.ProviderExecuted -} - // ApplyAnthropicProviderToolGuard fail-closes unsafe Anthropic provider-tool // history immediately before a provider request is issued. It returns a // sanitized prompt on success, or nil with ErrAnthropicProviderToolPromptUnsafe diff --git a/coderd/x/chatd/chatstate/concurrency_test.go b/coderd/x/chatd/chatstate/concurrency_test.go new file mode 100644 index 0000000000..881cb6431f --- /dev/null +++ b/coderd/x/chatd/chatstate/concurrency_test.go @@ -0,0 +1,219 @@ +package chatstate_test + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/testutil" +) + +// waitForChan returns true if c receives a value before ctx is done. +// Helper used in concurrency tests to avoid time.Sleep. +func waitForChan(ctx context.Context, c <-chan struct{}) bool { + select { + case <-c: + return true + case <-ctx.Done(): + return false + } +} + +// stillBlocked returns true if c has not received a value and has not +// been closed. The caller must already have established a happens-before +// ordering via another channel so this check is meaningful. +func stillBlocked(c <-chan struct{}) bool { + select { + case <-c: + return false + default: + return true + } +} + +// waitForWaitGroup returns true if wg completes before ctx is done. +func waitForWaitGroup(ctx context.Context, wg *sync.WaitGroup) bool { + done := make(chan struct{}) + go func() { + wg.Wait() + close(done) + }() + return waitForChan(ctx, done) +} + +type lockAttemptStore struct { + database.Store + + attempted chan struct{} + once *sync.Once +} + +func newLockAttemptStore(store database.Store, attempted chan struct{}) *lockAttemptStore { + return &lockAttemptStore{ + Store: store, + attempted: attempted, + once: new(sync.Once), + } +} + +func (s *lockAttemptStore) InTx(fn func(database.Store) error, opts *database.TxOptions) error { + return s.Store.InTx(func(tx database.Store) error { + return fn(&lockAttemptStore{ + Store: tx, + attempted: s.attempted, + once: s.once, + }) + }, opts) +} + +func (s *lockAttemptStore) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid.UUID) (database.Chat, error) { + s.once.Do(func() { close(s.attempted) }) + return s.Store.LockChatAndBumpSnapshotVersion(ctx, id) +} + +// TestLockLocksChatRow verifies that ChatMachine.Lock holds the chat +// row's FOR UPDATE lock until the callback returns, so a concurrent +// ChatMachine.Update cannot enter its callback until the Lock +// callback releases. +func TestLockLocksChatRow(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitMedium) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + updateLockAttempted := make(chan struct{}) + updateMachine := chatstate.NewChatMachine( + newLockAttemptStore(f.DB, updateLockAttempted), + f.Pub, + created.Chat.ID, + ) + + lockEntered := make(chan struct{}) + releaseLock := make(chan struct{}) + t.Cleanup(func() { + select { + case <-releaseLock: + default: + close(releaseLock) + } + }) + updateEntered := make(chan struct{}) + + // Goroutine A: hold a Lock and block. + var lockErr error + var lockWG sync.WaitGroup + lockWG.Go(func() { + lockErr = m.Lock(ctx, func(_ database.Store) error { + close(lockEntered) + select { + case <-releaseLock: + return nil + case <-ctx.Done(): + return ctx.Err() + } + }) + }) + + // Wait until A is inside its Lock callback (and therefore holds + // the FOR UPDATE lock). + require.True(t, waitForChan(ctx, lockEntered), "Lock callback never started") + + // Goroutine B: try to Update the same chat. It must block on + // LockChatAndBumpSnapshotVersion until A releases. + var updateErr error + var updateWG sync.WaitGroup + updateWG.Go(func() { + updateErr = updateMachine.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { + close(updateEntered) + return nil + }) + }) + + require.True(t, waitForChan(ctx, updateLockAttempted), + "Update never attempted to lock the chat row") + // Sleep to give a chance for the update to enter the callback. + // This isn't a deterministic solution - on a low resource, contended system + // it's possible that the Update won't call the callback even if the lock + // implementation is incorrect and doesn't block. But in most cases, this wait should be enough. + time.Sleep(50 * time.Millisecond) + require.True(t, stillBlocked(updateEntered), + "Update entered while Lock was still held") + + // Release Lock and confirm Update completes successfully. + close(releaseLock) + require.True(t, waitForChan(ctx, updateEntered), + "Update callback never started after Lock released") + require.True(t, waitForWaitGroup(ctx, &updateWG), "Update did not finish") + require.True(t, waitForWaitGroup(ctx, &lockWG), "Lock did not finish") + require.NoError(t, lockErr) + require.NoError(t, updateErr) +} + +// TestLockRollsBackCallbackError verifies that a Lock callback +// returning an error rolls back the surrounding transaction. +func TestLockRollsBackCallbackError(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + before := f.readChat(ctx, t, created.Chat.ID) + publishedBefore := len(f.Pub.channels) + + sentinel := xerrors.New("lock callback error") + err := m.Lock(ctx, func(store database.Store) error { + // Try a write that should be rolled back. + _, werr := store.UpdateChatByID(ctx, database.UpdateChatByIDParams{ + ID: created.Chat.ID, + Title: "rollback-me", + }) + require.NoError(t, werr) + return sentinel + }) + require.ErrorIs(t, err, sentinel) + + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.Title, after.Title, "Lock callback error rolls back writes") + require.Equal(t, publishedBefore, len(f.Pub.channels), "Lock publishes nothing on error") +} + +// TestConcurrentUpdatesSerializeOnChatRow verifies that two +// goroutines racing to Update the same chat both succeed but their +// effects serialize on the chat row lock: snapshot_version advances +// by exactly N (one per Update) and each transition observes the +// effects of the prior one. +func TestConcurrentUpdatesSerializeOnChatRow(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitMedium) + created := createTestChat(t, f) + before := f.readChat(ctx, t, created.Chat.ID) + + const updates = 8 + var wg sync.WaitGroup + wg.Add(updates) + errs := make([]error, updates) + for i := 0; i < updates; i++ { + i := i + go func() { + defer wg.Done() + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + errs[i] = m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil }) + }() + } + wg.Wait() + for i, err := range errs { + require.NoError(t, err, "concurrent update %d failed", i) + } + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.SnapshotVersion+int64(updates), after.SnapshotVersion, + "snapshot_version advanced by exactly one per update") +} diff --git a/coderd/x/chatd/chatstate/doc.go b/coderd/x/chatd/chatstate/doc.go new file mode 100644 index 0000000000..16cb4e548b --- /dev/null +++ b/coderd/x/chatd/chatstate/doc.go @@ -0,0 +1,31 @@ +// Package chatstate owns the durable execution-state transitions for +// the chatd subsystem. It implements the chat execution state model. +// +// The package exposes two top-level entry points: +// +// - [CreateChat] creates a brand new chat with its initial history +// in a single transaction. It is standalone because no chat-scoped +// state machine instance can exist before the chat row is written. +// - [ChatMachine] wraps an existing chat. Callers use it to apply +// one or more transitions atomically via [ChatMachine.Update], or +// to read related rows while holding the chat row lock via +// [ChatMachine.Lock]. +// +// Every successful [ChatMachine.Update] call locks the chat row, +// advances `snapshot_version` exactly once, applies transition methods +// in order, and (on commit) publishes a single typed `chat:update` +// pubsub message describing the post-transition snapshot. Optional +// `chat:ownership` hints are published only when the post-transition +// state is runnable and ownership is missing or stale. Stream side +// effects are handled by `chat:update` consumers, and ownership hints +// wake chat workers. +// +// Transition methods are explicit, typed wrappers around the durable +// mutations needed to move between states. Each transition reads the +// current chat row and queue cardinality, classifies the resulting +// execution state, validates it against the transition model, and +// rejects with an [*TransitionError] wrapping [ErrTransitionNotAllowed] +// when the transition is not legal from that state. The package owns +// transition validation, durable chat row and queue mutations, and +// post-commit pubsub publication. +package chatstate diff --git a/coderd/x/chatd/chatstate/errors.go b/coderd/x/chatd/chatstate/errors.go new file mode 100644 index 0000000000..757ec9f81b --- /dev/null +++ b/coderd/x/chatd/chatstate/errors.go @@ -0,0 +1,152 @@ +package chatstate + +import ( + "errors" + "fmt" + + "golang.org/x/xerrors" +) + +// Sentinel errors returned by chatstate transitions and helpers. +// Callers should use errors.Is to test for these. +var ( + // ErrTransitionNotAllowed is returned when a transition is applied + // to a chat whose current execution state does not permit it. The + // concrete error returned by transition methods is a + // *TransitionError that wraps this sentinel. + ErrTransitionNotAllowed = xerrors.New("chat state transition not allowed") + + // ErrInvalidState is returned when the chat row, queue, and + // archive flag together produce a combination outside the chat + // execution state model. + ErrInvalidState = xerrors.New("chat is in an invalid execution state") + + // ErrQueuedMessageNotFound is returned by queue-targeting + // transitions (delete, promote) when the supplied queued message + // ID does not match a row on the chat. + ErrQueuedMessageNotFound = xerrors.New("queued message not found") + + // ErrMessageNotFound is returned by [Tx.EditMessage] when the + // target chat_messages row is missing or belongs to another chat. + ErrMessageNotFound = xerrors.New("chat message not found") + + // ErrChatNotFound is returned when a non-create transition is + // applied to a chat row that does not exist (or has been deleted + // since the transition started). + ErrChatNotFound = xerrors.New("chat not found") + + // ErrChatNotRoot is returned by family-archive helpers when the + // supplied chat is not a root chat (its parent_chat_id is set). + ErrChatNotRoot = xerrors.New("chat is not a root chat") + + // ErrEditedMessageNotUser is returned by [Tx.EditMessage] when the + // targeted chat_messages row exists but its role is not user. + ErrEditedMessageNotUser = xerrors.New("only user messages can be edited") + + // ErrMessageQueueFull is returned by queue-appending transitions + // when the per-chat queue cap has been reached. The concrete + // error returned by transitions is a *MessageQueueFullError that + // wraps this sentinel. + ErrMessageQueueFull = xerrors.New("chat message queue is full") + + // ErrToolResultDuplicate is returned by [Tx.CompleteRequiresAction] + // when the same tool_call_id appears more than once in the + // submitted results. + ErrToolResultDuplicate = xerrors.New("duplicate tool result") + + // ErrToolResultUnexpected is returned by + // [Tx.CompleteRequiresAction] when a submitted tool_call_id does + // not correspond to a pending dynamic tool call. + ErrToolResultUnexpected = xerrors.New("unexpected tool result") + + // ErrToolResultMissing is returned by [Tx.CompleteRequiresAction] + // when a pending dynamic tool call has no submitted result. + ErrToolResultMissing = xerrors.New("missing tool result") + + // ErrToolResultInvalidJSON is returned by + // [Tx.CompleteRequiresAction] when a submitted tool result output + // is not valid JSON. + ErrToolResultInvalidJSON = xerrors.New("tool result output is not valid JSON") +) + +// MessageQueueFullError carries the per-chat queue cap so HTTP +// endpoints can include the cap in their response detail. It wraps +// [ErrMessageQueueFull] so callers can match it with errors.Is. +type MessageQueueFullError struct { + Max int64 +} + +// Error implements the error interface. +func (e *MessageQueueFullError) Error() string { + return fmt.Sprintf("chat message queue is full (max %d)", e.Max) +} + +// Unwrap returns [ErrMessageQueueFull] so callers can match the +// generic sentinel. +func (*MessageQueueFullError) Unwrap() error { return ErrMessageQueueFull } + +// ToolResultValidationError carries a structured tool-result +// validation failure. It always wraps a specific sentinel +// (ErrToolResultDuplicate, ErrToolResultMissing, +// ErrToolResultUnexpected, ErrToolResultInvalidJSON) so callers can +// match either the generic sentinel or the specific cause. +type ToolResultValidationError struct { + Cause error + ToolCallID string +} + +// Error implements the error interface. +func (e *ToolResultValidationError) Error() string { + if e.ToolCallID != "" { + return fmt.Sprintf("%s: %s", e.Cause.Error(), e.ToolCallID) + } + return e.Cause.Error() +} + +// Unwrap returns the specific cause so callers can match it. +func (e *ToolResultValidationError) Unwrap() error { return e.Cause } + +// TransitionError carries the structured detail for a rejected +// transition. It always wraps [ErrTransitionNotAllowed] so callers can +// match with errors.Is without losing context. When a specific +// chatstate sentinel is the proximate cause, Cause is set and +// errors.Is will match that sentinel too. +type TransitionError struct { + Transition Transition + From ExecutionState + Reason string + Cause error +} + +// Error implements the error interface. +func (e *TransitionError) Error() string { + if e.Reason == "" { + return fmt.Sprintf( + "chat state transition %s not allowed from state %s", + e.Transition, e.From, + ) + } + return fmt.Sprintf( + "chat state transition %s not allowed from state %s: %s", + e.Transition, e.From, e.Reason, + ) +} + +// Unwrap returns the error chain attached to this error. The chain +// always includes [ErrTransitionNotAllowed], and may include a more +// specific cause through errors.Join, so callers can use errors.Is +// without custom matching logic on TransitionError. +func (e *TransitionError) Unwrap() error { return e.Cause } + +// newTransitionError constructs a typed TransitionError. Returning the +// pointer type lets callers inspect the structured fields when needed. +func newTransitionError(t Transition, from ExecutionState, reason string) *TransitionError { + return &TransitionError{Transition: t, From: from, Reason: reason, Cause: ErrTransitionNotAllowed} +} + +// newTransitionErrorWithCause constructs a TransitionError carrying +// a specific underlying sentinel so callers can match the cause with +// errors.Is. +func newTransitionErrorWithCause(t Transition, from ExecutionState, cause error, reason string) *TransitionError { + return &TransitionError{Transition: t, From: from, Reason: reason, Cause: errors.Join(ErrTransitionNotAllowed, cause)} +} diff --git a/coderd/x/chatd/chatstate/family.go b/coderd/x/chatd/chatstate/family.go new file mode 100644 index 0000000000..fb22e05bae --- /dev/null +++ b/coderd/x/chatd/chatstate/family.go @@ -0,0 +1,130 @@ +package chatstate + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" +) + +// SetFamilyArchivedInput configures [SetFamilyArchived]. The struct +// shape avoids a boolean flag parameter at the API surface; callers +// build it explicitly with named fields for clarity. +type SetFamilyArchivedInput struct { + // RootID identifies the family root. SetFamilyArchived rejects + // calls for child chats with [ErrChatNotRoot] and unknown chats + // with [ErrChatNotFound]. + RootID uuid.UUID + // Archived is the desired post-call archived value for every + // family member. + Archived bool +} + +// SetFamilyArchived runs Update for every chat in the root chat's +// family inside one transaction, applying SetArchived when the chat's +// archived flag differs from the requested value. It owns its +// transaction lifecycle and its [PublishBuffer] lifecycle: pubsub +// publications are buffered while the transaction is open and +// flushed only after a successful commit; the deferred Discard +// suppresses every buffered publication on failure. +// +// On success SetFamilyArchived returns one [database.Chat] per +// family member in the order returned by GetChatFamilyIDsByRootID +// (root first, then children). +// +// Family members that are already in the [StateInvalid] execution +// state cause SetFamilyArchived to return [ErrInvalidState] and roll +// back the cascade even when their archived flag already matches the +// desired value; invalid-state detection is never bypassed. +// +// Family members that are valid and already match the desired +// archived value still run through Update, which increments their +// snapshot version and publishes a fresh snapshot without changing +// the archived flag. Advancing the snapshot version without a field +// change is safe, and it keeps publication behavior uniform while a +// partially archived family converges to the desired state. +func SetFamilyArchived( + ctx context.Context, + store database.Store, + publisher Publisher, + input SetFamilyArchivedInput, +) ([]database.Chat, error) { + if store == nil { + return nil, xerrors.New("chatstate: SetFamilyArchived called with nil store") + } + if publisher == nil { + return nil, xerrors.New("chatstate: SetFamilyArchived called with nil publisher") + } + + buffer := NewPublishBuffer(publisher) + defer buffer.Discard() + + var familyChats []database.Chat + err := store.InTx(func(tx database.Store) error { + // Lock the root chat first so concurrent archive races on the + // same family serialize on a stable row. + root, err := tx.GetChatByIDForUpdate(ctx, input.RootID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrChatNotFound + } + return xerrors.Errorf("lock root chat for archive: %w", err) + } + if root.ParentChatID.Valid { + return ErrChatNotRoot + } + ids, err := tx.GetChatFamilyIDsByRootID(ctx, input.RootID) + if err != nil { + return xerrors.Errorf("get chat family: %w", err) + } + if len(ids) == 0 { + return ErrChatNotFound + } + familyChats = make([]database.Chat, 0, len(ids)) + for _, id := range ids { + var chat database.Chat + machine := NewChatMachine(tx, buffer, id) + err := machine.Update(ctx, func(state *Tx, _ database.Store) error { + // Classify each member so any invalid execution state + // aborts and rolls back the whole family update, even + // when that member already has the requested archived + // value. + current, from, err := state.loadState() + if err != nil { + return err + } + if from == StateInvalid { + return ErrInvalidState + } + if current.Archived == input.Archived { + chat = current + return nil + } + if _, err := state.SetArchived(SetArchivedInput{Archived: input.Archived}); err != nil { + return err + } + chat, err = state.Store().GetChatByID(state.Ctx(), state.ChatID()) + if err != nil { + return xerrors.Errorf("reload archived chat: %w", err) + } + return nil + }) + if err != nil { + return err + } + familyChats = append(familyChats, chat) + } + return nil + }, nil) + if err != nil { + return nil, err + } + if err := buffer.Flush(); err != nil { + return familyChats, err + } + return familyChats, nil +} diff --git a/coderd/x/chatd/chatstate/family_test.go b/coderd/x/chatd/chatstate/family_test.go new file mode 100644 index 0000000000..b7781d83b6 --- /dev/null +++ b/coderd/x/chatd/chatstate/family_test.go @@ -0,0 +1,218 @@ +package chatstate_test + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// TestSetFamilyArchivedRejectsChildChat asserts the chatstate helper +// rejects calls that target a child chat. Family archive flows must +// always start at the root. +func TestSetFamilyArchivedRejectsChildChat(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + + root := dbgen.Chat(t, f.DB, database.Chat{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "root", + }) + child := dbgen.Chat(t, f.DB, database.Chat{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "child", + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + _, err := chatstate.SetFamilyArchived(ctx, f.DB, f.Pub, chatstate.SetFamilyArchivedInput{RootID: child.ID, Archived: true}) + require.ErrorIs(t, err, chatstate.ErrChatNotRoot) + + require.False(t, f.readChat(ctx, t, root.ID).Archived, + "failed family archive must not touch the root") + require.False(t, f.readChat(ctx, t, child.ID).Archived, + "failed family archive must not touch the child") +} + +// TestSetFamilyArchivedRollsBackWhenMemberCannotArchive verifies that +// SetFamilyArchived is atomic: when one family member is in a state +// that cannot satisfy the SetArchived transition, the whole cascade +// rolls back and no publications reach the inner publisher. +func TestSetFamilyArchivedRollsBackWhenMemberCannotArchive(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user, org, model := seedFamilyDeps(t, db) + + // Root chat: waiting is archive-eligible (state W). + root := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "root", + Status: database.ChatStatusWaiting, + }) + // Child chat: running with no queue is R0 and NOT archive + // eligible per the chatstate transition matrix. + child := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "child", + Status: database.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + pub := newRecordingPubsub() + _, err := chatstate.SetFamilyArchived(ctx, db, pub, chatstate.SetFamilyArchivedInput{RootID: root.ID, Archived: true}) + require.Error(t, err, "child in "+chatstate.StateR0.String()+" must reject SetArchived") + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + + rootAfter, err := db.GetChatByID(ctx, root.ID) + require.NoError(t, err) + require.False(t, rootAfter.Archived, "root archive must roll back when a child cannot archive") + childAfter, err := db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.False(t, childAfter.Archived, "child must not be archived in the rolled-back cascade") + + require.Empty(t, pub.channels, + "rolled-back family archive must publish nothing through the inner publisher") +} + +// TestSetFamilyArchivedRejectsInvalidStateEvenWhenAlreadyDesired +// verifies that invalid-state detection is never bypassed: a family +// member in StateInvalid causes the cascade to fail with +// ErrInvalidState even when that member's archived flag already +// matches the desired value. +func TestSetFamilyArchivedRejectsInvalidStateEvenWhenAlreadyDesired(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user, org, model := seedFamilyDeps(t, db) + + root := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "root", + Status: database.ChatStatusWaiting, + }) + child := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "child", + // status=waiting, archived=true; we will add a queued message + // to produce the chatstate-invalid combination (archived chat + // with a queued backlog is outside the valid state model). + Status: database.ChatStatusWaiting, + Archived: true, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + // Seed a queued message under the child to push it into the + // chatstate-invalid combination. + rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("queued"), + }) + require.NoError(t, err) + _, err = db.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ + ChatID: child.ID, + Content: rawContent.RawMessage, + ModelConfigID: uuid.NullUUID{}, + }) + require.NoError(t, err) + + pub := newRecordingPubsub() + _, err = chatstate.SetFamilyArchived(ctx, db, pub, chatstate.SetFamilyArchivedInput{ + RootID: root.ID, + Archived: true, + }) + require.ErrorIs(t, err, chatstate.ErrInvalidState, + "invalid-state child blocks the cascade even when archived flag already matches") + + // Root must not be archived because the cascade rolled back. + rootAfter, err := db.GetChatByID(ctx, root.ID) + require.NoError(t, err) + require.False(t, rootAfter.Archived, "root must roll back when a child is in StateInvalid") + + require.Empty(t, pub.channels, + "rolled-back cascade must not publish anything") +} + +// TestSetFamilyArchivedAcceptsAlreadyDesiredMembers verifies that an +// individually archived child does not block a root archive cascade. +// The cascade converges to the desired state even when some family +// members already match it. +func TestSetFamilyArchivedAcceptsAlreadyDesiredMembers(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + user, org, model := seedFamilyDeps(t, db) + + root := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "root", + Status: database.ChatStatusWaiting, + }) + child := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "child", + Status: database.ChatStatusWaiting, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + Archived: true, + }) + + pub := newRecordingPubsub() + family, err := chatstate.SetFamilyArchived(ctx, db, pub, chatstate.SetFamilyArchivedInput{RootID: root.ID, Archived: true}) + require.NoError(t, err, + "already archived members must not block the cascade") + require.Len(t, family, 2) + + rootAfter, err := db.GetChatByID(ctx, root.ID) + require.NoError(t, err) + require.True(t, rootAfter.Archived) + childAfter, err := db.GetChatByID(ctx, child.ID) + require.NoError(t, err) + require.True(t, childAfter.Archived) +} + +func seedFamilyDeps(t *testing.T, db database.Store) (database.User, database.Organization, database.ChatModelConfig) { + t.Helper() + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "openai", + BaseUrl: "http://example.invalid", + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + IsDefault: true, + }) + return user, org, model +} diff --git a/coderd/x/chatd/chatstate/helpers_test.go b/coderd/x/chatd/chatstate/helpers_test.go new file mode 100644 index 0000000000..efbc73a9fe --- /dev/null +++ b/coderd/x/chatd/chatstate/helpers_test.go @@ -0,0 +1,92 @@ +package chatstate_test + +import ( + "context" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/testutil" +) + +// ownershipPublishCount returns the number of `chat:ownership` messages +// recorded so far on the test publisher. Tests use it to assert that +// transitions do or do not publish an ownership hint. +func (r *recordingPubsub) ownershipPublishCount() int { + count := 0 + for _, c := range r.channels { + if c == coderdpubsub.ChatStateOwnershipChannel { + count++ + } + } + return count +} + +// sendQueuedMessage seeds one queued user message via SendMessage with +// BusyBehaviorQueue. The chat must already be in a state that allows +// SendMessage (typically R0, R1, or I*). +func sendQueuedMessage(t *testing.T, f *testFixture, m *chatstate.ChatMachine, body string) chatstate.SendMessageResult { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + var send chatstate.SendMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + send, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage(body, f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + return send +} + +// sendInterruptMessage seeds one queued user message via SendMessage +// with BusyBehaviorInterrupt. From R0/R1 this transitions the chat to +// `interrupting` and appends the new user message to the queue tail. +func sendInterruptMessage(t *testing.T, f *testFixture, m *chatstate.ChatMachine, body string) chatstate.SendMessageResult { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + var send chatstate.SendMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + send, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage(body, f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorInterrupt, + }) + return err + })) + return send +} + +// queuedIDsByPosition returns the queued-message IDs for the chat in +// queue order. +func queuedIDsByPosition(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID) []int64 { + t.Helper() + rows, err := f.DB.GetChatQueuedMessagesByPosition(ctx, chatID) + require.NoError(t, err) + ids := make([]int64, len(rows)) + for i, r := range rows { + ids[i] = r.ID + } + return ids +} + +// historyMessageIDs returns the chat history message IDs ordered by +// row id. Used to assert that PromoteQueuedMessage from R1/I1 does NOT +// insert any history rows. +func historyMessageIDs(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID) []int64 { + t.Helper() + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + }) + require.NoError(t, err) + out := make([]int64, len(msgs)) + for i, m := range msgs { + out[i] = m.ID + } + return out +} diff --git a/coderd/x/chatd/chatstate/machine.go b/coderd/x/chatd/chatstate/machine.go new file mode 100644 index 0000000000..afe85ae1da --- /dev/null +++ b/coderd/x/chatd/chatstate/machine.go @@ -0,0 +1,291 @@ +package chatstate + +import ( + "context" + "database/sql" + "errors" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" +) + +// HeartbeatStaleSeconds is the threshold chatstate uses when deciding +// whether to publish a `chat:ownership` hint for a runnable chat. A +// heartbeat older than this many seconds (by database time) counts +// as stale and triggers a hint so an idle worker can attempt a +// takeover. +const HeartbeatStaleSeconds = 30 + +// ChatMachine is a chat-scoped handle for state-machine operations on +// a single chat row. It captures the database store, the pubsub +// publisher, and the chat ID at construction time so callers do not +// have to thread them through Update, Lock, or any transition method. +// +// ChatMachine values are cheap. Create one per chat for the lifetime +// of a request or worker turn; do not cache mutable chat state across +// calls. +type ChatMachine struct { + store database.Store + publisher Publisher + chatID uuid.UUID +} + +// NewChatMachine constructs a chat-scoped state machine handle. The +// store may be the root database handle or an existing transaction +// handle; publisher is the pubsub used for `chat:update` and +// `chat:ownership` emissions. Both are required and captured for the +// lifetime of the returned machine. +func NewChatMachine( + store database.Store, + publisher Publisher, + chatID uuid.UUID, +) *ChatMachine { + return &ChatMachine{ + store: store, + publisher: publisher, + chatID: chatID, + } +} + +// ChatID returns the chat ID this machine is scoped to. +func (m *ChatMachine) ChatID() uuid.UUID { return m.chatID } + +// Tx is the per-transaction handle passed to [ChatMachine.Update] +// callbacks. It carries the active context, the transactional store, +// and the chat ID. Tx does not cache mutable chat state across calls: +// every transition method reads the chat row and queue cardinality +// from the database on entry, so a bundle of transitions inside one +// Update callback always validates against the latest committed state. +type Tx struct { + ctx context.Context + store database.Store + chatID uuid.UUID +} + +// Ctx returns the context the surrounding [ChatMachine.Update] call +// is using. +func (tx *Tx) Ctx() context.Context { return tx.ctx } + +// ChatID returns the chat ID this transaction is scoped to. +func (tx *Tx) ChatID() uuid.UUID { return tx.chatID } + +// Store exposes the active transaction store so callers can perform +// validation reads (for example loading the messages affected by an +// EditMessage transition) and metadata writes (for example updating +// title or labels) that must be atomic with the transition. +// +// Callers MUST NOT use Store to mutate execution-state tables +// (chats.status, chat_messages, chat_queued_messages, chat_heartbeats, +// or the version fields on chats). Those mutations belong to the +// transition methods and are validated against the state machine +// matrix. +func (tx *Tx) Store() database.Store { return tx.store } + +// loadState reads the current chat row and queue cardinality from the +// active transaction, classifies the execution state, and returns the +// inputs every transition method needs. Returns ErrChatNotFound if +// the chat row was deleted in this transaction (or never existed). +func (tx *Tx) loadState() (database.Chat, ExecutionState, error) { + chat, err := tx.store.GetChatByID(tx.ctx, tx.chatID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return database.Chat{}, StateN, ErrChatNotFound + } + return database.Chat{}, "", xerrors.Errorf("load chat: %w", err) + } + count, err := tx.store.CountChatQueuedMessages(tx.ctx, tx.chatID) + if err != nil { + return database.Chat{}, "", xerrors.Errorf("count queued messages: %w", err) + } + return chat, ClassifyExecutionState(chat, count > 0, true), nil +} + +// requireFromAllowed loads the current state and validates t against +// the transition matrix. Returns the loaded chat and execution state +// on success, [ErrInvalidState] when the chat is in an invalid state +// and t is not [TransitionReconcileInvalidState], and a typed +// *TransitionError otherwise. +func (tx *Tx) requireFromAllowed(t Transition) (database.Chat, ExecutionState, error) { + chat, from, err := tx.loadState() + if err != nil { + return chat, from, err + } + if from == StateInvalid && t != TransitionReconcileInvalidState { + return chat, from, ErrInvalidState + } + if err := requireExecutionTransition(t, from); err != nil { + return chat, from, err + } + return chat, from, nil +} + +// Update applies one or more transitions to the machine's chat. +// +// Update opens a transaction on the captured store, atomically locks +// the chat row with FOR UPDATE and increments `snapshot_version` +// exactly once, then runs fn against a fresh [*Tx] and the active +// transaction store. It constructs a [PublishBuffer], enqueues +// `chat:update` (and a `chat:ownership` hint +// when the post-transition state is worker-runnable and ownership is +// missing or stale) inside the transaction, and flushes the buffer only after +// the transaction function succeeds. If the transaction rolls back, +// the deferred Discard suppresses every buffered publication so +// subscribers never see uncommitted state. +// +// If Update is called with a store that is already in a transaction, +// [database.Store.InTx] reuses the active transaction. In that case, +// callers that need outer-transaction publication semantics can pass a +// [PublishBuffer] as the machine publisher. The inner buffer flushes +// into the outer buffer, and the outer owner remains responsible for +// publishing only after the outer transaction commits. +// +// If the chat row does not exist, Update returns [ErrChatNotFound] +// without mutating anything. +// +// Callbacks that return an error roll back the transaction (rolling +// back the automatic snapshot bump) and publish nothing. +func (m *ChatMachine) Update( + ctx context.Context, + fn func(*Tx, database.Store) error, +) error { + if m.store == nil { + return xerrors.New("chatstate: ChatMachine has nil store") + } + if m.publisher == nil { + return xerrors.New("chatstate: ChatMachine has nil publisher") + } + + buffer := NewPublishBuffer(m.publisher) + defer buffer.Discard() + + err := m.store.InTx(func(store database.Store) error { + if _, err := store.LockChatAndBumpSnapshotVersion(ctx, m.chatID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrChatNotFound + } + return xerrors.Errorf("lock chat and bump snapshot: %w", err) + } + tx := &Tx{ + ctx: ctx, + store: store, + chatID: m.chatID, + } + if err := fn(tx, store); err != nil { + return err + } + chat, state, err := tx.loadState() + if err != nil { + return err + } + if err := buffer.Publish( + coderdpubsub.ChatStateUpdateChannel(chat.ID), + buildChatUpdateMessage(chat), + ); err != nil { + return xerrors.Errorf("buffer chat update: %w", err) + } + if state.IsRunnable() { + stale, err := ownershipStaleOrMissing(ctx, store, chat, HeartbeatStaleSeconds) + if err != nil { + return xerrors.Errorf("evaluate ownership: %w", err) + } + if stale { + if err := buffer.Publish( + coderdpubsub.ChatStateOwnershipChannel, + buildChatOwnershipMessage(chat), + ); err != nil { + return xerrors.Errorf("buffer ownership hint: %w", err) + } + } + } + return nil + }, nil) + if err != nil { + return err + } + return buffer.Flush() +} + +// Lock locks the chat row with FOR UPDATE and runs fn in a +// transaction without advancing snapshot_version. It uses the store +// captured by [NewChatMachine]. Use it when the caller needs a +// consistent chat snapshot plus related rows such as messages or +// queued messages but is NOT applying a transition. +// +// Callers must not pass a store here; it belongs on the machine. +// +// Lock publishes nothing. Callback errors roll back the transaction +// and propagate to the caller. +func (m *ChatMachine) Lock( + ctx context.Context, + fn func(database.Store) error, +) error { + if m.store == nil { + return xerrors.New("chatstate: ChatMachine has nil store") + } + return m.store.InTx(func(store database.Store) error { + // GetChatByIDForUpdate locks the row WITHOUT bumping snapshot. + _, err := store.GetChatByIDForUpdate(ctx, m.chatID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrChatNotFound + } + return xerrors.Errorf("lock chat: %w", err) + } + return fn(store) + }, nil) +} + +// ReadLock takes a shared lock on the chat row with FOR SHARE and runs +// fn in a transaction without advancing snapshot_version. It uses the +// store captured by [NewChatMachine]. Use it when the caller needs a +// consistent chat snapshot plus related rows such as messages or queued +// messages but is NOT applying a transition and does NOT need to block +// concurrent readers. +// +// Unlike [ChatMachine.Lock], the FOR SHARE lock permits other shared +// lockers to proceed concurrently while still blocking writers that take +// FOR UPDATE (such as [ChatMachine.Update] and [ChatMachine.Lock]) until +// the transaction commits. +// +// Callers must not pass a store here; it belongs on the machine. +// +// ReadLock publishes nothing. Callback errors roll back the transaction +// and propagate to the caller. +func (m *ChatMachine) ReadLock( + ctx context.Context, + fn func(database.Store) error, +) error { + if m.store == nil { + return xerrors.New("chatstate: ChatMachine has nil store") + } + return m.store.InTx(func(store database.Store) error { + // GetChatByIDForShare takes a shared lock on the row WITHOUT + // bumping snapshot. + _, err := store.GetChatByIDForShare(ctx, m.chatID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrChatNotFound + } + return xerrors.Errorf("read lock chat: %w", err) + } + return fn(store) + }, nil) +} + +// ownershipStaleOrMissing reports whether the chat's current +// (chat_id, runner_id) lease is missing or stale. The staleSeconds +// threshold is forwarded to [database.IsChatHeartbeatStale] so the +// comparison runs against database time inside a single SQL query. +func ownershipStaleOrMissing(ctx context.Context, store database.Store, chat database.Chat, staleSeconds int32) (bool, error) { + if !chat.WorkerID.Valid || !chat.RunnerID.Valid { + return true, nil + } + return store.IsChatHeartbeatStale(ctx, database.IsChatHeartbeatStaleParams{ + ChatID: chat.ID, + RunnerID: chat.RunnerID.UUID, + StaleSeconds: staleSeconds, + }) +} diff --git a/coderd/x/chatd/chatstate/machine_test.go b/coderd/x/chatd/chatstate/machine_test.go new file mode 100644 index 0000000000..65e96b0f8a --- /dev/null +++ b/coderd/x/chatd/chatstate/machine_test.go @@ -0,0 +1,411 @@ +package chatstate_test + +import ( + "context" + "database/sql" + "encoding/json" + "slices" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// testFixture bundles the resources every integration test needs: +// a database, a publisher recorder, a user/org/model triple, and +// helper accessors. It is intentionally NOT a generic chatd test +// fixture; tests outside this package should not depend on it. +type testFixture struct { + DB database.Store + PubSub pubsub.Pubsub + Pub *recordingPubsub + User database.User + Org database.Organization + Model database.ChatModelConfig + APIKey database.APIKey +} + +// apiKeyID returns the fixture API key wrapped for the chatstate +// inputs that require a non-null api_key_id (for example EditMessage). +func (f *testFixture) apiKeyID() sql.NullString { + return sql.NullString{String: f.APIKey.ID, Valid: true} +} + +func newTestFixture(t *testing.T) *testFixture { + t.Helper() + db, ps := dbtestutil.NewDB(t) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "openai", + BaseUrl: "http://example.invalid", + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + IsDefault: true, + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + pub := newRecordingPubsub() + return &testFixture{ + DB: db, + PubSub: ps, + Pub: pub, + User: user, + Org: org, + Model: model, + APIKey: apiKey, + } +} + +// readChat re-reads the chat from the database. Tests use this to +// verify post-transition state because transition results no longer +// carry the chat snapshot. +func (f *testFixture) readChat(ctx context.Context, t *testing.T, chatID uuid.UUID) database.Chat { + t.Helper() + chat, err := f.DB.GetChatByID(ctx, chatID) + require.NoError(t, err) + return chat +} + +// classify reads the chat plus queue cardinality and returns the +// execution state. +func (f *testFixture) classify(ctx context.Context, t *testing.T, chatID uuid.UUID) chatstate.ExecutionState { + t.Helper() + chat := f.readChat(ctx, t, chatID) + count, err := f.DB.CountChatQueuedMessages(ctx, chatID) + require.NoError(t, err) + return chatstate.ClassifyExecutionState(chat, count > 0, true) +} + +// recordingPubsub captures every Publish call so tests can assert on +// the chatstate notifications without needing a live subscriber. The +// mutex makes it safe to use from concurrent tests that race multiple +// goroutines through the same publisher (see TestConcurrentUpdatesSerializeOnChatRow). +type recordingPubsub struct { + mu sync.Mutex + channels []string + payloads [][]byte +} + +func newRecordingPubsub() *recordingPubsub { return &recordingPubsub{} } + +func (r *recordingPubsub) Publish(channel string, payload []byte) error { + r.mu.Lock() + defer r.mu.Unlock() + r.channels = append(r.channels, channel) + r.payloads = append(r.payloads, slices.Clone(payload)) + return nil +} + +// expectChatUpdate finds the most recent chat:update message on the +// per-chat channel and asserts that it has snapshot_version == want. +func (r *recordingPubsub) expectChatUpdate(t *testing.T, chatID uuid.UUID, wantSnapshot int64) { + t.Helper() + channel := coderdpubsub.ChatStateUpdateChannel(chatID) + for i := len(r.channels) - 1; i >= 0; i-- { + if r.channels[i] != channel { + continue + } + var msg coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(r.payloads[i], &msg)) + require.Equal(t, wantSnapshot, msg.SnapshotVersion) + return + } + t.Fatalf("no chat:update on %s", channel) +} + +func (r *recordingPubsub) hasOwnership() bool { + for _, c := range r.channels { + if c == coderdpubsub.ChatStateOwnershipChannel { + return true + } + } + return false +} + +func userTextMessage(text string, createdBy uuid.UUID, modelConfigID uuid.UUID) chatstate.Message { + parts := []codersdk.ChatMessagePart{codersdk.ChatMessageText(text)} + raw, err := chatprompt.MarshalParts(parts) + if err != nil { + panic(err) + } + return chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + } +} + +// createTestChat is the standard "fresh R0 chat" helper used by other +// tests. It exercises CreateChat itself. +func createTestChat(t *testing.T, f *testFixture) chatstate.CreateChatResult { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + userTextMessage("hello", f.User.ID, f.Model.ID), + }, + }) + require.NoError(t, err) + return res +} + +func TestChatMachine_Update_RejectsMissingChat(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, uuid.New()) + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { return nil }) + require.ErrorIs(t, err, chatstate.ErrChatNotFound) + require.Empty(t, f.Pub.channels) +} + +func TestChatMachine_Lock_DoesNotBumpSnapshot(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + before := f.readChat(ctx, t, created.Chat.ID) + publishedBefore := len(f.Pub.channels) + + require.NoError(t, m.Lock(ctx, func(_ database.Store) error { + return nil + })) + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.SnapshotVersion, after.SnapshotVersion) + require.Equal(t, publishedBefore, len(f.Pub.channels), "Lock must not publish") +} + +func TestChatMachine_ReadLock_DoesNotBumpSnapshot(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + before := f.readChat(ctx, t, created.Chat.ID) + publishedBefore := len(f.Pub.channels) + + var called bool + require.NoError(t, m.ReadLock(ctx, func(_ database.Store) error { + called = true + return nil + })) + require.True(t, called, "ReadLock must invoke the callback") + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.SnapshotVersion, after.SnapshotVersion) + require.Equal(t, publishedBefore, len(f.Pub.channels), "ReadLock must not publish") +} + +func TestChatMachine_ReadLock_RejectsMissingChat(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, uuid.New()) + err := m.ReadLock(ctx, func(_ database.Store) error { + t.Fatal("callback must not run when the chat is missing") + return nil + }) + require.ErrorIs(t, err, chatstate.ErrChatNotFound) + require.Empty(t, f.Pub.channels) +} + +func TestChatMachine_UpdatePublishesAfterCommit(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + publishedBefore := len(f.Pub.channels) + // Run a no-op Update; snapshot bump still happens, one update message + // should follow the commit. + require.NoError(t, m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil })) + channel := coderdpubsub.ChatStateUpdateChannel(created.Chat.ID) + var found bool + for _, c := range f.Pub.channels[publishedBefore:] { + if c == channel { + found = true + break + } + } + require.True(t, found, "expected one chat:update message after commit") +} + +func TestChatMachine_FailedUpdate_PublishesNothing(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + before := f.readChat(ctx, t, created.Chat.ID) + channelsBefore := len(f.Pub.channels) + expected := newSentinel() + cbErr := m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return expected }) + require.ErrorIs(t, cbErr, expected) + require.Equal(t, channelsBefore, len(f.Pub.channels), "failed update should not publish") + // snapshot_version should not have advanced. + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.SnapshotVersion, after.SnapshotVersion) +} + +func TestMessageRevisionTrigger_AssignsRevisionFromSnapshot(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) // snapshot 1, history_version 1 via trigger + + // CommitStep an assistant message; it should land with revision = chat.snapshot_version after the bump. + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + assistant := userTextMessage("assistant", f.User.ID, f.Model.ID) + assistant.Role = database.ChatMessageRoleAssistant + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistant}, + }) + return err + })) + require.Len(t, step.InsertedMessages, 1) + after := f.readChat(ctx, t, created.Chat.ID) + // The Update call bumps snapshot_version once before the trigger + // runs, so the new revision should equal the bumped snapshot. + require.Equal(t, after.SnapshotVersion, step.InsertedMessages[0].Revision) + require.Equal(t, after.SnapshotVersion, after.HistoryVersion) + require.Equal(t, int64(0), after.GenerationAttempt, "trigger resets generation_attempt to 0") +} + +func TestQueueVersionTrigger_AdvancesOnInsert(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) // queue_version starts at 0 + + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("queue", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, after.SnapshotVersion, after.QueueVersion) + require.Greater(t, after.QueueVersion, int64(0)) +} + +func TestQueueVersionTrigger_StableForNonQueueMutations(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + assistant := userTextMessage("assistant", f.User.ID, f.Model.ID) + assistant.Role = database.ChatMessageRoleAssistant + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistant}, + }) + return err + })) + // queue_version must remain unchanged from initial 0. + require.Equal(t, int64(0), f.readChat(ctx, t, created.Chat.ID).QueueVersion) +} + +// TestUpdateFlushesBufferedPublicationsAfterCommit verifies that +// ChatMachine.Update owns the PublishBuffer lifecycle: nothing +// reaches the inner publisher until after the transaction commits, +// and at commit the buffered chat:update is forwarded. +func TestUpdateFlushesBufferedPublicationsAfterCommit(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + channel := coderdpubsub.ChatStateUpdateChannel(created.Chat.ID) + baseline := countChannel(f.Pub.channels, channel) + + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // During the callback, no new chat:update for this chat may have + // reached the inner publisher because the buffer holds it. + require.NoError(t, m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { + require.Equal(t, baseline, countChannel(f.Pub.channels, channel), + "inner publisher saw chat:update before transaction committed") + return nil + })) + + require.Equal(t, baseline+1, countChannel(f.Pub.channels, channel), + "exactly one new chat:update reached the inner publisher after commit") +} + +// TestUpdateDiscardsBufferedPublicationsOnCallbackError verifies the +// deferred Discard path: when the callback returns an error the +// transaction rolls back and no buffered messages reach the inner +// publisher. +func TestUpdateDiscardsBufferedPublicationsOnCallbackError(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + before := f.readChat(ctx, t, created.Chat.ID) + channelsBefore := len(f.Pub.channels) + + sentinel := xerrors.New("callback boom") + err := m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return sentinel }) + require.ErrorIs(t, err, sentinel) + + require.Equal(t, channelsBefore, len(f.Pub.channels), + "failed update must not flush any buffered publications") + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, before.SnapshotVersion, after.SnapshotVersion, + "snapshot bump rolled back when callback returns error") +} + +type sentinelError struct{ msg string } + +func (s *sentinelError) Error() string { return s.msg } + +func newSentinel() error { return &sentinelError{msg: "sentinel"} } + +func countChannel(channels []string, channel string) int { + c := 0 + for _, ch := range channels { + if ch == channel { + c++ + } + } + return c +} diff --git a/coderd/x/chatd/chatstate/messages.go b/coderd/x/chatd/chatstate/messages.go new file mode 100644 index 0000000000..ee92e0ea13 --- /dev/null +++ b/coderd/x/chatd/chatstate/messages.go @@ -0,0 +1,115 @@ +package chatstate + +import ( + "database/sql" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + + "github.com/coder/coder/v2/coderd/database" +) + +// Message is the durable message input shape used by chatstate +// transitions. It is intentionally lower level than the SDK message +// request types: callers must produce a fully materialized message +// (parsed parts, calculated cost, resolved model config) before +// passing it in. +// +// The state machine never reshapes a Message except to attach the +// runtime `chat_id`. +type Message struct { + Role database.ChatMessageRole + Content pqtype.NullRawMessage + Visibility database.ChatMessageVisibility + ModelConfigID uuid.NullUUID + CreatedBy uuid.NullUUID + ContentVersion int16 + Compressed bool + InputTokens sql.NullInt64 + OutputTokens sql.NullInt64 + TotalTokens sql.NullInt64 + ReasoningTokens sql.NullInt64 + CacheCreationTokens sql.NullInt64 + CacheReadTokens sql.NullInt64 + ContextLimit sql.NullInt64 + TotalCostMicros sql.NullInt64 + RuntimeMs sql.NullInt64 + ProviderResponseID sql.NullString + APIKeyID sql.NullString +} + +// toInsertParams converts a batch of Messages into the parallel-array +// shape required by `InsertChatMessages`. The returned struct has all +// arrays sized to len(messages). +// +// The chat ID is supplied by the caller because Message itself does +// not carry one (the chat machine already knows the chat). +func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMessagesParams { + n := len(messages) + params := database.InsertChatMessagesParams{ + ChatID: chatID, + CreatedBy: make([]uuid.UUID, n), + ModelConfigID: make([]uuid.UUID, n), + APIKeyID: make([]string, n), + Role: make([]database.ChatMessageRole, n), + Content: make([]string, n), + ContentVersion: make([]int16, n), + Visibility: make([]database.ChatMessageVisibility, n), + InputTokens: make([]int64, n), + OutputTokens: make([]int64, n), + TotalTokens: make([]int64, n), + ReasoningTokens: make([]int64, n), + CacheCreationTokens: make([]int64, n), + CacheReadTokens: make([]int64, n), + ContextLimit: make([]int64, n), + Compressed: make([]bool, n), + TotalCostMicros: make([]int64, n), + RuntimeMs: make([]int64, n), + ProviderResponseID: make([]string, n), + } + for i, m := range messages { + params.CreatedBy[i] = nullUUIDOrNil(m.CreatedBy) + params.ModelConfigID[i] = nullUUIDOrNil(m.ModelConfigID) + if m.APIKeyID.Valid { + params.APIKeyID[i] = m.APIKeyID.String + } + params.Role[i] = m.Role + if m.Content.Valid { + params.Content[i] = string(m.Content.RawMessage) + } else { + // Use the JSON null literal; UNNEST + ::jsonb requires a + // valid JSON value and the trigger leaves it untouched. + params.Content[i] = "null" + } + params.ContentVersion[i] = m.ContentVersion + params.Visibility[i] = m.Visibility + params.InputTokens[i] = nullInt64Or(m.InputTokens, 0) + params.OutputTokens[i] = nullInt64Or(m.OutputTokens, 0) + params.TotalTokens[i] = nullInt64Or(m.TotalTokens, 0) + params.ReasoningTokens[i] = nullInt64Or(m.ReasoningTokens, 0) + params.CacheCreationTokens[i] = nullInt64Or(m.CacheCreationTokens, 0) + params.CacheReadTokens[i] = nullInt64Or(m.CacheReadTokens, 0) + params.ContextLimit[i] = nullInt64Or(m.ContextLimit, 0) + params.Compressed[i] = m.Compressed + params.TotalCostMicros[i] = nullInt64Or(m.TotalCostMicros, 0) + params.RuntimeMs[i] = nullInt64Or(m.RuntimeMs, 0) + if m.ProviderResponseID.Valid { + params.ProviderResponseID[i] = m.ProviderResponseID.String + } + } + return params +} + +func nullUUIDOrNil(u uuid.NullUUID) uuid.UUID { + if u.Valid { + return u.UUID + } + return uuid.Nil +} + +func nullInt64Or(v sql.NullInt64, fallback int64) int64 { + if v.Valid { + return v.Int64 + } + return fallback +} diff --git a/coderd/x/chatd/chatstate/notify.go b/coderd/x/chatd/chatstate/notify.go new file mode 100644 index 0000000000..4b7c44eee3 --- /dev/null +++ b/coderd/x/chatd/chatstate/notify.go @@ -0,0 +1,166 @@ +package chatstate + +import ( + "encoding/json" + "errors" + "fmt" + "slices" + "sync" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" +) + +// Publisher is the minimal interface chatstate needs to publish +// pubsub messages. It is intentionally compatible with +// database/pubsub.Pubsub: real callers pass the live pubsub directly +// and tests pass a recording fake. +type Publisher interface { + Publish(event string, message []byte) error +} + +// PublishBuffer is a [Publisher] that records each Publish call in +// order without forwarding it until [PublishBuffer.Flush] is called. +// It is an internal primitive used by chatstate entry points to +// hold pubsub messages until the surrounding transaction commits, +// and by tests that need to observe buffered output. Normal callers +// do not construct a PublishBuffer themselves and do not invoke +// Flush or Discard; chatstate's entry points own that lifecycle. +type PublishBuffer struct { + inner Publisher + + mu sync.Mutex + pending []bufferedMessage + flushed bool + disabled bool +} + +type bufferedMessage struct { + Channel string + Payload []byte +} + +// NewPublishBuffer constructs a PublishBuffer that, when flushed, will +// forward messages in order to inner. +func NewPublishBuffer(inner Publisher) *PublishBuffer { + return &PublishBuffer{inner: inner} +} + +// Publish records a message. It never forwards to the inner publisher +// until [PublishBuffer.Flush] is called. Returns an error if Flush has +// already happened to make accidental reuse obvious. +func (b *PublishBuffer) Publish(channel string, payload []byte) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.flushed { + return xerrors.Errorf("publish buffer flushed; cannot accept message for %q", channel) + } + if b.disabled { + return nil + } + b.pending = append(b.pending, bufferedMessage{Channel: channel, Payload: slices.Clone(payload)}) + return nil +} + +// Flush forwards every pending message to the inner publisher in the +// order it was buffered, then marks the buffer flushed. Joined publish +// errors are returned with channel names annotated after every pending +// message has been attempted. +func (b *PublishBuffer) Flush() error { + b.mu.Lock() + defer b.mu.Unlock() + if b.flushed { + return nil + } + b.flushed = true + var errs []error + for _, msg := range b.pending { + if err := b.inner.Publish(msg.Channel, msg.Payload); err != nil { + errs = append(errs, xerrors.Errorf("publish %s: %w", msg.Channel, err)) + } + } + return errors.Join(errs...) +} + +// Discard clears the buffered messages without forwarding them. It +// is safe to call multiple times and is harmless after [PublishBuffer.Flush]: +// once Flush has marked the buffer flushed and forwarded its +// pending messages, a subsequent Discard simply clears the (now +// empty) pending slice and sets the buffer to drop any future +// Publish calls. This makes `defer buf.Discard()` a safe pattern +// after a successful flush, including the one chatstate entry +// points use to own the buffer lifecycle. +func (b *PublishBuffer) Discard() { + b.mu.Lock() + defer b.mu.Unlock() + b.pending = nil + b.disabled = true +} + +// pending returns a snapshot of the buffered messages, primarily for +// tests via [PublishBuffer.BufferedChannels]. The returned slice is a +// copy and safe to inspect without holding the buffer lock. +func (b *PublishBuffer) snapshotPending() []bufferedMessage { + b.mu.Lock() + defer b.mu.Unlock() + out := make([]bufferedMessage, len(b.pending)) + copy(out, b.pending) + return out +} + +// BufferedChannels returns just the channels of the pending messages +// in order. Primarily useful for assertions in tests. +func (b *PublishBuffer) BufferedChannels() []string { + pending := b.snapshotPending() + out := make([]string, len(pending)) + for i, m := range pending { + out[i] = m.Channel + } + return out +} + +// buildChatUpdateMessage produces the JSON payload for a +// `chat:update:{chat_id}` message describing the post-transition +// snapshot of chat. +func buildChatUpdateMessage(chat database.Chat) []byte { + msg := coderdpubsub.ChatStateUpdateMessage{ + SnapshotVersion: chat.SnapshotVersion, + HistoryVersion: chat.HistoryVersion, + QueueVersion: chat.QueueVersion, + RetryStateVersion: chat.RetryStateVersion, + GenerationAttempt: chat.GenerationAttempt, + Status: string(chat.Status), + Archived: chat.Archived, + } + if chat.WorkerID.Valid { + id := chat.WorkerID.UUID + msg.WorkerID = &id + } + if chat.RunnerID.Valid { + id := chat.RunnerID.UUID + msg.RunnerID = &id + } + payload, err := json.Marshal(msg) + if err != nil { + // json.Marshal on this struct is total; panic is acceptable + // because the only failure mode would be a bug in this + // package, not user input. + panic(fmt.Sprintf("marshal chat state update: %v", err)) + } + return payload +} + +// buildChatOwnershipMessage produces the JSON payload for the global +// `chat:ownership` ownership hint for chat. +func buildChatOwnershipMessage(chat database.Chat) []byte { + payload, err := json.Marshal(coderdpubsub.ChatStateOwnershipMessage{ + ChatID: chat.ID, + SnapshotVersion: chat.SnapshotVersion, + }) + if err != nil { + panic(fmt.Sprintf("marshal chat state ownership: %v", err)) + } + return payload +} diff --git a/coderd/x/chatd/chatstate/notify_integration_test.go b/coderd/x/chatd/chatstate/notify_integration_test.go new file mode 100644 index 0000000000..9030a56e71 --- /dev/null +++ b/coderd/x/chatd/chatstate/notify_integration_test.go @@ -0,0 +1,382 @@ +package chatstate_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/testutil" +) + +// publishedOn returns the indices into f.Pub.channels (and f.Pub.payloads) +// that match the given channel name, in order. +func publishedOn(f *testFixture, channel string) []int { + var idx []int + for i, c := range f.Pub.channels { + if c == channel { + idx = append(idx, i) + } + } + return idx +} + +// TestCreateChatPublishesAfterCommit asserts that a successful +// CreateChat call publishes exactly one chat:update message on the +// per-chat channel after the inner transaction commits. +func TestCreateChatPublishesAfterCommit(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + res := createTestChat(t, f) + + channel := coderdpubsub.ChatStateUpdateChannel(res.Chat.ID) + idx := publishedOn(f, channel) + require.Len(t, idx, 1, "exactly one chat:update for the new chat") + + var msg coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(f.Pub.payloads[idx[0]], &msg)) + require.Equal(t, res.Chat.SnapshotVersion, msg.SnapshotVersion) + require.Equal(t, string(database.ChatStatusRunning), msg.Status) +} + +// TestUpdatePublishesAfterCommit asserts that ChatMachine.Update +// publishes one chat:update on the per-chat channel after the inner +// transaction commits, even when the callback performs no transition. +func TestUpdatePublishesAfterCommit(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + createIdx := publishedOn(f, coderdpubsub.ChatStateUpdateChannel(created.Chat.ID)) + require.Len(t, createIdx, 1, "create published one chat:update") + + require.NoError(t, m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil })) + + updIdx := publishedOn(f, coderdpubsub.ChatStateUpdateChannel(created.Chat.ID)) + require.Len(t, updIdx, 2, "no-op Update still publishes a chat:update") + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + var msg coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(f.Pub.payloads[updIdx[1]], &msg)) + require.Equal(t, after.SnapshotVersion, msg.SnapshotVersion) +} + +// TestUpdatePublishesOneFinalChatUpdateForTransitionBundle bundles +// several transitions inside one Update callback and verifies the +// commit publishes exactly one chat:update on the per-chat channel +// (not one per transition). +func TestUpdatePublishesOneFinalChatUpdateForTransitionBundle(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + baseUpdates := len(publishedOn(f, coderdpubsub.ChatStateUpdateChannel(created.Chat.ID))) + + // chatstate.StateR0 -> chatstate.StateW (FinishTurn) -> + // chatstate.StateXW (SetArchived true) -> chatstate.StateW + // (SetArchived false). + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := tx.FinishTurn(chatstate.FinishTurnInput{}); err != nil { + return err + } + if _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true}); err != nil { + return err + } + if _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: false}); err != nil { + return err + } + return nil + })) + + updIdx := publishedOn(f, coderdpubsub.ChatStateUpdateChannel(created.Chat.ID)) + require.Equal(t, baseUpdates+1, len(updIdx), + "three-transition bundle publishes exactly one final chat:update") + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusWaiting, after.Status, + "ends in "+chatstate.StateW.String()) +} + +// TestUpdateAppliesTransitionBundleSequentially verifies that +// transitions chained inside a single Update callback see each +// other's effects: later transitions validate against the state +// produced by earlier ones (chatstate.StateR0 -> chatstate.StateW +// is rejected when called twice because the second call sees +// chatstate.StateW and FinishTurn is no longer allowed). +func TestUpdateAppliesTransitionBundleSequentially(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := tx.FinishTurn(chatstate.FinishTurnInput{}); err != nil { + return err + } + // Second FinishTurn should fail because state is now chatstate.StateW. + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + + // Failed bundle rolls back: state must not have advanced past + // chatstate.StateR0. + require.Equal(t, chatstate.StateR0, f.classify(ctx, t, created.Chat.ID), + "failed bundle rolls back the whole transaction") +} + +// TestFailedUpdatePublishesNothing verifies that a callback error +// rolls back the snapshot bump and publishes nothing. +func TestFailedUpdatePublishesNothing(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + publishedBefore := len(f.Pub.channels) + beforeChat := f.readChat(ctx, t, created.Chat.ID) + + sentinel := xerrors.New("forced failure") + err := m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return sentinel }) + require.ErrorIs(t, err, sentinel) + require.Equal(t, publishedBefore, len(f.Pub.channels), "failed update publishes nothing") + + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, beforeChat.SnapshotVersion, after.SnapshotVersion, + "failed update rolls back snapshot bump") +} + +// TestLockPublishesNothing verifies that Lock does not publish even +// though it locks the chat row. +func TestLockPublishesNothing(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + publishedBefore := len(f.Pub.channels) + require.NoError(t, m.Lock(ctx, func(_ database.Store) error { return nil })) + require.Equal(t, publishedBefore, len(f.Pub.channels), "Lock publishes nothing") +} + +// TestPublishBufferWithRolledBackOuterTransactionPublishesNothing +// wires a chatstate machine through a PublishBuffer and exercises +// the buffer primitive directly: when the caller discards before +// flushing, the inner publisher receives nothing. ChatMachine.Update +// uses the same primitive internally with a deferred Discard; +// callers no longer drive Flush or Discard themselves. +func TestPublishBufferWithRolledBackOuterTransactionPublishesNothing(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + // Run one normal Update to establish a stable baseline channel + // count. CreateChat plus this Update may publish chat:update + // and chat:ownership messages depending on ownership, so we + // take the snapshot after that activity settles. + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil })) + baseline := len(f.Pub.channels) + + // Now exercise the PublishBuffer rollback path explicitly. The + // outer transaction "rolls back": the caller buffers messages, + // discards them, then flushes. The inner publisher must see + // none of the buffered messages. + buf := chatstate.NewPublishBuffer(f.Pub) + require.NoError(t, buf.Publish("chat:update:bogus", []byte("payload"))) + require.NoError(t, buf.Publish("chat:ownership", []byte("payload"))) + buf.Discard() + require.NoError(t, buf.Flush()) + + require.Equal(t, baseline, len(f.Pub.channels), + "discarded buffer publishes nothing through the inner publisher") +} + +// TestChatUpdateMessagePayloadShape verifies the JSON shape of the +// chat:update payload contains every field consumers depend on: +// snapshot_version, history_version, queue_version, +// retry_state_version, generation_attempt, status, archived, and +// worker_id / runner_id, with explicit nulls when unowned. +func TestChatUpdateMessagePayloadShape(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + channel := coderdpubsub.ChatStateUpdateChannel(created.Chat.ID) + + // The create update is unowned and must still include explicit + // null ownership fields. + createIdx := publishedOn(f, channel) + require.NotEmpty(t, createIdx) + var createRaw map[string]json.RawMessage + require.NoError(t, json.Unmarshal(f.Pub.payloads[createIdx[0]], &createRaw)) + require.JSONEq(t, `null`, string(createRaw["worker_id"])) + require.JSONEq(t, `null`, string(createRaw["runner_id"])) + + // Acquire ownership so worker_id and runner_id are present. + worker := uuid.New() + runner := uuid.New() + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner}) + return err + })) + + // Find the last chat:update message. + idx := publishedOn(f, channel) + require.NotEmpty(t, idx) + last := f.Pub.payloads[idx[len(idx)-1]] + + // Strict-decode against the typed struct. + var typed coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(last, &typed)) + require.Greater(t, typed.SnapshotVersion, int64(0)) + require.NotNil(t, typed.WorkerID) + require.Equal(t, worker, *typed.WorkerID) + require.NotNil(t, typed.RunnerID) + require.Equal(t, runner, *typed.RunnerID) + require.Equal(t, string(database.ChatStatusRunning), typed.Status) + require.False(t, typed.Archived) + + // Permissive decode to assert exact JSON keys. + var raw map[string]json.RawMessage + require.NoError(t, json.Unmarshal(last, &raw)) + for _, key := range []string{ + "snapshot_version", + "history_version", + "queue_version", + "retry_state_version", + "generation_attempt", + "status", + "archived", + "worker_id", + "runner_id", + } { + _, ok := raw[key] + require.True(t, ok, "payload missing key %q", key) + } +} + +// TestChatOwnershipMessagePayloadShape verifies the JSON shape of +// chat:ownership: chat_id and snapshot_version. +func TestChatOwnershipMessagePayloadShape(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + // CreateChat publishes one ownership hint because the new chat is + // unowned and runnable. + created := createTestChat(t, f) + + idx := publishedOn(f, coderdpubsub.ChatStateOwnershipChannel) + require.NotEmpty(t, idx, "CreateChat publishes at least one chat:ownership hint") + + payload := f.Pub.payloads[idx[len(idx)-1]] + var typed coderdpubsub.ChatStateOwnershipMessage + require.NoError(t, json.Unmarshal(payload, &typed)) + require.Equal(t, created.Chat.ID, typed.ChatID) + require.Greater(t, typed.SnapshotVersion, int64(0)) + + var raw map[string]json.RawMessage + require.NoError(t, json.Unmarshal(payload, &raw)) + for _, key := range []string{"chat_id", "snapshot_version"} { + _, ok := raw[key] + require.True(t, ok, "ownership payload missing key %q", key) + } +} + +// TestOwnershipNotificationUsesDatabaseHeartbeatStaleness verifies +// that an ownership hint fires when the heartbeat is stale by the +// database's clock, regardless of what the local Go clock says. We +// rewrite the heartbeat row to a deterministically old timestamp via +// raw SQL and confirm the post-commit hint is sent on a subsequent +// runnable Update. +func TestOwnershipNotificationUsesDatabaseHeartbeatStaleness(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // Acquire ownership; this writes a fresh heartbeat. + worker := uuid.New() + runner := uuid.New() + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner}) + return err + })) + hb, err := f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: runner, + }) + require.NoError(t, err) + require.WithinDuration(t, time.Now(), hb.HeartbeatAt, time.Minute, + "Acquire wrote a fresh heartbeat") + + // Snapshot ownership-hint count before the test trigger. + ownershipBefore := f.Pub.ownershipPublishCount() + + // Force the heartbeat to a deterministically old time. + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_heartbeats + SET heartbeat_at = NOW() - INTERVAL '1 hour' + WHERE chat_id = $1 AND runner_id = $2 + `, created.Chat.ID, runner) + require.NoError(t, err) + + // Confirm database-side staleness check agrees. + stale, err := f.DB.IsChatHeartbeatStale(ctx, database.IsChatHeartbeatStaleParams{ + ChatID: created.Chat.ID, + RunnerID: runner, + StaleSeconds: chatstate.HeartbeatStaleSeconds, + }) + require.NoError(t, err) + require.True(t, stale, "heartbeat is stale per database time") + + // Run a no-op Update. The chat is runnable (chatstate.StateR0) + // and the heartbeat is stale, so post-commit logic must publish + // exactly one chat:ownership hint. + require.NoError(t, m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil })) + + ownershipAfter := f.Pub.ownershipPublishCount() + require.Equal(t, ownershipBefore+1, ownershipAfter, + "stale heartbeat triggers a fresh ownership hint") +} + +// TestUpdateContextCancellationPublishesNothing verifies that +// canceling the caller's context (between the inner commit and the +// publish loop's first call) does not corrupt state. We exercise the +// simpler observable contract: when the user cancels before Update +// gets to do anything, nothing is published. The strict before-publish +// race is exercised in concurrency tests with channel sync. +func TestUpdateContextCancellationPublishesNothing(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + publishedBefore := len(f.Pub.channels) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + err := m.Update(ctx, func(_ *chatstate.Tx, _ database.Store) error { return nil }) + require.Error(t, err) + require.Equal(t, publishedBefore, len(f.Pub.channels), + "caller-aborted update publishes nothing") +} diff --git a/coderd/x/chatd/chatstate/notify_internal_test.go b/coderd/x/chatd/chatstate/notify_internal_test.go new file mode 100644 index 0000000000..833f102295 --- /dev/null +++ b/coderd/x/chatd/chatstate/notify_internal_test.go @@ -0,0 +1,120 @@ +package chatstate + +import ( + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" +) + +type recordingPublisher struct { + calls []recordedCall + errOn map[string]error + failed map[string]int +} + +type recordedCall struct { + Channel string + Payload []byte +} + +func newRecordingPublisher() *recordingPublisher { + return &recordingPublisher{ + errOn: map[string]error{}, + failed: map[string]int{}, + } +} + +func (r *recordingPublisher) Publish(channel string, payload []byte) error { + r.calls = append(r.calls, recordedCall{Channel: channel, Payload: append([]byte(nil), payload...)}) + if err, ok := r.errOn[channel]; ok { + r.failed[channel]++ + return err + } + return nil +} + +func TestPublishBuffer_DefersPublishUntilFlush(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + buf := NewPublishBuffer(inner) + + require.NoError(t, buf.Publish("a", []byte("1"))) + require.NoError(t, buf.Publish("b", []byte("2"))) + + require.Empty(t, inner.calls, "inner publisher should not be called before flush") + require.Equal(t, []string{"a", "b"}, buf.BufferedChannels()) +} + +func TestPublishBuffer_FlushPublishesInOrder(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + buf := NewPublishBuffer(inner) + + require.NoError(t, buf.Publish("a", []byte("1"))) + require.NoError(t, buf.Publish("b", []byte("2"))) + require.NoError(t, buf.Publish("c", []byte("3"))) + + require.NoError(t, buf.Flush()) + require.Len(t, inner.calls, 3) + require.Equal(t, "a", inner.calls[0].Channel) + require.Equal(t, "b", inner.calls[1].Channel) + require.Equal(t, "c", inner.calls[2].Channel) + require.Equal(t, []byte("1"), inner.calls[0].Payload) +} + +func TestPublishBuffer_FlushReturnsJoinedErrors(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + errB := xerrors.New("broken b") + errC := xerrors.New("broken c") + inner.errOn["b"] = errB + inner.errOn["c"] = errC + buf := NewPublishBuffer(inner) + + require.NoError(t, buf.Publish("a", []byte("1"))) + require.NoError(t, buf.Publish("b", []byte("2"))) + require.NoError(t, buf.Publish("c", []byte("3"))) + require.NoError(t, buf.Publish("d", []byte("4"))) + + err := buf.Flush() + require.Error(t, err) + require.ErrorIs(t, err, errB) + require.ErrorIs(t, err, errC) + require.Contains(t, err.Error(), "publish b:") + require.Contains(t, err.Error(), "publish c:") + // Even after broken channels, later messages should still be + // attempted so the inner publisher sees them. + require.Len(t, inner.calls, 4) +} + +func TestPublishBuffer_PublishAfterFlushFails(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + buf := NewPublishBuffer(inner) + require.NoError(t, buf.Flush()) + require.Error(t, buf.Publish("x", []byte("y"))) +} + +func TestPublishBuffer_DiscardSuppressesPending(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + buf := NewPublishBuffer(inner) + require.NoError(t, buf.Publish("a", []byte("1"))) + buf.Discard() + require.NoError(t, buf.Flush()) + require.Empty(t, inner.calls) +} + +func TestPublishBuffer_DiscardBlocksLaterPublishes(t *testing.T) { + t.Parallel() + inner := newRecordingPublisher() + buf := NewPublishBuffer(inner) + buf.Discard() + // Discard sets disabled; subsequent Publish is a no-op (not an + // error) so callers using Discard before/around rollback paths + // do not have to special-case unwind. + require.NoError(t, buf.Publish("a", []byte("1"))) + require.NoError(t, buf.Flush()) + require.Empty(t, inner.calls) +} diff --git a/coderd/x/chatd/chatstate/state.go b/coderd/x/chatd/chatstate/state.go new file mode 100644 index 0000000000..4d39d27329 --- /dev/null +++ b/coderd/x/chatd/chatstate/state.go @@ -0,0 +1,167 @@ +package chatstate + +import ( + "github.com/coder/coder/v2/coderd/database" +) + +// ExecutionState identifies a chat's current execution state. Values +// outside the chat execution state model are represented by +// [StateInvalid]. +type ExecutionState string + +const ( + // StateN: chat does not exist. + StateN ExecutionState = "N" + // StateW: waiting, empty queue, not archived. + StateW ExecutionState = "W" + // StateE0: error, empty queue, not archived. + StateE0 ExecutionState = "E0" + // StateE1: error, non-empty queue, not archived. + StateE1 ExecutionState = "E1" + // StateR0: running, empty queue, not archived. + StateR0 ExecutionState = "R0" + // StateR1: running, non-empty queue, not archived. + StateR1 ExecutionState = "R1" + // StateI0: interrupting, empty queue, not archived. + StateI0 ExecutionState = "I0" + // StateI1: interrupting, non-empty queue, not archived. + StateI1 ExecutionState = "I1" + // StateA0: requires_action, empty queue, not archived. + StateA0 ExecutionState = "A0" + // StateA1: requires_action, non-empty queue, not archived. + StateA1 ExecutionState = "A1" + // StateXW: archived waiting, empty queue. + StateXW ExecutionState = "XW" + // StateXE0: archived error, empty queue. + StateXE0 ExecutionState = "XE0" + // StateXE1: archived error, non-empty queue. + StateXE1 ExecutionState = "XE1" + + // StateInvalid groups every status/archive/queue combination that + // is not one of the valid states above. The state machine refuses + // non-reconciliation transitions on invalid states and exposes the + // [Tx.ReconcileInvalidState] transition to recover. + StateInvalid ExecutionState = "Invalid" +) + +// String implements fmt.Stringer. +func (s ExecutionState) String() string { return string(s) } + +// AllExecutionStates is the canonical enumeration of every value the +// classifier can return. Tests rely on this list to iterate over every +// state when verifying transition coverage. +var AllExecutionStates = []ExecutionState{ + StateN, + StateW, + StateE0, + StateE1, + StateR0, + StateR1, + StateI0, + StateI1, + StateA0, + StateA1, + StateXW, + StateXE0, + StateXE1, + StateInvalid, +} + +// IsRunnable returns true for the execution states that the chat +// worker is allowed to acquire and drive forward: R0, R1, I0, I1, +// A0, and A1. Requires-action states need worker ownership for +// timeout processing. Other states are idle (W, E*, XW, XE*), absent +// (N), or invalid. +func (s ExecutionState) IsRunnable() bool { + switch s { + case StateR0, StateR1, StateI0, StateI1, StateA0, StateA1: + return true + default: + return false + } +} + +// IsArchived returns true for the three archived execution states. +func (s ExecutionState) IsArchived() bool { + switch s { + case StateXW, StateXE0, StateXE1: + return true + default: + return false + } +} + +// QueueNonEmpty returns true for execution states that require a +// non-empty queue. Useful when seeding test fixtures. +func (s ExecutionState) QueueNonEmpty() bool { + switch s { + case StateE1, StateR1, StateI1, StateA1, StateXE1: + return true + default: + return false + } +} + +// ClassifyExecutionState turns the chat row, queue cardinality, and +// whether the chat row exists into an [ExecutionState]. The caller is +// responsible for loading the chat under the row lock and reading the +// queue count in the same transaction. +// +// Callers that have no chat row (lookup returned sql.ErrNoRows) +// should pass exists=false; the chat, status, and archive arguments +// are then ignored. +// +// The classifier is a single flat switch over the valid (status, +// archived, queue) tuples in the chat execution state model. Anything +// outside that set (legacy pending/paused/completed statuses, archived +// busy states, waiting with a non-empty queue, future enum values) +// falls through to [StateInvalid]. +// +//nolint:revive // queueNonEmpty/exists are simple classifier inputs. +func ClassifyExecutionState(chat database.Chat, queueNonEmpty, exists bool) ExecutionState { + if !exists { + return StateN + } + switch { + case chat.Status == database.ChatStatusWaiting && !chat.Archived && !queueNonEmpty: + return StateW + case chat.Status == database.ChatStatusWaiting && chat.Archived && !queueNonEmpty: + return StateXW + case chat.Status == database.ChatStatusError && !chat.Archived && !queueNonEmpty: + return StateE0 + case chat.Status == database.ChatStatusError && !chat.Archived && queueNonEmpty: + return StateE1 + case chat.Status == database.ChatStatusError && chat.Archived && !queueNonEmpty: + return StateXE0 + case chat.Status == database.ChatStatusError && chat.Archived && queueNonEmpty: + return StateXE1 + case chat.Status == database.ChatStatusRunning && !chat.Archived && !queueNonEmpty: + return StateR0 + case chat.Status == database.ChatStatusRunning && !chat.Archived && queueNonEmpty: + return StateR1 + case chat.Status == database.ChatStatusInterrupting && !chat.Archived && !queueNonEmpty: + return StateI0 + case chat.Status == database.ChatStatusInterrupting && !chat.Archived && queueNonEmpty: + return StateI1 + case chat.Status == database.ChatStatusRequiresAction && !chat.Archived && !queueNonEmpty: + return StateA0 + case chat.Status == database.ChatStatusRequiresAction && !chat.Archived && queueNonEmpty: + return StateA1 + } + return StateInvalid +} + +// OwnershipState identifies whether a chat row is currently owned by a +// worker. The state machine treats execution and ownership as +// orthogonal. +type OwnershipState string + +const ( + // StateU: chat has no owner (worker_id IS NULL). + StateU OwnershipState = "U" + // StateO: chat has an owner (worker_id IS NOT NULL). + StateO OwnershipState = "O" +) + +// AllOwnershipStates is the canonical enumeration of ownership states. +var AllOwnershipStates = []OwnershipState{StateU, StateO} diff --git a/coderd/x/chatd/chatstate/state_internal_test.go b/coderd/x/chatd/chatstate/state_internal_test.go new file mode 100644 index 0000000000..3fe6318921 --- /dev/null +++ b/coderd/x/chatd/chatstate/state_internal_test.go @@ -0,0 +1,163 @@ +package chatstate + +import ( + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" +) + +func chatWithStatus(status database.ChatStatus, archived bool) database.Chat { + return database.Chat{ + ID: uuid.New(), + Status: status, + Archived: archived, + OwnerID: uuid.New(), + } +} + +// TestClassifyExecutionState_Valid covers every valid classification: +// N (missing chat) plus every valid existing-chat state. +func TestClassifyExecutionState_Valid(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status database.ChatStatus + archived bool + queueNonEmpty bool + exists bool + want ExecutionState + }{ + {name: "N", exists: false, want: StateN}, + {name: "W", status: database.ChatStatusWaiting, exists: true, want: StateW}, + {name: "E0", status: database.ChatStatusError, exists: true, want: StateE0}, + {name: "E1", status: database.ChatStatusError, queueNonEmpty: true, exists: true, want: StateE1}, + {name: "R0", status: database.ChatStatusRunning, exists: true, want: StateR0}, + {name: "R1", status: database.ChatStatusRunning, queueNonEmpty: true, exists: true, want: StateR1}, + {name: "I0", status: database.ChatStatusInterrupting, exists: true, want: StateI0}, + {name: "I1", status: database.ChatStatusInterrupting, queueNonEmpty: true, exists: true, want: StateI1}, + {name: "A0", status: database.ChatStatusRequiresAction, exists: true, want: StateA0}, + {name: "A1", status: database.ChatStatusRequiresAction, queueNonEmpty: true, exists: true, want: StateA1}, + {name: "XW", status: database.ChatStatusWaiting, archived: true, exists: true, want: StateXW}, + {name: "XE0", status: database.ChatStatusError, archived: true, exists: true, want: StateXE0}, + {name: "XE1", status: database.ChatStatusError, archived: true, queueNonEmpty: true, exists: true, want: StateXE1}, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + chat := database.Chat{} + if tc.exists { + chat = chatWithStatus(tc.status, tc.archived) + } + require.Equal(t, tc.want, ClassifyExecutionState(chat, tc.queueNonEmpty, tc.exists)) + }) + } +} + +// TestClassifyExecutionState_Invalid covers every documented invalid +// combination: legacy statuses, waiting-with-queue, and archived busy +// statuses. +func TestClassifyExecutionState_Invalid(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status database.ChatStatus + archived bool + queueNonEmpty bool + }{ + // Legacy statuses (pending/paused/completed) are invalid for + // the new state machine. + {name: "LegacyPending", status: "pending"}, + {name: "LegacyPaused", status: "paused"}, + {name: "LegacyCompleted", status: "completed"}, + + // Waiting must always have an empty queue. + {name: "WaitingWithQueue", status: database.ChatStatusWaiting, queueNonEmpty: true}, + {name: "WaitingArchivedWithQueue", status: database.ChatStatusWaiting, archived: true, queueNonEmpty: true}, + + // Archived busy statuses are invalid. + {name: "ArchivedRunning", status: database.ChatStatusRunning, archived: true}, + {name: "ArchivedInterrupting", status: database.ChatStatusInterrupting, archived: true}, + {name: "ArchivedRequiresAction", status: database.ChatStatusRequiresAction, archived: true}, + } + for _, tc := range cases { + tc := tc + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + got := ClassifyExecutionState(chatWithStatus(tc.status, tc.archived), tc.queueNonEmpty, true) + require.Equal(t, StateInvalid, got) + }) + } +} + +// TestClassifyExecutionState_RejectsAllUnlistedCombinations enumerates +// every (status, archived, queueNonEmpty) tuple for an existing chat +// and asserts exactly the expected valid tuples classify out of +// [StateInvalid]. Missing chats are handled separately via the N case +// in [TestClassifyExecutionState_Valid]. +func TestClassifyExecutionState_RejectsAllUnlistedCombinations(t *testing.T) { + t.Parallel() + allStatuses := []database.ChatStatus{ + database.ChatStatusWaiting, + database.ChatStatusError, + database.ChatStatusRunning, + database.ChatStatusInterrupting, + database.ChatStatusRequiresAction, + "pending", "paused", "completed", + } + validCount := 0 + for _, status := range allStatuses { + for _, archived := range []bool{false, true} { + for _, queueNonEmpty := range []bool{false, true} { + got := ClassifyExecutionState(chatWithStatus(status, archived), queueNonEmpty, true) + if got != StateInvalid { + validCount++ + } + } + } + } + wantValid := len(AllExecutionStates) - 2 // Exclude StateN and StateInvalid. + require.Equal(t, wantValid, validCount, "valid existing-chat (status, archived, queue) tuples") +} + +// TestAllExecutionStates_Enumeration verifies AllExecutionStates +// contains every declared execution state exactly once. +func TestAllExecutionStates_Enumeration(t *testing.T) { + t.Parallel() + want := map[ExecutionState]bool{ + StateN: true, StateW: true, StateE0: true, StateE1: true, + StateR0: true, StateR1: true, StateI0: true, StateI1: true, + StateA0: true, StateA1: true, StateXW: true, StateXE0: true, + StateXE1: true, StateInvalid: true, + } + require.Len(t, AllExecutionStates, len(want)) + seen := make(map[ExecutionState]bool, len(want)) + for _, s := range AllExecutionStates { + require.True(t, want[s], "unexpected state %s", s) + require.False(t, seen[s], "duplicate state %s", s) + seen[s] = true + } +} + +// TestExecutionState_Predicates covers IsRunnable and QueueNonEmpty +// for every declared execution state. +func TestExecutionState_Predicates(t *testing.T) { + t.Parallel() + + runnable := map[ExecutionState]bool{ + StateR0: true, StateR1: true, StateI0: true, StateI1: true, + StateA0: true, StateA1: true, + } + nonEmpty := map[ExecutionState]bool{ + StateE1: true, StateR1: true, StateI1: true, StateA1: true, StateXE1: true, + } + for _, s := range AllExecutionStates { + require.Equal(t, runnable[s], s.IsRunnable(), "IsRunnable(%s)", s) + require.Equal(t, nonEmpty[s], s.QueueNonEmpty(), "QueueNonEmpty(%s)", s) + } +} diff --git a/coderd/x/chatd/chatstate/synthetic_cancellation_test.go b/coderd/x/chatd/chatstate/synthetic_cancellation_test.go new file mode 100644 index 0000000000..75880aa2f7 --- /dev/null +++ b/coderd/x/chatd/chatstate/synthetic_cancellation_test.go @@ -0,0 +1,517 @@ +package chatstate_test + +import ( + "context" + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// nonDynamicAssistantToolCallMessage builds an assistant message that +// issues a single tool call against a tool that is NOT in the chat's +// dynamic_tools set. The send-message and edit-message paths use the +// "cancel every outstanding tool call regardless of source" variant +// (dynamicOnly=false), so the cancellation must still fire even for +// non-dynamic tools. +func nonDynamicAssistantToolCallMessage(t *testing.T, modelID uuid.UUID, callID string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: callID, + ToolName: "non_dynamic_tool", + Args: json.RawMessage(`{}`), + }}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + } +} + +// assertToolResultForCall asserts that msg is a tool-result message +// that resolves a tool call with id wantCallID and is_error=true. +func assertToolResultForCall(t *testing.T, msg database.ChatMessage, wantCallID string) { + t.Helper() + require.Equal(t, database.ChatMessageRoleTool, msg.Role) + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + require.NotEmpty(t, parts) + var found bool + for _, p := range parts { + if p.Type != codersdk.ChatMessagePartTypeToolResult { + continue + } + require.Equal(t, wantCallID, p.ToolCallID, "tool-call id matches") + require.True(t, p.IsError, "synthetic cancellation must be marked is_error=true") + found = true + } + require.True(t, found, "expected at least one tool-result part") +} + +// commitAssistantToolCall pushes an assistant message that calls +// `tool_name` with `callID` into history via CommitStep. Returns the +// inserted assistant ChatMessage. Use the dynamic-tools chat fixture +// (createTestChatWithDynamicTools) when dynamicOnly cancellation +// paths are exercised. +func commitAssistantToolCall( + t *testing.T, + f *testFixture, + m *chatstate.ChatMachine, + msg chatstate.Message, +) database.ChatMessage { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{msg}}) + return err + })) + require.Len(t, step.InsertedMessages, 1) + return step.InsertedMessages[0] +} + +// landInW puts a fresh R0 chat into state W (waiting) via FinishTurn. +func landInW(t *testing.T, f *testFixture, m *chatstate.ChatMachine) { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + require.Equal(t, chatstate.StateW, f.classify(ctx, t, m.ChatID())) +} + +// landInE0 puts a fresh R0 chat into state E0 (error, empty queue). +func landInE0(t *testing.T, f *testFixture, m *chatstate.ChatMachine) { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + require.Equal(t, chatstate.StateE0, f.classify(ctx, t, m.ChatID())) +} + +func TestSyntheticCancellation_SendMessageDirect(t *testing.T) { + t.Parallel() + + t.Run("waiting", func(t *testing.T) { + t.Parallel() + testSendMessageDirectWSynthesizesToolCancellations(t) + }) + t.Run("error", func(t *testing.T) { + t.Parallel() + testSendMessageDirectE0SynthesizesToolCancellations(t) + }) +} + +// testSendMessageDirectWSynthesizesToolCancellations verifies that +// from W, SendMessage inserts synthetic tool-result rows for every +// outstanding tool call on the last assistant message BEFORE the new +// user message, regardless of whether the tools are dynamic. +func testSendMessageDirectWSynthesizesToolCancellations(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + callID := "call_" + uuid.NewString() + assistant := commitAssistantToolCall(t, f, m, + nonDynamicAssistantToolCallMessage(t, f.Model.ID, callID)) + require.Equal(t, database.ChatMessageRoleAssistant, assistant.Role) + + // R0 -> W. + landInW(t, f, m) + + // SendMessage with a fresh user message. The direct-history path + // must insert a synthetic tool-result (for callID) followed by + // the new user message. + var send chatstate.SendMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + send, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("after-cancel", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + + require.Len(t, send.InsertedMessages, 2, "synthetic cancel + new user") + assertToolResultForCall(t, send.InsertedMessages[0], callID) + require.Equal(t, database.ChatMessageRoleUser, send.InsertedMessages[1].Role) + require.Less(t, send.InsertedMessages[0].ID, send.InsertedMessages[1].ID, + "synthetic cancel is inserted before the user message") +} + +// testSendMessageDirectE0SynthesizesToolCancellations exercises +// the same path from E0. +func testSendMessageDirectE0SynthesizesToolCancellations(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + callID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + nonDynamicAssistantToolCallMessage(t, f.Model.ID, callID)) + + // R0 -> E0. + landInE0(t, f, m) + + var send chatstate.SendMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + send, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("after-error", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + + require.Len(t, send.InsertedMessages, 2) + assertToolResultForCall(t, send.InsertedMessages[0], callID) + require.Equal(t, database.ChatMessageRoleUser, send.InsertedMessages[1].Role) +} + +func TestSyntheticCancellation_EditMessage(t *testing.T) { + t.Parallel() + + t.Run("replacement insertion", func(t *testing.T) { + t.Parallel() + testEditMessageSynthesizesToolCancellationsBeforeReplacement(t) + }) +} + +// testEditMessageSynthesizesToolCancellationsBeforeReplacement +// verifies that EditMessage from a state with an outstanding tool +// call before the edited user message inserts a synthetic +// tool-result before the replacement user message in history. +// +// The scenario is: +// - user message 1 (initial) +// - assistant tool-call (outstanding) +// - user message 2 (the one we will edit) +// +// EditMessage soft-deletes user message 2 and everything after it, +// then synthesizes cancellations for tool calls on the last +// surviving assistant message that have no matching tool-result. +func testEditMessageSynthesizesToolCancellationsBeforeReplacement(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // Build the history described above. CommitStep is happy to insert + // a mixed batch as long as it stays inside R0. + callID := "call_" + uuid.NewString() + assistantTC := nonDynamicAssistantToolCallMessage(t, f.Model.ID, callID) + secondUser := userTextMessage("second user", f.User.ID, f.Model.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistantTC, secondUser}, + }) + return err + })) + require.Len(t, step.InsertedMessages, 2) + secondUserID := step.InsertedMessages[1].ID + require.Equal(t, database.ChatMessageRoleUser, step.InsertedMessages[1].Role) + + var edit chatstate.EditMessageResult + editedContent := mustMarshalParts(t, []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("edited"), + }) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + edit, err = tx.EditMessage(chatstate.EditMessageInput{ + MessageID: secondUserID, + CreatedBy: f.User.ID, + Content: editedContent, + APIKeyID: f.apiKeyID(), + }) + return err + })) + + require.Len(t, edit.CancellationMessages, 1, "synthetic cancel inserted") + assertToolResultForCall(t, edit.CancellationMessages[0], callID) + require.Equal(t, database.ChatMessageRoleUser, edit.ReplacementMessage.Role) + require.Less(t, edit.CancellationMessages[0].ID, edit.ReplacementMessage.ID, + "cancellations are inserted before the replacement user message") +} + +func TestSyntheticCancellation_PromoteQueuedMessage(t *testing.T) { + t.Parallel() + + t.Run("error queued message", func(t *testing.T) { + t.Parallel() + testPromoteQueuedMessageE1SynthesizesToolCancellations(t) + }) + t.Run("requires action queued message", func(t *testing.T) { + t.Parallel() + testPromoteQueuedMessageA1SynthesizesDynamicToolCancellations(t) + }) +} + +// testPromoteQueuedMessageE1SynthesizesToolCancellations verifies +// that promoting a queued message from E1 inserts synthetic +// tool-result rows for outstanding tool calls before the promoted +// user message in history. +func testPromoteQueuedMessageE1SynthesizesToolCancellations(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + callID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + nonDynamicAssistantToolCallMessage(t, f.Model.ID, callID)) + + // Land in R1 with one queued message. + queued := sendQueuedMessage(t, f, m, "queued-for-promote") + require.NotNil(t, queued.QueuedMessage) + require.Equal(t, chatstate.StateR1, f.classify(ctx, t, created.Chat.ID)) + + // R1 -> E1 via FinishError. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + require.Equal(t, chatstate.StateE1, f.classify(ctx, t, created.Chat.ID)) + + var promote chatstate.PromoteQueuedMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + promote, err = tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: queued.QueuedMessage.ID, + }) + return err + })) + + require.Len(t, promote.CancellationMessages, 1) + assertToolResultForCall(t, promote.CancellationMessages[0], callID) + require.NotNil(t, promote.InsertedMessage) + require.Equal(t, database.ChatMessageRoleUser, promote.InsertedMessage.Role) + require.Less(t, promote.CancellationMessages[0].ID, promote.InsertedMessage.ID, + "cancel is inserted before the promoted user message") +} + +// testPromoteQueuedMessageA1SynthesizesDynamicToolCancellations +// verifies that the dynamic outstanding tool call is canceled when +// promoting from A1. +func testPromoteQueuedMessageA1SynthesizesDynamicToolCancellations(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + + toolName := "dyn_promote_a1" + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + dynCallID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + assistantToolCallMessage(t, f.Model.ID, toolName, dynCallID)) + + // Land in A0. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + require.Equal(t, chatstate.StateA0, f.classify(ctx, t, created.Chat.ID)) + + // A0 -> A1 with one queued user message. + queued := sendQueuedMessage(t, f, m, "queued-for-a1-promote") + require.NotNil(t, queued.QueuedMessage) + require.Equal(t, chatstate.StateA1, f.classify(ctx, t, created.Chat.ID)) + + var promote chatstate.PromoteQueuedMessageResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + promote, err = tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: queued.QueuedMessage.ID, + }) + return err + })) + + require.Len(t, promote.CancellationMessages, 1, "dynamic tool call canceled") + assertToolResultForCall(t, promote.CancellationMessages[0], dynCallID) + require.NotNil(t, promote.InsertedMessage) + require.Equal(t, database.ChatMessageRoleUser, promote.InsertedMessage.Role) +} + +func TestSyntheticCancellation_FinishTurn(t *testing.T) { + t.Parallel() + + t.Run("running queued message", func(t *testing.T) { + t.Parallel() + testFinishTurnR1SynthesizesToolCancellationsBeforePromotion(t) + }) +} + +// testFinishTurnR1SynthesizesToolCancellationsBeforePromotion +// verifies that finishing a turn while a queued message exists +// synthesizes outstanding tool cancellations before promoting the +// queue head into history. +func testFinishTurnR1SynthesizesToolCancellationsBeforePromotion(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + callID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + nonDynamicAssistantToolCallMessage(t, f.Model.ID, callID)) + + queued := sendQueuedMessage(t, f, m, "queued-for-finish") + require.NotNil(t, queued.QueuedMessage) + require.Equal(t, chatstate.StateR1, f.classify(ctx, t, created.Chat.ID)) + + beforeIDs := historyMessageIDs(ctx, t, f, created.Chat.ID) + + var finish chatstate.FinishTurnResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + finish, err = tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + require.NotNil(t, finish.PromotedMessage) + require.Equal(t, database.ChatMessageRoleUser, finish.PromotedMessage.Role) + + afterIDs := historyMessageIDs(ctx, t, f, created.Chat.ID) + require.Equal(t, len(beforeIDs)+2, len(afterIDs), + "finish inserts both a tool cancel and the promoted user") + + // The two newly inserted messages are tool-result then user. + newIDs := afterIDs[len(beforeIDs):] + cancel, err := f.DB.GetChatMessageByID(ctx, newIDs[0]) + require.NoError(t, err) + assertToolResultForCall(t, cancel, callID) + require.Equal(t, finish.PromotedMessage.ID, newIDs[1]) +} + +func TestSyntheticCancellation_FinishInterruption(t *testing.T) { + t.Parallel() + + t.Run("interrupting queued message", func(t *testing.T) { + t.Parallel() + testFinishInterruptionI1PromotesQueueHead(t) + }) + t.Run("rejects outstanding dynamic tool calls", func(t *testing.T) { + t.Parallel() + testFinishInterruptionRejectsOutstandingToolCalls(t) + }) +} + +// testFinishInterruptionI1PromotesQueueHead verifies that +// FinishInterruption from I1 with no outstanding tool calls +// promotes the queue head into history. +func testFinishInterruptionI1PromotesQueueHead(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // Reach R1 with one queued message. + queued := sendQueuedMessage(t, f, m, "queued-for-interruption") + require.NotNil(t, queued.QueuedMessage) + // R1 -> I1 via Interrupt. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Interrupt(chatstate.InterruptInput{Reason: "test"}) + return err + })) + require.Equal(t, chatstate.StateI1, f.classify(ctx, t, created.Chat.ID)) + + beforeIDs := historyMessageIDs(ctx, t, f, created.Chat.ID) + + var finish chatstate.FinishInterruptionResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + finish, err = tx.FinishInterruption(chatstate.FinishInterruptionInput{}) + return err + })) + require.NotNil(t, finish.PromotedMessage) + require.Equal(t, database.ChatMessageRoleUser, finish.PromotedMessage.Role) + + afterIDs := historyMessageIDs(ctx, t, f, created.Chat.ID) + require.Equal(t, len(beforeIDs)+1, len(afterIDs)) + require.Equal(t, chatstate.StateR0, f.classify(ctx, t, created.Chat.ID)) +} + +// testFinishInterruptionRejectsOutstandingToolCalls verifies that +// FinishInterruption fails (TransitionNotAllowed-shaped) when the +// chat still has an outstanding dynamic tool call after the partial +// commit. The chat must remain in its prior state. +func testFinishInterruptionRejectsOutstandingToolCalls(t *testing.T) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + + toolName := "dyn_finish_reject" + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + dynCallID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + assistantToolCallMessage(t, f.Model.ID, toolName, dynCallID)) + + // R0 -> I0 via Interrupt. Interrupt closes pending dynamic calls + // when transitioning from A0/A1, but from R0 it does NOT, so the + // chat keeps its outstanding dynamic call. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Interrupt(chatstate.InterruptInput{Reason: "test"}) + return err + })) + require.Equal(t, chatstate.StateI0, f.classify(ctx, t, created.Chat.ID)) + + stateBefore := f.classify(ctx, t, created.Chat.ID) + historyBefore := historyMessageIDs(ctx, t, f, created.Chat.ID) + publishedBefore := len(f.Pub.channels) + + // FinishInterruption with no partial commits should reject + // because the dynamic call is still outstanding. + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishInterruption(chatstate.FinishInterruptionInput{}) + return err + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + + require.Equal(t, stateBefore, f.classify(ctx, t, created.Chat.ID), "state unchanged") + require.Equal(t, historyBefore, historyMessageIDs(ctx, t, f, created.Chat.ID), + "history unchanged on rejected finish") + require.Equal(t, publishedBefore, len(f.Pub.channels), + "failed FinishInterruption publishes nothing") +} + +// ensure unused imports don't break the build if any helper is +// removed later. +var _ = context.Background diff --git a/coderd/x/chatd/chatstate/synthetics.go b/coderd/x/chatd/chatstate/synthetics.go new file mode 100644 index 0000000000..d442843d17 --- /dev/null +++ b/coderd/x/chatd/chatstate/synthetics.go @@ -0,0 +1,262 @@ +package chatstate + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" +) + +// synthesizePendingToolCancellations builds [Message] inserts that +// satisfy every outstanding tool call on the chat's last assistant +// message with a synthetic cancellation tool-result message. +// +// "Outstanding" means a tool call present on the last assistant +// message that does not yet have a matching tool-result message in +// the active history after it. The caller controls whether to limit +// to dynamic-tool calls (true) or close every outstanding tool call +// regardless of source (false). The dynamic-only variant is used by +// requires-action interrupts; the all-tool variant is used by any +// transition that needs to insert a new user message into history. +// +// The synthetic results use the supplied chat's last_model_config_id. +// Returns (nil, nil) when there is nothing to synthesize. +// +//nolint:revive // dynamicOnly is a domain flag, not a control flag. +func synthesizePendingToolCancellations( + ctx context.Context, + store database.Store, + chat database.Chat, + reason string, + dynamicOnly bool, +) ([]Message, error) { + var dynamicToolNames map[string]bool + if dynamicOnly { + var err error + dynamicToolNames, err = parseDynamicToolNamesFromRaw(chat.DynamicTools) + if err != nil { + return nil, xerrors.Errorf("parse dynamic tool names: %w", err) + } + if len(dynamicToolNames) == 0 { + return nil, nil + } + } + + lastAssistant, err := store.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + }) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return nil, nil + } + return nil, xerrors.Errorf("get last assistant message: %w", err) + } + assistantParts, err := chatprompt.ParseContent(lastAssistant) + if err != nil { + return nil, xerrors.Errorf("parse assistant message: %w", err) + } + afterMsgs, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: lastAssistant.ID, + }) + if err != nil { + return nil, xerrors.Errorf("get messages after assistant: %w", err) + } + handled := make(map[string]bool) + // Provider-executed tool results (e.g. web_search) are persisted + // inside the assistant message itself, not as tool-role messages + // after it. Count them as handled so their calls are not treated + // as outstanding. + for _, p := range assistantParts { + if p.Type == codersdk.ChatMessagePartTypeToolResult { + handled[p.ToolCallID] = true + } + } + for _, msg := range afterMsgs { + if msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + // Don't fail the whole cancellation just because one + // historical message is unparsable; treat its tool + // results as unknown. + continue + } + for _, p := range parts { + if p.Type == codersdk.ChatMessagePartTypeToolResult { + handled[p.ToolCallID] = true + } + } + } + out := make([]Message, 0) + for _, part := range assistantParts { + if part.Type != codersdk.ChatMessagePartTypeToolCall { + continue + } + // Provider-executed tool calls are handled server-side by the + // LLM provider. A synthetic client tool-result for them is + // invalid replay history: Anthropic rejects a plain tool_result + // block that references a server_tool_use ID. + if part.ProviderExecuted { + continue + } + if dynamicOnly && !dynamicToolNames[part.ToolName] { + continue + } + if handled[part.ToolCallID] { + continue + } + resultPart := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: part.ToolCallID, + ToolName: part.ToolName, + Result: json.RawMessage(fmt.Sprintf("%q", reason)), + IsError: true, + } + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{resultPart}) + if err != nil { + return nil, xerrors.Errorf("marshal synthetic tool result: %w", err) + } + out = append(out, Message{ + Role: database.ChatMessageRoleTool, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + }) + } + if len(out) == 0 { + return nil, nil + } + return out, nil +} + +// pendingDynamicToolCallIDs returns the dynamic tool-call IDs on the +// chat's last assistant message that do not yet have a matching +// tool-result message in active history. The returned map is keyed by +// tool-call ID and valued by tool name so callers can build matching +// result messages without re-parsing the assistant content. +func pendingDynamicToolCallIDs(ctx context.Context, store database.Store, chat database.Chat) (map[string]string, error) { + dynamic, err := parseDynamicToolNamesFromRaw(chat.DynamicTools) + if err != nil { + return nil, err + } + if len(dynamic) == 0 { + return map[string]string{}, nil + } + return outstandingToolCallIDs(ctx, store, chat, func(toolName string) bool { + return dynamic[toolName] + }) +} + +// pendingAllToolCallIDs returns the tool-call IDs of every outstanding +// tool call on the chat's last assistant message, regardless of +// whether the tool is dynamic. The returned map is keyed by tool-call +// ID and valued by tool name. Callers that must guarantee a valid +// LLM message history (e.g. before promoting a user message into +// active history, or after committing an interruption's partial +// messages) should use this variant so non-dynamic tool calls do not +// silently bypass the check. +func pendingAllToolCallIDs(ctx context.Context, store database.Store, chat database.Chat) (map[string]string, error) { + return outstandingToolCallIDs(ctx, store, chat, func(string) bool { return true }) +} + +// outstandingToolCallIDs walks the chat's last assistant message and +// returns the subset of its tool calls that have no matching +// tool-result message in the active history after it. The accept +// callback can be used to restrict the walk to a subset of tools +// (e.g. dynamic-only). +func outstandingToolCallIDs(ctx context.Context, store database.Store, chat database.Chat, accept func(toolName string) bool) (map[string]string, error) { + lastAssistant, err := store.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chat.ID, + Role: database.ChatMessageRoleAssistant, + }) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return map[string]string{}, nil + } + return nil, xerrors.Errorf("get last assistant: %w", err) + } + parts, err := chatprompt.ParseContent(lastAssistant) + if err != nil { + return nil, xerrors.Errorf("parse assistant: %w", err) + } + afterMsgs, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chat.ID, + AfterID: lastAssistant.ID, + }) + if err != nil { + return nil, xerrors.Errorf("get messages after assistant: %w", err) + } + handled := make(map[string]bool) + // Provider-executed tool results are persisted inside the + // assistant message itself; count them as handled. + for _, p := range parts { + if p.Type == codersdk.ChatMessagePartTypeToolResult { + handled[p.ToolCallID] = true + } + } + for _, msg := range afterMsgs { + if msg.Role != database.ChatMessageRoleTool { + continue + } + messageParts, err := chatprompt.ParseContent(msg) + if err != nil { + continue + } + for _, p := range messageParts { + if p.Type == codersdk.ChatMessagePartTypeToolResult { + handled[p.ToolCallID] = true + } + } + } + out := make(map[string]string) + for _, p := range parts { + if p.Type != codersdk.ChatMessagePartTypeToolCall { + continue + } + // Provider-executed tool calls are answered server-side by the + // LLM provider and must never be reported as outstanding. + if p.ProviderExecuted { + continue + } + if !accept(p.ToolName) { + continue + } + if handled[p.ToolCallID] { + continue + } + out[p.ToolCallID] = p.ToolName + } + return out, nil +} + +// parseDynamicToolNamesFromRaw is a private mirror of +// chatd.parseDynamicToolNames so chatstate does not pull a runtime +// dependency on the chatd package. It accepts a nullable raw JSON +// blob and returns a name set. +func parseDynamicToolNamesFromRaw(raw pqtype.NullRawMessage) (map[string]bool, error) { + if !raw.Valid || len(raw.RawMessage) == 0 { + return map[string]bool{}, nil + } + var tools []codersdk.DynamicTool + if err := json.Unmarshal(raw.RawMessage, &tools); err != nil { + return nil, err + } + out := make(map[string]bool, len(tools)) + for _, t := range tools { + out[t.Name] = true + } + return out, nil +} diff --git a/coderd/x/chatd/chatstate/transition.go b/coderd/x/chatd/chatstate/transition.go new file mode 100644 index 0000000000..d96a91fef9 --- /dev/null +++ b/coderd/x/chatd/chatstate/transition.go @@ -0,0 +1,226 @@ +package chatstate + +import "slices" + +// Transition is the enumeration of transitions implemented by the +// state machine. Values intentionally match the names of the public +// methods on [Tx] (and [CreateChat]). The transition matrix below +// declares the legal (from -> to) execution-state mappings used by +// each transition method for validation. +type Transition string + +const ( + TransitionCreateChat Transition = "CreateChat" + TransitionSetArchived Transition = "SetArchived" + TransitionSendMessage Transition = "SendMessage" + TransitionEditMessage Transition = "EditMessage" + TransitionDeleteQueuedMessage Transition = "DeleteQueuedMessage" + TransitionPromoteQueuedMessage Transition = "PromoteQueuedMessage" + TransitionInterrupt Transition = "Interrupt" + TransitionCompleteRequiresAction Transition = "CompleteRequiresAction" + TransitionAcquire Transition = "Acquire" + TransitionAbandon Transition = "Abandon" + TransitionRecordGenerationAttempt Transition = "RecordGenerationAttempt" + TransitionRecordRetryState Transition = "RecordRetryState" + TransitionCommitStep Transition = "CommitStep" + TransitionEnterRequiresAction Transition = "EnterRequiresAction" + TransitionFinishInterruption Transition = "FinishInterruption" + TransitionFinishTurn Transition = "FinishTurn" + TransitionFinishError Transition = "FinishError" + TransitionCancelRequiresAction Transition = "CancelRequiresAction" + TransitionReconcileInvalidState Transition = "ReconcileInvalidState" +) + +// String implements fmt.Stringer. +func (t Transition) String() string { return string(t) } + +// AllExecutionTransitions is the canonical enumeration of every +// execution-state transition that has an entry in the matrix below. +// Ownership transitions (Acquire, Abandon) are intentionally not part +// of this slice because they are validated independently and do not +// have a (from->to) execution mapping. +var AllExecutionTransitions = []Transition{ + TransitionCreateChat, + TransitionSetArchived, + TransitionSendMessage, + TransitionEditMessage, + TransitionDeleteQueuedMessage, + TransitionPromoteQueuedMessage, + TransitionInterrupt, + TransitionCompleteRequiresAction, + TransitionRecordGenerationAttempt, + TransitionRecordRetryState, + TransitionCommitStep, + TransitionEnterRequiresAction, + TransitionFinishInterruption, + TransitionFinishTurn, + TransitionFinishError, + TransitionCancelRequiresAction, + TransitionReconcileInvalidState, +} + +// transitionMatrix is the in-code representation of the chat execution +// state transition table. Each entry maps an input state to the set of +// allowed transitions together with the possible classified output +// states that the transition implementation may land in. Outputs may +// depend on the post-mutation queue cardinality (for example +// DeleteQueuedMessage from E1 lands in E0 when the deleted row was the +// last queued message, or stays in E1 otherwise), which is why several +// entries list more than one output. +// +// Ownership transitions (Acquire, Abandon) are intentionally not +// included; they are orthogonal to execution state. +var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ + StateN: { + TransitionCreateChat: {StateR0}, + }, + StateW: { + TransitionSetArchived: {StateXW}, + TransitionSendMessage: {StateR0}, + TransitionEditMessage: {StateR0}, + }, + StateE0: { + TransitionSetArchived: {StateXE0}, + TransitionSendMessage: {StateR0}, + TransitionEditMessage: {StateR0}, + }, + StateE1: { + TransitionSetArchived: {StateXE1}, + TransitionSendMessage: {StateR1}, + TransitionEditMessage: {StateR0}, + TransitionDeleteQueuedMessage: {StateE0, StateE1}, + TransitionPromoteQueuedMessage: {StateR0, StateR1}, + }, + StateR0: { + TransitionSendMessage: {StateR1, StateI1}, + TransitionEditMessage: {StateR0}, + TransitionInterrupt: {StateI0}, + TransitionRecordGenerationAttempt: {StateR0}, + TransitionRecordRetryState: {StateR0}, + TransitionCommitStep: {StateR0}, + TransitionEnterRequiresAction: {StateA0}, + TransitionFinishTurn: {StateW}, + TransitionFinishError: {StateE0}, + }, + StateR1: { + TransitionSendMessage: {StateR1, StateI1}, + TransitionEditMessage: {StateR0}, + TransitionDeleteQueuedMessage: {StateR0, StateR1}, + TransitionPromoteQueuedMessage: {StateI1}, + TransitionInterrupt: {StateI1}, + TransitionRecordGenerationAttempt: {StateR1}, + TransitionRecordRetryState: {StateR1}, + TransitionCommitStep: {StateR1}, + TransitionEnterRequiresAction: {StateA1}, + TransitionFinishTurn: {StateR0, StateR1}, + TransitionFinishError: {StateE1}, + }, + StateI0: { + TransitionSendMessage: {StateI1}, + TransitionEditMessage: {StateR0}, + TransitionFinishInterruption: {StateW}, + }, + StateI1: { + TransitionSendMessage: {StateI1}, + TransitionEditMessage: {StateR0}, + TransitionDeleteQueuedMessage: {StateI0, StateI1}, + TransitionPromoteQueuedMessage: {StateI1}, + TransitionFinishInterruption: {StateR0, StateR1}, + }, + StateA0: { + TransitionSendMessage: {StateA1, StateR1}, + TransitionEditMessage: {StateR0}, + TransitionInterrupt: {StateR0}, + TransitionCompleteRequiresAction: {StateR0}, + TransitionCancelRequiresAction: {StateR0}, + }, + StateA1: { + TransitionSendMessage: {StateA1, StateR1}, + TransitionEditMessage: {StateR0}, + TransitionDeleteQueuedMessage: {StateA0, StateA1}, + TransitionPromoteQueuedMessage: {StateR0, StateR1}, + TransitionInterrupt: {StateR1}, + TransitionCompleteRequiresAction: {StateR1}, + TransitionCancelRequiresAction: {StateR1}, + }, + StateXW: { + TransitionSetArchived: {StateW}, + }, + StateXE0: { + TransitionSetArchived: {StateE0}, + }, + StateXE1: { + TransitionSetArchived: {StateE1}, + }, + StateInvalid: { + TransitionReconcileInvalidState: {StateE0, StateE1}, + }, +} + +// isExecutionTransitionAllowed reports whether a transition is legal +// from the supplied input state per the matrix above. Ownership +// transitions are not stored in the matrix and always return false. +func isExecutionTransitionAllowed(t Transition, from ExecutionState) bool { + allowed, ok := transitionMatrix[from] + if !ok { + return false + } + _, ok = allowed[t] + return ok +} + +// requireExecutionTransition validates that t is legal from `from` +// and returns a typed *TransitionError otherwise. +func requireExecutionTransition(t Transition, from ExecutionState) error { + if isExecutionTransitionAllowed(t, from) { + return nil + } + return newTransitionError(t, from, "") +} + +// AllowedExecutionTransitionsFrom returns a deterministic slice of +// transitions legal from `from`. Mostly used by tests to enumerate the +// matrix without leaking the internal map. +func AllowedExecutionTransitionsFrom(from ExecutionState) []Transition { + allowed := transitionMatrix[from] + out := make([]Transition, 0, len(allowed)) + for _, t := range AllExecutionTransitions { + if _, ok := allowed[t]; ok { + out = append(out, t) + } + } + return out +} + +// AllowedInputStates returns a deterministic slice of execution states +// from which `tr` is legal per the matrix above. Mostly used by tests +// to enumerate the matrix without leaking the internal map. +func AllowedInputStates(tr Transition) []ExecutionState { + var out []ExecutionState + for _, from := range AllExecutionStates { + if isExecutionTransitionAllowed(tr, from) { + out = append(out, from) + } + } + return out +} + +// AllowedExecutionTransitionOutputs returns the set of classified +// post-states that the transition `tr` may produce from `from` per +// the matrix above. The returned slice is a copy so callers may mutate +// it without affecting the underlying matrix. +// +// When `tr` is not allowed from `from`, an empty (nil) slice is +// returned. Tests use this helper to enumerate the (transition, from, +// want) triples that must be exercised by the row-level matrix tests. +func AllowedExecutionTransitionOutputs(from ExecutionState, tr Transition) []ExecutionState { + allowed, ok := transitionMatrix[from] + if !ok { + return nil + } + outputs, ok := allowed[tr] + if !ok { + return nil + } + return slices.Clone(outputs) +} diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go new file mode 100644 index 0000000000..964610f84d --- /dev/null +++ b/coderd/x/chatd/chatstate/transitions.go @@ -0,0 +1,1454 @@ +package chatstate + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "time" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" +) + +// CreateChatInput configures [CreateChat]. +type CreateChatInput struct { + OrganizationID uuid.UUID + OwnerID uuid.UUID + WorkspaceID uuid.NullUUID + BuildID uuid.NullUUID + AgentID uuid.NullUUID + ParentChatID uuid.NullUUID + RootChatID uuid.NullUUID + LastModelConfigID uuid.UUID + Title string + Mode database.NullChatMode + PlanMode database.NullChatPlanMode + MCPServerIDs []uuid.UUID + Labels pqtype.NullRawMessage + DynamicTools pqtype.NullRawMessage + ClientType database.ChatClientType + InitialMessages []Message + LastInjectedContext pqtype.NullRawMessage +} + +// CreateChatResult is the value returned by [CreateChat]. It carries +// the new chat row and the inserted initial history. +type CreateChatResult struct { + Chat database.Chat + InitialMessages []database.ChatMessage +} + +// CreateChat creates a brand new chat with initial history in a single +// transaction. It is package-level rather than a method on [ChatMachine] +// because no chat-scoped machine can exist before the chat row is written. +// +// Validation: +// - InitialMessages must be non-empty. +// +// After commit CreateChat publishes a `chat:update` message describing +// the new chat snapshot. Because the new chat has no worker assigned, +// CreateChat also publishes an ownership hint so workers can race to +// acquire the runnable chat. +func CreateChat( + ctx context.Context, + store database.Store, + publisher Publisher, + input CreateChatInput, +) (CreateChatResult, error) { + if store == nil { + return CreateChatResult{}, xerrors.New("chatstate: CreateChat called with nil store") + } + if publisher == nil { + return CreateChatResult{}, xerrors.New("chatstate: CreateChat called with nil publisher") + } + if len(input.InitialMessages) == 0 { + return CreateChatResult{}, newTransitionError( + TransitionCreateChat, StateN, + "initial messages must include at least one message", + ) + } + var result CreateChatResult + buffer := NewPublishBuffer(publisher) + defer buffer.Discard() + err := store.InTx(func(store database.Store) error { + chat, err := store.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: input.OrganizationID, + OwnerID: input.OwnerID, + WorkspaceID: input.WorkspaceID, + BuildID: input.BuildID, + AgentID: input.AgentID, + ParentChatID: input.ParentChatID, + RootChatID: input.RootChatID, + LastModelConfigID: input.LastModelConfigID, + Title: input.Title, + Mode: input.Mode, + PlanMode: input.PlanMode, + Status: database.ChatStatusRunning, + MCPServerIDs: input.MCPServerIDs, + Labels: input.Labels, + DynamicTools: input.DynamicTools, + ClientType: input.ClientType, + }) + if err != nil { + return xerrors.Errorf("insert chat: %w", err) + } + // Insert the initial history under the new chat row. The + // message revision trigger advances `history_version` to the + // current `snapshot_version` (which is 1 for a brand new chat). + inserted, err := store.InsertChatMessages(ctx, toInsertParams(chat.ID, input.InitialMessages)) + if err != nil { + return xerrors.Errorf("insert initial messages: %w", err) + } + if input.LastInjectedContext.Valid { + if _, err := store.UpdateChatLastInjectedContext(ctx, database.UpdateChatLastInjectedContextParams{ + ID: chat.ID, + LastInjectedContext: input.LastInjectedContext, + }); err != nil { + return xerrors.Errorf("set last injected context: %w", err) + } + } + refreshed, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("reload chat after initial messages: %w", err) + } + result = CreateChatResult{ + Chat: refreshed, + InitialMessages: inserted, + } + if err := buffer.Publish( + coderdpubsub.ChatStateUpdateChannel(refreshed.ID), + buildChatUpdateMessage(refreshed), + ); err != nil { + return xerrors.Errorf("buffer chat update: %w", err) + } + if ClassifyExecutionState(refreshed, false, true).IsRunnable() { + if err := buffer.Publish( + coderdpubsub.ChatStateOwnershipChannel, + buildChatOwnershipMessage(refreshed), + ); err != nil { + return xerrors.Errorf("buffer ownership hint: %w", err) + } + } + return nil + }, nil) + if err != nil { + return CreateChatResult{}, err + } + if err := buffer.Flush(); err != nil { + return result, err + } + return result, nil +} + +// applyExecutionStateUpdate is a small adapter so transition methods +// do not have to repeat the UpdateChatExecutionState boilerplate. +// The state machine writes status, archived, last_error, ownership +// identifiers, and the requires-action deadline as one atomic update. +type executionStateUpdate struct { + Status database.ChatStatus + Archived bool + WorkerID uuid.NullUUID + RunnerID uuid.NullUUID + LastError pqtype.NullRawMessage + RequiresActionDeadlineAt sql.NullTime +} + +func (tx *Tx) applyExecutionState(u executionStateUpdate) (database.Chat, error) { + return tx.store.UpdateChatExecutionState(tx.ctx, database.UpdateChatExecutionStateParams{ + ID: tx.chatID, + Status: u.Status, + Archived: u.Archived, + WorkerID: u.WorkerID, + RunnerID: u.RunnerID, + LastError: u.LastError, + RequiresActionDeadlineAt: u.RequiresActionDeadlineAt, + }) +} + +// insertMessages inserts the given Message batch under the current +// chat. +func (tx *Tx) insertMessages(messages []Message) ([]database.ChatMessage, error) { + if len(messages) == 0 { + return nil, nil + } + inserted, err := tx.store.InsertChatMessages(tx.ctx, toInsertParams(tx.chatID, messages)) + if err != nil { + return nil, xerrors.Errorf("insert messages: %w", err) + } + return inserted, nil +} + +// clearQueue deletes all queued messages on the chat and returns the +// IDs that were deleted in queue order. +func (tx *Tx) clearQueue() ([]int64, error) { + queued, err := tx.store.GetChatQueuedMessagesByPosition(tx.ctx, tx.chatID) + if err != nil { + return nil, xerrors.Errorf("get queued for clear: %w", err) + } + if len(queued) == 0 { + return nil, nil + } + if _, err := tx.store.DeleteAllChatQueuedMessagesReturningCount(tx.ctx, tx.chatID); err != nil { + return nil, xerrors.Errorf("delete queued: %w", err) + } + ids := make([]int64, len(queued)) + for i, q := range queued { + ids[i] = q.ID + } + return ids, nil +} + +// MaxQueueSize is the maximum number of queued user messages per chat. +// Queue-appending transitions reject inserts that would exceed this +// cap with a *MessageQueueFullError that wraps [ErrMessageQueueFull]. +const MaxQueueSize = 20 + +// requireQueueCapacity rejects the call when the chat already has +// MaxQueueSize queued messages. Queue-appending transitions invoke +// this helper inside the transaction immediately before inserting a +// new queued message so the check is atomic with the insert. +func (tx *Tx) requireQueueCapacity() error { + count, err := tx.store.CountChatQueuedMessages(tx.ctx, tx.chatID) + if err != nil { + return xerrors.Errorf("count queued messages: %w", err) + } + if count >= MaxQueueSize { + return &MessageQueueFullError{Max: MaxQueueSize} + } + return nil +} + +// insertQueuedMessage inserts a queued user message. created_by falls +// back to chats.owner_id only when the message does not supply one. +func (tx *Tx) insertQueuedMessage(ownerFallback uuid.UUID, m Message) (database.ChatQueuedMessage, error) { + createdBy := ownerFallback + if m.CreatedBy.Valid { + createdBy = m.CreatedBy.UUID + } + rawContent := m.Content.RawMessage + if !m.Content.Valid || len(rawContent) == 0 { + rawContent = json.RawMessage("null") + } + if err := tx.requireQueueCapacity(); err != nil { + return database.ChatQueuedMessage{}, err + } + return tx.store.InsertChatQueuedMessageWithCreator(tx.ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: tx.chatID, + Content: rawContent, + ModelConfigID: m.ModelConfigID, + CreatedBy: createdBy, + APIKeyID: m.APIKeyID, + }) +} + +// messageFromQueuedRow synthesizes a Message from a stored queued row, +// suitable for promoting into active history. +func messageFromQueuedRow(q database.ChatQueuedMessage) Message { + return Message{ + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: q.Content, Valid: q.Content != nil}, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: q.ModelConfigID, + CreatedBy: uuid.NullUUID{UUID: q.CreatedBy, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: q.APIKeyID, + } +} + +// SetArchivedInput configures [Tx.SetArchived]. +type SetArchivedInput struct { + Archived bool +} + +// SetArchivedResult is returned by [Tx.SetArchived]. +type SetArchivedResult struct{} + +// SetArchived sets or clears the chat's archived marker. +func (tx *Tx) SetArchived(input SetArchivedInput) (SetArchivedResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionSetArchived) + if err != nil { + return SetArchivedResult{}, err + } + if input.Archived == chat.Archived { + // The matrix only allows SetArchived(true) from W/E0/E1 and + // SetArchived(false) from XW/XE0/XE1. A request whose Archived + // field already matches the chat's current archived flag is + // the wrong direction (or a no-op) and must be rejected so we + // do not silently roll the snapshot or publish a chat:update. + return SetArchivedResult{}, newTransitionError( + TransitionSetArchived, from, + "SetArchived input matches the current archived flag", + ) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: chat.Status, + Archived: input.Archived, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }); err != nil { + return SetArchivedResult{}, xerrors.Errorf("update archive: %w", err) + } + return SetArchivedResult{}, nil +} + +// BusyBehavior controls how SendMessage behaves when the chat is +// currently busy (R*/I*/A*). From idle/error states the two behaviors +// are equivalent. +type BusyBehavior string + +const ( + BusyBehaviorQueue BusyBehavior = "queue" + BusyBehaviorInterrupt BusyBehavior = "interrupt" +) + +// SendMessageInput configures [Tx.SendMessage]. +type SendMessageInput struct { + Message Message + BusyBehavior BusyBehavior +} + +// SendMessageResult is returned by [Tx.SendMessage]. +type SendMessageResult struct { + InsertedMessages []database.ChatMessage + QueuedMessage *database.ChatQueuedMessage +} + +// SendMessage admits a new user message. Depending on input state and +// BusyBehavior, the message lands directly in history, in the queue, +// or replaces the queue head as part of a running-state promotion. +func (tx *Tx) SendMessage(input SendMessageInput) (SendMessageResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionSendMessage) + if err != nil { + return SendMessageResult{}, err + } + if input.Message.Role != database.ChatMessageRoleUser { + return SendMessageResult{}, newTransitionError( + TransitionSendMessage, from, + "SendMessage requires a user message", + ) + } + switch input.BusyBehavior { + case BusyBehaviorQueue, BusyBehaviorInterrupt: + // ok + default: + // Reject unknown / empty BusyBehavior up front so an invalid + // value cannot fall through to the queue path on busy states + // or be silently ignored on idle states. The callers in chatd + // default empty to queue; chatstate is the lower-level API + // and refuses to guess. + return SendMessageResult{}, newTransitionError( + TransitionSendMessage, from, + "invalid BusyBehavior", + ) + } + switch from { + // Idle / empty-queue error: insert directly into history, clear + // last_error, leave queue alone. + case StateW, StateE0: + return tx.sendMessageDirect(chat, input.Message) + + // Error-with-queue: append to tail, promote previous head into + // history, clear last_error. + case StateE1: + return tx.sendMessageE1(chat, input.Message) + + // Running with no queue. + case StateR0: + if input.BusyBehavior == BusyBehaviorInterrupt { + return tx.sendMessageQueueAndSetStatus(chat, input.Message, database.ChatStatusInterrupting, chat.LastError, chat.RequiresActionDeadlineAt) + } + return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + + // Running with queue. + case StateR1: + if input.BusyBehavior == BusyBehaviorInterrupt { + return tx.sendMessageQueueAndSetStatus(chat, input.Message, database.ChatStatusInterrupting, chat.LastError, chat.RequiresActionDeadlineAt) + } + return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + + // Interrupting: queue regardless of busy behavior. + case StateI0, StateI1: + return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + + // Requires-action: queue keeps A*; interrupt cancels pending + // dynamic calls and resumes in running. + case StateA0, StateA1: + if input.BusyBehavior == BusyBehaviorInterrupt { + return tx.sendMessageInterruptRequiresAction(chat, input.Message) + } + return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + } + return SendMessageResult{}, newTransitionError(TransitionSendMessage, from, "unhandled state in SendMessage") +} + +func (tx *Tx) sendMessageDirect(chat database.Chat, m Message) (SendMessageResult, error) { + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by new user message", false) + if err != nil { + return SendMessageResult{}, err + } + inserted, err := tx.insertMessages(append(cancels, m)) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("insert direct user message: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: pqtype.NullRawMessage{}, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return SendMessageResult{}, xerrors.Errorf("set running: %w", err) + } + return SendMessageResult{ + InsertedMessages: inserted, + }, nil +} + +func (tx *Tx) sendMessageE1(chat database.Chat, m Message) (SendMessageResult, error) { + queued, err := tx.insertQueuedMessage(chat.OwnerID, m) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("insert queued: %w", err) + } + head, err := tx.store.GetChatQueuedMessageHead(tx.ctx, tx.chatID) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("get queue head: %w", err) + } + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by queued message promotion", false) + if err != nil { + return SendMessageResult{}, err + } + promoted := messageFromQueuedRow(head) + inserted, err := tx.insertMessages(append(cancels, promoted)) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("insert promoted queued head: %w", err) + } + if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: head.ID, + ChatID: tx.chatID, + }); err != nil { + return SendMessageResult{}, xerrors.Errorf("delete promoted queued head: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: pqtype.NullRawMessage{}, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return SendMessageResult{}, xerrors.Errorf("set running: %w", err) + } + return SendMessageResult{ + InsertedMessages: inserted, + QueuedMessage: &queued, + }, nil +} + +func (tx *Tx) sendMessageQueueAndSetStatus( + chat database.Chat, + m Message, + status database.ChatStatus, + lastError pqtype.NullRawMessage, + deadline sql.NullTime, +) (SendMessageResult, error) { + queued, err := tx.insertQueuedMessage(chat.OwnerID, m) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("insert queued: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: status, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: lastError, + RequiresActionDeadlineAt: deadline, + }); err != nil { + return SendMessageResult{}, xerrors.Errorf("update status: %w", err) + } + return SendMessageResult{ + QueuedMessage: &queued, + }, nil +} + +func (tx *Tx) sendMessageInterruptRequiresAction(chat database.Chat, m Message) (SendMessageResult, error) { + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by user message", true) + if err != nil { + return SendMessageResult{}, err + } + if _, err := tx.insertMessages(cancels); err != nil { + return SendMessageResult{}, xerrors.Errorf("insert requires-action cancellations: %w", err) + } + return tx.sendMessageQueueAndSetStatus(chat, m, database.ChatStatusRunning, chat.LastError, sql.NullTime{}) +} + +// EditMessageInput configures [Tx.EditMessage]. +type EditMessageInput struct { + MessageID int64 + CreatedBy uuid.UUID + Content pqtype.NullRawMessage + ModelConfigIDOverride uuid.NullUUID + APIKeyID sql.NullString +} + +// EditMessageResult is returned by [Tx.EditMessage]. +type EditMessageResult struct { + ReplacementMessage database.ChatMessage + DeletedMessageIDs []int64 + DeletedQueuedMessageIDs []int64 + CancellationMessages []database.ChatMessage +} + +// EditMessage replaces an earlier user message and discards the +// active-history suffix that followed it. +func (tx *Tx) EditMessage(input EditMessageInput) (EditMessageResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionEditMessage) + if err != nil { + return EditMessageResult{}, err + } + target, err := tx.store.GetChatMessageByID(tx.ctx, input.MessageID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return EditMessageResult{}, ErrMessageNotFound + } + return EditMessageResult{}, xerrors.Errorf("get target message: %w", err) + } + if target.ChatID != tx.chatID { + return EditMessageResult{}, ErrMessageNotFound + } + if target.Deleted { + return EditMessageResult{}, ErrMessageNotFound + } + if target.Role != database.ChatMessageRoleUser { + return EditMessageResult{}, newTransitionErrorWithCause( + TransitionEditMessage, from, + ErrEditedMessageNotUser, + "only user messages can be edited", + ) + } + + suffix, err := tx.store.GetChatMessagesByChatID(tx.ctx, database.GetChatMessagesByChatIDParams{ + ChatID: tx.chatID, + AfterID: target.ID - 1, // include target and everything after + }) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("get suffix messages: %w", err) + } + deletedIDs := make([]int64, 0, len(suffix)) + for _, m := range suffix { + if !m.Deleted { + deletedIDs = append(deletedIDs, m.ID) + } + } + + if err := tx.store.SoftDeleteChatMessageByID(tx.ctx, target.ID); err != nil { + return EditMessageResult{}, xerrors.Errorf("soft-delete target: %w", err) + } + if err := tx.store.SoftDeleteChatMessagesAfterID(tx.ctx, database.SoftDeleteChatMessagesAfterIDParams{ + ChatID: tx.chatID, + AfterID: target.ID, + }); err != nil { + return EditMessageResult{}, xerrors.Errorf("soft-delete suffix: %w", err) + } + + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by message edit", false) + if err != nil { + return EditMessageResult{}, err + } + cancellationMessages, err := tx.insertMessages(cancels) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("insert message edit cancellations: %w", err) + } + + modelConfig := target.ModelConfigID + if input.ModelConfigIDOverride.Valid { + modelConfig = input.ModelConfigIDOverride + } + apiKeyID := input.APIKeyID + if !apiKeyID.Valid { + return EditMessageResult{}, xerrors.Errorf("api_key_id is required") + } + replacement := Message{ + Role: database.ChatMessageRoleUser, + Content: input.Content, + Visibility: target.Visibility, + ModelConfigID: modelConfig, + CreatedBy: uuid.NullUUID{UUID: input.CreatedBy, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: apiKeyID, + } + insertedReplacement, err := tx.insertMessages([]Message{replacement}) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("insert replacement message: %w", err) + } + var replacementRow database.ChatMessage + if len(insertedReplacement) == 1 { + replacementRow = insertedReplacement[0] + } + + deletedQueuedIDs, err := tx.clearQueue() + if err != nil { + return EditMessageResult{}, err + } + + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: pqtype.NullRawMessage{}, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return EditMessageResult{}, xerrors.Errorf("set running: %w", err) + } + return EditMessageResult{ + ReplacementMessage: replacementRow, + DeletedMessageIDs: deletedIDs, + DeletedQueuedMessageIDs: deletedQueuedIDs, + CancellationMessages: cancellationMessages, + }, nil +} + +// DeleteQueuedMessageInput configures [Tx.DeleteQueuedMessage]. +type DeleteQueuedMessageInput struct { + QueuedMessageID int64 +} + +// DeleteQueuedMessageResult is returned by [Tx.DeleteQueuedMessage]. +type DeleteQueuedMessageResult struct { + DeletedQueuedMessage database.ChatQueuedMessage +} + +// DeleteQueuedMessage removes a single queued user message. +func (tx *Tx) DeleteQueuedMessage(input DeleteQueuedMessageInput) (DeleteQueuedMessageResult, error) { + _, _, err := tx.requireFromAllowed(TransitionDeleteQueuedMessage) + if err != nil { + return DeleteQueuedMessageResult{}, err + } + target, err := tx.store.GetChatQueuedMessageByID(tx.ctx, database.GetChatQueuedMessageByIDParams{ + ID: input.QueuedMessageID, + ChatID: tx.chatID, + }) + if errors.Is(err, sql.ErrNoRows) { + return DeleteQueuedMessageResult{}, ErrQueuedMessageNotFound + } + if err != nil { + return DeleteQueuedMessageResult{}, xerrors.Errorf("get queued: %w", err) + } + rows, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: input.QueuedMessageID, + ChatID: tx.chatID, + }) + if err != nil { + return DeleteQueuedMessageResult{}, xerrors.Errorf("delete queued: %w", err) + } + if rows == 0 { + return DeleteQueuedMessageResult{}, ErrQueuedMessageNotFound + } + return DeleteQueuedMessageResult{ + DeletedQueuedMessage: target, + }, nil +} + +// PromoteQueuedMessageInput configures [Tx.PromoteQueuedMessage]. +type PromoteQueuedMessageInput struct { + QueuedMessageID int64 +} + +// PromoteQueuedMessageResult is returned by [Tx.PromoteQueuedMessage]. +type PromoteQueuedMessageResult struct { + QueuedMessage database.ChatQueuedMessage + InsertedMessage *database.ChatMessage + ReorderedQueueOnly bool + CancellationMessages []database.ChatMessage +} + +// PromoteQueuedMessage promotes the target queued message to the +// queue head; from E1/A1 it also pops it into active history. +func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueuedMessageResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionPromoteQueuedMessage) + if err != nil { + return PromoteQueuedMessageResult{}, err + } + target, err := tx.store.GetChatQueuedMessageByID(tx.ctx, database.GetChatQueuedMessageByIDParams{ + ID: input.QueuedMessageID, + ChatID: tx.chatID, + }) + if errors.Is(err, sql.ErrNoRows) { + return PromoteQueuedMessageResult{}, ErrQueuedMessageNotFound + } + if err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("get queued: %w", err) + } + rows, err := tx.store.ReorderChatQueuedMessageToHead(tx.ctx, database.ReorderChatQueuedMessageToHeadParams{ + ID: input.QueuedMessageID, + ChatID: tx.chatID, + }) + if err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("reorder queue: %w", err) + } + reorderOnly := rows > 0 + + // R1/I1: leave the target at the queue head and transition to + // status `interrupting` so the worker can drain the in-flight + // generation before promoting the queue head into active history. + // No history row is inserted here and no queue rows are deleted. + if from == StateR1 || from == StateI1 { + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusInterrupting, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }); err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("set interrupting: %w", err) + } + return PromoteQueuedMessageResult{ + QueuedMessage: target, + ReorderedQueueOnly: reorderOnly, + }, nil + } + + // E1/A1: synthesize cancellations, pop the head, insert into + // history, set running. Both paths insert a queued user message + // into active history, so every outstanding tool call must be + // closed (not just dynamic ones) to keep the LLM history valid. + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by queued message promotion", false) + if err != nil { + return PromoteQueuedMessageResult{}, err + } + promotedMsg := messageFromQueuedRow(target) + inserted, err := tx.insertMessages(append(cancels, promotedMsg)) + if err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("insert promoted queued message: %w", err) + } + if len(inserted) != len(cancels)+1 { + return PromoteQueuedMessageResult{}, xerrors.Errorf( + "insert promoted queued message: expected %d rows, got %d", + len(cancels)+1, len(inserted), + ) + } + if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: target.ID, + ChatID: tx.chatID, + }); err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("delete promoted queued: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: pqtype.NullRawMessage{}, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("set running: %w", err) + } + cancellations := inserted[:len(inserted)-1] + insertedUserMsg := inserted[len(inserted)-1] + return PromoteQueuedMessageResult{ + QueuedMessage: target, + InsertedMessage: &insertedUserMsg, + CancellationMessages: cancellations, + ReorderedQueueOnly: reorderOnly, + }, nil +} + +// InterruptInput configures [Tx.Interrupt]. +type InterruptInput struct { + Reason string +} + +// InterruptResult is returned by [Tx.Interrupt]. +type InterruptResult struct { + CancellationMessages []database.ChatMessage +} + +// Interrupt requests interruption of an active or requires-action +// chat. +func (tx *Tx) Interrupt(input InterruptInput) (InterruptResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionInterrupt) + if err != nil { + return InterruptResult{}, err + } + switch from { + case StateR0, StateR1: + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusInterrupting, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }); err != nil { + return InterruptResult{}, xerrors.Errorf("set interrupting: %w", err) + } + return InterruptResult{}, nil + case StateA0, StateA1: + reason := input.Reason + if reason == "" { + reason = "Tool execution interrupted by user" + } + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, reason, true) + if err != nil { + return InterruptResult{}, err + } + inserted, err := tx.insertMessages(cancels) + if err != nil { + return InterruptResult{}, xerrors.Errorf("insert interrupt cancellations: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return InterruptResult{}, xerrors.Errorf("set running: %w", err) + } + return InterruptResult{ + CancellationMessages: inserted, + }, nil + default: + return InterruptResult{}, newTransitionError(TransitionInterrupt, from, "unhandled state in Interrupt") + } +} + +// ToolResultInput is one submitted dynamic-tool result. +type ToolResultInput struct { + ToolCallID string + Output json.RawMessage + IsError bool +} + +// CompleteRequiresActionInput configures [Tx.CompleteRequiresAction]. +type CompleteRequiresActionInput struct { + CreatedBy uuid.UUID + ModelConfigID uuid.UUID + Results []ToolResultInput +} + +// CompleteRequiresActionResult is returned by [Tx.CompleteRequiresAction]. +type CompleteRequiresActionResult struct { + InsertedMessages []database.ChatMessage +} + +// CompleteRequiresAction validates and stores user-submitted tool +// results that satisfy the chat's pending dynamic tool calls, then +// returns the chat to running. +func (tx *Tx) CompleteRequiresAction(input CompleteRequiresActionInput) (CompleteRequiresActionResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionCompleteRequiresAction) + if err != nil { + return CompleteRequiresActionResult{}, err + } + pending, err := pendingDynamicToolCallIDs(tx.ctx, tx.store, chat) + if err != nil { + return CompleteRequiresActionResult{}, err + } + submitted := make(map[string]ToolResultInput, len(input.Results)) + for _, r := range input.Results { + if _, dup := submitted[r.ToolCallID]; dup { + return CompleteRequiresActionResult{}, newTransitionErrorWithCause( + TransitionCompleteRequiresAction, from, + &ToolResultValidationError{Cause: ErrToolResultDuplicate, ToolCallID: r.ToolCallID}, + "duplicate tool_call_id submitted", + ) + } + if !json.Valid(r.Output) { + return CompleteRequiresActionResult{}, newTransitionErrorWithCause( + TransitionCompleteRequiresAction, from, + &ToolResultValidationError{Cause: ErrToolResultInvalidJSON, ToolCallID: r.ToolCallID}, + "tool result output is not valid JSON", + ) + } + submitted[r.ToolCallID] = r + } + for id := range pending { + if _, ok := submitted[id]; !ok { + return CompleteRequiresActionResult{}, newTransitionErrorWithCause( + TransitionCompleteRequiresAction, from, + &ToolResultValidationError{Cause: ErrToolResultMissing, ToolCallID: id}, + "submitted tool results do not match pending tool calls", + ) + } + } + for id := range submitted { + if _, ok := pending[id]; !ok { + return CompleteRequiresActionResult{}, newTransitionErrorWithCause( + TransitionCompleteRequiresAction, from, + &ToolResultValidationError{Cause: ErrToolResultUnexpected, ToolCallID: id}, + "submitted tool_call_id does not match a pending dynamic tool call", + ) + } + } + messages := make([]Message, 0, len(input.Results)) + for _, r := range input.Results { + part := codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: r.ToolCallID, + ToolName: pending[r.ToolCallID], + Result: r.Output, + IsError: r.IsError, + } + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part}) + if err != nil { + return CompleteRequiresActionResult{}, xerrors.Errorf("marshal tool result: %w", err) + } + messages = append(messages, Message{ + Role: database.ChatMessageRoleTool, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + CreatedBy: uuid.NullUUID{UUID: input.CreatedBy, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: input.ModelConfigID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + }) + } + inserted, err := tx.insertMessages(messages) + if err != nil { + return CompleteRequiresActionResult{}, xerrors.Errorf("insert tool results: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return CompleteRequiresActionResult{}, xerrors.Errorf("set running: %w", err) + } + return CompleteRequiresActionResult{ + InsertedMessages: inserted, + }, nil +} + +// AcquireInput configures [Tx.Acquire]. +type AcquireInput struct { + WorkerID uuid.UUID + RunnerID uuid.UUID +} + +// AcquireResult is returned by [Tx.Acquire]. +type AcquireResult struct{} + +// Acquire claims the chat for a worker/runner pair. Execution state +// is preserved. +// +// Acquire never inspects the chat's current ownership: it simply +// overwrites worker_id/runner_id with the supplied identifiers and +// upserts a fresh heartbeat. Detecting and recovering from stale +// leases is a worker-side fence concern outside the state machine. +// Callers that need to coordinate takeovers with the previous owner +// must arrange that out-of-band before calling Acquire. +func (tx *Tx) Acquire(input AcquireInput) (AcquireResult, error) { + chat, _, err := tx.loadState() + if err != nil { + return AcquireResult{}, err + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: chat.Status, + Archived: chat.Archived, + WorkerID: uuid.NullUUID{UUID: input.WorkerID, Valid: true}, + RunnerID: uuid.NullUUID{UUID: input.RunnerID, Valid: true}, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }); err != nil { + return AcquireResult{}, xerrors.Errorf("set ownership: %w", err) + } + if err := tx.store.UpsertChatHeartbeat(tx.ctx, database.UpsertChatHeartbeatParams{ + ChatID: tx.chatID, + RunnerID: input.RunnerID, + }); err != nil { + return AcquireResult{}, xerrors.Errorf("upsert heartbeat: %w", err) + } + // Acquire writes a fresh heartbeat itself, so the post-commit + // ownership-hint logic in Update will evaluate the heartbeat as + // fresh and skip publishing a `chat:ownership` hint. + return AcquireResult{}, nil +} + +// AbandonInput is intentionally empty. Ownership-fence checks belong +// outside the transition in caller code that reads the locked row before +// invoking Abandon. +type AbandonInput struct{} + +// AbandonResult is returned by [Tx.Abandon]. +type AbandonResult struct{} + +// Abandon clears worker_id and runner_id from the locked chat row. It +// rejects calls when the chat is not currently owned (worker_id IS NULL). +// Callers that need to verify their own identity before abandoning should +// read the locked row through the transactional store and compare values before +// invoking Abandon. +func (tx *Tx) Abandon(_ AbandonInput) (AbandonResult, error) { + chat, from, err := tx.loadState() + if err != nil { + return AbandonResult{}, err + } + if !chat.WorkerID.Valid { + return AbandonResult{}, newTransitionError(TransitionAbandon, from, "chat is not owned") + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: chat.Status, + Archived: chat.Archived, + WorkerID: uuid.NullUUID{}, + RunnerID: uuid.NullUUID{}, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }); err != nil { + return AbandonResult{}, xerrors.Errorf("clear ownership: %w", err) + } + return AbandonResult{}, nil +} + +// RecordGenerationAttemptInput is intentionally empty. +type RecordGenerationAttemptInput struct{} + +// RecordGenerationAttemptResult is returned by [Tx.RecordGenerationAttempt]. +type RecordGenerationAttemptResult struct { + GenerationAttempt int64 +} + +// RecordGenerationAttempt durably records that the worker is +// attempting another generation under the current history version. +func (tx *Tx) RecordGenerationAttempt(_ RecordGenerationAttemptInput) (RecordGenerationAttemptResult, error) { + _, _, err := tx.requireFromAllowed(TransitionRecordGenerationAttempt) + if err != nil { + return RecordGenerationAttemptResult{}, err + } + value, err := tx.store.IncrementChatGenerationAttempt(tx.ctx, tx.chatID) + if err != nil { + return RecordGenerationAttemptResult{}, xerrors.Errorf("increment generation attempt: %w", err) + } + return RecordGenerationAttemptResult{ + GenerationAttempt: value, + }, nil +} + +// RecordRetryStateInput configures [Tx.RecordRetryState]. +type RecordRetryStateInput struct { + RetryState pqtype.NullRawMessage +} + +// RecordRetryStateResult is returned by [Tx.RecordRetryState]. +type RecordRetryStateResult struct { + Chat database.Chat +} + +// RecordRetryState stores the client-visible retry payload for the +// current generation attempt. +func (tx *Tx) RecordRetryState(input RecordRetryStateInput) (RecordRetryStateResult, error) { + _, from, err := tx.requireFromAllowed(TransitionRecordRetryState) + if err != nil { + return RecordRetryStateResult{}, err + } + if !input.RetryState.Valid || len(input.RetryState.RawMessage) == 0 { + return RecordRetryStateResult{}, newTransitionError( + TransitionRecordRetryState, from, + "RecordRetryState requires a retry payload", + ) + } + if !json.Valid(input.RetryState.RawMessage) { + return RecordRetryStateResult{}, newTransitionError( + TransitionRecordRetryState, from, + "retry payload is not valid JSON", + ) + } + chat, err := tx.store.UpdateChatRetryState(tx.ctx, database.UpdateChatRetryStateParams{ + ID: tx.chatID, + RetryState: input.RetryState.RawMessage, + }) + if err != nil { + return RecordRetryStateResult{}, xerrors.Errorf("update retry state: %w", err) + } + return RecordRetryStateResult{Chat: chat}, nil +} + +// CommitStepInput configures [Tx.CommitStep]. +type CommitStepInput struct { + Messages []Message +} + +// CommitStepResult is returned by [Tx.CommitStep]. +type CommitStepResult struct { + InsertedMessages []database.ChatMessage +} + +// CommitStep stores one durable message suffix while remaining +// running. +func (tx *Tx) CommitStep(input CommitStepInput) (CommitStepResult, error) { + _, from, err := tx.requireFromAllowed(TransitionCommitStep) + if err != nil { + return CommitStepResult{}, err + } + if len(input.Messages) == 0 { + return CommitStepResult{}, newTransitionError( + TransitionCommitStep, from, + "CommitStep requires at least one message", + ) + } + inserted, err := tx.insertMessages(input.Messages) + if err != nil { + return CommitStepResult{}, xerrors.Errorf("insert commit step messages: %w", err) + } + return CommitStepResult{ + InsertedMessages: inserted, + }, nil +} + +// requiresActionTimeout is the time allowed for a client to submit +// required dynamic tool results before follow-up logic may consider +// the requires-action state expired. +const requiresActionTimeout = 5 * time.Minute + +// EnterRequiresActionInput is intentionally empty. +type EnterRequiresActionInput struct{} + +// EnterRequiresActionResult is returned by [Tx.EnterRequiresAction]. +type EnterRequiresActionResult struct { + RequiresActionDeadlineAt sql.NullTime +} + +// EnterRequiresAction parks the chat in requires_action with a +// database-time deadline of now() + requiresActionTimeout. +func (tx *Tx) EnterRequiresAction(_ EnterRequiresActionInput) (EnterRequiresActionResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionEnterRequiresAction) + if err != nil { + return EnterRequiresActionResult{}, err + } + pending, err := pendingDynamicToolCallIDs(tx.ctx, tx.store, chat) + if err != nil { + return EnterRequiresActionResult{}, err + } + if len(pending) == 0 { + return EnterRequiresActionResult{}, newTransitionError( + TransitionEnterRequiresAction, from, + "no pending dynamic tool calls", + ) + } + now, err := tx.store.GetDatabaseNow(tx.ctx) + if err != nil { + return EnterRequiresActionResult{}, xerrors.Errorf("get db now: %w", err) + } + deadline := sql.NullTime{Time: now.Add(requiresActionTimeout), Valid: true} + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRequiresAction, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: deadline, + }); err != nil { + return EnterRequiresActionResult{}, xerrors.Errorf("set requires_action: %w", err) + } + return EnterRequiresActionResult{ + RequiresActionDeadlineAt: deadline, + }, nil +} + +// FinishInterruptionInput configures [Tx.FinishInterruption]. +type FinishInterruptionInput struct { + PartialMessages []Message +} + +// FinishInterruptionResult is returned by [Tx.FinishInterruption]. +type FinishInterruptionResult struct { + InsertedMessages []database.ChatMessage + PromotedMessage *database.ChatMessage +} + +// FinishInterruption commits an optional partial assistant/tool suffix +// and lands the chat in waiting (I0) or running with the next queued +// message promoted (I1). +func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterruptionResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionFinishInterruption) + if err != nil { + return FinishInterruptionResult{}, err + } + insertedPartial, err := tx.insertMessages(input.PartialMessages) + if err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("insert interruption partial messages: %w", err) + } + pendingAll, err := pendingAllToolCallIDs(tx.ctx, tx.store, chat) + if err != nil { + return FinishInterruptionResult{}, err + } + if len(pendingAll) > 0 { + return FinishInterruptionResult{}, newTransitionError( + TransitionFinishInterruption, from, + "outstanding tool calls remain after partial commit", + ) + } + + if from == StateI0 { + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusWaiting, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("set waiting: %w", err) + } + return FinishInterruptionResult{ + InsertedMessages: insertedPartial, + }, nil + } + + // I1: promote queue head into history. + head, err := tx.store.GetChatQueuedMessageHead(tx.ctx, tx.chatID) + if err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("get queue head: %w", err) + } + promotedMsg := messageFromQueuedRow(head) + insertedHead, err := tx.insertMessages([]Message{promotedMsg}) + if err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("insert promoted queue head: %w", err) + } + if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: head.ID, + ChatID: tx.chatID, + }); err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("delete promoted head: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("set running: %w", err) + } + insertedPartial = append(insertedPartial, insertedHead...) + var promoted *database.ChatMessage + if len(insertedHead) == 1 { + promoted = &insertedHead[0] + } + return FinishInterruptionResult{ + InsertedMessages: insertedPartial, + PromotedMessage: promoted, + }, nil +} + +// FinishTurnInput is intentionally empty. +type FinishTurnInput struct{} + +// FinishTurnResult is returned by [Tx.FinishTurn]. +type FinishTurnResult struct { + Chat database.Chat + PromotedMessage *database.ChatMessage +} + +// FinishTurn completes a running turn. +func (tx *Tx) FinishTurn(_ FinishTurnInput) (FinishTurnResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionFinishTurn) + if err != nil { + return FinishTurnResult{}, err + } + if from == StateR0 { + updated, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusWaiting, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }) + if err != nil { + return FinishTurnResult{}, xerrors.Errorf("set waiting: %w", err) + } + return FinishTurnResult{Chat: updated}, nil + } + // R1. + head, err := tx.store.GetChatQueuedMessageHead(tx.ctx, tx.chatID) + if err != nil { + return FinishTurnResult{}, xerrors.Errorf("get queue head: %w", err) + } + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by queued message promotion", false) + if err != nil { + return FinishTurnResult{}, err + } + promotedMsg := messageFromQueuedRow(head) + inserted, err := tx.insertMessages(append(cancels, promotedMsg)) + if err != nil { + return FinishTurnResult{}, xerrors.Errorf("insert promoted queue head: %w", err) + } + if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: head.ID, + ChatID: tx.chatID, + }); err != nil { + return FinishTurnResult{}, xerrors.Errorf("delete promoted head: %w", err) + } + updated, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }) + if err != nil { + return FinishTurnResult{}, xerrors.Errorf("set running: %w", err) + } + var promoted *database.ChatMessage + if len(inserted) > 0 { + promoted = &inserted[len(inserted)-1] + } + return FinishTurnResult{ + Chat: updated, + PromotedMessage: promoted, + }, nil +} + +// FinishErrorInput configures [Tx.FinishError]. +type FinishErrorInput struct { + LastError pqtype.NullRawMessage +} + +// FinishErrorResult is returned by [Tx.FinishError]. +type FinishErrorResult struct{} + +// FinishError parks the chat in error with the supplied last_error. +func (tx *Tx) FinishError(input FinishErrorInput) (FinishErrorResult, error) { + chat, _, err := tx.requireFromAllowed(TransitionFinishError) + if err != nil { + return FinishErrorResult{}, err + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusError, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: input.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return FinishErrorResult{}, xerrors.Errorf("set error: %w", err) + } + return FinishErrorResult{}, nil +} + +// CancelRequiresActionInput configures [Tx.CancelRequiresAction]. +type CancelRequiresActionInput struct { + Reason string +} + +// CancelRequiresActionResult is returned by [Tx.CancelRequiresAction]. +type CancelRequiresActionResult struct { + CancellationMessages []database.ChatMessage +} + +// CancelRequiresAction synthesizes cancellation results for every +// pending dynamic tool call and returns the chat to running. +func (tx *Tx) CancelRequiresAction(input CancelRequiresActionInput) (CancelRequiresActionResult, error) { + chat, from, err := tx.requireFromAllowed(TransitionCancelRequiresAction) + if err != nil { + return CancelRequiresActionResult{}, err + } + reason := input.Reason + if reason == "" { + reason = "Tool execution timed out" + } + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, reason, true) + if err != nil { + return CancelRequiresActionResult{}, err + } + if len(cancels) == 0 { + return CancelRequiresActionResult{}, newTransitionError( + TransitionCancelRequiresAction, from, + "no pending dynamic tool calls to cancel", + ) + } + inserted, err := tx.insertMessages(cancels) + if err != nil { + return CancelRequiresActionResult{}, xerrors.Errorf("insert requires-action cancellations: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusRunning, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return CancelRequiresActionResult{}, xerrors.Errorf("set running: %w", err) + } + return CancelRequiresActionResult{ + CancellationMessages: inserted, + }, nil +} + +// ReconcileInvalidStateInput configures [Tx.ReconcileInvalidState]. +type ReconcileInvalidStateInput struct { + LastError pqtype.NullRawMessage + CancellationReason string +} + +// ReconcileInvalidStateResult is returned by [Tx.ReconcileInvalidState]. +type ReconcileInvalidStateResult struct { + CancellationMessages []database.ChatMessage +} + +// ReconcileInvalidState moves an invalid execution-state combination +// into a valid error state. Queued messages are preserved; pending +// dynamic-tool calls are closed with synthetic cancellation results. +func (tx *Tx) ReconcileInvalidState(input ReconcileInvalidStateInput) (ReconcileInvalidStateResult, error) { + chat, from, err := tx.loadState() + if err != nil { + return ReconcileInvalidStateResult{}, err + } + if from != StateInvalid { + return ReconcileInvalidStateResult{}, newTransitionError( + TransitionReconcileInvalidState, from, + "reconcile is only valid for invalid states", + ) + } + reason := input.CancellationReason + if reason == "" { + reason = "Tool execution canceled due to invalid chat state" + } + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, reason, true) + if err != nil { + return ReconcileInvalidStateResult{}, err + } + var inserted []database.ChatMessage + if len(cancels) > 0 { + inserted, err = tx.insertMessages(cancels) + if err != nil { + return ReconcileInvalidStateResult{}, xerrors.Errorf("insert invalid-state cancellations: %w", err) + } + } + lastErr := input.LastError + if !lastErr.Valid { + lastErr = pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"chat was in an invalid state; send a new message or edit history to continue"}`), + Valid: true, + } + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusError, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: lastErr, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return ReconcileInvalidStateResult{}, xerrors.Errorf("set error: %w", err) + } + return ReconcileInvalidStateResult{ + CancellationMessages: inserted, + }, nil +} diff --git a/coderd/x/chatd/chatstate/transitions_helpers_test.go b/coderd/x/chatd/chatstate/transitions_helpers_test.go new file mode 100644 index 0000000000..91cc2419f8 --- /dev/null +++ b/coderd/x/chatd/chatstate/transitions_helpers_test.go @@ -0,0 +1,894 @@ +package chatstate_test + +import ( + "context" + "encoding/json" + "fmt" + "testing" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// seededChat is the shared output of seedState. Some transition tests +// need extra context beyond the chat ID (for example, the queued +// message ID to delete, or the message ID to edit), so this struct +// surfaces what each state was seeded with. + +type seededChat struct { + chatID uuid.UUID + exists bool + initialUserMessageID int64 + assistantToolCallMsgID int64 + queuedMessageIDs []int64 + // queuedMessageBodies is parallel to queuedMessageIDs and records + // the text body each queued message was seeded with. Cases that + // promote queued messages into history use this to assert the + // promoted message content matches what was originally queued. + queuedMessageBodies []string + queuedMessageCreatedBy []uuid.UUID + dynamicToolName string + pendingToolCallID string + pendingToolCallIDs []string +} + +// dynamicToolJSON returns the canonical [{name,description,input_schema}] +// payload used to seed dynamic_tools on a chat. Tests that need +// pending dynamic tool calls (A0, A1) reuse this and reference the +// returned tool name in their assistant tool-call message. +func dynamicToolJSON(name string) []byte { + tools := []codersdk.DynamicTool{{ + Name: name, + Description: "test tool", + InputSchema: json.RawMessage(`{"type":"object","properties":{}}`), + }} + raw, err := json.Marshal(tools) + if err != nil { + panic(err) + } + return raw +} + +// assistantToolCallMessage builds a chatstate.Message for an +// assistant message that issues one tool call against the supplied +// dynamic tool name. The tool-call ID is unique per call so multiple +// messages do not collide. +func assistantToolCallMessage(t *testing.T, modelID uuid.UUID, toolName, callID string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: callID, + ToolName: toolName, + Args: json.RawMessage(`{}`), + }}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + } +} + +func mixedAssistantToolCallMessage(t *testing.T, modelID uuid.UUID, dynamicTool, dynCallID, nonDynCallID string) chatstate.Message { + t.Helper() + parts := []codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: dynCallID, + ToolName: dynamicTool, + Args: json.RawMessage(`{}`), + }, + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: nonDynCallID, + ToolName: "non_dynamic_tool", + Args: json.RawMessage(`{}`), + }, + } + raw, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelID, Valid: true}, + } +} + +// createTestChatWithDynamicTools mirrors createTestChat but seeds the +// chat with a non-empty dynamic_tools blob so EnterRequiresAction, +// CompleteRequiresAction, and CancelRequiresAction can find pending +// dynamic tool calls. +func createTestChatWithDynamicTools(t *testing.T, f *testFixture, toolName string) chatstate.CreateChatResult { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + DynamicTools: pqtype.NullRawMessage{ + RawMessage: dynamicToolJSON(toolName), + Valid: true, + }, + InitialMessages: []chatstate.Message{ + userTextMessage("hello", f.User.ID, f.Model.ID), + }, + }) + require.NoError(t, err) + return res +} + +// seedAOrA1 seeds a chat into A0 (queuedExtras=0) or A1 +// (queuedExtras>=1) with a real pending dynamic tool call. Used by +// cases that need A0 or A1 with a configurable queue cardinality. +func seedAOrA1(t *testing.T, f *testFixture, queuedExtras int, namePrefix string) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + toolName := namePrefix + callID := "call_" + uuid.NewString() + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + assistantToolCallMessage(t, f.Model.ID, toolName, callID), + }, + }) + return err + })) + require.Len(t, step.InsertedMessages, 1) + // R0 -> A0. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + var ( + queuedIDs []int64 + queuedBodies []string + ) + for i := 0; i < queuedExtras; i++ { + body := fmt.Sprintf("queued-%s-%d", namePrefix, i) + sm := sendQueuedMessage(t, f, m, body) + require.NotNil(t, sm.QueuedMessage) + queuedIDs = append(queuedIDs, sm.QueuedMessage.ID) + queuedBodies = append(queuedBodies, body) + } + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + assistantToolCallMsgID: step.InsertedMessages[0].ID, + queuedMessageIDs: queuedIDs, + queuedMessageBodies: queuedBodies, + dynamicToolName: toolName, + pendingToolCallID: callID, + } +} + +// seedState seeds a chat into the supplied execution state and +// returns identifying handles useful for downstream assertions. For +// [chatstate.StateN] the returned chatID is a fresh UUID that does +// not exist in the database. Multi-queued seeds (for E1, R1, I1, +// A1 with 2 queued messages, and Invalid with a non-empty queue) live in +// seedStateMultiQueued. +func seedState(t *testing.T, f *testFixture, state chatstate.ExecutionState) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + + switch state { + case chatstate.StateN: + return seededChat{chatID: uuid.New(), exists: false} + + case chatstate.StateR0: + created := createTestChat(t, f) + initial := firstUserMessageID(ctx, t, f, created.Chat.ID) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: initial, + } + + case chatstate.StateW: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + + case chatstate.StateE0: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + + case chatstate.StateE1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + // R0 -> R1 + queuedBody := "queued-for-E1" + queued := sendQueuedMessage(t, f, m, queuedBody) + require.NotNil(t, queued.QueuedMessage) + // R1 -> E1 + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{queued.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + } + + case chatstate.StateR1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + queuedBody := "queued-for-R1" + queued := sendQueuedMessage(t, f, m, queuedBody) + require.NotNil(t, queued.QueuedMessage) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{queued.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + } + + case chatstate.StateI0: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Interrupt(chatstate.InterruptInput{Reason: "seed"}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + + case chatstate.StateI1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + // R0 -> I1: SendMessage with interrupt behavior queues the + // message and sets status to interrupting. + queuedBody := "queued-for-I1" + sm := sendInterruptMessage(t, f, m, queuedBody) + require.NotNil(t, sm.QueuedMessage) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{sm.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + } + + case chatstate.StateA0: + return seedAOrA1(t, f, 0, "seed_tool_a0") + + case chatstate.StateA1: + return seedAOrA1(t, f, 1, "seed_tool_a1") + + case chatstate.StateXW: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + + case chatstate.StateXE0: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + + case chatstate.StateXE1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + queuedBody := "queued-for-XE1" + queued := sendQueuedMessage(t, f, m, queuedBody) + require.NotNil(t, queued.QueuedMessage) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{queued.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + } + + case chatstate.StateInvalid: + created := createTestChat(t, f) + // Force running + archived, a deliberately invalid + // combination per the classifier. + _, err := f.DB.UpdateChatExecutionState(ctx, database.UpdateChatExecutionStateParams{ + ID: created.Chat.ID, + Status: database.ChatStatusRunning, + Archived: true, + WorkerID: created.Chat.WorkerID, + RunnerID: created.Chat.RunnerID, + LastError: created.Chat.LastError, + RequiresActionDeadlineAt: created.Chat.RequiresActionDeadlineAt, + }) + require.NoError(t, err) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + } + } + t.Fatalf("seedState: unsupported execution state %s", state) + return seededChat{} +} + +// seedStateMultiQueued seeds a state with two queued messages. Used +// by cases that need the post-mutation queue to remain non-empty. +// Supported states: E1, R1, I1, A1. +func seedStateMultiQueued(t *testing.T, f *testFixture, state chatstate.ExecutionState) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + switch state { + case chatstate.StateE1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + firstBody := "queued-e1-a" + first := sendQueuedMessage(t, f, m, firstBody) + require.NotNil(t, first.QueuedMessage) + secondBody := "queued-e1-b" + second := sendQueuedMessage(t, f, m, secondBody) + require.NotNil(t, second.QueuedMessage) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{first.QueuedMessage.ID, second.QueuedMessage.ID}, + queuedMessageBodies: []string{firstBody, secondBody}, + } + + case chatstate.StateR1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + firstBody := "queued-r1-a" + first := sendQueuedMessage(t, f, m, firstBody) + require.NotNil(t, first.QueuedMessage) + secondBody := "queued-r1-b" + second := sendQueuedMessage(t, f, m, secondBody) + require.NotNil(t, second.QueuedMessage) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{first.QueuedMessage.ID, second.QueuedMessage.ID}, + queuedMessageBodies: []string{firstBody, secondBody}, + } + + case chatstate.StateI1: + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + firstBody := "queued-i1-a" + first := sendQueuedMessage(t, f, m, firstBody) + require.NotNil(t, first.QueuedMessage) + // R1 -> I1 via interrupt-mode SendMessage queues a second + // message and flips status to interrupting. + secondBody := "queued-i1-b" + second := sendInterruptMessage(t, f, m, secondBody) + require.NotNil(t, second.QueuedMessage) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + queuedMessageIDs: []int64{first.QueuedMessage.ID, second.QueuedMessage.ID}, + queuedMessageBodies: []string{firstBody, secondBody}, + } + + case chatstate.StateA1: + return seedAOrA1(t, f, 2, "seed_tool_a1_multi") + } + t.Fatalf("seedStateMultiQueued: unsupported execution state %s", state) + return seededChat{} +} + +// seedA1WithMixedOutstandingToolCalls seeds A1 with one queued message +// and one assistant message carrying both a dynamic and non-dynamic +// outstanding tool call. It is used by PromoteQueuedMessage(A1) to +// prove all tool calls are closed before inserting the promoted user. +func seedA1WithMixedOutstandingToolCalls(t *testing.T, f *testFixture, queuedExtras int, namePrefix string) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + toolName := namePrefix + dynCallID := "call_" + uuid.NewString() + nonDynCallID := "call_" + uuid.NewString() + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + mixedAssistantToolCallMessage(t, f.Model.ID, toolName, dynCallID, nonDynCallID), + }, + }) + return err + })) + require.Len(t, step.InsertedMessages, 1) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + var ( + queuedIDs []int64 + queuedBodies []string + queuedCreatedBy []uuid.UUID + ) + for i := range queuedExtras { + body := fmt.Sprintf("queued-%s-%d", namePrefix, i) + createdBy := uuid.New() + queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, body), + ModelConfigID: uuid.NullUUID{UUID: f.Model.ID, Valid: true}, + CreatedBy: createdBy, + }) + require.NoError(t, err) + queuedIDs = append(queuedIDs, queued.ID) + queuedBodies = append(queuedBodies, body) + queuedCreatedBy = append(queuedCreatedBy, createdBy) + } + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + assistantToolCallMsgID: step.InsertedMessages[0].ID, + queuedMessageIDs: queuedIDs, + queuedMessageBodies: queuedBodies, + queuedMessageCreatedBy: queuedCreatedBy, + dynamicToolName: toolName, + pendingToolCallID: dynCallID, + pendingToolCallIDs: []string{dynCallID, nonDynCallID}, + } +} + +// seedInvalidWithQueue seeds Invalid with a single queued message so +// ReconcileInvalidState lands in E1 (non-empty queue) instead of E0. +func seedInvalidWithQueue(t *testing.T, f *testFixture) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + queuedBody := "queued-invalid" + queued := sendQueuedMessage(t, f, m, queuedBody) + require.NotNil(t, queued.QueuedMessage) + // Force the deliberately invalid running + archived combo on + // top of the queue. + chat, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + _, err = f.DB.UpdateChatExecutionState(ctx, database.UpdateChatExecutionStateParams{ + ID: chat.ID, + Status: database.ChatStatusRunning, + Archived: true, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }) + require.NoError(t, err) + return seededChat{ + chatID: chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, chat.ID), + queuedMessageIDs: []int64{queued.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + } +} + +// firstUserMessageID returns the lowest-id non-deleted user message +// on the chat. Most transition tests reuse this when they need a +// user message to edit. +func firstUserMessageID(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID) int64 { + t.Helper() + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + }) + require.NoError(t, err) + for _, m := range msgs { + if m.Role == database.ChatMessageRoleUser && !m.Deleted { + return m.ID + } + } + t.Fatalf("firstUserMessageID: chat %s has no user messages", chatID) + return 0 +} + +func firstAssistantMessageID(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID) int64 { + t.Helper() + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + }) + require.NoError(t, err) + for _, m := range msgs { + if m.Role == database.ChatMessageRoleAssistant && !m.Deleted { + return m.ID + } + } + t.Fatalf("firstAssistantMessageID: chat %s has no assistant messages", chatID) + return 0 +} + +// seedForEnterRequiresAction extends seedState for R0 and R1 with a +// chat that has dynamic_tools plus an assistant tool-call message in +// history. EnterRequiresAction's precondition rejects R0/R1 without +// pending dynamic tool calls, so the generic seedState path will not +// do. Other states fall through to the default seedState. +func seedForEnterRequiresAction(t *testing.T, f *testFixture, state chatstate.ExecutionState) seededChat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + switch state { + case chatstate.StateR0: + toolName := "ra_tool_r0" + callID := "call_" + uuid.NewString() + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + assistantToolCallMessage(t, f.Model.ID, toolName, callID), + }, + }) + return err + })) + require.Len(t, step.InsertedMessages, 1) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + assistantToolCallMsgID: step.InsertedMessages[0].ID, + dynamicToolName: toolName, + pendingToolCallID: callID, + pendingToolCallIDs: []string{callID}, + } + case chatstate.StateR1: + toolName := "ra_tool_r1" + callID := "call_" + uuid.NewString() + created := createTestChatWithDynamicTools(t, f, toolName) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + var step chatstate.CommitStepResult + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var err error + step, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + assistantToolCallMessage(t, f.Model.ID, toolName, callID), + }, + }) + return err + })) + // R0 -> R1 with a queued message. + queuedBody := "queued-for-RA-r1" + sm := sendQueuedMessage(t, f, m, queuedBody) + require.NotNil(t, sm.QueuedMessage) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + assistantToolCallMsgID: step.InsertedMessages[0].ID, + queuedMessageIDs: []int64{sm.QueuedMessage.ID}, + queuedMessageBodies: []string{queuedBody}, + dynamicToolName: toolName, + pendingToolCallID: callID, + pendingToolCallIDs: []string{callID}, + } + } + return seedState(t, f, state) +} + +// activeHistoryIDs returns the ids of non-deleted history messages +// for the chat in row-id order. Useful for verifying CommitStep, +// EditMessage replacement, and PromoteQueuedMessage head insertion. +func activeHistoryIDs(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID) []int64 { + t.Helper() + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: chatID, + }) + require.NoError(t, err) + out := make([]int64, 0, len(msgs)) + for _, m := range msgs { + if !m.Deleted { + out = append(out, m.ID) + } + } + return out +} + +func requireChatMessageByID(ctx context.Context, t *testing.T, f *testFixture, id int64) database.ChatMessage { + t.Helper() + msg, err := f.DB.GetChatMessageByID(ctx, id) + require.NoError(t, err) + return msg +} + +func requireQueuedMessageByID(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, id int64) database.ChatQueuedMessage { + t.Helper() + msg, err := f.DB.GetChatQueuedMessageByID(ctx, database.GetChatQueuedMessageByIDParams{ + ID: id, + ChatID: chatID, + }) + require.NoError(t, err) + return msg +} + +func requireQueuedMessageDeleted(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, id int64) { + t.Helper() + _, err := f.DB.GetChatQueuedMessageByID(ctx, database.GetChatQueuedMessageByIDParams{ + ID: id, + ChatID: chatID, + }) + require.Error(t, err) +} + +func assertFetchedUserMessage(ctx context.Context, t *testing.T, f *testFixture, msg database.ChatMessage) database.ChatMessage { + t.Helper() + fetched := requireChatMessageByID(ctx, t, f, msg.ID) + require.Equal(t, msg.ChatID, fetched.ChatID) + require.Equal(t, database.ChatMessageRoleUser, fetched.Role) + require.True(t, fetched.CreatedBy.Valid) + require.Equal(t, f.User.ID, fetched.CreatedBy.UUID) + require.True(t, fetched.ModelConfigID.Valid) + require.Equal(t, f.Model.ID, fetched.ModelConfigID.UUID) + require.Equal(t, chatprompt.CurrentContentVersion, fetched.ContentVersion) + return fetched +} + +func assertFetchedQueuedMessage(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, queued database.ChatQueuedMessage) database.ChatQueuedMessage { + t.Helper() + fetched := requireQueuedMessageByID(ctx, t, f, chatID, queued.ID) + require.Equal(t, chatID, fetched.ChatID) + require.Equal(t, f.User.ID, fetched.CreatedBy) + require.True(t, fetched.ModelConfigID.Valid) + require.Equal(t, f.Model.ID, fetched.ModelConfigID.UUID) + require.NotEmpty(t, fetched.Content) + return fetched +} + +func newActiveMessageIDs(base snapshotBaseline, after []int64) []int64 { + seen := make(map[int64]struct{}, len(base.historyIDs)) + for _, id := range base.historyIDs { + seen[id] = struct{}{} + } + out := make([]int64, 0, len(after)) + for _, id := range after { + if _, ok := seen[id]; !ok { + out = append(out, id) + } + } + return out +} + +// assertToolResultForCallNoError asserts that msg is a tool-result +// message that resolves a tool call with id wantCallID, is_error=false, +// and that the result JSON matches wantResultJSON. Complements +// assertToolResultForCall in synthetic_cancellation_test.go which +// asserts is_error=true. +func assertToolResultForCallNoError(t *testing.T, msg database.ChatMessage, wantCallID, wantResultJSON string) { + t.Helper() + require.Equal(t, database.ChatMessageRoleTool, msg.Role) + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err) + require.NotEmpty(t, parts) + var found bool + for _, p := range parts { + if p.Type != codersdk.ChatMessagePartTypeToolResult { + continue + } + require.Equal(t, wantCallID, p.ToolCallID, "tool-call id matches") + require.False(t, p.IsError, "CompleteRequiresAction tool result must not be is_error") + require.JSONEq(t, wantResultJSON, string(p.Result), "CompleteRequiresAction tool result JSON matches submitted output") + found = true + } + require.True(t, found, "expected at least one tool-result part") +} + +// assertChatMessageText asserts that the persisted content of msg +// decodes to a single text part with the supplied body. Used by +// matrix cases that need to verify the actual text submitted via +// SendMessage / EditMessage / CommitStep, or the text that was +// promoted out of the queue into history. +func assertChatMessageText(t *testing.T, msg database.ChatMessage, want string) { + t.Helper() + parts, err := chatprompt.ParseContent(msg) + require.NoError(t, err, "parse chat message content") + require.Len(t, parts, 1, "expected exactly one content part") + require.Equal(t, codersdk.ChatMessagePartTypeText, parts[0].Type, + "expected a text content part") + require.Equal(t, want, parts[0].Text, "unexpected chat message text") +} + +// assertQueuedMessageText asserts that the JSON content of queued +// decodes to a single text part with the supplied body. Used by +// matrix cases that need to verify the body inserted into +// chat_queued_messages via SendMessage. +func assertQueuedMessageText(t *testing.T, queued database.ChatQueuedMessage, want string) { + t.Helper() + var parts []codersdk.ChatMessagePart + require.NoError(t, json.Unmarshal(queued.Content, &parts), "unmarshal queued content") + require.Len(t, parts, 1, "expected exactly one queued content part") + require.Equal(t, codersdk.ChatMessagePartTypeText, parts[0].Type, + "expected a text content part") + require.Equal(t, want, parts[0].Text, "unexpected queued message text") +} + +// assertQueueBodiesInOrder fetches the queued messages for the chat +// in queue order and asserts each row's text body matches the +// supplied bodies. Used by matrix cases that need to verify the +// remaining queue content after a promote / finish-turn / +// finish-interruption. +func assertQueueBodiesInOrder(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, want []string) { + t.Helper() + rows, err := f.DB.GetChatQueuedMessagesByPosition(ctx, chatID) + require.NoError(t, err) + require.Len(t, rows, len(want), "queue length must match expected bodies") + for i, r := range rows { + assertQueuedMessageText(t, r, want[i]) + } +} + +// snapshotBaseline records the chat's snapshot_version and the +// publisher's recorded channel count immediately before a transition +// runs. Tests use it to verify either a single snapshot bump and one +// chat:update on success, or zero mutation and zero publishes on +// failure. +type snapshotBaseline struct { + exists bool + chat database.Chat + snapshot int64 + historyVersion int64 + queueVersion int64 + retryStateVersion int64 + generationAttempt int64 + queueCount int64 + queueIDs []int64 + historyIDs []int64 + channels int +} + +func captureBaseline(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat) snapshotBaseline { + t.Helper() + base := snapshotBaseline{ + exists: seeded.exists, + channels: len(f.Pub.channels), + } + if !seeded.exists { + return base + } + chat, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + base.chat = chat + base.snapshot = chat.SnapshotVersion + base.historyVersion = chat.HistoryVersion + base.queueVersion = chat.QueueVersion + base.retryStateVersion = chat.RetryStateVersion + base.generationAttempt = chat.GenerationAttempt + base.queueIDs = queuedIDsByPosition(ctx, t, f, seeded.chatID) + count, err := f.DB.CountChatQueuedMessages(ctx, seeded.chatID) + require.NoError(t, err) + base.queueCount = count + base.historyIDs = activeHistoryIDs(ctx, t, f, seeded.chatID) + return base +} + +// assertSnapshotBumpedOnce asserts that one Update committed; that is, +// snapshot_version advanced by exactly one and the publisher saw at +// least one chat:update on the per-chat channel after the baseline. +func assertSnapshotBumpedOnce(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, base snapshotBaseline) { + t.Helper() + after, err := f.DB.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.Equal(t, base.snapshot+1, after.SnapshotVersion, "snapshot_version must bump exactly once") + channel := coderdpubsub.ChatStateUpdateChannel(chatID) + found := false + for _, c := range f.Pub.channels[base.channels:] { + if c == channel { + found = true + break + } + } + require.True(t, found, "expected one chat:update on %s after commit", channel) +} + +// assertNoMutationOrPublish asserts a failed transition rolled back +// the automatic snapshot bump and published nothing. +func assertNoMutationOrPublish(ctx context.Context, t *testing.T, f *testFixture, chatID uuid.UUID, base snapshotBaseline) { + t.Helper() + require.Equal(t, base.channels, len(f.Pub.channels), "failed transition must not publish") + if base.exists { + after, err := f.DB.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.Equal(t, base.snapshot, after.SnapshotVersion, "failed transition must not advance snapshot_version") + } +} diff --git a/coderd/x/chatd/chatstate/transitions_matrix_test.go b/coderd/x/chatd/chatstate/transitions_matrix_test.go new file mode 100644 index 0000000000..ef35090f7d --- /dev/null +++ b/coderd/x/chatd/chatstate/transitions_matrix_test.go @@ -0,0 +1,1845 @@ +package chatstate_test + +import ( + "context" + "encoding/json" + "fmt" + "slices" + "sync" + "testing" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// Matrix harness: spec types, scenario labels, appliers, case runners, +// and the single entry point that walks the production transition +// matrix to confirm every allowed combination has positive coverage +// and every disallowed combination surfaces the right sentinel error. + +// scenario is a typed, semantic label that distinguishes positive +// matrix cases that share the same (transition, from, want) key. +// Empty scenario is fine when no label is needed. The constants +// below enumerate every label used by matrixCases(). +type scenario string + +const ( + // scenarioQueue marks SendMessage cases driven by + // BusyBehaviorQueue. + scenarioQueue scenario = "queue" + // scenarioInterrupt marks SendMessage cases driven by + // BusyBehaviorInterrupt. + scenarioInterrupt scenario = "interrupt" + // scenarioMulti marks cases seeded with multiple queued + // messages so the post-mutation queue stays non-empty. + scenarioMulti scenario = "multi" + // scenarioHeadTarget marks multi-queued PromoteQueuedMessage + // cases that target the queue head. For R1/I1 head-target is + // reorder-only: no rows are updated, so queue order and + // queue_version are unchanged. For E1/A1 head-target still + // pops the head into history. + scenarioHeadTarget scenario = "head_target" + // scenarioNonHead marks multi-queued PromoteQueuedMessage cases + // that target a non-head queued message so the target moves to + // the head and queue_version advances. + scenarioNonHead scenario = "non_head" + // scenarioWithQueue marks ReconcileInvalidState cases seeded + // with a non-empty queue. + scenarioWithQueue scenario = "with_queue" + // scenarioRejectNonDynamicOutstandingToolCall marks the + // FinishInterruption case that exercises the precondition + // rejecting outstanding non-dynamic tool calls. + scenarioRejectNonDynamicOutstandingToolCall scenario = "reject_non_dynamic_outstanding_tool_call" +) + +func transitionAllowed(tr chatstate.Transition, from chatstate.ExecutionState) bool { + return slices.Contains(chatstate.AllowedExecutionTransitionsFrom(from), tr) +} + +// expectedErrorForDisallowed returns the sentinel chatstate package +// returns when a transition is attempted from a state where the +// matrix forbids it. N (missing chat) becomes ErrChatNotFound; +// Invalid becomes ErrInvalidState (except for ReconcileInvalidState +// which is allowed); everything else becomes ErrTransitionNotAllowed. +func expectedErrorForDisallowed(tr chatstate.Transition, from chatstate.ExecutionState) error { + switch from { + case chatstate.StateN: + if tr == chatstate.TransitionCreateChat { + // CreateChat is not exercised through ChatMachine.Update, + // so this branch is unused in practice. Returning the + // not-allowed sentinel keeps the helper total. + return chatstate.ErrTransitionNotAllowed + } + return chatstate.ErrChatNotFound + case chatstate.StateInvalid: + if tr == chatstate.TransitionReconcileInvalidState { + return nil + } + return chatstate.ErrInvalidState + } + return chatstate.ErrTransitionNotAllowed +} + +// Transition appliers +// +// Each transition has one default applier that exercises it with +// inputs derived from the seeded chat. Positive case specs reuse these +// appliers unless a case needs a different input shape (for example, +// SendMessage queue versus interrupt from the same source state). +// The disallowed coverage path also uses these defaults. + +func applySetArchived(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, from chatstate.ExecutionState, _ *transitionCaseResult) error { + t.Helper() + // Archived states unarchive, others archive. For disallowed + // states the value does not matter; the transition fails first. + archived := true + switch from { + case chatstate.StateXW, chatstate.StateXE0, chatstate.StateXE1: + archived = false + } + _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: archived}) + return err +} + +func applySendMessageQueue(t *testing.T, f *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.sendMessage, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("sm-queue", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err +} + +func applySendMessageInterrupt(t *testing.T, f *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.sendMessage, err = tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("sm-interrupt", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorInterrupt, + }) + return err +} + +func applyEditMessage(t *testing.T, f *testFixture, tx *chatstate.Tx, seeded seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + content := mustMarshalParts(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}) + var err error + result.editMessage, err = tx.EditMessage(chatstate.EditMessageInput{ + MessageID: seeded.initialUserMessageID, + CreatedBy: f.User.ID, + Content: content, + APIKeyID: f.apiKeyID(), + }) + return err +} + +func applyDeleteQueuedMessage(t *testing.T, _ *testFixture, tx *chatstate.Tx, seeded seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var targetQueueID int64 + if len(seeded.queuedMessageIDs) > 0 { + targetQueueID = seeded.queuedMessageIDs[0] + } + var err error + result.deleteQueuedMessage, err = tx.DeleteQueuedMessage(chatstate.DeleteQueuedMessageInput{ + QueuedMessageID: targetQueueID, + }) + return err +} + +func applyPromoteQueuedMessage(t *testing.T, _ *testFixture, tx *chatstate.Tx, seeded seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var targetQueueID int64 + if len(seeded.queuedMessageIDs) > 0 { + targetQueueID = seeded.queuedMessageIDs[0] + } + var err error + result.promoteQueuedMessage, err = tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: targetQueueID, + }) + return err +} + +func applyInterrupt(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.interrupt, err = tx.Interrupt(chatstate.InterruptInput{Reason: "test"}) + return err +} + +func applyCompleteRequiresAction(t *testing.T, f *testFixture, tx *chatstate.Tx, seeded seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var results []chatstate.ToolResultInput + if seeded.pendingToolCallID != "" { + results = []chatstate.ToolResultInput{{ + ToolCallID: seeded.pendingToolCallID, + Output: json.RawMessage(`{"ok":true}`), + IsError: false, + }} + } + var err error + result.completeRequiresAction, err = tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ + CreatedBy: f.User.ID, + ModelConfigID: f.Model.ID, + Results: results, + }) + return err +} + +func applyRecordGenerationAttempt(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.recordGenerationAttempt, err = tx.RecordGenerationAttempt(chatstate.RecordGenerationAttemptInput{}) + return err +} + +func applyRecordRetryState(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.recordRetryState, err = tx.RecordRetryState(chatstate.RecordRetryStateInput{ + RetryState: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`), + Valid: true, + }, + }) + return err +} + +func applyCommitStep(t *testing.T, f *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + assistant := userTextMessage("assistant", f.User.ID, f.Model.ID) + assistant.Role = database.ChatMessageRoleAssistant + var err error + result.commitStep, err = tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistant}, + }) + return err +} + +func applyEnterRequiresAction(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.enterRequiresAction, err = tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err +} + +func applyFinishInterruption(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.finishInterruption, err = tx.FinishInterruption(chatstate.FinishInterruptionInput{}) + return err +} + +func applyFinishTurn(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.finishTurn, err = tx.FinishTurn(chatstate.FinishTurnInput{}) + return err +} + +func applyFinishError(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.finishError, err = tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"finish-error"}`), + Valid: true, + }, + }) + return err +} + +func applyCancelRequiresAction(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.cancelRequiresAction, err = tx.CancelRequiresAction(chatstate.CancelRequiresActionInput{ + Reason: "cancel from test", + }) + return err +} + +func applyReconcileInvalidState(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.reconcileInvalidState, err = tx.ReconcileInvalidState(chatstate.ReconcileInvalidStateInput{}) + return err +} + +// defaultApplier returns the canonical applier for tr. Used by the +// disallowed coverage path where the input shape does not matter +// because the transition fails before the inputs are consumed. +func defaultApplier(tr chatstate.Transition) applierFn { + switch tr { + case chatstate.TransitionSetArchived: + return applySetArchived + case chatstate.TransitionSendMessage: + return applySendMessageQueue + case chatstate.TransitionEditMessage: + return applyEditMessage + case chatstate.TransitionDeleteQueuedMessage: + return applyDeleteQueuedMessage + case chatstate.TransitionPromoteQueuedMessage: + return applyPromoteQueuedMessage + case chatstate.TransitionInterrupt: + return applyInterrupt + case chatstate.TransitionCompleteRequiresAction: + return applyCompleteRequiresAction + case chatstate.TransitionRecordGenerationAttempt: + return applyRecordGenerationAttempt + case chatstate.TransitionRecordRetryState: + return applyRecordRetryState + case chatstate.TransitionCommitStep: + return applyCommitStep + case chatstate.TransitionEnterRequiresAction: + return applyEnterRequiresAction + case chatstate.TransitionFinishInterruption: + return applyFinishInterruption + case chatstate.TransitionFinishTurn: + return applyFinishTurn + case chatstate.TransitionFinishError: + return applyFinishError + case chatstate.TransitionCancelRequiresAction: + return applyCancelRequiresAction + case chatstate.TransitionReconcileInvalidState: + return applyReconcileInvalidState + } + return nil +} + +// mustMarshalParts is a tiny test helper that fails the test on +// marshal error rather than forcing every call site to handle it. +func mustMarshalParts(t *testing.T, parts []codersdk.ChatMessagePart) pqtype.NullRawMessage { + t.Helper() + raw, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return raw +} + +// Case-level transition matrix spec. +// +// Each entry in matrixCases is one positive (transition, from, want) +// triple. The coverage key is (transition, from, want); scenario is a +// readability-and-semantic suffix for the subtest name that +// distinguishes multiple cases sharing the same coverage key. +// Disallowed combinations are enumerated separately from +// AllowedExecutionTransitionsFrom and AllowedExecutionTransitionOutputs. + +type transitionCaseResult struct { + sendMessage chatstate.SendMessageResult + editMessage chatstate.EditMessageResult + deleteQueuedMessage chatstate.DeleteQueuedMessageResult + promoteQueuedMessage chatstate.PromoteQueuedMessageResult + interrupt chatstate.InterruptResult + completeRequiresAction chatstate.CompleteRequiresActionResult + recordGenerationAttempt chatstate.RecordGenerationAttemptResult + recordRetryState chatstate.RecordRetryStateResult + commitStep chatstate.CommitStepResult + enterRequiresAction chatstate.EnterRequiresActionResult + finishInterruption chatstate.FinishInterruptionResult + finishTurn chatstate.FinishTurnResult + finishError chatstate.FinishErrorResult + cancelRequiresAction chatstate.CancelRequiresActionResult + reconcileInvalidState chatstate.ReconcileInvalidStateResult +} + +type applierFn func(t *testing.T, f *testFixture, tx *chatstate.Tx, seeded seededChat, from chatstate.ExecutionState, result *transitionCaseResult) error + +type assertFn func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) + +// seederFn produces a seededChat for a case. Cases that omit a custom +// seeder use seedState by default. Custom seeders are required when +// the case needs more than one queued message, an Invalid chat with a +// non-empty queue, or a transition that needs a fresh A0/A1 seed. +type seederFn func(t *testing.T, f *testFixture, from chatstate.ExecutionState) seededChat + +type transitionCaseSpec struct { + transition chatstate.Transition + from chatstate.ExecutionState + want chatstate.ExecutionState + // scenario is a semantic label appended to the subtest name + // when the same (transition, from, want) key needs to run more + // than once. It is not part of the coverage key but is part of + // the duplicate-detection key. + scenario scenario + + seed seederFn + apply applierFn + assert assertFn + assertFailure func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, err error) +} + +// caseKey is the unit of coverage for positive cases. scenario is +// intentionally not part of the key so cases with different scenarios +// can still satisfy the same coverage cell. +type caseKey struct { + transition chatstate.Transition + from chatstate.ExecutionState + want chatstate.ExecutionState +} + +// fullCaseKey extends caseKey with scenario. Used for duplicate +// detection: two cases must not share the same full key. +type fullCaseKey struct { + transition chatstate.Transition + from chatstate.ExecutionState + want chatstate.ExecutionState + scenario scenario +} + +// queueShape selects the seed variant for transition case builders. +// A typed enum is used instead of a bool to avoid the revive +// flag-parameter rule and to make call sites self-documenting. +type queueShape int + +const ( + // queueShapeDefault routes through seedState, which produces the + // canonical single-queued seed for queue-bearing states and the + // empty queue for non-queue states. + queueShapeDefault queueShape = iota + // queueShapeMulti routes through seedStateMultiQueued (or + // seedInvalidWithQueue for ReconcileInvalidState) so the + // post-mutation queue can remain non-empty. + queueShapeMulti +) + +func (s queueShape) isMulti() bool { return s == queueShapeMulti } + +func (s transitionCaseSpec) key() caseKey { + return caseKey{transition: s.transition, from: s.from, want: s.want} +} + +func (s transitionCaseSpec) fullKey() fullCaseKey { + return fullCaseKey{ + transition: s.transition, + from: s.from, + want: s.want, + scenario: s.scenario, + } +} + +func (s transitionCaseSpec) subtestName() string { + name := fmt.Sprintf("%s/%s_to_%s", s.transition, s.from, s.want) + if s.scenario != "" { + name += "/" + string(s.scenario) + } + return name +} + +// disallowedCaseKey is the unit of coverage for negative cases. +type disallowedCaseKey struct { + transition chatstate.Transition + from chatstate.ExecutionState +} + +// remainingExcluding returns ids with the entry at exclude removed. +// The order of the surviving entries is preserved. +func remainingExcluding(ids []int64, exclude int) []int64 { + out := make([]int64, 0, len(ids)) + for i, id := range ids { + if i == exclude { + continue + } + out = append(out, id) + } + return out +} + +// remainingBodiesExcluding returns bodies with the entry at exclude +// removed. The order of the surviving entries is preserved. +func remainingBodiesExcluding(bodies []string, exclude int) []string { + out := make([]string, 0, len(bodies)) + for i, b := range bodies { + if i == exclude { + continue + } + out = append(out, b) + } + return out +} + +// Test runner + +// runPositiveCase seeds the chat, runs the transition, and asserts the +// post-state plus case-specific effects. +func runPositiveCase(t *testing.T, spec transitionCaseSpec) { + t.Helper() + require.NotNil(t, spec.apply, "case %s missing apply", spec.subtestName()) + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + + seeder := spec.seed + if seeder == nil { + seeder = seedState + } + seeded := seeder(t, f, spec.from) + if seeded.exists { + require.Equal(t, spec.from, f.classify(ctx, t, seeded.chatID), + "seed must land in %s", spec.from) + } + base := captureBaseline(ctx, t, f, seeded) + + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + var result transitionCaseResult + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + return spec.apply(t, f, tx, seeded, spec.from, &result) + }) + if spec.assertFailure != nil { + spec.assertFailure(ctx, t, f, seeded, base, err) + return + } + require.NoError(t, err, "%s from %s must succeed", spec.transition, spec.from) + assertSnapshotBumpedOnce(ctx, t, f, seeded.chatID, base) + require.Equal(t, spec.want, f.classify(ctx, t, seeded.chatID), + "%s: %s -> %s", spec.transition, spec.from, spec.want) + if spec.assert != nil { + spec.assert(ctx, t, f, seeded, base, result) + } +} + +// runDisallowedCase seeds the chat, runs the transition with default +// inputs, and asserts that the chatstate package surfaces the right +// sentinel error and rolled the snapshot bump back. +func runDisallowedCase(t *testing.T, tr chatstate.Transition, from chatstate.ExecutionState) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + seeded := seedState(t, f, from) + if seeded.exists { + require.Equal(t, from, f.classify(ctx, t, seeded.chatID), + "disallowed seed must land in %s", from) + } + base := captureBaseline(ctx, t, f, seeded) + + applier := defaultApplier(tr) + require.NotNil(t, applier, "no default applier for transition %s", tr) + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + var result transitionCaseResult + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + return applier(t, f, tx, seeded, from, &result) + }) + + if tr == chatstate.TransitionReconcileInvalidState && from != chatstate.StateN { + // ReconcileInvalidState does not use requireFromAllowed. + // It hits loadState successfully, sees the state is not + // Invalid, and returns a TransitionError directly. + require.Error(t, err) + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te, + "reconcile from non-invalid state must return TransitionError") + require.Equal(t, chatstate.TransitionReconcileInvalidState, te.Transition) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) + return + } + + expectErr := expectedErrorForDisallowed(tr, from) + require.Error(t, err) + require.ErrorIs(t, err, expectErr) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +// TestTransitionMatrix_AllCombinations is the single entry point for +// the case-level transition matrix coverage. Each positive case in +// matrixCases() is one (transition, from, want) triple with a focused +// effect assertion. Disallowed combinations are enumerated from +// transition.go to confirm every non-CreateChat (transition, from) +// pair outside the allowed set surfaces the right sentinel error. +// +// After all parallel subtests complete the test verifies that the +// positive coverage matches AllowedExecutionTransitionOutputs (no +// missing key, no unexpected key) and that every disallowed +// (transition, from) pair was exercised exactly once. +// +// Input-specific rejection tests live in TestTransitionInputValidation +// and are intentionally not part of this matrix entry point so the +// matrix focus stays on positive cases and generated disallowed cases. +func TestTransitionMatrix_AllCombinations(t *testing.T) { + t.Parallel() + + cases := matrixCases() + + // Detect duplicate full keys and duplicate subtest names. The + // coverage key intentionally ignores scenario, so two cases may + // share the same (transition, from, want) only when their + // scenarios differ. + seenFullKeys := make(map[fullCaseKey]string, len(cases)) + seenNames := make(map[string]struct{}, len(cases)) + for _, tc := range cases { + full := tc.fullKey() + name := tc.subtestName() + if prev, ok := seenFullKeys[full]; ok { + t.Fatalf("duplicate matrix case %+v: previous %s, new %s", full, prev, name) + } + seenFullKeys[full] = name + if _, ok := seenNames[name]; ok { + t.Fatalf("duplicate matrix subtest name %s", name) + } + seenNames[name] = struct{}{} + } + + // Build the expected positive set from the matrix in + // transition.go. CreateChat is intentionally excluded because + // it is not exercised via ChatMachine.Update. + expectedPositive := make(map[caseKey]struct{}) + for _, from := range chatstate.AllExecutionStates { + for _, tr := range chatstate.AllowedExecutionTransitionsFrom(from) { + if tr == chatstate.TransitionCreateChat { + continue + } + for _, to := range chatstate.AllowedExecutionTransitionOutputs(from, tr) { + expectedPositive[caseKey{transition: tr, from: from, want: to}] = struct{}{} + } + } + } + + // Build the expected disallowed set: for each non-CreateChat + // transition, every state where the transition is not allowed. + expectedDisallowed := make(map[disallowedCaseKey]struct{}) + for _, tr := range chatstate.AllExecutionTransitions { + if tr == chatstate.TransitionCreateChat { + continue + } + for _, from := range chatstate.AllExecutionStates { + if transitionAllowed(tr, from) { + continue + } + expectedDisallowed[disallowedCaseKey{transition: tr, from: from}] = struct{}{} + } + } + + // Validate that every case in matrixCases describes a + // (transition, from, want) combination that the matrix actually + // admits. This guards against typos in matrixCases wiring up + // nonsense cases that happen to compile. + for _, tc := range cases { + if tc.assertFailure != nil { + continue + } + key := tc.key() + _, ok := expectedPositive[key] + require.True(t, ok, + "case %s is not in the allowed (transition, from, want) set", tc.subtestName()) + } + + // actualPositive and actualDisallowed are mutated under mu from + // parallel subtests. The final comparison runs in t.Cleanup, + // which fires only after every parallel child finishes. + var mu sync.Mutex + actualPositive := make(map[caseKey]struct{}, len(expectedPositive)) + actualDisallowed := make(map[disallowedCaseKey]struct{}, len(expectedDisallowed)) + + t.Cleanup(func() { + mu.Lock() + defer mu.Unlock() + for k := range expectedPositive { + if _, ok := actualPositive[k]; !ok { + t.Errorf("matrix coverage: missing positive case %+v", k) + } + } + for k := range actualPositive { + if _, ok := expectedPositive[k]; !ok { + t.Errorf("matrix coverage: unexpected positive case %+v", k) + } + } + for k := range expectedDisallowed { + if _, ok := actualDisallowed[k]; !ok { + t.Errorf("matrix coverage: missing disallowed case %+v", k) + } + } + for k := range actualDisallowed { + if _, ok := expectedDisallowed[k]; !ok { + t.Errorf("matrix coverage: unexpected disallowed case %+v", k) + } + } + }) + + // Positive cases: one parallel subtest per case. + t.Run("positive", func(t *testing.T) { + t.Parallel() + for _, tc := range cases { + tc := tc + t.Run(tc.subtestName(), func(t *testing.T) { + t.Parallel() + if tc.assertFailure == nil { + mu.Lock() + actualPositive[tc.key()] = struct{}{} + mu.Unlock() + } + runPositiveCase(t, tc) + }) + } + }) + + // Negative cases: one parallel subtest per (transition, from) + // pair where the transition is not allowed. Iterate over + // transitions in canonical order, and within each transition + // iterate states in canonical AllExecutionStates order, so + // subtest names are stable. + t.Run("disallowed", func(t *testing.T) { + t.Parallel() + // Sort disallowed keys for deterministic subtest names. + // AllExecutionTransitions and AllExecutionStates are + // already canonical, so iterate in their order. + for _, tr := range chatstate.AllExecutionTransitions { + tr := tr + if tr == chatstate.TransitionCreateChat { + continue + } + t.Run(string(tr), func(t *testing.T) { + t.Parallel() + for _, from := range chatstate.AllExecutionStates { + from := from + if transitionAllowed(tr, from) { + continue + } + t.Run(string(from), func(t *testing.T) { + t.Parallel() + mu.Lock() + actualDisallowed[disallowedCaseKey{transition: tr, from: from}] = struct{}{} + mu.Unlock() + runDisallowedCase(t, tr, from) + }) + } + }) + } + }) +} + +// Positive case specs. +// +// Each case asserts (at minimum) the resulting classified post-state +// matches want, plus one transition-specific effect. Helpers reused +// from other tests handle the snapshot bump and the chat:update +// publish; per-case assertions focus on what the transition meant to +// change. + +func matrixCases() []transitionCaseSpec { + return []transitionCaseSpec{ + // SetArchived cases: each archived/unarchived pair flips the + // archived flag, preserves status, history and last_error, + // and does not insert anything new. + setArchivedCase(chatstate.StateW, chatstate.StateXW, database.ChatStatusWaiting), + setArchivedCase(chatstate.StateE0, chatstate.StateXE0, database.ChatStatusError), + setArchivedCase(chatstate.StateE1, chatstate.StateXE1, database.ChatStatusError), + setArchivedCase(chatstate.StateXW, chatstate.StateW, database.ChatStatusWaiting), + setArchivedCase(chatstate.StateXE0, chatstate.StateE0, database.ChatStatusError), + setArchivedCase(chatstate.StateXE1, chatstate.StateE1, database.ChatStatusError), + + // SendMessage(queue) cases: idle states insert directly, + // busy states append to the queue tail. + sendMessageQueueCase(chatstate.StateW, chatstate.StateR0, true, 0), + sendMessageQueueCase(chatstate.StateE0, chatstate.StateR0, true, 0), + // E1 promotes the queue head and queues the new tail, so + // the net queue delta is zero. + sendMessageQueueCase(chatstate.StateE1, chatstate.StateR1, false, 0), + sendMessageQueueCase(chatstate.StateR0, chatstate.StateR1, false, +1), + sendMessageQueueCase(chatstate.StateR1, chatstate.StateR1, false, +1), + sendMessageQueueCase(chatstate.StateI0, chatstate.StateI1, false, +1), + sendMessageQueueCase(chatstate.StateI1, chatstate.StateI1, false, +1), + sendMessageQueueCase(chatstate.StateA0, chatstate.StateA1, false, +1), + sendMessageQueueCase(chatstate.StateA1, chatstate.StateA1, false, +1), + + // SendMessage(interrupt) cases. The interrupt applier runs + // with body "sm-interrupt" so the assertion can prove the + // interrupt input path was taken. From W/E0/E1/I0/I1 the + // resulting (transition, from, want) coverage key is + // identical to the queue case, but we still exercise the + // interrupt entry point to guard against a future bug where + // it stops routing through the correct direct-insert / + // queue-tail / promotion paths. From the busy R0/R1/A0/A1 + // states the interrupt destination differs from the queue + // destination so the scenario label is the only case for that key. + sendMessageInterruptCase(chatstate.StateW, chatstate.StateR0), + sendMessageInterruptCase(chatstate.StateE0, chatstate.StateR0), + sendMessageInterruptCase(chatstate.StateE1, chatstate.StateR1), + sendMessageInterruptCase(chatstate.StateR0, chatstate.StateI1), + sendMessageInterruptCase(chatstate.StateR1, chatstate.StateI1), + sendMessageInterruptCase(chatstate.StateI0, chatstate.StateI1), + sendMessageInterruptCase(chatstate.StateI1, chatstate.StateI1), + sendMessageInterruptCase(chatstate.StateA0, chatstate.StateR1), + sendMessageInterruptCase(chatstate.StateA1, chatstate.StateR1), + + // EditMessage cases: every allowed source state lands in R0 + // with the queue cleared, last_error reset, and a + // replacement user message in active history. + editMessageCase(chatstate.StateW), + editMessageCase(chatstate.StateE0), + editMessageCase(chatstate.StateE1), + editMessageCase(chatstate.StateR0), + editMessageCase(chatstate.StateR1), + editMessageCase(chatstate.StateI0), + editMessageCase(chatstate.StateI1), + editMessageCase(chatstate.StateA0), + editMessageCase(chatstate.StateA1), + + // DeleteQueuedMessage cases. Empty-tail want collapses the + // classified state (E1->E0, R1->R0, I1->I0, A1->A0). The + // non-empty-tail cases need a multi-queued seed. + deleteQueuedCase(chatstate.StateE1, chatstate.StateE0, queueShapeDefault), + deleteQueuedCase(chatstate.StateE1, chatstate.StateE1, queueShapeMulti), + deleteQueuedCase(chatstate.StateR1, chatstate.StateR0, queueShapeDefault), + deleteQueuedCase(chatstate.StateR1, chatstate.StateR1, queueShapeMulti), + deleteQueuedCase(chatstate.StateI1, chatstate.StateI0, queueShapeDefault), + deleteQueuedCase(chatstate.StateI1, chatstate.StateI1, queueShapeMulti), + deleteQueuedCase(chatstate.StateA1, chatstate.StateA0, queueShapeDefault), + deleteQueuedCase(chatstate.StateA1, chatstate.StateA1, queueShapeMulti), + + // PromoteQueuedMessage cases. E1/A1 pop the head into + // history; R1/I1 only reorder the queue without + // inserting history. R1/I1 has both a head-target + // scenario (zero rows updated, queue_version unchanged) + // and a non-head scenario (target moves to head, + // queue_version advances). + promoteQueuedCase(chatstate.StateE1, chatstate.StateR0, queueShapeDefault, 0), + promoteQueuedCase(chatstate.StateE1, chatstate.StateR1, queueShapeMulti, 0), + promoteQueuedCase(chatstate.StateR1, chatstate.StateI1, queueShapeMulti, 0), + promoteQueuedCase(chatstate.StateR1, chatstate.StateI1, queueShapeMulti, 1), + promoteQueuedCase(chatstate.StateI1, chatstate.StateI1, queueShapeMulti, 1), + promoteQueuedCase(chatstate.StateA1, chatstate.StateR0, queueShapeDefault, 0), + promoteQueuedCase(chatstate.StateA1, chatstate.StateR1, queueShapeMulti, 0), + + // Interrupt cases. + interruptCase(chatstate.StateR0, chatstate.StateI0), + interruptCase(chatstate.StateR1, chatstate.StateI1), + interruptCase(chatstate.StateA0, chatstate.StateR0), + interruptCase(chatstate.StateA1, chatstate.StateR1), + + // CompleteRequiresAction cases: A0->R0, A1->R1. + completeRequiresActionCase(chatstate.StateA0, chatstate.StateR0), + completeRequiresActionCase(chatstate.StateA1, chatstate.StateR1), + + // CancelRequiresAction cases: A0->R0, A1->R1. + cancelRequiresActionCase(chatstate.StateA0, chatstate.StateR0), + cancelRequiresActionCase(chatstate.StateA1, chatstate.StateR1), + + // RecordGenerationAttempt cases: from-state preserved. + recordGenerationAttemptCase(chatstate.StateR0), + recordGenerationAttemptCase(chatstate.StateR1), + + // RecordRetryState cases: from-state preserved. + recordRetryStateCase(chatstate.StateR0), + recordRetryStateCase(chatstate.StateR1), + + // CommitStep cases: from-state preserved, history grows by + // one message. + commitStepCase(chatstate.StateR0), + commitStepCase(chatstate.StateR1), + + // EnterRequiresAction cases. R0/R1 need a pending tool call + // seeded; use seedForEnterRequiresAction so the precondition + // is met. + enterRequiresActionCase(chatstate.StateR0, chatstate.StateA0), + enterRequiresActionCase(chatstate.StateR1, chatstate.StateA1), + + // FinishInterruption cases: I0->W, I1->R0 (head promoted into + // history when only one queued), I1->R1 (with more than one + // queued, the head is promoted but the queue stays + // non-empty). + finishInterruptionCase(chatstate.StateI0, chatstate.StateW, queueShapeDefault), + finishInterruptionRejectsOutstandingToolCallCase(), + finishInterruptionCase(chatstate.StateI1, chatstate.StateR0, queueShapeDefault), + finishInterruptionCase(chatstate.StateI1, chatstate.StateR1, queueShapeMulti), + + // FinishTurn cases. + finishTurnCase(chatstate.StateR0, chatstate.StateW, queueShapeDefault), + finishTurnCase(chatstate.StateR1, chatstate.StateR0, queueShapeDefault), + finishTurnCase(chatstate.StateR1, chatstate.StateR1, queueShapeMulti), + + // FinishError cases. + finishErrorCase(chatstate.StateR0, chatstate.StateE0), + finishErrorCase(chatstate.StateR1, chatstate.StateE1), + + // ReconcileInvalidState cases: Invalid with empty queue + // lands in E0; Invalid with non-empty queue lands in E1. + reconcileInvalidStateCase(chatstate.StateE0, queueShapeDefault), + reconcileInvalidStateCase(chatstate.StateE1, queueShapeMulti), + } +} + +func setArchivedCase(from, want chatstate.ExecutionState, wantStatus database.ChatStatus) transitionCaseSpec { + wantArchived := false + switch want { + case chatstate.StateXW, chatstate.StateXE0, chatstate.StateXE1: + wantArchived = true + } + return transitionCaseSpec{ + transition: chatstate.TransitionSetArchived, + from: from, + want: want, + apply: applySetArchived, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + _ = result + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, wantArchived, after.Archived, + "SetArchived must set archived=%v", wantArchived) + require.Equal(t, wantStatus, after.Status, + "SetArchived preserves chat status") + require.Equal(t, base.chat.LastError, after.LastError, + "SetArchived preserves last_error") + require.Equal(t, base.historyVersion, after.HistoryVersion, + "SetArchived does not insert history") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "SetArchived leaves history messages unchanged") + require.Equal(t, base.queueVersion, after.QueueVersion, + "SetArchived does not mutate queued messages") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "SetArchived leaves queued messages unchanged") + }, + } +} + +func sendMessageQueueCase(from, want chatstate.ExecutionState, directInsert bool, queueDelta int64) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionSendMessage, + from: from, + want: want, + scenario: scenarioQueue, + apply: applySendMessageQueue, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterQueue, err := f.DB.CountChatQueuedMessages(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + + require.Equal(t, base.queueCount+queueDelta, afterQueue, + "SendMessage(queue): unexpected queue count delta") + + switch { + case directInsert: + // W/E0: insert directly into history, no queue + // mutation. result.InsertedMessages contains exactly + // the new user message. + require.Len(t, result.sendMessage.InsertedMessages, 1, + "SendMessage(queue) into W/E0 inserts exactly one history message") + require.Nil(t, result.sendMessage.QueuedMessage, + "SendMessage(queue) into W/E0 does not queue") + inserted := assertFetchedUserMessage(ctx, t, f, result.sendMessage.InsertedMessages[0]) + require.Equal(t, seeded.chatID, inserted.ChatID) + assertChatMessageText(t, inserted, "sm-queue") + require.False(t, after.LastError.Valid, + "SendMessage(queue) clears last_error when transitioning out of an error state") + require.Equal(t, database.ChatStatusRunning, after.Status, + "SendMessage(queue) into W/E0 lands in running") + require.Equal(t, base.queueIDs, afterQueueIDs, + "SendMessage(queue) into W/E0 must not touch queued messages") + require.Equal(t, base.queueVersion, after.QueueVersion, + "SendMessage(queue) into W/E0 must not bump queue_version") + require.Equal(t, append([]int64{}, base.historyIDs...), afterHistory[:len(base.historyIDs)], + "SendMessage(queue) into W/E0 leaves the existing history prefix intact") + require.Equal(t, []int64{inserted.ID}, newActiveMessageIDs(base, afterHistory), + "SendMessage(queue) into W/E0 appends exactly the new user message") + + case from == chatstate.StateE1: + // E1: the previous head is promoted into history + // and replaced by the new tail. Net queue size + // unchanged. + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(queue) from E1 returns the new queued tail") + require.Len(t, result.sendMessage.InsertedMessages, 1, + "SendMessage(queue) from E1 promotes the previous head into history") + promoted := assertFetchedUserMessage(ctx, t, f, result.sendMessage.InsertedMessages[0]) + require.Equal(t, seeded.chatID, promoted.ChatID) + require.NotEmpty(t, seeded.queuedMessageBodies, + chatstate.StateE1.String()+" seed must record the queue head body") + assertChatMessageText(t, promoted, seeded.queuedMessageBodies[0]) + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-queue") + // Previous head queued message is gone from the + // queue and now lives in history. + require.NotEmpty(t, base.queueIDs, + chatstate.StateE1.String()+" seed must have a queue head") + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, base.queueIDs[0]) + require.Equal(t, []int64{newQueued.ID}, afterQueueIDs, + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + ": queue must end with only the new tail") + require.False(t, after.LastError.Valid, + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " clears last_error") + require.Equal(t, database.ChatStatusRunning, after.Status) + require.Equal(t, []int64{promoted.ID}, newActiveMessageIDs(base, afterHistory), + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " inserts only the promoted user message") + require.Greater(t, after.QueueVersion, base.queueVersion, + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " advances queue_version") + + default: + // Busy states: the new user message is appended at + // the queue tail; history is untouched. + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(queue) from busy states returns the queued message") + require.Empty(t, result.sendMessage.InsertedMessages, + "SendMessage(queue) from busy states does not insert history") + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-queue") + wantQueue := append(append([]int64{}, base.queueIDs...), newQueued.ID) + require.Equal(t, wantQueue, afterQueueIDs, + "SendMessage(queue) from busy states appends to the queue tail") + require.Equal(t, base.historyIDs, afterHistory, + "SendMessage(queue) from busy states does not change history") + require.Greater(t, after.QueueVersion, base.queueVersion, + "SendMessage(queue) from busy states advances queue_version") + switch from { + case chatstate.StateA0, chatstate.StateA1: + require.True(t, after.RequiresActionDeadlineAt.Valid, + "SendMessage(queue) from A* preserves requires_action_deadline_at") + require.Equal(t, base.chat.RequiresActionDeadlineAt, after.RequiresActionDeadlineAt, + "SendMessage(queue) from A* preserves the deadline value") + require.Equal(t, database.ChatStatusRequiresAction, after.Status) + case chatstate.StateI0, chatstate.StateI1: + require.Equal(t, database.ChatStatusInterrupting, after.Status) + case chatstate.StateR0, chatstate.StateR1: + require.Equal(t, database.ChatStatusRunning, after.Status) + } + } + }, + } +} + +func sendMessageInterruptCase(from, want chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionSendMessage, + from: from, + want: want, + scenario: scenarioInterrupt, + apply: applySendMessageInterrupt, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterQueue, err := f.DB.CountChatQueuedMessages(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + + switch from { + case chatstate.StateW, chatstate.StateE0: + // W/E0 with interrupt-mode behaves like the direct + // insert from queue-mode: the new user message lands + // directly in history, the queue is left untouched, + // last_error is cleared, and the chat lands in R0. + require.Equal(t, base.queueCount, afterQueue, + "SendMessage(interrupt) into W/E0 must not queue") + require.Nil(t, result.sendMessage.QueuedMessage, + "SendMessage(interrupt) into W/E0 does not return a queued message") + require.Len(t, result.sendMessage.InsertedMessages, 1, + "SendMessage(interrupt) into W/E0 inserts exactly one history message") + inserted := assertFetchedUserMessage(ctx, t, f, result.sendMessage.InsertedMessages[0]) + require.Equal(t, seeded.chatID, inserted.ChatID) + assertChatMessageText(t, inserted, "sm-interrupt") + require.False(t, after.LastError.Valid, + "SendMessage(interrupt) into W/E0 clears last_error") + require.Equal(t, database.ChatStatusRunning, after.Status, + "SendMessage(interrupt) into W/E0 lands in running") + require.Equal(t, base.queueIDs, afterQueueIDs, + "SendMessage(interrupt) into W/E0 must not touch queued messages") + require.Equal(t, []int64{inserted.ID}, newActiveMessageIDs(base, afterHistory), + "SendMessage(interrupt) into W/E0 appends exactly the new user message") + + case chatstate.StateE1: + // E1 with interrupt-mode mirrors queue-mode: the + // previous head is promoted into history and the new + // tail replaces it in the queue. Net queue size + // unchanged, last_error cleared. + require.Equal(t, base.queueCount, afterQueue, + "SendMessage(interrupt) from E1 leaves queue size unchanged") + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(interrupt) from E1 returns the new queued tail") + require.Len(t, result.sendMessage.InsertedMessages, 1, + "SendMessage(interrupt) from E1 promotes the previous head into history") + promoted := assertFetchedUserMessage(ctx, t, f, result.sendMessage.InsertedMessages[0]) + require.Equal(t, seeded.chatID, promoted.ChatID) + require.NotEmpty(t, seeded.queuedMessageBodies, + chatstate.StateE1.String()+" seed must record queue head body") + assertChatMessageText(t, promoted, seeded.queuedMessageBodies[0]) + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-interrupt") + require.NotEmpty(t, base.queueIDs, + chatstate.StateE1.String()+" seed must have a queue head") + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, base.queueIDs[0]) + require.Equal(t, []int64{newQueued.ID}, afterQueueIDs, + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " interrupt: queue must end with only the new tail") + require.False(t, after.LastError.Valid) + require.Equal(t, database.ChatStatusRunning, after.Status) + require.Equal(t, []int64{promoted.ID}, newActiveMessageIDs(base, afterHistory), + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " interrupt inserts only the promoted user message") + require.Greater(t, after.QueueVersion, base.queueVersion, + chatstate.StateE1.String()+" -> "+chatstate.StateR1.String()+ + " interrupt advances queue_version") + + case chatstate.StateI0, chatstate.StateI1: + // I*: append to queue tail, history untouched, status + // stays interrupting. + require.Equal(t, base.queueCount+1, afterQueue, + "SendMessage(interrupt) from I* appends one queued message") + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(interrupt) from I* returns the queued tail") + require.Empty(t, result.sendMessage.InsertedMessages, + "SendMessage(interrupt) from I* does not insert history") + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-interrupt") + wantQueue := append(append([]int64{}, base.queueIDs...), newQueued.ID) + require.Equal(t, wantQueue, afterQueueIDs, + "SendMessage(interrupt) from I* appends to the queue tail") + require.Equal(t, base.historyIDs, afterHistory, + "SendMessage(interrupt) from I* must not touch history") + require.Equal(t, database.ChatStatusInterrupting, after.Status, + "SendMessage(interrupt) from I* keeps status interrupting") + require.Greater(t, after.QueueVersion, base.queueVersion, + "SendMessage(interrupt) from I* advances queue_version") + + case chatstate.StateR0, chatstate.StateR1: + require.Equal(t, base.queueCount+1, afterQueue, + "SendMessage(interrupt) from R* appends one queued message") + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(interrupt) from R* returns the queued tail") + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-interrupt") + wantQueue := append(append([]int64{}, base.queueIDs...), newQueued.ID) + require.Equal(t, wantQueue, afterQueueIDs, + "SendMessage(interrupt) from R* appends to the queue tail") + require.Greater(t, after.QueueVersion, base.queueVersion, + "SendMessage(interrupt) from R* advances queue_version") + require.Equal(t, database.ChatStatusInterrupting, after.Status, + "R* -> I1 sets status interrupting") + require.Equal(t, base.historyIDs, afterHistory, + "SendMessage(interrupt) from R* must not touch history") + + case chatstate.StateA0, chatstate.StateA1: + require.Equal(t, base.queueCount+1, afterQueue, + "SendMessage(interrupt) from A* appends one queued message") + require.NotNil(t, result.sendMessage.QueuedMessage, + "SendMessage(interrupt) from A* returns the queued tail") + newQueued := assertFetchedQueuedMessage(ctx, t, f, seeded.chatID, *result.sendMessage.QueuedMessage) + assertQueuedMessageText(t, newQueued, "sm-interrupt") + wantQueue := append(append([]int64{}, base.queueIDs...), newQueued.ID) + require.Equal(t, wantQueue, afterQueueIDs, + "SendMessage(interrupt) from A* appends to the queue tail") + require.Greater(t, after.QueueVersion, base.queueVersion, + "SendMessage(interrupt) from A* advances queue_version") + require.Equal(t, database.ChatStatusRunning, after.Status, + "A* -> R1 cancels pending dynamic calls and resumes running") + require.False(t, after.RequiresActionDeadlineAt.Valid, + "A* -> R1 clears requires_action_deadline_at") + // Cancellation messages for the pending dynamic + // tool call should land in active history. They are + // not returned via SendMessageResult, so we fetch + // them by diffing the active history set. + newIDs := newActiveMessageIDs(base, afterHistory) + require.Len(t, newIDs, 1, + "SendMessage(interrupt) from A* synthesizes exactly one tool-result cancellation") + cancel := requireChatMessageByID(ctx, t, f, newIDs[0]) + assertToolResultForCall(t, cancel, seeded.pendingToolCallID) + } + }, + } +} + +func editMessageCase(from chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionEditMessage, + from: from, + want: chatstate.StateR0, + apply: applyEditMessage, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, after.Status, + "EditMessage always lands in running") + require.False(t, after.Archived, "EditMessage clears archived") + require.False(t, after.LastError.Valid, + "EditMessage clears last_error") + count, err := f.DB.CountChatQueuedMessages(ctx, seeded.chatID) + require.NoError(t, err) + require.Zero(t, count, "EditMessage clears the queue") + require.Empty(t, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "EditMessage leaves no queued messages") + + // Replacement message must be a fresh user message that + // replaces the original target and lives in active history. + require.NotZero(t, result.editMessage.ReplacementMessage.ID, + "EditMessage returns the replacement message") + replacement := assertFetchedUserMessage(ctx, t, f, result.editMessage.ReplacementMessage) + require.Equal(t, seeded.chatID, replacement.ChatID) + require.NotEqual(t, seeded.initialUserMessageID, replacement.ID, + "EditMessage inserts a new replacement message") + assertChatMessageText(t, replacement, "edited") + + // Every history message from the edited message onward, + // inclusive, must be soft-deleted. base.historyIDs is the + // active history in order before the transition, so the + // expected deleted suffix is everything from the target's + // position to the end of that slice. GetChatMessageByID + // filters deleted=false, so it must return an error for + // each deleted ID. + require.NotEmpty(t, result.editMessage.DeletedMessageIDs, + "EditMessage deletes at least the target user message") + targetIdx := slices.Index(base.historyIDs, seeded.initialUserMessageID) + require.GreaterOrEqual(t, targetIdx, 0, + "baseline active history must contain the edited message") + wantDeleted := append([]int64{}, base.historyIDs[targetIdx:]...) + require.Equal(t, wantDeleted, result.editMessage.DeletedMessageIDs, + "EditMessage soft-deletes the edited message and every later active history message in order") + for _, id := range result.editMessage.DeletedMessageIDs { + _, err := f.DB.GetChatMessageByID(ctx, id) + require.Error(t, err, + "EditMessage: deleted message %d must not be active", id) + } + // Every deleted queued message must be gone from the queue. + for _, id := range result.editMessage.DeletedQueuedMessageIDs { + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, id) + } + for _, id := range base.queueIDs { + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, id) + } + }, + } +} + +func deleteQueuedCase(from, want chatstate.ExecutionState, shape queueShape) transitionCaseSpec { + spec := transitionCaseSpec{ + transition: chatstate.TransitionDeleteQueuedMessage, + from: from, + want: want, + apply: applyDeleteQueuedMessage, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + afterQueue, err := f.DB.CountChatQueuedMessages(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, base.queueCount-1, afterQueue, + "DeleteQueuedMessage removes exactly one queued message") + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Greater(t, after.QueueVersion, base.queueVersion, + "DeleteQueuedMessage advances queue_version") + + // The target queued message is the seeded head. It must + // be returned in DeletedQueuedMessage, and it must no + // longer be fetchable. + require.NotEmpty(t, seeded.queuedMessageIDs) + targetID := seeded.queuedMessageIDs[0] + require.Equal(t, targetID, result.deleteQueuedMessage.DeletedQueuedMessage.ID, + "DeletedQueuedMessage returns the targeted queued message") + require.Equal(t, seeded.chatID, result.deleteQueuedMessage.DeletedQueuedMessage.ChatID) + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, targetID) + + // Remaining queue IDs are the baseline tail. + wantRemaining := append([]int64{}, base.queueIDs[1:]...) + require.Equal(t, wantRemaining, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "DeleteQueuedMessage preserves remaining queue order") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "DeleteQueuedMessage does not touch history") + }, + } + if shape.isMulti() { + spec.scenario = scenarioMulti + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + return seedStateMultiQueued(t, f, from) + } + } + return spec +} + +func promoteQueuedCase(from, want chatstate.ExecutionState, shape queueShape, targetIdx int) transitionCaseSpec { + var sc scenario + if shape.isMulti() { + switch targetIdx { + case 0: + sc = scenarioHeadTarget + default: + sc = scenarioNonHead + } + } + apply := func(t *testing.T, _ *testFixture, tx *chatstate.Tx, seeded seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + require.Less(t, targetIdx, len(seeded.queuedMessageIDs), "promote target index out of range") + var err error + result.promoteQueuedMessage, err = tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: seeded.queuedMessageIDs[targetIdx], + }) + return err + } + spec := transitionCaseSpec{ + transition: chatstate.TransitionPromoteQueuedMessage, + from: from, + want: want, + scenario: sc, + apply: apply, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + afterQueue := int64(len(afterQueueIDs)) + + require.NotEmpty(t, seeded.queuedMessageIDs) + require.Less(t, targetIdx, len(seeded.queuedMessageIDs)) + targetID := seeded.queuedMessageIDs[targetIdx] + require.Equal(t, targetID, result.promoteQueuedMessage.QueuedMessage.ID, + "PromoteQueuedMessage returns the targeted queued message") + + switch from { + case chatstate.StateE1, chatstate.StateA1: + // Head is popped into history. + require.Equal(t, base.queueCount-1, afterQueue, + "E1/A1 promote pops the head into history") + require.Equal(t, database.ChatStatusRunning, after.Status, + "E1/A1 promote lands in running") + require.False(t, after.LastError.Valid, + "E1/A1 promote clears last_error") + require.False(t, after.RequiresActionDeadlineAt.Valid, + "E1/A1 promote clears requires_action_deadline_at") + require.NotNil(t, result.promoteQueuedMessage.InsertedMessage, + "E1/A1 promote inserts a user history message") + inserted := requireChatMessageByID(ctx, t, f, result.promoteQueuedMessage.InsertedMessage.ID) + require.Equal(t, seeded.chatID, inserted.ChatID) + require.Equal(t, database.ChatMessageRoleUser, inserted.Role) + require.True(t, inserted.ModelConfigID.Valid) + require.Equal(t, f.Model.ID, inserted.ModelConfigID.UUID) + require.Equal(t, chatprompt.CurrentContentVersion, inserted.ContentVersion) + require.True(t, inserted.CreatedBy.Valid) + require.Equal(t, result.promoteQueuedMessage.QueuedMessage.CreatedBy, inserted.CreatedBy.UUID, + "promoted history message preserves queued created_by") + if len(seeded.queuedMessageCreatedBy) > targetIdx { + require.Equal(t, seeded.queuedMessageCreatedBy[targetIdx], inserted.CreatedBy.UUID, + "promoted history message preserves non-owner queued creator") + } + require.NotEmpty(t, seeded.queuedMessageBodies, + "E1/A1 seed must record queued message bodies") + assertChatMessageText(t, inserted, seeded.queuedMessageBodies[targetIdx]) + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, targetID) + wantRemaining := remainingExcluding(base.queueIDs, targetIdx) + require.Equal(t, wantRemaining, afterQueueIDs, + "E1/A1 promote leaves the remaining queue order intact") + assertQueueBodiesInOrder(ctx, t, f, seeded.chatID, + remainingBodiesExcluding(seeded.queuedMessageBodies, targetIdx)) + // New active history adds exactly the inserted + // user message plus any synthetic cancellations. + newIDs := newActiveMessageIDs(base, afterHistory) + require.Contains(t, newIDs, inserted.ID, + "newly-active history contains the promoted user message") + if from == chatstate.StateA1 { + // A1: every outstanding tool call must be + // canceled before the promoted user message. + require.Len(t, result.promoteQueuedMessage.CancellationMessages, len(seeded.pendingToolCallIDs), + "A1 promote synthesizes one tool-result cancellation per outstanding call") + gotIDs := make(map[string]bool) + for _, cancelMsg := range result.promoteQueuedMessage.CancellationMessages { + cancel := requireChatMessageByID(ctx, t, f, cancelMsg.ID) + require.Less(t, cancel.ID, inserted.ID, + "A1 promote inserts cancellations before the promoted user message") + parts, err := chatprompt.ParseContent(cancel) + require.NoError(t, err) + for _, part := range parts { + if part.Type != codersdk.ChatMessagePartTypeToolResult { + continue + } + require.True(t, part.IsError, + "A1 promote synthetic cancellation is marked as an error") + gotIDs[part.ToolCallID] = true + } + } + for _, callID := range seeded.pendingToolCallIDs { + require.True(t, gotIDs[callID], + "A1 promote cancels outstanding tool call %s", callID) + } + } else { + require.Empty(t, result.promoteQueuedMessage.CancellationMessages, + "E1 promote has no synthetic cancellations") + } + case chatstate.StateR1, chatstate.StateI1: + // Reorder-only: status flips to interrupting, no + // history insert, queue cardinality unchanged. + require.Equal(t, base.queueCount, afterQueue, + "R1/I1 promote leaves queue cardinality unchanged") + require.Equal(t, database.ChatStatusInterrupting, after.Status, + "R1/I1 promote lands in interrupting") + require.Nil(t, result.promoteQueuedMessage.InsertedMessage, + "R1/I1 promote must not insert a history message") + require.Empty(t, result.promoteQueuedMessage.CancellationMessages, + "R1/I1 promote has no synthetic cancellations") + require.Equal(t, base.historyIDs, afterHistory, + "R1/I1 promote leaves history unchanged") + // Target must still be present and now at the head. + queued := requireQueuedMessageByID(ctx, t, f, seeded.chatID, targetID) + require.Equal(t, targetID, queued.ID) + require.NotEmpty(t, afterQueueIDs) + require.Equal(t, targetID, afterQueueIDs[0], + "R1/I1 promote brings the target to the queue head") + require.NotEmpty(t, seeded.queuedMessageBodies, + "R1/I1 seed must record queued message bodies") + if targetIdx == 0 { + // Head-target: zero rows updated, so the + // queue order is unchanged and queue_version + // stays put. + require.Equal(t, base.queueIDs, afterQueueIDs, + "head-target promote preserves queue order") + require.Equal(t, base.queueVersion, after.QueueVersion, + "head-target promote leaves queue_version unchanged") + assertQueueBodiesInOrder(ctx, t, f, seeded.chatID, seeded.queuedMessageBodies) + } else { + // Non-head: target moves to the head, the rest + // of the original order is preserved. + wantQueue := append([]int64{targetID}, remainingExcluding(base.queueIDs, targetIdx)...) + require.Equal(t, wantQueue, afterQueueIDs, + "non-head promote moves the target to the head and preserves the rest") + require.Greater(t, after.QueueVersion, base.queueVersion, + "non-head promote advances queue_version") + wantBodies := append([]string{seeded.queuedMessageBodies[targetIdx]}, + remainingBodiesExcluding(seeded.queuedMessageBodies, targetIdx)...) + assertQueueBodiesInOrder(ctx, t, f, seeded.chatID, wantBodies) + } + } + }, + } + if from == chatstate.StateA1 { + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + queuedExtras := 1 + if shape.isMulti() { + queuedExtras = 2 + } + return seedA1WithMixedOutstandingToolCalls(t, f, queuedExtras, "seed_tool_a1_promote") + } + } else if shape.isMulti() { + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + return seedStateMultiQueued(t, f, from) + } + } + return spec +} + +func interruptCase(from, want chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionInterrupt, + from: from, + want: want, + apply: applyInterrupt, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + require.Equal(t, base.queueIDs, afterQueueIDs, + "Interrupt does not touch queued messages") + + switch from { + case chatstate.StateR0, chatstate.StateR1: + require.Equal(t, database.ChatStatusInterrupting, after.Status, + "Interrupt from R* sets status interrupting") + require.Equal(t, base.historyIDs, afterHistory, + "Interrupt from R* leaves history unchanged") + require.Empty(t, result.interrupt.CancellationMessages, + "Interrupt from R* does not synthesize tool cancellations") + case chatstate.StateA0, chatstate.StateA1: + require.Equal(t, database.ChatStatusRunning, after.Status, + "Interrupt from A* cancels pending dynamic calls and resumes running") + require.False(t, after.RequiresActionDeadlineAt.Valid, + "Interrupt from A* clears requires_action_deadline_at") + require.Len(t, result.interrupt.CancellationMessages, 1, + "Interrupt from A* synthesizes one tool-result cancellation") + cancel := requireChatMessageByID(ctx, t, f, + result.interrupt.CancellationMessages[0].ID) + assertToolResultForCall(t, cancel, seeded.pendingToolCallID) + } + }, + } +} + +func completeRequiresActionCase(from, want chatstate.ExecutionState) transitionCaseSpec { + // Re-seed A0/A1 fresh per case so the pending tool call ID is + // available on the seeded chat. + return transitionCaseSpec{ + transition: chatstate.TransitionCompleteRequiresAction, + from: from, + want: want, + apply: applyCompleteRequiresAction, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, after.Status, + "CompleteRequiresAction sets status running") + require.False(t, after.RequiresActionDeadlineAt.Valid, + "CompleteRequiresAction clears requires_action_deadline_at") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "CompleteRequiresAction preserves queued messages") + + // The user-submitted tool result must be inserted as a + // tool-role message that references the seeded + // pendingToolCallID with is_error=false. + require.Len(t, result.completeRequiresAction.InsertedMessages, 1, + "CompleteRequiresAction inserts one tool-result message per pending call") + inserted := requireChatMessageByID(ctx, t, f, + result.completeRequiresAction.InsertedMessages[0].ID) + assertToolResultForCallNoError(t, inserted, seeded.pendingToolCallID, `{"ok":true}`) + }, + } +} + +func cancelRequiresActionCase(from, want chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionCancelRequiresAction, + from: from, + want: want, + apply: applyCancelRequiresAction, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, after.Status, + "CancelRequiresAction sets status running") + require.False(t, after.RequiresActionDeadlineAt.Valid, + "CancelRequiresAction clears requires_action_deadline_at") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "CancelRequiresAction preserves queued messages") + + // One synthetic tool-result cancellation per pending call. + require.Len(t, result.cancelRequiresAction.CancellationMessages, 1, + "CancelRequiresAction synthesizes one tool-result per pending call") + cancel := requireChatMessageByID(ctx, t, f, + result.cancelRequiresAction.CancellationMessages[0].ID) + assertToolResultForCall(t, cancel, seeded.pendingToolCallID) + }, + } +} + +func recordGenerationAttemptCase(from chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionRecordGenerationAttempt, + from: from, + want: from, // state preserved + apply: applyRecordGenerationAttempt, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, int64(1), after.GenerationAttempt, + "RecordGenerationAttempt increments generation_attempt by one") + require.Equal(t, result.recordGenerationAttempt.GenerationAttempt, after.GenerationAttempt, + "RecordGenerationAttempt result mirrors the persisted value") + require.Equal(t, base.historyVersion, after.HistoryVersion, + "RecordGenerationAttempt does not change history_version") + require.Equal(t, base.queueVersion, after.QueueVersion, + "RecordGenerationAttempt does not change queue_version") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "RecordGenerationAttempt does not change queue order") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "RecordGenerationAttempt does not change history messages") + }, + } +} + +func recordRetryStateCase(from chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionRecordRetryState, + from: from, + want: from, // state preserved + apply: applyRecordRetryState, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.True(t, after.RetryState.Valid, + "RecordRetryState stores retry_state") + require.JSONEq(t, + string(result.recordRetryState.Chat.RetryState.RawMessage), + string(after.RetryState.RawMessage), + "RecordRetryState result mirrors persisted retry_state") + require.JSONEq(t, + `{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`, + string(after.RetryState.RawMessage), + "RecordRetryState stores the expected payload") + require.Equal(t, after.SnapshotVersion, after.RetryStateVersion, + "RecordRetryState sets retry_state_version to snapshot_version") + require.Greater(t, after.RetryStateVersion, base.retryStateVersion, + "RecordRetryState advances retry_state_version") + require.Equal(t, base.historyVersion, after.HistoryVersion, + "RecordRetryState does not change history_version") + require.Equal(t, base.queueVersion, after.QueueVersion, + "RecordRetryState does not change queue_version") + require.Equal(t, base.generationAttempt, after.GenerationAttempt, + "RecordRetryState does not change generation_attempt") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "RecordRetryState does not change queue order") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "RecordRetryState does not change history messages") + }, + } +} + +func commitStepCase(from chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionCommitStep, + from: from, + want: from, // state preserved + apply: applyCommitStep, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + require.Equal(t, len(base.historyIDs)+1, len(afterHistory), + "CommitStep appends exactly one history message") + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Greater(t, after.HistoryVersion, base.historyVersion, + "CommitStep advances history_version") + + require.Len(t, result.commitStep.InsertedMessages, 1, + "CommitStep returns the inserted assistant message") + inserted := requireChatMessageByID(ctx, t, f, + result.commitStep.InsertedMessages[0].ID) + require.Equal(t, seeded.chatID, inserted.ChatID) + require.Equal(t, database.ChatMessageRoleAssistant, inserted.Role, + "CommitStep inserts an assistant-role message") + assertChatMessageText(t, inserted, "assistant") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "CommitStep does not change queue order") + require.Equal(t, base.queueVersion, after.QueueVersion, + "CommitStep does not change queue_version") + }, + } +} + +func enterRequiresActionCase(from, want chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionEnterRequiresAction, + from: from, + want: want, + seed: seedForEnterRequiresAction, + apply: applyEnterRequiresAction, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRequiresAction, after.Status, + "EnterRequiresAction sets status requires_action") + require.True(t, after.RequiresActionDeadlineAt.Valid, + "EnterRequiresAction populates requires_action_deadline_at") + require.True(t, result.enterRequiresAction.RequiresActionDeadlineAt.Valid, + "EnterRequiresAction returns the deadline") + require.Equal(t, result.enterRequiresAction.RequiresActionDeadlineAt, after.RequiresActionDeadlineAt, + "EnterRequiresAction returned deadline matches the persisted value") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "EnterRequiresAction preserves queued messages") + require.Equal(t, base.queueVersion, after.QueueVersion, + "EnterRequiresAction does not bump queue_version") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "EnterRequiresAction does not insert history") + }, + } +} + +func finishInterruptionRejectsOutstandingToolCallCase() transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionFinishInterruption, + from: chatstate.StateI0, + want: chatstate.StateI0, + scenario: scenarioRejectNonDynamicOutstandingToolCall, + seed: func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + nonDynCallID := "call_" + uuid.NewString() + commitAssistantToolCall(t, f, m, + nonDynamicAssistantToolCallMessage(t, f.Model.ID, nonDynCallID)) + + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Interrupt(chatstate.InterruptInput{Reason: "test"}) + return err + })) + return seededChat{ + chatID: created.Chat.ID, + exists: true, + initialUserMessageID: firstUserMessageID(ctx, t, f, created.Chat.ID), + assistantToolCallMsgID: firstAssistantMessageID(ctx, t, f, created.Chat.ID), + pendingToolCallIDs: []string{nonDynCallID}, + } + }, + apply: applyFinishInterruption, + assertFailure: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, err error) { + require.Error(t, err, "FinishInterruption must reject an outstanding tool call") + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed, + "rejection must wrap ErrTransitionNotAllowed") + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te, + "FinishInterruption must return a typed TransitionError") + require.Equal(t, chatstate.TransitionFinishInterruption, te.Transition) + require.Equal(t, chatstate.StateI0, te.From) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) + }, + } +} + +func finishInterruptionCase(from, want chatstate.ExecutionState, shape queueShape) transitionCaseSpec { + spec := transitionCaseSpec{ + transition: chatstate.TransitionFinishInterruption, + from: from, + want: want, + apply: applyFinishInterruption, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + switch from { + case chatstate.StateI0: + require.Equal(t, database.ChatStatusWaiting, after.Status, + "FinishInterruption from I0 lands in waiting") + require.Nil(t, result.finishInterruption.PromotedMessage, + "FinishInterruption from I0 promotes nothing") + require.Equal(t, base.queueIDs, afterQueueIDs, + "FinishInterruption from I0 leaves queued messages unchanged") + require.Equal(t, base.historyIDs, afterHistory, + "FinishInterruption from I0 with no partial messages leaves history unchanged") + case chatstate.StateI1: + require.Equal(t, database.ChatStatusRunning, after.Status, + "FinishInterruption from I1 lands in running") + require.NotNil(t, result.finishInterruption.PromotedMessage, + "FinishInterruption from I1 promotes the head into history") + promoted := assertFetchedUserMessage(ctx, t, f, + *result.finishInterruption.PromotedMessage) + require.Equal(t, seeded.chatID, promoted.ChatID) + require.Contains(t, newActiveMessageIDs(base, afterHistory), promoted.ID, + "FinishInterruption from I1 inserts the promoted user message") + require.NotEmpty(t, seeded.queuedMessageBodies, + "I1 seed must record queued message bodies") + assertChatMessageText(t, promoted, seeded.queuedMessageBodies[0]) + require.NotEmpty(t, base.queueIDs) + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, base.queueIDs[0]) + wantRemaining := append([]int64{}, base.queueIDs[1:]...) + require.Equal(t, wantRemaining, afterQueueIDs, + "FinishInterruption from I1 preserves the queue tail order") + assertQueueBodiesInOrder(ctx, t, f, seeded.chatID, + seeded.queuedMessageBodies[1:]) + } + }, + } + if shape.isMulti() { + spec.scenario = scenarioMulti + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + return seedStateMultiQueued(t, f, from) + } + } + return spec +} + +func finishTurnCase(from, want chatstate.ExecutionState, shape queueShape) transitionCaseSpec { + spec := transitionCaseSpec{ + transition: chatstate.TransitionFinishTurn, + from: from, + want: want, + apply: applyFinishTurn, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + afterHistory := activeHistoryIDs(ctx, t, f, seeded.chatID) + afterQueueIDs := queuedIDsByPosition(ctx, t, f, seeded.chatID) + switch from { + case chatstate.StateR0: + require.Equal(t, database.ChatStatusWaiting, after.Status, + "FinishTurn from R0 lands in waiting") + require.Nil(t, result.finishTurn.PromotedMessage, + "FinishTurn from R0 promotes nothing") + require.Equal(t, base.queueIDs, afterQueueIDs, + "FinishTurn from R0 leaves queued messages unchanged") + require.Equal(t, base.historyIDs, afterHistory, + "FinishTurn from R0 leaves history unchanged") + case chatstate.StateR1: + require.Equal(t, database.ChatStatusRunning, after.Status, + "FinishTurn from R1 lands in running") + require.NotNil(t, result.finishTurn.PromotedMessage, + "FinishTurn from R1 promotes the head into history") + promoted := assertFetchedUserMessage(ctx, t, f, + *result.finishTurn.PromotedMessage) + require.Equal(t, seeded.chatID, promoted.ChatID) + require.Contains(t, newActiveMessageIDs(base, afterHistory), promoted.ID, + "FinishTurn from R1 inserts the promoted user message") + require.NotEmpty(t, seeded.queuedMessageBodies, + "R1 seed must record queued message bodies") + assertChatMessageText(t, promoted, seeded.queuedMessageBodies[0]) + require.NotEmpty(t, base.queueIDs) + requireQueuedMessageDeleted(ctx, t, f, seeded.chatID, base.queueIDs[0]) + wantRemaining := append([]int64{}, base.queueIDs[1:]...) + require.Equal(t, wantRemaining, afterQueueIDs, + "FinishTurn from R1 preserves the queue tail order") + assertQueueBodiesInOrder(ctx, t, f, seeded.chatID, + seeded.queuedMessageBodies[1:]) + } + }, + } + if shape.isMulti() { + spec.scenario = scenarioMulti + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + return seedStateMultiQueued(t, f, from) + } + } + return spec +} + +func finishErrorCase(from, want chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionFinishError, + from: from, + want: want, + apply: applyFinishError, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + _ = result + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusError, after.Status, + "FinishError sets status error") + require.True(t, after.LastError.Valid, + "FinishError stores last_error") + require.JSONEq(t, `{"message":"finish-error"}`, string(after.LastError.RawMessage), + "FinishError persists the input last_error JSON") + require.Equal(t, base.historyVersion, after.HistoryVersion, + "FinishError does not change history_version") + require.Equal(t, base.queueVersion, after.QueueVersion, + "FinishError does not change queue_version") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "FinishError preserves queued messages") + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID), + "FinishError preserves history messages") + }, + } +} + +func reconcileInvalidStateCase(want chatstate.ExecutionState, shape queueShape) transitionCaseSpec { + spec := transitionCaseSpec{ + transition: chatstate.TransitionReconcileInvalidState, + from: chatstate.StateInvalid, + want: want, + apply: applyReconcileInvalidState, + assert: func(ctx context.Context, t *testing.T, f *testFixture, seeded seededChat, base snapshotBaseline, result transitionCaseResult) { + after, err := f.DB.GetChatByID(ctx, seeded.chatID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusError, after.Status, + "ReconcileInvalidState lands in error") + require.False(t, after.Archived, + "ReconcileInvalidState clears archived") + require.True(t, after.LastError.Valid, + "ReconcileInvalidState sets a default last_error") + require.Equal(t, base.queueIDs, queuedIDsByPosition(ctx, t, f, seeded.chatID), + "ReconcileInvalidState preserves queued messages") + // For the current invalid seeds there are no pending + // dynamic tool calls, so no cancellation messages are + // expected. Still, if any are returned we fetch them + // to verify they were persisted as tool-role messages. + for _, c := range result.reconcileInvalidState.CancellationMessages { + msg := requireChatMessageByID(ctx, t, f, c.ID) + require.Equal(t, database.ChatMessageRoleTool, msg.Role) + } + }, + } + if shape.isMulti() { + spec.scenario = scenarioWithQueue + spec.seed = func(t *testing.T, f *testFixture, _ chatstate.ExecutionState) seededChat { + return seedInvalidWithQueue(t, f) + } + } + return spec +} diff --git a/coderd/x/chatd/chatstate/transitions_test.go b/coderd/x/chatd/chatstate/transitions_test.go new file mode 100644 index 0000000000..c4df476d7b --- /dev/null +++ b/coderd/x/chatd/chatstate/transitions_test.go @@ -0,0 +1,743 @@ +package chatstate_test + +import ( + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// CreateChat tests. +// +// CreateChat is the only transition that originates from StateN and it +// is not exercised through ChatMachine.Update, so it lives outside +// TestTransitionMatrix_AllCombinations. + +// TestTransitionCreate_NToR0 verifies that CreateChat lands a fresh +// chat in R0 with snapshot_version 1, the initial user message +// recorded at revision 1, queue_version still 0, and the post-commit +// publish requesting an ownership hint plus a chat:update. +func TestTransitionCreate_NToR0(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + res := createTestChat(t, f) + + require.Equal(t, database.ChatStatusRunning, res.Chat.Status) + require.False(t, res.Chat.Archived) + require.Equal(t, int64(1), res.Chat.SnapshotVersion, "snapshot_version starts at 1") + require.Equal(t, int64(1), res.Chat.HistoryVersion, "history_version set by trigger after initial insert") + require.Equal(t, int64(0), res.Chat.QueueVersion, "queue_version stays 0 when no queue rows") + require.Equal(t, int64(0), res.Chat.GenerationAttempt) + require.NotEmpty(t, res.InitialMessages) + require.Equal(t, int64(1), res.InitialMessages[0].Revision) + require.Equal(t, chatstate.StateR0, f.classify(ctx, t, res.Chat.ID)) + require.True(t, f.Pub.hasOwnership(), "newly created chat is runnable and unowned") + f.Pub.expectChatUpdate(t, res.Chat.ID, 1) +} + +// TestCreateChat_RejectsEmptyInitialMessages verifies that CreateChat +// rejects an empty InitialMessages slice with ErrTransitionNotAllowed +// and does not publish anything. +func TestCreateChat_RejectsEmptyInitialMessages(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + _, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + ClientType: database.ChatClientTypeApi, + Title: "t", + InitialMessages: nil, + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + require.Empty(t, f.Pub.channels, "rejected create must not publish") +} + +func TestCreateChat_AllowsNoUserMessages(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + assistant := userTextMessage("oops", f.User.ID, f.Model.ID) + assistant.Role = database.ChatMessageRoleAssistant + res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "t", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{assistant}, + }) + require.NoError(t, err) + require.Len(t, res.InitialMessages, 1) +} + +func TestCreateChat_AllowsNonFinalUserMessage(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{ + OrganizationID: f.Org.ID, + OwnerID: f.User.ID, + LastModelConfigID: f.Model.ID, + Title: "t", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + userTextMessage("context user", f.User.ID, f.Model.ID), + userTextMessage("final user", f.User.ID, f.Model.ID), + }, + }) + require.NoError(t, err) + require.Len(t, res.InitialMessages, 2) +} + +// Input-specific rejection cases. +// +// These tests cover the same matrix rows as TestTransitionMatrix_AllCombinations +// but exercise legal source states with invalid transition inputs. They are +// intentionally outside the matrix entry point so the matrix focus stays on +// positive cases and generated disallowed cases. + +type setArchivedWrongDirectionCase struct { + from chatstate.ExecutionState + wantArchive bool + label string +} + +func setArchivedWrongDirectionCases() []setArchivedWrongDirectionCase { + return []setArchivedWrongDirectionCase{ + // Non-archived states with archived=false: no-op. + {from: chatstate.StateW, wantArchive: false, label: "W_to_W"}, + {from: chatstate.StateE0, wantArchive: false, label: "E0_to_E0"}, + {from: chatstate.StateE1, wantArchive: false, label: "E1_to_E1"}, + // Archived states with archived=true: no-op. + {from: chatstate.StateXW, wantArchive: true, label: "XW_to_XW"}, + {from: chatstate.StateXE0, wantArchive: true, label: "XE0_to_XE0"}, + {from: chatstate.StateXE1, wantArchive: true, label: "XE1_to_XE1"}, + } +} + +var invalidBusyBehaviors = []chatstate.BusyBehavior{ + chatstate.BusyBehavior(""), + chatstate.BusyBehavior("not-a-real-mode"), +} + +func runSetArchivedWrongDirectionCase(t *testing.T, tc setArchivedWrongDirectionCase) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + seeded := seedState(t, f, tc.from) + require.Equal(t, tc.from, f.classify(ctx, t, seeded.chatID)) + + base := captureBaseline(ctx, t, f, seeded) + + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, serr := tx.SetArchived(chatstate.SetArchivedInput{Archived: tc.wantArchive}) + return serr + }) + require.Error(t, err, "SetArchived must reject when Archived matches the current value") + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed, + "SetArchived must wrap ErrTransitionNotAllowed") + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te, + "SetArchived must return a typed TransitionError") + require.Equal(t, chatstate.TransitionSetArchived, te.Transition) + require.Equal(t, tc.from, te.From, "TransitionError records the loaded from-state") + + require.Equal(t, tc.from, f.classify(ctx, t, seeded.chatID), + "rejected SetArchived must leave the chat in the same state") + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +func runInvalidBusyBehaviorCase(t *testing.T, from chatstate.ExecutionState, bb chatstate.BusyBehavior) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + seeded := seedState(t, f, from) + require.Equal(t, from, f.classify(ctx, t, seeded.chatID), + "seed must land in %s", from) + base := captureBaseline(ctx, t, f, seeded) + + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, serr := tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("invalid-bb", f.User.ID, f.Model.ID), + BusyBehavior: bb, + }) + return serr + }) + require.Error(t, err, "SendMessage must reject invalid BusyBehavior") + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed, + "SendMessage rejection must wrap ErrTransitionNotAllowed") + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te, + "SendMessage must return a typed TransitionError") + require.Equal(t, chatstate.TransitionSendMessage, te.Transition) + require.Equal(t, from, te.From, + "TransitionError records the source state") + + require.Equal(t, from, f.classify(ctx, t, seeded.chatID), + "rejected SendMessage must leave the chat in the same state") + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +type completeRequiresActionRejectCase struct { + name string + results func(seeded seededChat) []chatstate.ToolResultInput +} + +type recordRetryStateRejectCase struct { + name string + retryState pqtype.NullRawMessage +} + +func completeRequiresActionRejectCases() []completeRequiresActionRejectCase { + valid := func(id string) chatstate.ToolResultInput { + return chatstate.ToolResultInput{ + ToolCallID: id, + Output: json.RawMessage(`{"ok":true}`), + } + } + return []completeRequiresActionRejectCase{ + { + name: "missing_required_tool_result", + results: func(seeded seededChat) []chatstate.ToolResultInput { return nil }, + }, + { + name: "extra_tool_result", + results: func(seeded seededChat) []chatstate.ToolResultInput { + return []chatstate.ToolResultInput{valid(seeded.pendingToolCallID), valid("call_extra")} + }, + }, + { + name: "duplicate_tool_call_id", + results: func(seeded seededChat) []chatstate.ToolResultInput { + return []chatstate.ToolResultInput{valid(seeded.pendingToolCallID), valid(seeded.pendingToolCallID)} + }, + }, + { + name: "mismatched_tool_call_id", + results: func(seeded seededChat) []chatstate.ToolResultInput { + return []chatstate.ToolResultInput{valid("call_mismatch")} + }, + }, + { + name: "invalid_json_output", + results: func(seeded seededChat) []chatstate.ToolResultInput { + return []chatstate.ToolResultInput{{ToolCallID: seeded.pendingToolCallID, Output: json.RawMessage(`{`)}} + }, + }, + } +} + +func recordRetryStateRejectCases() []recordRetryStateRejectCase { + return []recordRetryStateRejectCase{ + { + name: "sql_null_payload", + }, + { + name: "empty_payload", + retryState: pqtype.NullRawMessage{RawMessage: json.RawMessage(``), Valid: true}, + }, + { + name: "invalid_json_payload", + retryState: pqtype.NullRawMessage{RawMessage: json.RawMessage(`{`), Valid: true}, + }, + } +} + +func runCompleteRequiresActionRejectCase(t *testing.T, tc completeRequiresActionRejectCase) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + seeded := seedAOrA1(t, f, 0, "reject_complete_requires_action") + require.Equal(t, chatstate.StateA0, f.classify(ctx, t, seeded.chatID)) + base := captureBaseline(ctx, t, f, seeded) + + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, cerr := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ + CreatedBy: f.User.ID, + ModelConfigID: f.Model.ID, + Results: tc.results(seeded), + }) + return cerr + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te) + require.Equal(t, chatstate.TransitionCompleteRequiresAction, te.Transition) + require.Equal(t, chatstate.StateA0, te.From) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +func runRecordRetryStateRejectCase(t *testing.T, tc recordRetryStateRejectCase) { + t.Helper() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + seeded := seedState(t, f, chatstate.StateR0) + require.Equal(t, chatstate.StateR0, f.classify(ctx, t, seeded.chatID)) + base := captureBaseline(ctx, t, f, seeded) + + m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID) + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, rerr := tx.RecordRetryState(chatstate.RecordRetryStateInput{ + RetryState: tc.retryState, + }) + return rerr + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te) + require.Equal(t, chatstate.TransitionRecordRetryState, te.Transition) + require.Equal(t, chatstate.StateR0, te.From) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +// TestTransitionInputValidation groups every input-specific rejection +// test. The matrix coverage entry point in +// TestTransitionMatrix_AllCombinations intentionally focuses on +// positive cases and generated disallowed cases; rejection cases that +// exercise legal matrix rows with invalid inputs live here so the +// matrix entry point stays focused. +func TestTransitionInputValidation(t *testing.T) { + t.Parallel() + + t.Run("SetArchived_wrong_direction", func(t *testing.T) { + t.Parallel() + for _, tc := range setArchivedWrongDirectionCases() { + t.Run(tc.label, func(t *testing.T) { + t.Parallel() + runSetArchivedWrongDirectionCase(t, tc) + }) + } + }) + + t.Run("SendMessage_invalid_busy_behavior", func(t *testing.T) { + t.Parallel() + for _, from := range chatstate.AllowedInputStates(chatstate.TransitionSendMessage) { + for _, bb := range invalidBusyBehaviors { + label := from.String() + "/" + string(bb) + if bb == "" { + label = from.String() + "/empty" + } + t.Run(label, func(t *testing.T) { + t.Parallel() + runInvalidBusyBehaviorCase(t, from, bb) + }) + } + } + }) + + t.Run("CompleteRequiresAction_invalid_results", func(t *testing.T) { + t.Parallel() + for _, tc := range completeRequiresActionRejectCases() { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + runCompleteRequiresActionRejectCase(t, tc) + }) + } + }) + + t.Run("RecordRetryState_invalid_payload", func(t *testing.T) { + t.Parallel() + for _, tc := range recordRetryStateRejectCases() { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + runRecordRetryStateRejectCase(t, tc) + }) + } + }) +} + +// TestSendMessageQueueCapRejectsQueueAppend seeds a chat with the +// maximum queued messages and asserts that the next SendMessage in +// a queue-appending state returns chatstate.ErrMessageQueueFull and +// rolls back without persisting another queued row. +func TestSendMessageQueueCapRejectsQueueAppend(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // createTestChat lands the chat in R0; SendMessage in R0 with + // BusyBehaviorQueue queues. Fill the queue to MaxQueueSize. + for i := 0; i < chatstate.MaxQueueSize; i++ { + sendQueuedMessage(t, f, m, "filler") + } + count, err := f.DB.CountChatQueuedMessages(ctx, created.Chat.ID) + require.NoError(t, err) + require.EqualValues(t, chatstate.MaxQueueSize, count) + chatBefore := f.readChat(ctx, t, created.Chat.ID) + + // The next queue append must fail with ErrMessageQueueFull and a + // typed wrapper that exposes the cap. + err = m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, serr := tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage("overflow", f.User.ID, f.Model.ID), + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return serr + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrMessageQueueFull, + "queue-append over the cap returns ErrMessageQueueFull") + var typed *chatstate.MessageQueueFullError + require.ErrorAs(t, err, &typed, "ErrMessageQueueFull is carried as a typed error") + require.EqualValues(t, chatstate.MaxQueueSize, typed.Max) + + // The transaction rolled back: queue size, snapshot version, + // and queue version are unchanged. + countAfter, err := f.DB.CountChatQueuedMessages(ctx, created.Chat.ID) + require.NoError(t, err) + require.EqualValues(t, chatstate.MaxQueueSize, countAfter, + "queue size must not change when the cap rejects the append") + chatAfter := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, chatBefore.SnapshotVersion, chatAfter.SnapshotVersion, + "failed queue append must not bump snapshot_version") + require.Equal(t, chatBefore.QueueVersion, chatAfter.QueueVersion, + "failed queue append must not bump queue_version") +} + +// TestEditMessageNonUserReturnsSentinel asserts that editing a +// non-user message returns chatstate.ErrEditedMessageNotUser via +// the TransitionError cause chain, and still matches the generic +// transition sentinel. +func TestEditMessageNonUserReturnsSentinel(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + // Insert an assistant message via CommitStep so we have a + // non-user message to target. + var assistantID int64 + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + assistant := userTextMessage("assistant", f.User.ID, f.Model.ID) + assistant.Role = database.ChatMessageRoleAssistant + step, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistant}, + }) + if err != nil { + return err + } + require.Len(t, step.InsertedMessages, 1) + assistantID = step.InsertedMessages[0].ID + return nil + })) + + rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("new content"), + }) + require.NoError(t, err) + + editErr := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, eerr := tx.EditMessage(chatstate.EditMessageInput{ + MessageID: assistantID, + CreatedBy: f.User.ID, + Content: rawContent, + }) + return eerr + }) + require.Error(t, editErr) + require.ErrorIs(t, editErr, chatstate.ErrEditedMessageNotUser, + "non-user edit returns ErrEditedMessageNotUser via TransitionError cause") + require.ErrorIs(t, editErr, chatstate.ErrTransitionNotAllowed, + "ErrEditedMessageNotUser still matches the generic transition sentinel") +} + +// TestTransitionAbandon_RejectsUnowned verifies that calling Abandon +// on a chat the runner does not own returns ErrTransitionNotAllowed +// wrapped in a TransitionError that records the loaded from-state, +// without mutating chat state or publishing anything. +func TestTransitionAbandon_RejectsUnowned(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + seeded := seededChat{chatID: created.Chat.ID, exists: true} + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + base := captureBaseline(ctx, t, f, seeded) + + err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, aerr := tx.Abandon(chatstate.AbandonInput{}) + return aerr + }) + require.Error(t, err) + require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed) + var te *chatstate.TransitionError + require.ErrorAs(t, err, &te) + require.Equal(t, chatstate.TransitionAbandon, te.Transition) + // createTestChat lands the chat in R0; Abandon's precondition + // rejects an unowned chat there. + require.Equal(t, chatstate.StateR0, te.From) + assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base) +} + +// TestTransitionAbandon_ClearsOwnership verifies the Acquire/Abandon +// round-trip: after Acquire the chat carries a worker+runner and a +// fresh heartbeat row exists, and after Abandon both ownership fields +// are cleared. The heartbeat row is not deleted by Abandon; heartbeat +// cleanup is a separate concern. +func TestTransitionAbandon_ClearsOwnership(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + worker := uuid.New() + runner := uuid.New() + + // Acquire writes ownership and a fresh heartbeat row. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner}) + return err + })) + owned := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, worker, owned.WorkerID.UUID) + require.Equal(t, runner, owned.RunnerID.UUID) + hb, err := f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: runner, + }) + require.NoError(t, err, "Acquire writes a fresh heartbeat row") + require.Equal(t, runner, hb.RunnerID) + + // Abandon clears ownership but leaves the heartbeat row intact. + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Abandon(chatstate.AbandonInput{}) + return err + })) + hb, err = f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: runner, + }) + require.NoError(t, err, "Abandon does not delete the heartbeat row") + abandoned := f.readChat(ctx, t, created.Chat.ID) + require.False(t, abandoned.WorkerID.Valid, "Abandon clears worker_id") + require.False(t, abandoned.RunnerID.Valid, "Abandon clears runner_id") +} + +// TestTransitionAcquire_OverwritesFreshOwnership verifies that Acquire +// is an unconditional ownership handoff: a second worker calling +// Acquire on a chat that was *just* acquired by another worker +// successfully replaces ownership without inspecting heartbeat +// freshness. It also asserts that Acquire itself does not request an +// ownership hint, so the post-commit publish stays quiet on +// `chat:ownership` when the resulting heartbeat is fresh. +func TestTransitionAcquire_OverwritesFreshOwnership(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + + firstWorker := uuid.New() + firstRunner := uuid.New() + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: firstWorker, RunnerID: firstRunner}) + return err + })) + + // The chat is now owned with a fresh (chat_id, firstRunner) + // heartbeat written by the first Acquire. + firstChat := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, firstWorker, firstChat.WorkerID.UUID) + require.Equal(t, firstRunner, firstChat.RunnerID.UUID) + _, err := f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: firstRunner, + }) + require.NoError(t, err, "first Acquire wrote a fresh heartbeat") + // Sanity check: heartbeat is not stale by the same threshold the + // machine uses for ownership-hint decisions. + stale, err := f.DB.IsChatHeartbeatStale(ctx, database.IsChatHeartbeatStaleParams{ + ChatID: created.Chat.ID, + RunnerID: firstRunner, + StaleSeconds: chatstate.HeartbeatStaleSeconds, + }) + require.NoError(t, err) + require.False(t, stale, "first runner's heartbeat is fresh before the second Acquire") + + // Snapshot publish counts before the takeover so we can assert + // Acquire does not publish an ownership hint itself. + ownershipBefore := f.Pub.ownershipPublishCount() + beforeChat := f.readChat(ctx, t, created.Chat.ID) + + secondWorker := uuid.New() + secondRunner := uuid.New() + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: secondWorker, RunnerID: secondRunner}) + return err + })) + + after := f.readChat(ctx, t, created.Chat.ID) + require.Equal(t, secondWorker, after.WorkerID.UUID, "ownership replaced") + require.Equal(t, secondRunner, after.RunnerID.UUID, "runner replaced") + require.Equal(t, beforeChat.SnapshotVersion+1, after.SnapshotVersion, "snapshot bumps exactly once") + f.Pub.expectChatUpdate(t, created.Chat.ID, after.SnapshotVersion) + + // The new (chat_id, secondRunner) heartbeat exists. The old + // (chat_id, firstRunner) row may or may not exist; Acquire is not + // responsible for cleaning it up. + _, err = f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: created.Chat.ID, + RunnerID: secondRunner, + }) + require.NoError(t, err, "second Acquire wrote a heartbeat for the new runner") + + // Acquire does not publish an ownership hint when it writes a fresh + // heartbeat. The post-commit ownership-hint logic in Update stays + // quiet because the new heartbeat is fresh, so no `chat:ownership` + // notification fires. + require.Equal(t, ownershipBefore, f.Pub.ownershipPublishCount(), + "Acquire must not publish an ownership hint when the resulting heartbeat is fresh") +} + +// TestTransitionAcquire_ExecutionStateOrthogonal verifies that Acquire +// preserves every execution-state field on the chat across +// representative valid execution states, including idle, runnable, and +// archived states. The transition only mutates ownership. +func TestTransitionAcquire_ExecutionStateOrthogonal(t *testing.T) { + t.Parallel() + + // Each setup leaves the chat in the named state and returns the + // chat ID for downstream assertions. + cases := []struct { + name string + state chatstate.ExecutionState + setup func(t *testing.T, f *testFixture) uuid.UUID + }{ + { + name: "R0", + state: chatstate.StateR0, + setup: func(t *testing.T, f *testFixture) uuid.UUID { + return createTestChat(t, f).Chat.ID + }, + }, + { + name: "W", + state: chatstate.StateW, + setup: func(t *testing.T, f *testFixture) uuid.UUID { + created := createTestChat(t, f) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + return created.Chat.ID + }, + }, + { + name: "E0", + state: chatstate.StateE0, + setup: func(t *testing.T, f *testFixture) uuid.UUID { + created := createTestChat(t, f) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(`{"message":"boom"}`), + Valid: true, + }, + }) + return err + })) + return created.Chat.ID + }, + }, + { + name: "I0", + state: chatstate.StateI0, + setup: func(t *testing.T, f *testFixture) uuid.UUID { + created := createTestChat(t, f) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Interrupt(chatstate.InterruptInput{Reason: "test"}) + return err + })) + return created.Chat.ID + }, + }, + { + name: "XW", + state: chatstate.StateXW, + setup: func(t *testing.T, f *testFixture) uuid.UUID { + created := createTestChat(t, f) + ctx := testutil.Context(t, testutil.WaitShort) + m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true}) + return err + })) + return created.Chat.ID + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + chatID := tc.setup(t, f) + require.Equal(t, tc.state, f.classify(ctx, t, chatID), "test setup must leave chat in %s", tc.state) + + before := f.readChat(ctx, t, chatID) + queueBefore, err := f.DB.CountChatQueuedMessages(ctx, chatID) + require.NoError(t, err) + historyBefore := historyMessageIDs(ctx, t, f, chatID) + + worker := uuid.New() + runner := uuid.New() + m := chatstate.NewChatMachine(f.DB, f.Pub, chatID) + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner}) + return err + })) + + after := f.readChat(ctx, t, chatID) + // Ownership updated. + require.Equal(t, worker, after.WorkerID.UUID) + require.Equal(t, runner, after.RunnerID.UUID) + // Execution state preserved. + require.Equal(t, before.Status, after.Status, "status preserved") + require.Equal(t, before.Archived, after.Archived, "archived flag preserved") + require.Equal(t, before.RequiresActionDeadlineAt, after.RequiresActionDeadlineAt, "requires-action deadline preserved") + require.Equal(t, before.LastError, after.LastError, "last_error preserved") + require.Equal(t, before.HistoryVersion, after.HistoryVersion, "history_version preserved") + require.Equal(t, before.QueueVersion, after.QueueVersion, "queue_version preserved") + require.Equal(t, before.GenerationAttempt, after.GenerationAttempt, "generation_attempt preserved") + // Classified state unchanged. + require.Equal(t, tc.state, f.classify(ctx, t, chatID), "execution state preserved by Acquire") + // Queue and history rows untouched. + queueAfter, err := f.DB.CountChatQueuedMessages(ctx, chatID) + require.NoError(t, err) + require.Equal(t, queueBefore, queueAfter, "queue cardinality preserved") + require.Equal(t, historyBefore, historyMessageIDs(ctx, t, f, chatID), "history preserved") + }) + } +} diff --git a/coderd/x/chatd/chatstate/trigger_test.go b/coderd/x/chatd/chatstate/trigger_test.go new file mode 100644 index 0000000000..dc31651ac2 --- /dev/null +++ b/coderd/x/chatd/chatstate/trigger_test.go @@ -0,0 +1,627 @@ +package chatstate_test + +import ( + "database/sql" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +// triggerFixture is a slim variant of testFixture that also exposes a +// raw *sql.DB so the trigger tests can run UPDATE/INSERT statements +// that bypass the typed sqlc layer. Tests that only need the typed +// store should keep using newTestFixture. +type triggerFixture struct { + f *testFixture + sqlDB *sql.DB +} + +func newTriggerFixture(t *testing.T) *triggerFixture { + t.Helper() + db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "openai", + BaseUrl: "http://example.invalid", + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + IsDefault: true, + }) + f := &testFixture{ + DB: db, + PubSub: ps, + Pub: newRecordingPubsub(), + User: user, + Org: org, + Model: model, + } + return &triggerFixture{f: f, sqlDB: sqlDB} +} + +// userMessageContent returns a marshaled user message body suitable +// for raw INSERT into chat_messages. +func userMessageContent(t *testing.T, text string) []byte { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + return raw.RawMessage +} + +// TestMessageInsertAssignsRevisionAndHistoryVersion verifies that +// inserting a chat message via the legacy InsertChatMessages query +// assigns NEW.revision from chats.snapshot_version (BEFORE trigger) +// and bumps chats.history_version + resets generation_attempt (AFTER +// STATEMENT trigger). +func TestMessageInsertAssignsRevisionAndHistoryVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + + created := createTestChat(t, f) + require.Equal(t, int64(1), created.Chat.SnapshotVersion) + require.Equal(t, int64(1), created.Chat.HistoryVersion) + + // Force generation_attempt > 0 so we can prove the trigger + // resets it on a new history change. + _, err := f.DB.IncrementChatGenerationAttempt(ctx, created.Chat.ID) + require.NoError(t, err) + before, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, int64(1), before.GenerationAttempt) + + // Bump snapshot_version directly to simulate a transition having + // taken the row lock. + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, before.SnapshotVersion+1, bumped.SnapshotVersion) + + // Insert a new assistant message via raw SQL so we know the + // BEFORE+AFTER triggers (and only those) decide revision and + // history_version. + content := userMessageContent(t, "hello-after-bump") + _, err = tf.sqlDB.ExecContext(ctx, ` + INSERT INTO chat_messages (chat_id, role, content, content_version, visibility) + VALUES ($1, 'assistant', $2::jsonb, $3, 'both') + `, created.Chat.ID, string(content), int(chatprompt.CurrentContentVersion)) + require.NoError(t, err) + + // History version equals snapshot_version, generation_attempt resets. + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, after.HistoryVersion) + require.Equal(t, int64(0), after.GenerationAttempt) + + // The inserted message picked up revision = bumped snapshot. + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: created.Chat.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, msgs) + last := msgs[len(msgs)-1] + require.Equal(t, database.ChatMessageRoleAssistant, last.Role) + require.Equal(t, bumped.SnapshotVersion, last.Revision) +} + +// TestMessageUpdateAssignsNewRevisionAndHistoryVersion verifies that +// updating a chat message's content advances NEW.revision to the +// current chats.snapshot_version and that chats.history_version +// bumps to match. +func TestMessageUpdateAssignsNewRevisionAndHistoryVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: created.Chat.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, msgs) + target := msgs[0] + originalRevision := target.Revision + + // Bump the snapshot so the trigger sees a new revision target. + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + require.Greater(t, bumped.SnapshotVersion, originalRevision) + + newContent := userMessageContent(t, "edited content") + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_messages SET content = $1::jsonb WHERE id = $2 + `, string(newContent), target.ID) + require.NoError(t, err) + + reloaded, err := f.DB.GetChatMessageByID(ctx, target.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, reloaded.Revision, + "updated message picks up the current snapshot version") + + chatAfter, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, chatAfter.HistoryVersion) + require.Equal(t, int64(0), chatAfter.GenerationAttempt, + "history change resets generation_attempt") +} + +// TestMessageRevisionCannotBeSetByRuntimeCode verifies the BEFORE +// trigger rejects explicit revision values on INSERT and rejects +// revision changes on UPDATE. +func TestMessageRevisionCannotBeSetByRuntimeCode(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + content := userMessageContent(t, "explicit revision") + _, err := tf.sqlDB.ExecContext(ctx, ` + INSERT INTO chat_messages (chat_id, role, content, content_version, visibility, revision) + VALUES ($1, 'user', $2::jsonb, $3, 'both', 999) + `, created.Chat.ID, string(content), int(chatprompt.CurrentContentVersion)) + require.Error(t, err, "INSERT with explicit revision must be rejected") + require.Contains(t, err.Error(), "revision must be assigned by trigger") + + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: created.Chat.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, msgs) + target := msgs[0] + + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_messages SET revision = revision + 100 WHERE id = $1 + `, target.ID) + require.Error(t, err, "UPDATE that changes revision must be rejected") + require.Contains(t, err.Error(), "revision must be assigned by trigger") +} + +// TestMessageChatIDCannotChange verifies the BEFORE trigger rejects +// updates that change chat_messages.chat_id. +func TestMessageChatIDCannotChange(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + first := createTestChat(t, f) + second := createTestChat(t, f) + + firstMsgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: first.Chat.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, firstMsgs) + target := firstMsgs[0] + + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_messages SET chat_id = $1 WHERE id = $2 + `, second.Chat.ID, target.ID) + require.Error(t, err, "UPDATE that changes chat_id must be rejected") + require.Contains(t, err.Error(), "chat_id is immutable") +} + +// TestNoopMessageUpdateDoesNotAdvanceHistoryVersion verifies that a +// no-op UPDATE on a chat_messages row (one whose OLD and NEW are +// indistinguishable) does NOT advance chats.history_version even +// when the snapshot was previously bumped. This guards against the +// AFTER UPDATE STATEMENT trigger naively reacting to every touched +// row id regardless of whether the row actually changed. +func TestNoopMessageUpdateDoesNotAdvanceHistoryVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + msgs, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: created.Chat.ID, + }) + require.NoError(t, err) + require.NotEmpty(t, msgs) + target := msgs[0] + originalRevision := target.Revision + + // Bump snapshot so the AFTER STATEMENT guard + // (history_version != snapshot_version) is now true. + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + require.NotEqual(t, bumped.SnapshotVersion, bumped.HistoryVersion, + "snapshot bump leaves history_version trailing") + + // No-op UPDATE: SET content = content. OLD IS NOT DISTINCT FROM NEW. + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_messages SET content = content WHERE id = $1 + `, target.ID) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.HistoryVersion, after.HistoryVersion, + "no-op update must NOT advance history_version") + + // And the row's revision is untouched. + reloaded, err := f.DB.GetChatMessageByID(ctx, target.ID) + require.NoError(t, err) + require.Equal(t, originalRevision, reloaded.Revision, + "no-op update must NOT advance message revision") +} + +// Queue version triggers + +// TestQueueInsertUpdatesQueueVersion verifies that an INSERT into +// chat_queued_messages bumps chats.queue_version to the current +// snapshot_version. +func TestQueueInsertUpdatesQueueVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + before, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, int64(0), before.QueueVersion) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + + content := userMessageContent(t, "queued") + _, err = f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: content, + CreatedBy: f.User.ID, + }) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, after.QueueVersion, + "INSERT into chat_queued_messages bumps queue_version") +} + +// TestQueuedMessageCreatedByIsRequired verifies the database enforces +// creator metadata for every queued message row. +func TestQueuedMessageCreatedByIsRequired(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + content := userMessageContent(t, "queued-without-creator") + _, err := tf.sqlDB.ExecContext(ctx, ` + INSERT INTO chat_queued_messages (chat_id, content, model_config_id, created_by) + VALUES ($1, $2::jsonb, NULL, NULL) + `, created.Chat.ID, string(content)) + require.Error(t, err) + require.Contains(t, err.Error(), "created_by") +} + +func TestLegacyQueuedMessageInsertUsesChatOwnerAsCreator(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + queued, err := f.DB.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "legacy-queued"), + }) + require.NoError(t, err) + require.Equal(t, created.Chat.OwnerID, queued.CreatedBy) +} + +// TestQueueUpdateContentUpdatesQueueVersion verifies that an UPDATE +// of chat_queued_messages.content bumps queue_version. The +// AFTER UPDATE trigger explicitly listens for content changes. +func TestQueueUpdateContentUpdatesQueueVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "initial"), + CreatedBy: f.User.ID, + }) + require.NoError(t, err) + + before, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + require.Greater(t, bumped.SnapshotVersion, before.QueueVersion) + + updated := userMessageContent(t, "updated") + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_queued_messages SET content = $1::jsonb WHERE id = $2 + `, string(updated), queued.ID) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, after.QueueVersion, + "UPDATE of queued content bumps queue_version") +} + +// TestQueueUpdatePositionUpdatesQueueVersion verifies that an UPDATE +// of chat_queued_messages.position (such as the reorder-to-head +// path) bumps queue_version. +func TestQueueUpdatePositionUpdatesQueueVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + q1, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "first"), + CreatedBy: f.User.ID, + }) + require.NoError(t, err) + q2, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "second"), + CreatedBy: f.User.ID, + }) + require.NoError(t, err) + require.NotEqual(t, q1.ID, q2.ID) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + + // Move q2 to head by setting its position to q1.position - 1. + _, err = tf.sqlDB.ExecContext(ctx, ` + UPDATE chat_queued_messages SET position = $1 WHERE id = $2 + `, q1.Position-1, q2.ID) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, after.QueueVersion, + "UPDATE of queued position bumps queue_version") +} + +// TestQueueDeleteUpdatesQueueVersion verifies that DELETE from +// chat_queued_messages bumps queue_version. +func TestQueueDeleteUpdatesQueueVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + queued, err := f.DB.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: created.Chat.ID, + Content: userMessageContent(t, "to delete"), + CreatedBy: f.User.ID, + }) + require.NoError(t, err) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + + rows, err := f.DB.DeleteChatQueuedMessageReturningCount(ctx, database.DeleteChatQueuedMessageReturningCountParams{ + ID: queued.ID, + ChatID: created.Chat.ID, + }) + require.NoError(t, err) + require.Equal(t, int64(1), rows) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, bumped.SnapshotVersion, after.QueueVersion, + "DELETE from queue bumps queue_version") +} + +// TestNonQueueUpdateDoesNotUpdateQueueVersion verifies that mutations +// on other chat-related tables do NOT bump queue_version. The +// canonical case is inserting a chat message: it must update +// history_version but leave queue_version untouched. +func TestNonQueueUpdateDoesNotUpdateQueueVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + before, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + + content := userMessageContent(t, "non-queue mutation") + _, err = tf.sqlDB.ExecContext(ctx, ` + INSERT INTO chat_messages (chat_id, role, content, content_version, visibility) + VALUES ($1, 'assistant', $2::jsonb, $3, 'both') + `, created.Chat.ID, string(content), int(chatprompt.CurrentContentVersion)) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, before.QueueVersion, after.QueueVersion, + "chat_messages INSERT must not bump queue_version") + // Sanity: history_version DID move. + require.Equal(t, bumped.SnapshotVersion, after.HistoryVersion) +} + +// Retry state triggers + +func TestRetryStateDefaults(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + chat, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.False(t, chat.RetryState.Valid) + require.Equal(t, int64(0), chat.RetryStateVersion) +} + +func TestRetryStateUpdateSetsRetryStateVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + + after, err := f.DB.UpdateChatRetryState(ctx, database.UpdateChatRetryStateParams{ + ID: created.Chat.ID, + RetryState: []byte(`{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`), + }) + require.NoError(t, err) + require.True(t, after.RetryState.Valid) + require.JSONEq(t, + `{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`, + string(after.RetryState.RawMessage)) + require.Equal(t, bumped.SnapshotVersion, after.RetryStateVersion) +} + +func TestRetryStateSameValueDoesNotUpdateRetryStateVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + payload := []byte(`{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`) + _, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + first, err := f.DB.UpdateChatRetryState(ctx, database.UpdateChatRetryStateParams{ + ID: created.Chat.ID, + RetryState: payload, + }) + require.NoError(t, err) + + _, err = f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + second, err := f.DB.UpdateChatRetryState(ctx, database.UpdateChatRetryStateParams{ + ID: created.Chat.ID, + RetryState: payload, + }) + require.NoError(t, err) + require.Equal(t, first.RetryStateVersion, second.RetryStateVersion, + "same retry_state payload must not update retry_state_version") +} + +func TestGenerationAttemptClearsRetryState(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + _, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + withRetry, err := f.DB.UpdateChatRetryState(ctx, database.UpdateChatRetryStateParams{ + ID: created.Chat.ID, + RetryState: []byte(`{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`), + }) + require.NoError(t, err) + require.True(t, withRetry.RetryState.Valid) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + attempt, err := f.DB.IncrementChatGenerationAttempt(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, int64(1), attempt) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.False(t, after.RetryState.Valid) + require.Equal(t, bumped.SnapshotVersion, after.RetryStateVersion, + "clearing retry_state on generation attempt bumps retry_state_version") +} + +func TestGenerationAttemptWithNullRetryStateDoesNotUpdateRetryStateVersion(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + before, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.False(t, before.RetryState.Valid) + + _, err = f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + _, err = f.DB.IncrementChatGenerationAttempt(ctx, created.Chat.ID) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.False(t, after.RetryState.Valid) + require.Equal(t, before.RetryStateVersion, after.RetryStateVersion, + "generation attempt with null retry_state leaves retry_state_version unchanged") +} + +func TestRetryStateVersionCannotBeSetByRuntimeCode(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + _, err := tf.sqlDB.ExecContext(ctx, ` + UPDATE chats SET retry_state_version = retry_state_version + 1 WHERE id = $1 + `, created.Chat.ID) + require.Error(t, err) + require.Contains(t, err.Error(), "retry_state_version must be assigned by trigger") +} + +func TestHistoryChangeClearsRetryState(t *testing.T) { + t.Parallel() + tf := newTriggerFixture(t) + f := tf.f + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + + _, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + _, err = f.DB.IncrementChatGenerationAttempt(ctx, created.Chat.ID) + require.NoError(t, err) + _, err = f.DB.UpdateChatRetryState(ctx, database.UpdateChatRetryStateParams{ + ID: created.Chat.ID, + RetryState: []byte(`{"attempt":1,"delay_ms":250,"error":"retry","retrying_at":"2026-05-29T00:00:00Z"}`), + }) + require.NoError(t, err) + + bumped, err := f.DB.LockChatAndBumpSnapshotVersion(ctx, created.Chat.ID) + require.NoError(t, err) + content := userMessageContent(t, "history clears retry state") + _, err = tf.sqlDB.ExecContext(ctx, ` + INSERT INTO chat_messages (chat_id, role, content, content_version, visibility) + VALUES ($1, 'assistant', $2::jsonb, $3, 'both') + `, created.Chat.ID, string(content), int(chatprompt.CurrentContentVersion)) + require.NoError(t, err) + + after, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, int64(0), after.GenerationAttempt) + require.False(t, after.RetryState.Valid) + require.Equal(t, bumped.SnapshotVersion, after.RetryStateVersion, + "history reset of generation_attempt clears retry_state") +} diff --git a/coderd/x/chatd/chatstate_bridge.go b/coderd/x/chatd/chatstate_bridge.go new file mode 100644 index 0000000000..2a6f394d4d --- /dev/null +++ b/coderd/x/chatd/chatstate_bridge.go @@ -0,0 +1,53 @@ +package chatd + +import ( + "database/sql" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" +) + +// newChatMachine constructs a chat-scoped state machine handle bound to +// the server's database and pubsub. +func (p *Server) newChatMachine(chatID uuid.UUID) *chatstate.ChatMachine { + return chatstate.NewChatMachine(p.db, p.pubsub, chatID) +} + +// systemMessage builds a chatstate.Message representing a system +// prompt entry for the initial-history slice of CreateChat. +func systemMessage(rawContent pqtype.NullRawMessage, modelConfigID uuid.UUID) chatstate.Message { + return chatstate.Message{ + Role: database.ChatMessageRoleSystem, + Content: rawContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + } +} + +func userMessageWithAPIKeyID(rawContent pqtype.NullRawMessage, modelConfigID, createdBy uuid.UUID, apiKeyID string) chatstate.Message { + return chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: rawContent, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: createdBy != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + } +} + +// busyBehaviorToChatState converts the public busy-behavior enum used +// by the server API to the chatstate variant. +func busyBehaviorToChatState(b SendMessageBusyBehavior) chatstate.BusyBehavior { + switch b { + case SendMessageBusyBehaviorInterrupt: + return chatstate.BusyBehaviorInterrupt + default: + return chatstate.BusyBehaviorQueue + } +} diff --git a/coderd/x/chatd/chattest/anthropic.go b/coderd/x/chatd/chattest/anthropic.go index cb5ffe5dc5..ba23571b8d 100644 --- a/coderd/x/chatd/chattest/anthropic.go +++ b/coderd/x/chatd/chattest/anthropic.go @@ -26,7 +26,9 @@ type AnthropicResponse struct { type AnthropicRequest struct { *http.Request // Embed http.Request Model string `json:"model"` + System json.RawMessage `json:"system,omitempty"` Messages []AnthropicRequestMessage `json:"messages"` + Tools []AnthropicRequestTool `json:"tools,omitempty"` Stream bool `json:"stream,omitempty"` MaxTokens int `json:"max_tokens,omitempty"` // TODO: encoding/json ignores inline tags. Add custom UnmarshalJSON to capture unknown keys. @@ -40,6 +42,11 @@ type AnthropicRequestMessage struct { Content json.RawMessage `json:"content"` } +// AnthropicRequestTool represents a tool in an Anthropic request. +type AnthropicRequestTool struct { + Name string `json:"name"` +} + // AnthropicMessage represents a message in an Anthropic response. type AnthropicMessage struct { ID string `json:"id,omitempty"` @@ -59,6 +66,13 @@ type AnthropicUsage struct { CacheReadInputTokens int `json:"cache_read_input_tokens,omitempty"` } +// AnthropicReasoningBlock describes one Anthropic thinking block for a +// streaming test response. +type AnthropicReasoningBlock struct { + Text string + Signature string +} + // AnthropicChunk represents a streaming chunk from Anthropic. type AnthropicChunk struct { Type string `json:"type"` @@ -83,17 +97,22 @@ type AnthropicChunkMessage struct { // AnthropicContentBlock represents a content block in a chunk. type AnthropicContentBlock struct { - Type string `json:"type"` - Text string `json:"text,omitempty"` - ID string `json:"id,omitempty"` - Name string `json:"name,omitempty"` - Input json.RawMessage `json:"input,omitempty"` + Type string `json:"type"` + Text string `json:"text,omitempty"` + Thinking string `json:"thinking,omitempty"` + ID string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Input json.RawMessage `json:"input,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + Content any `json:"content,omitempty"` } // AnthropicDeltaBlock represents a delta block in a chunk. type AnthropicDeltaBlock struct { Type string `json:"type"` Text string `json:"text,omitempty"` + Thinking string `json:"thinking,omitempty"` + Signature string `json:"signature,omitempty"` PartialJSON string `json:"partial_json,omitempty"` } @@ -176,8 +195,7 @@ func (s *anthropicServer) writeResponse(w http.ResponseWriter, req *AnthropicReq } } -func (s *anthropicServer) writeStreamingResponse(w http.ResponseWriter, chunks <-chan AnthropicChunk) { - _ = s // receiver unused but kept for consistency +func (*anthropicServer) writeStreamingResponse(w http.ResponseWriter, chunks <-chan AnthropicChunk) { w.Header().Set("Content-Type", "text/event-stream") w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") @@ -424,6 +442,95 @@ func AnthropicTextChunksWithCacheUsage(usage AnthropicUsage, deltas ...string) [ return chunks } +// AnthropicReasoningTextChunks creates a streaming response with one or more +// thinking blocks followed by one text block. +func AnthropicReasoningTextChunks(reasoning []AnthropicReasoningBlock, text string) []AnthropicChunk { + messageID := fmt.Sprintf("msg-%s", uuid.New().String()[:8]) + model := "claude-3-opus-20240229" + + chunks := []AnthropicChunk{ + { + Type: "message_start", + Message: AnthropicChunkMessage{ + ID: messageID, + Type: "message", + Role: "assistant", + Model: model, + }, + }, + } + + for i, block := range reasoning { + chunks = append(chunks, + AnthropicChunk{ + Type: "content_block_start", + Index: i, + ContentBlock: AnthropicContentBlock{ + Type: "thinking", + Thinking: "", + }, + }, + AnthropicChunk{ + Type: "content_block_delta", + Index: i, + Delta: AnthropicDeltaBlock{ + Type: "thinking_delta", + Thinking: block.Text, + }, + }, + ) + if block.Signature != "" { + chunks = append(chunks, AnthropicChunk{ + Type: "content_block_delta", + Index: i, + Delta: AnthropicDeltaBlock{ + Type: "signature_delta", + Signature: block.Signature, + }, + }) + } + chunks = append(chunks, AnthropicChunk{ + Type: "content_block_stop", + Index: i, + }) + } + + textIndex := len(reasoning) + chunks = append(chunks, + AnthropicChunk{ + Type: "content_block_start", + Index: textIndex, + ContentBlock: AnthropicContentBlock{ + Type: "text", + Text: "", + }, + }, + AnthropicChunk{ + Type: "content_block_delta", + Index: textIndex, + Delta: AnthropicDeltaBlock{ + Type: "text_delta", + Text: text, + }, + }, + AnthropicChunk{ + Type: "content_block_stop", + Index: textIndex, + }, + AnthropicChunk{ + Type: "message_delta", + StopReason: "end_turn", + Usage: AnthropicUsage{ + InputTokens: 10, + OutputTokens: 5, + }, + }, + AnthropicChunk{Type: "message_stop"}, + ) + + return chunks +} + // AnthropicToolCallChunks creates a complete streaming response for a tool call. // Input JSON can be split across multiple deltas, matching Anthropic's // input_json_delta streaming behavior. diff --git a/coderd/x/chatd/chattest/errors.go b/coderd/x/chatd/chattest/errors.go index 2c84339600..b9b3f5d759 100644 --- a/coderd/x/chatd/chattest/errors.go +++ b/coderd/x/chatd/chattest/errors.go @@ -43,17 +43,6 @@ func AnthropicErrorResponse(statusCode int, errorType, message string) Anthropic } } -// AnthropicOverloadedResponse returns a 529 "overloaded" error matching -// Anthropic's overloaded response format. -func AnthropicOverloadedResponse() AnthropicResponse { - return AnthropicErrorResponse(529, "overloaded_error", "Overloaded") -} - -// AnthropicRateLimitResponse returns a 429 rate limit error. -func AnthropicRateLimitResponse() AnthropicResponse { - return AnthropicErrorResponse(http.StatusTooManyRequests, "rate_limit_error", "Rate limited") -} - // OpenAIErrorResponse returns an OpenAIResponse that causes the // test server to respond with the given HTTP status code and error. func OpenAIErrorResponse(statusCode int, errorType, message string) OpenAIResponse { diff --git a/coderd/x/chatd/chattool/createworkspace.go b/coderd/x/chatd/chattool/createworkspace.go index f65247fd02..168eacb186 100644 --- a/coderd/x/chatd/chattool/createworkspace.go +++ b/coderd/x/chatd/chattool/createworkspace.go @@ -637,11 +637,10 @@ func waitForAgentReady( var lastErr error for { attemptCtx, attemptCancel := context.WithTimeout(agentCtx, agentAttemptTimeout) - conn, release, err := agentConnFn(attemptCtx, agentID) + _, release, err := agentConnFn(attemptCtx, agentID) attemptCancel() if err == nil { release() - _ = conn break } lastErr = err diff --git a/coderd/x/chatd/context_helpers.go b/coderd/x/chatd/context_helpers.go new file mode 100644 index 0000000000..be684c7c6d --- /dev/null +++ b/coderd/x/chatd/context_helpers.go @@ -0,0 +1,82 @@ +package chatd + +import ( + "bytes" + "encoding/json" + + "github.com/google/uuid" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/codersdk" +) + +// agentChatContextSentinelPath marks the synthetic empty context-file +// part used to record an attempted workspace-context fetch when no +// AGENTS.md content is available. It mirrors the constant of the same +// value in the chatd package so the worker can recognize sentinel +// parts without importing chatd (which would be a cycle). +const agentChatContextSentinelPath = ".coder/agent-chat-context-sentinel" + +// contextFileAgentIDFromMessages returns the most recent workspace +// agent ID stamped on a persisted context-file part, ignoring the +// skill-only sentinel. Returns uuid.Nil, false when no stamped +// non-sentinel context-file parts exist. +// +// This mirrors chatd.contextFileAgentID. It is duplicated here as a +// small pure helper so chatworker can decide whether workspace +// context is current without importing chatd. +func contextFileAgentIDFromMessages(messages []database.ChatMessage) (uuid.UUID, bool) { + var lastID uuid.UUID + found := false + for _, msg := range messages { + if !msg.Content.Valid || !bytes.Contains(msg.Content.RawMessage, []byte(`"context-file"`)) { + continue + } + var parts []codersdk.ChatMessagePart + if err := json.Unmarshal(msg.Content.RawMessage, &parts); err != nil { + continue + } + for _, p := range parts { + if p.Type != codersdk.ChatMessagePartTypeContextFile || + !p.ContextFileAgentID.Valid || + p.ContextFilePath == agentChatContextSentinelPath { + continue + } + lastID = p.ContextFileAgentID.UUID + found = true + break + } + } + return lastID, found +} + +// hasPersistedContextFileForAgent reports whether messages include +// any persisted context-file marker for the given agent, including +// the skill-only sentinel. This is true once the +// persist_workspace_context action has committed at least one +// context-file row for the agent (with or without content), so a +// subsequent decision pass will not loop on the same agent. +func hasPersistedContextFileForAgent(messages []database.ChatMessage, agentID uuid.UUID) bool { + if agentID == uuid.Nil { + return false + } + for _, msg := range messages { + if !msg.Content.Valid || !bytes.Contains(msg.Content.RawMessage, []byte(`"context-file"`)) { + continue + } + var parts []codersdk.ChatMessagePart + if err := json.Unmarshal(msg.Content.RawMessage, &parts); err != nil { + continue + } + for _, p := range parts { + if p.Type != codersdk.ChatMessagePartTypeContextFile || + !p.ContextFileAgentID.Valid { + continue + } + if p.ContextFileAgentID.UUID == agentID { + return true + } + } + } + return false +} diff --git a/coderd/x/chatd/export_test.go b/coderd/x/chatd/export_test.go deleted file mode 100644 index 519ed0dcad..0000000000 --- a/coderd/x/chatd/export_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package chatd - -import ( - "context" - - "github.com/sqlc-dev/pqtype" - - "cdr.dev/slog/v3" - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/codersdk" -) - -// FinishActiveChatForTest exposes the unexported cleanup TX so tests -// can drive the post-run state machine deterministically. Returns the -// resulting chat, the promoted message (if any), the synthetic -// tool-result rows the cleanup TX inserted (if any), and the cleanup -// error. The lastError string is encoded into a structured payload -// the same way runChat does, so callers do not need to know about -// the structured-error wrapper. -func FinishActiveChatForTest( - ctx context.Context, - server *Server, - chat database.Chat, - status database.ChatStatus, - lastError string, -) (database.Chat, *database.ChatMessage, []database.ChatMessage, error) { - logger := server.logger.With(slog.F("chat_id", chat.ID)) - var encoded pqtype.NullRawMessage - if lastError != "" { - var err error - encoded, err = encodeChatLastErrorPayload(&codersdk.ChatError{ - Message: lastError, - }) - if err != nil { - return database.Chat{}, nil, nil, err - } - } - result, err := server.finishActiveChat(ctx, logger, chat, status, encoded) - if err != nil { - return database.Chat{}, nil, nil, err - } - return result.updatedChat, result.promotedMessage, result.syntheticToolResults, nil -} - -// RecoverStaleChatsForTest exposes the unexported stale-recovery loop -// so tests can assert the recovery state machine without waiting for -// the periodic ticker. -func RecoverStaleChatsForTest(ctx context.Context, server *Server) { - server.recoverStaleChats(ctx) -} - -// InsertSyntheticToolResultsTxForTest exposes the unexported helper -// so tests can verify the dedup path against pre-existing tool -// results. -func InsertSyntheticToolResultsTxForTest( - ctx context.Context, - store database.Store, - chat database.Chat, - reason string, -) ([]database.ChatMessage, error) { - return insertSyntheticToolResultsTx(ctx, store, chat, reason) -} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go new file mode 100644 index 0000000000..47b17c8e45 --- /dev/null +++ b/coderd/x/chatd/generation.go @@ -0,0 +1,1109 @@ +package chatd + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "strings" + "time" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "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/chatloop" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatretry" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" +) + +// generationPrepareInput contains the committed state used to prepare one +// generation action. +type generationPrepareInput struct { + Chat database.Chat + Messages []database.ChatMessage + ChainModeDisabled bool +} + +// generationPrepared contains the side-effect inputs for a generation task. +type generationPrepared struct { + Chat database.Chat + Messages []database.ChatMessage + + Model fantasy.LanguageModel + Prompt []fantasy.Message + Tools []fantasy.AgentTool + ActiveTools []string + ProviderTools []chatloop.ProviderTool + ProviderKeys chatprovider.ProviderAPIKeys + ModelRoute resolvedModelRoute + ModelBuildOptions modelBuildOptions + + ModelConfigID uuid.UUID + ModelConfig codersdk.ChatModelCallConfig + ProviderOptions fantasy.ProviderOptions + ContextLimitFallback int64 + + DynamicToolNames map[string]bool + StopAfterTools map[string]struct{} + ExclusiveToolNames map[string]bool + BuiltinToolNames map[string]bool + ToolNameToConfigID map[string]uuid.UUID + + MaxSteps int + Compaction *generationCompaction + // Cleanup is always non-nil when prepareGeneration succeeds. + Cleanup func() + + Debug *generationDebug + + // WorkspaceContextEligible reports whether the current turn is allowed + // by policy to inject workspace context. The decision helper combines + // this fact with committed chat metadata and history to decide whether + // the persist_workspace_context action should run. + WorkspaceContextEligible bool +} + +// generationCompaction contains compaction inputs prepared for generation. +type generationCompaction struct { + Required bool + Options chatloop.GenerateCompactionOptions +} + +type generationDebug struct { + Enabled bool + Service *chatdebug.Service + Provider string + Model string + TriggerMessageID int64 + HistoryTipMessageID int64 + TriggerLabel string + ModelConfig database.ChatModelConfig +} + +type workspaceContextBuildInput struct { + Chat database.Chat + Messages []database.ChatMessage + ActiveAPIKeyID string +} + +type workspaceContextBuildResult struct { + Messages []chatstate.Message +} + +// generationOutcome describes a completed generation outcome. +type generationOutcome struct { + Chat database.Chat + Kind runnerActionKind + WatchEventKind codersdk.ChatWatchEventKind + LastError string + PromotedMessageID int64 + InsertedMessages []runnerActionMessage +} + +type generationActionKind string + +const ( + generationActionExecuteLocalTools generationActionKind = "execute_local_tools" + generationActionEnterRequiresAction generationActionKind = "enter_requires_action" + generationActionFinishTurn generationActionKind = "finish_turn" + generationActionCompact generationActionKind = "compact" + generationActionGenerateAssistant generationActionKind = "generate_assistant" + generationActionPersistWorkspaceContext generationActionKind = "persist_workspace_context" +) + +type generationFinishReason string + +const ( + generationFinishReasonStopAfterTool generationFinishReason = "stop_after_tool" + generationFinishReasonComplete generationFinishReason = "complete" + generationFinishReasonMaxSteps generationFinishReason = "max_steps" +) + +type compactionTrigger string + +const ( + compactionTriggerRequired compactionTrigger = "required" + compactionTriggerAlreadyCompacted compactionTrigger = "already_compacted" +) + +var errCompactionStillOverLimit = xerrors.New("compaction left the chat above the compaction limit") + +type generationDecision struct { + kind generationActionKind + localToolCalls []fantasy.ToolCallContent + pendingDynamicToolCalls []pendingDynamicToolCall + finishReason generationFinishReason + compactionTrigger compactionTrigger + promotedMessageID int64 +} + +type generationRetryDecision struct { + retry bool + generationAttempt int64 + delay time.Duration +} + +var errRetryStateDecisionOnly = xerrors.New("retry state decision only") + +// errTerminalGeneration marks a prepare or decide failure as terminal: a +// deterministic error where retrying cannot help. The generation loop +// finishes the turn with an error instead of retrying when an error +// unwraps to this sentinel. +var errTerminalGeneration = xerrors.New("terminal generation error") + +type terminalGenerationError struct{ err error } + +func (e terminalGenerationError) Error() string { return e.err.Error() } + +func (e terminalGenerationError) Unwrap() error { return errors.Join(errTerminalGeneration, e.err) } + +// terminalGeneration wraps err so the prepare/decide retry loop stops +// immediately and finishes the turn with an error. +func terminalGeneration(err error) error { + if err == nil { + return nil + } + return terminalGenerationError{err: err} +} + +func isTerminalGeneration(err error) bool { + return errors.Is(err, errTerminalGeneration) +} + +type generationDecisionInput struct { + chat database.Chat + messages []database.ChatMessage + dynamicToolNames map[string]bool + exclusiveToolNames map[string]bool + stopAfterTools map[string]struct{} + maxSteps int + compactionEnabled bool + compactionNeeded bool + workspaceContextEligible bool +} + +// shouldPersistWorkspaceContext reports whether the committed chat +// state and history indicate that the persistWorkspaceContext +// generation action should run before the next assistant call. The +// decision uses two facts: +// - chat metadata says a workspace and selected agent are attached; +// - committed history either has no context-file marker for the +// currently selected workspace agent, or the latest non-sentinel +// marker points to a different agent. +// +// The decision is intentionally pure so generation can choose the +// action without dialing the workspace. Once the action commits a +// context-file marker for the agent (with or without content), this +// helper returns false on the next pass and the loop is broken. +func shouldPersistWorkspaceContext(chat database.Chat, messages []database.ChatMessage) bool { + if !chat.WorkspaceID.Valid || !chat.AgentID.Valid { + return false + } + if hasPersistedContextFileForAgent(messages, chat.AgentID.UUID) { + return false + } + persistedAgentID, found := contextFileAgentIDFromMessages(messages) + if !found { + return true + } + return persistedAgentID != chat.AgentID.UUID +} + +func decideGenerationAction(input generationDecisionInput) (generationDecision, error) { + localCalls, dynamicCalls, err := unresolvedToolCallsFromHistory(input.messages, input.dynamicToolNames) + if err != nil { + return generationDecision{}, err + } + if len(localCalls) > 0 { + if len(dynamicCalls) > 0 && hasExclusiveToolCall(localCalls, input.exclusiveToolNames) { + for _, dynamicCall := range dynamicCalls { + localCalls = append(localCalls, fantasy.ToolCallContent{ + ToolCallID: dynamicCall.ToolCallID, + ToolName: dynamicCall.ToolName, + Input: dynamicCall.Args, + }) + } + dynamicCalls = nil + } + return generationDecision{kind: generationActionExecuteLocalTools, localToolCalls: localCalls, pendingDynamicToolCalls: dynamicCalls}, nil + } + if len(dynamicCalls) > 0 { + return generationDecision{kind: generationActionEnterRequiresAction, pendingDynamicToolCalls: dynamicCalls}, nil + } + + stopAfter, err := historyHasStopAfterToolResult(input.messages, input.stopAfterTools) + if err != nil { + return generationDecision{}, err + } + if stopAfter { + return generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonStopAfterTool}, nil + } + complete, err := currentHistoryComplete(input.messages) + if err != nil { + return generationDecision{}, err + } + if complete { + return generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, nil + } + if input.maxSteps > 0 && currentTurnStepCount(input.messages) >= input.maxSteps { + return generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonMaxSteps}, nil + } + if input.workspaceContextEligible && shouldPersistWorkspaceContext(input.chat, input.messages) { + return generationDecision{kind: generationActionPersistWorkspaceContext}, nil + } + compactionRequirement := compactionRequirementNotNeeded + if input.compactionEnabled && input.compactionNeeded { + compactionRequirement = compactionRequirementNeeded + } + switch compactionStatusFromHistory(input.messages, compactionRequirement) { + case compactionStatusNeeded: + return generationDecision{kind: generationActionCompact, compactionTrigger: compactionTriggerRequired}, nil + case compactionStatusAfterCompaction: + return generationDecision{kind: generationActionGenerateAssistant, compactionTrigger: compactionTriggerAlreadyCompacted}, nil + case compactionStatusStillOverLimit: + return generationDecision{}, terminalGeneration(errCompactionStillOverLimit) + case compactionStatusNotNeeded: + return generationDecision{kind: generationActionGenerateAssistant}, nil + default: + return generationDecision{}, terminalGeneration(xerrors.New("unknown compaction status")) + } +} + +func unresolvedToolCallsFromHistory( + messages []database.ChatMessage, + dynamicToolNames map[string]bool, +) ([]fantasy.ToolCallContent, []pendingDynamicToolCall, error) { + assistantIndex := lastMessageIndex(messages, func(msg database.ChatMessage) bool { + return msg.Role == database.ChatMessageRoleAssistant + }) + if assistantIndex == -1 { + return nil, nil, nil + } + assistantParts, err := chatprompt.ParseContent(messages[assistantIndex]) + if err != nil { + return nil, nil, xerrors.Errorf("parse assistant message: %w", err) + } + handled, err := handledToolCallIDs(messages[assistantIndex+1:]) + if err != nil { + return nil, nil, err + } + localCalls := make([]fantasy.ToolCallContent, 0) + dynamicCalls := make([]pendingDynamicToolCall, 0) + for _, part := range assistantParts { + if part.Type != codersdk.ChatMessagePartTypeToolCall || part.ProviderExecuted || handled[part.ToolCallID] { + continue + } + if dynamicToolNames[part.ToolName] { + dynamicCalls = append(dynamicCalls, pendingDynamicToolCall{ + ToolCallID: part.ToolCallID, + ToolName: part.ToolName, + Args: string(part.Args), + }) + continue + } + localCalls = append(localCalls, fantasy.ToolCallContent{ + ToolCallID: part.ToolCallID, + ToolName: part.ToolName, + Input: string(part.Args), + ProviderExecuted: part.ProviderExecuted, + }) + } + return localCalls, dynamicCalls, nil +} + +func hasExclusiveToolCall(toolCalls []fantasy.ToolCallContent, exclusiveToolNames map[string]bool) bool { + if len(exclusiveToolNames) == 0 { + return false + } + for _, toolCall := range toolCalls { + if exclusiveToolNames[toolCall.ToolName] { + return true + } + } + return false +} + +func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskStartInput) error { + if s.server == nil { + return xerrors.New("chatworker: server is required") + } + machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID) + chainModeDisabled := false + for { + locked, messages, err := loadGenerationState(ctx, machine, input) + if err != nil { + return err + } + prepareInput := generationPrepareInput{ + Chat: locked, + Messages: messages, + ChainModeDisabled: chainModeDisabled, + } + prepared, err := retryGenerationPhase(ctx, s.waitGenerationPhaseBackoff, func() (generationPrepared, error) { + return s.server.prepareGeneration(ctx, prepareInput) + }) + if err != nil { + if errors.Is(err, errTaskExpectedExit) { + return errTaskExpectedExit + } + return s.finishGenerationError(ctx, machine, input, 0, err, generationAttemptNotRequired) + } + cleanup := prepared.Cleanup + decision, err := retryGenerationPhase(ctx, s.waitGenerationPhaseBackoff, func() (generationDecision, error) { + return decideGenerationAction(generationDecisionInput{ + chat: prepared.Chat, + messages: prepared.Messages, + dynamicToolNames: prepared.DynamicToolNames, + exclusiveToolNames: prepared.ExclusiveToolNames, + stopAfterTools: prepared.StopAfterTools, + maxSteps: prepared.MaxSteps, + compactionEnabled: prepared.Compaction != nil, + compactionNeeded: prepared.Compaction != nil && prepared.Compaction.Required, + workspaceContextEligible: prepared.WorkspaceContextEligible, + }) + }) + if err != nil { + cleanup() + if errors.Is(err, errTaskExpectedExit) { + return errTaskExpectedExit + } + return s.finishGenerationError(ctx, machine, input, 0, err, generationAttemptNotRequired) + } + + var actionErr error + switch decision.kind { + case generationActionEnterRequiresAction: + cleanup() + return s.enterRequiresAction(ctx, machine, input) + case generationActionFinishTurn: + cleanup() + return s.finishGenerationTurn(ctx, machine, input, 0, decision, generationAttemptNotRequired) + case generationActionGenerateAssistant: + actionErr = s.generateAssistant(ctx, machine, input, prepared, decision) + case generationActionExecuteLocalTools: + actionErr = s.executeLocalTools(ctx, machine, input, prepared, decision) + case generationActionCompact: + actionErr = s.generateCompaction(ctx, machine, input, prepared) + case generationActionPersistWorkspaceContext: + actionErr = s.persistWorkspaceContext(ctx, machine, input, prepared.Chat) + default: + return s.finishGenerationError(ctx, machine, input, 0, xerrors.Errorf("unknown generation action %q", decision.kind), generationAttemptNotRequired) + } + cleanup() + if actionErr == nil { + return nil + } + if errors.Is(actionErr, errTaskExpectedExit) || errors.Is(actionErr, chatloop.ErrInterrupted) { + return nil + } + if errors.Is(actionErr, context.Canceled) && ctx.Err() != nil { + return nil + } + classified := chaterror.Classify(actionErr) + if classified.Retryable { + decision, err := s.recordGenerationRetry(ctx, machine, input, classified) + if err != nil { + return err + } + if decision.retry { + if classified.ChainBroken { + chainModeDisabled = true + } + if err := s.waitGenerationRetry(ctx, decision.delay); err != nil { + return err + } + continue + } + return s.finishGenerationError(ctx, machine, input, decision.generationAttempt, actionErr, generationAttemptRequired) + } + return s.finishGenerationError(ctx, machine, input, 0, actionErr, generationAttemptNotRequired) + } +} + +func loadGenerationState( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, +) (database.Chat, []database.ChatMessage, error) { + var locked database.Chat + var messages []database.ChatMessage + err := machine.ReadLock(ctx, func(store database.Store) error { + chat, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load locked chat: %w", err) + } + if err := verifyTaskFence(chat, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + loaded, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: input.ChatID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("load chat messages: %w", err) + } + locked = chat + messages = loaded + return nil + }) + if err != nil { + return database.Chat{}, nil, normalizeTaskInfrastructureError(err, "lock chat for generation") + } + return locked, messages, nil +} + +func (*taskStarter) recordGenerationRetry( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + classified chaterror.ClassifiedError, +) (generationRetryDecision, error) { + var decision generationRetryDecision + var payload *codersdk.ChatStreamRetry + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + decision.generationAttempt = locked.GenerationAttempt + if locked.GenerationAttempt <= 0 || locked.GenerationAttempt >= int64(chatretry.MaxAttempts) { + decision.retry = false + return errRetryStateDecisionOnly + } + + attempt := int(locked.GenerationAttempt) + delay := chatretry.Delay(attempt - 1) + if classified.RetryAfter > delay { + delay = classified.RetryAfter + } + decision.retry = true + decision.delay = delay + + payload = chaterror.StreamRetryPayload(attempt, delay, classified) + if payload == nil { + return errRetryStateDecisionOnly + } + encoded, err := json.Marshal(payload) + if err != nil { + return xerrors.Errorf("marshal retry state: %w", err) + } + _, err = tx.RecordRetryState(chatstate.RecordRetryStateInput{ + RetryState: pqtype.NullRawMessage{RawMessage: encoded, Valid: true}, + }) + return err + }) + if errors.Is(err, errRetryStateDecisionOnly) { + return decision, nil + } + if err != nil { + return generationRetryDecision{}, normalizeTaskTransitionError(err, "record retry state") + } + return decision, nil +} + +func (s *taskStarter) waitGenerationRetry(ctx context.Context, delay time.Duration) error { + timer := s.opts.Clock.NewTimer(delay, "chatworker", "generation-retry") + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return errTaskExpectedExit + } +} + +const ( + // generationPhaseMaxAttempts bounds how many times prepareGeneration + // and decideGenerationAction run before the turn finishes with an + // error. Both phases are retried because prepareGeneration performs + // I/O (DB reads, MCP connects, workspace dials) that can fail + // transiently. + generationPhaseMaxAttempts = 3 + // generationPhaseBaseBackoff is the delay before the first retry. It + // doubles on each subsequent attempt. + generationPhaseBaseBackoff = 200 * time.Millisecond +) + +func generationPhaseBackoff(attempt int) time.Duration { + d := generationPhaseBaseBackoff + for range attempt { + d *= 2 + } + return d +} + +// retryGenerationPhase runs fn up to generationPhaseMaxAttempts times. It +// returns early on success or on a terminal error (see terminalGeneration). +// Non-terminal errors are retried with exponential backoff. Context +// cancellation returns errTaskExpectedExit so shutdown does not write an +// error state. When every attempt fails, the last error is returned. +func retryGenerationPhase[T any]( + ctx context.Context, + wait func(context.Context, time.Duration) error, + fn func() (T, error), +) (T, error) { + var zero T + var lastErr error + for attempt := 0; attempt < generationPhaseMaxAttempts; attempt++ { + result, err := fn() + if err == nil { + return result, nil + } + if isTerminalGeneration(err) { + return zero, err + } + if ctx.Err() != nil { + return zero, errTaskExpectedExit + } + lastErr = err + if attempt < generationPhaseMaxAttempts-1 { + if waitErr := wait(ctx, generationPhaseBackoff(attempt)); waitErr != nil { + return zero, waitErr + } + } + } + return zero, lastErr +} + +func (s *taskStarter) waitGenerationPhaseBackoff(ctx context.Context, delay time.Duration) error { + timer := s.opts.Clock.NewTimer(delay, "chatworker", "generation-phase-retry") + defer timer.Stop() + select { + case <-timer.C: + return nil + case <-ctx.Done(): + return errTaskExpectedExit + } +} + +func (s *taskStarter) generateAssistant( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + prepared generationPrepared, + decision generationDecision, +) error { + attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input) + if err != nil { + return err + } + defer closeEpisode() + runCtx := input.DebugTurn.Ensure(ctx, prepared.Chat, prepared.Debug) + outcome, err := chatloop.GenerateAssistant(runCtx, chatloop.GenerateAssistantOptions{ + Model: prepared.Model, + Messages: prepared.Prompt, + Tools: prepared.Tools, + ActiveTools: prepared.ActiveTools, + ProviderTools: prepared.ProviderTools, + ContextLimitFallback: prepared.ContextLimitFallback, + ModelConfig: prepared.ModelConfig, + ProviderOptions: prepared.ProviderOptions, + PublishMessagePart: publish, + Logger: s.opts.Logger, + Clock: s.opts.Clock, + Metrics: s.server.metrics, + }) + if err != nil { + return err + } + if decision.compactionTrigger == compactionTriggerAlreadyCompacted && + shouldCompactPromptUsage(outcome.Step.Usage, prepared.ContextLimitFallback, prepared.Compaction.Options.ThresholdPercent) { + err := errCompactionStillOverLimit + s.server.metrics.RecordCompaction(compactionProvider(prepared.Compaction.Options), compactionModel(prepared.Compaction.Options), false, err) + return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired) + } + if len(outcome.Step.Content) == 0 { + return s.finishGenerationTurn(ctx, machine, input, attempt, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, generationAttemptRequired) + } + messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: prepared.ModelConfigID, + modelCallConfig: prepared.ModelConfig, + step: stepDataFromPersisted(outcome.Step), + toolNameToConfigID: prepared.ToolNameToConfigID, + logger: s.opts.Logger, + contentVersion: chatprompt.CurrentContentVersion, + }) + if err != nil { + return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired) + } + return s.commitGenerationStep(ctx, machine, input, attempt, generationActionGenerateAssistant, messages) +} + +func (s *taskStarter) executeLocalTools( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + prepared generationPrepared, + decision generationDecision, +) error { + attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input) + if err != nil { + return err + } + defer closeEpisode() + provider := "" + modelName := "" + if prepared.Model != nil { + provider = prepared.Model.Provider() + modelName = prepared.Model.Model() + } + // Local tool callbacks (e.g. spawn_agent, message_agent) read the + // active turn's delegated API key ID from the context to route + // subagent traffic through the AI Gateway. prepareGeneration sets it + // only on its own context, so re-derive it here for tool execution. + toolCtx := withActiveTurnAPIKeyID(ctx, prepared.ModelBuildOptions) + outcome, err := chatloop.ExecuteLocalTools(toolCtx, chatloop.ExecuteLocalToolsOptions{ + Tools: prepared.Tools, + ActiveTools: prepared.ActiveTools, + ProviderTools: prepared.ProviderTools, + ToolCalls: decision.localToolCalls, + ExclusiveToolNames: prepared.ExclusiveToolNames, + BuiltinToolNames: prepared.BuiltinToolNames, + ModelProvider: provider, + ModelName: modelName, + PublishMessagePart: publish, + Logger: s.opts.Logger, + Metrics: s.server.metrics, + Clock: s.opts.Clock, + }) + if err != nil { + return err + } + messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: prepared.ModelConfigID, + modelCallConfig: prepared.ModelConfig, + step: stepDataFromPersisted(outcome.Step), + toolNameToConfigID: prepared.ToolNameToConfigID, + logger: s.opts.Logger, + contentVersion: chatprompt.CurrentContentVersion, + }) + if err != nil { + return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired) + } + return s.commitGenerationStep(ctx, machine, input, attempt, generationActionExecuteLocalTools, messages) +} + +func (s *taskStarter) generateCompaction( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + prepared generationPrepared, +) error { + attempt, _, publish, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input) + if err != nil { + return err + } + defer closeEpisode() + if prepared.Compaction == nil { + return s.finishGenerationError(ctx, machine, input, attempt, xerrors.New("compaction action missing options"), generationAttemptRequired) + } + compactionOpts := prepared.Compaction.Options + compactionOpts.PublishMessagePart = publish + outcome, err := chatloop.GenerateCompaction(ctx, compactionOpts) + if err != nil { + s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err) + return err + } + if strings.TrimSpace(outcome.SystemSummary) == "" || strings.TrimSpace(outcome.SummaryReport) == "" { + err := xerrors.New("compaction produced no summary") + s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err) + return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired) + } + messages, err := buildCompactionMessages(buildCompactionMessagesInput{ + modelConfigID: prepared.ModelConfigID, + activeAPIKeyID: prepared.ModelBuildOptions.ActiveAPIKeyID, + toolCallID: compactionOpts.ToolCallID, + toolName: compactionOpts.ToolName, + compaction: compactionOutcome(outcome), + contentVersion: chatprompt.CurrentContentVersion, + }) + if err != nil { + s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), false, err) + return s.finishGenerationError(ctx, machine, input, attempt, err, generationAttemptRequired) + } + err = s.commitGenerationStep(ctx, machine, input, attempt, generationActionCompact, stepMessagesForCommit{ + Messages: messages.Messages, + VisibleIndexes: visibleMessageIndexes(messages.Messages), + }) + s.server.metrics.RecordCompaction(compactionProvider(compactionOpts), compactionModel(compactionOpts), err == nil, err) + return err +} + +func compactionProvider(opts chatloop.GenerateCompactionOptions) string { + if opts.Model == nil { + return "" + } + return opts.Model.Provider() +} + +func compactionModel(opts chatloop.GenerateCompactionOptions) string { + if opts.Model == nil { + return "" + } + return opts.Model.Model() +} + +// persistWorkspaceContext is the generation action that commits durable +// workspace context messages (e.g. AGENTS.md, workspace skills) into +// chat history. It records a generation attempt, calls the injected +// workspace context builder without holding the DB lock, then commits +// the returned messages fenced to the attempt. If the builder returns +// no messages, the action exits as expected and the next worker task +// re-reads the chat. +func (s *taskStarter) persistWorkspaceContext( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + locked database.Chat, +) error { + if s.server == nil { + return errTaskExpectedExit + } + messages, err := s.opts.Store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: input.ChatID, + AfterID: 0, + }) + if err != nil { + return taskRetryableError{err: xerrors.Errorf("load chat messages for workspace context: %w", err)} + } + attempt, _, _, closeEpisode, err := s.beginGenerationAttempt(ctx, machine, input) + if err != nil { + return err + } + defer closeEpisode() + modelOpts := modelBuildOptionsFromMessages(messages) + result, err := s.server.buildWorkspaceContext(ctx, workspaceContextBuildInput{ + Chat: locked, + Messages: messages, + ActiveAPIKeyID: modelOpts.ActiveAPIKeyID, + }) + if err != nil { + if errors.Is(err, errWorkspaceContextUnavailable) { + // Builder reported nothing durable to commit (workspace or + // agent missing, unreachable, etc.). Exit the action without + // committing so the next worker task can re-read the chat. + return errTaskExpectedExit + } + return err + } + return s.commitGenerationStep(ctx, machine, input, attempt, generationActionPersistWorkspaceContext, stepMessagesForCommit{ + Messages: result.Messages, + VisibleIndexes: visibleMessageIndexes(result.Messages), + }) +} + +func (s *taskStarter) beginGenerationAttempt( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, +) (int64, messagepartbuffer.Key, func(codersdk.ChatMessageRole, codersdk.ChatMessagePart), func(), error) { + var attempt int64 + var committed database.Chat + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + result, err := tx.RecordGenerationAttempt(chatstate.RecordGenerationAttemptInput{}) + if err != nil { + return err + } + attempt = result.GenerationAttempt + committed, err = store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("load committed chat: %w", err) + } + return nil + }) + if err != nil { + return 0, messagepartbuffer.Key{}, nil, nil, normalizeTaskTransitionError(err, "record generation attempt") + } + key := messagepartbuffer.Key{ + ChatID: input.ChatID, + HistoryVersion: committed.HistoryVersion, + GenerationAttempt: attempt, + } + if err := s.opts.MessagePartBuffer.CreateEpisode(key); err != nil && ctx.Err() == nil { + return 0, messagepartbuffer.Key{}, nil, nil, taskRetryableError{err: xerrors.Errorf("create message part episode: %w", err)} + } + publish := func(role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) { + _ = s.opts.MessagePartBuffer.AddPart(key, role, part) + } + closeEpisode := func() { + _ = s.opts.MessagePartBuffer.CloseEpisode(key) + } + return attempt, key, publish, closeEpisode, nil +} + +func (s *taskStarter) commitGenerationStep( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + attempt int64, + kind generationActionKind, + messages stepMessagesForCommit, +) error { + if len(messages.Messages) == 0 { + return s.finishGenerationTurn(ctx, machine, input, attempt, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, generationAttemptRequired) + } + var committed database.Chat + insertedMessages := []runnerActionMessage{} + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyGenerationFence(locked, input, attempt); err != nil { + return err + } + commitResult, err := tx.CommitStep(chatstate.CommitStepInput{Messages: messages.Messages}) + if err != nil { + return err + } + insertedMessages = make([]runnerActionMessage, 0, len(commitResult.InsertedMessages)) + for _, msg := range commitResult.InsertedMessages { + insertedMessages = append(insertedMessages, runnerActionMessage{ID: msg.ID, Role: codersdk.ChatMessageRole(msg.Role)}) + } + committed, err = store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("load committed chat: %w", err) + } + return nil + }) + if err != nil { + return normalizeTaskTransitionError(err, "commit generation step") + } + s.routeStateHint(ctx, stateUpdateFromChat(committed)) + return s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKind(kind), + InsertedMessages: insertedMessages, + }) +} + +func (s *taskStarter) enterRequiresAction( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, +) error { + var committed database.Chat + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + if _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}); err != nil { + return err + } + committed, err = store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("load committed chat: %w", err) + } + return nil + }) + if err != nil { + return normalizeTaskTransitionError(err, "enter requires action") + } + if err := s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindActionRequired); err != nil { + return err + } + return s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKindEnterRequiresAction, + WatchEventKind: codersdk.ChatWatchEventKindActionRequired, + }) +} + +type generationAttemptFence int + +const ( + generationAttemptNotRequired generationAttemptFence = iota + generationAttemptRequired +) + +func (s *taskStarter) finishGenerationTurn( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + attempt int64, + decision generationDecision, + attemptFence generationAttemptFence, +) error { + var committed database.Chat + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if attemptFence == generationAttemptRequired { + if err := verifyGenerationFence(locked, input, attempt); err != nil { + return err + } + } else if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + finishResult, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + if err != nil { + return err + } + if finishResult.PromotedMessage != nil { + decision.promotedMessageID = finishResult.PromotedMessage.ID + } + committed = finishResult.Chat + return nil + }) + if err != nil { + return normalizeTaskTransitionError(err, "finish generation turn") + } + input.DebugTurn.RecordOutcome(chatdebug.StatusCompleted) + watchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postCommitWatchPublishTimeout) + defer cancel() + if err := s.publishWatchWithRetry(watchCtx, committed, codersdk.ChatWatchEventKindStatusChange); err != nil { + return err + } + if err := s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKindFinishTurn, + WatchEventKind: codersdk.ChatWatchEventKindStatusChange, + PromotedMessageID: decision.promotedMessageID, + }); err != nil { + return err + } + s.routeStateHint(ctx, stateUpdateFromChat(committed)) + return nil +} + +func (s *taskStarter) finishGenerationError( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + attempt int64, + cause error, + attemptFence generationAttemptFence, +) error { + lastError, message := generationLastError(cause) + var committed database.Chat + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if attemptFence == generationAttemptRequired { + if err := verifyGenerationFence(locked, input, attempt); err != nil { + return err + } + } else if err := verifyTaskFence(locked, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + if _, err := tx.FinishError(chatstate.FinishErrorInput{LastError: lastError}); err != nil { + return err + } + committed, err = store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("load committed chat: %w", err) + } + return nil + }) + if err != nil { + return normalizeTaskTransitionError(err, "finish generation error") + } + input.DebugTurn.RecordOutcome(chatdebug.StatusError) + if err := s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindStatusChange); err != nil { + return err + } + return s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKindFinishError, + WatchEventKind: codersdk.ChatWatchEventKindStatusChange, + LastError: message, + }) +} + +func generationLastError(err error) (pqtype.NullRawMessage, string) { + if err == nil { + return pqtype.NullRawMessage{}, "" + } + classified := chaterror.Classify(err) + payload := chaterror.TerminalErrorPayload(classified) + if payload == nil { + payload = &codersdk.ChatError{Message: err.Error()} + } + encoded, marshalErr := json.Marshal(payload) + if marshalErr != nil { + return pqtype.NullRawMessage{}, payload.Message + } + return pqtype.NullRawMessage{RawMessage: encoded, Valid: true}, payload.Message +} + +func (s *taskStarter) afterGenerationOutcome(ctx context.Context, outcome generationOutcome) error { + if s.server == nil { + return nil + } + if err := s.server.afterGenerationOutcome(ctx, outcome); err != nil { + return taskRetryableError{err: xerrors.Errorf("generation post-outcome side effects: %w", err)} + } + return nil +} + +func verifyGenerationFence(chat database.Chat, input chatWorkerTaskStartInput, attempt int64) error { + if err := verifyTaskFence(chat, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + if chat.GenerationAttempt != attempt { + return errTaskExpectedExit + } + return nil +} + +func stepDataFromPersisted(step chatloop.PersistedStep) stepData { + return stepData{ + Content: step.Content, + Usage: step.Usage, + ContextLimit: step.ContextLimit, + ProviderResponseID: step.ProviderResponseID, + Runtime: step.Runtime, + ToolCallCreatedAt: step.ToolCallCreatedAt, + ToolResultCreatedAt: step.ToolResultCreatedAt, + ReasoningStartedAt: step.ReasoningStartedAt, + ReasoningCompletedAt: step.ReasoningCompletedAt, + } +} diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go new file mode 100644 index 0000000000..459457f524 --- /dev/null +++ b/coderd/x/chatd/generation_preparer.go @@ -0,0 +1,765 @@ +package chatd + +import ( + "context" + "encoding/json" + "slices" + "strings" + "sync" + + "charm.land/fantasy" + "github.com/google/uuid" + "golang.org/x/sync/errgroup" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatadvisor" + "github.com/coder/coder/v2/coderd/x/chatd/chatloop" + "github.com/coder/coder/v2/coderd/x/chatd/chatopenai" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatsanitize" + "github.com/coder/coder/v2/coderd/x/chatd/chattool" + "github.com/coder/coder/v2/coderd/x/chatd/mcpclient" + skillspkg "github.com/coder/coder/v2/coderd/x/skills" + "github.com/coder/coder/v2/codersdk" +) + +func (server *Server) prepareGeneration( + ctx context.Context, + input generationPrepareInput, +) (generationPrepared, error) { + chat := input.Chat + logger := server.logger.With( + slog.F("chat_id", chat.ID), + slog.F("owner_id", chat.OwnerID), + ) + + var ( + model fantasy.LanguageModel + modelConfig database.ChatModelConfig + providerKeys chatprovider.ProviderAPIKeys + modelRoute resolvedModelRoute + modelOpts modelBuildOptions + callConfig codersdk.ChatModelCallConfig + promptRows []database.ChatMessage + mcpConfigs []database.MCPServerConfig + mcpTokens []database.MCPServerUserToken + debugEnabled bool + debugProvider string + debugModel string + ) + + var g errgroup.Group + g.Go(func() error { + var err error + promptRows, err = server.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("get chat messages for prompt: %w", err) + } + return nil + }) + if len(chat.MCPServerIDs) > 0 { + g.Go(func() error { + var err error + mcpConfigs, err = server.db.GetMCPServerConfigsByIDs(ctx, chat.MCPServerIDs) + if err != nil { + logger.Warn(ctx, "failed to load MCP server configs", slog.Error(err)) + } + return nil + }) + g.Go(func() error { + var err error + mcpTokens, err = server.db.GetMCPServerUserTokensByUserID(ctx, chat.OwnerID) + if err != nil { + logger.Warn(ctx, "failed to load MCP user tokens", slog.Error(err)) + } + return nil + }) + } + if err := g.Wait(); err != nil { + return generationPrepared{}, err + } + + modelOpts = modelBuildOptionsFromMessages(promptRows) + ctx = withActiveTurnAPIKeyID(ctx, modelOpts) + + var err error + model, modelConfig, providerKeys, modelRoute, debugEnabled, debugProvider, debugModel, err = server.resolveChatModel(ctx, chat, modelOpts) + if err != nil { + return generationPrepared{}, err + } + if len(modelConfig.Options) > 0 { + if err := json.Unmarshal(modelConfig.Options, &callConfig); err != nil { + return generationPrepared{}, xerrors.Errorf("parse model call config: %w", err) + } + } + + if callConfig.MaxOutputTokens == nil { + maxOutputTokens := int64(32_000) + callConfig.MaxOutputTokens = &maxOutputTokens + } + + currentPlanMode := chat.PlanMode + isPlanModeTurn := currentPlanMode.Valid && currentPlanMode.ChatPlanMode == database.ChatPlanModePlan + isExploreSubagent := isExploreSubagentMode(chat.Mode) + isRootChat := !chat.ParentChatID.Valid + + mcpConnectConfigs, approvedPlanMCPConfigIDs := filterExternalMCPConfigsForTurn( + mcpConfigs, + currentPlanMode, + chat.ParentChatID, + ) + if isExploreSubagent && isRootChat { + mcpConnectConfigs = nil + approvedPlanMCPConfigIDs = map[uuid.UUID]struct{}{} + } + + planModeInstructions := server.loadPlanModeInstructions(ctx, currentPlanMode, logger) + advisorCfg := server.loadAdvisorConfig(ctx, logger) + + var advisorRuntime *chatadvisor.Runtime + if advisorCfg.Enabled && isRootChat && !isPlanModeTurn && !isExploreSubagent { + var advisorErr error + advisorRuntime, advisorErr = server.newAdvisorRuntime( + ctx, + chat, + advisorCfg, + model, + callConfig, + providerKeys, + modelOpts, + logger, + ) + if advisorErr != nil { + return generationPrepared{}, advisorErr + } + } + + var advisorPromptSnapshot []fantasy.Message + setAdvisorPromptSnapshot := func(msgs []fantasy.Message) { + if advisorRuntime == nil { + return + } + advisorPromptSnapshot = slices.Clone(msgs) + } + + currentChat := chat + loadChatSnapshot := func(loadCtx context.Context, chatID uuid.UUID) (database.Chat, error) { + return server.db.GetChatByID(loadCtx, chatID) + } + var chatStateMu sync.Mutex + var workspaceMu sync.Mutex + workspaceCtx := turnWorkspaceContext{ + server: server, + chatStateMu: &chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: loadChatSnapshot, + } + cleanup := func() { + workspaceCtx.close() + } + + planPathFn := func(ctx context.Context) (string, string, error) { + conn, err := workspaceCtx.getWorkspaceConn(ctx) + if err != nil { + return "", "", err + } + home, err := chattool.ResolveWorkspaceHome(ctx, conn) + if err != nil { + return "", "", err + } + return chattool.PlanPathForChat(home, chat.ID), home, nil + } + resolvePlanPathForTools := func(ctx context.Context) (string, string, error) { + planCtx, cancel := context.WithTimeout(ctx, planPathLookupTimeout) + defer cancel() + return planPathFn(planCtx) + } + resolvePlanPathBlock := func(resolveCtx context.Context) string { + if chat.ParentChatID.Valid { + return "" + } + + planCtx, cancel := context.WithTimeout(resolveCtx, planPathLookupTimeout) + defer cancel() + + if _, _, err := workspaceCtx.workspaceAgentIDForConn(planCtx); err != nil { + logger.Debug(resolveCtx, "plan path instruction: agent not reachable", + slog.Error(err), + slog.F("chat_id", chat.ID), + ) + return "" + } + + planPath, home, err := planPathFn(planCtx) + if err != nil { + logger.Debug(resolveCtx, "plan path instruction: failed to resolve plan path", + slog.Error(err), + slog.F("chat_id", chat.ID), + ) + return "" + } + return formatPlanPathBlock(planPath, home) + } + + var ( + prompt []fantasy.Message + instruction string + mcpTools []fantasy.AgentTool + mcpCleanup func() + workspaceMCPTools []fantasy.AgentTool + workspaceSkills []chattool.SkillMeta + personalSkills []skillspkg.Skill + resolvedUserPrompt string + ) + + persistedSkills := skillsFromParts(promptRows) + hasContextFiles := false + if chat.WorkspaceID.Valid { + // Resolve the workspace agent so the chat row's AgentID and + // BuildID bindings are up to date before the chatworker + // decision helper inspects them. ensureWorkspaceAgent does a + // DB lookup and lazily calls persistBuildAgentBinding when + // the bound agent has changed, so this is a cheap metadata + // refresh, not a workspace dial. It must not insert chat + // history; only metadata is mutated here. + _, _ = workspaceCtx.getWorkspaceAgent(ctx) + _, found := contextFileAgentID(promptRows) + hasContextFiles = found + } + + var g2 errgroup.Group + g2.Go(func() error { + var err error + prompt, err = chatprompt.ConvertMessagesWithFiles(ctx, promptRows, server.chatFileResolver(modelConfig.Provider), logger) + if err != nil { + return xerrors.Errorf("build chat prompt: %w", err) + } + return nil + }) + if hasContextFiles { + instruction = instructionFromContextFiles(promptRows) + workspaceSkills = persistedSkills + } + g2.Go(func() error { + personalSkills = server.fetchPersonalSkillMetadata(ctx, chat.OwnerID, logger) + return nil + }) + g2.Go(func() error { + resolvedUserPrompt = server.resolveUserPrompt(ctx, chat.OwnerID) + return nil + }) + if len(mcpConnectConfigs) > 0 { + g2.Go(func() error { + mcpTokens = server.refreshExpiredMCPTokens(ctx, logger, mcpConnectConfigs, mcpTokens) + mcpTools, mcpCleanup = mcpclient.ConnectAll( + ctx, + logger, + mcpConnectConfigs, + mcpTokens, + chat.OwnerID, + server.oidcTokenSource, + chatprovider.CoderHeaders(chat), + ) + return nil + }) + } + if chat.WorkspaceID.Valid && !isPlanModeTurn && !isExploreSubagent { + g2.Go(func() error { + workspaceMCPTools = server.discoverWorkspaceMCPTools(ctx, logger, chat.ID, &workspaceCtx) + return nil + }) + } + if err := g2.Wait(); err != nil { + cleanup() + return generationPrepared{}, err + } + + if mcpCleanup != nil { + previousCleanup := cleanup + cleanup = func() { + mcpCleanup() + previousCleanup() + } + } + + prompt, sanitizeStats := chatsanitize.SanitizeAnthropicProviderToolHistory(model.Provider(), prompt) + chatsanitize.LogAnthropicProviderToolSanitization( + ctx, + logger, + "persisted_history_replay", + model.Provider(), + model.Model(), + sanitizeStats, + ) + + subagentInstruction := "" + if !isRootChat { + subagentInstruction = defaultSubagentInstruction + } + resolvedSkillsFor := func(workspaceSkills []chattool.SkillMeta) []skillspkg.ResolvedSkill { + return mergeTurnSkills(personalSkills, workspaceSkills) + } + resolveSkillAlias := func(alias string) (skillspkg.ResolvedSkill, error) { + return skillspkg.Lookup(resolvedSkillsFor(workspaceSkills), alias) + } + initialResolvedSkills := resolvedSkillsFor(workspaceSkills) + + prompt = buildSystemPrompt( + prompt, + subagentInstruction, + instruction, + initialResolvedSkills, + resolvedUserPrompt, + systemPromptBehaviorContext{ + planMode: currentPlanMode, + chatMode: chat.Mode, + planModeInstructions: planModeInstructions, + isRootChat: isRootChat, + }, + ) + if advisorRuntime != nil { + prompt = chatprompt.InsertSystem(prompt, chatadvisor.ParentGuidanceBlock) + } + prompt = renderPlanPathPrompt(prompt, resolvePlanPathBlock(ctx)) + setAdvisorPromptSnapshot(prompt) + + storeChatAttachment := server.newStoreChatAttachmentFunc(&workspaceCtx) + tools := []fantasy.AgentTool{ + chattool.ReadFile(chattool.ReadFileOptions{GetWorkspaceConn: workspaceCtx.getWorkspaceConn}), + chattool.WriteFile(chattool.WriteFileOptions{ + GetWorkspaceConn: workspaceCtx.getWorkspaceConn, + ResolvePlanPath: resolvePlanPathForTools, + IsPlanTurn: isPlanModeTurn, + }), + chattool.EditFiles(chattool.EditFilesOptions{ + GetWorkspaceConn: workspaceCtx.getWorkspaceConn, + ResolvePlanPath: resolvePlanPathForTools, + IsPlanTurn: isPlanModeTurn, + }), + chattool.AttachFile(chattool.AttachFileOptions{ + GetWorkspaceConn: workspaceCtx.getWorkspaceConn, + StoreFile: storeChatAttachment, + }), + chattool.Execute(chattool.ExecuteOptions{GetWorkspaceConn: workspaceCtx.getWorkspaceConn}), + chattool.ProcessOutput(chattool.ProcessToolOptions{GetWorkspaceConn: workspaceCtx.getWorkspaceConn}), + chattool.ProcessList(chattool.ProcessToolOptions{GetWorkspaceConn: workspaceCtx.getWorkspaceConn}), + chattool.ProcessSignal(chattool.ProcessToolOptions{GetWorkspaceConn: workspaceCtx.getWorkspaceConn}), + } + if isPlanModeTurn && isRootChat { + tools = append(tools, chattool.NewAskUserQuestionTool()) + } + if isRootChat { + tools = server.appendRootChatTools(ctx, tools, rootChatToolsOptions{ + chat: chat, + modelConfigID: modelConfig.ID, + workspaceCtx: &workspaceCtx, + workspaceMu: &workspaceMu, + resolvePlanPath: resolvePlanPathForTools, + storeFile: storeChatAttachment, + isPlanModeTurn: isPlanModeTurn, + primerCtx: ctx, + }) + } + + skillOpts := chattool.ReadSkillOptions{ + GetWorkspaceConn: workspaceCtx.getWorkspaceConn, + GetSkills: func() []chattool.SkillMeta { + return workspaceSkills + }, + ResolveAlias: resolveSkillAlias, + LoadPersonalSkillBody: func(ctx context.Context, name string) (skillspkg.ParsedSkill, error) { + return server.loadPersonalSkillBody(ctx, chat.OwnerID, name) + }, + } + appendCurrentSkillTools := func(current []fantasy.AgentTool) ([]fantasy.AgentTool, bool) { + if len(personalSkills) == 0 && len(workspaceSkills) == 0 { + return current, false + } + updated := current + changed := false + appendTool := func(tool fantasy.AgentTool) { + name := tool.Info().Name + if slices.ContainsFunc(current, func(existing fantasy.AgentTool) bool { + return existing.Info().Name == name + }) { + return + } + if !changed { + updated = slices.Clone(current) + changed = true + } + updated = append(updated, tool) + } + appendTool(chattool.ReadSkill(skillOpts)) + if len(workspaceSkills) > 0 { + appendTool(chattool.ReadSkillFile(skillOpts)) + } + return updated, changed + } + tools, _ = appendCurrentSkillTools(tools) + if advisorRuntime != nil { + tools = append(tools, chatadvisor.Tool(chatadvisor.ToolOptions{ + Runtime: advisorRuntime, + GetConversationSnapshot: func() []fantasy.Message { + return stripAdvisorGuidanceBlock(slices.Clone(advisorPromptSnapshot)) + }, + })) + } + + var exclusiveToolNames map[string]bool + if advisorRuntime != nil { + exclusiveToolNames = map[string]bool{chatadvisor.ToolName: true} + } + + builtinToolNames := make(map[string]bool, len(tools)) + for _, t := range tools { + builtinToolNames[t.Info().Name] = true + } + + tools = append(tools, mcpTools...) + if !isExploreSubagent { + tools = append(tools, workspaceMCPTools...) + } + tools = filterToolsForTurn(tools, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) + + tools, dynamicToolNames, err := appendDynamicTools(ctx, logger, tools, chat.DynamicTools, currentPlanMode, chat.Mode) + if err != nil { + cleanup() + return generationPrepared{}, err + } + + var providerTools []chatloop.ProviderTool + if !isPlanModeTurn && callConfig.ProviderOptions != nil { + providerTools = buildProviderTools(callConfig.ProviderOptions) + if isExploreSubagent { + if !chat.ParentChatID.Valid { + providerTools = nil + } else { + providerTools = slices.DeleteFunc(providerTools, func(tool chatloop.ProviderTool) bool { + return tool.Definition.GetName() != "web_search" + }) + } + } + } + + isComputerUse := chat.Mode.Valid && chat.Mode.ChatMode == database.ChatModeComputerUse + if isComputerUse { + computerUseProvider, computerUseModelProvider, computerUseModelName, err := server.computerUseProviderAndModelFromConfig(ctx) + if err != nil { + cleanup() + return generationPrepared{}, xerrors.Errorf("resolve computer use provider and model: %w", err) + } + computerUseRoute, keyErr := server.resolveModelRouteForProviderType(ctx, chat.OwnerID, computerUseModelProvider) + if keyErr != nil { + cleanup() + return generationPrepared{}, xerrors.Errorf("resolve computer use provider route: %w", keyErr) + } + modelRoute = computerUseRoute + providerKeys = computerUseRoute.directProviderKeys() + cuModel, cuDebugEnabled, resolvedProvider, resolvedModel, cuErr := server.resolveComputerUseModel( + ctx, + chat, + computerUseRoute, + computerUseProvider, + computerUseModelProvider, + computerUseModelName, + modelOpts, + ) + if cuErr != nil { + cleanup() + return generationPrepared{}, cuErr + } + model = cuModel + debugEnabled = cuDebugEnabled + debugProvider = resolvedProvider + debugModel = resolvedModel + providerTools, err = appendComputerUseProviderTool(providerTools, computerUseProviderToolOptions{ + provider: computerUseProvider, + isPlanModeTurn: isPlanModeTurn, + isComputerUse: isComputerUse, + getWorkspaceConn: workspaceCtx.getWorkspaceConn, + storeFile: storeChatAttachment, + clock: server.clock, + logger: server.logger.Named("computer_use"), + }) + if err != nil { + cleanup() + return generationPrepared{}, xerrors.Errorf("register computer use provider tool for provider %q: %w", computerUseProvider, err) + } + } else { + providerTools, err = appendComputerUseProviderTool(providerTools, computerUseProviderToolOptions{ + isPlanModeTurn: isPlanModeTurn, + isComputerUse: false, + }) + if err != nil { + cleanup() + return generationPrepared{}, err + } + } + + providerOptions := chatprovider.ProviderOptionsFromChatModelConfig(model, callConfig.ProviderOptions) + chainInfo := chatopenai.ResolveChainMode(promptRows) + if !input.ChainModeDisabled && chatopenai.ShouldActivateChainMode( + providerOptions, + chainInfo, + modelConfig.ID, + isPlanModeTurn, + ) { + providerOptions = chatopenai.WithPreviousResponseID(providerOptions, chainInfo.PreviousResponseID()) + prompt = chatopenai.FilterPromptForChainMode(prompt, chainInfo) + } + + activeToolNames := activeToolNamesForTurn(tools, currentPlanMode, chat.ParentChatID, approvedPlanMCPConfigIDs) + if isExploreSubagent { + activeToolNames = allowedExploreToolNames(tools) + } + + toolNameToConfigID := make(map[string]uuid.UUID) + for _, t := range tools { + if mcpTool, ok := t.(mcpclient.MCPToolIdentifier); ok { + toolNameToConfigID[t.Info().Name] = mcpTool.MCPServerConfigID() + } + } + + triggerMessageID, historyTipMessageID, triggerLabel := deriveChatDebugSeed(promptRows) + debugSvc := server.existingDebugService() + var debug *generationDebug + if debugEnabled { + if debugSvc == nil { + cleanup() + return generationPrepared{}, xerrors.New("chat debug service missing after enablement check") + } + debug = &generationDebug{ + Enabled: true, + Service: debugSvc, + Provider: debugProvider, + Model: debugModel, + TriggerMessageID: triggerMessageID, + HistoryTipMessageID: historyTipMessageID, + TriggerLabel: triggerLabel, + ModelConfig: modelConfig, + } + } + + compactionToolCallID := "chat_summarized_" + uuid.NewString() + effectiveThreshold := modelConfig.CompressionThreshold + if override, ok := server.resolveUserCompactionThreshold(ctx, chat.OwnerID, modelConfig.ID); ok { + effectiveThreshold = override + } + compactionOptions := chatloop.GenerateCompactionOptions{ + Model: model, + Messages: prompt, + ThresholdPercent: effectiveThreshold, + ContextLimit: modelConfig.ContextLimit, + ContextLimitFallback: modelConfig.ContextLimit, + ToolCallID: compactionToolCallID, + ToolName: "chat_summarized", + DebugSvc: debugSvc, + ChatID: chat.ID, + HistoryTipMessageID: historyTipMessageID, + } + compactionOptions.StepUsage = latestPromptUsage(promptRows) + compactionNeeded := shouldCompactPromptUsage(compactionOptions.StepUsage, modelConfig.ContextLimit, effectiveThreshold) + + workspaceContextEligible := chat.WorkspaceID.Valid && isRootChat && !isPlanModeTurn && !isExploreSubagent + + // workspaceCtx.currentChatSnapshot may carry a freshly persisted + // AgentID/BuildID binding from the getWorkspaceAgent call above. + // Return that snapshot so the chatworker decision helper sees + // the up-to-date metadata when deciding whether to run + // persist_workspace_context. + refreshedChat := workspaceCtx.currentChatSnapshot() + if refreshedChat.ID == uuid.Nil { + refreshedChat = chat + } + + return generationPrepared{ + Chat: refreshedChat, + Messages: input.Messages, + Model: model, + Prompt: prompt, + Tools: tools, + ActiveTools: activeToolNames, + ProviderTools: providerTools, + ProviderKeys: providerKeys, + ModelRoute: modelRoute, + ModelBuildOptions: modelOpts, + ModelConfigID: modelConfig.ID, + ModelConfig: callConfig, + ProviderOptions: providerOptions, + ContextLimitFallback: modelConfig.ContextLimit, + DynamicToolNames: dynamicToolNames, + StopAfterTools: stopAfterBehaviorTools(currentPlanMode, chat.Mode, chat.ParentChatID), + ExclusiveToolNames: exclusiveToolNames, + BuiltinToolNames: builtinToolNames, + ToolNameToConfigID: toolNameToConfigID, + MaxSteps: maxChatSteps, + Compaction: &generationCompaction{ + Required: compactionNeeded, + Options: compactionOptions, + }, + Cleanup: cleanup, + Debug: debug, + WorkspaceContextEligible: workspaceContextEligible, + }, nil +} + +func latestPromptUsage(messages []database.ChatMessage) fantasy.Usage { + for i := len(messages) - 1; i >= 0; i-- { + usage := fantasy.Usage{ + InputTokens: messages[i].InputTokens.Int64, + OutputTokens: messages[i].OutputTokens.Int64, + TotalTokens: messages[i].TotalTokens.Int64, + ReasoningTokens: messages[i].ReasoningTokens.Int64, + CacheCreationTokens: messages[i].CacheCreationTokens.Int64, + CacheReadTokens: messages[i].CacheReadTokens.Int64, + } + if usage != (fantasy.Usage{}) { + return usage + } + } + return fantasy.Usage{} +} + +func shouldCompactPromptUsage(usage fantasy.Usage, contextLimit int64, thresholdPercent int32) bool { + if thresholdPercent >= 100 || contextLimit <= 0 { + return false + } + contextTokens := contextTokensFromUsage(usage) + if contextTokens <= 0 { + return false + } + usagePercent := (float64(contextTokens) / float64(contextLimit)) * 100 + return usagePercent >= float64(thresholdPercent) +} + +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 +} + +func (server *Server) afterInterruptionOutcome( + ctx context.Context, + outcome interruptionOutcome, +) error { + chat := outcome.Chat + logger := server.logger.With(slog.F("chat_id", chat.ID), slog.F("owner_id", chat.OwnerID)) + + if outcome.Kind == runnerActionKindFinishInterruption { + server.maybeClearLastTurnSummaryAsync(context.WithoutCancel(ctx), chat, logger) + } + return nil +} + +func (server *Server) afterGenerationOutcome( + ctx context.Context, + outcome generationOutcome, +) error { + chat := outcome.Chat + logger := server.logger.With(slog.F("chat_id", chat.ID), slog.F("owner_id", chat.OwnerID)) + + switch outcome.Kind { + case runnerActionKindFinishTurn: + finalizeCtx := context.WithoutCancel(ctx) + runResult := server.deriveFinalTurnRunResult(finalizeCtx, chat, logger) + statusLabel := server.generateFinalTurnStatusLabel(finalizeCtx, chat, chat.Status, runResult, logger) + server.updateLastTurnSummary(finalizeCtx, chat, chat.HistoryVersion, statusLabel, logger) + server.dispatchSuccessfulTurnPush(finalizeCtx, chat, statusLabel, logger) + case runnerActionKindFinishError: + server.maybeFinalizeTurnStatusLabelAndPush(context.WithoutCancel(ctx), chat, chat.Status, outcome.LastError, runChatResult{}, logger) + case runnerActionKindEnterRequiresAction: + server.maybeFinalizeTurnStatusLabelAndPush(context.WithoutCancel(ctx), chat, chat.Status, "", runChatResult{}, logger) + } + return nil +} + +// deriveFinalTurnRunResult rebuilds the inputs needed to generate the +// end-of-turn status label directly from persisted state. +func (server *Server) deriveFinalTurnRunResult( + ctx context.Context, + chat database.Chat, + logger slog.Logger, +) runChatResult { + // generateFinalTurnStatusLabel only produces a model-generated label for + // the Waiting status, so skip the model resolution and history read + // otherwise. + if chat.Status != database.ChatStatusWaiting { + return runChatResult{} + } + + promptRows, err := server.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + if err != nil { + logger.Warn(ctx, "derive final turn status label: load prompt rows", slog.Error(err)) + return runChatResult{} + } + triggerMessageID, historyTipMessageID, _ := deriveChatDebugSeed(promptRows) + finalAssistantText := latestAssistantText(promptRows) + if finalAssistantText == "" { + return runChatResult{} + } + + // resolvedProvider/resolvedModel describe the model the fallback handle was + // built from; they only feed the status-label fallback candidate's labels. + modelOpts := modelBuildOptionsFromMessages(promptRows) + ctx = withActiveTurnAPIKeyID(ctx, modelOpts) + model, _, providerKeys, modelRoute, _, resolvedProvider, resolvedModel, err := server.resolveChatModel(ctx, chat, modelOpts) + if err != nil { + // Return what we have; generateFinalTurnStatusLabel falls back to a + // generic label when StatusLabelModel is nil. + logger.Warn(ctx, "derive final turn status label: resolve model", slog.Error(err)) + return runChatResult{ + FinalAssistantText: finalAssistantText, + TriggerMessageID: triggerMessageID, + HistoryTipMessageID: historyTipMessageID, + } + } + + return runChatResult{ + FinalAssistantText: finalAssistantText, + StatusLabelModel: model, + ProviderKeys: providerKeys, + FallbackProvider: resolvedProvider, + FallbackRoute: modelRoute, + FallbackModel: resolvedModel, + ModelBuildOptions: modelOpts, + TriggerMessageID: triggerMessageID, + HistoryTipMessageID: historyTipMessageID, + } +} + +// latestAssistantText returns the trimmed text of the most recent assistant +// message. It mirrors the FinalAssistantText that buildCommitStepMessages +// produced from the freshly generated step, making persisted history the +// single source of truth for the turn status label input. +func latestAssistantText(messages []database.ChatMessage) string { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Role != database.ChatMessageRoleAssistant { + continue + } + parts, err := chatprompt.ParseContent(messages[i]) + if err != nil { + return "" + } + return strings.TrimSpace(textFromParts(parts)) + } + return "" +} diff --git a/coderd/x/chatd/generation_preparer_internal_test.go b/coderd/x/chatd/generation_preparer_internal_test.go new file mode 100644 index 0000000000..c3c5ed0b7f --- /dev/null +++ b/coderd/x/chatd/generation_preparer_internal_test.go @@ -0,0 +1,277 @@ +package chatd //nolint:testpackage // Exercises unexported re-derivation helpers. + +import ( + "database/sql" + "encoding/json" + "testing" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" +) + +func mustMarshalText(t *testing.T, parts ...string) pqtype.NullRawMessage { + t.Helper() + messageParts := make([]codersdk.ChatMessagePart, 0, len(parts)) + for _, p := range parts { + messageParts = append(messageParts, codersdk.ChatMessageText(p)) + } + content, err := chatprompt.MarshalParts(messageParts) + require.NoError(t, err) + return content +} + +func textMessage(t *testing.T, id int64, role database.ChatMessageRole, parts ...string) database.ChatMessage { + t.Helper() + return database.ChatMessage{ + ID: id, + Role: role, + Content: mustMarshalText(t, parts...), + ContentVersion: chatprompt.CurrentContentVersion, + } +} + +func TestLatestAssistantText(t *testing.T) { + t.Parallel() + + t.Run("ReturnsMostRecentAssistantMessage", func(t *testing.T) { + t.Parallel() + messages := []database.ChatMessage{ + textMessage(t, 1, database.ChatMessageRoleUser, "hi"), + textMessage(t, 2, database.ChatMessageRoleAssistant, "first answer"), + textMessage(t, 3, database.ChatMessageRoleTool, "tool result"), + textMessage(t, 4, database.ChatMessageRoleAssistant, " final answer "), + } + require.Equal(t, "final answer", latestAssistantText(messages)) + }) + + t.Run("ConcatenatesTextParts", func(t *testing.T) { + t.Parallel() + messages := []database.ChatMessage{ + textMessage(t, 1, database.ChatMessageRoleAssistant, "foo", "bar"), + } + require.Equal(t, "foobar", latestAssistantText(messages)) + }) + + t.Run("NoAssistantMessage", func(t *testing.T) { + t.Parallel() + messages := []database.ChatMessage{ + textMessage(t, 1, database.ChatMessageRoleUser, "hi"), + textMessage(t, 2, database.ChatMessageRoleTool, "tool result"), + } + require.Empty(t, latestAssistantText(messages)) + }) + + t.Run("EmptyAssistantText", func(t *testing.T) { + t.Parallel() + messages := []database.ChatMessage{ + textMessage(t, 1, database.ChatMessageRoleAssistant, " "), + } + require.Empty(t, latestAssistantText(messages)) + }) + + t.Run("EmptyHistory", func(t *testing.T) { + t.Parallel() + require.Empty(t, latestAssistantText(nil)) + }) +} + +// TestDeriveFinalTurnRunResult exercises the re-derivation path that replaces +// the old in-memory generationSideEffects stash. The server here never ran +// prepareGeneration, so a passing test proves the finish-turn inputs are +// rebuilt purely from persisted state. +func TestDeriveFinalTurnRunResult(t *testing.T) { + t.Parallel() + + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + + setup := func(t *testing.T) (*Server, database.Chat) { + t.Helper() + db, ps := dbtestutil.NewDB(t) + ctx := chatdTestContext(t) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + APIKey: "test-key", + Enabled: true, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + }) + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "gpt-4o-mini", + DisplayName: "gpt-4o-mini", + Options: json.RawMessage(`{}`), + }, func(p *database.InsertChatModelConfigParams) { + p.Enabled = true + p.IsDefault = true + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: modelCfg.ID, + Title: "derive-chat", + ClientType: database.ChatClientTypeUi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: mustMarshalText(t, "what is the answer?"), + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + return server, created.Chat + } + + commitAssistant := func(t *testing.T, server *Server, chat database.Chat, text string) { + t.Helper() + ctx := chatdTestContext(t) + machine := chatstate.NewChatMachine(server.db, server.pubsub, chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: mustMarshalText(t, text), + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: chat.LastModelConfigID, Valid: true}, + }, + }, + }) + return err + })) + } + + t.Run("WaitingDerivesFromHistory", func(t *testing.T) { + t.Parallel() + server, chat := setup(t) + ctx := chatdTestContext(t) + commitAssistant(t, server, chat, "the answer is 42") + + rows, err := server.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + require.NotEmpty(t, rows) + var lastUserID int64 + for _, row := range rows { + if row.Role == database.ChatMessageRoleUser { + lastUserID = row.ID + } + } + tipID := rows[len(rows)-1].ID + + chat.Status = database.ChatStatusWaiting + result := server.deriveFinalTurnRunResult(ctx, chat, logger) + + require.Equal(t, "the answer is 42", result.FinalAssistantText) + require.Equal(t, lastUserID, result.TriggerMessageID) + require.Equal(t, tipID, result.HistoryTipMessageID) + require.NotNil(t, result.StatusLabelModel) + require.Equal(t, "openai", result.FallbackProvider) + require.Equal(t, "gpt-4o-mini", result.FallbackModel) + require.False(t, result.ProviderKeys.Empty()) + }) + + t.Run("NonWaitingReturnsEmpty", func(t *testing.T) { + t.Parallel() + server, chat := setup(t) + ctx := chatdTestContext(t) + commitAssistant(t, server, chat, "the answer is 42") + + chat.Status = database.ChatStatusError + result := server.deriveFinalTurnRunResult(ctx, chat, logger) + require.Equal(t, runChatResult{}, result) + }) + + t.Run("WaitingWithoutAssistantReturnsEmpty", func(t *testing.T) { + t.Parallel() + server, chat := setup(t) + ctx := chatdTestContext(t) + + // No assistant message was committed, so there is nothing to label. + chat.Status = database.ChatStatusWaiting + result := server.deriveFinalTurnRunResult(ctx, chat, logger) + require.Equal(t, runChatResult{}, result) + }) + + t.Run("ModelResolveErrorKeepsTextAndIDs", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := chatdTestContext(t) + + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + // A disabled AI provider makes resolveChatModel fail, exercising the + // degraded path that still returns the re-derived text and IDs. + provider := insertInternalAIProvider(t, db, database.AiProviderTypeOpenai, "provider-api-key", false) + modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "gpt-4o-mini", + DisplayName: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: modelCfg.ID, + Title: "derive-chat-error", + ClientType: database.ChatClientTypeUi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: mustMarshalText(t, "what is the answer?"), + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + chat := created.Chat + + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + commitAssistant(t, server, chat, "the answer is 42") + + chat.Status = database.ChatStatusWaiting + result := server.deriveFinalTurnRunResult(ctx, chat, logger) + + require.Equal(t, "the answer is 42", result.FinalAssistantText) + require.NotZero(t, result.TriggerMessageID) + require.NotZero(t, result.HistoryTipMessageID) + require.Nil(t, result.StatusLabelModel) + require.Empty(t, result.FallbackProvider) + require.Empty(t, result.FallbackModel) + }) +} diff --git a/coderd/x/chatd/generation_retry_internal_test.go b/coderd/x/chatd/generation_retry_internal_test.go new file mode 100644 index 0000000000..87ef475ca5 --- /dev/null +++ b/coderd/x/chatd/generation_retry_internal_test.go @@ -0,0 +1,148 @@ +package chatd //nolint:testpackage // Exercises unexported generation retry helpers. + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" +) + +func TestTerminalGeneration(t *testing.T) { + t.Parallel() + + require.Nil(t, terminalGeneration(nil)) + + cause := xerrors.New("boom") + wrapped := terminalGeneration(cause) + require.True(t, isTerminalGeneration(wrapped)) + require.ErrorIs(t, wrapped, cause) + require.ErrorIs(t, wrapped, errTerminalGeneration) + require.Equal(t, cause.Error(), wrapped.Error()) + + require.False(t, isTerminalGeneration(cause)) + require.False(t, isTerminalGeneration(nil)) +} + +func TestGenerationPhaseBackoff(t *testing.T) { + t.Parallel() + + require.Equal(t, generationPhaseBaseBackoff, generationPhaseBackoff(0)) + require.Equal(t, 2*generationPhaseBaseBackoff, generationPhaseBackoff(1)) + require.Equal(t, 4*generationPhaseBaseBackoff, generationPhaseBackoff(2)) +} + +func TestRetryGenerationPhase(t *testing.T) { + t.Parallel() + + noopWait := func(context.Context, time.Duration) error { return nil } + + t.Run("SuccessFirstTry", func(t *testing.T) { + t.Parallel() + calls := 0 + waits := 0 + wait := func(context.Context, time.Duration) error { + waits++ + return nil + } + got, err := retryGenerationPhase(context.Background(), wait, func() (int, error) { + calls++ + return 42, nil + }) + require.NoError(t, err) + require.Equal(t, 42, got) + require.Equal(t, 1, calls) + require.Equal(t, 0, waits) + }) + + t.Run("RetryThenSuccess", func(t *testing.T) { + t.Parallel() + calls := 0 + waits := 0 + var delays []time.Duration + wait := func(_ context.Context, d time.Duration) error { + waits++ + delays = append(delays, d) + return nil + } + got, err := retryGenerationPhase(context.Background(), wait, func() (string, error) { + calls++ + if calls < 2 { + return "", xerrors.New("transient") + } + return "ok", nil + }) + require.NoError(t, err) + require.Equal(t, "ok", got) + require.Equal(t, 2, calls) + require.Equal(t, 1, waits) + require.Equal(t, []time.Duration{generationPhaseBackoff(0)}, delays) + }) + + t.Run("ExhaustsAndReturnsLastError", func(t *testing.T) { + t.Parallel() + calls := 0 + waits := 0 + wait := func(context.Context, time.Duration) error { + waits++ + return nil + } + _, err := retryGenerationPhase(context.Background(), wait, func() (int, error) { + calls++ + return 0, xerrors.Errorf("attempt %d", calls) + }) + require.EqualError(t, err, "attempt 3") + require.Equal(t, generationPhaseMaxAttempts, calls) + require.Equal(t, generationPhaseMaxAttempts-1, waits) + }) + + t.Run("TerminalShortCircuits", func(t *testing.T) { + t.Parallel() + calls := 0 + waits := 0 + wait := func(context.Context, time.Duration) error { + waits++ + return nil + } + cause := xerrors.New("deterministic") + _, err := retryGenerationPhase(context.Background(), wait, func() (int, error) { + calls++ + return 0, terminalGeneration(cause) + }) + require.ErrorIs(t, err, cause) + require.True(t, isTerminalGeneration(err)) + require.Equal(t, 1, calls) + require.Equal(t, 0, waits) + }) + + t.Run("ContextCanceledExitsCleanly", func(t *testing.T) { + t.Parallel() + ctx, cancel := context.WithCancel(context.Background()) + cancel() + calls := 0 + _, err := retryGenerationPhase(ctx, noopWait, func() (int, error) { + calls++ + return 0, xerrors.New("transient") + }) + require.ErrorIs(t, err, errTaskExpectedExit) + require.Equal(t, 1, calls) + }) + + t.Run("WaitCancellationExitsCleanly", func(t *testing.T) { + t.Parallel() + calls := 0 + waits := 0 + wait := func(context.Context, time.Duration) error { + waits++ + return errTaskExpectedExit + } + _, err := retryGenerationPhase(context.Background(), wait, func() (int, error) { + calls++ + return 0, xerrors.New("transient") + }) + require.ErrorIs(t, err, errTaskExpectedExit) + require.Equal(t, 1, calls) + require.Equal(t, 1, waits) + }) +} diff --git a/coderd/x/chatd/helpers_test.go b/coderd/x/chatd/helpers_test.go new file mode 100644 index 0000000000..352392f26c --- /dev/null +++ b/coderd/x/chatd/helpers_test.go @@ -0,0 +1,537 @@ +package chatd //nolint:testpackage // Uses unexported chatworker helpers. + +import ( + "context" + "database/sql" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +func testAPIKeyID(t testing.TB, db database.Store, userID uuid.UUID) string { + t.Helper() + key, _ := dbgen.APIKey(t, db, database.APIKey{ID: uuid.NewString(), UserID: userID}) + return key.ID +} + +type workerTestFixture struct { + db database.Store + pubsub dbpubsub.Pubsub + sqlDB *sql.DB + user database.User + org database.Organization + model database.ChatModelConfig + apiKey database.APIKey +} + +type publishedEvent struct { + channel string + payload []byte +} + +type recordingPubsub struct { + inner dbpubsub.Pubsub + mu sync.Mutex + events []publishedEvent +} + +func newRecordingPubsub(inner dbpubsub.Pubsub) *recordingPubsub { + return &recordingPubsub{inner: inner} +} + +func (p *recordingPubsub) Publish(channel string, payload []byte) error { + p.mu.Lock() + p.events = append(p.events, publishedEvent{ + channel: channel, + payload: append([]byte(nil), payload...), + }) + p.mu.Unlock() + return p.inner.Publish(channel, payload) +} + +func (p *recordingPubsub) SubscribeWithErr(channel string, listener dbpubsub.ListenerWithErr) (func(), error) { + return p.inner.SubscribeWithErr(channel, listener) +} + +func (p *recordingPubsub) ownershipMessages(t *testing.T) []coderdpubsub.ChatStateOwnershipMessage { + t.Helper() + p.mu.Lock() + defer p.mu.Unlock() + messages := make([]coderdpubsub.ChatStateOwnershipMessage, 0) + for _, event := range p.events { + if event.channel != coderdpubsub.ChatStateOwnershipChannel { + continue + } + var msg coderdpubsub.ChatStateOwnershipMessage + require.NoError(t, json.Unmarshal(event.payload, &msg)) + messages = append(messages, msg) + } + return messages +} + +func (p *recordingPubsub) watchEvents(t *testing.T) []codersdk.ChatWatchEvent { + t.Helper() + p.mu.Lock() + defer p.mu.Unlock() + events := make([]codersdk.ChatWatchEvent, 0) + for _, event := range p.events { + var msg codersdk.ChatWatchEvent + if err := json.Unmarshal(event.payload, &msg); err != nil { + continue + } + if event.channel != coderdpubsub.ChatWatchEventChannel(msg.Chat.OwnerID) { + continue + } + events = append(events, msg) + } + return events +} + +func (p *recordingPubsub) stateUpdateMessages(t *testing.T, chatID uuid.UUID) []coderdpubsub.ChatStateUpdateMessage { + t.Helper() + p.mu.Lock() + defer p.mu.Unlock() + messages := make([]coderdpubsub.ChatStateUpdateMessage, 0) + for _, event := range p.events { + if event.channel != coderdpubsub.ChatStateUpdateChannel(chatID) { + continue + } + var msg coderdpubsub.ChatStateUpdateMessage + require.NoError(t, json.Unmarshal(event.payload, &msg)) + messages = append(messages, msg) + } + return messages +} + +func newWorkerTestFixture(t *testing.T) *workerTestFixture { + t.Helper() + db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "openai", + BaseUrl: "http://example.invalid", + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + IsDefault: true, + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + return &workerTestFixture{db: db, pubsub: ps, sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey} +} + +func (f *workerTestFixture) createRunningChat(t *testing.T) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + res, err := chatstate.CreateChat(ctx, f.db, f.pubsub, chatstate.CreateChatInput{ + OrganizationID: f.org.ID, + OwnerID: f.user.ID, + LastModelConfigID: f.model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + userTextMessage(t, "hello", f.user.ID, f.model.ID, f.apiKey.ID), + }, + }) + require.NoError(t, err) + return res.Chat +} + +func (f *workerTestFixture) createRequiresActionChat(t *testing.T) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + toolName := "dynamic_" + uuid.NewString() + dynamicTools, err := json.Marshal([]codersdk.DynamicTool{{ + Name: toolName, + Description: "test tool", + InputSchema: json.RawMessage(`{"type":"object"}`), + }}) + require.NoError(t, err) + res, err := chatstate.CreateChat(ctx, f.db, f.pubsub, chatstate.CreateChatInput{ + OrganizationID: f.org.ID, + OwnerID: f.user.ID, + LastModelConfigID: f.model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + DynamicTools: pqtype.NullRawMessage{ + RawMessage: dynamicTools, + Valid: true, + }, + InitialMessages: []chatstate.Message{ + userTextMessage(t, "hello", f.user.ID, f.model.ID, f.apiKey.ID), + }, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(f.db, f.pubsub, res.Chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + assistantToolCallMessage(t, f.model.ID, toolName), + }, + }) + return err + })) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + chat, err := f.db.GetChatByID(ctx, res.Chat.ID) + require.NoError(t, err) + return chat +} + +func userTextMessage(t *testing.T, text string, createdBy uuid.UUID, modelConfigID uuid.UUID, apiKeyID string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + } +} + +func assistantTextMessage(t *testing.T, text string, modelConfigID uuid.UUID) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + } +} + +func assistantToolCallMessage(t *testing.T, modelConfigID uuid.UUID, toolName string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "call_" + uuid.NewString(), + ToolName: toolName, + Args: json.RawMessage(`{}`), + }}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + } +} + +func testOptions(t *testing.T, f *workerTestFixture, starter chatWorkerTaskStarter) chatWorkerOptions { + t.Helper() + if starter == nil { + starter = newRecordingTaskStarter() + } + return chatWorkerOptions{ + WorkerID: uuid.New(), + Store: f.db, + Pubsub: f.pubsub, + Logger: testutil.Logger(t), + TaskStarter: starter, + AcquisitionInterval: time.Hour, + AcquisitionBatchSize: 10, + RunnerSyncInterval: time.Hour, + HeartbeatInterval: time.Hour, + HeartbeatCleanupInterval: time.Hour, + HeartbeatStaleSeconds: 30, + StateChannelSize: 16, + RunnerManagerChannelSize: 16, + AcquisitionWakeChannelSize: 1, + } +} + +func startWorker(t *testing.T, opts chatWorkerOptions) *chatWorker { + t.Helper() + worker, err := newChatWorker(nil, opts) + require.NoError(t, err) + require.NoError(t, worker.Start(context.Background())) + t.Cleanup(func() { require.NoError(t, worker.Close()) }) + return worker +} + +type taskCall struct { + kind taskKind + input chatWorkerTaskStartInput + ctx context.Context +} + +type releaseGate struct { + once sync.Once + ch chan struct{} +} + +type recordingTaskStarter struct { + mu sync.Mutex + calls []taskCall + callCh chan taskCall + releases []*releaseGate + block bool + ignoreCancel bool +} + +func newRecordingTaskStarter() *recordingTaskStarter { + return &recordingTaskStarter{callCh: make(chan taskCall, 128)} +} + +func newBlockingTaskStarter(ignoreCancel bool) *recordingTaskStarter { + return &recordingTaskStarter{ + callCh: make(chan taskCall, 128), + block: true, + ignoreCancel: ignoreCancel, + } +} + +func (s *recordingTaskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskStartInput) error { + return s.start(ctx, taskKindGeneration, input) +} + +func (s *recordingTaskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskStartInput) error { + return s.start(ctx, taskKindInterrupt, input) +} + +func (s *recordingTaskStarter) StartRequiresActionTimeout(ctx context.Context, input chatWorkerTaskStartInput) error { + return s.start(ctx, taskKindRequiresActionTimeout, input) +} + +func (s *recordingTaskStarter) StartAbandon(ctx context.Context, input chatWorkerTaskStartInput) error { + return s.start(ctx, taskKindAbandon, input) +} + +func (s *recordingTaskStarter) start(ctx context.Context, kind taskKind, input chatWorkerTaskStartInput) error { + call := taskCall{kind: kind, input: input, ctx: ctx} + var gate *releaseGate + s.mu.Lock() + if s.block { + gate = &releaseGate{ch: make(chan struct{})} + s.releases = append(s.releases, gate) + } + s.calls = append(s.calls, call) + s.mu.Unlock() + s.callCh <- call + if gate == nil { + return nil + } + if s.ignoreCancel { + <-gate.ch + return nil + } + select { + case <-gate.ch: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +func (s *recordingTaskStarter) waitCall(t *testing.T, kind taskKind, chatID uuid.UUID) taskCall { + t.Helper() + deadline := time.After(testutil.WaitLong) + for { + select { + case call := <-s.callCh: + if (kind == "" || call.kind == kind) && (chatID == uuid.Nil || call.input.ChatID == chatID) { + return call + } + case <-deadline: + t.Fatalf("timed out waiting for task call kind=%q chat_id=%s", kind, chatID) + return taskCall{} + } + } +} + +func (s *recordingTaskStarter) assertNoCall(t *testing.T) { + t.Helper() + select { + case call := <-s.callCh: + t.Fatalf("unexpected task call: %s for chat %s", call.kind, call.input.ChatID) + case <-time.After(100 * time.Millisecond): + } +} + +func (s *recordingTaskStarter) release(t *testing.T, index int) { + t.Helper() + s.mu.Lock() + defer s.mu.Unlock() + require.Less(t, index, len(s.releases)) + s.releases[index].once.Do(func() { close(s.releases[index].ch) }) +} + +func (s *recordingTaskStarter) releaseAll() { + s.mu.Lock() + defer s.mu.Unlock() + for _, gate := range s.releases { + gate.once.Do(func() { close(gate.ch) }) + } +} + +func finishTurn(t *testing.T, f *workerTestFixture, chatID uuid.UUID) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + chat, err := f.db.GetChatByID(ctx, chatID) + require.NoError(t, err) + return chat +} + +func commitAssistantStep(t *testing.T, f *workerTestFixture, chatID uuid.UUID, text string) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{assistantTextMessage(t, text, f.model.ID)}, + }) + return err + })) + chat, err := f.db.GetChatByID(ctx, chatID) + require.NoError(t, err) + return chat +} + +func interruptChat(t *testing.T, f *workerTestFixture, chatID uuid.UUID) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: userTextMessage(t, "interrupt", f.user.ID, f.model.ID, f.apiKey.ID), + BusyBehavior: chatstate.BusyBehaviorInterrupt, + }) + return err + })) + chat, err := f.db.GetChatByID(ctx, chatID) + require.NoError(t, err) + return chat +} + +func acquireChat(t *testing.T, f *workerTestFixture, chatID uuid.UUID, workerID uuid.UUID, runnerID uuid.UUID) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: workerID, RunnerID: runnerID}) + return err + })) + chat, err := f.db.GetChatByID(ctx, chatID) + require.NoError(t, err) + return chat +} + +func forceExecutionState( + t *testing.T, + f *workerTestFixture, + chatID uuid.UUID, + status database.ChatStatus, + archived bool, +) database.Chat { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + var updated database.Chat + require.NoError(t, f.db.InTx(func(store database.Store) error { + if _, err := store.LockChatAndBumpSnapshotVersion(ctx, chatID); err != nil { + return err + } + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return err + } + updated, err = store.UpdateChatExecutionState(ctx, database.UpdateChatExecutionStateParams{ + ID: chat.ID, + Status: status, + Archived: archived, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + }) + return err + }, nil)) + return updated +} + +func forceExecutionStateAndPublish( + t *testing.T, + f *workerTestFixture, + chatID uuid.UUID, + status database.ChatStatus, + archived bool, +) database.Chat { + t.Helper() + updated := forceExecutionState(t, f, chatID, status, archived) + publishChatUpdate(t, f, updated) + return updated +} + +func publishChatUpdate(t *testing.T, f *workerTestFixture, chat database.Chat) { + t.Helper() + msg := coderdpubsub.ChatStateUpdateMessage{ + SnapshotVersion: chat.SnapshotVersion, + HistoryVersion: chat.HistoryVersion, + QueueVersion: chat.QueueVersion, + RetryStateVersion: chat.RetryStateVersion, + GenerationAttempt: chat.GenerationAttempt, + Status: string(chat.Status), + Archived: chat.Archived, + } + if chat.WorkerID.Valid { + id := chat.WorkerID.UUID + msg.WorkerID = &id + } + if chat.RunnerID.Valid { + id := chat.RunnerID.UUID + msg.RunnerID = &id + } + payload, err := json.Marshal(msg) + require.NoError(t, err) + require.NoError(t, f.pubsub.Publish(coderdpubsub.ChatStateUpdateChannel(chat.ID), payload)) +} + +func makeHeartbeatStale(t *testing.T, f *workerTestFixture, chatID uuid.UUID, runnerID uuid.UUID) time.Time { + t.Helper() + _, err := f.sqlDB.ExecContext( + testutil.Context(t, testutil.WaitShort), + `UPDATE chat_heartbeats SET heartbeat_at = NOW() - INTERVAL '1 hour' WHERE chat_id = $1 AND runner_id = $2`, + chatID, + runnerID, + ) + require.NoError(t, err) + heartbeat, err := f.db.GetChatHeartbeat(testutil.Context(t, testutil.WaitShort), database.GetChatHeartbeatParams{ + ChatID: chatID, + RunnerID: runnerID, + }) + require.NoError(t, err) + return heartbeat.HeartbeatAt +} diff --git a/coderd/x/chatd/instruction.go b/coderd/x/chatd/instruction.go index 02f6dc675a..05476ed6f0 100644 --- a/coderd/x/chatd/instruction.go +++ b/coderd/x/chatd/instruction.go @@ -128,80 +128,6 @@ func instructionFromContextFiles( return formatSystemInstructions(os, dir, contextParts) } -// hasPersistedInstructionFiles reports whether messages include a -// persisted context-file part that should suppress another baseline -// instruction-file lookup. The workspace-agent skill-only sentinel is -// ignored so default instructions still load on fresh chats. -func hasPersistedInstructionFiles( - messages []database.ChatMessage, -) bool { - for _, msg := range messages { - if !msg.Content.Valid || - !bytes.Contains(msg.Content.RawMessage, []byte(`"context-file"`)) { - continue - } - var parts []codersdk.ChatMessagePart - if err := json.Unmarshal(msg.Content.RawMessage, &parts); err != nil { - continue - } - for _, part := range parts { - if part.Type != codersdk.ChatMessagePartTypeContextFile || - !part.ContextFileAgentID.Valid || - part.ContextFilePath == AgentChatContextSentinelPath { - continue - } - return true - } - } - return false -} - -func mergeSkillMetas( - persisted []chattool.SkillMeta, - discovered []chattool.SkillMeta, -) []chattool.SkillMeta { - if len(persisted) == 0 { - return discovered - } - if len(discovered) == 0 { - return persisted - } - - seen := make(map[string]struct{}, len(persisted)+len(discovered)) - merged := make([]chattool.SkillMeta, 0, len(persisted)+len(discovered)) - appendUnique := func(skill chattool.SkillMeta) { - if _, ok := seen[skill.Name]; ok { - return - } - seen[skill.Name] = struct{}{} - merged = append(merged, skill) - } - for _, skill := range discovered { - appendUnique(skill) - } - for _, skill := range persisted { - appendUnique(skill) - } - return merged -} - -// selectSkillMetasForInstructionRefresh chooses which skill metadata -// should be injected on a turn that refreshes instruction files. -func selectSkillMetasForInstructionRefresh( - persisted []chattool.SkillMeta, - discovered []chattool.SkillMeta, - currentAgentID uuid.NullUUID, - latestInjectedAgentID uuid.NullUUID, -) []chattool.SkillMeta { - if currentAgentID.Valid && latestInjectedAgentID.Valid && latestInjectedAgentID.UUID == currentAgentID.UUID { - return mergeSkillMetas(persisted, discovered) - } - if !currentAgentID.Valid && len(discovered) == 0 { - return persisted - } - return discovered -} - // skillsFromParts reconstructs skill metadata from persisted // skill parts. This is analogous to instructionFromContextFiles // so the skill index can be re-injected after compaction without @@ -238,19 +164,3 @@ func skillsFromParts( } return skills } - -// filterSkillParts returns stripped copies of skill-type parts from -// the given slice. Internal fields are removed so the result is safe -// for the cache column. Returns nil when no skill parts exist. -func filterSkillParts(parts []codersdk.ChatMessagePart) []codersdk.ChatMessagePart { - var out []codersdk.ChatMessagePart - for _, p := range parts { - if p.Type != codersdk.ChatMessagePartTypeSkill { - continue - } - cp := p - cp.StripInternal() - out = append(out, cp) - } - return out -} diff --git a/coderd/x/chatd/integration_responses_test.go b/coderd/x/chatd/integration_responses_test.go index 97e1f0a076..e6ca40b1c8 100644 --- a/coderd/x/chatd/integration_responses_test.go +++ b/coderd/x/chatd/integration_responses_test.go @@ -78,6 +78,7 @@ func TestOpenAIResponsesNoStaleWebSearchReplay(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: uniqueResponsesTitle(t, "no-stale"), ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -92,6 +93,7 @@ func TestOpenAIResponsesNoStaleWebSearchReplay(t *testing.T) { _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ModelConfigID: model.ID, Content: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("summarize the result without searching again"), @@ -167,6 +169,7 @@ func TestOpenAIResponsesFullReplayPairsReasoningAndWebSearch(t *testing.T) { chat, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: uniqueResponsesTitle(t, "full-replay"), ModelConfigID: firstModel.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -181,6 +184,7 @@ func TestOpenAIResponsesFullReplayPairsReasoningAndWebSearch(t *testing.T) { _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ModelConfigID: secondModel.ID, Content: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("summarize the result without searching again"), @@ -253,6 +257,7 @@ func TestOpenAIResponsesChainModeSkipsWhenLocalCallPending(t *testing.T) { _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ModelConfigID: model.ID, Content: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("continue after that tool call"), @@ -341,6 +346,7 @@ func TestOpenAIResponsesChainModeStillFiresForProviderExecutedOnly(t *testing.T) _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ModelConfigID: model.ID, Content: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("what did it find"), diff --git a/coderd/x/chatd/integration_test.go b/coderd/x/chatd/integration_test.go index 0203eced36..249d72ec3d 100644 --- a/coderd/x/chatd/integration_test.go +++ b/coderd/x/chatd/integration_test.go @@ -113,7 +113,7 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { }) require.NoError(t, err) - // --- Step 1: Send a message that triggers web_search --- + // Step 1: Send a message that triggers web_search. t.Log("Creating chat with web search query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: user.OrganizationID, @@ -172,7 +172,7 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { } } - // --- Step 2: Send a follow-up message --- + // Step 2: Send a follow-up message. // This is the critical test: if PE tool results were lost during // persistence, the reconstructed conversation will be rejected // by Anthropic because server_tool_use has no matching @@ -374,7 +374,7 @@ func TestOpenAIReasoningRoundTrip(t *testing.T) { }) require.NoError(t, err) - // --- Step 1: Send a message that triggers reasoning --- + // Step 1: Send a message that triggers reasoning. t.Log("Creating chat with reasoning query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: user.OrganizationID, @@ -418,7 +418,7 @@ func TestOpenAIReasoningRoundTrip(t *testing.T) { require.Contains(t, partTypes, codersdk.ChatMessagePartTypeText, "assistant message should contain a text part") - // --- Step 2: Send a follow-up message --- + // Step 2: Send a follow-up message. // This is the critical test: if reasoning items are sent back // without their required following item, the API will reject // the request with: @@ -524,7 +524,7 @@ func TestOpenAIReasoningRoundTripStoreFalse(t *testing.T) { }) require.NoError(t, err) - // --- Step 1: Send a message that triggers reasoning --- + // Step 1: Send a message that triggers reasoning. t.Log("Creating chat with reasoning query...") chat, err := expClient.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: user.OrganizationID, @@ -568,7 +568,7 @@ func TestOpenAIReasoningRoundTripStoreFalse(t *testing.T) { require.Contains(t, partTypes, codersdk.ChatMessagePartTypeText, "assistant message should contain a text part") - // --- Step 2: Send a follow-up message --- + // Step 2: Send a follow-up message. // This is the critical test: when Store is false, item IDs are // ephemeral and cannot be looked up from OpenAI later. t.Log("Sending follow-up message...") diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go new file mode 100644 index 0000000000..b587836151 --- /dev/null +++ b/coderd/x/chatd/message_conversion.go @@ -0,0 +1,850 @@ +package chatd + +import ( + "cmp" + "context" + "database/sql" + "encoding/json" + "slices" + "strings" + "time" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatcost" + "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/chattool" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" +) + +const interruptedToolResultErrorMessage = "tool call was interrupted before it produced a result" + +type buildCommitStepMessagesInput struct { + modelConfigID uuid.UUID + modelCallConfig codersdk.ChatModelCallConfig + step stepData + toolNameToConfigID map[string]uuid.UUID + logger slog.Logger + contentVersion int16 +} + +type stepMessagesForCommit struct { + Messages []chatstate.Message + VisibleIndexes []int +} + +func buildCommitStepMessages(input buildCommitStepMessagesInput) (stepMessagesForCommit, error) { + contentVersion := input.contentVersion + if contentVersion == 0 { + contentVersion = chatprompt.CurrentContentVersion + } + + assistantBlocks, toolResults := splitStepContent(input.step.Content) + assistantParts := buildAssistantParts(input.logger, assistantBlocks, toolResults, input.step, input.toolNameToConfigID) + + messages := make([]chatstate.Message, 0, 1+len(toolResults)) + if len(assistantParts) > 0 { + assistantContent, err := chatprompt.MarshalParts(assistantParts) + if err != nil { + return stepMessagesForCommit{}, xerrors.Errorf("marshal assistant content: %w", err) + } + messages = append(messages, assistantMessage(input.modelConfigID, contentVersion, assistantContent, input.step, input.modelCallConfig)) + } + + for _, toolResult := range toolResults { + part := chatprompt.PartFromContentWithLogger(context.Background(), input.logger, toolResult) + applyToolMetadata(&part, input.toolNameToConfigID) + if part.ToolCallID != "" && input.step.ToolResultCreatedAt != nil { + if ts, ok := input.step.ToolResultCreatedAt[part.ToolCallID]; ok { + part.CreatedAt = &ts + } + } + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part}) + if err != nil { + return stepMessagesForCommit{}, xerrors.Errorf("marshal tool result: %w", err) + } + messages = append(messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, content)) + } + + return stepMessagesForCommit{ + Messages: messages, + VisibleIndexes: visibleMessageIndexes(messages), + }, nil +} + +func splitStepContent(content []fantasy.Content) ([]fantasy.Content, []fantasy.ToolResultContent) { + assistantBlocks := make([]fantasy.Content, 0, len(content)) + toolResults := make([]fantasy.ToolResultContent, 0) + for _, block := range content { + if tr, ok := asToolResultContent(block); ok && !tr.ProviderExecuted { + toolResults = append(toolResults, tr) + continue + } + assistantBlocks = append(assistantBlocks, block) + } + return assistantBlocks, toolResults +} + +func asToolResultContent(block fantasy.Content) (fantasy.ToolResultContent, bool) { + if tr, ok := fantasy.AsContentType[fantasy.ToolResultContent](block); ok { + return tr, true + } + if tr, ok := fantasy.AsContentType[*fantasy.ToolResultContent](block); ok && tr != nil { + return *tr, true + } + return fantasy.ToolResultContent{}, false +} + +func buildAssistantParts( + logger slog.Logger, + assistantBlocks []fantasy.Content, + toolResults []fantasy.ToolResultContent, + step stepData, + toolNameToConfigID map[string]uuid.UUID, +) []codersdk.ChatMessagePart { + parts := make([]codersdk.ChatMessagePart, 0, len(assistantBlocks)+len(toolResults)) + reasoningIdx := 0 + for _, block := range assistantBlocks { + part := chatprompt.PartFromContentWithLogger(context.Background(), logger, block) + applyToolMetadata(&part, toolNameToConfigID) + switch part.Type { + case codersdk.ChatMessagePartTypeToolCall: + if part.ToolCallID != "" && step.ToolCallCreatedAt != nil { + if ts, ok := step.ToolCallCreatedAt[part.ToolCallID]; ok { + part.CreatedAt = &ts + } + } + case codersdk.ChatMessagePartTypeToolResult: + if part.ToolCallID != "" && step.ToolResultCreatedAt != nil { + if ts, ok := step.ToolResultCreatedAt[part.ToolCallID]; ok { + part.CreatedAt = &ts + } + } + case codersdk.ChatMessagePartTypeReasoning: + if reasoningIdx < len(step.ReasoningStartedAt) { + if ts := step.ReasoningStartedAt[reasoningIdx]; !ts.IsZero() { + part.CreatedAt = &ts + } + } + if reasoningIdx < len(step.ReasoningCompletedAt) { + if ts := step.ReasoningCompletedAt[reasoningIdx]; !ts.IsZero() { + part.CompletedAt = &ts + } + } + reasoningIdx++ + } + if part.Type != "" { + parts = append(parts, part) + } + } + for _, tr := range toolResults { + attachments, err := chattool.AttachmentsFromMetadata(tr.ClientMetadata) + if err != nil { + logger.Warn(context.Background(), "skipping malformed tool attachment metadata", + slog.F("tool_name", tr.ToolName), + slog.F("tool_call_id", tr.ToolCallID), + slog.Error(err), + ) + continue + } + for _, attachment := range attachments { + parts = append(parts, codersdk.ChatMessageFile(attachment.FileID, attachment.MediaType, attachment.Name)) + } + } + return parts +} + +func applyToolMetadata(part *codersdk.ChatMessagePart, toolNameToConfigID map[string]uuid.UUID) { + if part.ToolName == "" || len(toolNameToConfigID) == 0 { + return + } + if configID, ok := toolNameToConfigID[part.ToolName]; ok { + part.MCPServerConfigID = uuid.NullUUID{UUID: configID, Valid: true} + } +} + +func assistantMessage( + modelConfigID uuid.UUID, + contentVersion int16, + content pqtype.NullRawMessage, + step stepData, + modelCallConfig codersdk.ChatModelCallConfig, +) chatstate.Message { + msg := baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, modelConfigID, contentVersion, content) + if step.Usage != (fantasy.Usage{}) { + msg.InputTokens = nullInt64IfNonZero(step.Usage.InputTokens) + msg.OutputTokens = nullInt64IfNonZero(step.Usage.OutputTokens) + msg.TotalTokens = nullInt64IfNonZero(step.Usage.TotalTokens) + msg.ReasoningTokens = nullInt64IfNonZero(step.Usage.ReasoningTokens) + msg.CacheCreationTokens = nullInt64IfNonZero(step.Usage.CacheCreationTokens) + msg.CacheReadTokens = nullInt64IfNonZero(step.Usage.CacheReadTokens) + usage := codersdk.ChatMessageUsage{ + InputTokens: int64PtrIfNonZero(step.Usage.InputTokens), + OutputTokens: int64PtrIfNonZero(step.Usage.OutputTokens), + ReasoningTokens: int64PtrIfNonZero(step.Usage.ReasoningTokens), + CacheCreationTokens: int64PtrIfNonZero(step.Usage.CacheCreationTokens), + CacheReadTokens: int64PtrIfNonZero(step.Usage.CacheReadTokens), + } + if totalCost := chatcost.CalculateTotalCostMicros(usage, modelCallConfig.Cost); totalCost != nil { + msg.TotalCostMicros = sql.NullInt64{Int64: *totalCost, Valid: true} + } + } + msg.ContextLimit = step.ContextLimit + if step.Runtime > 0 { + msg.RuntimeMs = sql.NullInt64{Int64: step.Runtime.Milliseconds(), Valid: true} + } + if step.ProviderResponseID != "" { + msg.ProviderResponseID = sql.NullString{String: step.ProviderResponseID, Valid: true} + } + return msg +} + +func baseMessage( + role database.ChatMessageRole, + visibility database.ChatMessageVisibility, + modelConfigID uuid.UUID, + contentVersion int16, + content pqtype.NullRawMessage, +) chatstate.Message { + return chatstate.Message{ + Role: role, + Content: content, + Visibility: visibility, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ContentVersion: contentVersion, + } +} + +func nullInt64IfNonZero(value int64) sql.NullInt64 { + if value == 0 { + return sql.NullInt64{} + } + return sql.NullInt64{Int64: value, Valid: true} +} + +func int64PtrIfNonZero(value int64) *int64 { + if value == 0 { + return nil + } + return &value +} + +func visibleMessageIndexes(messages []chatstate.Message) []int { + indexes := make([]int, 0, len(messages)) + for i, msg := range messages { + if msg.Visibility == database.ChatMessageVisibilityBoth || msg.Visibility == database.ChatMessageVisibilityUser { + indexes = append(indexes, i) + } + } + return indexes +} + +func textFromParts(parts []codersdk.ChatMessagePart) string { + var builder strings.Builder + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeText { + _, _ = builder.WriteString(part.Text) + } + } + return builder.String() +} + +type buildCompactionMessagesInput struct { + modelConfigID uuid.UUID + activeAPIKeyID string + toolCallID string + toolName string + compaction compactionOutcome + contentVersion int16 +} + +type compactionMessagesForCommit struct { + Messages []chatstate.Message + HiddenCount int +} + +func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMessagesForCommit, error) { + contentVersion := input.contentVersion + if contentVersion == 0 { + contentVersion = chatprompt.CurrentContentVersion + } + toolName := input.toolName + if toolName == "" { + toolName = "chat_summarized" + } + + systemContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(input.compaction.SystemSummary)}) + if err != nil { + return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction system summary: %w", err) + } + args, err := json.Marshal(map[string]any{ + "source": "automatic", + "threshold_percent": input.compaction.ThresholdPercent, + }) + if err != nil { + return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction args: %w", err) + } + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolCall(input.toolCallID, toolName, args), + }) + if err != nil { + return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction tool call: %w", err) + } + summaryResult, err := json.Marshal(map[string]any{ + "summary": input.compaction.SummaryReport, + "source": "automatic", + "threshold_percent": input.compaction.ThresholdPercent, + "usage_percent": input.compaction.UsagePercent, + "context_tokens": input.compaction.ContextTokens, + "context_limit_tokens": input.compaction.ContextLimit, + }) + if err != nil { + return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction result: %w", err) + } + toolContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageToolResult(input.toolCallID, toolName, summaryResult, false, false), + }) + if err != nil { + return compactionMessagesForCommit{}, xerrors.Errorf("marshal compaction tool result: %w", err) + } + + messages := []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: systemContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: uuid.NullUUID{UUID: input.modelConfigID, Valid: input.modelConfigID != uuid.Nil}, + ContentVersion: contentVersion, + APIKeyID: sql.NullString{String: input.activeAPIKeyID, Valid: input.activeAPIKeyID != ""}, + }, + baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityUser, input.modelConfigID, contentVersion, assistantContent), + baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, input.modelConfigID, contentVersion, toolContent), + } + for i := range messages { + messages[i].Compressed = true + } + return compactionMessagesForCommit{Messages: messages, HiddenCount: 1}, nil +} + +func currentTurnStepCount(messages []database.ChatMessage) int { + latestUser := -1 + for i, msg := range messages { + if msg.Deleted || msg.Compressed { + continue + } + if msg.Role == database.ChatMessageRoleUser { + latestUser = i + } + } + count := 0 + for i := latestUser + 1; i < len(messages); i++ { + msg := messages[i] + if msg.Deleted || msg.Compressed { + continue + } + if msg.Role == database.ChatMessageRoleAssistant { + count++ + } + } + return count +} + +type compactionRequirement int + +const ( + compactionRequirementNotNeeded compactionRequirement = iota + compactionRequirementNeeded +) + +func compactionStatusFromHistory(messages []database.ChatMessage, requirement compactionRequirement) compactionStatus { + boundaryIndex := latestCompactionBoundaryIndex(messages) + if requirement == compactionRequirementNeeded { + if boundaryIndex == -1 { + return compactionStatusNeeded + } + if hasUncompressedAssistantAfter(messages, boundaryIndex) { + return compactionStatusStillOverLimit + } + return compactionStatusAfterCompaction + } + if boundaryIndex != -1 && !hasUncompressedAssistantAfter(messages, boundaryIndex) { + return compactionStatusAfterCompaction + } + return compactionStatusNotNeeded +} + +func latestCompactionBoundaryIndex(messages []database.ChatMessage) int { + for i := len(messages) - 1; i >= 0; i-- { + if isCompactionBoundaryMessage(messages[i]) { + return i + } + } + return -1 +} + +func isCompactionBoundaryMessage(msg database.ChatMessage) bool { + if msg.Deleted || !msg.Compressed { + return false + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + return false + } + for _, part := range parts { + if part.ToolName == "chat_summarized" && + (part.Type == codersdk.ChatMessagePartTypeToolCall || part.Type == codersdk.ChatMessagePartTypeToolResult) { + return true + } + } + return false +} + +func hasUncompressedAssistantAfter(messages []database.ChatMessage, index int) bool { + for i := index + 1; i < len(messages); i++ { + msg := messages[i] + if msg.Deleted || msg.Compressed { + continue + } + if msg.Role == database.ChatMessageRoleAssistant { + return true + } + } + return false +} + +func historyHasStopAfterToolResult(messages []database.ChatMessage, stopAfterTools map[string]struct{}) (bool, error) { + if len(stopAfterTools) == 0 { + return false, nil + } + start := 0 + for i, msg := range messages { + if msg.Deleted || msg.Compressed { + continue + } + if msg.Role == database.ChatMessageRoleUser { + start = i + 1 + } + } + for _, msg := range messages[start:] { + if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + return false, xerrors.Errorf("parse tool message: %w", err) + } + for _, part := range parts { + if part.Type != codersdk.ChatMessagePartTypeToolResult || part.IsError { + continue + } + if _, ok := stopAfterTools[part.ToolName]; ok { + return true, nil + } + } + } + return false, nil +} + +func currentHistoryComplete(messages []database.ChatMessage) (bool, error) { + idx := lastMessageIndex(messages, func(database.ChatMessage) bool { return true }) + if idx == -1 || messages[idx].Role != database.ChatMessageRoleAssistant { + return false, nil + } + parts, err := chatprompt.ParseContent(messages[idx]) + if err != nil { + return false, xerrors.Errorf("parse latest assistant message: %w", err) + } + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolCall && !part.ProviderExecuted { + return false, nil + } + } + return true, nil +} + +func lastMessageIndex(messages []database.ChatMessage, accept func(database.ChatMessage) bool) int { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].Deleted || messages[i].Compressed { + continue + } + if accept(messages[i]) { + return i + } + } + return -1 +} + +func handledToolCallIDs(messages []database.ChatMessage) (map[string]bool, error) { + handled := make(map[string]bool) + for _, msg := range messages { + if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + return nil, xerrors.Errorf("parse tool message: %w", err) + } + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.ToolCallID != "" { + handled[part.ToolCallID] = true + } + } + } + return handled, nil +} + +type bufferedPartsToPartialMessagesInput struct { + parts []messagepartbuffer.Part + modelConfigID uuid.UUID + contentVersion int16 + logger slog.Logger + interruptedAt time.Time +} + +type partialToolCall struct { + part codersdk.ChatMessagePart + index int + argsDelta strings.Builder + valid bool + durable bool +} + +type partialToolResult struct { + part codersdk.ChatMessagePart + resultDelta strings.Builder + completed bool +} + +func bufferedPartsToPartialMessages(input bufferedPartsToPartialMessagesInput) ([]chatstate.Message, error) { + contentVersion := input.contentVersion + if contentVersion == 0 { + contentVersion = chatprompt.CurrentContentVersion + } + parts := slices.Clone(input.parts) + slices.SortFunc(parts, func(a, b messagepartbuffer.Part) int { + return cmp.Compare(a.Seq, b.Seq) + }) + + state := partialMessageConversionState{ + input: input, + contentVersion: contentVersion, + toolCalls: make(map[string]*partialToolCall), + toolResults: make(map[string]*partialToolResult), + answered: make(map[string]bool), + } + for _, buffered := range parts { + if err := state.consume(buffered); err != nil { + return nil, err + } + } + if err := state.finalizeToolCallPlaceholders(); err != nil { + return nil, err + } + if err := state.flushAssistant(); err != nil { + return nil, err + } + if err := state.flushAccumulatedToolResults(); err != nil { + return nil, err + } + if err := state.appendSyntheticInterruptionResults(); err != nil { + return nil, err + } + return state.messages, nil +} + +type partialMessageConversionState struct { + input bufferedPartsToPartialMessagesInput + contentVersion int16 + + messages []chatstate.Message + assistantParts []codersdk.ChatMessagePart + toolCalls map[string]*partialToolCall + toolCallOrder []string + toolResults map[string]*partialToolResult + toolResultOrder []string + answered map[string]bool +} + +func (s *partialMessageConversionState) consume(buffered messagepartbuffer.Part) error { + switch buffered.Role { + case codersdk.ChatMessageRoleAssistant: + s.consumeAssistantPart(buffered) + case codersdk.ChatMessageRoleTool: + return s.consumeToolPart(buffered) + default: + s.logSkippedPart(buffered, "unsupported buffered part role") + } + return nil +} + +func (s *partialMessageConversionState) consumeAssistantPart(buffered messagepartbuffer.Part) { + part := buffered.MessagePart + if part.Type == "" { + s.logSkippedPart(buffered, "empty buffered assistant part type") + return + } + if part.Type != codersdk.ChatMessagePartTypeToolCall { + if part.Type == codersdk.ChatMessagePartTypeReasoning && + !s.input.interruptedAt.IsZero() { + interruptedAt := s.input.interruptedAt + if part.CreatedAt == nil { + part.CreatedAt = &interruptedAt + } + if part.CompletedAt == nil { + part.CompletedAt = &interruptedAt + } + } + s.assistantParts = append(s.assistantParts, part) + return + } + if part.ToolCallID == "" { + s.logSkippedPart(buffered, "tool call part missing tool call ID") + return + } + call := s.toolCall(part.ToolCallID) + call.part.Type = codersdk.ChatMessagePartTypeToolCall + call.part.ToolCallID = part.ToolCallID + if part.ToolName != "" { + call.part.ToolName = part.ToolName + } + if part.MCPServerConfigID.Valid { + call.part.MCPServerConfigID = part.MCPServerConfigID + } + if part.CreatedAt != nil { + call.part.CreatedAt = part.CreatedAt + } + call.part.ProviderExecuted = call.part.ProviderExecuted || part.ProviderExecuted + + if part.ArgsDelta != "" { + if call.durable { + s.logSkippedPart(buffered, "tool call args delta arrived after full tool call") + return + } + _, _ = call.argsDelta.WriteString(part.ArgsDelta) + return + } + + durable := part + durable.ArgsDelta = "" + if len(durable.Args) > 0 && !json.Valid(durable.Args) { + call.valid = false + s.assistantParts[call.index] = codersdk.ChatMessagePart{} + s.logSkippedPart(buffered, "tool call part has invalid durable args") + return + } + if call.durable { + s.logSkippedPart(buffered, "duplicate durable tool call part") + } + call.part = durable + call.valid = true + call.durable = true + s.assistantParts[call.index] = durable +} + +func (s *partialMessageConversionState) consumeToolPart(buffered messagepartbuffer.Part) error { + part := buffered.MessagePart + if part.Type != codersdk.ChatMessagePartTypeToolResult { + s.logSkippedPart(buffered, "non tool-result part with tool role") + return nil + } + if part.ToolCallID == "" { + s.logSkippedPart(buffered, "tool result part missing tool call ID") + return nil + } + if part.ResultReset { + result := s.toolResult(part.ToolCallID) + result.part.ToolCallID = part.ToolCallID + result.part.ToolName = part.ToolName + result.resultDelta.Reset() + s.logSkippedPart(buffered, "streaming tool result reset is not durable") + return nil + } + if part.ResultDelta != "" { + result := s.toolResult(part.ToolCallID) + result.part.ToolCallID = part.ToolCallID + if part.ToolName != "" { + result.part.ToolName = part.ToolName + } + if part.MCPServerConfigID.Valid { + result.part.MCPServerConfigID = part.MCPServerConfigID + } + if part.CreatedAt != nil { + result.part.CreatedAt = part.CreatedAt + } + result.part.ProviderExecuted = result.part.ProviderExecuted || part.ProviderExecuted + _, _ = result.resultDelta.WriteString(part.ResultDelta) + return nil + } + if err := s.finalizeToolCallPlaceholders(); err != nil { + return err + } + if !s.toolCallDurable(part.ToolCallID) { + s.logSkippedPart(buffered, "tool result has no matching durable tool call") + return nil + } + if len(part.Result) == 0 || !json.Valid(part.Result) { + s.logSkippedPart(buffered, "tool result part has invalid durable result") + return nil + } + if s.answered[part.ToolCallID] { + s.logSkippedPart(buffered, "duplicate durable tool result part") + return nil + } + part.ResultDelta = "" + part.ResultReset = false + if err := s.flushAssistant(); err != nil { + return err + } + if err := s.appendToolResult(part); err != nil { + return err + } + s.answered[part.ToolCallID] = true + return nil +} + +func (s *partialMessageConversionState) toolCall(id string) *partialToolCall { + call := s.toolCalls[id] + if call != nil { + return call + } + call = &partialToolCall{index: len(s.assistantParts), valid: true} + s.toolCalls[id] = call + s.toolCallOrder = append(s.toolCallOrder, id) + s.assistantParts = append(s.assistantParts, codersdk.ChatMessagePart{}) + return call +} + +func (s *partialMessageConversionState) toolResult(id string) *partialToolResult { + result := s.toolResults[id] + if result != nil { + return result + } + result = &partialToolResult{} + s.toolResults[id] = result + s.toolResultOrder = append(s.toolResultOrder, id) + return result +} + +func (s *partialMessageConversionState) finalizeToolCallPlaceholders() error { + for _, id := range s.toolCallOrder { + call := s.toolCalls[id] + if call == nil || call.durable || !call.valid { + continue + } + args := json.RawMessage(call.argsDelta.String()) + if len(args) == 0 || !json.Valid(args) { + s.assistantParts[call.index] = codersdk.ChatMessagePart{} + call.valid = false + s.logSkippedPart(messagepartbuffer.Part{ + Role: codersdk.ChatMessageRoleAssistant, + MessagePart: call.part, + }, "tool call args delta did not form durable JSON") + continue + } + call.part.Args = args + call.part.ArgsDelta = "" + call.durable = true + s.assistantParts[call.index] = call.part + } + return nil +} + +func (s *partialMessageConversionState) flushAssistant() error { + if len(s.assistantParts) == 0 { + return nil + } + durable := make([]codersdk.ChatMessagePart, 0, len(s.assistantParts)) + for _, part := range s.assistantParts { + if part.Type == "" { + continue + } + part.ArgsDelta = "" + part.ResultDelta = "" + part.ResultReset = false + durable = append(durable, part) + } + s.assistantParts = nil + if len(durable) == 0 { + return nil + } + content, err := chatprompt.MarshalParts(durable) + if err != nil { + return xerrors.Errorf("marshal partial assistant: %w", err) + } + s.messages = append(s.messages, baseMessage(database.ChatMessageRoleAssistant, database.ChatMessageVisibilityBoth, s.input.modelConfigID, s.contentVersion, content)) + return nil +} + +func (s *partialMessageConversionState) flushAccumulatedToolResults() error { + for _, id := range s.toolResultOrder { + if s.answered[id] { + continue + } + result := s.toolResults[id] + if result == nil || result.completed { + continue + } + if result.resultDelta.Len() == 0 { + continue + } + s.logSkippedPart(messagepartbuffer.Part{Role: codersdk.ChatMessageRoleTool, MessagePart: result.part}, "streaming tool result delta is not durable") + } + return nil +} + +func (s *partialMessageConversionState) appendToolResult(part codersdk.ChatMessagePart) error { + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{part}) + if err != nil { + return xerrors.Errorf("marshal partial tool result: %w", err) + } + s.messages = append(s.messages, baseMessage(database.ChatMessageRoleTool, database.ChatMessageVisibilityBoth, s.input.modelConfigID, s.contentVersion, content)) + return nil +} + +func (s *partialMessageConversionState) appendSyntheticInterruptionResults() error { + for _, id := range s.toolCallOrder { + if s.answered[id] { + continue + } + call := s.toolCalls[id] + if call == nil || !call.valid || !call.durable || call.part.ProviderExecuted { + continue + } + result, err := json.Marshal(map[string]string{"error": interruptedToolResultErrorMessage}) + if err != nil { + return xerrors.Errorf("marshal synthetic interruption result: %w", err) + } + part := codersdk.ChatMessageToolResult(call.part.ToolCallID, call.part.ToolName, result, true, false) + part.MCPServerConfigID = call.part.MCPServerConfigID + if !s.input.interruptedAt.IsZero() { + part.CreatedAt = &s.input.interruptedAt + } + if err := s.appendToolResult(part); err != nil { + return xerrors.Errorf("marshal synthetic interruption message: %w", err) + } + s.answered[id] = true + } + return nil +} + +func (s *partialMessageConversionState) toolCallDurable(id string) bool { + call := s.toolCalls[id] + return call != nil && call.valid && call.durable +} + +func (s *partialMessageConversionState) logSkippedPart(buffered messagepartbuffer.Part, reason string) { + s.input.logger.Warn(context.Background(), "skipping buffered chat message part", + slog.F("reason", reason), + slog.F("role", buffered.Role), + slog.F("part_type", buffered.MessagePart.Type), + slog.F("tool_call_id", buffered.MessagePart.ToolCallID), + slog.F("tool_name", buffered.MessagePart.ToolName), + ) +} diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go new file mode 100644 index 0000000000..5469901e10 --- /dev/null +++ b/coderd/x/chatd/message_conversion_test.go @@ -0,0 +1,532 @@ +package chatd //nolint:testpackage // Uses unexported chatworker helpers. + +import ( + "context" + "database/sql" + "encoding/json" + "sync" + "testing" + "time" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/shopspring/decimal" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3" + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" +) + +func TestBuildCommitStepMessages_AssistantTextAndReasoning(t *testing.T) { + t.Parallel() + + modelConfigID := uuid.New() + startedAt := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC) + completedAt := startedAt.Add(2 * time.Second) + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: modelConfigID, + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{ + Content: []fantasy.Content{ + fantasy.ReasoningContent{Text: "thinking"}, + fantasy.TextContent{Text: "hello"}, + }, + ReasoningStartedAt: []time.Time{startedAt}, + ReasoningCompletedAt: []time.Time{completedAt}, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 1) + require.Equal(t, []int{0}, got.VisibleIndexes) + + msg := got.Messages[0] + require.Equal(t, database.ChatMessageRoleAssistant, msg.Role) + require.Equal(t, database.ChatMessageVisibilityBoth, msg.Visibility) + require.Equal(t, uuid.NullUUID{UUID: modelConfigID, Valid: true}, msg.ModelConfigID) + require.Equal(t, chatprompt.CurrentContentVersion, msg.ContentVersion) + parts := parseMessageParts(t, msg.Role, msg.Content) + require.Len(t, parts, 2) + require.Equal(t, codersdk.ChatMessagePartTypeReasoning, parts[0].Type) + require.Equal(t, "thinking", parts[0].Text) + require.Equal(t, startedAt, requireNotNilTime(t, parts[0].CreatedAt)) + require.Equal(t, completedAt, requireNotNilTime(t, parts[0].CompletedAt)) + require.Equal(t, codersdk.ChatMessagePartTypeText, parts[1].Type) + require.Equal(t, "hello", parts[1].Text) +} + +func TestBuildCommitStepMessages_LocalToolResultsBecomeToolMessages(t *testing.T) { + t.Parallel() + + modelConfigID := uuid.New() + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: modelConfigID, + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{Content: []fantasy.Content{ + fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "execute", Input: `{"cmd":"pwd"}`}, + fantasy.ToolResultContent{ + ToolCallID: "call-1", + ToolName: "execute", + Result: fantasy.ToolResultOutputContentText{Text: `{"stdout":"/tmp"}`}, + }, + }}, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + require.Equal(t, []int{0, 1}, got.VisibleIndexes) + + assistantParts := parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content) + require.Len(t, assistantParts, 1) + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, assistantParts[0].Type) + require.Equal(t, "call-1", assistantParts[0].ToolCallID) + require.Equal(t, "execute", assistantParts[0].ToolName) + + toolParts := parseMessageParts(t, got.Messages[1].Role, got.Messages[1].Content) + require.Len(t, toolParts, 1) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, toolParts[0].Type) + require.Equal(t, "call-1", toolParts[0].ToolCallID) + require.Equal(t, "execute", toolParts[0].ToolName) + require.JSONEq(t, `{"stdout":"/tmp"}`, string(toolParts[0].Result)) +} + +func TestBuildCommitStepMessages_ProviderExecutedResultsStayAssistantContent(t *testing.T) { + t.Parallel() + + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + step: stepData{Content: []fantasy.Content{ + fantasy.ToolCallContent{ + ToolCallID: "web-1", + ToolName: "web_search", + ProviderExecuted: true, + }, + fantasy.ToolResultContent{ + ToolCallID: "web-1", + ToolName: "web_search", + ProviderExecuted: true, + Result: fantasy.ToolResultOutputContentText{Text: `{"ok":true}`}, + }, + }}, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 1) + parts := parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content) + require.Len(t, parts, 2) + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, parts[0].Type) + require.True(t, parts[0].ProviderExecuted) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, parts[1].Type) + require.True(t, parts[1].ProviderExecuted) +} + +func TestBuildCommitStepMessages_UsageCostRuntimeProviderResponseID(t *testing.T) { + t.Parallel() + + inputPrice := decimal.NewFromFloat(2.5) + outputPrice := decimal.NewFromFloat(7.5) + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + modelCallConfig: codersdk.ChatModelCallConfig{ + Cost: &codersdk.ModelCostConfig{ + InputPricePerMillionTokens: &inputPrice, + OutputPricePerMillionTokens: &outputPrice, + }, + }, + step: stepData{ + Content: []fantasy.Content{fantasy.TextContent{Text: "usage"}}, + Usage: fantasy.Usage{InputTokens: 100, OutputTokens: 20, TotalTokens: 120, ReasoningTokens: 3, CacheCreationTokens: 4, CacheReadTokens: 5}, + ContextLimit: sql.NullInt64{Int64: 4096, Valid: true}, + ProviderResponseID: "resp-123", + Runtime: 1500 * time.Millisecond, + }, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 1) + msg := got.Messages[0] + require.Equal(t, sql.NullInt64{Int64: 100, Valid: true}, msg.InputTokens) + require.Equal(t, sql.NullInt64{Int64: 20, Valid: true}, msg.OutputTokens) + require.Equal(t, sql.NullInt64{Int64: 120, Valid: true}, msg.TotalTokens) + require.Equal(t, sql.NullInt64{Int64: 3, Valid: true}, msg.ReasoningTokens) + require.Equal(t, sql.NullInt64{Int64: 4, Valid: true}, msg.CacheCreationTokens) + require.Equal(t, sql.NullInt64{Int64: 5, Valid: true}, msg.CacheReadTokens) + require.Equal(t, sql.NullInt64{Int64: 4096, Valid: true}, msg.ContextLimit) + require.Equal(t, sql.NullInt64{Int64: 1500, Valid: true}, msg.RuntimeMs) + require.Equal(t, sql.NullString{String: "resp-123", Valid: true}, msg.ProviderResponseID) + require.True(t, msg.TotalCostMicros.Valid) + require.Greater(t, msg.TotalCostMicros.Int64, int64(0)) +} + +func TestBuildCommitStepMessages_ToolTimestampsAndMCPConfigIDs(t *testing.T) { + t.Parallel() + + callAt := time.Date(2026, 2, 3, 4, 5, 6, 0, time.UTC) + resultAt := callAt.Add(3 * time.Second) + configID := uuid.New() + got, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + toolNameToConfigID: map[string]uuid.UUID{ + "mcp_tool": configID, + }, + step: stepData{Content: []fantasy.Content{ + fantasy.ToolCallContent{ToolCallID: "call-1", ToolName: "mcp_tool", Input: `{}`}, + fantasy.ToolResultContent{ToolCallID: "call-1", ToolName: "mcp_tool", Result: fantasy.ToolResultOutputContentText{Text: `{"ok":true}`}}, + }, ToolCallCreatedAt: map[string]time.Time{ + "call-1": callAt, + }, ToolResultCreatedAt: map[string]time.Time{ + "call-1": resultAt, + }}, + }) + require.NoError(t, err) + require.Len(t, got.Messages, 2) + callPart := parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content)[0] + resultPart := parseMessageParts(t, got.Messages[1].Role, got.Messages[1].Content)[0] + require.Equal(t, uuid.NullUUID{UUID: configID, Valid: true}, callPart.MCPServerConfigID) + require.Equal(t, callAt, requireNotNilTime(t, callPart.CreatedAt)) + require.Equal(t, uuid.NullUUID{UUID: configID, Valid: true}, resultPart.MCPServerConfigID) + require.Equal(t, resultAt, requireNotNilTime(t, resultPart.CreatedAt)) +} + +func TestBuildCompactionMessages_CompressedSummaryToolCallAndResult(t *testing.T) { + t.Parallel() + + modelConfigID := uuid.New() + got, err := buildCompactionMessages(buildCompactionMessagesInput{ + modelConfigID: modelConfigID, + contentVersion: chatprompt.CurrentContentVersion, + toolCallID: "summary-1", + toolName: "chat_summarized", + compaction: compactionOutcome{ + SystemSummary: "system summary", + SummaryReport: "user report", + ThresholdPercent: 70, + UsagePercent: 81.5, + ContextTokens: 815, + ContextLimit: 1000, + }, + }) + require.NoError(t, err) + require.Equal(t, 1, got.HiddenCount) + require.Len(t, got.Messages, 3) + + require.Equal(t, database.ChatMessageRoleUser, got.Messages[0].Role) + require.Equal(t, database.ChatMessageVisibilityModel, got.Messages[0].Visibility) + require.True(t, got.Messages[0].Compressed) + require.Equal(t, uuid.NullUUID{UUID: modelConfigID, Valid: true}, got.Messages[0].ModelConfigID) + require.Equal(t, "system summary", parseMessageParts(t, got.Messages[0].Role, got.Messages[0].Content)[0].Text) + + require.Equal(t, database.ChatMessageRoleAssistant, got.Messages[1].Role) + require.Equal(t, database.ChatMessageVisibilityUser, got.Messages[1].Visibility) + require.True(t, got.Messages[1].Compressed) + callPart := parseMessageParts(t, got.Messages[1].Role, got.Messages[1].Content)[0] + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, callPart.Type) + require.Equal(t, "summary-1", callPart.ToolCallID) + require.JSONEq(t, `{"source":"automatic","threshold_percent":70}`, string(callPart.Args)) + + require.Equal(t, database.ChatMessageRoleTool, got.Messages[2].Role) + require.Equal(t, database.ChatMessageVisibilityBoth, got.Messages[2].Visibility) + require.True(t, got.Messages[2].Compressed) + resultPart := parseMessageParts(t, got.Messages[2].Role, got.Messages[2].Content)[0] + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, resultPart.Type) + require.Equal(t, "summary-1", resultPart.ToolCallID) + require.JSONEq(t, `{"summary":"user report","source":"automatic","threshold_percent":70,"usage_percent":81.5,"context_tokens":815,"context_limit_tokens":1000}`, string(resultPart.Result)) +} + +func TestCurrentTurnStepCount_ExcludesCompressedCompactionMessages(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("start")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("first")), + dbMessage(t, 3, database.ChatMessageRoleUser, true, codersdk.ChatMessageText("compressed summary")), + dbMessage(t, 4, database.ChatMessageRoleAssistant, true, codersdk.ChatMessageToolCall("summary", "chat_summarized", nil)), + dbMessage(t, 5, database.ChatMessageRoleTool, true, codersdk.ChatMessageToolResult("summary", "chat_summarized", json.RawMessage(`{}`), false, false)), + dbMessage(t, 6, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("second")), + } + got := currentTurnStepCount(messages) + require.Equal(t, 2, got) +} + +func TestCurrentTurnStepCount_CountsAssistantMessagesAfterLatestUser(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("old")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("old answer")), + dbMessage(t, 3, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("new")), + dbMessage(t, 4, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("one")), + dbMessage(t, 5, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("call", "tool", json.RawMessage(`{}`), false, false)), + dbMessage(t, 6, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("two")), + } + got := currentTurnStepCount(messages) + require.Equal(t, 2, got) +} + +func TestDecisionDetectsStopAfterToolFromCommittedHistory(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("plan")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("plan-1", "propose_plan", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("plan-1", "propose_plan", json.RawMessage(`{"ok":true}`), false, false)), + } + got, err := historyHasStopAfterToolResult(messages, map[string]struct{}{"propose_plan": {}}) + require.NoError(t, err) + require.True(t, got) + + messages[2] = dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("plan-1", "propose_plan", json.RawMessage(`{"error":"no"}`), true, false)) + got, err = historyHasStopAfterToolResult(messages, map[string]struct{}{"propose_plan": {}}) + require.NoError(t, err) + require.False(t, got) +} + +func TestDecisionDetectsCurrentHistoryCompletion(t *testing.T) { + t.Parallel() + + complete, err := currentHistoryComplete([]database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hello")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("done")), + }) + require.NoError(t, err) + require.True(t, complete) + + complete, err = currentHistoryComplete([]database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hello")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "execute", json.RawMessage(`{}`))), + }) + require.NoError(t, err) + require.False(t, complete) + + complete, err = currentHistoryComplete([]database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hello")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "execute", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("call-1", "execute", json.RawMessage(`{"ok":true}`), false, false)), + }) + require.NoError(t, err) + require.False(t, complete) +} + +func TestBufferedPartsToPartialMessages_NormalizesToolCallDeltasBeforeFinal(t *testing.T) { + t.Parallel() + + createdAt := time.Date(2026, 3, 4, 5, 6, 7, 0, time.UTC) + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageText("partial ")}, + {Seq: 2, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: "call-1", ToolName: "execute", ArgsDelta: `{"cmd":`}}, + {Seq: 3, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: "call-1", ToolName: "execute", ArgsDelta: `"ignored"}`}}, + {Seq: 4, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageToolCall("call-1", "execute", json.RawMessage(`{"cmd":"pwd"}`))}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + interruptedAt: createdAt, + }) + require.NoError(t, err) + require.Len(t, got, 2) + assistantParts := parseMessageParts(t, got[0].Role, got[0].Content) + require.Len(t, assistantParts, 2) + require.Equal(t, codersdk.ChatMessagePartTypeText, assistantParts[0].Type) + call := assistantParts[1] + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, call.Type) + require.Equal(t, "call-1", call.ToolCallID) + require.Empty(t, call.ArgsDelta) + require.JSONEq(t, `{"cmd":"pwd"}`, string(call.Args)) + syntheticParts := parseMessageParts(t, got[1].Role, got[1].Content) + require.Len(t, syntheticParts, 1) + require.Equal(t, "call-1", syntheticParts[0].ToolCallID) +} + +func TestBufferedPartsToPartialMessages_MergesToolCallDeltasWithoutFinal(t *testing.T) { + t.Parallel() + + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: "call-1", ToolName: "execute", ArgsDelta: `{"cmd":`}}, + {Seq: 2, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: "call-1", ToolName: "execute", ArgsDelta: `"pwd"}`}}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + }) + require.NoError(t, err) + require.Len(t, got, 2) + assistantParts := parseMessageParts(t, got[0].Role, got[0].Content) + require.Len(t, assistantParts, 1) + require.Empty(t, assistantParts[0].ArgsDelta) + require.JSONEq(t, `{"cmd":"pwd"}`, string(assistantParts[0].Args)) + syntheticParts := parseMessageParts(t, got[1].Role, got[1].Content) + require.Len(t, syntheticParts, 1) + require.Equal(t, "call-1", syntheticParts[0].ToolCallID) +} + +func TestBufferedPartsToPartialMessages_DeltaOnlyToolResultDoesNotAnswer(t *testing.T) { + t.Parallel() + + logSink := &partialConversionLogSink{} + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink) + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageToolCall("call-1", "advisor", json.RawMessage(`{}`))}, + {Seq: 2, Role: codersdk.ChatMessageRoleTool, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolResult, ToolCallID: "call-1", ToolName: "advisor", ResultDelta: `{"type":"advice"}`}}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: logger, + }) + require.NoError(t, err) + require.Len(t, got, 2) + toolParts := parseMessageParts(t, got[1].Role, got[1].Content) + require.Len(t, toolParts, 1) + require.Equal(t, "call-1", toolParts[0].ToolCallID) + require.True(t, toolParts[0].IsError) + require.Empty(t, toolParts[0].ResultDelta) + require.JSONEq(t, `{"error":"tool call was interrupted before it produced a result"}`, string(toolParts[0].Result)) + require.NotEmpty(t, logSink.entriesAtLevelWithMessage(slog.LevelWarn, "skipping buffered chat message part")) +} + +func TestBufferedPartsToPartialMessages_LogsMalformedSkippedParts(t *testing.T) { + t.Parallel() + + logSink := &partialConversionLogSink{} + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink) + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleSystem, MessagePart: codersdk.ChatMessageText("bad role")}, + {Seq: 2, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{}}, + {Seq: 3, Role: codersdk.ChatMessageRoleTool, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolResult, ToolName: "execute", Result: json.RawMessage(`{"ok":true}`)}}, + {Seq: 4, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessagePart{Type: codersdk.ChatMessagePartTypeToolCall, ToolCallID: "bad-args", ToolName: "execute", ArgsDelta: `{"cmd":`}}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: uuid.New(), + contentVersion: chatprompt.CurrentContentVersion, + logger: logger, + }) + require.NoError(t, err) + require.Empty(t, got) + require.GreaterOrEqual(t, len(logSink.entriesAtLevelWithMessage(slog.LevelWarn, "skipping buffered chat message part")), 4) +} + +func TestBufferedPartsToPartialMessages_SynthesizesMissingToolResults(t *testing.T) { + t.Parallel() + + modelConfigID := uuid.New() + createdAt := time.Date(2026, 3, 4, 5, 6, 7, 0, time.UTC) + reasoningStartedAt := createdAt.Add(-2 * time.Second) + reasoningPart := codersdk.ChatMessageReasoning("partial thought") + reasoningPart.CreatedAt = &reasoningStartedAt + parts := []messagepartbuffer.Part{ + {Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageText("partial ")}, + {Seq: 2, Role: codersdk.ChatMessageRoleAssistant, MessagePart: reasoningPart}, + {Seq: 3, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageToolCall("call-1", "execute", json.RawMessage(`{}`))}, + {Seq: 4, Role: codersdk.ChatMessageRoleAssistant, MessagePart: codersdk.ChatMessageToolCall("call-2", "read_file", json.RawMessage(`{}`))}, + {Seq: 5, Role: codersdk.ChatMessageRoleTool, MessagePart: withCreatedAt(codersdk.ChatMessageToolResult("call-2", "read_file", json.RawMessage(`{"ok":true}`), false, false), createdAt)}, + } + got, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: modelConfigID, + contentVersion: chatprompt.CurrentContentVersion, + logger: slog.Make(), + interruptedAt: createdAt, + }) + require.NoError(t, err) + require.Len(t, got, 3) + require.Equal(t, database.ChatMessageRoleAssistant, got[0].Role) + assistantParts := parseMessageParts(t, got[0].Role, got[0].Content) + require.Len(t, assistantParts, 4) + require.Equal(t, codersdk.ChatMessagePartTypeReasoning, assistantParts[1].Type) + require.Equal(t, "partial thought", assistantParts[1].Text) + require.Equal(t, reasoningStartedAt, requireNotNilTime(t, assistantParts[1].CreatedAt)) + require.Equal(t, createdAt, requireNotNilTime(t, assistantParts[1].CompletedAt)) + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, assistantParts[2].Type) + require.Equal(t, codersdk.ChatMessagePartTypeToolCall, assistantParts[3].Type) + + require.Equal(t, database.ChatMessageRoleTool, got[1].Role) + toolParts := parseMessageParts(t, got[1].Role, got[1].Content) + require.Equal(t, "call-2", toolParts[0].ToolCallID) + require.Equal(t, createdAt, requireNotNilTime(t, toolParts[0].CreatedAt)) + + require.Equal(t, database.ChatMessageRoleTool, got[2].Role) + syntheticParts := parseMessageParts(t, got[2].Role, got[2].Content) + require.Len(t, syntheticParts, 1) + require.Equal(t, "call-1", syntheticParts[0].ToolCallID) + require.Equal(t, "execute", syntheticParts[0].ToolName) + require.True(t, syntheticParts[0].IsError) + require.JSONEq(t, `{"error":"tool call was interrupted before it produced a result"}`, string(syntheticParts[0].Result)) + require.Equal(t, createdAt, requireNotNilTime(t, syntheticParts[0].CreatedAt)) + require.Equal(t, uuid.NullUUID{UUID: modelConfigID, Valid: true}, got[2].ModelConfigID) +} + +func parseMessageParts(t *testing.T, role database.ChatMessageRole, raw pqtype.NullRawMessage) []codersdk.ChatMessagePart { + t.Helper() + parts, err := chatprompt.ParseContent(database.ChatMessage{ + Role: role, + Content: raw, + }) + require.NoError(t, err) + return parts +} + +func dbMessage(t *testing.T, id int64, role database.ChatMessageRole, compressed bool, parts ...codersdk.ChatMessagePart) database.ChatMessage { + t.Helper() + raw, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return database.ChatMessage{ + ID: id, + Role: role, + Content: raw, + ContentVersion: chatprompt.CurrentContentVersion, + Visibility: database.ChatMessageVisibilityBoth, + Compressed: compressed, + } +} + +func requireNotNilTime(t *testing.T, value *time.Time) time.Time { + t.Helper() + require.NotNil(t, value) + return *value +} + +func withCreatedAt(part codersdk.ChatMessagePart, createdAt time.Time) codersdk.ChatMessagePart { + part.CreatedAt = &createdAt + return part +} + +type partialConversionLogSink struct { + mu sync.Mutex + entries []slog.SinkEntry +} + +func (s *partialConversionLogSink) LogEntry(_ context.Context, entry slog.SinkEntry) { + s.mu.Lock() + defer s.mu.Unlock() + s.entries = append(s.entries, entry) +} + +func (*partialConversionLogSink) Sync() {} + +func (s *partialConversionLogSink) entriesAtLevelWithMessage(level slog.Level, message string) []slog.SinkEntry { + s.mu.Lock() + defer s.mu.Unlock() + + entries := make([]slog.SinkEntry, 0, len(s.entries)) + for _, entry := range s.entries { + if entry.Level == level && entry.Message == message { + entries = append(entries, entry) + } + } + return entries +} diff --git a/coderd/x/chatd/messagepartbuffer/export_test.go b/coderd/x/chatd/messagepartbuffer/export_test.go new file mode 100644 index 0000000000..0f157f5743 --- /dev/null +++ b/coderd/x/chatd/messagepartbuffer/export_test.go @@ -0,0 +1,9 @@ +package messagepartbuffer + +// EpisodeCount returns the number of tracked episodes so tests can assert +// that episode state is reclaimed and does not leak. +func (b *Buffer) EpisodeCount() int { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.episodes) +} diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go new file mode 100644 index 0000000000..a41c91c293 --- /dev/null +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer.go @@ -0,0 +1,494 @@ +package messagepartbuffer + +import ( + "container/heap" + "context" + "encoding/json" + "sync" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk" + "github.com/coder/quartz" +) + +const ( + defaultMaxEpisodeBytes = int64(1024 * 1024) + defaultClosedEpisodeRetention = 15 * time.Second + defaultSubscriberSendTimeout = 10 * time.Second + defaultSubscriberChannelSize = 16 +) + +var ( + // ErrEpisodeExists means the episode already exists. + ErrEpisodeExists = xerrors.New("message part episode already exists") + // ErrEpisodeNotFound means the episode has not been created. + ErrEpisodeNotFound = xerrors.New("message part episode not found") + // ErrEpisodeClosed means the episode no longer accepts parts. + ErrEpisodeClosed = xerrors.New("message part episode closed") + // ErrEpisodeFull means the episode byte limit would be exceeded. + ErrEpisodeFull = xerrors.New("message part episode full") + // ErrMessagePartBufferClosed means the whole buffer is closed. + ErrMessagePartBufferClosed = xerrors.New("message part buffer closed") +) + +// Key identifies a buffered message part episode. +type Key struct { + ChatID uuid.UUID + HistoryVersion int64 + GenerationAttempt int64 +} + +// Part is a buffered chat message part with its sequence number. +type Part struct { + Seq int64 + Role codersdk.ChatMessageRole + MessagePart codersdk.ChatMessagePart +} + +type partJSON struct { + Seq int64 `json:"seq"` + Role codersdk.ChatMessageRole `json:"role"` + Part codersdk.ChatMessagePart `json:"part"` +} + +func (p Part) jsonValue() partJSON { + return partJSON{ + Seq: p.Seq, + Role: p.Role, + Part: p.MessagePart, + } +} + +// Options configures a Buffer. +type Options struct { + MaxEpisodeBytes int64 + ClosedEpisodeRetention time.Duration + SubscriberSendTimeout time.Duration + SubscriberChannelSize int + Clock quartz.Clock +} + +// Buffer stores streamed message parts by episode. +type Buffer struct { + mu sync.Mutex + opts Options + episodes map[Key]*episodeState + closedEpisodes closedEpisodeHeap + closed bool + done chan struct{} +} + +type episodeState struct { + created bool + closed bool + closedAt time.Time + closedHeapItem *closedEpisodeItem + parts []Part + bytes int64 + subscribers map[*episodeSubscriber]struct{} +} + +type closedEpisodeItem struct { + key Key + closedAt time.Time +} + +type closedEpisodeHeap []*closedEpisodeItem + +func (h closedEpisodeHeap) Len() int { + return len(h) +} + +func (h closedEpisodeHeap) Less(i, j int) bool { + return h[i].closedAt.Before(h[j].closedAt) +} + +func (h closedEpisodeHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] +} + +func (h *closedEpisodeHeap) Push(value any) { + item, ok := value.(*closedEpisodeItem) + if !ok { + // The reason we panic here instead of returning an error is that + // closedEpisodeHeap implements the https://pkg.go.dev/container/heap interface. + // We must accept an any type and we must not return an error. + panic("closed episode heap received invalid item") + } + *h = append(*h, item) +} + +func (h *closedEpisodeHeap) Pop() any { + old := *h + last := old[len(old)-1] + old[len(old)-1] = nil + *h = old[:len(old)-1] + return last +} + +type episodeSubscriber struct { + out chan Part + notifyCh chan struct{} + stopCh chan struct{} + next int + stopOnce sync.Once +} + +// New returns a message part buffer. +func New(options Options) *Buffer { + if options.MaxEpisodeBytes <= 0 { + options.MaxEpisodeBytes = defaultMaxEpisodeBytes + } + if options.ClosedEpisodeRetention <= 0 { + options.ClosedEpisodeRetention = defaultClosedEpisodeRetention + } + if options.SubscriberSendTimeout <= 0 { + options.SubscriberSendTimeout = defaultSubscriberSendTimeout + } + if options.SubscriberChannelSize <= 0 { + options.SubscriberChannelSize = defaultSubscriberChannelSize + } + if options.Clock == nil { + options.Clock = quartz.NewReal() + } + buffer := &Buffer{ + opts: options, + episodes: make(map[Key]*episodeState), + done: make(chan struct{}), + } + buffer.startCleanupLoop() + return buffer +} + +// CreateEpisode creates a new episode. +func (b *Buffer) CreateEpisode(key Key) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "create")) + episode := b.episodeLocked(key) + if episode.created { + return ErrEpisodeExists + } + episode.created = true + return nil +} + +// AddPart appends a part to an existing episode. +func (b *Buffer) AddPart(key Key, role codersdk.ChatMessageRole, part codersdk.ChatMessagePart) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode := b.episodes[key] + if episode == nil || !episode.created { + return ErrEpisodeNotFound + } + if episode.closed { + return ErrEpisodeClosed + } + buffered := Part{ + Seq: int64(len(episode.parts) + 1), + Role: role, + MessagePart: part, + } + sizeBytes, err := serializedPartBytes(buffered) + if err != nil { + return err + } + if episode.bytes+sizeBytes > b.opts.MaxEpisodeBytes { + return ErrEpisodeFull + } + episode.parts = append(episode.parts, buffered) + episode.bytes += sizeBytes + for subscriber := range episode.subscribers { + notifySubscriber(subscriber) + } + return nil +} + +// CloseEpisode marks an episode closed and closes its subscribers. +func (b *Buffer) CloseEpisode(key Key) error { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return ErrMessagePartBufferClosed + } + episode := b.episodeLocked(key) + episode.created = true + if episode.closed { + return nil + } + episode.closed = true + episode.closedAt = b.opts.Clock.Now("message-part-buffer", "close") + b.queueClosedEpisodeLocked(key, episode) + for subscriber := range episode.subscribers { + notifySubscriber(subscriber) + } + return nil +} + +// GetParts returns a snapshot of buffered parts for an episode. +func (b *Buffer) GetParts(key Key) ([]Part, error) { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return nil, ErrMessagePartBufferClosed + } + b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "get")) + episode := b.episodes[key] + if episode == nil || !episode.created { + return nil, ErrEpisodeNotFound + } + return append([]Part(nil), episode.parts...), nil +} + +// SubscribeToEpisode replays existing parts and streams new parts. +func (b *Buffer) SubscribeToEpisode(ctx context.Context, key Key) (<-chan Part, func(), error) { + b.mu.Lock() + if b.closed { + b.mu.Unlock() + return nil, nil, ErrMessagePartBufferClosed + } + episode := b.episodeLocked(key) + subscriber := &episodeSubscriber{ + out: make(chan Part), + notifyCh: make(chan struct{}, 1), + stopCh: make(chan struct{}), + } + if episode.subscribers == nil { + episode.subscribers = make(map[*episodeSubscriber]struct{}) + } + episode.subscribers[subscriber] = struct{}{} + notifySubscriber(subscriber) + b.mu.Unlock() + + go b.deliverSubscriber(ctx, key, subscriber) + cancel := func() { + b.cancelSubscriber(key, subscriber) + } + return subscriber.out, cancel, nil +} + +// Close closes the buffer and all pending subscriptions. +func (b *Buffer) Close() { + b.mu.Lock() + if b.closed { + b.mu.Unlock() + return + } + b.closed = true + close(b.done) + for _, episode := range b.episodes { + for subscriber := range episode.subscribers { + b.stopSubscriberLocked(episode, subscriber) + } + } + b.mu.Unlock() +} + +func (b *Buffer) startCleanupLoop() { + ticker := b.opts.Clock.NewTicker(b.opts.ClosedEpisodeRetention, "message-part-buffer", "cleanup") + go func() { + defer ticker.Stop() + for { + select { + case <-ticker.C: + b.mu.Lock() + if b.closed { + b.mu.Unlock() + return + } + b.gcClosedEpisodesLocked(b.opts.Clock.Now("message-part-buffer", "cleanup")) + b.mu.Unlock() + case <-b.done: + return + } + } + }() +} + +func (b *Buffer) gcClosedEpisodesLocked(now time.Time) { + cutoff := now.Add(-b.opts.ClosedEpisodeRetention) + type retainedEpisode struct { + key Key + episode *episodeState + } + retained := make([]retainedEpisode, 0) + for b.closedEpisodes.Len() > 0 { + item := b.closedEpisodes[0] + if item.closedAt.After(cutoff) { + break + } + popped, ok := heap.Pop(&b.closedEpisodes).(*closedEpisodeItem) + if !ok || popped != item { + continue + } + episode := b.episodes[item.key] + if episode == nil || episode.closedHeapItem != item || !episode.closed { + continue + } + episode.closedHeapItem = nil + if len(episode.subscribers) > 0 { + retained = append(retained, retainedEpisode{key: item.key, episode: episode}) + continue + } + delete(b.episodes, item.key) + } + for _, item := range retained { + if b.episodes[item.key] != item.episode || !item.episode.closed || item.episode.closedHeapItem != nil { + continue + } + b.queueClosedEpisodeLocked(item.key, item.episode) + } +} + +func (b *Buffer) queueClosedEpisodeLocked(key Key, episode *episodeState) { + if episode.closedHeapItem != nil { + return + } + item := &closedEpisodeItem{key: key, closedAt: episode.closedAt} + episode.closedHeapItem = item + heap.Push(&b.closedEpisodes, item) +} + +func (b *Buffer) episodeLocked(key Key) *episodeState { + episode := b.episodes[key] + if episode != nil { + return episode + } + episode = &episodeState{} + b.episodes[key] = episode + return episode +} + +func (b *Buffer) subscriberParts(key Key, subscriber *episodeSubscriber) (parts []Part, closed bool, ok bool) { + b.mu.Lock() + defer b.mu.Unlock() + if b.closed { + return nil, false, false + } + episode := b.episodes[key] + if episode == nil { + return nil, false, false + } + if !episode.created { + return nil, false, true + } + if subscriber.next > len(episode.parts) { + return nil, false, false + } + parts = append([]Part(nil), episode.parts[subscriber.next:]...) + subscriber.next = len(episode.parts) + return parts, episode.closed && subscriber.next == len(episode.parts), true +} + +func (b *Buffer) deliverSubscriber(ctx context.Context, key Key, subscriber *episodeSubscriber) { + defer close(subscriber.out) + defer b.removeSubscriber(key, subscriber) + for { + parts, closed, ok := b.subscriberParts(key, subscriber) + if !ok { + return + } + for _, part := range parts { + if !b.sendSubscriberPart(ctx, subscriber, part) { + return + } + } + if closed { + return + } + select { + case <-subscriber.notifyCh: + case <-subscriber.stopCh: + return + case <-ctx.Done(): + return + case <-b.done: + return + } + } +} + +func (b *Buffer) sendSubscriberPart(ctx context.Context, subscriber *episodeSubscriber, part Part) bool { + timer := b.opts.Clock.NewTimer(b.opts.SubscriberSendTimeout, "message-part-buffer", "subscriber-send") + defer timer.Stop() + select { + case subscriber.out <- part: + return true + case <-timer.C: + return false + case <-subscriber.stopCh: + return false + case <-ctx.Done(): + return false + case <-b.done: + return false + } +} + +func (b *Buffer) cancelSubscriber(key Key, subscriber *episodeSubscriber) { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode != nil { + b.stopSubscriberLocked(episode, subscriber) + return + } + subscriber.stop() +} + +func (b *Buffer) removeSubscriber(key Key, subscriber *episodeSubscriber) { + b.mu.Lock() + defer b.mu.Unlock() + episode := b.episodes[key] + if episode == nil { + return + } + delete(episode.subscribers, subscriber) + if len(episode.subscribers) != 0 { + return + } + switch { + case episode.closed: + b.queueClosedEpisodeLocked(key, episode) + case !episode.created: + // SubscribeToEpisode inserts a placeholder state for unknown keys so + // that CreateEpisode can adopt subscribers that arrive early. Once the + // last subscriber leaves a still-uncreated episode, no CreateEpisode or + // CloseEpisode call will ever reclaim it, so delete it here to avoid + // leaking the map entry for the lifetime of the buffer. + delete(b.episodes, key) + } +} + +func (*Buffer) stopSubscriberLocked(episode *episodeState, subscriber *episodeSubscriber) { + delete(episode.subscribers, subscriber) + subscriber.stop() +} + +func notifySubscriber(subscriber *episodeSubscriber) { + select { + case subscriber.notifyCh <- struct{}{}: + default: + } +} + +func (s *episodeSubscriber) stop() { + s.stopOnce.Do(func() { close(s.stopCh) }) +} + +func serializedPartBytes(part Part) (int64, error) { + data, err := json.Marshal(part.jsonValue()) + if err != nil { + return 0, err + } + return int64(len(data)), nil +} diff --git a/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go new file mode 100644 index 0000000000..f1fcf300b2 --- /dev/null +++ b/coderd/x/chatd/messagepartbuffer/message_part_buffer_test.go @@ -0,0 +1,395 @@ +package messagepartbuffer_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestBuffer_CreateEpisodeRejectsDuplicate(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.ErrorIs(t, buffer.CreateEpisode(key), messagepartbuffer.ErrEpisodeExists) +} + +func TestBuffer_AddPartAndGetParts(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello"))) + + parts, err := buffer.GetParts(key) + require.NoError(t, err) + require.Len(t, parts, 1) + require.Equal(t, int64(1), parts[0].Seq) + require.Equal(t, codersdk.ChatMessageRoleAssistant, parts[0].Role) + require.Equal(t, codersdk.ChatMessageText("hello"), parts[0].MessagePart) +} + +func TestBuffer_AddPartMissingEpisodeReturnsNotFound(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + err := buffer.AddPart(testEpisodeKey(), codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello")) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) +} + +func TestBuffer_GetPartsMissingEpisodeReturnsNotFound(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + _, err := buffer.GetParts(testEpisodeKey()) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) +} + +func TestBuffer_AddPartFullEpisodeReturnsFull(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{MaxEpisodeBytes: 1}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + err := buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("hello")) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeFull) + parts, getErr := buffer.GetParts(key) + require.NoError(t, getErr) + require.Empty(t, parts) +} + +func TestBuffer_CloseEpisodeMissingCreatesClosedEpisode(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CloseEpisode(key)) + parts, err := buffer.GetParts(key) + require.NoError(t, err) + require.Empty(t, parts) + err = buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("tail")) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeClosed) +} + +func TestBuffer_CloseEpisodeIdempotent(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.CloseEpisode(key)) + require.NoError(t, buffer.CloseEpisode(key)) +} + +func TestBuffer_SubscribeExistingReplaysThenStreamsLiveParts(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("before"))) + + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + require.Equal(t, "before", receivePart(t, ch).MessagePart.Text) + + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("after"))) + require.Equal(t, "after", receivePart(t, ch).MessagePart.Text) +} + +func TestBuffer_SubscribeClosedEpisodeReplaysThenCloses(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("before"))) + require.NoError(t, buffer.CloseEpisode(key)) + + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + require.Equal(t, "before", receivePart(t, ch).MessagePart.Text) + assertChannelClosed(t, ch) +} + +func TestBuffer_SubscribeBeforeCreateReturnsAndWaitsWithoutNotFound(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + + select { + case part := <-ch: + t.Fatalf("received part before episode create: %+v", part) + default: + } + + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("live"))) + require.Equal(t, "live", receivePart(t, ch).MessagePart.Text) +} + +func TestBuffer_AddPartAssignsContiguousSeq(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + for i := range 3 { + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText(string(rune('a'+i))))) + } + parts, err := buffer.GetParts(key) + require.NoError(t, err) + require.Equal(t, []int64{1, 2, 3}, []int64{parts[0].Seq, parts[1].Seq, parts[2].Seq}) +} + +func TestBuffer_EpisodeByteLimitUsesJSONAccounting(t *testing.T) { + t.Parallel() + + part := codersdk.ChatMessageText("hello") + limit := serializedPartBytes(t, messagepartbuffer.Part{Seq: 1, Role: codersdk.ChatMessageRoleAssistant, MessagePart: part}) + buffer := messagepartbuffer.New(messagepartbuffer.Options{MaxEpisodeBytes: limit}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, part)) + err := buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("too much")) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeFull) +} + +func TestBuffer_GCClosedEpisodeAfterGraceAndNoSubscribers(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send") + defer trap.Close() + buffer := messagepartbuffer.New(messagepartbuffer.Options{ + Clock: clock, + ClosedEpisodeRetention: time.Minute, + SubscriberSendTimeout: 10 * time.Minute, + }) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("held"))) + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + require.NoError(t, buffer.CloseEpisode(key)) + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(time.Minute).MustWait(ctx) + clock.Advance(time.Second).MustWait(ctx) + _, err = buffer.GetParts(key) + require.NoError(t, err) + + cancel() + drainUntilClosed(t, ch) + _, err = buffer.GetParts(key) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) +} + +func TestBuffer_GCRetainedSubscribedEpisodeDoesNotBlockOtherExpiredEpisodes(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send") + defer trap.Close() + buffer := messagepartbuffer.New(messagepartbuffer.Options{ + Clock: clock, + ClosedEpisodeRetention: time.Minute, + SubscriberSendTimeout: 10 * time.Minute, + }) + retainedKey := testEpisodeKey() + collectedKey := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(retainedKey)) + require.NoError(t, buffer.AddPart(retainedKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("held"))) + require.NoError(t, buffer.CreateEpisode(collectedKey)) + require.NoError(t, buffer.AddPart(collectedKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("collect me"))) + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, retainedKey) + require.NoError(t, err) + defer cancel() + require.NoError(t, buffer.CloseEpisode(retainedKey)) + require.NoError(t, buffer.CloseEpisode(collectedKey)) + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(time.Minute).MustWait(ctx) + clock.Advance(time.Second).MustWait(ctx) + + _, err = buffer.GetParts(retainedKey) + require.NoError(t, err) + _, err = buffer.GetParts(collectedKey) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) + + cancel() + drainUntilClosed(t, ch) + _, err = buffer.GetParts(retainedKey) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) +} + +func TestBuffer_SlowSubscriberClosed(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + trap := clock.Trap().NewTimer("message-part-buffer", "subscriber-send") + defer trap.Close() + stopTrap := clock.Trap().TimerStop() + defer stopTrap.Close() + buffer := messagepartbuffer.New(messagepartbuffer.Options{ + Clock: clock, + SubscriberSendTimeout: time.Second, + }) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("blocked"))) + call := trap.MustWait(ctx) + call.MustRelease(ctx) + clock.Advance(time.Second).MustWait(ctx) + stopCall := stopTrap.MustWait(ctx) + stopCall.MustRelease(ctx) + assertChannelClosed(t, ch) +} + +func TestBuffer_BurstyOutputDoesNotCloseSubscriberBeforeSendTimeout(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{SubscriberChannelSize: 1}) + key := testEpisodeKey() + require.NoError(t, buffer.CreateEpisode(key)) + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + + for i := range 8 { + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText(string(rune('a'+i))))) + } + for i := range 8 { + part := receivePart(t, ch) + require.Equal(t, string(rune('a'+i)), part.MessagePart.Text) + } +} + +func TestBuffer_SubscribeCanceledBeforeCreateCanCreateEpisode(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + ctx, cancel := context.WithCancel(context.Background()) + ch, cancelSub, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + cancel() + drainUntilClosed(t, ch) + cancelSub() + require.NoError(t, buffer.CreateEpisode(key)) +} + +func TestBuffer_SubscribeCanceledWithoutCreateReclaimsEpisode(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancelSub, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + cancelSub() + // The subscriber goroutine removes itself from the episode before closing + // the output channel, so cleanup is complete once the channel is closed. + drainUntilClosed(t, ch) + + _, err = buffer.GetParts(key) + require.ErrorIs(t, err, messagepartbuffer.ErrEpisodeNotFound) + require.Equal(t, 0, buffer.EpisodeCount()) +} + +func TestBuffer_CloseClosesPendingSubscriptionAndRejectsOperations(t *testing.T) { + t.Parallel() + + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := testEpisodeKey() + ctx := testutil.Context(t, testutil.WaitLong) + ch, cancel, err := buffer.SubscribeToEpisode(ctx, key) + require.NoError(t, err) + defer cancel() + buffer.Close() + assertChannelClosed(t, ch) + require.ErrorIs(t, buffer.CreateEpisode(key), messagepartbuffer.ErrMessagePartBufferClosed) +} + +func testEpisodeKey() messagepartbuffer.Key { + return messagepartbuffer.Key{ChatID: uuid.New(), HistoryVersion: 1, GenerationAttempt: 1} +} + +func receivePart(t *testing.T, ch <-chan messagepartbuffer.Part) messagepartbuffer.Part { + t.Helper() + select { + case part, ok := <-ch: + require.True(t, ok) + return part + case <-time.After(testutil.WaitLong): + t.Fatal("timed out waiting for buffered part") + return messagepartbuffer.Part{} + } +} + +func assertChannelClosed[T any](t *testing.T, ch <-chan T) { + t.Helper() + select { + case _, ok := <-ch: + require.False(t, ok) + case <-time.After(testutil.WaitLong): + t.Fatal("timed out waiting for channel close") + } +} + +func drainUntilClosed[T any](t *testing.T, ch <-chan T) { + t.Helper() + for { + select { + case _, ok := <-ch: + if !ok { + return + } + case <-time.After(testutil.WaitLong): + t.Fatal("timed out waiting for channel close") + } + } +} + +func serializedPartBytes(t *testing.T, part messagepartbuffer.Part) int64 { + t.Helper() + data, err := json.Marshal(struct { + Seq int64 `json:"seq"` + Role codersdk.ChatMessageRole `json:"role"` + Part codersdk.ChatMessagePart `json:"part"` + }{ + Seq: part.Seq, + Role: part.Role, + Part: part.MessagePart, + }) + require.NoError(t, err) + return int64(len(data)) +} diff --git a/coderd/x/chatd/model_routing.go b/coderd/x/chatd/model_routing.go index c5fa7129db..e94db9af55 100644 --- a/coderd/x/chatd/model_routing.go +++ b/coderd/x/chatd/model_routing.go @@ -8,6 +8,7 @@ import ( "github.com/google/uuid" "golang.org/x/xerrors" + "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" ) @@ -29,6 +30,17 @@ func modelBuildOptionsFromMessages(messages []database.ChatMessage) modelBuildOp return modelBuildOptions{ActiveAPIKeyID: apiKeyID} } +// withActiveTurnAPIKeyID augments ctx with the active turn's delegated API +// key ID when one is known. AI Gateway routing and subagent tool callbacks +// read this value from the context to attribute requests to the correct +// turn. When no key is known, ctx is returned unchanged. +func withActiveTurnAPIKeyID(ctx context.Context, opts modelBuildOptions) context.Context { + if opts.ActiveAPIKeyID == "" { + return ctx + } + return aibridge.WithDelegatedAPIKeyID(ctx, opts.ActiveAPIKeyID) +} + type modelRouteKind int const ( diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go new file mode 100644 index 0000000000..ff3dbdd3d9 --- /dev/null +++ b/coderd/x/chatd/options.go @@ -0,0 +1,159 @@ +package chatd + +import ( + "context" + "database/sql" + "sync/atomic" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/audit" + "github.com/coder/coder/v2/coderd/database" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" + "github.com/coder/coder/v2/coderd/notifications" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/quartz" +) + +const ( + defaultAcquisitionInterval = 30 * time.Second + defaultAcquisitionBatchSize = int32(10) + defaultRunnerSyncInterval = 15 * time.Second + defaultHeartbeatInterval = 9 * time.Second + defaultHeartbeatCleanupEvery = 30 * time.Second + defaultHeartbeatStaleSeconds = int32(30) + // The archive cutoff is based on UTC start-of-day and only moves + // once per day, so hourly runs are more than enough to keep up + // while still catching chats that cross the threshold shortly + // after midnight. + defaultArchiveInterval = time.Hour + defaultArchiveBatchSize = int32(1000) + defaultStateChannelSize = 64 + defaultTaskRetryInitialBackoff = 100 * time.Millisecond + defaultTaskRetryMaxBackoff = 5 * time.Second +) + +// chatWorkerPubsub is the chat worker pubsub dependency. +type chatWorkerPubsub interface { + Publish(event string, message []byte) error + SubscribeWithErr(event string, listener dbpubsub.ListenerWithErr) (func(), error) +} + +// chatWorkerTaskStarter starts runner-owned side-effect tasks. +type chatWorkerTaskStarter interface { + StartGeneration(context.Context, chatWorkerTaskStartInput) error + StartInterrupt(context.Context, chatWorkerTaskStartInput) error + StartRequiresActionTimeout(context.Context, chatWorkerTaskStartInput) error + StartAbandon(context.Context, chatWorkerTaskStartInput) error +} + +// chatWorkerTaskStartInput describes one runner task invocation. +type chatWorkerTaskStartInput struct { + TaskID uuid.UUID + ChatID uuid.UUID + WorkerID uuid.UUID + RunnerID uuid.UUID + HistoryVersion int64 + GenerationAttempt int64 + Status database.ChatStatus + RequiresActionDeadlineAt sql.NullTime + DebugTurn *runnerDebugTurn +} + +// chatWorkerOptions configures a chatWorker. +type chatWorkerOptions struct { + WorkerID uuid.UUID + + Store database.Store + Pubsub chatWorkerPubsub + Logger slog.Logger + Clock quartz.Clock + TaskStarter chatWorkerTaskStarter + MessagePartBuffer *messagepartbuffer.Buffer + + NotificationsEnqueuer notifications.Enqueuer + Auditor *atomic.Pointer[audit.Auditor] + AutoArchiveRecords prometheus.Counter + + AcquisitionInterval time.Duration + AcquisitionBatchSize int32 + ArchiveInterval time.Duration + ArchiveBatchSize int32 + RunnerSyncInterval time.Duration + HeartbeatInterval time.Duration + HeartbeatCleanupInterval time.Duration + HeartbeatStaleSeconds int32 + StateChannelSize int + RunnerManagerChannelSize int + AcquisitionWakeChannelSize int + TaskRetryInitialBackoff time.Duration + TaskRetryMaxBackoff time.Duration +} + +func (o chatWorkerOptions) withDefaults() (chatWorkerOptions, error) { + if o.Store == nil { + return chatWorkerOptions{}, xerrors.New("chatworker: store is required") + } + if o.Pubsub == nil { + return chatWorkerOptions{}, xerrors.New("chatworker: pubsub is required") + } + if o.TaskStarter == nil && o.MessagePartBuffer == nil { + return chatWorkerOptions{}, xerrors.New("chatworker: task starter or message part buffer is required") + } + if o.WorkerID == uuid.Nil { + return chatWorkerOptions{}, xerrors.New("chatworker: worker ID is required") + } + if o.Clock == nil { + o.Clock = quartz.NewReal() + } + if o.AcquisitionInterval <= 0 { + o.AcquisitionInterval = defaultAcquisitionInterval + } + if o.AcquisitionBatchSize <= 0 { + o.AcquisitionBatchSize = defaultAcquisitionBatchSize + } + if o.ArchiveInterval <= 0 { + o.ArchiveInterval = defaultArchiveInterval + } + if o.ArchiveBatchSize <= 0 { + o.ArchiveBatchSize = defaultArchiveBatchSize + } + if o.NotificationsEnqueuer == nil { + o.NotificationsEnqueuer = notifications.NewNoopEnqueuer() + } + if o.RunnerSyncInterval <= 0 { + o.RunnerSyncInterval = defaultRunnerSyncInterval + } + if o.HeartbeatInterval <= 0 { + o.HeartbeatInterval = defaultHeartbeatInterval + } + if o.HeartbeatCleanupInterval <= 0 { + o.HeartbeatCleanupInterval = defaultHeartbeatCleanupEvery + } + if o.HeartbeatStaleSeconds <= 0 { + o.HeartbeatStaleSeconds = defaultHeartbeatStaleSeconds + } + if o.StateChannelSize <= 0 { + o.StateChannelSize = defaultStateChannelSize + } + if o.RunnerManagerChannelSize <= 0 { + o.RunnerManagerChannelSize = defaultStateChannelSize + } + if o.AcquisitionWakeChannelSize <= 0 { + o.AcquisitionWakeChannelSize = 1 + } + if o.TaskRetryInitialBackoff <= 0 { + o.TaskRetryInitialBackoff = defaultTaskRetryInitialBackoff + } + if o.TaskRetryMaxBackoff <= 0 { + o.TaskRetryMaxBackoff = defaultTaskRetryMaxBackoff + } + if o.TaskRetryMaxBackoff < o.TaskRetryInitialBackoff { + o.TaskRetryMaxBackoff = o.TaskRetryInitialBackoff + } + return o, nil +} diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index 774e02d107..ec9dc01356 100644 --- a/coderd/x/chatd/quickgen.go +++ b/coderd/x/chatd/quickgen.go @@ -136,6 +136,66 @@ type generatedTurnStatusLabel struct { Label string `json:"label" description:"Compact 2-5 word current chat status label"` } +// GenerateChatTitleAsync fires a best-effort, automatic title-generation +// pass for a freshly created chat. It is intended to be called from the +// chat-creation endpoint right after the chat and its initial user +// message are persisted. +// +// The work runs in a tracked goroutine with a detached context so it +// neither blocks the HTTP response nor is canceled when the request +// completes. It resolves the chat's model and provider keys, then +// delegates to maybeGenerateChatTitle, which only acts on the first user +// turn (see titleInput) and is otherwise a no-op. Errors are logged and +// swallowed. +func (p *Server) GenerateChatTitleAsync(ctx context.Context, chat database.Chat) { + logger := p.logger.With( + slog.F("chat_id", chat.ID), + slog.F("owner_id", chat.OwnerID), + ) + // Snapshot the messages synchronously so the first-turn eligibility + // check (titleInput) is evaluated against creation-time state. Loading + // inside the goroutine would race the chat worker's first assistant + // reply and could skip title generation. + messages, err := p.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + if err != nil { + logger.Debug(ctx, "failed to load messages for automatic title generation", + slog.Error(err), + ) + return + } + if _, ok := titleInput(chat, messages); !ok { + return + } + // Detach from the request lifetime so title generation can finish + // even after the create response is written. + titleCtx := context.WithoutCancel(ctx) + p.inflight.Go(func() { + modelOpts := modelBuildOptionsFromMessages(messages) + titleCtx = withActiveTurnAPIKeyID(titleCtx, modelOpts) + model, modelConfig, keys, route, _, _, _, err := p.resolveChatModel(titleCtx, chat, modelOpts) + if err != nil { + logger.Debug(titleCtx, "failed to resolve model for automatic title generation", + slog.Error(err), + ) + return + } + p.maybeGenerateChatTitle( + titleCtx, + chat, + messages, + modelConfig.Provider, + modelConfig.Model, + model, + route, + keys, + modelOpts, + &generatedChatTitle{}, + logger, + p.existingDebugService(), + ) + }) +} + // maybeGenerateChatTitle generates an AI title for the chat when // appropriate (first user message, no assistant reply yet, and the // current title is either empty or still the fallback truncation). @@ -284,10 +344,6 @@ func (p *Server) maybeGenerateChatTitle( chat.Title = title generatedTitle.Store(title) p.publishChatPubsubEvent(chat, codersdk.ChatWatchEventKindTitleChange, nil) - - // AcquireChats uses SKIP LOCKED; re-wake so a wake racing this - // UPDATE's row lock does not strand a freshly-pending chat. - p.signalWake() return } diff --git a/coderd/x/chatd/recording_internal_test.go b/coderd/x/chatd/recording_internal_test.go index 24bdf3cf76..0ad40a97b9 100644 --- a/coderd/x/chatd/recording_internal_test.go +++ b/coderd/x/chatd/recording_internal_test.go @@ -1057,7 +1057,7 @@ func TestStopAndStoreRecording_UnknownPartIgnored(t *testing.T) { } // TestStopAndStoreRecording_MalformedContentType verifies that a -// response with an unparseable Content-Type returns an empty result. +// response with an unparsable Content-Type returns an empty result. func TestStopAndStoreRecording_MalformedContentType(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/runner.go b/coderd/x/chatd/runner.go new file mode 100644 index 0000000000..5498d4d769 --- /dev/null +++ b/coderd/x/chatd/runner.go @@ -0,0 +1,341 @@ +package chatd + +import ( + "context" + "errors" + "sync" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" +) + +type taskKind string + +const ( + taskKindGeneration taskKind = "generation" + taskKindInterrupt taskKind = "interrupt" + taskKindRequiresActionTimeout taskKind = "requires_action_timeout" + taskKindAbandon taskKind = "abandon" +) + +type taskInstanceID uuid.UUID + +type localWorkKey struct { + historyVersion int64 + status database.ChatStatus +} + +type taskIndexKey struct { + kind taskKind + key localWorkKey +} + +type taskRecord struct { + id taskInstanceID + kind taskKind + localKey localWorkKey + cancel context.CancelFunc + done <-chan struct{} +} + +type runner struct { + ctx context.Context + mgr *runnerManager + rec *runnerRecord + opts chatWorkerOptions + + lastSnapshotVersion int64 + hasAcceptedState bool + latestState runnerStateUpdate + + activeTaskID taskInstanceID + activeTaskSet bool + tasks map[taskInstanceID]*taskRecord + tasksByIndex map[taskIndexKey]taskInstanceID + localLocks *localLockSet + debugTurn *runnerDebugTurn +} + +func newRunner(ctx context.Context, mgr *runnerManager, rec *runnerRecord, opts chatWorkerOptions) *runner { + return &runner{ + ctx: ctx, + mgr: mgr, + rec: rec, + opts: opts, + tasks: make(map[taskInstanceID]*taskRecord), + tasksByIndex: make(map[taskIndexKey]taskInstanceID), + localLocks: newLocalLockSet(), + debugTurn: newRunnerDebugTurn(ctx, opts.Logger), + } +} + +func (r *runner) run() { + if !r.bootstrap() { + return + } + for { + select { + case state := <-r.rec.stateCh: + r.processState(state) + case <-r.ctx.Done(): + r.cancelActiveTask() + r.waitForTasks() + r.closeDebugTurn() + return + } + } +} + +func (r *runner) bootstrap() bool { + channel := coderdpubsub.ChatStateUpdateChannel(r.rec.key.ChatID) + unsubscribe, err := r.opts.Pubsub.SubscribeWithErr(channel, coderdpubsub.HandleChatStateUpdate( + func(ctx context.Context, payload coderdpubsub.ChatStateUpdateMessage, err error) { + if err != nil { + r.opts.Logger.Warn(ctx, "chatworker state update decode failed", slogError(err)) + return + } + r.mgr.RouteStateHint(ctx, stateUpdateFromPubsub(r.rec.key.ChatID, payload)) + }, + )) + if err != nil { + r.mgr.requestCleanup(r.ctx, r.rec.key) + return false + } + if !r.rec.setUnsubscribe(unsubscribe) { + return false + } + chat, err := r.opts.Store.GetChatByID(r.ctx, r.rec.key.ChatID) + if err != nil { + r.opts.Logger.Warn(r.ctx, "chatworker runner bootstrap failed", slogError(err)) + r.mgr.requestCleanup(r.ctx, r.rec.key) + return false + } + r.mgr.RouteStateHint(r.ctx, stateUpdateFromChat(chat)) + return true +} + +func stateUpdateFromPubsub(chatID uuid.UUID, payload coderdpubsub.ChatStateUpdateMessage) runnerStateUpdate { + return runnerStateUpdate{ + ChatID: chatID, + WorkerID: payload.WorkerID, + RunnerID: payload.RunnerID, + SnapshotVersion: payload.SnapshotVersion, + HistoryVersion: payload.HistoryVersion, + QueueVersion: payload.QueueVersion, + GenerationAttempt: payload.GenerationAttempt, + Status: database.ChatStatus(payload.Status), + Archived: payload.Archived, + } +} + +func (r *runner) processState(state runnerStateUpdate) { + if state.SnapshotVersion <= r.lastSnapshotVersion { + return + } + + r.removeFinishedTasks() + + if !uuidPtrEqual(state.WorkerID, r.rec.workerID) || !uuidPtrEqual(state.RunnerID, r.rec.key.RunnerID) { + r.acceptState(state) + r.mgr.requestCleanup(r.ctx, r.rec.key) + return + } + + changed := !r.hasAcceptedState || + r.latestState.HistoryVersion != state.HistoryVersion || + r.latestState.Status != state.Status || + r.latestState.Archived != state.Archived + if !changed { + r.acceptState(state) + return + } + if r.hasAcceptedState && r.activeTaskSet { + r.cancelActiveTask() + } + + r.spawnForState(state) + r.acceptState(state) +} + +func (r *runner) acceptState(state runnerStateUpdate) { + r.hasAcceptedState = true + r.latestState = state + r.lastSnapshotVersion = state.SnapshotVersion +} + +func (r *runner) spawnForState(state runnerStateUpdate) { + if state.Archived { + r.spawnTaskIfNeeded(taskKindAbandon, state) + return + } + switch state.Status { + case database.ChatStatusRunning: + r.spawnTaskIfNeeded(taskKindGeneration, state) + case database.ChatStatusInterrupting: + r.spawnTaskIfNeeded(taskKindInterrupt, state) + case database.ChatStatusRequiresAction: + r.spawnTaskIfNeeded(taskKindRequiresActionTimeout, state) + case database.ChatStatusWaiting, database.ChatStatusError: + r.spawnTaskIfNeeded(taskKindAbandon, state) + default: + r.spawnTaskIfNeeded(taskKindAbandon, state) + } +} + +func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { + key := localWorkKey{historyVersion: state.HistoryVersion, status: state.Status} + idx := taskIndexKey{kind: kind, key: key} + if r.activeTaskSet && r.tasksByIndex[idx] == r.activeTaskID { + return + } + + id := taskInstanceID(uuid.New()) + taskCtx, cancel := context.WithCancel(r.ctx) + done := make(chan struct{}) + record := &taskRecord{ + id: id, + kind: kind, + localKey: key, + cancel: cancel, + done: done, + } + r.tasks[id] = record + r.tasksByIndex[idx] = id + r.activeTaskID = id + r.activeTaskSet = true + + input := chatWorkerTaskStartInput{ + TaskID: uuid.UUID(id), + ChatID: r.rec.key.ChatID, + WorkerID: r.rec.workerID, + RunnerID: r.rec.key.RunnerID, + HistoryVersion: state.HistoryVersion, + GenerationAttempt: state.GenerationAttempt, + Status: state.Status, + RequiresActionDeadlineAt: state.RequiresActionDeadlineAt, + DebugTurn: r.debugTurn, + } + go r.runTask(taskCtx, kind, key, input, done) +} + +func (r *runner) runTask( + ctx context.Context, + kind taskKind, + key localWorkKey, + input chatWorkerTaskStartInput, + done chan<- struct{}, +) { + defer close(done) + err := runTaskWithRetry(ctx, r.opts.retryOptions(), kind, func(ctx context.Context) error { + unlock, ok := r.localLocks.acquire(ctx, key) + if !ok { + return errTaskExpectedExit + } + defer unlock() + if ctx.Err() != nil { + return errTaskExpectedExit + } + + switch kind { + case taskKindGeneration: + return r.opts.TaskStarter.StartGeneration(ctx, input) + case taskKindInterrupt: + return r.opts.TaskStarter.StartInterrupt(ctx, input) + case taskKindRequiresActionTimeout: + return r.opts.TaskStarter.StartRequiresActionTimeout(ctx, input) + case taskKindAbandon: + return r.opts.TaskStarter.StartAbandon(ctx, input) + default: + return errors.Join(errTaskExpectedExit, xerrors.Errorf("unknown task kind %q", kind)) + } + }) + if err != nil && ctx.Err() == nil { + r.opts.Logger.Warn(ctx, "chatworker task failed", slogError(err)) + } +} + +func (r *runner) cancelActiveTask() { + if !r.activeTaskSet { + return + } + id := r.activeTaskID + r.activeTaskSet = false + if record := r.tasks[id]; record != nil { + record.cancel() + } +} + +func (r *runner) waitForTasks() { + for _, record := range r.tasks { + <-record.done + } +} + +func (r *runner) closeDebugTurn() { + if r.debugTurn == nil { + return + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(r.ctx), debugFinalizeTimeout) + defer cancel() + r.debugTurn.Finalize(ctx) +} + +func (r *runner) removeFinishedTasks() { + for id, record := range r.tasks { + select { + case <-record.done: + delete(r.tasks, id) + idx := taskIndexKey{kind: record.kind, key: record.localKey} + if r.tasksByIndex[idx] == id { + delete(r.tasksByIndex, idx) + } + if r.activeTaskSet && r.activeTaskID == id { + r.activeTaskSet = false + } + default: + } + } +} + +func uuidPtrEqual(got *uuid.UUID, want uuid.UUID) bool { + return got != nil && *got == want +} + +type localLockSet struct { + mu sync.Mutex + locked map[localWorkKey]chan struct{} +} + +func newLocalLockSet() *localLockSet { + return &localLockSet{locked: make(map[localWorkKey]chan struct{})} +} + +func (l *localLockSet) acquire(ctx context.Context, key localWorkKey) (func(), bool) { + for { + l.mu.Lock() + wait, ok := l.locked[key] + if !ok { + released := make(chan struct{}) + l.locked[key] = released + l.mu.Unlock() + return func() { + l.mu.Lock() + if l.locked[key] == released { + delete(l.locked, key) + close(released) + } + l.mu.Unlock() + }, true + } + l.mu.Unlock() + + select { + case <-wait: + case <-ctx.Done(): + return nil, false + } + } +} diff --git a/coderd/x/chatd/runner_manager.go b/coderd/x/chatd/runner_manager.go new file mode 100644 index 0000000000..dc9737c8d6 --- /dev/null +++ b/coderd/x/chatd/runner_manager.go @@ -0,0 +1,530 @@ +package chatd + +import ( + "context" + "database/sql" + "encoding/json" + "sync" + "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" +) + +const shutdownCleanupTimeout = 5 * time.Second + +type runnerKey struct { + ChatID uuid.UUID + RunnerID uuid.UUID +} + +type runnerStateUpdate struct { + ChatID uuid.UUID + WorkerID *uuid.UUID + RunnerID *uuid.UUID + SnapshotVersion int64 + HistoryVersion int64 + QueueVersion int64 + GenerationAttempt int64 + Status database.ChatStatus + Archived bool + RequiresActionDeadlineAt sql.NullTime +} + +type spawnRunnerRequest struct { + ChatID uuid.UUID + WorkerID uuid.UUID + RunnerID uuid.UUID +} + +type runnerRecord struct { + key runnerKey + workerID uuid.UUID + cancel context.CancelFunc + done <-chan struct{} + stateCh chan runnerStateUpdate + + mu sync.Mutex + unsubscribe func() + cleanupStarted bool +} + +func (r *runnerRecord) setUnsubscribe(unsubscribe func()) bool { + r.mu.Lock() + if r.cleanupStarted { + r.mu.Unlock() + if unsubscribe != nil { + unsubscribe() + } + return false + } + r.unsubscribe = unsubscribe + r.mu.Unlock() + return true +} + +func (r *runnerRecord) startCleanup() { + r.mu.Lock() + if r.cleanupStarted { + r.mu.Unlock() + return + } + r.cleanupStarted = true + unsubscribe := r.unsubscribe + r.unsubscribe = nil + r.mu.Unlock() + if unsubscribe != nil { + unsubscribe() + } + r.cancel() +} + +type runnerManager struct { + server *Server + opts chatWorkerOptions + ctx context.Context + + closed bool + spawnMu sync.Mutex + + mu sync.Mutex + spawnCh chan spawnRunnerRequest + cleanupReqCh chan runnerKey + cleanupDoneCh chan runnerKey + runners map[runnerKey]*runnerRecord + runnersByChat map[uuid.UUID]map[uuid.UUID]*runnerRecord + cleaning map[runnerKey]*runnerRecord + + wg sync.WaitGroup +} + +func newRunnerManager(ctx context.Context, server *Server, opts chatWorkerOptions) *runnerManager { + return &runnerManager{ + server: server, + opts: opts, + ctx: ctx, + spawnCh: make(chan spawnRunnerRequest, opts.RunnerManagerChannelSize), + cleanupReqCh: make(chan runnerKey, opts.RunnerManagerChannelSize), + cleanupDoneCh: make(chan runnerKey, opts.RunnerManagerChannelSize), + runners: make(map[runnerKey]*runnerRecord), + runnersByChat: make(map[uuid.UUID]map[uuid.UUID]*runnerRecord), + cleaning: make(map[runnerKey]*runnerRecord), + } +} + +func (m *runnerManager) start() { + m.wg.Go(m.run) + m.wg.Go(m.databaseSyncLoop) + m.wg.Go(m.heartbeatLoop) + m.wg.Go(m.heartbeatCleanupLoop) +} + +func (m *runnerManager) wait() { + m.wg.Wait() +} + +func (m *runnerManager) idle() bool { + m.mu.Lock() + defer m.mu.Unlock() + return len(m.runners) == 0 && len(m.cleaning) == 0 +} + +func (m *runnerManager) Spawn(ctx context.Context, req spawnRunnerRequest) error { + m.spawnMu.Lock() + defer m.spawnMu.Unlock() + if m.closed { + return xerrors.New("chatworker: runner manager closed") + } + + select { + case m.spawnCh <- req: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-m.ctx.Done(): + return m.ctx.Err() + } +} + +func (m *runnerManager) requestCleanup(ctx context.Context, key runnerKey) { + select { + case m.cleanupReqCh <- key: + case <-ctx.Done(): + case <-m.ctx.Done(): + } +} + +func (m *runnerManager) RouteStateHint(ctx context.Context, state runnerStateUpdate) { + m.mu.Lock() + byRunner := m.runnersByChat[state.ChatID] + targets := make([]*runnerRecord, 0, len(byRunner)) + for _, rec := range byRunner { + targets = append(targets, rec) + } + m.mu.Unlock() + + for _, rec := range targets { + select { + case rec.stateCh <- state: + case <-rec.done: + // Only this runner exited; keep fanning out to the rest. + continue + case <-ctx.Done(): + return + case <-m.ctx.Done(): + return + default: + // stateCh is full; drop the hint for this runner. + } + } +} + +func (m *runnerManager) run() { + for { + select { + case req := <-m.spawnCh: + m.handleSpawn(req) + case key := <-m.cleanupReqCh: + m.handleCleanupRequest(key) + case key := <-m.cleanupDoneCh: + m.handleCleanupDone(key) + case <-m.ctx.Done(): + queued := m.closeAndDrainQueues() + m.cancelAll() + m.releaseOwnershipHints(queued) + return + } + } +} + +func (m *runnerManager) handleSpawn(req spawnRunnerRequest) { + key := runnerKey{ChatID: req.ChatID, RunnerID: req.RunnerID} + m.mu.Lock() + if _, ok := m.runners[key]; ok { + // A duplicate spawn for a live runner indicates a logic error + // in the sync loop. + m.opts.Logger.Error(m.ctx, "invalid spawn request: chat runner already spawned", slog.F("key", key)) + m.mu.Unlock() + return + } + if _, ok := m.cleaning[key]; ok { + // A duplicate spawn for a live runner indicates a logic error + // in the sync loop. + m.opts.Logger.Error(m.ctx, "invalid spawn request: chat runner in cleanup", slog.F("key", key)) + m.mu.Unlock() + return + } + runnerCtx, cancel := context.WithCancel(m.ctx) + done := make(chan struct{}) + rec := &runnerRecord{ + key: key, + workerID: req.WorkerID, + cancel: cancel, + done: done, + stateCh: make(chan runnerStateUpdate, m.opts.StateChannelSize), + } + m.runners[key] = rec + if m.runnersByChat[req.ChatID] == nil { + m.runnersByChat[req.ChatID] = make(map[uuid.UUID]*runnerRecord) + } + m.runnersByChat[req.ChatID][req.RunnerID] = rec + m.mu.Unlock() + + r := newRunner(runnerCtx, m, rec, m.opts) + m.wg.Go(func() { + defer close(done) + r.run() + }) +} + +func (m *runnerManager) closeAndDrainQueues() []runnerKey { + m.spawnMu.Lock() + defer m.spawnMu.Unlock() + + m.closed = true + return m.drainQueues() +} + +func (m *runnerManager) drainQueues() []runnerKey { + queued := make([]runnerKey, 0) + for { + select { + case req := <-m.spawnCh: + queued = append(queued, runnerKey{ChatID: req.ChatID, RunnerID: req.RunnerID}) + case key := <-m.cleanupReqCh: + m.handleCleanupRequest(key) + case key := <-m.cleanupDoneCh: + m.handleCleanupDone(key) + default: + return queued + } + } +} + +func (m *runnerManager) handleCleanupRequest(key runnerKey) { + m.mu.Lock() + rec, ok := m.runners[key] + if !ok { + m.mu.Unlock() + return + } + delete(m.runners, key) + if byChat := m.runnersByChat[key.ChatID]; byChat != nil { + delete(byChat, key.RunnerID) + if len(byChat) == 0 { + delete(m.runnersByChat, key.ChatID) + } + } + m.cleaning[key] = rec + m.mu.Unlock() + + rec.startCleanup() + m.registerCleanupWaiter(key, rec) +} + +func (m *runnerManager) registerCleanupWaiter(key runnerKey, rec *runnerRecord) { + m.wg.Go(func() { + <-rec.done + if m.ctx.Err() != nil { + m.mu.Lock() + delete(m.cleaning, key) + m.mu.Unlock() + return + } + select { + case m.cleanupDoneCh <- key: + case <-m.ctx.Done(): + m.mu.Lock() + delete(m.cleaning, key) + m.mu.Unlock() + } + }) +} + +func (m *runnerManager) handleCleanupDone(key runnerKey) { + m.mu.Lock() + delete(m.cleaning, key) + m.mu.Unlock() +} + +func (m *runnerManager) cancelAll() { + type cleanupTarget struct { + key runnerKey + rec *runnerRecord + } + + m.mu.Lock() + active := make([]cleanupTarget, 0, len(m.runners)) + cleaning := make([]*runnerRecord, 0, len(m.cleaning)) + for _, rec := range m.cleaning { + cleaning = append(cleaning, rec) + } + for key, rec := range m.runners { + delete(m.runners, key) + m.cleaning[key] = rec + active = append(active, cleanupTarget{key: key, rec: rec}) + } + clear(m.runnersByChat) + m.mu.Unlock() + + keys := make([]runnerKey, 0, len(cleaning)+len(active)) + for _, rec := range cleaning { + rec.startCleanup() + keys = append(keys, rec.key) + } + for _, target := range active { + target.rec.startCleanup() + m.registerCleanupWaiter(target.key, target.rec) + keys = append(keys, target.key) + } + m.releaseOwnershipHints(keys) +} + +func (m *runnerManager) releaseOwnershipHints(keys []runnerKey) { + if len(keys) == 0 { + return + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(m.ctx), shutdownCleanupTimeout) + defer cancel() + + chatIDs := make([]uuid.UUID, 0, len(keys)) + runnerIDs := make([]uuid.UUID, 0, len(keys)) + uniqueChatIDs := make(map[uuid.UUID]struct{}, len(keys)) + for _, key := range keys { + chatIDs = append(chatIDs, key.ChatID) + runnerIDs = append(runnerIDs, key.RunnerID) + uniqueChatIDs[key.ChatID] = struct{}{} + } + if _, err := m.opts.Store.BatchDeleteChatHeartbeats(ctx, database.BatchDeleteChatHeartbeatsParams{ + ChatIds: chatIDs, + RunnerIds: runnerIDs, + }); err != nil { + m.opts.Logger.Warn(ctx, "chatworker shutdown heartbeat cleanup failed", slogError(err)) + } + + syncIDs := make([]uuid.UUID, 0, len(uniqueChatIDs)) + for id := range uniqueChatIDs { + syncIDs = append(syncIDs, id) + } + chats, err := m.opts.Store.GetChatsByIDsForRunnerSync(ctx, syncIDs) + if err != nil { + m.opts.Logger.Warn(ctx, "chatworker shutdown ownership lookup failed", slogError(err)) + } + snapshotByChat := make(map[uuid.UUID]int64, len(chats)) + for _, chat := range chats { + snapshotByChat[chat.ID] = chat.SnapshotVersion + } + for _, key := range keys { + payload, err := json.Marshal(coderdpubsub.ChatStateOwnershipMessage{ + ChatID: key.ChatID, + SnapshotVersion: snapshotByChat[key.ChatID], + }) + if err != nil { + m.opts.Logger.Warn(ctx, "chatworker shutdown ownership marshal failed", slogError(err)) + continue + } + if err := m.opts.Pubsub.Publish(coderdpubsub.ChatStateOwnershipChannel, payload); err != nil { + m.opts.Logger.Warn(ctx, "chatworker shutdown ownership publish failed", slogError(err)) + } + } +} + +func (m *runnerManager) snapshotRunnerKeys() []runnerKey { + m.mu.Lock() + defer m.mu.Unlock() + keys := make([]runnerKey, 0, len(m.runners)) + for key := range m.runners { + keys = append(keys, key) + } + return keys +} + +func (m *runnerManager) databaseSyncLoop() { + ticker := m.opts.Clock.NewTicker(m.opts.RunnerSyncInterval, "chatworker", "runner-sync") + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := m.syncOnce(m.ctx); err != nil { + m.opts.Logger.Warn(m.ctx, "chatworker runner sync failed", slogError(err)) + } + case <-m.ctx.Done(): + return + } + } +} + +func (m *runnerManager) syncOnce(ctx context.Context) error { + keys := m.snapshotRunnerKeys() + if len(keys) == 0 { + return nil + } + idsByChat := make(map[uuid.UUID]struct{}, len(keys)) + for _, key := range keys { + idsByChat[key.ChatID] = struct{}{} + } + chatIDs := make([]uuid.UUID, 0, len(idsByChat)) + for id := range idsByChat { + chatIDs = append(chatIDs, id) + } + chats, err := m.opts.Store.GetChatsByIDsForRunnerSync(ctx, chatIDs) + if err != nil { + return xerrors.Errorf("get chats for runner sync: %w", err) + } + seen := make(map[uuid.UUID]struct{}, len(chats)) + for _, chat := range chats { + seen[chat.ID] = struct{}{} + m.RouteStateHint(ctx, stateUpdateFromChat(chat)) + } + for _, key := range keys { + if _, ok := seen[key.ChatID]; !ok { + m.requestCleanup(ctx, key) + } + } + return nil +} + +func (m *runnerManager) heartbeatLoop() { + ticker := m.opts.Clock.NewTicker(m.opts.HeartbeatInterval, "chatworker", "heartbeat") + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := m.heartbeatOnce(m.ctx); err != nil { + m.opts.Logger.Warn(m.ctx, "chatworker heartbeat failed", slogError(err)) + } + case <-m.ctx.Done(): + return + } + } +} + +func (m *runnerManager) heartbeatOnce(ctx context.Context) error { + keys := m.snapshotRunnerKeys() + if len(keys) == 0 { + return nil + } + chatIDs := make([]uuid.UUID, 0, len(keys)) + runnerIDs := make([]uuid.UUID, 0, len(keys)) + for _, key := range keys { + chatIDs = append(chatIDs, key.ChatID) + runnerIDs = append(runnerIDs, key.RunnerID) + } + return m.opts.Store.BatchUpsertChatHeartbeats(ctx, database.BatchUpsertChatHeartbeatsParams{ + ChatIds: chatIDs, + RunnerIds: runnerIDs, + }) +} + +func (m *runnerManager) heartbeatCleanupLoop() { + ticker := m.opts.Clock.NewTicker(m.opts.HeartbeatCleanupInterval, "chatworker", "heartbeat-cleanup") + defer ticker.Stop() + for { + select { + case <-ticker.C: + if err := m.heartbeatCleanupOnce(m.ctx); err != nil { + m.opts.Logger.Warn(m.ctx, "chatworker heartbeat cleanup failed", slogError(err)) + } + case <-m.ctx.Done(): + return + } + } +} + +func (m *runnerManager) heartbeatCleanupOnce(ctx context.Context) error { + _, err := m.opts.Store.DeleteStaleChatHeartbeats(ctx, m.opts.HeartbeatStaleSeconds) + return err +} + +func stateUpdateFromChat(chat database.Chat) runnerStateUpdate { + var workerID *uuid.UUID + if chat.WorkerID.Valid { + id := chat.WorkerID.UUID + workerID = &id + } + var runnerID *uuid.UUID + if chat.RunnerID.Valid { + id := chat.RunnerID.UUID + runnerID = &id + } + return runnerStateUpdate{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + SnapshotVersion: chat.SnapshotVersion, + HistoryVersion: chat.HistoryVersion, + QueueVersion: chat.QueueVersion, + GenerationAttempt: chat.GenerationAttempt, + Status: chat.Status, + Archived: chat.Archived, + RequiresActionDeadlineAt: chat.RequiresActionDeadlineAt, + } +} + +func slogError(err error) slog.Field { + return slog.Error(err) +} diff --git a/coderd/x/chatd/runner_test.go b/coderd/x/chatd/runner_test.go new file mode 100644 index 0000000000..eca1df9a26 --- /dev/null +++ b/coderd/x/chatd/runner_test.go @@ -0,0 +1,137 @@ +package chatd //nolint:testpackage // Uses unexported chatworker helpers. + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestRunner_IgnoresDuplicateStateNotifications(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + startWorker(t, testOptions(t, f, starter)) + starter.waitCall(t, taskKindGeneration, chat.ID) + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + + publishChatUpdate(t, f, latest) + publishChatUpdate(t, f, latest) + starter.assertNoCall(t) +} + +func TestRunner_CancelsActiveTaskWhenHistoryChanges(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + updated := commitAssistantStep(t, f, chat.ID, "first step") + require.Greater(t, updated.HistoryVersion, first.input.HistoryVersion) + requireTaskCanceled(t, first) + second := starter.waitCall(t, taskKindGeneration, chat.ID) + require.Equal(t, updated.HistoryVersion, second.input.HistoryVersion) +} + +func TestRunner_CancelsActiveTaskWhenStatusChanges(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + updated := interruptChat(t, f, chat.ID) + require.Equal(t, database.ChatStatusInterrupting, updated.Status) + requireTaskCanceled(t, first) + second := starter.waitCall(t, taskKindInterrupt, chat.ID) + require.Equal(t, updated.HistoryVersion, second.input.HistoryVersion) +} + +func TestRunner_CleansUpOnOwnershipTakeover(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + acquireChat(t, f, chat.ID, uuid.New(), uuid.New()) + requireTaskCanceled(t, first) + starter.assertNoCall(t) +} + +func TestRunner_SerializesReplacementTasksForSameHistoryAndStatus(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(true) + defer starter.releaseAll() + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + forceExecutionStateAndPublish(t, f, chat.ID, database.ChatStatusInterrupting, false) + starter.waitCall(t, taskKindInterrupt, chat.ID) + forceExecutionStateAndPublish(t, f, chat.ID, database.ChatStatusRunning, false) + starter.assertNoCall(t) + + starter.release(t, 0) + replacement := starter.waitCall(t, taskKindGeneration, chat.ID) + require.Equal(t, first.input.HistoryVersion, replacement.input.HistoryVersion) +} + +func TestRunner_AllowsReplacementForDifferentHistoryOrStatus(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(true) + defer starter.releaseAll() + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + updated := commitAssistantStep(t, f, chat.ID, "different history") + second := starter.waitCall(t, taskKindGeneration, chat.ID) + require.Greater(t, second.input.HistoryVersion, first.input.HistoryVersion) + require.Equal(t, updated.HistoryVersion, second.input.HistoryVersion) +} + +func TestWorker_RoutesDatabaseSyncStateToActiveRunner(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + clock := quartz.NewMock(t) + starter := newBlockingTaskStarter(false) + opts := testOptions(t, f, starter) + opts.Clock = clock + opts.RunnerSyncInterval = time.Minute + startWorker(t, opts) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + forceExecutionState(t, f, chat.ID, database.ChatStatusInterrupting, false) + clock.Advance(time.Minute).MustWait(testutil.Context(t, testutil.WaitLong)) + requireTaskCanceled(t, first) + starter.waitCall(t, taskKindInterrupt, chat.ID) +} + +func TestWorker_CleanupStopsRoutingAndCancelsTasks(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + startWorker(t, testOptions(t, f, starter)) + first := starter.waitCall(t, taskKindGeneration, chat.ID) + + latest := acquireChat(t, f, chat.ID, uuid.New(), uuid.New()) + requireTaskCanceled(t, first) + publishChatUpdate(t, f, latest) + starter.assertNoCall(t) +} diff --git a/coderd/x/chatd/stream_loop.go b/coderd/x/chatd/stream_loop.go new file mode 100644 index 0000000000..5004e8a490 --- /dev/null +++ b/coderd/x/chatd/stream_loop.go @@ -0,0 +1,450 @@ +package chatd + +import ( + "context" + "encoding/json" + "net/http" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/db2sdk" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" +) + +type streamLoop struct { + chatID uuid.UUID + db database.Store + logger slog.Logger + state streamLocalState +} + +type streamLocalState struct { + snapshotVersion int64 + historyVersion int64 + queueVersion int64 + retryVersion int64 + + knownMessages map[int64]int64 + + status database.ChatStatus + + errorHistoryVersion int64 + actionRequiredHistoryVersion int64 + + workerID uuid.NullUUID + generationAttempt int64 + lastPartSeq int64 + + afterMessageID int64 + initialMessageSyncDone bool +} + +type streamSyncHint struct { + snapshotVersion int64 + historyVersion int64 + queueVersion int64 + retryVersion int64 + status database.ChatStatus + workerID uuid.NullUUID + generationAttempt int64 +} + +type streamDBSnapshot struct { + chat database.Chat + + historyChanged bool + changedMessages []database.ChatMessage + historyReset bool + fullHistory []database.ChatMessage + + queueChanged bool + queue []database.ChatQueuedMessage + + actionRequired *codersdk.ChatStreamActionRequired +} + +func newStreamLoop(chat database.Chat, db database.Store, logger slog.Logger, afterMessageID int64) *streamLoop { + return &streamLoop{ + chatID: chat.ID, + db: db, + logger: logger, + state: streamLocalState{ + knownMessages: make(map[int64]int64), + afterMessageID: afterMessageID, + }, + } +} + +func streamSyncHintFromUpdate(update coderdpubsub.ChatStateUpdateMessage) streamSyncHint { + hint := streamSyncHint{ + snapshotVersion: update.SnapshotVersion, + historyVersion: update.HistoryVersion, + queueVersion: update.QueueVersion, + retryVersion: update.RetryStateVersion, + status: database.ChatStatus(update.Status), + generationAttempt: update.GenerationAttempt, + } + if update.WorkerID != nil { + hint.workerID = uuid.NullUUID{UUID: *update.WorkerID, Valid: true} + } + return hint +} + +func (l *streamLoop) sync(ctx context.Context, hint streamSyncHint) ([]codersdk.ChatStreamEvent, streamRelayTarget, bool, error) { + if !l.shouldFetch(hint) { + return nil, l.currentRelayTarget(), false, nil + } + return l.syncDB(ctx) +} + +func (l *streamLoop) syncDB(ctx context.Context) ([]codersdk.ChatStreamEvent, streamRelayTarget, bool, error) { + snapshot, err := l.loadDBSnapshot(ctx) + if err != nil { + return nil, l.currentRelayTarget(), false, err + } + if snapshot.chat.SnapshotVersion <= l.state.snapshotVersion { + return nil, l.currentRelayTarget(), false, nil + } + return l.applyDBSnapshot(snapshot), l.currentRelayTarget(), true, nil +} + +func (l *streamLoop) shouldFetch(hint streamSyncHint) bool { + if hint.snapshotVersion <= l.state.snapshotVersion { + return false + } + if hint.historyVersion > l.state.historyVersion { + return true + } + if hint.queueVersion > l.state.queueVersion { + return true + } + if hint.retryVersion > l.state.retryVersion { + return true + } + if hint.status != l.state.status { + return true + } + if !sameNullUUID(hint.workerID, l.state.workerID) { + return true + } + if hint.generationAttempt != l.state.generationAttempt { + return true + } + return false +} + +func (l *streamLoop) loadDBSnapshot(ctx context.Context) (streamDBSnapshot, error) { + var snapshot streamDBSnapshot + machine := chatstate.NewChatMachine(l.db, nil, l.chatID) + err := machine.ReadLock(ctx, func(tx database.Store) error { + chat, err := tx.GetChatByID(ctx, l.chatID) + if err != nil { + return xerrors.Errorf("get chat for stream: %w", err) + } + snapshot.chat = chat + + if chat.HistoryVersion > l.state.historyVersion { + snapshot.historyChanged = true + snapshot.changedMessages, err = tx.GetChatMessagesByRevisionForStream(ctx, database.GetChatMessagesByRevisionForStreamParams{ + ChatID: l.chatID, + AfterRevision: l.state.historyVersion, + }) + if err != nil { + return xerrors.Errorf("get changed chat messages: %w", err) + } + for _, msg := range snapshot.changedMessages { + if msg.Deleted { + snapshot.historyReset = true + break + } + } + if snapshot.historyReset { + snapshot.fullHistory, err = tx.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: l.chatID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("get full chat history: %w", err) + } + } + } + + if chat.QueueVersion > l.state.queueVersion { + snapshot.queueChanged = true + snapshot.queue, err = tx.GetChatQueuedMessages(ctx, l.chatID) + if err != nil { + return xerrors.Errorf("get chat queue: %w", err) + } + } + + if chat.Status == database.ChatStatusRequiresAction { + history := snapshot.fullHistory + if len(history) == 0 { + history, err = tx.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: l.chatID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("get requires_action history: %w", err) + } + } + actionRequired, err := l.actionRequiredFromHistory(chat, history) + if err != nil { + return err + } + snapshot.actionRequired = actionRequired + } + return nil + }) + if err != nil { + return streamDBSnapshot{}, err + } + return snapshot, nil +} + +func (*streamLoop) actionRequiredFromHistory(chat database.Chat, messages []database.ChatMessage) (*codersdk.ChatStreamActionRequired, error) { + dynamicToolNames, err := parseDynamicToolNames(chat.DynamicTools) + if err != nil { + return nil, xerrors.Errorf("parse dynamic tools for stream: %w", err) + } + _, pending, err := unresolvedToolCallsFromHistory(messages, dynamicToolNames) + if err != nil { + return nil, xerrors.Errorf("derive pending dynamic tool calls: %w", err) + } + toolCalls := make([]codersdk.ChatStreamToolCall, 0, len(pending)) + for _, call := range pending { + toolCalls = append(toolCalls, codersdk.ChatStreamToolCall{ + ToolCallID: call.ToolCallID, + ToolName: call.ToolName, + Args: call.Args, + }) + } + return &codersdk.ChatStreamActionRequired{ToolCalls: toolCalls}, nil +} + +func (l *streamLoop) applyDBSnapshot(snapshot streamDBSnapshot) []codersdk.ChatStreamEvent { + chat := snapshot.chat + events := make([]codersdk.ChatStreamEvent, 0) + historyChanged := chat.HistoryVersion > l.state.historyVersion + generationChanged := chat.GenerationAttempt != l.state.generationAttempt + + if historyChanged { + events = append(events, l.messageEvents(snapshot)...) + } + if !l.state.initialMessageSyncDone { + l.state.initialMessageSyncDone = true + } + + if chat.QueueVersion > l.state.queueVersion { + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeQueueUpdate, + ChatID: l.chatID, + QueuedMessages: db2sdk.ChatQueuedMessages(snapshot.queue), + }) + } + + if chat.Status != l.state.status { + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeStatus, + ChatID: l.chatID, + Status: &codersdk.ChatStreamStatus{Status: codersdk.ChatStatus(chat.Status)}, + }) + } + + if chat.Status == database.ChatStatusError && chat.HistoryVersion > l.state.errorHistoryVersion { + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeError, + ChatID: l.chatID, + Error: l.chatError(chat), + }) + l.state.errorHistoryVersion = chat.HistoryVersion + } + + if chat.Status == database.ChatStatusRequiresAction && chat.HistoryVersion > l.state.actionRequiredHistoryVersion { + actionRequired := snapshot.actionRequired + if actionRequired == nil { + actionRequired = &codersdk.ChatStreamActionRequired{} + } + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeActionRequired, + ChatID: l.chatID, + ActionRequired: actionRequired, + }) + l.state.actionRequiredHistoryVersion = chat.HistoryVersion + } + + if chat.RetryStateVersion > l.state.retryVersion { + if retry := l.retryEvent(chat); retry != nil { + events = append(events, *retry) + } + } + + if historyChanged || (generationChanged && chat.GenerationAttempt != 0) { + l.state.lastPartSeq = 0 + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypePreviewReset, + ChatID: l.chatID, + }) + } + + l.state.snapshotVersion = chat.SnapshotVersion + l.state.historyVersion = chat.HistoryVersion + l.state.queueVersion = chat.QueueVersion + l.state.retryVersion = chat.RetryStateVersion + l.state.status = chat.Status + l.state.workerID = chat.WorkerID + l.state.generationAttempt = chat.GenerationAttempt + return events +} + +func (l *streamLoop) messageEvents(snapshot streamDBSnapshot) []codersdk.ChatStreamEvent { + if snapshot.historyReset { + events := []codersdk.ChatStreamEvent{{ + Type: codersdk.ChatStreamEventTypeHistoryReset, + ChatID: l.chatID, + }} + clear(l.state.knownMessages) + for _, msg := range snapshot.fullHistory { + l.state.knownMessages[msg.ID] = msg.Revision + sdkMsg := db2sdk.ChatMessage(msg) + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeMessage, + ChatID: l.chatID, + Message: &sdkMsg, + }) + } + return events + } + + events := make([]codersdk.ChatStreamEvent, 0, len(snapshot.changedMessages)) + for _, msg := range snapshot.changedMessages { + knownRevision := l.state.knownMessages[msg.ID] + if knownRevision >= msg.Revision { + continue + } + l.state.knownMessages[msg.ID] = msg.Revision + if !l.state.initialMessageSyncDone && msg.ID <= l.state.afterMessageID { + continue + } + sdkMsg := db2sdk.ChatMessage(msg) + events = append(events, codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeMessage, + ChatID: l.chatID, + Message: &sdkMsg, + }) + } + return events +} + +func (l *streamLoop) chatError(chat database.Chat) *codersdk.ChatError { + if !chat.LastError.Valid || len(chat.LastError.RawMessage) == 0 { + return &codersdk.ChatError{ + Message: "The chat request failed unexpectedly.", + Kind: codersdk.ChatErrorKindGeneric, + } + } + var payload codersdk.ChatError + if err := json.Unmarshal(chat.LastError.RawMessage, &payload); err != nil { + l.logger.Warn(context.Background(), "failed to parse chat stream last_error", + slog.F("chat_id", l.chatID), + slog.Error(err), + ) + return &codersdk.ChatError{ + Message: "The chat request failed unexpectedly.", + Kind: codersdk.ChatErrorKindGeneric, + } + } + if payload.Message == "" { + payload.Message = "The chat request failed unexpectedly." + } + if payload.Kind == "" { + payload.Kind = codersdk.ChatErrorKindGeneric + } + return &payload +} + +func (l *streamLoop) retryEvent(chat database.Chat) *codersdk.ChatStreamEvent { + if !chat.RetryState.Valid || len(chat.RetryState.RawMessage) == 0 { + return nil + } + var retry codersdk.ChatStreamRetry + if err := json.Unmarshal(chat.RetryState.RawMessage, &retry); err != nil { + l.logger.Warn(context.Background(), "failed to parse chat stream retry_state", + slog.F("chat_id", l.chatID), + slog.Error(err), + ) + return nil + } + return &codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeRetry, + ChatID: l.chatID, + Retry: &retry, + } +} + +func (l *streamLoop) part(part streamPart) (event codersdk.ChatStreamEvent, accepted bool, err error) { + if part.HistoryVersion != l.state.historyVersion || part.GenerationAttempt != l.state.generationAttempt { + return codersdk.ChatStreamEvent{}, false, nil + } + if part.Seq <= l.state.lastPartSeq { + return codersdk.ChatStreamEvent{}, false, nil + } + if part.Seq != l.state.lastPartSeq+1 { + err := xerrors.Errorf( + "chat stream message part sequence gap: got %d after %d", + part.Seq, + l.state.lastPartSeq, + ) + l.logger.Error(context.Background(), "chat stream message part sequence gap", + slog.F("chat_id", l.chatID), + slog.F("history_version", part.HistoryVersion), + slog.F("generation_attempt", part.GenerationAttempt), + slog.F("last_seq", l.state.lastPartSeq), + slog.F("seq", part.Seq), + slog.Error(err), + ) + return codersdk.ChatStreamEvent{}, false, err + } + l.state.lastPartSeq = part.Seq + return codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeMessagePart, + ChatID: l.chatID, + MessagePart: &codersdk.ChatStreamMessagePart{ + Role: part.Role, + Part: part.Part, + HistoryVersion: part.HistoryVersion, + GenerationAttempt: part.GenerationAttempt, + Seq: part.Seq, + }, + }, true, nil +} + +func (l *streamLoop) currentRelayTarget() streamRelayTarget { + return streamRelayTarget{ + workerID: l.state.workerID, + historyVersion: l.state.historyVersion, + generationAttempt: l.state.generationAttempt, + } +} + +func sameNullUUID(a, b uuid.NullUUID) bool { + if a.Valid != b.Valid { + return false + } + if !a.Valid { + return true + } + return a.UUID == b.UUID +} + +func cloneHeader(header http.Header) http.Header { + if header == nil { + return nil + } + return header.Clone() +} diff --git a/coderd/x/chatd/stream_loop_internal_test.go b/coderd/x/chatd/stream_loop_internal_test.go new file mode 100644 index 0000000000..eebd6d0c97 --- /dev/null +++ b/coderd/x/chatd/stream_loop_internal_test.go @@ -0,0 +1,351 @@ +package chatd + +import ( + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + "go.uber.org/mock/gomock" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" +) + +func TestStreamLoopSyncHintDecision(t *testing.T) { + t.Parallel() + + workerA := uuid.New() + workerB := uuid.New() + loop := &streamLoop{ + state: streamLocalState{ + snapshotVersion: 5, + historyVersion: 2, + queueVersion: 3, + retryVersion: 4, + status: database.ChatStatusRunning, + workerID: uuid.NullUUID{UUID: workerA, Valid: true}, + generationAttempt: 1, + }, + } + + for _, tt := range []struct { + name string + hint streamSyncHint + want bool + }{ + { + name: "stale snapshot ignored even with higher history", + hint: streamSyncHint{snapshotVersion: 5, historyVersion: 3}, + }, + { + name: "duplicate snapshot ignored", + hint: streamSyncHint{snapshotVersion: 5}, + }, + { + name: "new snapshot with no changed fields is ignored", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 3, retryVersion: 4, status: database.ChatStatusRunning, workerID: uuid.NullUUID{UUID: workerA, Valid: true}, generationAttempt: 1}, + }, + { + name: "new history fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 3}, + want: true, + }, + { + name: "new queue fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 4}, + want: true, + }, + { + name: "new retry fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 3, retryVersion: 5}, + want: true, + }, + { + name: "new status fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 3, retryVersion: 4, status: database.ChatStatusWaiting, workerID: uuid.NullUUID{UUID: workerA, Valid: true}, generationAttempt: 1}, + want: true, + }, + { + name: "new worker fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 3, retryVersion: 4, status: database.ChatStatusRunning, workerID: uuid.NullUUID{UUID: workerB, Valid: true}, generationAttempt: 1}, + want: true, + }, + { + name: "new generation attempt fetches", + hint: streamSyncHint{snapshotVersion: 6, historyVersion: 2, queueVersion: 3, retryVersion: 4, status: database.ChatStatusRunning, workerID: uuid.NullUUID{UUID: workerA, Valid: true}, generationAttempt: 2}, + want: true, + }, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + require.Equal(t, tt.want, loop.shouldFetch(tt.hint)) + }) + } +} + +func TestStreamLoopMessageSyncAfterIDAndEdits(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 1) + initial := streamDBSnapshot{ + chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 1, + HistoryVersion: 1, + }, + changedMessages: []database.ChatMessage{ + streamMessage(t, chatID, 1, 1, database.ChatMessageRoleUser, "already seen", false), + streamMessage(t, chatID, 2, 1, database.ChatMessageRoleAssistant, "new", false), + }, + } + + events := loop.applyDBSnapshot(initial) + requireEventTypes(t, events, + codersdk.ChatStreamEventTypeMessage, + codersdk.ChatStreamEventTypeStatus, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, int64(2), events[0].Message.ID) + + edited := streamDBSnapshot{ + chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 2, + HistoryVersion: 2, + }, + changedMessages: []database.ChatMessage{ + streamMessage(t, chatID, 1, 2, database.ChatMessageRoleUser, "edited", false), + }, + } + events = loop.applyDBSnapshot(edited) + requireEventTypes(t, events, + codersdk.ChatStreamEventTypeMessage, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, int64(1), events[0].Message.ID) +} + +func TestStreamLoopHistoryReset(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 0) + loop.state.snapshotVersion = 1 + loop.state.historyVersion = 1 + loop.state.status = database.ChatStatusRunning + loop.state.initialMessageSyncDone = true + loop.state.knownMessages[1] = 1 + loop.state.knownMessages[2] = 1 + + events := loop.applyDBSnapshot(streamDBSnapshot{ + chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRunning, + SnapshotVersion: 2, + HistoryVersion: 2, + }, + changedMessages: []database.ChatMessage{ + streamMessage(t, chatID, 1, 2, database.ChatMessageRoleUser, "deleted", true), + }, + historyReset: true, + fullHistory: []database.ChatMessage{ + streamMessage(t, chatID, 3, 2, database.ChatMessageRoleUser, "replacement", false), + }, + }) + + requireEventTypes(t, events, + codersdk.ChatStreamEventTypeHistoryReset, + codersdk.ChatStreamEventTypeMessage, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, int64(3), events[1].Message.ID) + require.Equal(t, map[int64]int64{3: 2}, loop.state.knownMessages) +} + +func TestStreamLoopQueueStatusRetryErrorActionRequiredAndPreviewReset(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + retry := codersdk.ChatStreamRetry{Attempt: 2, DelayMs: 100, Error: "retrying", RetryingAt: time.Now()} + retryRaw, err := json.Marshal(retry) + require.NoError(t, err) + chatError := codersdk.ChatError{Message: "provider failed", Kind: codersdk.ChatErrorKindConfig} + errorRaw, err := json.Marshal(chatError) + require.NoError(t, err) + + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 0) + loop.state.snapshotVersion = 1 + loop.state.historyVersion = 1 + loop.state.queueVersion = 1 + loop.state.retryVersion = 1 + loop.state.generationAttempt = 1 + loop.state.status = database.ChatStatusRunning + + events := loop.applyDBSnapshot(streamDBSnapshot{ + chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusError, + SnapshotVersion: 2, + HistoryVersion: 2, + QueueVersion: 2, + RetryStateVersion: 2, + GenerationAttempt: 2, + LastError: pqtype.NullRawMessage{RawMessage: errorRaw, Valid: true}, + RetryState: pqtype.NullRawMessage{RawMessage: retryRaw, Valid: true}, + }, + queue: []database.ChatQueuedMessage{}, + }) + + requireEventTypes(t, events, + codersdk.ChatStreamEventTypeQueueUpdate, + codersdk.ChatStreamEventTypeStatus, + codersdk.ChatStreamEventTypeError, + codersdk.ChatStreamEventTypeRetry, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, chatError.Message, events[2].Error.Message) + require.Equal(t, retry.Attempt, events[3].Retry.Attempt) + + actionLoop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 0) + actionEvents := actionLoop.applyDBSnapshot(streamDBSnapshot{ + chat: database.Chat{ + ID: chatID, + Status: database.ChatStatusRequiresAction, + SnapshotVersion: 1, + HistoryVersion: 1, + }, + actionRequired: &codersdk.ChatStreamActionRequired{ToolCalls: []codersdk.ChatStreamToolCall{{ToolCallID: "call-1", ToolName: "browser"}}}, + }) + requireEventTypes(t, actionEvents, + codersdk.ChatStreamEventTypeStatus, + codersdk.ChatStreamEventTypeActionRequired, + codersdk.ChatStreamEventTypePreviewReset, + ) + require.Equal(t, "call-1", actionEvents[1].ActionRequired.ToolCalls[0].ToolCallID) +} + +func TestStreamLoopActionRequiredFromHistory(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + toolDefs, err := json.Marshal([]codersdk.DynamicTool{{Name: "browser"}}) + require.NoError(t, err) + assistant := streamMessageParts(t, chatID, 1, 1, database.ChatMessageRoleAssistant, []codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "call-1", + ToolName: "browser", + Args: json.RawMessage(`{"url":"https://example.com"}`), + }}, false) + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, nil), 0) + action, err := loop.actionRequiredFromHistory(database.Chat{ + ID: chatID, + DynamicTools: pqtype.NullRawMessage{RawMessage: toolDefs, Valid: true}, + }, []database.ChatMessage{assistant}) + require.NoError(t, err) + require.Len(t, action.ToolCalls, 1) + require.Equal(t, "call-1", action.ToolCalls[0].ToolCallID) + require.Equal(t, "browser", action.ToolCalls[0].ToolName) +} + +func TestStreamLoopPartValidation(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + loop := newStreamLoop(database.Chat{ID: chatID}, nil, slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), 0) + loop.state.historyVersion = 7 + loop.state.generationAttempt = 3 + + event, accepted, err := loop.part(StreamPart{HistoryVersion: 7, GenerationAttempt: 3, Seq: 1, Role: codersdk.ChatMessageRoleAssistant, Part: codersdk.ChatMessageText("a")}) + require.NoError(t, err) + require.True(t, accepted) + require.Equal(t, codersdk.ChatStreamEventTypeMessagePart, event.Type) + require.Equal(t, int64(7), event.MessagePart.HistoryVersion) + require.Equal(t, int64(3), event.MessagePart.GenerationAttempt) + require.Equal(t, int64(1), event.MessagePart.Seq) + + _, accepted, err = loop.part(StreamPart{HistoryVersion: 6, GenerationAttempt: 3, Seq: 2, Part: codersdk.ChatMessageText("old history")}) + require.NoError(t, err) + require.False(t, accepted) + _, accepted, err = loop.part(StreamPart{HistoryVersion: 7, GenerationAttempt: 2, Seq: 2, Part: codersdk.ChatMessageText("old attempt")}) + require.NoError(t, err) + require.False(t, accepted) + _, accepted, err = loop.part(StreamPart{HistoryVersion: 7, GenerationAttempt: 3, Seq: 1, Part: codersdk.ChatMessageText("dup")}) + require.NoError(t, err) + require.False(t, accepted) + _, accepted, err = loop.part(StreamPart{HistoryVersion: 7, GenerationAttempt: 3, Seq: 3, Part: codersdk.ChatMessageText("gap")}) + require.Error(t, err) + require.False(t, accepted) +} + +func TestStreamLoopInitialSyncRecoversWithoutHint(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + tx := dbmock.NewMockStore(ctrl) + chatID := uuid.New() + loop := newStreamLoop(database.Chat{ID: chatID}, db, slogtest.Make(t, nil), 0) + loop.state.snapshotVersion = 1 + loop.state.status = database.ChatStatusRunning + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(tx) }, + ) + tx.EXPECT().GetChatByIDForShare(gomock.Any(), chatID).Return(database.Chat{ + ID: chatID, + Status: database.ChatStatusWaiting, + SnapshotVersion: 2, + }, nil) + tx.EXPECT().GetChatByID(gomock.Any(), chatID).Return(database.Chat{ + ID: chatID, + Status: database.ChatStatusWaiting, + SnapshotVersion: 2, + }, nil) + + events, _, changed, err := loop.syncDB(ctx) + require.NoError(t, err) + require.True(t, changed) + requireEventTypes(t, events, codersdk.ChatStreamEventTypeStatus) + require.Equal(t, codersdk.ChatStatusWaiting, events[0].Status.Status) +} + +func requireEventTypes(t *testing.T, events []codersdk.ChatStreamEvent, types ...codersdk.ChatStreamEventType) { + t.Helper() + require.Len(t, events, len(types)) + for i, typ := range types { + require.Equal(t, typ, events[i].Type, "event %d", i) + } +} + +func streamMessage(t *testing.T, chatID uuid.UUID, id int64, revision int64, role database.ChatMessageRole, text string, deleted bool) database.ChatMessage { + t.Helper() + return streamMessageParts(t, chatID, id, revision, role, []codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}, deleted) +} + +func streamMessageParts(t *testing.T, chatID uuid.UUID, id int64, revision int64, role database.ChatMessageRole, parts []codersdk.ChatMessagePart, deleted bool) database.ChatMessage { + t.Helper() + content, err := chatprompt.MarshalParts(parts) + require.NoError(t, err) + return database.ChatMessage{ + ID: id, + ChatID: chatID, + CreatedAt: time.Unix(id, 0), + Role: role, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + Deleted: deleted, + Revision: revision, + } +} diff --git a/coderd/x/chatd/stream_parts.go b/coderd/x/chatd/stream_parts.go new file mode 100644 index 0000000000..cd1879ae28 --- /dev/null +++ b/coderd/x/chatd/stream_parts.go @@ -0,0 +1,190 @@ +package chatd + +import ( + "context" + "net/http" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/websocket" +) + +type streamPartsControl struct { + HistoryVersion int64 `json:"history_version"` + GenerationAttempt int64 `json:"generation_attempt"` +} + +type streamPartsEndpoint struct { + chatID uuid.UUID + buffer *messagepartbuffer.Buffer + logger slog.Logger +} + +// ServeStreamPartsAuthorized serves the internal episode-selected parts stream +// for an already authorized chat route. +func (p *Server) ServeStreamPartsAuthorized(rw http.ResponseWriter, r *http.Request, chat database.Chat) error { + if p == nil || p.messagePartBuffer == nil { + return xerrors.New("message part buffer is not configured") + } + endpoint := streamPartsEndpoint{ + chatID: chat.ID, + buffer: p.messagePartBuffer, + logger: p.logger.Named("chat_stream_parts").With(slog.F("chat_id", chat.ID)), + } + return endpoint.serveWebSocket(rw, r) +} + +func (e streamPartsEndpoint) serveWebSocket(rw http.ResponseWriter, r *http.Request) error { + ctx := r.Context() + conn, err := websocket.Accept(rw, r, nil) + if err != nil { + return xerrors.Errorf("accept parts websocket: %w", err) + } + transport := streamPartsWebSocketServerTransport{conn: conn} + defer func() { + _ = transport.Close() + }() + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + go httpapi.HeartbeatClose(ctx, e.logger, cancel, conn) + + return e.serve(ctx, transport) +} + +func (e streamPartsEndpoint) serve(ctx context.Context, transport streamPartsServerTransport) error { + if e.buffer == nil { + return xerrors.New("message part buffer is not configured") + } + if transport == nil { + return xerrors.New("stream parts transport is not configured") + } + + ctx, cancel := context.WithCancel(ctx) + defer cancel() + + controlCh := make(chan streamPartsControl, 1) + errCh := make(chan error, 1) + go func() { + for { + control, err := transport.ReadControl(ctx) + if err != nil { + select { + case errCh <- err: + case <-ctx.Done(): + } + return + } + select { + case controlCh <- control: + case <-ctx.Done(): + return + } + } + }() + + var ( + parts <-chan messagepartbuffer.Part + partCancel func() + partCancelFn context.CancelFunc + selected streamPartsControl + lastSeq int64 + ) + defer func() { + if partCancel != nil { + partCancel() + } + if partCancelFn != nil { + partCancelFn() + } + }() + + selectEpisode := func(control streamPartsControl) error { + if partCancel != nil { + partCancel() + partCancel = nil + } + if partCancelFn != nil { + partCancelFn() + partCancelFn = nil + } + parts = nil + selected = control + lastSeq = 0 + partCtx, cancel := context.WithCancel(ctx) + ch, cancelSub, err := e.buffer.SubscribeToEpisode(partCtx, messagepartbuffer.Key{ + ChatID: e.chatID, + HistoryVersion: control.HistoryVersion, + GenerationAttempt: control.GenerationAttempt, + }) + if err != nil { + cancel() + return err + } + partCancelFn = cancel + partCancel = cancelSub + parts = ch + return nil + } + + for { + select { + case <-ctx.Done(): + return nil + case err := <-errCh: + if ctx.Err() != nil || streamPartsExpectedTransportClose(err) { + return nil + } + return err + case control := <-controlCh: + if err := selectEpisode(control); err != nil { + return err + } + case part, ok := <-parts: + if !ok { + parts = nil + continue + } + if part.Seq != lastSeq+1 { + return xerrors.Errorf("message part sequence gap: got %d after %d", part.Seq, lastSeq) + } + lastSeq = part.Seq + event := codersdk.ChatStreamEvent{ + Type: codersdk.ChatStreamEventTypeMessagePart, + ChatID: e.chatID, + MessagePart: &codersdk.ChatStreamMessagePart{ + Role: part.Role, + Part: part.MessagePart, + HistoryVersion: selected.HistoryVersion, + GenerationAttempt: selected.GenerationAttempt, + Seq: part.Seq, + }, + } + if err := transport.WriteEvents(ctx, []codersdk.ChatStreamEvent{event}); err != nil { + if ctx.Err() != nil || streamPartsExpectedTransportClose(err) { + return nil + } + return err + } + } + } +} + +func StreamPartFromEvent(event codersdk.ChatStreamEvent) (StreamPart, bool) { + if event.Type != codersdk.ChatStreamEventTypeMessagePart || event.MessagePart == nil { + return StreamPart{}, false + } + return StreamPart{ + HistoryVersion: event.MessagePart.HistoryVersion, + GenerationAttempt: event.MessagePart.GenerationAttempt, + Seq: event.MessagePart.Seq, + Role: event.MessagePart.Role, + Part: event.MessagePart.Part, + }, true +} diff --git a/coderd/x/chatd/stream_parts_dialer.go b/coderd/x/chatd/stream_parts_dialer.go new file mode 100644 index 0000000000..eaada8220d --- /dev/null +++ b/coderd/x/chatd/stream_parts_dialer.go @@ -0,0 +1,60 @@ +package chatd + +import ( + "context" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" +) + +// LocalStreamPartsDialerConfig configures an in-process stream parts dialer. +type LocalStreamPartsDialerConfig struct { + Buffer *messagepartbuffer.Buffer + Logger slog.Logger +} + +// NewLocalStreamPartsDialer returns a dialer that streams message parts through +// in-process channels while using the same stream serving loop as WebSockets. +func NewLocalStreamPartsDialer(cfg LocalStreamPartsDialerConfig) StreamPartsDialer { + return func(ctx context.Context, input StreamPartsDialInput) (StreamPartsSession, error) { + if cfg.Buffer == nil { + return nil, xerrors.New("message part buffer is not configured") + } + serverTransport, clientTransport := newStreamPartsChannelTransportPair() + logger := cfg.Logger.Named("chat_stream_parts").With(slog.F("chat_id", input.ChatID)) + endpoint := streamPartsEndpoint{ + chatID: input.ChatID, + buffer: cfg.Buffer, + logger: logger, + } + serveCtx, cancel := context.WithCancel(ctx) + go func() { + defer cancel() + defer func() { + _ = serverTransport.Close() + }() + if err := endpoint.serve(serveCtx, serverTransport); err != nil && !streamPartsExpectedTransportClose(err) { + logger.Debug(serveCtx, "chat stream parts closed", slog.Error(err)) + } + }() + return newStreamPartsTransportSession(serveCtx, clientTransport), nil + } +} + +func streamPartsDialerForServer(workerID uuid.UUID, local StreamPartsDialer, remote StreamPartsDialer) StreamPartsDialer { + return func(ctx context.Context, input StreamPartsDialInput) (StreamPartsSession, error) { + if local == nil && remote == nil { + return nil, xerrors.New("stream parts dialer is not configured") + } + if remote == nil || input.WorkerID == uuid.Nil || input.WorkerID == workerID { + if local == nil { + return nil, xerrors.New("local stream parts dialer is not configured") + } + return local(ctx, input) + } + return remote(ctx, input) + } +} diff --git a/coderd/x/chatd/stream_parts_internal_test.go b/coderd/x/chatd/stream_parts_internal_test.go new file mode 100644 index 0000000000..cc855daddd --- /dev/null +++ b/coderd/x/chatd/stream_parts_internal_test.go @@ -0,0 +1,354 @@ +package chatd + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" +) + +func TestStreamPartsEndpointReplayLiveAndReselect(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + serverTransport, clientTransport := newStreamPartsChannelTransportPair() + serveDone := serveStreamPartsEndpoint(ctx, t, endpoint, serverTransport) + defer func() { + require.NoError(t, clientTransport.Close()) + <-serveDone + }() + + firstKey := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 1, GenerationAttempt: 1} + secondKey := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 2, GenerationAttempt: 1} + require.NoError(t, buffer.CreateEpisode(firstKey)) + require.NoError(t, buffer.AddPart(firstKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("replayed"))) + require.NoError(t, buffer.CreateEpisode(secondKey)) + require.NoError(t, buffer.AddPart(secondKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("second"))) + + require.NoError(t, clientTransport.WriteControl(ctx, streamPartsControl{HistoryVersion: 1, GenerationAttempt: 1})) + got := readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Len(t, got, 1) + require.Equal(t, "replayed", got[0].MessagePart.Part.Text) + require.Equal(t, int64(1), got[0].MessagePart.Seq) + require.Equal(t, int64(1), got[0].MessagePart.HistoryVersion) + require.Equal(t, int64(1), got[0].MessagePart.GenerationAttempt) + + require.NoError(t, buffer.AddPart(firstKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("live"))) + got = readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Len(t, got, 1) + require.Equal(t, "live", got[0].MessagePart.Part.Text) + require.Equal(t, int64(2), got[0].MessagePart.Seq) + + require.NoError(t, clientTransport.WriteControl(ctx, streamPartsControl{HistoryVersion: 2, GenerationAttempt: 1})) + got = readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Len(t, got, 1) + require.Equal(t, "second", got[0].MessagePart.Part.Text) + require.Equal(t, int64(2), got[0].MessagePart.HistoryVersion) + + require.NoError(t, buffer.AddPart(firstKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("ignored"))) + select { + case <-ctx.Done(): + t.Fatal("timed out waiting to verify previous episode was canceled") + default: + } + require.NoError(t, buffer.AddPart(secondKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("second-live"))) + got = readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Equal(t, "second-live", got[0].MessagePart.Part.Text) +} + +func TestStreamPartsEndpointWaitsForMissingEpisode(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + serverTransport, clientTransport := newStreamPartsChannelTransportPair() + serveDone := serveStreamPartsEndpoint(ctx, t, endpoint, serverTransport) + defer func() { + require.NoError(t, clientTransport.Close()) + <-serveDone + }() + + key := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 9, GenerationAttempt: 2} + require.NoError(t, clientTransport.WriteControl(ctx, streamPartsControl{HistoryVersion: key.HistoryVersion, GenerationAttempt: key.GenerationAttempt})) + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("eventual"))) + + got := readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Len(t, got, 1) + require.Equal(t, "eventual", got[0].MessagePart.Part.Text) + require.Equal(t, int64(9), got[0].MessagePart.HistoryVersion) + require.Equal(t, int64(2), got[0].MessagePart.GenerationAttempt) +} + +func TestStreamPartsEndpointReselectsWhileEpisodeMissing(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + serverTransport, clientTransport := newStreamPartsChannelTransportPair() + serveDone := serveStreamPartsEndpoint(ctx, t, endpoint, serverTransport) + defer func() { + require.NoError(t, clientTransport.Close()) + <-serveDone + }() + + missingKey := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 10, GenerationAttempt: 1} + selectedKey := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 11, GenerationAttempt: 1} + require.NoError(t, clientTransport.WriteControl(ctx, streamPartsControl{HistoryVersion: missingKey.HistoryVersion, GenerationAttempt: missingKey.GenerationAttempt})) + require.NoError(t, clientTransport.WriteControl(ctx, streamPartsControl{HistoryVersion: selectedKey.HistoryVersion, GenerationAttempt: selectedKey.GenerationAttempt})) + require.NoError(t, buffer.CreateEpisode(selectedKey)) + require.NoError(t, buffer.AddPart(selectedKey, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("selected"))) + + got := readStreamPartsTransportBatch(ctx, t, clientTransport) + require.Len(t, got, 1) + require.Equal(t, "selected", got[0].MessagePart.Part.Text) + require.Equal(t, selectedKey.HistoryVersion, got[0].MessagePart.HistoryVersion) +} + +func TestStreamPartsEndpointClientDisconnectCancels(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + serverTransport, clientTransport := newStreamPartsChannelTransportPair() + serveDone := serveStreamPartsEndpoint(ctx, t, endpoint, serverTransport) + require.NoError(t, clientTransport.Close()) + + select { + case <-serveDone: + case <-ctx.Done(): + t.Fatal("stream parts endpoint did not exit after client disconnect") + } +} + +func TestStreamPartsEndpointWebSocket(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + _ = endpoint.serveWebSocket(rw, r) + })) + t.Cleanup(server.Close) + + key := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 1, GenerationAttempt: 1} + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("websocket"))) + + conn, resp, err := websocket.Dial(ctx, server.URL, nil) + require.NoError(t, err) + if resp != nil && resp.Body != nil { + require.NoError(t, resp.Body.Close()) + } + defer conn.Close(websocket.StatusNormalClosure, "") + + require.NoError(t, wsjson.Write(ctx, conn, streamPartsControl{HistoryVersion: 1, GenerationAttempt: 1})) + got := readStreamPartsWebSocketBatch(ctx, t, conn) + require.Len(t, got, 1) + require.Equal(t, "websocket", got[0].MessagePart.Part.Text) +} + +func TestStreamPartsWebSocketSession(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + endpoint := streamPartsEndpoint{ + chatID: chatID, + buffer: buffer, + logger: slogtest.Make(t, nil), + } + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + _ = endpoint.serveWebSocket(rw, r) + })) + t.Cleanup(server.Close) + + key := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 4, GenerationAttempt: 2} + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("session"))) + + conn, resp, err := websocket.Dial(ctx, server.URL, nil) + require.NoError(t, err) + if resp != nil && resp.Body != nil { + require.NoError(t, resp.Body.Close()) + } + session := NewStreamPartsJSONSession(ctx, conn) + defer session.Close() + + require.NoError(t, session.SelectEpisode(ctx, key.HistoryVersion, key.GenerationAttempt)) + part := readStreamPart(ctx, t, session.Parts()) + require.Equal(t, key.HistoryVersion, part.HistoryVersion) + require.Equal(t, key.GenerationAttempt, part.GenerationAttempt) + require.Equal(t, int64(1), part.Seq) + require.Equal(t, "session", part.Part.Text) +} + +func TestLocalStreamPartsDialerReplayLiveAndClose(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + chatID := uuid.New() + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + dialer := NewLocalStreamPartsDialer(LocalStreamPartsDialerConfig{ + Buffer: buffer, + Logger: slogtest.Make(t, nil), + }) + key := messagepartbuffer.Key{ChatID: chatID, HistoryVersion: 3, GenerationAttempt: 1} + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("replayed"))) + + session, err := dialer(ctx, StreamPartsDialInput{ChatID: chatID, WorkerID: uuid.New()}) + require.NoError(t, err) + require.NoError(t, session.SelectEpisode(ctx, key.HistoryVersion, key.GenerationAttempt)) + + part := readStreamPart(ctx, t, session.Parts()) + require.Equal(t, int64(1), part.Seq) + require.Equal(t, "replayed", part.Part.Text) + + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("live"))) + part = readStreamPart(ctx, t, session.Parts()) + require.Equal(t, int64(2), part.Seq) + require.Equal(t, "live", part.Part.Text) + + require.NoError(t, session.Close()) + select { + case _, ok := <-session.Parts(): + require.False(t, ok) + case <-ctx.Done(): + t.Fatal("stream parts session did not close") + } +} + +func TestStreamPartsDialerForServer(t *testing.T) { + t.Parallel() + + serverWorkerID := uuid.New() + remoteWorkerID := uuid.New() + + cases := []struct { + name string + remote bool + workerID uuid.UUID + want string + }{ + {name: "no remote uses local", workerID: remoteWorkerID, want: "local"}, + {name: "same worker uses local", remote: true, workerID: serverWorkerID, want: "local"}, + {name: "different worker uses remote", remote: true, workerID: remoteWorkerID, want: "remote"}, + {name: "nil worker uses local", remote: true, workerID: uuid.Nil, want: "local"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + called := make(chan string, 1) + local := func(context.Context, StreamPartsDialInput) (StreamPartsSession, error) { + called <- "local" + return nil, xerrors.New("local") + } + var remote StreamPartsDialer + if tc.remote { + remote = func(context.Context, StreamPartsDialInput) (StreamPartsSession, error) { + called <- "remote" + return nil, xerrors.New("remote") + } + } + dialer := streamPartsDialerForServer(serverWorkerID, local, remote) + _, _ = dialer(ctx, StreamPartsDialInput{WorkerID: tc.workerID}) + require.Equal(t, tc.want, <-called) + }) + } +} + +func serveStreamPartsEndpoint(ctx context.Context, t *testing.T, endpoint streamPartsEndpoint, transport streamPartsServerTransport) <-chan struct{} { + t.Helper() + done := make(chan struct{}) + go func() { + defer close(done) + err := endpoint.serve(ctx, transport) + if err != nil && !streamPartsExpectedTransportClose(err) { + require.NoError(t, err) + } + }() + return done +} + +func readStreamPartsTransportBatch(ctx context.Context, t *testing.T, transport streamPartsClientTransport) []codersdk.ChatStreamEvent { + t.Helper() + got, err := transport.ReadEvents(ctx) + require.NoError(t, err) + assertStreamPartsBatch(t, got) + return got +} + +func readStreamPartsWebSocketBatch(ctx context.Context, t *testing.T, conn *websocket.Conn) []codersdk.ChatStreamEvent { + t.Helper() + var got []codersdk.ChatStreamEvent + require.NoError(t, wsjson.Read(ctx, conn, &got)) + assertStreamPartsBatch(t, got) + return got +} + +func assertStreamPartsBatch(t *testing.T, got []codersdk.ChatStreamEvent) { + t.Helper() + for _, event := range got { + require.Equal(t, codersdk.ChatStreamEventTypeMessagePart, event.Type) + require.NotNil(t, event.MessagePart) + } +} + +func readStreamPart(ctx context.Context, t *testing.T, parts <-chan StreamPart) StreamPart { + t.Helper() + select { + case part, ok := <-parts: + require.True(t, ok) + return part + case <-ctx.Done(): + t.Fatal("timed out waiting for stream part") + return StreamPart{} + } +} diff --git a/coderd/x/chatd/stream_parts_transport.go b/coderd/x/chatd/stream_parts_transport.go new file mode 100644 index 0000000000..461a950f6b --- /dev/null +++ b/coderd/x/chatd/stream_parts_transport.go @@ -0,0 +1,267 @@ +package chatd + +import ( + "context" + "errors" + "net" + "sync" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk" + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" +) + +var errStreamPartsTransportClosed = xerrors.New("stream parts transport closed") + +type streamPartsServerTransport interface { + ReadControl(context.Context) (streamPartsControl, error) + WriteEvents(context.Context, []codersdk.ChatStreamEvent) error + Close() error +} + +type streamPartsClientTransport interface { + WriteControl(context.Context, streamPartsControl) error + ReadEvents(context.Context) ([]codersdk.ChatStreamEvent, error) + Close() error +} + +type streamPartsWebSocketServerTransport struct { + conn *websocket.Conn +} + +func (t streamPartsWebSocketServerTransport) ReadControl(ctx context.Context) (streamPartsControl, error) { + var control streamPartsControl + if err := wsjson.Read(ctx, t.conn, &control); err != nil { + return streamPartsControl{}, err + } + return control, nil +} + +func (t streamPartsWebSocketServerTransport) WriteEvents(ctx context.Context, events []codersdk.ChatStreamEvent) error { + return wsjson.Write(ctx, t.conn, events) +} + +func (t streamPartsWebSocketServerTransport) Close() error { + return t.conn.Close(websocket.StatusNormalClosure, "") +} + +type streamPartsWebSocketClientTransport struct { + conn *websocket.Conn +} + +func (t streamPartsWebSocketClientTransport) WriteControl(ctx context.Context, control streamPartsControl) error { + return wsjson.Write(ctx, t.conn, control) +} + +func (t streamPartsWebSocketClientTransport) ReadEvents(ctx context.Context) ([]codersdk.ChatStreamEvent, error) { + var batch []codersdk.ChatStreamEvent + if err := wsjson.Read(ctx, t.conn, &batch); err != nil { + return nil, err + } + return batch, nil +} + +func (t streamPartsWebSocketClientTransport) Close() error { + return t.conn.Close(websocket.StatusNormalClosure, "") +} + +type streamPartsChannelPipe struct { + controlCh chan streamPartsControl + eventsCh chan []codersdk.ChatStreamEvent + done chan struct{} + closeOnce sync.Once +} + +type streamPartsChannelServerTransport struct { + pipe *streamPartsChannelPipe +} + +type streamPartsChannelClientTransport struct { + pipe *streamPartsChannelPipe +} + +func newStreamPartsChannelTransportPair() (streamPartsServerTransport, streamPartsClientTransport) { + pipe := &streamPartsChannelPipe{ + controlCh: make(chan streamPartsControl, 1), + eventsCh: make(chan []codersdk.ChatStreamEvent, 128), + done: make(chan struct{}), + } + return streamPartsChannelServerTransport{pipe: pipe}, streamPartsChannelClientTransport{pipe: pipe} +} + +func (t streamPartsChannelServerTransport) ReadControl(ctx context.Context) (streamPartsControl, error) { + select { + case <-ctx.Done(): + return streamPartsControl{}, ctx.Err() + case <-t.pipe.done: + return streamPartsControl{}, errStreamPartsTransportClosed + default: + } + select { + case control := <-t.pipe.controlCh: + return control, nil + case <-ctx.Done(): + return streamPartsControl{}, ctx.Err() + case <-t.pipe.done: + return streamPartsControl{}, errStreamPartsTransportClosed + } +} + +func (t streamPartsChannelServerTransport) WriteEvents(ctx context.Context, events []codersdk.ChatStreamEvent) error { + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.pipe.done: + return errStreamPartsTransportClosed + default: + } + select { + case t.pipe.eventsCh <- events: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-t.pipe.done: + return errStreamPartsTransportClosed + } +} + +func (t streamPartsChannelServerTransport) Close() error { + return t.pipe.close() +} + +func (t streamPartsChannelClientTransport) WriteControl(ctx context.Context, control streamPartsControl) error { + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.pipe.done: + return errStreamPartsTransportClosed + default: + } + select { + case t.pipe.controlCh <- control: + return nil + case <-ctx.Done(): + return ctx.Err() + case <-t.pipe.done: + return errStreamPartsTransportClosed + } +} + +func (t streamPartsChannelClientTransport) ReadEvents(ctx context.Context) ([]codersdk.ChatStreamEvent, error) { + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-t.pipe.done: + return nil, errStreamPartsTransportClosed + default: + } + select { + case events := <-t.pipe.eventsCh: + return events, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-t.pipe.done: + return nil, errStreamPartsTransportClosed + } +} + +func (t streamPartsChannelClientTransport) Close() error { + return t.pipe.close() +} + +func (p *streamPartsChannelPipe) close() error { + p.closeOnce.Do(func() { + close(p.done) + }) + return nil +} + +type streamPartsTransportSession struct { + ctx context.Context + cancel context.CancelFunc + transport streamPartsClientTransport + parts chan StreamPart + closeOnce sync.Once + closeErr error +} + +func newStreamPartsTransportSession(ctx context.Context, transport streamPartsClientTransport) *streamPartsTransportSession { + sessionCtx, cancel := context.WithCancel(ctx) + session := &streamPartsTransportSession{ + ctx: sessionCtx, + cancel: cancel, + transport: transport, + parts: make(chan StreamPart, 128), + } + go session.readLoop() + return session +} + +func (s *streamPartsTransportSession) SelectEpisode(ctx context.Context, historyVersion, generationAttempt int64) error { + return s.transport.WriteControl(ctx, streamPartsControl{ + HistoryVersion: historyVersion, + GenerationAttempt: generationAttempt, + }) +} + +func (s *streamPartsTransportSession) Parts() <-chan StreamPart { + return s.parts +} + +func (s *streamPartsTransportSession) Close() error { + s.closeOnce.Do(func() { + s.cancel() + s.closeErr = s.transport.Close() + }) + return s.closeErr +} + +func (s *streamPartsTransportSession) readLoop() { + defer close(s.parts) + for { + batch, err := s.transport.ReadEvents(s.ctx) + if err != nil { + return + } + for _, event := range batch { + part, ok := StreamPartFromEvent(event) + if !ok { + continue + } + select { + case s.parts <- part: + case <-s.ctx.Done(): + return + } + } + } +} + +type StreamPartsJSONSession struct { + *streamPartsTransportSession +} + +func NewStreamPartsJSONSession(ctx context.Context, conn *websocket.Conn) *StreamPartsJSONSession { + return &StreamPartsJSONSession{ + streamPartsTransportSession: newStreamPartsTransportSession(ctx, streamPartsWebSocketClientTransport{conn: conn}), + } +} + +func streamPartsExpectedTransportClose(err error) bool { + if err == nil { + return true + } + if errors.Is(err, errStreamPartsTransportClosed) || + errors.Is(err, context.Canceled) || + errors.Is(err, net.ErrClosed) { + return true + } + switch websocket.CloseStatus(err) { + case websocket.StatusNormalClosure, websocket.StatusGoingAway: + return true + default: + return false + } +} diff --git a/coderd/x/chatd/stream_relay.go b/coderd/x/chatd/stream_relay.go new file mode 100644 index 0000000000..b2d42e069a --- /dev/null +++ b/coderd/x/chatd/stream_relay.go @@ -0,0 +1,248 @@ +package chatd + +import ( + "context" + "errors" + "net/http" + "sync" + "time" + + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/quartz" +) + +const ( + streamRelayRetryInitialBackoff = 100 * time.Millisecond + streamRelayRetryMaxBackoff = 5 * time.Second +) + +type streamRelayForwarder struct { + chatID uuid.UUID + requestHeader http.Header + dialer StreamPartsDialer + clock quartz.Clock + logger slog.Logger + + parts chan StreamPart + + ctx context.Context + cancel context.CancelFunc + done chan struct{} + + configure chan streamRelayTarget + closeOnce sync.Once +} + +func newStreamRelayForwarder( + chatID uuid.UUID, + requestHeader http.Header, + dialer StreamPartsDialer, + clock quartz.Clock, + logger slog.Logger, +) *streamRelayForwarder { + if clock == nil { + clock = quartz.NewReal() + } + ctx, cancel := context.WithCancel(context.Background()) + f := &streamRelayForwarder{ + chatID: chatID, + requestHeader: cloneHeader(requestHeader), + dialer: dialer, + clock: clock, + logger: logger, + parts: make(chan StreamPart, 128), + ctx: ctx, + cancel: cancel, + done: make(chan struct{}), + configure: make(chan streamRelayTarget, 1), + } + go f.loop() + return f +} + +func (f *streamRelayForwarder) Parts() <-chan StreamPart { + return f.parts +} + +func (f *streamRelayForwarder) Configure(ctx context.Context, target streamRelayTarget) { + if f == nil { + return + } + // Drop any pending target so the buffered channel always holds the most + // recent configuration. + select { + case <-f.configure: + default: + } + select { + case f.configure <- target: + case <-f.ctx.Done(): + case <-ctx.Done(): + } +} + +func (f *streamRelayForwarder) Close() { + if f == nil { + return + } + f.closeOnce.Do(func() { + f.cancel() + <-f.done + }) +} + +func (f *streamRelayForwarder) loop() { + defer close(f.done) + defer close(f.parts) + var ( + target streamRelayTarget + connected streamRelayTarget + session StreamPartsSession + sessionParts <-chan StreamPart + retryTimer *quartz.Timer + retryC <-chan time.Time + retryBackoff = streamRelayRetryInitialBackoff + ) + stopRetry := func() { + if retryTimer != nil { + retryTimer.Stop() + retryTimer = nil + retryC = nil + } + } + defer stopRetry() + closeSession := func() { + if session != nil { + _ = session.Close() + } + session = nil + sessionParts = nil + connected = streamRelayTarget{} + } + defer closeSession() + scheduleRetry := func() { + if !target.needsRelay() || f.dialer == nil || retryTimer != nil { + return + } + retryTimer = f.clock.NewTimer(retryBackoff, "chatd", "stream-relay-retry") + retryC = retryTimer.C + if retryBackoff < streamRelayRetryMaxBackoff { + retryBackoff *= 2 + if retryBackoff > streamRelayRetryMaxBackoff { + retryBackoff = streamRelayRetryMaxBackoff + } + } + } + connect := func(ctx context.Context) { + stopRetry() + if !target.needsRelay() { + closeSession() + return + } + if f.dialer == nil { + return + } + if session != nil && connected.workerID.Valid && sameNullUUID(connected.workerID, target.workerID) { + if err := session.SelectEpisode(ctx, target.historyVersion, target.generationAttempt); err != nil { + f.logger.Warn(ctx, "failed to select stream parts episode", + slog.F("chat_id", f.chatID), + slog.F("history_version", target.historyVersion), + slog.F("generation_attempt", target.generationAttempt), + slog.Error(err), + ) + closeSession() + scheduleRetry() + return + } + connected = target + retryBackoff = streamRelayRetryInitialBackoff + return + } + closeSession() + newSession, err := f.dialer(ctx, StreamPartsDialInput{ + ChatID: f.chatID, + WorkerID: target.workerID.UUID, + RequestHeader: cloneHeader(f.requestHeader), + }) + if err != nil { + f.logger.Warn(ctx, "failed to dial stream parts relay", + slog.F("chat_id", f.chatID), + slog.F("worker_id", target.workerID.UUID), + slog.Error(err), + ) + // Unrecoverable dial errors (e.g. auth failures) will not + // succeed on retry with the same inputs, so wait for the next + // configuration instead of scheduling a retry. + if !streamPartsDialUnrecoverable(err) { + scheduleRetry() + } + return + } + session = newSession + sessionParts = newSession.Parts() + connected = streamRelayTarget{workerID: target.workerID} + if err := session.SelectEpisode(ctx, target.historyVersion, target.generationAttempt); err != nil { + f.logger.Warn(ctx, "failed to select stream parts episode", + slog.F("chat_id", f.chatID), + slog.F("history_version", target.historyVersion), + slog.F("generation_attempt", target.generationAttempt), + slog.Error(err), + ) + closeSession() + scheduleRetry() + return + } + connected = target + retryBackoff = streamRelayRetryInitialBackoff + } + + for { + select { + case <-f.ctx.Done(): + return + case nextTarget := <-f.configure: + target = nextTarget + connect(f.ctx) + case <-retryC: + retryTimer = nil + retryC = nil + connect(f.ctx) + case part, ok := <-sessionParts: + if !ok { + closeSession() + scheduleRetry() + continue + } + if !connected.sameEpisode(target) || + part.HistoryVersion != target.historyVersion || + part.GenerationAttempt != target.generationAttempt { + continue + } + select { + case f.parts <- part: + case <-f.ctx.Done(): + return + } + } + } +} + +func (t streamRelayTarget) needsRelay() bool { + return t.workerID.Valid && t.generationAttempt > 0 +} + +// streamPartsDialUnrecoverable reports whether a dial error signals that +// retrying with the same inputs is futile, such as an auth failure. Dialers +// opt in by returning errors that implement IsUnrecoverable. +func streamPartsDialUnrecoverable(err error) bool { + var unrecoverable interface{ IsUnrecoverable() bool } + return errors.As(err, &unrecoverable) && unrecoverable.IsUnrecoverable() +} + +func (t streamRelayTarget) sameEpisode(other streamRelayTarget) bool { + return sameNullUUID(t.workerID, other.workerID) && + t.historyVersion == other.historyVersion && + t.generationAttempt == other.generationAttempt +} diff --git a/coderd/x/chatd/stream_relay_internal_test.go b/coderd/x/chatd/stream_relay_internal_test.go new file mode 100644 index 0000000000..cc334c9f41 --- /dev/null +++ b/coderd/x/chatd/stream_relay_internal_test.go @@ -0,0 +1,25 @@ +package chatd + +import ( + "testing" + + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" +) + +type fakeDialError struct { + unrecoverable bool +} + +func (fakeDialError) Error() string { return "fake dial error" } +func (e fakeDialError) IsUnrecoverable() bool { return e.unrecoverable } + +func TestStreamPartsDialUnrecoverable(t *testing.T) { + t.Parallel() + + require.False(t, streamPartsDialUnrecoverable(nil)) + require.False(t, streamPartsDialUnrecoverable(xerrors.New("plain error"))) + require.False(t, streamPartsDialUnrecoverable(fakeDialError{unrecoverable: false})) + require.True(t, streamPartsDialUnrecoverable(fakeDialError{unrecoverable: true})) + require.True(t, streamPartsDialUnrecoverable(xerrors.Errorf("wrapped: %w", fakeDialError{unrecoverable: true}))) +} diff --git a/coderd/x/chatd/stream_subscribe.go b/coderd/x/chatd/stream_subscribe.go new file mode 100644 index 0000000000..9545b707c1 --- /dev/null +++ b/coderd/x/chatd/stream_subscribe.go @@ -0,0 +1,255 @@ +package chatd + +import ( + "context" + "net/http" + "time" + + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" +) + +const ( + streamSyncRetryInitialBackoff = 100 * time.Millisecond + streamSyncRetryMaxBackoff = time.Second + streamSyncRetryMaxAttempts = 5 +) + +func (p *Server) subscribeStreamLoop( + ctx context.Context, + chat database.Chat, + requestHeader http.Header, + afterMessageID int64, +) ([]codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), bool) { + if p == nil || p.db == nil || p.pubsub == nil { + return nil, nil, nil, false + } + if p.messagePartBuffer == nil { + p.messagePartBuffer = messagepartbuffer.New(messagepartbuffer.Options{Clock: p.clock}) + } + chatID := chat.ID + streamCtx, streamCancel := context.WithCancel(ctx) + events := make(chan codersdk.ChatStreamEvent, 128) + logger := p.logger.With(slog.F("chat_id", chatID)) + + updateCh := make(chan streamSyncHint, 32) + pubsubCancel, err := p.pubsub.SubscribeWithErr( + coderdpubsub.ChatStateUpdateChannel(chatID), + coderdpubsub.HandleChatStateUpdate(func(_ context.Context, payload coderdpubsub.ChatStateUpdateMessage, err error) { + if err != nil { + logger.Warn(streamCtx, "chat stream pubsub error", slog.Error(err)) + return + } + select { + case updateCh <- streamSyncHintFromUpdate(payload): + case <-streamCtx.Done(): + } + }), + ) + if err != nil { + logger.Warn(ctx, "failed to subscribe to chat state updates", slog.Error(err)) + streamCancel() + return subscribeWithInitialError(chatID, "failed to subscribe to chat updates") + } + + pollerCh, unregisterPoller := p.streamSyncPoller.Register(chatID) + loop := newStreamLoop(chat, p.db, logger, afterMessageID) + // The immediate sync builds the initial snapshot returned to the caller + // and the relay target for the forwarder. Hints only fire on state + // changes, so without it an idle chat would never deliver a snapshot and + // an actively streaming chat would not relay parts until the next hint. + //nolint:gocritic // The HTTP route authorizes the chat before subscribing; the stream loop needs chatd-scoped reads for one consistent snapshot. + initial, target, _, err := loop.syncDB(dbauthz.AsChatd(ctx)) + if err != nil { + logger.Error(ctx, "failed to load initial chat stream snapshot", slog.Error(err)) + unregisterPoller() + pubsubCancel() + streamCancel() + return subscribeWithInitialError(chatID, "failed to load initial snapshot") + } + + relay := newStreamRelayForwarder( + chatID, + requestHeader, + p.streamPartsDialer, + p.clock, + logger, + ) + relay.Configure(streamCtx, target) + + done := make(chan struct{}) + go func() { + defer close(done) + defer close(events) + defer relay.Close() + defer unregisterPoller() + for { + select { + case <-streamCtx.Done(): + return + case hint := <-updateCh: + if !p.runStreamSync(streamCtx, loop, relay, events, hint) { + return + } + case hint, ok := <-pollerCh: + if !ok { + return + } + if !p.runStreamSync(streamCtx, loop, relay, events, hint) { + return + } + case part, ok := <-relay.Parts(): + if !ok { + return + } + event, accepted, err := loop.part(part) + if err != nil { + logger.Error(streamCtx, "chat stream invariant violation", slog.Error(err)) + return + } + if accepted { + sendStreamEvent(streamCtx, events, event) + } + } + } + }() + + cancel := func() { + streamCancel() + pubsubCancel() + <-done + } + return initial, events, cancel, true +} + +func (p *Server) runStreamSync( + ctx context.Context, + loop *streamLoop, + relay *streamRelayForwarder, + events chan<- codersdk.ChatStreamEvent, + hint streamSyncHint, +) bool { + syncEvents, target, changed, err := p.syncStreamWithRetry(ctx, loop, hint) + if err != nil { + p.logger.Error(ctx, "failed to sync chat stream after retries", slog.Error(err)) + return false + } + for _, event := range syncEvents { + if !sendStreamEvent(ctx, events, event) { + return false + } + } + if changed { + relay.Configure(ctx, target) + } + return true +} + +func (p *Server) syncStreamWithRetry( + ctx context.Context, + loop *streamLoop, + hint streamSyncHint, +) ([]codersdk.ChatStreamEvent, streamRelayTarget, bool, error) { + var ( + syncEvents []codersdk.ChatStreamEvent + target streamRelayTarget + changed bool + err error + ) + for attempt := 1; attempt <= streamSyncRetryMaxAttempts; attempt++ { + //nolint:gocritic // The subscriber was authorized before the loop started; follow-up syncs need chatd-scoped reads for consistency. + syncEvents, target, changed, err = loop.sync(dbauthz.AsChatd(ctx), hint) + if err == nil || ctx.Err() != nil { + return syncEvents, target, changed, err + } + p.logger.Warn(ctx, "failed to sync chat stream", + slog.F("attempt", attempt), + slog.Error(err), + ) + if attempt == streamSyncRetryMaxAttempts { + break + } + if !p.waitBeforeStreamSyncRetry(ctx, attempt) { + return nil, loop.currentRelayTarget(), false, ctx.Err() + } + } + return nil, loop.currentRelayTarget(), false, err +} + +func (p *Server) waitBeforeStreamSyncRetry(ctx context.Context, attempt int) bool { + delay := streamSyncRetryInitialBackoff + for range attempt - 1 { + delay *= 2 + if delay >= streamSyncRetryMaxBackoff { + delay = streamSyncRetryMaxBackoff + break + } + } + timer := p.clock.NewTimer(delay, "chatd", "stream-sync-retry") + defer timer.Stop() + select { + case <-timer.C: + return true + case <-ctx.Done(): + return false + } +} + +func sendStreamEvent(ctx context.Context, ch chan<- codersdk.ChatStreamEvent, event codersdk.ChatStreamEvent) bool { + select { + case ch <- event: + return true + case <-ctx.Done(): + return false + } +} + +func (p *Server) Subscribe( + ctx context.Context, + chatID uuid.UUID, + requestHeader http.Header, + afterMessageID int64, +) ( + []codersdk.ChatStreamEvent, + <-chan codersdk.ChatStreamEvent, + func(), + bool, +) { + if p == nil { + return nil, nil, nil, false + } + + chat, err := p.db.GetChatByID(ctx, chatID) + if err != nil { + if dbauthz.IsNotAuthorizedError(err) { + return nil, nil, nil, false + } + p.logger.Warn(ctx, "failed to load chat for stream subscription", + slog.F("chat_id", chatID), + slog.Error(err), + ) + return subscribeWithInitialError(chatID, "failed to load initial snapshot") + } + return p.SubscribeAuthorized(ctx, chat, requestHeader, afterMessageID) +} + +// SubscribeAuthorized subscribes an already-authorized chat to stream updates. +func (p *Server) SubscribeAuthorized( + ctx context.Context, + chat database.Chat, + requestHeader http.Header, + afterMessageID int64, +) ( + []codersdk.ChatStreamEvent, + <-chan codersdk.ChatStreamEvent, + func(), + bool, +) { + return p.subscribeStreamLoop(ctx, chat, requestHeader, afterMessageID) +} diff --git a/coderd/x/chatd/stream_sync_poller.go b/coderd/x/chatd/stream_sync_poller.go new file mode 100644 index 0000000000..11e9171687 --- /dev/null +++ b/coderd/x/chatd/stream_sync_poller.go @@ -0,0 +1,167 @@ +package chatd + +import ( + "context" + "sync" + "time" + + "github.com/google/uuid" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/quartz" +) + +const streamSyncInterval = 10 * time.Second + +type streamSyncPoller struct { + ctx context.Context + cancel context.CancelFunc + db database.Store + clock quartz.Clock + logger slog.Logger + + mu sync.Mutex + subscribers map[uuid.UUID]map[*streamSyncPollerSubscriber]struct{} +} + +type streamSyncPollerSubscriber struct { + chatID uuid.UUID + hints chan streamSyncHint +} + +func newStreamSyncPoller( + ctx context.Context, + db database.Store, + clock quartz.Clock, + logger slog.Logger, +) *streamSyncPoller { + if clock == nil { + clock = quartz.NewReal() + } + //nolint:gocritic // The poller is internal chatd infrastructure. Each + // registered stream was already authorized before subscription, and this + // batch query only fetches synchronization metadata for subscribed chats. + pollerCtx, cancel := context.WithCancel(dbauthz.AsChatd(ctx)) + return &streamSyncPoller{ + ctx: pollerCtx, + cancel: cancel, + db: db, + clock: clock, + logger: logger, + subscribers: make(map[uuid.UUID]map[*streamSyncPollerSubscriber]struct{}), + } +} + +func (p *streamSyncPoller) Start() { + if p == nil { + return + } + go p.loop() +} + +func (p *streamSyncPoller) Close() { + if p == nil { + return + } + p.cancel() +} + +func (p *streamSyncPoller) Register(chatID uuid.UUID) (<-chan streamSyncHint, func()) { + if p == nil { + ch := make(chan streamSyncHint) + close(ch) + return ch, func() {} + } + subscriber := &streamSyncPollerSubscriber{ + chatID: chatID, + hints: make(chan streamSyncHint, 1), + } + p.mu.Lock() + if p.subscribers[chatID] == nil { + p.subscribers[chatID] = make(map[*streamSyncPollerSubscriber]struct{}) + } + p.subscribers[chatID][subscriber] = struct{}{} + p.mu.Unlock() + + return subscriber.hints, func() { + p.unregister(subscriber) + } +} + +func (p *streamSyncPoller) unregister(subscriber *streamSyncPollerSubscriber) { + p.mu.Lock() + defer p.mu.Unlock() + chatSubscribers := p.subscribers[subscriber.chatID] + if chatSubscribers == nil { + return + } + delete(chatSubscribers, subscriber) + if len(chatSubscribers) == 0 { + delete(p.subscribers, subscriber.chatID) + } + close(subscriber.hints) +} + +func (p *streamSyncPoller) loop() { + ticker := p.clock.NewTicker(streamSyncInterval, "chatd", "stream-sync-poller") + defer ticker.Stop() + for { + select { + case <-p.ctx.Done(): + return + case <-ticker.C: + p.pollOnce() + } + } +} + +func (p *streamSyncPoller) pollOnce() { + chatIDs, subscribers := p.snapshotSubscribers() + if len(chatIDs) == 0 { + return + } + rows, err := p.db.GetChatStreamSyncRows(p.ctx, chatIDs) + if err != nil { + if p.ctx.Err() == nil { + p.logger.Warn(p.ctx, "failed to poll chat streams", slog.Error(err)) + } + return + } + for _, row := range rows { + hint := streamSyncHintFromPollRow(row) + for _, subscriber := range subscribers[row.ID] { + select { + case subscriber.hints <- hint: + default: + } + } + } +} + +func (p *streamSyncPoller) snapshotSubscribers() ([]uuid.UUID, map[uuid.UUID][]*streamSyncPollerSubscriber) { + p.mu.Lock() + defer p.mu.Unlock() + chatIDs := make([]uuid.UUID, 0, len(p.subscribers)) + subscribers := make(map[uuid.UUID][]*streamSyncPollerSubscriber, len(p.subscribers)) + for chatID, chatSubscribers := range p.subscribers { + chatIDs = append(chatIDs, chatID) + for subscriber := range chatSubscribers { + subscribers[chatID] = append(subscribers[chatID], subscriber) + } + } + return chatIDs, subscribers +} + +func streamSyncHintFromPollRow(row database.GetChatStreamSyncRowsRow) streamSyncHint { + return streamSyncHint{ + snapshotVersion: row.SnapshotVersion, + historyVersion: row.HistoryVersion, + queueVersion: row.QueueVersion, + retryVersion: row.RetryStateVersion, + status: row.Status, + workerID: row.WorkerID, + generationAttempt: row.GenerationAttempt, + } +} diff --git a/coderd/x/chatd/stream_types.go b/coderd/x/chatd/stream_types.go new file mode 100644 index 0000000000..d413079ff9 --- /dev/null +++ b/coderd/x/chatd/stream_types.go @@ -0,0 +1,44 @@ +package chatd + +import ( + "context" + "net/http" + + "github.com/google/uuid" + + "github.com/coder/coder/v2/codersdk" +) + +// StreamPartsDialer dials an episode-aware source of message parts. +type StreamPartsDialer func(ctx context.Context, input StreamPartsDialInput) (StreamPartsSession, error) + +// StreamPartsDialInput carries the metadata needed to dial a parts source. +type StreamPartsDialInput struct { + ChatID uuid.UUID + WorkerID uuid.UUID + RequestHeader http.Header +} + +// StreamPartsSession streams message parts for selected episodes. +type StreamPartsSession interface { + SelectEpisode(ctx context.Context, historyVersion, generationAttempt int64) error + Parts() <-chan StreamPart + Close() error +} + +// StreamPart is a live preview part scoped to one chat history episode. +type StreamPart struct { + HistoryVersion int64 + GenerationAttempt int64 + Seq int64 + Role codersdk.ChatMessageRole + Part codersdk.ChatMessagePart +} + +type streamPart = StreamPart + +type streamRelayTarget struct { + workerID uuid.NullUUID + historyVersion int64 + generationAttempt int64 +} diff --git a/coderd/x/chatd/streamcollector_internal_test.go b/coderd/x/chatd/streamcollector_internal_test.go deleted file mode 100644 index 81dae5f133..0000000000 --- a/coderd/x/chatd/streamcollector_internal_test.go +++ /dev/null @@ -1,216 +0,0 @@ -package chatd - -import ( - "sync" - "testing" - "time" - - "github.com/google/uuid" - "github.com/prometheus/client_golang/prometheus" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" - - "cdr.dev/slog/v3" - "github.com/coder/coder/v2/coderd/x/chatd/chatloop" - "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/testutil" - "github.com/coder/quartz" -) - -// TestStreamStateCollector exercises the four gauges emitted by -// streamStateCollector against representative map states. -func TestStreamStateCollector(t *testing.T) { - t.Parallel() - - t.Run("EmptyMap", func(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - server := &Server{} - reg.MustRegister(&streamStateCollector{server: server}) - - assertGauges(t, reg, gaugeExpectations{ - active: 0, - bufferMax: 0, - bufferTotal: 0, - subscribers: 0, - }) - }) - - t.Run("PopulatedMap", func(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - server := &Server{} - - server.chatStreams.Store(uuid.New(), &chatStreamState{ - buffer: make([]bufferedStreamPart, 10), - subscribers: newSubscribers(t, 2), - }) - server.chatStreams.Store(uuid.New(), &chatStreamState{ - buffer: make([]bufferedStreamPart, 25), - subscribers: map[uuid.UUID]chan codersdk.ChatStreamEvent{}, - }) - server.chatStreams.Store(uuid.New(), &chatStreamState{ - buffer: nil, - subscribers: newSubscribers(t, 1), - }) - - reg.MustRegister(&streamStateCollector{server: server}) - - assertGauges(t, reg, gaugeExpectations{ - active: 3, - bufferMax: 25, - bufferTotal: 35, - subscribers: 3, - }) - }) - - t.Run("SkipsWrongType", func(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - server := &Server{} - - server.chatStreams.Store(uuid.New(), "garbage") - server.chatStreams.Store(uuid.New(), &chatStreamState{ - buffer: make([]bufferedStreamPart, 5), - subscribers: newSubscribers(t, 1), - }) - - reg.MustRegister(&streamStateCollector{server: server}) - - // The non-matching entry is silently skipped. Only the - // valid chatStreamState counts. - assertGauges(t, reg, gaugeExpectations{ - active: 1, - bufferMax: 5, - bufferTotal: 5, - subscribers: 1, - }) - }) - - // Runs Collect concurrently with state.mu mutations; catches - // missing lock acquisition under `go test -race`. - t.Run("LockContentionSmoke", func(t *testing.T) { - t.Parallel() - - server := &Server{} - state := &chatStreamState{ - buffer: make([]bufferedStreamPart, 0, 100), - subscribers: newSubscribers(t, 1), - } - server.chatStreams.Store(uuid.New(), state) - collector := &streamStateCollector{server: server} - - const iterations = 100 - var wg sync.WaitGroup - - // Mutator: grows and shrinks the buffer under state.mu. - wg.Go(func() { - for range iterations { - state.mu.Lock() - state.buffer = append(state.buffer, bufferedStreamPart{}) - if len(state.buffer) > 50 { - state.buffer = state.buffer[10:] - } - state.mu.Unlock() - } - }) - - // Scraper: repeatedly invokes Collect into a discard - // channel. A panic or race here fails the test. - wg.Go(func() { - ctx := testutil.Context(t, 10*time.Second) - for range iterations { - ch := make(chan prometheus.Metric, 4) - collector.Collect(ch) - // Drain all metrics the collector wrote. - for range 4 { - testutil.SoftTryReceive(ctx, t, ch) - } - } - }) - - wg.Wait() - }) -} - -type gaugeExpectations struct { - active float64 - bufferMax float64 - bufferTotal float64 - subscribers float64 -} - -func assertGauges(t *testing.T, reg *prometheus.Registry, want gaugeExpectations) { - t.Helper() - families, err := reg.Gather() - require.NoError(t, err) - - got := map[string]float64{} - for _, f := range families { - require.Len(t, f.GetMetric(), 1, "metric %q should have exactly one sample", f.GetName()) - got[f.GetName()] = f.GetMetric()[0].GetGauge().GetValue() - } - - assert.Equal(t, want.active, got["coderd_chatd_streams_active"], "streams_active") - assert.Equal(t, want.bufferMax, got["coderd_chatd_stream_buffer_size_max"], "buffer_size_max") - assert.Equal(t, want.bufferTotal, got["coderd_chatd_stream_buffer_events"], "buffer_events") - assert.Equal(t, want.subscribers, got["coderd_chatd_stream_subscribers"], "subscribers") -} - -func newSubscribers(t *testing.T, n int) map[uuid.UUID]chan codersdk.ChatStreamEvent { - t.Helper() - subs := make(map[uuid.UUID]chan codersdk.ChatStreamEvent, n) - for range n { - subs[uuid.New()] = make(chan codersdk.ChatStreamEvent, 1) - } - return subs -} - -// TestStreamStateCollector_BufferDroppedIncrementsOnCapacity pre-fills -// a buffer to capacity and asserts stream_buffer_dropped_total -// increments on each subsequent publishToStream drop. -func TestStreamStateCollector_BufferDroppedIncrementsOnCapacity(t *testing.T) { - t.Parallel() - - reg := prometheus.NewRegistry() - server := &Server{ - logger: slog.Make(), - clock: quartz.NewMock(t), - metrics: chatloop.NewMetrics(reg), - } - - chatID := uuid.New() - server.chatStreams.Store(chatID, &chatStreamState{ - buffering: true, - buffer: make([]bufferedStreamPart, maxStreamBufferSize), - }) - - partEvent := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{}, - } - - server.publishToStream(chatID, partEvent) - assert.Equal(t, float64(1), counterValue(t, reg, "coderd_chatd_stream_buffer_dropped_total")) - - server.publishToStream(chatID, partEvent) - assert.Equal(t, float64(2), counterValue(t, reg, "coderd_chatd_stream_buffer_dropped_total")) -} - -func counterValue(t *testing.T, reg *prometheus.Registry, name string) float64 { - t.Helper() - families, err := reg.Gather() - require.NoError(t, err) - for _, f := range families { - if f.GetName() != name { - continue - } - require.Len(t, f.GetMetric(), 1, "counter %q should have exactly one sample", name) - return f.GetMetric()[0].GetCounter().GetValue() - } - t.Fatalf("counter %q not registered", name) - return 0 -} diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 450397416b..a735f0b3f4 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -20,9 +20,11 @@ import ( "github.com/coder/coder/v2/coderd/aibridge" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbauthz" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/workspacesdk" ) @@ -909,10 +911,10 @@ func (p *Server) resolveExploreToolSnapshot( return inheritedMCPServerIDs, nil } -func (p *Server) delegatedAPIKeyIDForSubagent(ctx context.Context) (string, error) { +func (*Server) delegatedAPIKeyIDForSubagent(ctx context.Context) (string, error) { apiKeyID, ok := aibridge.DelegatedAPIKeyIDFromContext(ctx) - if !ok && p.shouldUseAIGatewayRouting() { - return "", xerrors.New("AI Gateway routing requires the active turn API key ID for subagent messages") + if !ok || apiKeyID == "" { + return "", xerrors.New("active turn API key ID is required for subagent messages") } return apiKeyID, nil } @@ -987,147 +989,114 @@ func (p *Server) createChildSubagentChatWithOptions( // for another pool checkout. deploymentPrompt := p.resolveDeploymentSystemPrompt(ctx) - var child database.Chat - txErr := p.db.InTx(func(tx database.Store) error { - if limitErr := p.checkUsageLimit(ctx, tx, parent.OwnerID, uuid.NullUUID{UUID: parent.OrganizationID, Valid: true}); limitErr != nil { - return limitErr - } - - insertedChat, err := tx.InsertChat(ctx, database.InsertChatParams{ - OrganizationID: parent.OrganizationID, - OwnerID: parent.OwnerID, - WorkspaceID: parent.WorkspaceID, - BuildID: parent.BuildID, - AgentID: parent.AgentID, - ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, - RootChatID: uuid.NullUUID{UUID: rootChatID, Valid: true}, - LastModelConfigID: modelConfigID, - Title: title, - Mode: opts.chatMode, - PlanMode: childPlanMode, - ClientType: parent.ClientType, - Status: database.ChatStatusPending, - MCPServerIDs: mcpServerIDs, - Labels: pqtype.NullRawMessage{ - RawMessage: labelsJSON, - Valid: true, - }, - DynamicTools: pqtype.NullRawMessage{}, - }) - if err != nil { - return xerrors.Errorf("insert child chat: %w", err) - } - - workspaceAwareness := workspaceDetachedNoCreateAwareness - if insertedChat.WorkspaceID.Valid { - workspaceAwareness = workspaceAttachedAwareness - } - workspaceAwarenessContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(workspaceAwareness), - }) - if err != nil { - return xerrors.Errorf("marshal workspace awareness: %w", err) - } - userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)}) - if err != nil { - return xerrors.Errorf("marshal initial user content: %w", err) - } - - systemParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendChatMessage. - ChatID: insertedChat.ID, - } - if deploymentPrompt != "" { - deploymentContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(deploymentPrompt), - }) - if err != nil { - return xerrors.Errorf("marshal deployment system prompt: %w", err) - } - appendChatMessage(&systemParams, newChatMessage( - database.ChatMessageRoleSystem, - deploymentContent, - database.ChatMessageVisibilityModel, - modelConfigID, - chatprompt.CurrentContentVersion, - )) - } - if childSystemPrompt != "" { - childSystemPromptContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ - codersdk.ChatMessageText(childSystemPrompt), - }) - if err != nil { - return xerrors.Errorf("marshal child system prompt: %w", err) - } - appendChatMessage(&systemParams, newChatMessage( - database.ChatMessageRoleSystem, - childSystemPromptContent, - database.ChatMessageVisibilityModel, - modelConfigID, - chatprompt.CurrentContentVersion, - )) - } - appendChatMessage(&systemParams, newChatMessage( - database.ChatMessageRoleSystem, - workspaceAwarenessContent, - database.ChatMessageVisibilityModel, - modelConfigID, - chatprompt.CurrentContentVersion, - )) - if _, err := tx.InsertChatMessages(ctx, systemParams); err != nil { - return xerrors.Errorf("insert initial child system messages: %w", err) - } - - child = insertedChat - - // Copy persisted context before the initial child prompt so the - // child cannot be acquired until its inherited context is in - // place. signalWake runs only after commit. - copiedContextParts, err := copyParentContextMessages(ctx, p.logger, tx, parent, child) - if err != nil { - return xerrors.Errorf("copy parent context messages: %w", err) - } - if err := updateChildLastInjectedContext(ctx, p.logger, tx, child.ID, copiedContextParts); err != nil { - return xerrors.Errorf("update child injected context: %w", err) - } - - userParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by appendUserChatMessage. - ChatID: insertedChat.ID, - } - childUserMsg := newUserChatMessage( - childAPIKeyID, - userContent, - database.ChatMessageVisibilityBoth, - modelConfigID, - chatprompt.CurrentContentVersion, - ) - childUserMsg = childUserMsg.withCreatedBy(parent.OwnerID) - appendUserChatMessage(&userParams, childUserMsg) - if _, err := tx.InsertChatMessages(ctx, userParams); err != nil { - return xerrors.Errorf("insert initial child user message: %w", err) - } - - return nil - }, nil) - if txErr != nil { - return database.Chat{}, xerrors.Errorf("create child chat: %w", txErr) + if limitErr := p.checkUsageLimit(ctx, p.db, parent.OwnerID, uuid.NullUUID{UUID: parent.OrganizationID, Valid: true}); limitErr != nil { + return database.Chat{}, limitErr } + workspaceAwareness := workspaceDetachedNoCreateAwareness + if parent.WorkspaceID.Valid { + workspaceAwareness = workspaceAttachedAwareness + } + workspaceAwarenessContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(workspaceAwareness), + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal workspace awareness: %w", err) + } + userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)}) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal initial user content: %w", err) + } + + initialMessages := make([]chatstate.Message, 0, 4) + if deploymentPrompt != "" { + deploymentContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(deploymentPrompt), + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal deployment system prompt: %w", err) + } + initialMessages = append(initialMessages, systemMessage(deploymentContent, modelConfigID)) + } + if childSystemPrompt != "" { + childSystemPromptContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText(childSystemPrompt), + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal child system prompt: %w", err) + } + initialMessages = append(initialMessages, systemMessage(childSystemPromptContent, modelConfigID)) + } + initialMessages = append(initialMessages, systemMessage(workspaceAwarenessContent, modelConfigID)) + + copiedContextParts, err := copyParentContextMessages(ctx, p.logger, p.db, parent) + if err != nil { + return database.Chat{}, xerrors.Errorf("copy parent context messages: %w", err) + } + var lastInjectedContext pqtype.NullRawMessage + if len(copiedContextParts) > 0 { + filteredContent, err := chatprompt.MarshalParts(copiedContextParts) + if err != nil { + return database.Chat{}, xerrors.Errorf("marshal copied context parts: %w", err) + } + initialMessages = append(initialMessages, userMessageWithAPIKeyID( + filteredContent, + modelConfigID, + parent.OwnerID, + childAPIKeyID, + )) + lastInjectedContext, err = BuildLastInjectedContext(FilterContextPartsToLatestAgent(copiedContextParts)) + if err != nil { + return database.Chat{}, xerrors.Errorf("build inherited injected context: %w", err) + } + } + initialMessages = append(initialMessages, userMessageWithAPIKeyID(userContent, modelConfigID, parent.OwnerID, childAPIKeyID)) + + publisher := p.pubsub + if publisher == nil { + publisher = dbpubsub.NewInMemory() + } + result, err := chatstate.CreateChat(ctx, p.db, publisher, chatstate.CreateChatInput{ + OrganizationID: parent.OrganizationID, + OwnerID: parent.OwnerID, + WorkspaceID: parent.WorkspaceID, + BuildID: parent.BuildID, + AgentID: parent.AgentID, + ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: rootChatID, Valid: true}, + LastModelConfigID: modelConfigID, + Title: title, + Mode: opts.chatMode, + PlanMode: childPlanMode, + MCPServerIDs: mcpServerIDs, + Labels: pqtype.NullRawMessage{ + RawMessage: labelsJSON, + Valid: true, + }, + DynamicTools: pqtype.NullRawMessage{}, + ClientType: parent.ClientType, + InitialMessages: initialMessages, + LastInjectedContext: lastInjectedContext, + }) + if err != nil { + return database.Chat{}, xerrors.Errorf("create child chat: %w", err) + } + + child := result.Chat + p.publishChatPubsubEvent(child, codersdk.ChatWatchEventKindCreated, nil) - p.signalWake() return child, nil } // copyParentContextMessages reads persisted context-file and skill -// messages from the parent chat and inserts copies into the child -// chat. This ensures sub-agents inherit the same instruction and -// skill context as their parent without independently re-fetching -// from the agent. +// messages from the parent chat. This ensures sub-agents inherit the +// same instruction and skill context as their parent without +// independently re-fetching from the agent. func copyParentContextMessages( ctx context.Context, logger slog.Logger, store database.Store, parent database.Chat, - child database.Chat, ) ([]codersdk.ChatMessagePart, error) { parentMessages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: parent.ID, @@ -1137,12 +1106,7 @@ func copyParentContextMessages( return nil, xerrors.Errorf("get parent messages: %w", err) } - var ( - copiedParts []codersdk.ChatMessagePart - copiedRole database.ChatMessageRole - copiedVisibility database.ChatMessageVisibility - copiedVersion int16 - ) + var copiedParts []codersdk.ChatMessagePart for _, msg := range parentMessages { if !msg.Content.Valid { continue @@ -1161,11 +1125,6 @@ func copyParentContextMessages( if len(messageContextParts) == 0 { continue } - if copiedParts == nil { - copiedRole = msg.Role - copiedVisibility = msg.Visibility - copiedVersion = msg.ContentVersion - } copiedParts = append(copiedParts, messageContextParts...) } if len(copiedParts) == 0 { @@ -1173,69 +1132,10 @@ func copyParentContextMessages( } copiedParts = FilterContextPartsToLatestAgent(copiedParts) - filteredContent, err := chatprompt.MarshalParts(copiedParts) - if err != nil { - return nil, xerrors.Errorf("marshal filtered context parts: %w", err) - } - - msgParams := database.InsertChatMessagesParams{ //nolint:exhaustruct // Fields populated by append[User]ChatMessage. - ChatID: child.ID, - } - if copiedRole == database.ChatMessageRoleUser { - copiedAPIKeyID, _ := aibridge.DelegatedAPIKeyIDFromContext(ctx) - appendUserChatMessage(&msgParams, newUserChatMessage( - copiedAPIKeyID, - filteredContent, - copiedVisibility, - child.LastModelConfigID, - copiedVersion, - )) - } else { - appendChatMessage(&msgParams, newChatMessage( - copiedRole, - filteredContent, - copiedVisibility, - child.LastModelConfigID, - copiedVersion, - )) - } - if _, err := store.InsertChatMessages(ctx, msgParams); err != nil { - return nil, xerrors.Errorf("insert context message: %w", err) - } return copiedParts, nil } -func updateChildLastInjectedContext( - ctx context.Context, - logger slog.Logger, - store database.Store, - chatID uuid.UUID, - parts []codersdk.ChatMessagePart, -) error { - parts = FilterContextPartsToLatestAgent(parts) - param, err := BuildLastInjectedContext(parts) - if err != nil { - logger.Warn(ctx, "failed to marshal inherited injected context", - slog.F("chat_id", chatID), - slog.Error(err), - ) - return xerrors.Errorf("marshal inherited injected context: %w", err) - } - if _, err := store.UpdateChatLastInjectedContext(ctx, database.UpdateChatLastInjectedContextParams{ - ID: chatID, - LastInjectedContext: param, - }); err != nil { - logger.Warn(ctx, "failed to update inherited injected context", - slog.F("chat_id", chatID), - slog.Error(err), - ) - return xerrors.Errorf("update inherited injected context: %w", err) - } - - return nil -} - func (p *Server) sendSubagentMessage( ctx context.Context, parentChatID uuid.UUID, @@ -1310,34 +1210,29 @@ func (p *Server) awaitSubagentCompletion( timer := p.clock.NewTimer(timeout, "chatd", "subagent_await") defer timer.Stop() - // When pubsub is available, subscribe for fast status - // notifications and use a less aggressive fallback poll. - // Without pubsub (single-instance / in-memory) fall back - // to the original 200ms polling. - pollInterval := subagentAwaitPollInterval - var notifyCh <-chan struct{} - if p.pubsub != nil { - pollInterval = subagentAwaitFallbackPoll - ch := make(chan struct{}, 1) - notifyCh = ch - cancel, subErr := p.pubsub.SubscribeWithErr( - coderdpubsub.ChatStreamNotifyChannel(targetChatID), - func(_ context.Context, _ []byte, _ error) { - // Non-blocking send so we never stall the - // pubsub dispatch goroutine. - select { - case ch <- struct{}{}: - default: - } - }, - ) - if subErr == nil { - defer cancel() - } else { - // Subscription failed; fall back to fast polling. - pollInterval = subagentAwaitPollInterval - notifyCh = nil - } + // Subscribe for fast status notifications and use a less + // aggressive fallback poll. If subscription fails, fall back to + // the original 200ms polling. + pollInterval := subagentAwaitFallbackPoll + ch := make(chan struct{}, 1) + notifyCh := (<-chan struct{})(ch) + cancel, subErr := p.pubsub.SubscribeWithErr( + coderdpubsub.ChatStateUpdateChannel(targetChatID), + func(_ context.Context, _ []byte, _ error) { + // Non-blocking send so we never stall the + // pubsub dispatch goroutine. + select { + case ch <- struct{}{}: + default: + } + }, + ) + if subErr == nil { + defer cancel() + } else { + // Subscription failed; fall back to fast polling. + pollInterval = subagentAwaitPollInterval + notifyCh = nil } ticker := p.clock.NewTicker(pollInterval, "chatd", "subagent_poll") @@ -1401,10 +1296,18 @@ func (p *Server) closeSubagent( return targetChat, nil } - updatedChat := p.InterruptChat(ctx, targetChat) - if updatedChat.Status != database.ChatStatusWaiting { - return database.Chat{}, xerrors.New("set target chat waiting") + updatedChat, err := p.InterruptChat(ctx, targetChat) + if err != nil { + // Idle / archived chats no longer satisfy the + // chatstate.Interrupt precondition. Surface the error + // so the caller can decide whether the parent expected + // the subagent to already be waiting. + return database.Chat{}, xerrors.Errorf("interrupt subagent chat: %w", err) } + // chatstate.Interrupt lands active runs in `interrupting` + // and requires-action chats in `running`. Workers finalize + // the transition; accept either non-active status as long as + // the transition committed. return updatedChat, nil } diff --git a/coderd/x/chatd/subagent_context_internal_test.go b/coderd/x/chatd/subagent_context_internal_test.go index 5ccab312d6..d56bdbbcb0 100644 --- a/coderd/x/chatd/subagent_context_internal_test.go +++ b/coderd/x/chatd/subagent_context_internal_test.go @@ -155,6 +155,7 @@ func createParentChatWithInheritedContext( parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-with-context", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -324,6 +325,7 @@ func createParentChatWithRotatedInheritedContext( parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-with-rotated-context", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -400,6 +402,7 @@ func TestCreateChildSubagentChatCopiesOnlyLatestAgentContext(t *testing.T) { ctx := chatdTestContext(t) parentChat := createParentChatWithRotatedInheritedContext(ctx, t, db, server) + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) child, err := server.createChildSubagentChat(ctx, parentChat, "inspect bindings", "") require.NoError(t, err) @@ -493,6 +496,7 @@ func TestSpawnComputerUseAgentInheritsContext(t *testing.T) { // before the Anthropic provider was inserted. server.configCache.InvalidateProviders() + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, db, parentChat.OwnerID)) tools := server.subagentTools(ctx, func() database.Chat { return parentChat }, parentChat.LastModelConfigID) tool := findToolByName(tools, spawnAgentToolName) require.NotNil(t, tool) diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index ce860f1249..f575fe2383 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -4,6 +4,8 @@ import ( "context" "database/sql" "encoding/json" + "net/http" + "net/http/httptest" "sync" "testing" "time" @@ -13,6 +15,7 @@ import ( "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/xerrors" "cdr.dev/slog/v3" "cdr.dev/slog/v3/sloggers/slogtest" @@ -122,11 +125,10 @@ func newInternalTestServerWithLoggerAndClock( ) *Server { t.Helper() - server := New(Config{ + server := New(ps, Config{ Logger: logger, Database: db, ReplicaID: uuid.New(), - Pubsub: ps, Clock: clk, // Use a very long interval so the background loop // does not interfere with test assertions. @@ -140,6 +142,18 @@ func newInternalTestServerWithLoggerAndClock( return server } +type subscribeFailingPubsub struct { + pubsub.Pubsub +} + +func (subscribeFailingPubsub) Subscribe(_ string, _ pubsub.Listener) (func(), error) { + return nil, xerrors.New("subscribe disabled") +} + +func (subscribeFailingPubsub) SubscribeWithErr(_ string, _ pubsub.ListenerWithErr) (func(), error) { + return nil, xerrors.New("subscribe disabled") +} + type subagentTestLogSink struct { mu sync.Mutex entries []slog.SinkEntry @@ -180,6 +194,7 @@ func seedInternalChatDeps( t.Helper() user := dbgen.User(t, db, database.User{}) + _ = testAPIKeyID(t, db, user.ID) org := dbgen.Organization(t, db, database.Organization{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -289,6 +304,7 @@ func TestSendSubagentMessagePropagatesActiveTurnAPIKeyID(t *testing.T) { child, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, Title: "child-send-subagent-key", @@ -345,7 +361,7 @@ func TestCreateChildSubagentChatRequiresActiveTurnAPIKeyIDForAIGateway(t *testin aiGatewayRoutingEnabled: true, } _, err := server.createChildSubagentChat(ctx, parent, "inspect the workspace", "") - require.ErrorContains(t, err, "AI Gateway routing requires the active turn API key ID for subagent messages") + require.ErrorContains(t, err, "active turn API key ID is required for subagent messages") } func TestSendSubagentMessageRequiresActiveTurnAPIKeyIDForAIGateway(t *testing.T) { @@ -360,6 +376,7 @@ func TestSendSubagentMessageRequiresActiveTurnAPIKeyIDForAIGateway(t *testing.T) parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-send-subagent-missing-key", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -368,6 +385,7 @@ func TestSendSubagentMessageRequiresActiveTurnAPIKeyIDForAIGateway(t *testing.T) child, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true}, Title: "child-send-subagent-missing-key", @@ -386,7 +404,66 @@ func TestSendSubagentMessageRequiresActiveTurnAPIKeyIDForAIGateway(t *testing.T) "follow up", SendMessageBusyBehaviorInterrupt, ) - require.ErrorContains(t, err, "AI Gateway routing requires the active turn API key ID for subagent messages") + require.ErrorContains(t, err, "active turn API key ID is required for subagent messages") +} + +// TestSpawnAgentUsesActiveTurnAPIKeyIDFromContext verifies that, with AI +// Gateway routing enabled, the spawn_agent tool succeeds when the active +// turn's delegated API key ID is present on the context and fails without +// it. The generation worker supplies that key by enriching the tool +// execution context with withActiveTurnAPIKeyID, derived from the prompt +// rows' model build options. This guards the regression where +// executeLocalTools passed an un-enriched context to tool callbacks, +// breaking subagent spawning under AI Gateway routing. +func TestSpawnAgentUsesActiveTurnAPIKeyIDFromContext(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + server.aiGatewayRoutingEnabled = true + + ctx := chatdTestContext(t) + user, org, model := seedInternalChatDeps(t, db) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + + parent, err := server.CreateChat(ctx, CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "parent-active-turn-key", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + APIKeyID: apiKey.ID, + }) + require.NoError(t, err) + parentChat, err := db.GetChatByID(ctx, parent.ID) + require.NoError(t, err) + + // The generation worker derives model build options from the prompt + // rows; this is the source executeLocalTools uses to enrich the tool + // execution context. + promptRows, err := server.db.GetChatMessagesForPromptByChatID(ctx, parentChat.ID) + require.NoError(t, err) + modelOpts := modelBuildOptionsFromMessages(promptRows) + require.Equal(t, apiKey.ID, modelOpts.ActiveAPIKeyID) + + // Without the delegated key on the context the spawn fails, matching + // the original un-enriched executeLocalTools behavior. + resp := runSpawnAgentTool(ctx, t, server, parentChat, spawnAgentArgs{ + Type: subagentTypeGeneral, + Prompt: "delegate work", + }) + require.True(t, resp.IsError, "expected error without active turn key, got: %s", resp.Content) + require.Contains(t, resp.Content, "active turn API key ID is required for subagent messages") + + // With the key on the context (as withActiveTurnAPIKeyID supplies in + // executeLocalTools), the spawn succeeds. + enrichedCtx := withActiveTurnAPIKeyID(ctx, modelOpts) + resp = runSpawnAgentTool(enrichedCtx, t, server, parentChat, spawnAgentArgs{ + Type: subagentTypeGeneral, + Prompt: "delegate work", + }) + result := requireSpawnAgentResponse(t, resp) + require.Equal(t, subagentTypeGeneral, result.SubagentType) } func TestResolveUserProviderAPIKeys_AIProvider(t *testing.T) { @@ -768,6 +845,7 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), WorkspaceID: uuid.NullUUID{ UUID: workspace.ID, Valid: true, @@ -789,7 +867,8 @@ func TestCreateChildSubagentChatInheritsWorkspaceBinding(t *testing.T) { parentChat, err := db.GetChatByID(ctx, parent.ID) require.NoError(t, err) - child, err := server.createChildSubagentChat(ctx, parentChat, "inspect bindings", "") + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) + child, err := server.createChildSubagentChatWithOptions(ctx, parentChat, "inspect bindings", "", childSubagentChatOptions{}) require.NoError(t, err) childChat, err := db.GetChatByID(ctx, child.ID) @@ -815,6 +894,7 @@ func createInternalParentChat( parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: orgID, OwnerID: userID, + APIKeyID: testAPIKeyID(t, db, userID), Title: title, ModelConfigID: modelConfigID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -837,6 +917,11 @@ func runSubagentTool( args any, ) fantasy.ToolResponse { t.Helper() + if !server.shouldUseAIGatewayRouting() { + if apiKeyID, ok := aibridge.DelegatedAPIKeyIDFromContext(ctx); !ok || apiKeyID == "" { + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) + } + } tools := server.subagentTools( ctx, @@ -939,6 +1024,7 @@ func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-parent", ModelConfigID: model.ID, PlanMode: planMode, @@ -952,7 +1038,8 @@ func TestCreateChildSubagentChatCopiesPlanMode(t *testing.T) { require.NoError(t, err) require.Equal(t, planMode, parentChat.PlanMode) - child, err := server.createChildSubagentChat(ctx, parentChat, "inspect bindings", "") + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) + child, err := server.createChildSubagentChatWithOptions(ctx, parentChat, "inspect bindings", "", childSubagentChatOptions{}) require.NoError(t, err) childChat, err := db.GetChatByID(ctx, child.ID) @@ -1244,6 +1331,7 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenCredentialsUnavailable(t parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-general-credentials-fallback", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -1313,6 +1401,7 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenProviderDisabled(t *testi parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-general-disabled-provider-fallback", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ @@ -1408,6 +1497,7 @@ func TestCreateChildSubagentChat_OverrideWorksWhenParentHasNoModel(t *testing.T) // The chats table enforces a foreign key for last_model_config_id, so // use a synthetic parent value here to exercise the override path. parentChat.LastModelConfigID = uuid.Nil + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) child, err := server.createChildSubagentChatWithOptions( ctx, parentChat, @@ -1716,6 +1806,7 @@ func TestCreateChat_ExploreRootStartsWithoutMCPSnapshot(t *testing.T) { root, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root-explore", ModelConfigID: model.ID, ChatMode: database.NullChatMode{ @@ -1834,6 +1925,7 @@ func TestCreateChildSubagentChatWithOptions_ExplorePersistsMCPSnapshot(t *testin t, db, user.ID, "snapshot-"+uuid.NewString(), false, ) + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) child, err := server.createChildSubagentChatWithOptions( ctx, parentChat, @@ -1872,6 +1964,7 @@ func TestSpawnAgent_ExploreSnapshotsTurnStateParentState(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-turn-state-snapshot", ModelConfigID: model.ID, MCPServerIDs: []uuid.UUID{turnStartConfig.ID}, @@ -1884,6 +1977,7 @@ func TestSpawnAgent_ExploreSnapshotsTurnStateParentState(t *testing.T) { turnParent, err := db.GetChatByID(ctx, parent.ID) require.NoError(t, err) + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, db, user.ID)) tools := server.subagentTools( ctx, func() database.Chat { return turnParent }, @@ -2144,6 +2238,7 @@ func TestSpawnAgent_PlanModeDescriptionOmitsComputerUse(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-parent-description", ModelConfigID: model.ID, PlanMode: database.NullChatPlanMode{ @@ -2184,6 +2279,7 @@ func TestSpawnAgent_PlanModeRejectsComputerUse(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "plan-parent-computer-use-reject", ModelConfigID: model.ID, PlanMode: database.NullChatPlanMode{ @@ -2294,6 +2390,7 @@ func TestSpawnAgent_ComputerUseRejectsMissingConfiguredProvider(t *testing.T) { server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) user := dbgen.User(t, db, database.User{}) + _ = testAPIKeyID(t, db, user.ID) org := dbgen.Organization(t, db, database.Organization{}) dbgen.OrganizationMember(t, db, database.OrganizationMember{ UserID: user.ID, @@ -2476,6 +2573,7 @@ func TestSpawnAgent_NotAvailableForExploreChats(t *testing.T) { exploreChat, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root-explore", ModelConfigID: model.ID, ChatMode: database.NullChatMode{ @@ -2601,6 +2699,7 @@ func TestSubagentLifecycleToolErrorsIncludePersistedSubagentType(t *testing.T) { unrelated, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "unrelated-lifecycle-parent", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("other")}, @@ -2672,6 +2771,7 @@ func TestSpawnAgent_ComputerUseUsesComputerUseModelNotParent(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, BuildID: uuid.NullUUID{UUID: build.ID, Valid: true}, AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true}, @@ -2738,6 +2838,7 @@ func TestSpawnAgent_ComputerUseInheritsMCPServerIDs(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-cu-mcp", ModelConfigID: model.ID, MCPServerIDs: parentMCPIDs, @@ -2798,6 +2899,7 @@ func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-with-mcp", ModelConfigID: model.ID, MCPServerIDs: parentMCPIDs, @@ -2812,11 +2914,13 @@ func TestCreateChildSubagentChat_InheritsMCPServerIDs(t *testing.T) { "parent chat must have the MCP server IDs we set") // Spawn a child subagent chat. - child, err := server.createChildSubagentChat( + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) + child, err := server.createChildSubagentChatWithOptions( ctx, parentChat, "do some work", "child-task", + childSubagentChatOptions{}, ) require.NoError(t, err) @@ -2840,6 +2944,7 @@ func TestCreateChildSubagentChat_NoMCPServersStaysEmpty(t *testing.T) { parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent-no-mcp", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -2850,11 +2955,13 @@ func TestCreateChildSubagentChat_NoMCPServersStaysEmpty(t *testing.T) { require.NoError(t, err) // Spawn a child. - child, err := server.createChildSubagentChat( + ctx = aibridge.WithDelegatedAPIKeyID(ctx, testAPIKeyID(t, server.db, parentChat.OwnerID)) + child, err := server.createChildSubagentChatWithOptions( ctx, parentChat, "do some work", "child-no-mcp", + childSubagentChatOptions{}, ) require.NoError(t, err) @@ -2877,6 +2984,7 @@ func TestIsSubagentDescendant(t *testing.T) { root, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("root")}, @@ -2886,6 +2994,7 @@ func TestIsSubagentDescendant(t *testing.T) { child, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ParentChatID: uuid.NullUUID{ UUID: root.ID, Valid: true, @@ -2903,6 +3012,7 @@ func TestIsSubagentDescendant(t *testing.T) { grandchild, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ParentChatID: uuid.NullUUID{ UUID: child.ID, Valid: true, @@ -2921,6 +3031,7 @@ func TestIsSubagentDescendant(t *testing.T) { unrelated, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "unrelated-root", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("unrelated")}, @@ -2930,6 +3041,7 @@ func TestIsSubagentDescendant(t *testing.T) { unrelatedChild, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), ParentChatID: uuid.NullUUID{ UUID: unrelated.ID, Valid: true, @@ -3020,6 +3132,7 @@ func createParentChildChats( parent, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, server.db, user.ID), Title: "parent-" + t.Name(), ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -3029,6 +3142,7 @@ func createParentChildChats( child, err = server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, server.db, user.ID), ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -3271,6 +3385,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { unrelated, err := server.CreateChat(ctx, CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "unrelated", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("other")}, @@ -3335,33 +3450,17 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Run("CompletesViaPoll", func(t *testing.T) { t.Parallel() - // Use nil pubsub so awaitSubagentCompletion falls back to - // the fast 200ms poll interval. + // Force subscription failure so awaitSubagentCompletion + // falls back to the fast 200ms poll interval. db, _ := dbtestutil.NewDB(t) mClock := quartz.NewMock(t) - server := newInternalTestServerWithClock(t, db, nil, chatprovider.ProviderAPIKeys{}, mClock) + ps := subscribeFailingPubsub{Pubsub: pubsub.NewInMemory()} + server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock) ctx := chatdTestContext(t) user, org, model := seedInternalChatDeps(t, db) parent, child := createParentChildChats(ctx, t, server, user, org, model) - // signalWake from CreateChat triggers background processing. Wait - // for those runs to finish, then reset both chats so this test owns - // the state transition observed by the poll loop. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - parentChat, err := db.GetChatByID(ctx, parent.ID) - if err != nil { - return false - } - childChat, err := db.GetChatByID(ctx, child.ID) - if err != nil { - return false - } - return parentChat.Status != database.ChatStatusPending && - parentChat.Status != database.ChatStatusRunning && - childChat.Status != database.ChatStatusPending && - childChat.Status != database.ChatStatusRunning - }, testutil.IntervalFast) setChatStatus(ctx, t, db, parent.ID, database.ChatStatusRunning, "") setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "") @@ -3449,7 +3548,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { probeCh := make(chan struct{}, 1) cancelProbe, err := ps.SubscribeWithErr( - coderdpubsub.ChatStreamNotifyChannel(child.ID), + coderdpubsub.ChatStateUpdateChannel(child.ID), func(_ context.Context, _ []byte, _ error) { select { case probeCh <- struct{}{}: @@ -3479,7 +3578,7 @@ func TestAwaitSubagentCompletion(t *testing.T) { assert.Equal(c, "pubsub result", report) }, testutil.WaitMedium, testutil.IntervalFast) require.NoError(t, ps.Publish( - coderdpubsub.ChatStreamNotifyChannel(child.ID), + coderdpubsub.ChatStateUpdateChannel(child.ID), []byte("done"), )) testutil.RequireReceive(ctx, t, probeCh) @@ -3552,26 +3651,44 @@ func TestAwaitSubagentCompletion(t *testing.T) { t.Run("ContextCanceled", func(t *testing.T) { t.Parallel() + + providerCalled := make(chan struct{}, 1) + providerReleased := make(chan struct{}) + providerServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + select { + case providerCalled <- struct{}{}: + default: + } + + select { + case <-r.Context().Done(): + case <-providerReleased: + } + })) + t.Cleanup(func() { + close(providerReleased) + providerServer.Close() + }) + + db, ps := dbtestutil.NewDB(t) + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) ctx := chatdTestContext(t) + user, org, _ := seedInternalChatDeps(t, db) + provider := dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "OpenAI", + BaseUrl: providerServer.URL, + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Provider: "openai", + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + }) parent, child := createParentChildChats(ctx, t, server, user, org, model) - // signalWake from CreateChat triggers background - // processing. drainInflight waits for in-flight goroutines - // but can't guarantee a pending DB row has been acquired - // yet — the child chat may still be pending if the second - // wake signal hasn't been consumed. Poll until the child - // reaches a terminal DB state so processChat has fully - // finished, then reset to running for the cancellation - // test. - testutil.Eventually(ctx, t, func(ctx context.Context) bool { - c, err := db.GetChatByID(ctx, child.ID) - if err != nil { - return false - } - return c.Status != database.ChatStatusPending && c.Status != database.ChatStatusRunning - }, testutil.IntervalFast) - setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "") + testutil.RequireReceive(ctx, t, providerCalled) + // Use a short-lived context instead of goroutine + sleep. shortCtx, cancel := context.WithTimeout(ctx, testutil.IntervalMedium) defer cancel() diff --git a/coderd/x/chatd/subagent_test.go b/coderd/x/chatd/subagent_test.go index a768f3487e..f621fa579e 100644 --- a/coderd/x/chatd/subagent_test.go +++ b/coderd/x/chatd/subagent_test.go @@ -27,6 +27,7 @@ func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) { parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -40,6 +41,7 @@ func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) { child, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: parent.OwnerID, + APIKeyID: testAPIKeyID(t, db, parent.OwnerID), ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -82,6 +84,7 @@ func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) { parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -94,6 +97,7 @@ func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) { child, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: parent.OwnerID, + APIKeyID: testAPIKeyID(t, db, parent.OwnerID), ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -141,6 +145,7 @@ func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) { parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "parent", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -152,6 +157,7 @@ func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) { child, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: parent.OwnerID, + APIKeyID: testAPIKeyID(t, db, parent.OwnerID), ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, @@ -187,6 +193,7 @@ func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) { parent, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: user.ID, + APIKeyID: testAPIKeyID(t, db, user.ID), Title: "root-parent", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, @@ -198,6 +205,7 @@ func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) { child, err := server.CreateChat(ctx, chatd.CreateOptions{ OrganizationID: org.ID, OwnerID: parent.OwnerID, + APIKeyID: testAPIKeyID(t, db, parent.OwnerID), ParentChatID: uuid.NullUUID{ UUID: parent.ID, Valid: true, diff --git a/coderd/x/chatd/subscribe_out_of_order_internal_test.go b/coderd/x/chatd/subscribe_out_of_order_internal_test.go deleted file mode 100644 index a6bb083785..0000000000 --- a/coderd/x/chatd/subscribe_out_of_order_internal_test.go +++ /dev/null @@ -1,212 +0,0 @@ -package chatd - -import ( - "testing" - "time" - - "github.com/google/uuid" - "github.com/stretchr/testify/require" - "go.uber.org/mock/gomock" - - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/dbmock" - coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" - "github.com/coder/coder/v2/codersdk" - "github.com/coder/coder/v2/testutil" -) - -// TestSubscribeDeliversOutOfOrderDurableMessage tests that a -// late-arriving lower-ID durable message is delivered when a -// higher-ID was already cached and sent. -func TestSubscribeDeliversOutOfOrderDurableMessage(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitMedium) - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRequiresAction} - initialUser := database.ChatMessage{ID: 3, ChatID: chatID, Role: database.ChatMessageRoleUser} - initialAssistant := database.ChatMessage{ID: 4, ChatID: chatID, Role: database.ChatMessageRoleAssistant} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 0, - }).Return([]database.ChatMessage{initialUser, initialAssistant}, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - // Notify-driven catch-up queries return nothing so the test only - // exercises the cache delivery path. - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - - server := newSubscribeTestServer(t, db) - - toolResult := codersdk.ChatMessage{ID: 5, ChatID: chatID, Role: codersdk.ChatMessageRoleTool} - resumed := codersdk.ChatMessage{ID: 7, ChatID: chatID, Role: codersdk.ChatMessageRoleAssistant} - promoted := codersdk.ChatMessage{ID: 6, ChatID: chatID, Role: codersdk.ChatMessageRoleUser} - - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, - Message: &codersdk.ChatMessage{ID: 4, ChatID: chatID, Role: codersdk.ChatMessageRoleAssistant}, - }) - - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 0) - require.True(t, ok) - defer cancel() - - // Cache id=5 and id=7, but not id=6, then emit the notify for - // id=5. The merge goroutine drains [5, 7] from the cache. - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, Message: &toolResult, - }) - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, Message: &resumed, - }) - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{AfterMessageID: 4}) - - first := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, first.Type) - require.NotNil(t, first.Message) - require.Equal(t, int64(5), first.Message.ID) - second := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, second.Type) - require.NotNil(t, second.Message) - require.Equal(t, int64(7), second.Message.ID) - - // Cache id=6 after the merge goroutine has already advanced - // lastMessageID to 7, then emit the notify for id=6. - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, Message: &promoted, - }) - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{AfterMessageID: 5}) - - third := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, third.Type) - require.NotNil(t, third.Message) - require.Equal(t, int64(6), third.Message.ID) - - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -// TestSubscribeRespectsAfterMessageIDOnLateNotify tests that -// lookupAfter never drops below afterMessageID, preventing -// re-emission of messages the client already has via REST. -func TestSubscribeRespectsAfterMessageIDOnLateNotify(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitMedium) - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 100, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - ) - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes() - - server := newSubscribeTestServer(t, db) - - // Seed the cache with messages the client claims to already have - // (id<=100) plus one new message (id=101). - for _, id := range []int64{96, 97, 98, 99, 100, 101} { - msg := &codersdk.ChatMessage{ID: id, ChatID: chatID, Role: codersdk.ChatMessageRoleAssistant} - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, Message: msg, - }) - } - - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 100) - require.True(t, ok) - defer cancel() - - // A stale notify with AfterMessageID=95 would naively pull - // id=96..101 back from the cache; only id=101 should reach the - // live stream because the client already has 96-100. - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{AfterMessageID: 95}) - - ev := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, ev.Type) - require.NotNil(t, ev.Message) - require.Equal(t, int64(101), ev.Message.ID, - "messages at or below afterMessageID must not be re-emitted") - - requireNoStreamEvent(t, events, 200*time.Millisecond) -} - -// TestSubscribeRunsDBFallbackWhenCacheDeliversUnrelatedMessage tests -// that the DB fallback runs even when the cache delivers, so -// cross-replica messages are not dropped. -func TestSubscribeRunsDBFallbackWhenCacheDeliversUnrelatedMessage(t *testing.T) { - t.Parallel() - - ctx := testutil.Context(t, testutil.WaitMedium) - - ctrl := gomock.NewController(t) - db := dbmock.NewMockStore(ctrl) - - chatID := uuid.New() - chat := database.Chat{ID: chatID, Status: database.ChatStatusRunning} - crossReplica := database.ChatMessage{ID: 6, ChatID: chatID, Role: database.ChatMessageRoleUser} - - gomock.InOrder( - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - db.EXPECT().GetChatByID(gomock.Any(), chatID).Return(chat, nil), - // Snapshot: nothing above the client's afterMessageID=5 yet. - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 5, - }).Return(nil, nil), - db.EXPECT().GetChatQueuedMessages(gomock.Any(), chatID).Return(nil, nil), - // Notify catchup: the cross-replica message lives only in the - // DB on this replica. - db.EXPECT().GetChatMessagesByChatID(gomock.Any(), database.GetChatMessagesByChatIDParams{ - ChatID: chatID, - AfterID: 5, - }).Return([]database.ChatMessage{crossReplica}, nil), - ) - - server := newSubscribeTestServer(t, db) - - // Cache a locally-published higher-ID message so the cache pass - // has something to deliver without covering id=6. - localOnly := codersdk.ChatMessage{ID: 8, ChatID: chatID, Role: codersdk.ChatMessageRoleAssistant} - server.cacheDurableMessage(chatID, codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessage, ChatID: chatID, Message: &localOnly, - }) - - _, events, cancel, ok := server.Subscribe(ctx, chatID, nil, 5) - require.True(t, ok) - defer cancel() - - server.publishChatStreamNotify(chatID, coderdpubsub.ChatStreamNotifyMessage{AfterMessageID: 5}) - - // The cache pass delivers id=8; the DB pass must still run and - // deliver id=6. Order between them is set by cache iteration vs - // DB query, so accept either ordering. - first := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, first.Type) - require.NotNil(t, first.Message) - second := testutil.RequireReceive(ctx, t, events) - require.Equal(t, codersdk.ChatStreamEventTypeMessage, second.Type) - require.NotNil(t, second.Message) - - got := map[int64]bool{first.Message.ID: true, second.Message.ID: true} - require.True(t, got[6], "cross-replica DB message id=6 must be delivered") - require.True(t, got[8], "locally-cached message id=8 must be delivered") - - requireNoStreamEvent(t, events, 200*time.Millisecond) -} diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go new file mode 100644 index 0000000000..e627556d3d --- /dev/null +++ b/coderd/x/chatd/tasks.go @@ -0,0 +1,644 @@ +package chatd + +import ( + "context" + "database/sql" + "encoding/json" + "errors" + "strings" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/db2sdk" + 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 + +var ( + errTaskExpectedExit = xerrors.New("chatworker task expected exit") + errTaskRetryable = xerrors.New("chatworker task retryable error") +) + +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 + initialDelay time.Duration + maxDelay time.Duration +} + +func runTaskWithRetry( + ctx context.Context, + opts retryWrapperOptions, + kind taskKind, + 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 { + err := executeTaskSafely(ctx, fn) + switch { + case err == nil: + return nil + case errors.Is(err, errTaskExpectedExit): + return nil + case ctx.Err() != nil: + return nil + } + + 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 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 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, + 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 { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load locked chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + chat = locked + return nil + }) + if err != nil { + return normalizeTaskInfrastructureError(err, "lock chat for interrupt") + } + + key := messagepartbuffer.Key{ + ChatID: input.ChatID, + HistoryVersion: input.HistoryVersion, + GenerationAttempt: chat.GenerationAttempt, + } + if err := s.opts.MessagePartBuffer.CloseEpisode(key); err != nil { + if ctx.Err() != nil { + return errTaskExpectedExit + } + 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 errTaskExpectedExit + } + return taskRetryableError{err: xerrors.Errorf("get message part episode: %w", err)} + } + partialMessages, err := bufferedPartsToPartialMessages(bufferedPartsToPartialMessagesInput{ + parts: parts, + modelConfigID: chat.LastModelConfigID, + contentVersion: chatprompt.CurrentContentVersion, + logger: s.opts.Logger, + interruptedAt: s.opts.Clock.Now("chatworker", "interrupt"), + }) + 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 { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusInterrupting, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + messages := partialMessages + committedCancels, err := committedPendingLocalToolCancellationMessages(ctx, store, locked, s.opts.Clock.Now("chatworker", "interrupt")) + if err != nil { + return err + } + if len(committedCancels) > 0 { + messages = append(append([]chatstate.Message{}, partialMessages...), committedCancels...) + } + if _, err := tx.FinishInterruption(chatstate.FinishInterruptionInput{PartialMessages: messages}); err != nil { + return 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 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 && s.server != 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 err + } + if decision.cancel { + return s.cancelRequiresAction(ctx, machine, input, decision.reason) + } + if !decision.waitUntil.Valid { + return errTaskExpectedExit + } + if err := s.waitUntil(ctx, decision.waitUntil.Time); err != nil { + return 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 { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load locked chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + if !locked.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(locked.RequiresActionDeadlineAt.Time) { + decision.waitUntil = locked.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 errTaskExpectedExit + } +} + +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 { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if err := verifyTaskFence(locked, input, database.ChatStatusRequiresAction, taskFenceOptions{requireHistory: true}); err != nil { + return err + } + if locked.RequiresActionDeadlineAt.Valid { + now, err := store.GetDatabaseNow(ctx) + if err != nil { + return xerrors.Errorf("get database time: %w", err) + } + if now.Before(locked.RequiresActionDeadlineAt.Time) { + return errTaskExpectedExit + } + } + if _, err := tx.CancelRequiresAction(chatstate.CancelRequiresActionInput{Reason: reason}); err != nil { + return 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 { + locked, err := store.GetChatByID(ctx, input.ChatID) + if errors.Is(err, sql.ErrNoRows) { + mismatch = true + return errTaskExpectedExit + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if !ownedByTask(locked, input) { + mismatch = true + return errTaskExpectedExit + } + if err := verifyTaskFence(locked, input, input.Status, taskFenceOptions{requireHistory: true, allowArchived: true}); err != nil { + return err + } + if _, err := tx.Abandon(chatstate.AbandonInput{}); err != nil { + return 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 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 errTaskExpectedExit + } + timer := s.opts.Clock.NewTimer(delay, "chatworker", "watch-publish-retry") + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return errTaskExpectedExit + } + 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: db2sdk.Chat(chat, nil, 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 +} + +func verifyTaskFence( + chat database.Chat, + input chatWorkerTaskStartInput, + status database.ChatStatus, + opts taskFenceOptions, +) error { + if !ownedByTask(chat, input) { + return errTaskExpectedExit + } + if chat.Status != status { + return errTaskExpectedExit + } + if !opts.allowArchived && chat.Archived { + return errTaskExpectedExit + } + if opts.requireHistory && chat.HistoryVersion != input.HistoryVersion { + return errTaskExpectedExit + } + 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) || errors.Is(err, chatstate.ErrChatNotFound) || errors.Is(err, sql.ErrNoRows) || errors.Is(err, context.Canceled) { + return errTaskExpectedExit + } + 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) || errors.Is(err, chatstate.ErrChatNotFound) || errors.Is(err, sql.ErrNoRows) || errors.Is(err, context.Canceled) { + return errTaskExpectedExit + } + 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 +} diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go new file mode 100644 index 0000000000..6ca34a4b27 --- /dev/null +++ b/coderd/x/chatd/tasks_test.go @@ -0,0 +1,1129 @@ +//nolint:testpackage // These tests exercise package-private task seams. +package chatd + +import ( + "context" + "database/sql" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" + "github.com/stretchr/testify/require" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "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/chatstate" + "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestRetryWrapper_ExpectedExitsDoNotRetry(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + calls := 0 + err := runTaskWithRetry(ctx, retryWrapperOptions{ + clock: quartz.NewMock(t), + initialDelay: time.Second, + maxDelay: time.Second, + }, taskKindInterrupt, func(context.Context) error { + calls++ + return errTaskExpectedExit + }) + require.NoError(t, err) + require.Equal(t, 1, calls) +} + +func TestRetryWrapper_UnexpectedErrorsRetry(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + trap := clock.Trap().NewTimer("chatworker", "task-retry-requires_action_timeout") + defer trap.Close() + ctx := testutil.Context(t, testutil.WaitLong) + calls := 0 + done := make(chan error, 1) + go func() { + done <- runTaskWithRetry(ctx, retryWrapperOptions{ + clock: clock, + initialDelay: time.Minute, + maxDelay: time.Minute, + }, taskKindRequiresActionTimeout, func(context.Context) error { + calls++ + if calls == 1 { + return xerrors.New("database unavailable") + } + return nil + }) + }() + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(time.Minute).MustWait(ctx) + require.NoError(t, <-done) + require.Equal(t, 2, calls) +} + +func TestRetryWrapper_PanicsRetry(t *testing.T) { + t.Parallel() + + clock := quartz.NewMock(t) + trap := clock.Trap().NewTimer("chatworker", "task-retry-generation") + defer trap.Close() + ctx := testutil.Context(t, testutil.WaitLong) + calls := 0 + done := make(chan error, 1) + go func() { + done <- runTaskWithRetry(ctx, retryWrapperOptions{ + clock: clock, + initialDelay: time.Minute, + maxDelay: time.Minute, + }, taskKindGeneration, func(context.Context) error { + calls++ + if calls == 1 { + panic("database unavailable") + } + return nil + }) + }() + + trap.MustWait(ctx).MustRelease(ctx) + clock.Advance(time.Minute).MustWait(ctx) + require.NoError(t, <-done) + require.Equal(t, 2, calls) +} + +func TestInterruptTask_FinishInterruptionOnly(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := messagepartbuffer.Key{ + ChatID: chat.ID, + HistoryVersion: acquired.HistoryVersion, + GenerationAttempt: acquired.GenerationAttempt, + } + require.NoError(t, buffer.CreateEpisode(key)) + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessageText("partial answer"))) + interrupting := f.interruptChat(t, chat.ID) + require.Equal(t, database.ChatStatusInterrupting, interrupting.Status) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, buffer, recorder) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, latest.Status) + recorder.requireStateHint(t, chat.ID, latest.SnapshotVersion, database.ChatStatusRunning) + recorder.requireInterruptionOutcome(t, chat.ID, database.ChatStatusRunning) + recorder.requireCleanupCount(t, 0) + f.requireWatchEvent(t, chat.ID, codersdk.ChatWatchEventKindStatusChange) + + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.GreaterOrEqual(t, len(messages), 3) + parts, err := chatprompt.ParseContent(messages[len(messages)-2]) + require.NoError(t, err) + require.Equal(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("partial answer")}, parts) + require.Equal(t, database.ChatMessageRoleUser, messages[len(messages)-1].Role) +} + +func TestInterruptTask_StaleFenceExits(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + f.acquireChat(t, chat.ID, workerID, runnerID) + interrupting := f.interruptChat(t, chat.ID) + otherWorkerID := uuid.New() + otherRunnerID := uuid.New() + f.acquireChat(t, chat.ID, otherWorkerID, otherRunnerID) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.ErrorIs(t, err, errTaskExpectedExit) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusInterrupting, latest.Status) + require.Equal(t, otherWorkerID, latest.WorkerID.UUID) + require.Equal(t, otherRunnerID, latest.RunnerID.UUID) + recorder.requireStateHintCount(t, 0) + f.requireNoWatchEvents(t) +} + +func TestInterruptTask_MissingEpisodePersistsNilPartials(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + f.acquireChat(t, chat.ID, workerID, runnerID) + interrupting := f.forceExecutionState(t, chat.ID, database.ChatStatusInterrupting, false, sql.NullTime{}) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusWaiting, latest.Status) + recorder.requireInterruptionOutcome(t, chat.ID, database.ChatStatusWaiting) + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Len(t, messages, 1) + recorder.requireStateHint(t, chat.ID, latest.SnapshotVersion, database.ChatStatusWaiting) +} + +func TestInterruptTask_BufferedPartsBecomePartialMessages(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + buffer := messagepartbuffer.New(messagepartbuffer.Options{}) + key := messagepartbuffer.Key{ChatID: chat.ID, HistoryVersion: acquired.HistoryVersion, GenerationAttempt: acquired.GenerationAttempt} + require.NoError(t, buffer.CreateEpisode(key)) + callID := "call_" + uuid.NewString() + require.NoError(t, buffer.AddPart(key, codersdk.ChatMessageRoleAssistant, codersdk.ChatMessagePart{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: callID, + ToolName: "local_tool", + Args: json.RawMessage(`{"value":1}`), + })) + interrupting := f.interruptChat(t, chat.ID) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, buffer, recorder) + + err := starter.StartInterrupt(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: interrupting.HistoryVersion, + GenerationAttempt: interrupting.GenerationAttempt, + Status: database.ChatStatusInterrupting, + }) + require.NoError(t, err) + + messages, err := f.db.GetChatMessagesByChatID(testutil.Context(t, testutil.WaitShort), database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.GreaterOrEqual(t, len(messages), 4) + assistant := messages[len(messages)-3] + tool := messages[len(messages)-2] + require.Equal(t, database.ChatMessageRoleAssistant, assistant.Role) + require.Equal(t, database.ChatMessageRoleTool, tool.Role) + toolParts, err := chatprompt.ParseContent(tool) + require.NoError(t, err) + require.Len(t, toolParts, 1) + require.Equal(t, codersdk.ChatMessagePartTypeToolResult, toolParts[0].Type) + require.Equal(t, callID, toolParts[0].ToolCallID) + require.True(t, toolParts[0].IsError) +} + +func TestRequiresActionTimeout_ExpiredCancelsOnly(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRequiresActionChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + expired := f.setRequiresActionDeadline(t, chat.ID, sql.NullTime{Time: time.Now().Add(-time.Minute), Valid: true}) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartRequiresActionTimeout(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRequiresAction, + RequiresActionDeadlineAt: expired.RequiresActionDeadlineAt, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, latest.Status) + require.False(t, latest.RequiresActionDeadlineAt.Valid) + recorder.requireStateHint(t, chat.ID, latest.SnapshotVersion, database.ChatStatusRunning) + f.requireWatchEvent(t, chat.ID, codersdk.ChatWatchEventKindStatusChange) +} + +func TestRequiresActionTimeout_NullDeadlineCancelsImmediately(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRequiresActionChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + nullDeadline := f.setRequiresActionDeadline(t, chat.ID, sql.NullTime{}) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartRequiresActionTimeout(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRequiresAction, + RequiresActionDeadlineAt: nullDeadline.RequiresActionDeadlineAt, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, latest.Status) + recorder.requireStateHint(t, chat.ID, latest.SnapshotVersion, database.ChatStatusRunning) +} + +func TestRequiresActionTimeout_StaleFenceExitsAfterToolResult(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRequiresActionChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + expired := f.setRequiresActionDeadline(t, chat.ID, sql.NullTime{Time: time.Now().Add(-time.Minute), Valid: true}) + f.forceExecutionState(t, chat.ID, database.ChatStatusRunning, false, sql.NullTime{}) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartRequiresActionTimeout(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRequiresAction, + RequiresActionDeadlineAt: expired.RequiresActionDeadlineAt, + }) + require.ErrorIs(t, err, errTaskExpectedExit) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, latest.Status) + recorder.requireStateHintCount(t, 0) + f.requireNoWatchEvents(t) +} + +func TestAbandonTask_AbandonOnly(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartAbandon(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.False(t, latest.WorkerID.Valid) + require.False(t, latest.RunnerID.Valid) + recorder.requireCleanup(t, chat.ID, runnerID) + recorder.requireStateHintCount(t, 0) + f.requireNoWatchEvents(t) +} + +func TestAbandonTask_OwnershipMismatchRequestsCleanup(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + f.acquireChat(t, chat.ID, workerID, runnerID) + otherWorkerID := uuid.New() + otherRunnerID := uuid.New() + latestOwner := f.acquireChat(t, chat.ID, otherWorkerID, otherRunnerID) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartAbandon(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: latestOwner.HistoryVersion, + Status: database.ChatStatusRunning, + }) + require.NoError(t, err) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, otherWorkerID, latest.WorkerID.UUID) + require.Equal(t, otherRunnerID, latest.RunnerID.UUID) + recorder.requireCleanup(t, chat.ID, runnerID) +} + +func TestAbandonTask_StaleStatusFenceExits(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + f.forceExecutionState(t, chat.ID, database.ChatStatusInterrupting, false, sql.NullTime{}) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + err := starter.StartAbandon(testutil.Context(t, testutil.WaitLong), chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusWaiting, + }) + require.ErrorIs(t, err, errTaskExpectedExit) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.True(t, latest.WorkerID.Valid) + require.True(t, latest.RunnerID.Valid) + require.Equal(t, database.ChatStatusInterrupting, latest.Status) + recorder.requireCleanupCount(t, 0) +} + +func TestGenerationTask_RecordRetryState(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + recorder := newTaskSideEffectRecorder() + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), recorder) + + attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt( + testutil.Context(t, testutil.WaitLong), + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + ) + require.NoError(t, err) + closeEpisode() + require.Equal(t, int64(1), attempt) + before, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.False(t, before.RetryState.Valid) + + decision, err := starter.recordGenerationRetry( + testutil.Context(t, testutil.WaitLong), + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + chaterror.ClassifiedError{ + Message: "OpenAI is rate limiting requests.", + Kind: codersdk.ChatErrorKindRateLimit, + Provider: "openai", + Retryable: true, + StatusCode: 429, + }, + ) + require.NoError(t, err) + require.True(t, decision.retry) + require.Equal(t, int64(1), decision.generationAttempt) + require.Equal(t, chatretry.Delay(0), decision.delay) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.True(t, latest.RetryState.Valid) + require.Equal(t, latest.SnapshotVersion, latest.RetryStateVersion) + require.Greater(t, latest.RetryStateVersion, before.RetryStateVersion) + require.Equal(t, before.GenerationAttempt, latest.GenerationAttempt) + recorder.requireStateHintCount(t, 0) + + var retryPayload codersdk.ChatStreamRetry + require.NoError(t, json.Unmarshal(latest.RetryState.RawMessage, &retryPayload)) + require.Equal(t, 1, retryPayload.Attempt) + require.Equal(t, chatretry.Delay(0).Milliseconds(), retryPayload.DelayMs) + require.Equal(t, "OpenAI is rate limiting requests.", retryPayload.Error) + require.Equal(t, codersdk.ChatErrorKindRateLimit, retryPayload.Kind) + require.Equal(t, "openai", retryPayload.Provider) + require.Equal(t, 429, retryPayload.StatusCode) + require.False(t, retryPayload.RetryingAt.IsZero()) +} + +func TestGenerationTask_RecordRetryStateUsesDurableGenerationAttempt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), newTaskSideEffectRecorder()) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID) + + for range 3 { + attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt( + testutil.Context(t, testutil.WaitLong), + machine, + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + ) + require.NoError(t, err) + closeEpisode() + require.Positive(t, attempt) + } + + decision, err := starter.recordGenerationRetry( + testutil.Context(t, testutil.WaitLong), + machine, + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + chaterror.ClassifiedError{ + Message: "OpenAI is temporarily unavailable.", + Kind: codersdk.ChatErrorKindTimeout, + Provider: "openai", + Retryable: true, + }, + ) + require.NoError(t, err) + require.True(t, decision.retry) + require.Equal(t, int64(3), decision.generationAttempt) + require.Equal(t, chatretry.Delay(2), decision.delay) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + var retryPayload codersdk.ChatStreamRetry + require.NoError(t, json.Unmarshal(latest.RetryState.RawMessage, &retryPayload)) + require.Equal(t, 3, retryPayload.Attempt) + require.Equal(t, chatretry.Delay(2).Milliseconds(), retryPayload.DelayMs) +} + +func TestGenerationTask_RecordRetryStateClearedByNextAttempt(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), newTaskSideEffectRecorder()) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID) + input := chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + } + + attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input) + require.NoError(t, err) + closeEpisode() + require.Equal(t, int64(1), attempt) + _, err = starter.recordGenerationRetry( + testutil.Context(t, testutil.WaitLong), + machine, + input, + chaterror.ClassifiedError{ + Message: "OpenAI is temporarily unavailable.", + Kind: codersdk.ChatErrorKindTimeout, + Provider: "openai", + Retryable: true, + }, + ) + require.NoError(t, err) + withRetry, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.True(t, withRetry.RetryState.Valid) + + attempt, _, _, closeEpisode, err = starter.beginGenerationAttempt(testutil.Context(t, testutil.WaitLong), machine, input) + require.NoError(t, err) + closeEpisode() + require.Equal(t, int64(2), attempt) + after, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.False(t, after.RetryState.Valid) + require.Equal(t, after.SnapshotVersion, after.RetryStateVersion) + require.Greater(t, after.RetryStateVersion, withRetry.RetryStateVersion) +} + +func TestGenerationTask_RecordRetryStateStaleFenceExits(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + acquired := f.acquireChat(t, chat.ID, workerID, runnerID) + starter := newTestTaskStarter(t, f, messagepartbuffer.New(messagepartbuffer.Options{}), newTaskSideEffectRecorder()) + machine := chatstate.NewChatMachine(f.db, f.pubsub, chat.ID) + attempt, _, _, closeEpisode, err := starter.beginGenerationAttempt( + testutil.Context(t, testutil.WaitLong), + machine, + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + ) + require.NoError(t, err) + closeEpisode() + require.Equal(t, int64(1), attempt) + + otherWorkerID := uuid.New() + otherRunnerID := uuid.New() + f.acquireChat(t, chat.ID, otherWorkerID, otherRunnerID) + _, err = starter.recordGenerationRetry( + testutil.Context(t, testutil.WaitLong), + machine, + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: acquired.HistoryVersion, + Status: database.ChatStatusRunning, + }, + chaterror.ClassifiedError{ + Message: "OpenAI is temporarily unavailable.", + Kind: codersdk.ChatErrorKindTimeout, + Provider: "openai", + Retryable: true, + }, + ) + require.ErrorIs(t, err, errTaskExpectedExit) + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.False(t, latest.RetryState.Valid) + require.Equal(t, otherWorkerID, latest.WorkerID.UUID) + require.Equal(t, otherRunnerID, latest.RunnerID.UUID) +} + +func TestRunner_StartsRealInterruptTask(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + worker := startRealTaskWorker(t, f, messagepartbuffer.New(messagepartbuffer.Options{})) + waitOwnedChat(t, f, chat.ID, worker.chatWorkerID()) + + interrupting := f.interruptChat(t, chat.ID) + require.Equal(t, database.ChatStatusInterrupting, interrupting.Status) + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + latest, err := f.db.GetChatByID(ctx, chat.ID) + return err == nil && latest.Status == database.ChatStatusRunning + }, testutil.IntervalFast) + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, worker.chatWorkerID(), latest.WorkerID.UUID) + f.requireWatchEvent(t, chat.ID, codersdk.ChatWatchEventKindStatusChange) +} + +func TestRunner_StartsRealRequiresActionTimeoutTask(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRequiresActionChat(t) + f.setRequiresActionDeadline(t, chat.ID, sql.NullTime{Time: time.Now().Add(-time.Minute), Valid: true}) + worker := startRealTaskWorker(t, f, messagepartbuffer.New(messagepartbuffer.Options{})) + + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + latest, err := f.db.GetChatByID(ctx, chat.ID) + return err == nil && latest.Status == database.ChatStatusRunning && latest.WorkerID.Valid && latest.WorkerID.UUID == worker.chatWorkerID() + }, testutil.IntervalFast) + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.True(t, latest.RunnerID.Valid) + f.requireWatchEvent(t, chat.ID, codersdk.ChatWatchEventKindStatusChange) +} + +func TestRunner_StartsRealAbandonTask(t *testing.T) { + t.Parallel() + + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + worker := startRealTaskWorker(t, f, messagepartbuffer.New(messagepartbuffer.Options{})) + waitOwnedChat(t, f, chat.ID, worker.chatWorkerID()) + + updated := f.forceExecutionState(t, chat.ID, database.ChatStatusError, false, sql.NullTime{}) + f.publishChatUpdate(t, updated) + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + latest, err := f.db.GetChatByID(ctx, chat.ID) + return err == nil && !latest.WorkerID.Valid && !latest.RunnerID.Valid + }, testutil.IntervalFast) +} + +type taskTestFixture struct { + db database.Store + pubsub *taskRecordingPubsub + sqlDB *sql.DB + user database.User + org database.Organization + model database.ChatModelConfig + apiKey database.APIKey +} + +func newTaskTestFixture(t *testing.T) *taskTestFixture { + t.Helper() + db, ps, sqlDB := dbtestutil.NewDBWithSQLDB(t) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + dbgen.ChatProvider(t, db, database.ChatProvider{ + Provider: "openai", + DisplayName: "openai", + BaseUrl: "http://example.invalid", + }) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{Provider: "openai", IsDefault: true}) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + return &taskTestFixture{db: db, pubsub: newTaskRecordingPubsub(ps), sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey} +} + +func (f *taskTestFixture) createRunningChat(t *testing.T) database.Chat { + t.Helper() + res, err := chatstate.CreateChat(testutil.Context(t, testutil.WaitShort), f.db, f.pubsub, chatstate.CreateChatInput{ + OrganizationID: f.org.ID, + OwnerID: f.user.ID, + LastModelConfigID: f.model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{taskUserTextMessage(t, "hello", f.user.ID, f.model.ID, f.apiKey.ID)}, + }) + require.NoError(t, err) + f.pubsub.clear() + return res.Chat +} + +func (f *taskTestFixture) createRequiresActionChat(t *testing.T) database.Chat { + t.Helper() + toolName := "dynamic_" + uuid.NewString() + dynamicTools, err := json.Marshal([]codersdk.DynamicTool{{ + Name: toolName, + Description: "test tool", + InputSchema: json.RawMessage(`{"type":"object"}`), + }}) + require.NoError(t, err) + res, err := chatstate.CreateChat(testutil.Context(t, testutil.WaitShort), f.db, f.pubsub, chatstate.CreateChatInput{ + OrganizationID: f.org.ID, + OwnerID: f.user.ID, + LastModelConfigID: f.model.ID, + Title: "test", + ClientType: database.ChatClientTypeApi, + DynamicTools: pqtype.NullRawMessage{RawMessage: dynamicTools, Valid: true}, + InitialMessages: []chatstate.Message{taskUserTextMessage(t, "hello", f.user.ID, f.model.ID, f.apiKey.ID)}, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(f.db, f.pubsub, res.Chat.ID) + require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{taskAssistantToolCallMessage(t, f.model.ID, toolName)}}) + return err + })) + require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}) + return err + })) + chat, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), res.Chat.ID) + require.NoError(t, err) + f.pubsub.clear() + return chat +} + +func (f *taskTestFixture) acquireChat(t *testing.T, chatID uuid.UUID, workerID uuid.UUID, runnerID uuid.UUID) database.Chat { + t.Helper() + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.Acquire(chatstate.AcquireInput{WorkerID: workerID, RunnerID: runnerID}) + return err + })) + chat, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chatID) + require.NoError(t, err) + f.pubsub.clear() + return chat +} + +func (f *taskTestFixture) interruptChat(t *testing.T, chatID uuid.UUID) database.Chat { + t.Helper() + machine := chatstate.NewChatMachine(f.db, f.pubsub, chatID) + require.NoError(t, machine.Update(testutil.Context(t, testutil.WaitShort), func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.SendMessage(chatstate.SendMessageInput{ + Message: taskUserTextMessage(t, "interrupt", f.user.ID, f.model.ID, f.apiKey.ID), + BusyBehavior: chatstate.BusyBehaviorInterrupt, + }) + return err + })) + chat, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chatID) + require.NoError(t, err) + f.pubsub.clear() + return chat +} + +func (f *taskTestFixture) forceExecutionState(t *testing.T, chatID uuid.UUID, status database.ChatStatus, archived bool, deadline sql.NullTime) database.Chat { + t.Helper() + var updated database.Chat + require.NoError(t, f.db.InTx(func(store database.Store) error { + if _, err := store.LockChatAndBumpSnapshotVersion(testutil.Context(t, testutil.WaitShort), chatID); err != nil { + return err + } + chat, err := store.GetChatByID(testutil.Context(t, testutil.WaitShort), chatID) + if err != nil { + return err + } + updated, err = store.UpdateChatExecutionState(testutil.Context(t, testutil.WaitShort), database.UpdateChatExecutionStateParams{ + ID: chat.ID, + Status: status, + Archived: archived, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: chat.LastError, + RequiresActionDeadlineAt: deadline, + }) + return err + }, nil)) + f.pubsub.clear() + return updated +} + +func (f *taskTestFixture) setRequiresActionDeadline(t *testing.T, chatID uuid.UUID, deadline sql.NullTime) database.Chat { + t.Helper() + chat, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chatID) + require.NoError(t, err) + return f.forceExecutionState(t, chatID, chat.Status, chat.Archived, deadline) +} + +func (f *taskTestFixture) publishChatUpdate(t *testing.T, chat database.Chat) { + t.Helper() + msg := coderdpubsub.ChatStateUpdateMessage{ + SnapshotVersion: chat.SnapshotVersion, + HistoryVersion: chat.HistoryVersion, + QueueVersion: chat.QueueVersion, + RetryStateVersion: chat.RetryStateVersion, + GenerationAttempt: chat.GenerationAttempt, + Status: string(chat.Status), + Archived: chat.Archived, + } + if chat.WorkerID.Valid { + id := chat.WorkerID.UUID + msg.WorkerID = &id + } + if chat.RunnerID.Valid { + id := chat.RunnerID.UUID + msg.RunnerID = &id + } + payload, err := json.Marshal(msg) + require.NoError(t, err) + require.NoError(t, f.pubsub.Publish(coderdpubsub.ChatStateUpdateChannel(chat.ID), payload)) +} + +func (f *taskTestFixture) requireWatchEvent(t *testing.T, chatID uuid.UUID, kind codersdk.ChatWatchEventKind) { + t.Helper() + events := f.pubsub.watchEvents(t) + for _, event := range events { + if event.Kind == kind && event.Chat.ID == chatID { + return + } + } + t.Fatalf("missing watch event kind=%s chat_id=%s events=%v", kind, chatID, events) +} + +func (f *taskTestFixture) requireNoWatchEvents(t *testing.T) { + t.Helper() + require.Empty(t, f.pubsub.watchEvents(t)) +} + +func taskUserTextMessage(t *testing.T, text string, createdBy uuid.UUID, modelConfigID uuid.UUID, apiKeyID string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: createdBy, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + APIKeyID: sql.NullString{String: apiKeyID, Valid: apiKeyID != ""}, + } +} + +func taskAssistantToolCallMessage(t *testing.T, modelConfigID uuid.UUID, toolName string) chatstate.Message { + t.Helper() + raw, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "call_" + uuid.NewString(), + ToolName: toolName, + Args: json.RawMessage(`{}`), + }}) + require.NoError(t, err) + return chatstate.Message{ + Role: database.ChatMessageRoleAssistant, + Content: raw, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + } +} + +type taskPublishedEvent struct { + channel string + payload []byte +} + +type taskRecordingPubsub struct { + inner dbpubsub.Pubsub + mu sync.Mutex + sent []taskPublishedEvent +} + +func newTaskRecordingPubsub(inner dbpubsub.Pubsub) *taskRecordingPubsub { + return &taskRecordingPubsub{inner: inner} +} + +func (p *taskRecordingPubsub) Publish(channel string, payload []byte) error { + p.mu.Lock() + p.sent = append(p.sent, taskPublishedEvent{channel: channel, payload: append([]byte(nil), payload...)}) + p.mu.Unlock() + return p.inner.Publish(channel, payload) +} + +func (p *taskRecordingPubsub) SubscribeWithErr(channel string, listener dbpubsub.ListenerWithErr) (func(), error) { + return p.inner.SubscribeWithErr(channel, listener) +} + +func (p *taskRecordingPubsub) clear() { + p.mu.Lock() + p.sent = nil + p.mu.Unlock() +} + +func (p *taskRecordingPubsub) events() []taskPublishedEvent { + p.mu.Lock() + defer p.mu.Unlock() + return append([]taskPublishedEvent(nil), p.sent...) +} + +func (p *taskRecordingPubsub) watchEvents(t *testing.T) []codersdk.ChatWatchEvent { + t.Helper() + events := p.events() + out := make([]codersdk.ChatWatchEvent, 0) + for _, event := range events { + var payload codersdk.ChatWatchEvent + if err := json.Unmarshal(event.payload, &payload); err != nil { + continue + } + if event.channel != coderdpubsub.ChatWatchEventChannel(payload.Chat.OwnerID) { + continue + } + out = append(out, payload) + } + return out +} + +func startRealTaskWorker(t *testing.T, f *taskTestFixture, buffer *messagepartbuffer.Buffer) *chatWorker { + t.Helper() + worker, err := newChatWorker(nil, chatWorkerOptions{ + WorkerID: uuid.New(), + Store: f.db, + Pubsub: f.pubsub, + Logger: slog.Make(), + MessagePartBuffer: buffer, + AcquisitionInterval: time.Hour, + AcquisitionBatchSize: 10, + RunnerSyncInterval: time.Hour, + HeartbeatInterval: time.Hour, + HeartbeatCleanupInterval: time.Hour, + HeartbeatStaleSeconds: 30, + StateChannelSize: 16, + RunnerManagerChannelSize: 16, + AcquisitionWakeChannelSize: 1, + TaskRetryInitialBackoff: time.Millisecond, + TaskRetryMaxBackoff: time.Millisecond, + }) + require.NoError(t, err) + require.NoError(t, worker.Start(context.Background())) + t.Cleanup(func() { require.NoError(t, worker.Close()) }) + return worker +} + +func waitOwnedChat(t *testing.T, f *taskTestFixture, chatID uuid.UUID, workerID uuid.UUID) database.Chat { + t.Helper() + var latest database.Chat + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + chat, err := f.db.GetChatByID(ctx, chatID) + if err != nil { + return false + } + latest = chat + return chat.WorkerID.Valid && chat.WorkerID.UUID == workerID && chat.RunnerID.Valid + }, testutil.IntervalFast) + return latest +} + +type taskSideEffectRecorder struct { + mu sync.Mutex + hints []runnerStateUpdate + cleanups []runnerKey + interrupts []interruptionOutcome +} + +func newTaskSideEffectRecorder() *taskSideEffectRecorder { + return &taskSideEffectRecorder{} +} + +func (r *taskSideEffectRecorder) routeStateHint(_ context.Context, state runnerStateUpdate) { + r.mu.Lock() + r.hints = append(r.hints, state) + r.mu.Unlock() +} + +func (r *taskSideEffectRecorder) requestCleanup(_ context.Context, key runnerKey) { + r.mu.Lock() + r.cleanups = append(r.cleanups, key) + r.mu.Unlock() +} + +func (r *taskSideEffectRecorder) afterInterruptionOutcome(_ context.Context, outcome interruptionOutcome) error { + r.mu.Lock() + r.interrupts = append(r.interrupts, outcome) + r.mu.Unlock() + return nil +} + +func (r *taskSideEffectRecorder) requireStateHint(t *testing.T, chatID uuid.UUID, snapshot int64, status database.ChatStatus) { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + for _, hint := range r.hints { + if hint.ChatID == chatID && hint.SnapshotVersion == snapshot && hint.Status == status { + return + } + } + t.Fatalf("missing state hint chat_id=%s snapshot=%d status=%s hints=%v", chatID, snapshot, status, r.hints) +} + +func (r *taskSideEffectRecorder) requireStateHintCount(t *testing.T, count int) { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + require.Len(t, r.hints, count) +} + +func (r *taskSideEffectRecorder) requireCleanup(t *testing.T, chatID uuid.UUID, runnerID uuid.UUID) { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + for _, cleanup := range r.cleanups { + if cleanup.ChatID == chatID && cleanup.RunnerID == runnerID { + return + } + } + t.Fatalf("missing cleanup chat_id=%s runner_id=%s cleanups=%v", chatID, runnerID, r.cleanups) +} + +func (r *taskSideEffectRecorder) requireCleanupCount(t *testing.T, count int) { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + require.Len(t, r.cleanups, count) +} + +func (r *taskSideEffectRecorder) requireInterruptionOutcome(t *testing.T, chatID uuid.UUID, status database.ChatStatus) { + t.Helper() + r.mu.Lock() + defer r.mu.Unlock() + for _, outcome := range r.interrupts { + if outcome.Chat.ID == chatID && outcome.Chat.Status == status { + return + } + } + t.Fatalf("missing interruption outcome chat_id=%s status=%s outcomes=%v", chatID, status, r.interrupts) +} + +func newTestTaskStarter(t *testing.T, f *taskTestFixture, buffer *messagepartbuffer.Buffer, recorder *taskSideEffectRecorder) *taskStarter { + t.Helper() + starter, err := newTaskStarter(nil, chatWorkerOptions{ + Store: f.db, + Pubsub: f.pubsub, + Logger: slog.Make(), + Clock: quartz.NewReal(), + MessagePartBuffer: buffer, + TaskRetryInitialBackoff: time.Millisecond, + TaskRetryMaxBackoff: time.Millisecond, + }, recorder.routeStateHint, recorder.requestCleanup) + require.NoError(t, err) + starter.afterInterruptionOutcome = recorder.afterInterruptionOutcome + return starter +} diff --git a/coderd/x/chatd/testhooks.go b/coderd/x/chatd/testhooks.go index 7c7177b88b..c1356ee3d0 100644 --- a/coderd/x/chatd/testhooks.go +++ b/coderd/x/chatd/testhooks.go @@ -1,9 +1,20 @@ package chatd +import ( + "context" + "time" +) + // WaitUntilIdleForTest waits for background chat work tracked by the server to // finish without shutting the server down. Tests use this to assert final // database state only after asynchronous chat processing has completed. // Close waits for the same tracked work, but also stops the server. func WaitUntilIdleForTest(server *Server) { server.drainInflight() + if server.chatWorker == nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + _ = server.chatWorker.WaitIdle(ctx) } diff --git a/coderd/x/chatd/turn_summary_internal_test.go b/coderd/x/chatd/turn_summary_internal_test.go index be3a595799..3ca6b0741f 100644 --- a/coderd/x/chatd/turn_summary_internal_test.go +++ b/coderd/x/chatd/turn_summary_internal_test.go @@ -6,7 +6,6 @@ import ( "encoding/json" "sync/atomic" "testing" - "time" "charm.land/fantasy" "github.com/google/uuid" @@ -16,6 +15,8 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "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/chattest" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" @@ -24,7 +25,7 @@ import ( func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { t.Parallel() - db, _ := dbtestutil.NewDB(t) + db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitMedium) owner := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) @@ -55,44 +56,71 @@ func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { }) require.NoError(t, err) - chat, err := db.InsertChat(ctx, database.InsertChatParams{ + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("hello"), + }) + require.NoError(t, err) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: owner.ID}) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ OrganizationID: org.ID, - Status: database.ChatStatusWaiting, - ClientType: database.ChatClientTypeUi, OwnerID: owner.ID, LastModelConfigID: modelCfg.ID, Title: "summary-chat", + ClientType: database.ChatClientTypeUi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: owner.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, }) require.NoError(t, err) + chat := created.Chat logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := &Server{db: db} - server.updateLastTurnSummary(ctx, chat, chat.UpdatedAt, "fresh summary", logger) + server := &Server{db: db, pubsub: ps} + server.updateLastTurnSummary(ctx, chat, chat.HistoryVersion, "fresh summary", logger) fetched, err := db.GetChatByID(ctx, chat.ID) require.NoError(t, err) require.Equal(t, sql.NullString{String: "fresh summary", Valid: true}, fetched.LastTurnSummary) - advancedUpdatedAt := chat.UpdatedAt.Add(time.Second) - _, err = db.UpdateChatStatusPreserveUpdatedAt(ctx, database.UpdateChatStatusPreserveUpdatedAtParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - UpdatedAt: advancedUpdatedAt, + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + codersdk.ChatMessageText("assistant response"), }) require.NoError(t, err) + machine := chatstate.NewChatMachine(db, ps, chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: modelCfg.ID, Valid: true}, + }, + }, + }) + return err + })) - server.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.UpdatedAt, "stale summary", logger) + server.updateLastTurnSummary(context.WithoutCancel(ctx), chat, chat.HistoryVersion, "stale summary", logger) fetched, err = db.GetChatByID(ctx, chat.ID) require.NoError(t, err) require.Equal(t, sql.NullString{String: "fresh summary", Valid: true}, fetched.LastTurnSummary) - require.Equal(t, advancedUpdatedAt, fetched.UpdatedAt) } func TestPendingChatPersistsSummaryButSkipsWebPush(t *testing.T) { t.Parallel() - db, _ := dbtestutil.NewDB(t) + db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitMedium) owner := dbgen.User(t, db, database.User{}) org := dbgen.Organization(t, db, database.Organization{}) @@ -150,7 +178,7 @@ func TestPendingChatPersistsSummaryButSkipsWebPush(t *testing.T) { dispatcher := &recordingWebpushDispatcher{} logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := &Server{db: db, webpushDispatcher: dispatcher} + server := &Server{db: db, pubsub: ps, webpushDispatcher: dispatcher} server.maybeFinalizeTurnStatusLabelAndPush( context.WithoutCancel(ctx), chat, diff --git a/coderd/x/chatd/worker.go b/coderd/x/chatd/worker.go new file mode 100644 index 0000000000..7b3e8d5666 --- /dev/null +++ b/coderd/x/chatd/worker.go @@ -0,0 +1,314 @@ +package chatd + +import ( + "context" + "database/sql" + "errors" + "sync" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd/database" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" +) + +// chatWorker owns chat acquisition and runner lifecycle for one process. +type chatWorker struct { + server *Server + opts chatWorkerOptions + + mu sync.Mutex + started bool + ctx context.Context + cancel context.CancelFunc + manager *runnerManager + unsubscribe func() + wakeCh chan struct{} + wg sync.WaitGroup +} + +// newChatWorker constructs a chat worker. The worker is idle until Start is +// called. +func newChatWorker(server *Server, opts chatWorkerOptions) (*chatWorker, error) { + withDefaults, err := opts.withDefaults() + if err != nil { + return nil, err + } + return &chatWorker{server: server, opts: withDefaults}, nil +} + +// chatWorkerID returns this worker's configured worker ID. +func (w *chatWorker) chatWorkerID() uuid.UUID { + return w.opts.WorkerID +} + +// Start starts the acquisition and runner manager loops. +func (w *chatWorker) Start(ctx context.Context) error { + w.mu.Lock() + defer w.mu.Unlock() + if w.started { + return xerrors.New("chatworker: worker already started") + } + workerID := w.opts.WorkerID + workerCtx, cancel := context.WithCancel(ctx) + manager := newRunnerManager(workerCtx, w.server, w.opts) + if manager.opts.TaskStarter == nil { + starter, err := newTaskStarter(manager.server, manager.opts, manager.RouteStateHint, manager.requestCleanup) + if err != nil { + cancel() + return err + } + manager.opts.TaskStarter = starter + } + wakeCh := make(chan struct{}, w.opts.AcquisitionWakeChannelSize) + + unsubscribe, err := w.opts.Pubsub.SubscribeWithErr( + coderdpubsub.ChatStateOwnershipChannel, + coderdpubsub.HandleChatStateOwnership(func(ctx context.Context, _ coderdpubsub.ChatStateOwnershipMessage, err error) { + if err != nil { + w.opts.Logger.Warn(ctx, "chatworker ownership hint decode failed", slogError(err)) + return + } + wake(wakeCh) + }), + ) + if err != nil { + cancel() + return xerrors.Errorf("subscribe ownership hints: %w", err) + } + + w.started = true + w.ctx = workerCtx + w.cancel = cancel + w.manager = manager + w.unsubscribe = unsubscribe + w.wakeCh = wakeCh + + manager.start() + w.wg.Go(func() { + w.acquisitionLoop(workerCtx, workerID, manager, wakeCh) + }) + w.wg.Go(func() { + w.archiveLoop(workerCtx) + }) + wake(wakeCh) + return nil +} + +// Wake requests an immediate acquisition pass. +func (w *chatWorker) Wake() { + w.mu.Lock() + wakeCh := w.wakeCh + w.mu.Unlock() + if wakeCh != nil { + wake(wakeCh) + } +} + +// WaitIdle waits until the worker has no active or cleaning runners. +func (w *chatWorker) WaitIdle(ctx context.Context) error { + for { + w.mu.Lock() + manager := w.manager + w.mu.Unlock() + if manager == nil || manager.idle() { + return nil + } + timer := w.opts.Clock.NewTimer(10*time.Millisecond, "chatworker", "wait-idle") + select { + case <-timer.C: + case <-ctx.Done(): + timer.Stop() + return ctx.Err() + } + timer.Stop() + } +} + +// Close stops the worker and waits for its loops to exit. +func (w *chatWorker) Close() error { + w.mu.Lock() + if !w.started { + w.mu.Unlock() + return nil + } + cancel := w.cancel + unsubscribe := w.unsubscribe + manager := w.manager + w.started = false + w.cancel = nil + w.unsubscribe = nil + w.manager = nil + w.wakeCh = nil + w.mu.Unlock() + + if unsubscribe != nil { + unsubscribe() + } + cancel() + w.wg.Wait() + if manager != nil { + manager.wait() + } + return nil +} + +func wake(ch chan<- struct{}) { + select { + case ch <- struct{}{}: + default: + } +} + +func (w *chatWorker) acquisitionLoop( + ctx context.Context, + workerID uuid.UUID, + manager *runnerManager, + wakeCh <-chan struct{}, +) { + ticker := w.opts.Clock.NewTicker(w.opts.AcquisitionInterval, "chatworker", "acquisition") + defer ticker.Stop() + for { + select { + case <-wakeCh: + w.acquireOnce(ctx, workerID, manager) + case <-ticker.C: + w.acquireOnce(ctx, workerID, manager) + case <-ctx.Done(): + return + } + } +} + +func (w *chatWorker) acquireOnce(ctx context.Context, workerID uuid.UUID, manager *runnerManager) { + attempted := make(map[uuid.UUID]struct{}) + for { + rows, err := w.opts.Store.GetChatWorkerAcquisitionCandidates(ctx, database.GetChatWorkerAcquisitionCandidatesParams{ + StaleSeconds: w.opts.HeartbeatStaleSeconds, + LimitCount: w.opts.AcquisitionBatchSize, + }) + if err != nil { + if ctx.Err() == nil { + w.opts.Logger.Warn(ctx, "chatworker acquisition query failed", slogError(err)) + } + return + } + if len(rows) == 0 { + return + } + newRows := 0 + for _, row := range rows { + if _, ok := attempted[row.ID]; ok { + continue + } + attempted[row.ID] = struct{}{} + newRows++ + if err := w.acquireCandidateSafely(ctx, workerID, manager, row.ID); err != nil { + if ctx.Err() != nil { + return + } + w.opts.Logger.Warn(ctx, "chatworker acquisition candidate failed", slogError(err)) + } + } + if len(rows) < int(w.opts.AcquisitionBatchSize) || newRows == 0 { + return + } + } +} + +var errSkipAcquire = xerrors.New("skip acquire") + +func (w *chatWorker) acquireCandidateSafely( + ctx context.Context, + workerID uuid.UUID, + manager *runnerManager, + chatID uuid.UUID, +) (err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = xerrors.Errorf("chatworker acquisition panic: %v", recovered) + } + }() + return w.acquireCandidate(ctx, workerID, manager, chatID) +} + +func (w *chatWorker) acquireCandidate( + ctx context.Context, + workerID uuid.UUID, + manager *runnerManager, + chatID uuid.UUID, +) error { + runnerID := uuid.New() + machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + chat, err := store.GetChatByID(ctx, chatID) + if errors.Is(err, sql.ErrNoRows) { + return errSkipAcquire + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + queueCount, err := store.CountChatQueuedMessages(ctx, chatID) + if err != nil { + return xerrors.Errorf("count queue: %w", err) + } + if !chatstate.ClassifyExecutionState(chat, queueCount > 0, true).IsRunnable() || chat.Archived { + return errSkipAcquire + } + if chat.WorkerID.Valid && chat.RunnerID.Valid { + stale, err := store.IsChatHeartbeatStale(ctx, database.IsChatHeartbeatStaleParams{ + ChatID: chat.ID, + RunnerID: chat.RunnerID.UUID, + StaleSeconds: w.opts.HeartbeatStaleSeconds, + }) + if err != nil { + return xerrors.Errorf("check heartbeat stale: %w", err) + } + if !stale { + return errSkipAcquire + } + } + _, err = tx.Acquire(chatstate.AcquireInput{WorkerID: workerID, RunnerID: runnerID}) + return err + }) + if errors.Is(err, errSkipAcquire) || errors.Is(err, chatstate.ErrChatNotFound) { + return nil + } + if err != nil { + return err + } + if err := manager.Spawn(ctx, spawnRunnerRequest{ChatID: chatID, WorkerID: workerID, RunnerID: runnerID}); err != nil { + if errAbandon := w.abandonAcquiredChat(ctx, workerID, runnerID, chatID); errAbandon != nil { + return errors.Join(err, errAbandon) + } + return err + } + return nil +} + +func (w *chatWorker) abandonAcquiredChat(ctx context.Context, workerID uuid.UUID, runnerID uuid.UUID, chatID uuid.UUID) error { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), shutdownCleanupTimeout) + defer cancel() + machine := chatstate.NewChatMachine(w.opts.Store, w.opts.Pubsub, chatID) + err := machine.Update(cleanupCtx, func(tx *chatstate.Tx, store database.Store) error { + chat, err := store.GetChatByID(cleanupCtx, chatID) + if errors.Is(err, sql.ErrNoRows) { + return errSkipAcquire + } + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if !chat.WorkerID.Valid || chat.WorkerID.UUID != workerID || !chat.RunnerID.Valid || chat.RunnerID.UUID != runnerID { + return errSkipAcquire + } + _, err = tx.Abandon(chatstate.AbandonInput{}) + return err + }) + if errors.Is(err, errSkipAcquire) || errors.Is(err, chatstate.ErrChatNotFound) { + return nil + } + return err +} diff --git a/coderd/x/chatd/worker_internal_test.go b/coderd/x/chatd/worker_internal_test.go new file mode 100644 index 0000000000..6a635d8418 --- /dev/null +++ b/coderd/x/chatd/worker_internal_test.go @@ -0,0 +1,315 @@ +package chatd //nolint:testpackage // Tests unexported chat worker internals. + +import ( + "context" + "database/sql" + "errors" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +func TestWorker_NewRequiresTaskStarterOrMessagePartBuffer(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + _, err := newChatWorker(nil, chatWorkerOptions{WorkerID: uuid.New(), Store: f.db, Pubsub: f.pubsub}) + require.ErrorContains(t, err, "task starter or message part buffer is required") +} + +func TestWorker_NewRequiresWorkerID(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + opts := testOptions(t, f, newRecordingTaskStarter()) + opts.WorkerID = uuid.Nil + _, err := newChatWorker(nil, opts) + require.ErrorContains(t, err, "worker ID is required") +} + +func TestWorker_UsesConfiguredWorkerID(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + workerID := opts.WorkerID + worker, err := newChatWorker(nil, opts) + require.NoError(t, err) + require.Equal(t, workerID, worker.chatWorkerID()) + require.NoError(t, worker.Start(context.Background())) + require.Equal(t, workerID, worker.chatWorkerID()) + require.NoError(t, worker.Close()) +} + +func TestWorker_AcquiresRunnableChatFromOwnershipHint(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newRecordingTaskStarter() + worker := startWorker(t, testOptions(t, f, starter)) + + call := starter.waitCall(t, taskKindGeneration, chat.ID) + require.Equal(t, worker.chatWorkerID(), call.input.WorkerID) + require.Equal(t, database.ChatStatusRunning, call.input.Status) + require.NotEqual(t, uuid.Nil, call.input.RunnerID) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, worker.chatWorkerID(), latest.WorkerID.UUID) + require.Equal(t, call.input.RunnerID, latest.RunnerID.UUID) + _, err = f.db.GetChatHeartbeat(testutil.Context(t, testutil.WaitShort), database.GetChatHeartbeatParams{ + ChatID: chat.ID, + RunnerID: call.input.RunnerID, + }) + require.NoError(t, err) +} + +func TestWorker_AcquiresRequiresActionChatFromOwnershipHint(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRequiresActionChat(t) + starter := newRecordingTaskStarter() + startWorker(t, testOptions(t, f, starter)) + + call := starter.waitCall(t, taskKindRequiresActionTimeout, chat.ID) + require.Equal(t, database.ChatStatusRequiresAction, call.input.Status) + require.True(t, call.input.RequiresActionDeadlineAt.Valid) +} + +func TestWorker_SkipsFreshlyOwnedChat(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + otherWorker := uuid.New() + otherRunner := uuid.New() + acquireChat(t, f, chat.ID, otherWorker, otherRunner) + starter := newRecordingTaskStarter() + worker := startWorker(t, testOptions(t, f, starter)) + worker.Wake() + + starter.assertNoCall(t) + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.Equal(t, otherWorker, latest.WorkerID.UUID) + require.Equal(t, otherRunner, latest.RunnerID.UUID) +} + +func TestWorker_TwoWorkersRaceSingleOwner(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + firstStarter := newRecordingTaskStarter() + secondStarter := newRecordingTaskStarter() + first := startWorker(t, testOptions(t, f, firstStarter)) + second := startWorker(t, testOptions(t, f, secondStarter)) + + call := waitAnyTaskCall(t, firstStarter, secondStarter, taskKindGeneration, chat.ID) + require.Contains(t, []uuid.UUID{first.chatWorkerID(), second.chatWorkerID()}, call.input.WorkerID) + firstStarter.assertNoCall(t) + secondStarter.assertNoCall(t) + + latest, err := f.db.GetChatByID(testutil.Context(t, testutil.WaitShort), chat.ID) + require.NoError(t, err) + require.True(t, latest.WorkerID.Valid) + require.True(t, latest.RunnerID.Valid) + require.Equal(t, call.input.WorkerID, latest.WorkerID.UUID) + require.Equal(t, call.input.RunnerID, latest.RunnerID.UUID) +} + +func TestWorker_DrainsMultipleRunnableChatsOnWake(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + first := f.createRunningChat(t) + second := f.createRunningChat(t) + third := f.createRunningChat(t) + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + opts.AcquisitionBatchSize = 1 + startWorker(t, opts) + + want := map[uuid.UUID]bool{first.ID: true, second.ID: true, third.ID: true} + for range 3 { + call := starter.waitCall(t, taskKindGeneration, uuid.Nil) + delete(want, call.input.ChatID) + } + require.Empty(t, want) +} + +func TestWorker_DoesNotAcquireIdleOrArchivedChats(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + waiting := f.createRunningChat(t) + finishTurn(t, f, waiting.ID) + errorChat := f.createRunningChat(t) + forceExecutionStateAndPublish(t, f, errorChat.ID, database.ChatStatusError, false) + archived := f.createRunningChat(t) + forceExecutionStateAndPublish(t, f, archived.ID, database.ChatStatusRunning, true) + starter := newRecordingTaskStarter() + worker := startWorker(t, testOptions(t, f, starter)) + worker.Wake() + + starter.assertNoCall(t) +} + +func TestWorker_HeartbeatLoopRefreshesActiveRunnerHeartbeat(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + clock := quartz.NewMock(t) + heartbeatTrap := clock.Trap().NewTicker("chatworker", "heartbeat") + defer heartbeatTrap.Close() + starter := newBlockingTaskStarter(false) + opts := testOptions(t, f, starter) + opts.Clock = clock + opts.HeartbeatInterval = time.Minute + startWorker(t, opts) + heartbeatTrap.MustWait(testutil.Context(t, testutil.WaitLong)).MustRelease(testutil.Context(t, testutil.WaitLong)) + call := starter.waitCall(t, taskKindGeneration, chat.ID) + oldHeartbeat := makeHeartbeatStale(t, f, chat.ID, call.input.RunnerID) + + clock.Advance(time.Minute).MustWait(testutil.Context(t, testutil.WaitLong)) + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + heartbeat, err := f.db.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: chat.ID, + RunnerID: call.input.RunnerID, + }) + return err == nil && heartbeat.HeartbeatAt.After(oldHeartbeat) + }, testutil.IntervalFast, "heartbeat should be refreshed") +} + +func TestWorker_HeartbeatCleanupDeletesStaleRows(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + finishTurn(t, f, chat.ID) + runnerID := uuid.New() + require.NoError(t, f.db.UpsertChatHeartbeat(testutil.Context(t, testutil.WaitShort), database.UpsertChatHeartbeatParams{ + ChatID: chat.ID, + RunnerID: runnerID, + })) + makeHeartbeatStale(t, f, chat.ID, runnerID) + clock := quartz.NewMock(t) + cleanupTrap := clock.Trap().NewTicker("chatworker", "heartbeat-cleanup") + defer cleanupTrap.Close() + starter := newRecordingTaskStarter() + opts := testOptions(t, f, starter) + opts.Clock = clock + opts.HeartbeatCleanupInterval = time.Minute + startWorker(t, opts) + cleanupTrap.MustWait(testutil.Context(t, testutil.WaitLong)).MustRelease(testutil.Context(t, testutil.WaitLong)) + + clock.Advance(time.Minute).MustWait(testutil.Context(t, testutil.WaitLong)) + testutil.Eventually(testutil.Context(t, testutil.WaitLong), t, func(ctx context.Context) bool { + _, err := f.db.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{ + ChatID: chat.ID, + RunnerID: runnerID, + }) + return errors.Is(err, sql.ErrNoRows) + }, testutil.IntervalFast) +} + +func TestWorker_CloseDeletesOwnedHeartbeatsAndPublishesOwnershipHints(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + first := f.createRunningChat(t) + second := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + pubsub := newRecordingPubsub(f.pubsub) + opts := testOptions(t, f, starter) + opts.Pubsub = pubsub + worker := startWorker(t, opts) + callsByChat := make(map[uuid.UUID]taskCall) + for range 2 { + call := starter.waitCall(t, taskKindGeneration, uuid.Nil) + callsByChat[call.input.ChatID] = call + } + require.Contains(t, callsByChat, first.ID) + require.Contains(t, callsByChat, second.ID) + + require.NoError(t, worker.Close()) + for _, call := range callsByChat { + _, err := f.db.GetChatHeartbeat(testutil.Context(t, testutil.WaitShort), database.GetChatHeartbeatParams{ + ChatID: call.input.ChatID, + RunnerID: call.input.RunnerID, + }) + require.ErrorIs(t, err, sql.ErrNoRows) + } + + messages := pubsub.ownershipMessages(t) + seen := make(map[uuid.UUID]bool) + for _, msg := range messages { + seen[msg.ChatID] = true + require.NotZero(t, msg.SnapshotVersion) + } + require.True(t, seen[first.ID], "expected ownership hint for first runner") + require.True(t, seen[second.ID], "expected ownership hint for second runner") +} + +func TestWorker_CloseIsIdempotentAndDoesNotBlock(t *testing.T) { + t.Parallel() + f := newWorkerTestFixture(t) + chat := f.createRunningChat(t) + starter := newBlockingTaskStarter(false) + worker := startWorker(t, testOptions(t, f, starter)) + call := starter.waitCall(t, taskKindGeneration, chat.ID) + + closed := make(chan error, 1) + go func() { + if err := worker.Close(); err != nil { + closed <- err + return + } + closed <- worker.Close() + }() + select { + case err := <-closed: + require.NoError(t, err) + case <-time.After(testutil.WaitLong): + t.Fatal("worker close did not return") + } + select { + case <-call.ctx.Done(): + case <-time.After(testutil.WaitLong): + t.Fatal("active task was not canceled") + } +} + +func waitAnyTaskCall( + t *testing.T, + first *recordingTaskStarter, + second *recordingTaskStarter, + kind taskKind, + chatID uuid.UUID, +) taskCall { + t.Helper() + deadline := time.After(testutil.WaitLong) + for { + select { + case call := <-first.callCh: + if call.kind == kind && call.input.ChatID == chatID { + return call + } + case call := <-second.callCh: + if call.kind == kind && call.input.ChatID == chatID { + return call + } + case <-deadline: + t.Fatal("timed out waiting for either worker to start task") + return taskCall{} + } + } +} + +func requireTaskCanceled(t *testing.T, call taskCall) { + t.Helper() + select { + case <-call.ctx.Done(): + require.True(t, errors.Is(call.ctx.Err(), context.Canceled)) + case <-time.After(testutil.WaitLong): + t.Fatal("task context was not canceled") + } +} diff --git a/coderd/x/chatd/workspace_context_builder.go b/coderd/x/chatd/workspace_context_builder.go new file mode 100644 index 0000000000..9f2aac93a5 --- /dev/null +++ b/coderd/x/chatd/workspace_context_builder.go @@ -0,0 +1,149 @@ +package chatd + +import ( + "context" + "database/sql" + "sync" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/workspacesdk" +) + +// errWorkspaceContextUnavailable is returned by buildWorkspaceContext +// when there is nothing safe to persist for the current committed +// metadata, e.g. the chat has no bound workspace agent or the agent is +// no longer resolvable. Callers treat it as an expected exit. +var errWorkspaceContextUnavailable = xerrors.New("workspace context unavailable") + +// buildWorkspaceContext fetches workspace context for the chat's +// bound workspace agent and returns durable chatstate.Message values +// for the generation action to commit. It returns +// errWorkspaceContextUnavailable when there is nothing safe to +// persist for the current committed metadata. +func (server *Server) buildWorkspaceContext( + ctx context.Context, + input workspaceContextBuildInput, +) (workspaceContextBuildResult, error) { + chat := input.Chat + if !chat.WorkspaceID.Valid || !chat.AgentID.Valid { + return workspaceContextBuildResult{}, errWorkspaceContextUnavailable + } + logger := server.logger.With( + slog.F("chat_id", chat.ID), + slog.F("owner_id", chat.OwnerID), + ) + + // Build a per-call workspace context with the latest committed + // chat snapshot so getWorkspaceAgent and getWorkspaceConn dial + // the agent we actually want to fetch context from. + currentChat := chat + var chatStateMu sync.Mutex + wsCtx := turnWorkspaceContext{ + server: server, + chatStateMu: &chatStateMu, + currentChat: ¤tChat, + loadChatSnapshot: server.db.GetChatByID, + } + defer wsCtx.close() + + parts, expectedAgentID := server.fetchContextForBuild(ctx, chat, &wsCtx, logger) + // If the workspace or agent is gone, fall back to no-op so the + // generation action exits without committing stale context. + if expectedAgentID == uuid.Nil { + return workspaceContextBuildResult{}, errWorkspaceContextUnavailable + } + + hasContent := false + hasContextFilePart := false + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeContextFile { + hasContextFilePart = true + if part.ContextFileContent != "" { + hasContent = true + } + } + } + + agentID := uuid.NullUUID{UUID: expectedAgentID, Valid: true} + + // If we have no content but the agent is known, commit a blank + // context-file marker (sentinel) so subsequent turns skip the + // workspace-agent dial and the decision helper observes the + // attempt in committed history. This applies whether the + // workspace connection succeeded but returned no AGENTS.md, or + // the agent's context config fetch failed: in both cases we + // have a known agent and committing a sentinel breaks the + // otherwise-infinite decision loop. + if !hasContent { + if !hasContextFilePart { + parts = append([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeContextFile, + ContextFileAgentID: agentID, + }}, parts...) + } + } + + content, err := chatprompt.MarshalParts(parts) + if err != nil { + return workspaceContextBuildResult{}, xerrors.Errorf("marshal workspace context parts: %w", err) + } + + modelConfigID := chat.LastModelConfigID + msg := chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: sql.NullString{String: input.ActiveAPIKeyID, Valid: input.ActiveAPIKeyID != ""}, + } + + // Update the cache column so subsequent turns can read the last + // injected context without scanning messages. This is a + // best-effort write that does not mutate chat history; the + // generation action separately commits the durable message + // below. + stripped := make([]codersdk.ChatMessagePart, len(parts)) + copy(stripped, parts) + for i := range stripped { + stripped[i].StripInternal() + } + server.updateLastInjectedContext(ctx, chat.ID, stripped) + + return workspaceContextBuildResult{Messages: []chatstate.Message{msg}}, nil +} + +// fetchContextForBuild fetches workspace context parts from the +// agent, returning the parts to persist. expectedAgentID is the agent +// ID the fetch was bound to, or uuid.Nil if the agent could not be +// resolved. +func (server *Server) fetchContextForBuild( + ctx context.Context, + chat database.Chat, + wsCtx *turnWorkspaceContext, + logger slog.Logger, +) (parts []codersdk.ChatMessagePart, expectedAgentID uuid.UUID) { + agent, agentParts, _, _ := server.fetchWorkspaceContext( + ctx, chat, wsCtx.getWorkspaceAgent, + func(instructionCtx context.Context) (workspacesdk.AgentConn, error) { + if _, _, err := wsCtx.workspaceAgentIDForConn(instructionCtx); err != nil { + return nil, err + } + return wsCtx.getWorkspaceConn(instructionCtx) + }, + ) + if agent == nil { + // fetchWorkspaceContext returns nil for the agent when the + // chat has no valid workspace or the agent lookup fails. + logger.Debug(ctx, "workspace context build: workspace agent not resolvable") + return nil, uuid.Nil + } + return agentParts, agent.ID +} diff --git a/codersdk/chats.go b/codersdk/chats.go index 4fce8a32c2..f07c1889be 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -93,6 +93,7 @@ const ( ChatStatusCompleted ChatStatus = "completed" ChatStatusError ChatStatus = "error" ChatStatusRequiresAction ChatStatus = "requires_action" + ChatStatusInterrupting ChatStatus = "interrupting" ) // ChatClientType indicates whether a chat was created from the @@ -1502,6 +1503,8 @@ const ( ChatStreamEventTypeQueueUpdate ChatStreamEventType = "queue_update" ChatStreamEventTypeRetry ChatStreamEventType = "retry" ChatStreamEventTypeActionRequired ChatStreamEventType = "action_required" + ChatStreamEventTypePreviewReset ChatStreamEventType = "preview_reset" + ChatStreamEventTypeHistoryReset ChatStreamEventType = "history_reset" ) // ChatQueuedMessage represents a queued message waiting to be processed. @@ -1515,8 +1518,11 @@ type ChatQueuedMessage struct { // ChatStreamMessagePart is a streamed message part update. type ChatStreamMessagePart struct { - Role ChatMessageRole `json:"role,omitempty"` - Part ChatMessagePart `json:"part"` + Role ChatMessageRole `json:"role,omitempty"` + Part ChatMessagePart `json:"part"` + HistoryVersion int64 `json:"history_version,omitempty"` + GenerationAttempt int64 `json:"generation_attempt,omitempty"` + Seq int64 `json:"seq,omitempty"` } // ChatStreamStatus represents an updated chat status. @@ -3250,6 +3256,22 @@ func (c *ExperimentalClient) InterruptChat(ctx context.Context, chatID uuid.UUID return chat, json.NewDecoder(res.Body).Decode(&chat) } +// ReconcileInvalidChatState recovers a chat stuck in an invalid +// execution state, moving it into an error state from which the caller +// can send a new message or edit history to continue. +func (c *ExperimentalClient) ReconcileInvalidChatState(ctx context.Context, chatID uuid.UUID) (Chat, error) { + res, err := c.Request(ctx, http.MethodPost, fmt.Sprintf("/api/experimental/chats/%s/reconcile-invalid", chatID), nil) + if err != nil { + return Chat{}, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return Chat{}, ReadBodyAsError(res) + } + var chat Chat + return chat, json.NewDecoder(res.Body).Decode(&chat) +} + // RegenerateChatTitle requests the server to regenerate the chat's // title using richer conversation context. func (c *ExperimentalClient) RegenerateChatTitle(ctx context.Context, chatID uuid.UUID) (Chat, error) { diff --git a/codersdk/chats_test.go b/codersdk/chats_test.go index 5c6201ac7a..a21b5ae7e2 100644 --- a/codersdk/chats_test.go +++ b/codersdk/chats_test.go @@ -172,6 +172,42 @@ func TestChatErrorKind_JSONRoundTrip(t *testing.T) { require.Equal(t, codersdk.ChatErrorKindUsageLimit, decodedRetry.Kind) } +func TestChatStreamEvent_JSONRoundTripIncludesResetTypesAndPartMetadata(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + events := []codersdk.ChatStreamEvent{ + {Type: codersdk.ChatStreamEventTypePreviewReset, ChatID: chatID}, + {Type: codersdk.ChatStreamEventTypeHistoryReset, ChatID: chatID}, + { + Type: codersdk.ChatStreamEventTypeMessagePart, + ChatID: chatID, + MessagePart: &codersdk.ChatStreamMessagePart{ + Role: codersdk.ChatMessageRoleAssistant, + Part: codersdk.ChatMessageText("partial"), + HistoryVersion: 12, + GenerationAttempt: 3, + Seq: 4, + }, + }, + } + data, err := json.Marshal(events) + require.NoError(t, err) + require.Contains(t, string(data), `"type":"preview_reset"`) + require.Contains(t, string(data), `"type":"history_reset"`) + require.Contains(t, string(data), `"history_version":12`) + require.Contains(t, string(data), `"generation_attempt":3`) + require.Contains(t, string(data), `"seq":4`) + + var decoded []codersdk.ChatStreamEvent + require.NoError(t, json.Unmarshal(data, &decoded)) + require.Equal(t, codersdk.ChatStreamEventTypePreviewReset, decoded[0].Type) + require.Equal(t, codersdk.ChatStreamEventTypeHistoryReset, decoded[1].Type) + require.Equal(t, int64(12), decoded[2].MessagePart.HistoryVersion) + require.Equal(t, int64(3), decoded[2].MessagePart.GenerationAttempt) + require.Equal(t, int64(4), decoded[2].MessagePart.Seq) +} + func TestChatMessagePart_StripInternal(t *testing.T) { t.Parallel() diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index 479c670bfd..92fbc1d812 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -212,11 +212,7 @@ deployment. They will always be available from the agent. | `coderd_chatd_prompt_size_bytes` | histogram | Estimated byte size of the prompt per LLM request. | `model` `provider` | | `coderd_chatd_steps_total` | counter | Total agentic loop steps across all chats. | `model` `provider` | | `coderd_chatd_stream_buffer_dropped_total` | counter | Number of chat stream buffer events dropped due to the per-chat buffer cap. | | -| `coderd_chatd_stream_buffer_events` | gauge | Sum of current buffer lengths across all chat streams. | | -| `coderd_chatd_stream_buffer_size_max` | gauge | Maximum current buffer length across all chat streams. | | | `coderd_chatd_stream_retries_total` | counter | Total LLM stream retries. | `chain_broken` `kind` `model` `provider` | -| `coderd_chatd_stream_subscribers` | gauge | Current number of chat stream subscribers across all chat streams. | | -| `coderd_chatd_streams_active` | gauge | Current number of chat stream state entries (in-flight plus retained). | | | `coderd_chatd_tool_errors_total` | counter | Total tool calls that returned an error result. | `model` `provider` `tool_name` | | `coderd_chatd_tool_result_size_bytes` | histogram | Size in bytes of each tool execution result. | `model` `provider` `tool_name` | | `coderd_chatd_ttft_seconds` | histogram | Time-to-first-token: wall time from LLM request to first streamed chunk. | `model` `provider` | diff --git a/docs/admin/security/audit-logs.md b/docs/admin/security/audit-logs.md index e7bc259b73..c572565d34 100644 --- a/docs/admin/security/audit-logs.md +++ b/docs/admin/security/audit-logs.md @@ -25,7 +25,7 @@ We track the following resources: | AuditableGroupAiBudget
write, delete | |
FieldTracked
created_atfalse
group_idfalse
group_namefalse
spend_limittrue
spend_limit_microsfalse
updated_atfalse
| | AuditableOrganizationMember
| |
FieldTracked
created_attrue
organization_idfalse
rolestrue
updated_attrue
user_idtrue
usernametrue
| | AuditableUserAiBudgetOverride
write, delete | |
FieldTracked
created_atfalse
group_idtrue
group_nametrue
spend_limittrue
spend_limit_microsfalse
updated_atfalse
user_idfalse
usernamefalse
| -| Chat
create, write | |
FieldTracked
agent_idfalse
archivedtrue
build_idfalse
client_typefalse
created_atfalse
dynamic_toolsfalse
group_acltrue
heartbeat_atfalse
idtrue
labelstrue
last_errorfalse
last_injected_contextfalse
last_model_config_idfalse
last_read_message_idfalse
last_turn_summaryfalse
mcp_server_idstrue
modetrue
organization_idfalse
owner_idtrue
owner_namefalse
owner_usernamefalse
parent_chat_idfalse
pin_ordertrue
plan_modefalse
root_chat_idfalse
started_atfalse
statusfalse
titletrue
updated_atfalse
user_acltrue
worker_idfalse
workspace_idtrue
| +| Chat
create, write | |
FieldTracked
agent_idfalse
archivedtrue
build_idfalse
client_typefalse
created_atfalse
dynamic_toolsfalse
generation_attemptfalse
group_acltrue
heartbeat_atfalse
history_versionfalse
idtrue
labelstrue
last_errorfalse
last_injected_contextfalse
last_model_config_idfalse
last_read_message_idfalse
last_turn_summaryfalse
mcp_server_idstrue
modetrue
organization_idfalse
owner_idtrue
owner_namefalse
owner_usernamefalse
parent_chat_idfalse
pin_ordertrue
plan_modefalse
queue_versionfalse
requires_action_deadline_atfalse
retry_statefalse
retry_state_versionfalse
root_chat_idfalse
runner_idfalse
snapshot_versionfalse
started_atfalse
statusfalse
titletrue
updated_atfalse
user_acltrue
worker_idfalse
workspace_idtrue
| | CustomRole
| |
FieldTracked
created_atfalse
display_nametrue
idfalse
is_systemfalse
member_permissionstrue
nametrue
org_permissionstrue
organization_idfalse
site_permissionstrue
updated_atfalse
user_permissionstrue
| | GitSSHKey
create | |
FieldTracked
created_atfalse
private_keytrue
private_key_key_idfalse
public_keytrue
updated_atfalse
user_idtrue
| | GroupSyncSettings
| |
FieldTracked
auto_create_missing_groupstrue
fieldtrue
legacy_group_name_mappingfalse
mappingtrue
regex_filtertrue
| diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index d66d4b3669..7047299d3f 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -300,7 +300,7 @@ Status Code **200** | `kind` | `auth`, `config`, `generic`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `stream_silence_timeout`, `timeout`, `usage_limit` | | `type` | `context-file`, `file`, `file-reference`, `reasoning`, `skill`, `source`, `text`, `tool-call`, `tool-result` | | `plan_mode` | `plan` | -| `status` | `completed`, `error`, `paused`, `pending`, `requires_action`, `running`, `waiting` | +| `status` | `completed`, `error`, `interrupting`, `paused`, `pending`, `requires_action`, `running`, `waiting` | To perform this operation, you must be authenticated. [Learn more](authentication.md). @@ -2279,6 +2279,319 @@ message in the chat. To perform this operation, you must be authenticated. [Learn more](authentication.md). +## Reconcile invalid chat state + +### Code samples + +```shell +# Example request using curl +curl -X POST http://coder-server:8080/api/experimental/chats/{chat}/reconcile-invalid \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`POST /api/experimental/chats/{chat}/reconcile-invalid` + +Experimental: this endpoint is subject to change. + +### Parameters + +| Name | In | Type | Required | Description | +|--------|------|--------------|----------|-------------| +| `chat` | path | string(uuid) | true | Chat ID | + +### Example responses + +> 200 Response + +```json +{ + "agent_id": "2b1e3b65-2c04-4fa2-a2d7-467901e98978", + "archived": true, + "build_id": "bfb1f3fa-bf7b-43a5-9e0b-26cc050e44cb", + "children": [ + { + "agent_id": "2b1e3b65-2c04-4fa2-a2d7-467901e98978", + "archived": true, + "build_id": "bfb1f3fa-bf7b-43a5-9e0b-26cc050e44cb", + "children": [], + "client_type": "ui", + "created_at": "2019-08-24T14:15:22Z", + "diff_status": { + "additions": 0, + "approved": true, + "author_avatar_url": "string", + "author_login": "string", + "base_branch": "string", + "changed_files": 0, + "changes_requested": true, + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "commits": 0, + "deletions": 0, + "head_branch": "string", + "pr_number": 0, + "pull_request_draft": true, + "pull_request_state": "string", + "pull_request_title": "string", + "refreshed_at": "2019-08-24T14:15:22Z", + "reviewer_count": 0, + "stale_at": "2019-08-24T14:15:22Z", + "url": "string" + }, + "files": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "mime_type": "string", + "name": "string", + "organization_id": "7c60d51f-b44e-4682-87d6-449835ea4de6", + "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" + } + ], + "has_unread": true, + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "labels": { + "property1": "string", + "property2": "string" + }, + "last_error": { + "detail": "string", + "kind": "generic", + "message": "string", + "provider": "string", + "retryable": true, + "status_code": 0 + }, + "last_injected_context": [ + { + "args": [ + 0 + ], + "args_delta": "string", + "completed_at": "2019-08-24T14:15:22Z", + "content": "string", + "context_file_agent_id": { + "uuid": "string", + "valid": true + }, + "context_file_content": "string", + "context_file_directory": "string", + "context_file_os": "string", + "context_file_path": "string", + "context_file_skill_meta_file": "string", + "context_file_truncated": true, + "created_at": "2019-08-24T14:15:22Z", + "data": [ + 0 + ], + "end_line": 0, + "file_id": { + "uuid": "string", + "valid": true + }, + "file_name": "string", + "is_error": true, + "is_media": true, + "mcp_server_config_id": { + "uuid": "string", + "valid": true + }, + "media_type": "string", + "name": "string", + "parsed_commands": [ + [ + "string" + ] + ], + "provider_executed": true, + "provider_metadata": [ + 0 + ], + "result": [ + 0 + ], + "result_delta": "string", + "result_reset": true, + "signature": "string", + "skill_description": "string", + "skill_dir": "string", + "skill_name": "string", + "source_id": "string", + "start_line": 0, + "text": "string", + "title": "string", + "tool_call_id": "string", + "tool_name": "string", + "type": "text", + "url": "string" + } + ], + "last_model_config_id": "30ebb95f-c255-4759-9429-89aa4ec1554c", + "last_turn_summary": "string", + "mcp_server_ids": [ + "497f6eca-6276-4993-bfeb-53cbbbba6f08" + ], + "organization_id": "7c60d51f-b44e-4682-87d6-449835ea4de6", + "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05", + "owner_name": "string", + "owner_username": "string", + "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", + "pin_order": 0, + "plan_mode": "plan", + "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", + "shared": true, + "status": "waiting", + "title": "string", + "updated_at": "2019-08-24T14:15:22Z", + "warnings": [ + "string" + ], + "workspace_id": "0967198e-ec7b-4c6b-b4d3-f71244cadbe9" + } + ], + "client_type": "ui", + "created_at": "2019-08-24T14:15:22Z", + "diff_status": { + "additions": 0, + "approved": true, + "author_avatar_url": "string", + "author_login": "string", + "base_branch": "string", + "changed_files": 0, + "changes_requested": true, + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "commits": 0, + "deletions": 0, + "head_branch": "string", + "pr_number": 0, + "pull_request_draft": true, + "pull_request_state": "string", + "pull_request_title": "string", + "refreshed_at": "2019-08-24T14:15:22Z", + "reviewer_count": 0, + "stale_at": "2019-08-24T14:15:22Z", + "url": "string" + }, + "files": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "mime_type": "string", + "name": "string", + "organization_id": "7c60d51f-b44e-4682-87d6-449835ea4de6", + "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05" + } + ], + "has_unread": true, + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "labels": { + "property1": "string", + "property2": "string" + }, + "last_error": { + "detail": "string", + "kind": "generic", + "message": "string", + "provider": "string", + "retryable": true, + "status_code": 0 + }, + "last_injected_context": [ + { + "args": [ + 0 + ], + "args_delta": "string", + "completed_at": "2019-08-24T14:15:22Z", + "content": "string", + "context_file_agent_id": { + "uuid": "string", + "valid": true + }, + "context_file_content": "string", + "context_file_directory": "string", + "context_file_os": "string", + "context_file_path": "string", + "context_file_skill_meta_file": "string", + "context_file_truncated": true, + "created_at": "2019-08-24T14:15:22Z", + "data": [ + 0 + ], + "end_line": 0, + "file_id": { + "uuid": "string", + "valid": true + }, + "file_name": "string", + "is_error": true, + "is_media": true, + "mcp_server_config_id": { + "uuid": "string", + "valid": true + }, + "media_type": "string", + "name": "string", + "parsed_commands": [ + [ + "string" + ] + ], + "provider_executed": true, + "provider_metadata": [ + 0 + ], + "result": [ + 0 + ], + "result_delta": "string", + "result_reset": true, + "signature": "string", + "skill_description": "string", + "skill_dir": "string", + "skill_name": "string", + "source_id": "string", + "start_line": 0, + "text": "string", + "title": "string", + "tool_call_id": "string", + "tool_name": "string", + "type": "text", + "url": "string" + } + ], + "last_model_config_id": "30ebb95f-c255-4759-9429-89aa4ec1554c", + "last_turn_summary": "string", + "mcp_server_ids": [ + "497f6eca-6276-4993-bfeb-53cbbbba6f08" + ], + "organization_id": "7c60d51f-b44e-4682-87d6-449835ea4de6", + "owner_id": "8826ee2e-7933-4665-aef2-2393f84a0d05", + "owner_name": "string", + "owner_username": "string", + "parent_chat_id": "c3609ee6-3b11-4a93-b9ae-e4fabcc99359", + "pin_order": 0, + "plan_mode": "plan", + "root_chat_id": "2898031c-fdce-4e3e-8c53-4481dd42fcd7", + "shared": true, + "status": "waiting", + "title": "string", + "updated_at": "2019-08-24T14:15:22Z", + "warnings": [ + "string" + ], + "workspace_id": "0967198e-ec7b-4c6b-b4d3-f71244cadbe9" +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.Chat](schemas.md#codersdkchat) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). + ## Stream chat events via WebSockets ### Code samples @@ -2406,6 +2719,8 @@ Experimental: this endpoint is subject to change. } }, "message_part": { + "generation_attempt": 0, + "history_version": 0, "part": { "args": [ 0 @@ -2468,7 +2783,8 @@ Experimental: this endpoint is subject to change. "type": "text", "url": "string" }, - "role": "system" + "role": "system", + "seq": 0 }, "queued_messages": [ { diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 6ffdefb5f9..220d0c3e8f 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -3256,9 +3256,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in #### Enumerated Values -| Value(s) | -|------------------------------------------------------------------------------------| -| `completed`, `error`, `paused`, `pending`, `requires_action`, `running`, `waiting` | +| Value(s) | +|----------------------------------------------------------------------------------------------------| +| `completed`, `error`, `interrupting`, `paused`, `pending`, `requires_action`, `running`, `waiting` | ## codersdk.ChatStreamActionRequired @@ -3384,6 +3384,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in } }, "message_part": { + "generation_attempt": 0, + "history_version": 0, "part": { "args": [ 0 @@ -3446,7 +3448,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "type": "text", "url": "string" }, - "role": "system" + "role": "system", + "seq": 0 }, "queued_messages": [ { @@ -3560,14 +3563,16 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in #### Enumerated Values -| Value(s) | -|------------------------------------------------------------------------------------------| -| `action_required`, `error`, `message`, `message_part`, `queue_update`, `retry`, `status` | +| Value(s) | +|----------------------------------------------------------------------------------------------------------------------------| +| `action_required`, `error`, `history_reset`, `message`, `message_part`, `preview_reset`, `queue_update`, `retry`, `status` | ## codersdk.ChatStreamMessagePart ```json { + "generation_attempt": 0, + "history_version": 0, "part": { "args": [ 0 @@ -3630,16 +3635,20 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "type": "text", "url": "string" }, - "role": "system" + "role": "system", + "seq": 0 } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|--------|------------------------------------------------------|----------|--------------|-------------| -| `part` | [codersdk.ChatMessagePart](#codersdkchatmessagepart) | false | | | -| `role` | [codersdk.ChatMessageRole](#codersdkchatmessagerole) | false | | | +| Name | Type | Required | Restrictions | Description | +|----------------------|------------------------------------------------------|----------|--------------|-------------| +| `generation_attempt` | integer | false | | | +| `history_version` | integer | false | | | +| `part` | [codersdk.ChatMessagePart](#codersdkchatmessagepart) | false | | | +| `role` | [codersdk.ChatMessageRole](#codersdkchatmessagerole) | false | | | +| `seq` | integer | false | | | ## codersdk.ChatStreamRetry diff --git a/enterprise/audit/table.go b/enterprise/audit/table.go index 7b050584ae..bcbc2b469b 100644 --- a/enterprise/audit/table.go +++ b/enterprise/audit/table.go @@ -436,38 +436,46 @@ var auditableResourcesTypes = map[any]map[string]Action{ "deleted_at": ActionIgnore, // Changes, but is implicit when a delete event is fired. }, &database.Chat{}: { - "id": ActionTrack, - "owner_id": ActionTrack, - "owner_username": ActionIgnore, - "owner_name": ActionIgnore, - "organization_id": ActionIgnore, // Never changes after creation. - "workspace_id": ActionTrack, - "build_id": ActionIgnore, // Internal lifecycle. - "agent_id": ActionIgnore, // Internal lifecycle. - "title": ActionSecret, // May contain sensitive content. - "status": ActionIgnore, // Churns every message. - "worker_id": ActionIgnore, // Internal. - "started_at": ActionIgnore, - "heartbeat_at": ActionIgnore, // Internal. - "created_at": ActionIgnore, // Never changes. - "updated_at": ActionIgnore, // Bumped on every mutation. - "parent_chat_id": ActionIgnore, // Immutable after creation. - "root_chat_id": ActionIgnore, // Immutable after creation. - "last_model_config_id": ActionIgnore, // Churns every message. - "archived": ActionTrack, - "last_error": ActionIgnore, // Internal. - "last_turn_summary": ActionIgnore, // Internal cached display text. - "mode": ActionTrack, - "mcp_server_ids": ActionTrack, - "labels": ActionTrack, - "user_acl": ActionTrack, - "group_acl": ActionTrack, - "pin_order": ActionTrack, - "last_read_message_id": ActionIgnore, // User-scoped read cursor. - "last_injected_context": ActionIgnore, // Internal lifecycle. - "dynamic_tools": ActionIgnore, // Internal lifecycle. - "plan_mode": ActionIgnore, // Can flip back and forth during a session. - "client_type": ActionIgnore, // Set at creation. + "id": ActionTrack, + "owner_id": ActionTrack, + "owner_username": ActionIgnore, + "owner_name": ActionIgnore, + "organization_id": ActionIgnore, // Never changes after creation. + "workspace_id": ActionTrack, + "build_id": ActionIgnore, // Internal lifecycle. + "agent_id": ActionIgnore, // Internal lifecycle. + "title": ActionSecret, // May contain sensitive content. + "status": ActionIgnore, // Churns every message. + "worker_id": ActionIgnore, // Internal. + "started_at": ActionIgnore, + "heartbeat_at": ActionIgnore, // Internal. + "created_at": ActionIgnore, // Never changes. + "updated_at": ActionIgnore, // Bumped on every mutation. + "parent_chat_id": ActionIgnore, // Immutable after creation. + "root_chat_id": ActionIgnore, // Immutable after creation. + "last_model_config_id": ActionIgnore, // Churns every message. + "archived": ActionTrack, + "last_error": ActionIgnore, // Internal. + "last_turn_summary": ActionIgnore, // Internal cached display text. + "mode": ActionTrack, + "mcp_server_ids": ActionTrack, + "labels": ActionTrack, + "user_acl": ActionTrack, + "group_acl": ActionTrack, + "pin_order": ActionTrack, + "last_read_message_id": ActionIgnore, // User-scoped read cursor. + "last_injected_context": ActionIgnore, // Internal lifecycle. + "dynamic_tools": ActionIgnore, // Internal lifecycle. + "plan_mode": ActionIgnore, // Can flip back and forth during a session. + "client_type": ActionIgnore, // Set at creation. + "snapshot_version": ActionIgnore, // Internal state machine version. + "history_version": ActionIgnore, // Internal state machine version. + "queue_version": ActionIgnore, // Internal state machine version. + "retry_state": ActionIgnore, // Internal transient retry UI state. + "retry_state_version": ActionIgnore, // Internal state machine version. + "generation_attempt": ActionIgnore, // Internal retry counter. + "runner_id": ActionIgnore, // Internal ownership identifier. + "requires_action_deadline_at": ActionIgnore, // Internal pending-action deadline. }, &database.UserSkill{}: { "id": ActionTrack, diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 941544f526..f82fb430ae 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -159,10 +159,14 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { } var replicaManagerPtr atomic.Pointer[replicasync.Manager] + var api *API resolveReplicaAddress := func( _ context.Context, replicaID uuid.UUID, ) (string, bool) { + if api != nil && api.AGPL != nil && replicaID == api.AGPL.ID && api.AGPL.AccessURL != nil { + return api.AGPL.AccessURL.String(), true + } manager := replicaManagerPtr.Load() if manager == nil { return "", false @@ -180,7 +184,7 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { return "", false } - api := &API{ + api = &API{ ctx: ctx, cancel: cancelFunc, Options: options, @@ -207,17 +211,13 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { replicaHTTPClient = http.DefaultClient } // Use a closure that captures api by reference so it can access - // api.AGPL.ID after coderd.New is called. The SubscribeFn is - // only invoked from Subscribe, which happens after init. - options.Options.ChatSubscribeFn = entchatd.NewMultiReplicaSubscribeFn(entchatd.MultiReplicaSubscribeConfig{ + // api.AGPL.ID after coderd.New is called. The parts dialer is + // only invoked from stream subscriptions, which happen after init. + options.Options.ChatStreamPartsDialer = entchatd.NewStreamPartsDialer(entchatd.StreamPartsDialerConfig{ ResolveReplicaAddress: resolveReplicaAddress, ReplicaHTTPClient: replicaHTTPClient, ReplicaIDFn: func() uuid.UUID { - id := api.AGPL.ID - if id == uuid.Nil { - return uuid.New() - } - return id + return api.AGPL.ID }, }) diff --git a/enterprise/coderd/x/chatd/chatd.go b/enterprise/coderd/x/chatd/chatd.go index 8301e1d191..d3b2d91fa3 100644 --- a/enterprise/coderd/x/chatd/chatd.go +++ b/enterprise/coderd/x/chatd/chatd.go @@ -2,25 +2,16 @@ package chatd import ( "context" - "errors" - "fmt" "net/http" "net/url" - "strconv" "strings" - "time" "github.com/google/uuid" "golang.org/x/xerrors" - "cdr.dev/slog/v3" - "github.com/coder/coder/v2/coderd/database" osschatd "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/codersdk" - "github.com/coder/quartz" - "github.com/coder/retry" "github.com/coder/websocket" - "github.com/coder/websocket/wsjson" ) // RelaySourceHeader marks replica-relayed stream requests. @@ -29,28 +20,10 @@ const RelaySourceHeader = "X-Coder-Relay-Source-Replica" const ( authorizationHeader = "Authorization" cookieHeader = "Cookie" - - // relayDrainTimeout is how long an established relay is - // kept open after the chat leaves running state, giving - // buffered snapshot events time to be forwarded before - // the relay is torn down. - relayDrainTimeout = 200 * time.Millisecond - - // Retry knobs for the cross-replica relay handshake. Uses the - // github.com/coder/retry defaults (φ-growth, no jitter) but drives - // the delay manually because retry.Retrier.Wait uses time.After, - // which isn't compatible with quartz.Clock determinism in tests. - relayRetryFloor = 500 * time.Millisecond // first retry matches old fixed delay - relayRetryCeil = 15 * time.Second // cap stall before tear-down - // After this many reconnect retries the relay leg is torn down. - // Total dial attempts = 1 initial dial + relayMaxRetries. - relayMaxRetries = 6 ) // RelayDialError wraps a failed relay handshake. HTTPStatus is 0 -// when the failure happened before a response (DNS, TCP, TLS, -// timeout, context cancel); otherwise it carries the peer's status -// code for the reconnect loop to classify. +// when the failure happened before a response. type RelayDialError struct { HTTPStatus int Err error @@ -60,661 +33,59 @@ func (e *RelayDialError) Error() string { return e.Err.Error() } func (e *RelayDialError) Unwrap() error { return e.Err } // IsUnrecoverable reports whether retrying with the same captured -// session token is futile. Only 401/403 qualify - the token is dead -// or the peer won't authorize it. 5xx, 429, network, and context -// errors fall through to backoff. +// session token is futile. func (e *RelayDialError) IsUnrecoverable() bool { return e.HTTPStatus == http.StatusUnauthorized || e.HTTPStatus == http.StatusForbidden } -// MultiReplicaSubscribeConfig holds the dependencies for multi-replica chat -// subscription. ReplicaIDFn is called lazily because the -// replica ID may not be known at construction time. -// -// DialerFn, when set, overrides the default WebSocket relay -// dialer. This is used in tests to inject mock relay behavior -// without requiring real HTTP servers. -type MultiReplicaSubscribeConfig struct { +// StreamPartsDialerConfig holds dependencies for multi-replica stream parts. +type StreamPartsDialerConfig struct { ResolveReplicaAddress func(context.Context, uuid.UUID) (string, bool) ReplicaHTTPClient *http.Client ReplicaIDFn func() uuid.UUID - DialerFn func( - ctx context.Context, - chatID uuid.UUID, - workerID uuid.UUID, - requestHeader http.Header, - ) ( - snapshot []codersdk.ChatStreamEvent, - parts <-chan codersdk.ChatStreamEvent, - cancel func(), - err error, - ) - // Clock is used for creating timers. In production use - // quartz.NewReal(); in tests use quartz.NewMock(t) to - // control reconnect timing deterministically. - Clock quartz.Clock + DialerFn func(context.Context, osschatd.StreamPartsDialInput) (osschatd.StreamPartsSession, error) } -// dial returns the configured dialer, preferring DialerFn (tests) -// over the real dialRelay. Returns nil when relay is not configured. -func (c MultiReplicaSubscribeConfig) dial() func( +// NewStreamPartsDialer returns a dialer for the owning replica's parts endpoint. +func NewStreamPartsDialer(cfg StreamPartsDialerConfig) osschatd.StreamPartsDialer { + return func(ctx context.Context, input osschatd.StreamPartsDialInput) (osschatd.StreamPartsSession, error) { + if cfg.DialerFn != nil { + return cfg.DialerFn(ctx, input) + } + return dialRelayParts(ctx, input, cfg) + } +} + +func dialRelayParts( ctx context.Context, - chatID uuid.UUID, - workerID uuid.UUID, - requestHeader http.Header, -) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, -) { - if c.DialerFn != nil { - return c.DialerFn + input osschatd.StreamPartsDialInput, + cfg StreamPartsDialerConfig, +) (osschatd.StreamPartsSession, error) { + if cfg.ResolveReplicaAddress == nil { + return nil, &RelayDialError{Err: xerrors.New("dial relay stream parts: resolver not configured")} } - if c.ResolveReplicaAddress == nil { - return nil - } - return func( - ctx context.Context, - chatID uuid.UUID, - workerID uuid.UUID, - requestHeader http.Header, - ) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, - ) { - return dialRelay(ctx, chatID, workerID, requestHeader, c, c.clock()) - } -} - -// clock returns the quartz.Clock to use. Defaults to a real clock -// when not set. -func (c MultiReplicaSubscribeConfig) clock() quartz.Clock { - if c.Clock != nil { - return c.Clock - } - return quartz.NewReal() -} - -// NewMultiReplicaSubscribeFn returns a SubscribeFn that manages -// relay connections to remote replicas and returns relay -// message_part events only. OSS handles pubsub subscription, -// message catch-up, queue updates, status forwarding, and local -// parts merging. -// -//nolint:gocognit // Complexity is inherent to the multi-source merge loop. -func NewMultiReplicaSubscribeFn( - cfg MultiReplicaSubscribeConfig, -) osschatd.SubscribeFn { - return func(ctx context.Context, params osschatd.SubscribeFnParams) <-chan codersdk.ChatStreamEvent { - chatID := params.ChatID - requestHeader := params.RequestHeader - logger := params.Logger - - var relayCancel func() - var relayParts <-chan codersdk.ChatStreamEvent - - // If the chat is currently running on a different worker - // and we have a remote parts provider, open an initial - // relay synchronously so the caller gets in-flight - // message_part events right away. - var initialRelaySnapshot []codersdk.ChatStreamEvent - if params.Chat.Status == database.ChatStatusRunning && - params.Chat.WorkerID.Valid && - params.Chat.WorkerID.UUID != params.WorkerID && - cfg.dial() != nil { - snapshot, parts, cancel, err := cfg.dial()(ctx, chatID, params.Chat.WorkerID.UUID, requestHeader) - if err == nil { - relayCancel = cancel - relayParts = parts - // Collect relay message_parts to forward at the - // start of the merge goroutine. - for _, event := range snapshot { - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - initialRelaySnapshot = append(initialRelaySnapshot, event) - } - } - } else { - logger.Warn(ctx, "failed to open initial relay for chat stream", - slog.F("chat_id", chatID), - slog.Error(err), - ) - } - } - - // Merge all event sources. - mergedEvents := make(chan codersdk.ChatStreamEvent, 128) - // Channel for async relay establishment. - type relayResult struct { - parts <-chan codersdk.ChatStreamEvent - cancel func() - workerID uuid.UUID // the worker this dial targeted - // err and parts are mutually exclusive: success sets - // parts; failure sets err (unwrap to *RelayDialError - // for classification). - err error - } - relayReadyCh := make(chan relayResult, 4) - - // Reset on successful dial or when the relay target - // changes, so a fresh target starts at the floor delay. - retryState := newRelayRetryState() - // Per-dial context so in-flight dials can be canceled when - // a new dial is initiated or the relay is closed. - var dialCancel context.CancelFunc - - // expectedWorkerID tracks which replica we expect the next - // relay result to target. Stale results are discarded. - var expectedWorkerID uuid.UUID - - // Reconnect timer state. - var reconnectTimer *quartz.Timer - var reconnectCh <-chan time.Time - - // drainAndClose is set when the chat transitions away - // from running while a relay dial is still in progress. - // Instead of canceling the dial immediately, we let it - // complete so the snapshot of buffered message_parts - // can be forwarded to the subscriber. - var drainAndClose bool - - // Drain timer state. When the relay connects in - // drain-and-close mode, a short timer is started. - // During this window the normal relayPartsCh case - // forwards buffered snapshot events. When the timer - // fires the relay is torn down. - var drainTimer *quartz.Timer - var drainTimerCh <-chan time.Time - - // Helper to close relay and stop any pending reconnect - // timer. - closeRelay := func() { - // Cancel any in-flight dial goroutine first. - if dialCancel != nil { - dialCancel() - dialCancel = nil - } - // Drain all buffered relay results from canceled dials. - for { - select { - case result := <-relayReadyCh: - if result.cancel != nil { - result.cancel() - } - default: - goto drained - } - } - drained: - expectedWorkerID = uuid.Nil - if relayCancel != nil { - relayCancel() - relayCancel = nil - } - relayParts = nil - if reconnectTimer != nil { - reconnectTimer.Stop() - reconnectTimer = nil - reconnectCh = nil - } - if drainTimer != nil { - drainTimer.Stop() - drainTimer = nil - drainTimerCh = nil - } - drainAndClose = false - } - - // openRelayAsync dials the remote replica in a background - // goroutine and delivers the result on relayReadyCh so the - // main select loop is never blocked by network I/O. - openRelayAsync := func(workerID uuid.UUID) { - if cfg.dial() == nil { - return - } - // Scoped here (not in closeRelay) so repeated dials - // against the same worker keep the attempt counter and - // correctly trip the cap. - if workerID != expectedWorkerID { - retryState.reset() - } - closeRelay() - // Create a per-dial context so this goroutine is - // canceled if closeRelay() or openRelayAsync() is - // called again before the dial completes. - var dialCtx context.Context - dialCtx, dialCancel = context.WithCancel(ctx) - expectedWorkerID = workerID - go func() { - snapshot, parts, cancel, err := cfg.dial()(dialCtx, chatID, workerID, requestHeader) - if err != nil { - // Don't log context-canceled errors - // since they are expected when a dial is - // superseded by a newer one. - if dialCtx.Err() == nil { - fields := []slog.Field{ - slog.F("chat_id", chatID), - slog.F("worker_id", workerID), - slog.Error(err), - } - // Surface the peer's HTTP status (when we - // got one) as a structured field so - // operators can filter 401/403 spam - // separately from 5xx/network warnings. - var dialErr *RelayDialError - if errors.As(err, &dialErr) && dialErr.HTTPStatus != 0 { - fields = append(fields, slog.F("http_status", dialErr.HTTPStatus)) - } - logger.Warn(ctx, "failed to open relay for message parts", fields...) - } - // Hand the error to the merge loop, which will - // classify it and either back off or tear down. - select { - case relayReadyCh <- relayResult{workerID: workerID, err: err}: - case <-dialCtx.Done(): - } - return - } - // Discard stale dials so we don't start a - // wrappedParts goroutine on a canceled connection. - if dialCtx.Err() != nil { - cancel() - return - } - // Wrap the relay channel so snapshot parts - // are delivered through the same channel as - // live parts. This goroutine only forwards - // events - it does not own the relay - // lifecycle. When dialCtx is canceled it - // simply returns, closing wrappedParts via - // its defer. The cancel() is called by - // whoever canceled dialCtx (closeRelay or - // the send-fallback select below). - wrappedParts := make(chan codersdk.ChatStreamEvent, 128) - go func() { - defer close(wrappedParts) - for _, event := range snapshot { - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - select { - case wrappedParts <- event: - case <-dialCtx.Done(): - return - } - } - } - for { - select { - case event, ok := <-parts: - if !ok { - return - } - select { - case wrappedParts <- event: - case <-dialCtx.Done(): - return - } - case <-dialCtx.Done(): - return - } - } - }() - select { - case relayReadyCh <- relayResult{parts: wrappedParts, cancel: cancel, workerID: workerID}: - case <-dialCtx.Done(): - cancel() - } - }() - } - - // scheduleRelayReconnect arms a timer so the select loop - // can re-check chat status and reopen the relay. Callers - // pass the delay from retryState so the failed-dial branch - // gets backoff while transient branches stay at the floor. - scheduleRelayReconnect := func(delay time.Duration) { - if cfg.dial() == nil { - return - } - if reconnectTimer != nil { - reconnectTimer.Stop() - } - reconnectTimer = cfg.clock().NewTimer(delay, "reconnect") - reconnectCh = reconnectTimer.C - } - - // sendRelayTerminalError enqueues one error event for the - // subscriber; callers return afterwards so the deferred - // close(mergedEvents) fires and the OSS merge loop tears - // the relay leg down while pubsub/local sources keep going. - sendRelayTerminalError := func(msg string) { - select { - case mergedEvents <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeError, - ChatID: chatID, - Error: &codersdk.ChatError{Message: msg}, - }: - case <-ctx.Done(): - } - } - statusNotifications := params.StatusNotifications - go func() { - defer close(mergedEvents) - defer closeRelay() - - // Forward any initial relay snapshot parts - // collected synchronously above. - for _, event := range initialRelaySnapshot { - select { - case <-ctx.Done(): - return - case mergedEvents <- event: - } - } - - for { - relayPartsCh := relayParts - select { - case <-ctx.Done(): - return - case result := <-relayReadyCh: - // Discard stale relay results from a - // previous dial that was superseded. - if result.workerID != expectedWorkerID { - if result.cancel != nil { - result.cancel() - } - continue - } - // A nil parts channel signals the dial - // failed - classify the error to decide - // whether to schedule a backoff retry, emit a - // terminal error and tear the relay leg down - // (unrecoverable / cap reached), or simply - // drop the stale drain. - if result.parts == nil { - if drainAndClose { - // Dial failed and we were only - // waiting to drain - nothing to do. - drainAndClose = false - continue - } - var dialErr *RelayDialError - if errors.As(result.err, &dialErr) && dialErr.IsUnrecoverable() { - logger.Warn(ctx, "relay dial unrecoverable; tearing down relay leg", - slog.F("chat_id", chatID), - slog.F("worker_id", result.workerID), - slog.F("http_status", dialErr.HTTPStatus), - ) - sendRelayTerminalError(fmt.Sprintf( - "relay authentication failed (status %d)", - dialErr.HTTPStatus, - )) - return - } - delay, giveUp := retryState.next() - if giveUp { - logger.Warn(ctx, "relay dial retry cap reached; tearing down relay leg", - slog.F("chat_id", chatID), - slog.F("worker_id", result.workerID), - slog.F("max_retries", relayMaxRetries), - ) - sendRelayTerminalError(fmt.Sprintf( - "relay connection failed after %d retries", - relayMaxRetries, - )) - return - } - scheduleRelayReconnect(delay) - continue - } - // An async relay dial completed. Swap in the - // new relay channel. We deliberately do NOT - // reset the retry counter here: a peer that - // accepts the handshake and immediately drops - // the stream would otherwise keep reconnecting - // forever, since each success would zero the - // counter before the next drop re-incremented - // it. The counter only resets when the target - // worker changes (see openRelayAsync). - if relayCancel != nil { - relayCancel() - relayCancel = nil - } - relayParts = result.parts - relayCancel = result.cancel - if drainAndClose { - // The chat is no longer running on - // the remote worker, but the dial - // completed. Verify no new worker - // has claimed the chat before we - // drain stale parts. - currentChat, dbErr := params.DB.GetChatByID(ctx, chatID) - if dbErr != nil { - logger.Warn(ctx, "failed to check chat status for relay drain", - slog.F("chat_id", chatID), - slog.Error(dbErr), - ) - } - if dbErr == nil && currentChat.Status == database.ChatStatusRunning && - currentChat.WorkerID.Valid && - currentChat.WorkerID.UUID != params.WorkerID { - // A new worker picked up the chat; - // discard the stale relay and let - // openRelayAsync handle the new one. - closeRelay() - } else { - // Chat is still idle - drain the - // buffered snapshot before closing. - if drainTimer != nil { - drainTimer.Stop() - } - drainTimer = cfg.clock().NewTimer(relayDrainTimeout, "drain") - drainTimerCh = drainTimer.C - drainAndClose = false - } - } - case <-reconnectCh: - reconnectCh = nil - // Re-check whether the chat is still - // running on a remote worker before - // reconnecting. - currentChat, chatErr := params.DB.GetChatByID(ctx, chatID) - if chatErr != nil { - logger.Warn(ctx, "failed to get chat for relay reconnect", - slog.F("chat_id", chatID), - slog.Error(chatErr), - ) - // Retry on transient DB errors to - // avoid permanently stalling the - // stream. The same retry state - // bounds the DB-error loop too so a - // persistently broken DB eventually - // tears the relay down instead of - // spinning forever. - delay, giveUp := retryState.next() - if giveUp { - logger.Warn(ctx, "relay reconnect retry cap reached; tearing down relay leg", - slog.F("chat_id", chatID), - slog.F("max_retries", relayMaxRetries), - ) - sendRelayTerminalError(fmt.Sprintf( - "relay connection failed after %d retries", - relayMaxRetries, - )) - return - } - scheduleRelayReconnect(delay) - continue - } - if currentChat.Status == database.ChatStatusRunning && - currentChat.WorkerID.Valid && currentChat.WorkerID.UUID != params.WorkerID { - openRelayAsync(currentChat.WorkerID.UUID) - } - case sn, ok := <-statusNotifications: - if !ok { - statusNotifications = nil - continue - } - if sn.Status == database.ChatStatusRunning && sn.WorkerID != uuid.Nil && sn.WorkerID != params.WorkerID { - openRelayAsync(sn.WorkerID) - } else { - switch { - case dialCancel != nil && relayParts == nil: - // In-progress dial: let it complete - // so its snapshot can be forwarded. - drainAndClose = true - case relayParts != nil: - // Active relay: give it a short - // window to deliver any remaining - // buffered parts before closing. - if drainTimer != nil { - drainTimer.Stop() - } - drainTimer = cfg.clock().NewTimer(relayDrainTimeout, "drain") - drainTimerCh = drainTimer.C - default: - closeRelay() - } - } - case <-drainTimerCh: - drainTimerCh = nil - drainTimer = nil - closeRelay() - case event, ok := <-relayPartsCh: - if !ok { - if relayCancel != nil { - relayCancel() - relayCancel = nil - } - relayParts = nil - // Reuse the retry state so a relay that - // repeatedly drops eventually tears down. - delay, giveUp := retryState.next() - if giveUp { - logger.Warn(ctx, "relay drop retry cap reached; tearing down relay leg", - slog.F("chat_id", chatID), - slog.F("max_retries", relayMaxRetries), - ) - sendRelayTerminalError(fmt.Sprintf( - "relay connection failed after %d retries", - relayMaxRetries, - )) - return - } - scheduleRelayReconnect(delay) - continue - } - // Only forward message_part events from - // relay. - if event.Type == codersdk.ChatStreamEventTypeMessagePart { - select { - case <-ctx.Done(): - return - case mergedEvents <- event: - } - } - } - } - }() - - // Cleanup is driven by ctx cancellation: the merge - // goroutine owns all relay state (reconnectTimer, - // relayCancel, dialCancel, etc.) and tears it down - // via defer closeRelay() when ctx is done. - return mergedEvents - } -} - -// relayRetryState drives the retry policy for the relay reconnect -// loop. Wraps github.com/coder/retry to reuse its φ-growth defaults -// but computes the delay without blocking so the merge loop can -// schedule its own quartz.Clock timer. -// -// Not safe for concurrent use. -type relayRetryState struct { - retrier *retry.Retrier - attempts int -} - -func newRelayRetryState() *relayRetryState { - return &relayRetryState{ - retrier: retry.New(relayRetryFloor, relayRetryCeil), - } -} - -// next returns the delay before the next dial and sets giveUp once -// attempts exceed relayMaxRetries. Adapts the math from -// retry.Retrier.Wait (github.com/coder/retry/retrier.go) without -// blocking: the library's Wait returns 0 on the first call and sets -// Delay to Floor only after the sleep, so we clamp to Floor up -// front. -func (s *relayRetryState) next() (delay time.Duration, giveUp bool) { - s.attempts++ - if s.attempts > relayMaxRetries { - return 0, true - } - r := s.retrier - d := time.Duration(float64(r.Delay) * r.Rate) - if d > r.Ceil { - d = r.Ceil - } - if d < r.Floor { - d = r.Floor - } - r.Delay = d - return d, false -} - -// reset returns the state to the floor delay and zero attempts. -// Called after a successful dial or a relay target change. -func (s *relayRetryState) reset() { - s.retrier.Reset() - s.attempts = 0 -} - -// dialRelay opens a WebSocket to the replica owning chatID and -// returns any buffered message_part snapshot plus a live channel of -// subsequent events. Handshake failures return an error unwrapping -// to *RelayDialError so callers can classify via IsUnrecoverable. -// -// websocket.Dial is called directly (not via the SDK wrapper) so we -// can read *http.Response.StatusCode for classification. -func dialRelay( - ctx context.Context, - chatID uuid.UUID, - workerID uuid.UUID, - requestHeader http.Header, - cfg MultiReplicaSubscribeConfig, - clk quartz.Clock, -) ( - snapshot []codersdk.ChatStreamEvent, - parts <-chan codersdk.ChatStreamEvent, - cancel func(), - err error, -) { - address, ok := cfg.ResolveReplicaAddress(ctx, workerID) + address, ok := cfg.ResolveReplicaAddress(ctx, input.WorkerID) if !ok { - return nil, nil, nil, &RelayDialError{ - Err: xerrors.New("dial relay stream: worker replica not found"), - } + return nil, &RelayDialError{Err: xerrors.New("dial relay stream parts: worker replica not found")} } - - wsURL, err := buildRelayURL(address, chatID) + wsURL, err := buildRelayURL(address, input.ChatID) if err != nil { - return nil, nil, nil, &RelayDialError{ - Err: xerrors.Errorf("dial relay stream: %w", err), - } + return nil, &RelayDialError{Err: xerrors.Errorf("dial relay stream parts: %w", err)} } + if cfg.ReplicaIDFn == nil { + return nil, &RelayDialError{Err: xerrors.New("dial relay stream parts: replica ID function not configured")} + } replicaID := cfg.ReplicaIDFn() + if replicaID == uuid.Nil { + return nil, &RelayDialError{Err: xerrors.New("dial relay stream parts: replica ID is nil")} + } headers := make(http.Header, 2) - headers.Set(codersdk.SessionTokenHeader, extractSessionToken(requestHeader)) + headers.Set(codersdk.SessionTokenHeader, extractSessionToken(input.RequestHeader)) headers.Set(RelaySourceHeader, replicaID.String()) - relayCtx, relayCancel := context.WithCancel(ctx) - conn, resp, dialErr := websocket.Dial(relayCtx, wsURL, &websocket.DialOptions{ + conn, resp, dialErr := websocket.Dial(ctx, wsURL, &websocket.DialOptions{ HTTPClient: cfg.ReplicaHTTPClient, HTTPHeader: headers, CompressionMode: websocket.CompressionDisabled, @@ -722,118 +93,22 @@ func dialRelay( status := 0 if resp != nil { status = resp.StatusCode - // The websocket library closes resp.Body on success; on - // failure we close it ourselves so we don't leak the TCP - // connection. if dialErr != nil && resp.Body != nil { _ = resp.Body.Close() } } if dialErr != nil { - relayCancel() - return nil, nil, nil, &RelayDialError{ + return nil, &RelayDialError{ HTTPStatus: status, - Err: xerrors.Errorf("dial relay stream: %w", dialErr), + Err: xerrors.Errorf("dial relay stream parts: %w", dialErr), } } - // Match the server's 4 MiB read limit in codersdk.StreamChat so - // large message_part batches don't trip the default 32 KiB cap. conn.SetReadLimit(1 << 22) - - snapshot = make([]codersdk.ChatStreamEvent, 0, 100) - - // sourceEvents is the flattened batch→event channel. A small - // goroutine reads batches off the websocket and fans them out; - // callers see a single event stream identical to the shape the - // old SDK call produced. - sourceEvents := make(chan codersdk.ChatStreamEvent, 128) - go func() { - defer close(sourceEvents) - for { - var batch []codersdk.ChatStreamEvent - if readErr := wsjson.Read(relayCtx, conn, &batch); readErr != nil { - return - } - for _, event := range batch { - select { - case sourceEvents <- event: - case <-relayCtx.Done(): - return - } - } - } - }() - - closeSource := func() { - relayCancel() - _ = conn.Close(websocket.StatusNormalClosure, "") - } - - // Wait briefly for the first event to handle the common - // case where the remote side has buffered parts but hasn't - // flushed them to the WebSocket yet. - const drainTimeout = time.Second - drainTimer := clk.NewTimer(drainTimeout, "drain") - defer drainTimer.Stop() - -drainInitial: - for len(snapshot) < cap(snapshot) { - select { - case <-relayCtx.Done(): - closeSource() - return nil, nil, nil, &RelayDialError{ - Err: xerrors.Errorf("dial relay stream: %w", relayCtx.Err()), - } - case event, ok := <-sourceEvents: - if !ok { - break drainInitial - } - if event.Type != codersdk.ChatStreamEventTypeMessagePart { - continue - } - snapshot = append(snapshot, event) - // After getting the first event, switch to - // non-blocking drain for remaining buffered events. - drainTimer.Stop() - drainTimer.Reset(0) - case <-drainTimer.C: - break drainInitial - } - } - - events := make(chan codersdk.ChatStreamEvent, 128) - - go func() { - defer close(events) - defer closeSource() - - // No need to re-send snapshot events - they're - // returned to the caller directly. - for { - select { - case <-relayCtx.Done(): - return - case event, ok := <-sourceEvents: - if !ok { - return - } - if event.Type != codersdk.ChatStreamEventTypeMessagePart { - continue - } - select { - case events <- event: - case <-relayCtx.Done(): - return - } - } - } - }() - - return snapshot, events, closeSource, nil + return osschatd.NewStreamPartsJSONSession(ctx, conn), nil } -// buildRelayURL builds the websocket URL for the chat stream -// endpoint on a peer replica. It maps http(s) schemes to ws(s). +// buildRelayURL builds the websocket URL for the chat stream parts endpoint on +// a peer replica. It maps http(s) schemes to ws(s). func buildRelayURL(address string, chatID uuid.UUID) (string, error) { u, err := url.Parse(address) if err != nil { @@ -845,40 +120,30 @@ func buildRelayURL(address string, chatID uuid.UUID) (string, error) { case "https": u.Scheme = "wss" case "ws", "wss": - // already a websocket URL, leave as-is. default: return "", xerrors.Errorf("unsupported relay address scheme %q", u.Scheme) } - u.Path = fmt.Sprintf("/api/experimental/chats/%s/stream", chatID) - q := u.Query() - // Relays only need live message_part events, not the full - // history; pass the relay sentinel so the peer skips its - // durable DB snapshot and delivers in-flight parts only. - q.Set("after_id", strconv.FormatInt(osschatd.RelaySentinelAfterID, 10)) - u.RawQuery = q.Encode() + u.Path = "/api/experimental/chats/" + chatID.String() + "/stream/parts" + u.RawQuery = "" return u.String(), nil } -// extractSessionToken returns the session token carried by the -// given request headers. It mirrors the priority order used by -// apiKeyMiddleware: cookie, then Coder-Session-Token header, then -// Authorization: Bearer header. +// extractSessionToken returns the session token carried by the given request +// headers. It mirrors the priority order used by apiKeyMiddleware: cookie, +// then Coder-Session-Token header, then Authorization: Bearer header. func extractSessionToken(header http.Header) string { if header == nil { return "" } - // Cookie (browser WebSocket upgrade - most common relay case). if raw := header.Get(cookieHeader); raw != "" { r := &http.Request{Header: http.Header{cookieHeader: {raw}}} if c, err := r.Cookie(codersdk.SessionTokenCookie); err == nil && c.Value != "" { return c.Value } } - // Coder-Session-Token header (SDK / CLI callers). if v := header.Get(codersdk.SessionTokenHeader); v != "" { return v } - // Authorization: Bearer . if v := header.Get(authorizationHeader); len(v) > 7 && strings.EqualFold(v[:7], "bearer ") { return strings.TrimSpace(v[7:]) } diff --git a/enterprise/coderd/x/chatd/chatd_retry_test.go b/enterprise/coderd/x/chatd/chatd_retry_test.go deleted file mode 100644 index d21a15b9ba..0000000000 --- a/enterprise/coderd/x/chatd/chatd_retry_test.go +++ /dev/null @@ -1,796 +0,0 @@ -package chatd_test - -import ( - "context" - "database/sql" - "encoding/json" - "io" - "math" - "net/http" - "net/http/httptest" - "regexp" - "sync/atomic" - "testing" - "time" - - "github.com/google/uuid" - "github.com/stretchr/testify/require" - "golang.org/x/xerrors" - - "cdr.dev/slog/v3/sloggers/slogtest" - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/dbtestutil" - dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" - coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" - osschatd "github.com/coder/coder/v2/coderd/x/chatd" - "github.com/coder/coder/v2/codersdk" - entchatd "github.com/coder/coder/v2/enterprise/coderd/x/chatd" - "github.com/coder/coder/v2/testutil" - "github.com/coder/quartz" -) - -// mulPhi multiplies a duration by math.Phi to compute the next -// step in retry.Retrier's φ-growth backoff sequence. If -// TestRelayReconnectUsesExponentialBackoff starts failing after a -// retry library bump, check whether the growth factor has changed. -func mulPhi(d time.Duration) time.Duration { - return time.Duration(float64(d) * math.Phi) -} - -// setChatRunningAndPublish marks the chat row as running on workerID -// and publishes a matching status notification. It keeps the DB row -// and pubsub notification in sync so the async reconnect loop -// re-dials on each timer fire (the reconnect branch re-checks DB -// status before calling openRelayAsync). -func setChatRunningAndPublish( - ctx context.Context, - t *testing.T, - db database.Store, - ps dbpubsub.Pubsub, - chatID, workerID uuid.UUID, -) { - t.Helper() - now := time.Now() - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chatID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, - StartedAt: sql.NullTime{Time: now, Valid: true}, - HeartbeatAt: sql.NullTime{Time: now, Valid: true}, - }) - require.NoError(t, err) - payload, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: workerID.String(), - }) - require.NoError(t, err) - require.NoError(t, ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chatID), payload)) -} - -// TestRelayDialErrorIsUnrecoverable locks the classification policy. -// Adding a new HTTP status to the unrecoverable set should force a -// test edit too. -func TestRelayDialErrorIsUnrecoverable(t *testing.T) { - t.Parallel() - - cases := []struct { - name string - status int - want bool - }{ - {"unauthorized", http.StatusUnauthorized, true}, - {"forbidden", http.StatusForbidden, true}, - {"internal_server", http.StatusInternalServerError, false}, - {"bad_gateway", http.StatusBadGateway, false}, - {"service_unavailable", http.StatusServiceUnavailable, false}, - {"too_many_requests", http.StatusTooManyRequests, false}, - {"pre_response", 0, false}, - {"bad_request", http.StatusBadRequest, false}, - {"not_found", http.StatusNotFound, false}, - } - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - e := &entchatd.RelayDialError{HTTPStatus: tc.status, Err: io.EOF} - require.Equal(t, tc.want, e.IsUnrecoverable(), - "status=%d", tc.status) - }) - } -} - -// TestRelayReconnectUsesExponentialBackoff asserts that the reconnect -// timer follows the φ-growth sequence produced by -// github.com/coder/retry's defaults, floored at relayRetryFloor. -func TestRelayReconnectUsesExponentialBackoff(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var failCount atomic.Int32 - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - failCount.Add(1) - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: http.StatusBadGateway, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-backoff") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Kick the async relay loop and keep the DB row in sync so - // each reconnect timer fire triggers another dial. - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) - // Expected sequence from retry.Retrier math: - // attempt 1 → floor (500ms) - // attempt n → prev × φ (capped at ceil) - floor := 500 * time.Millisecond - expected := []time.Duration{ - floor, - mulPhi(floor), - mulPhi(mulPhi(floor)), - mulPhi(mulPhi(mulPhi(floor))), - mulPhi(mulPhi(mulPhi(mulPhi(floor)))), - } - - for i, want := range expected { - call := trapReconnect.MustWait(ctx) - require.Equal(t, want, call.Duration, - "attempt %d: want %v got %v", i+1, want, call.Duration) - call.MustRelease(ctx) - mclk.Advance(want).MustWait(ctx) - } - - // We expect 1 initial attempt + 5 reconnects fired by the - // trapped timer = 6 dials before the cap-check runs. Use - // Eventually so we don't race the final dial goroutine that - // the last Advance kicked off. - require.Eventually(t, func() bool { - return failCount.Load() >= 6 - }, testutil.WaitShort, testutil.IntervalFast, - "expected 6 dials, got %d", failCount.Load()) - - // The events channel must remain open - we're still under the - // cap. - select { - case ev, open := <-events: - if !open { - t.Fatalf("events channel closed prematurely; retries should continue below cap") - } - // Allow through events that might have been queued; just - // confirm it's not a terminal error. - if ev.Type == codersdk.ChatStreamEventTypeError { - t.Fatalf("unexpected terminal error: %v", ev.Error) - } - default: - } -} - -// TestRelayReconnectResetsOnSuccess exercises the path where a -// successful dial resets the retry state so the next failure starts -// over at the floor delay. -// TestRelayRepeatedDropsHitCap verifies the cap covers a peer that -// accepts the handshake and immediately drops it. Without a proper -// cap, such a peer would produce one reconnect per floor delay -// forever. The retry counter must accumulate across dial-success / -// parts-close cycles so the cap trips. -func TestRelayRepeatedDropsHitCap(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - opened := make(chan chan codersdk.ChatStreamEvent, 32) - var call atomic.Int32 - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call.Add(1) - ch := make(chan codersdk.ChatStreamEvent, 1) - opened <- ch - return nil, ch, func() {}, nil - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-drops") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Kick off the first async dial. - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) - - // Close the first dial's parts channel so the merge loop - // schedules a reconnect. Then advance 6 reconnect timers, - // closing the parts channel each time so the cycle is: - // dial -> success -> parts-close -> next() -> reconnect. - // 1 initial dial + 6 timer-driven dials = 7 total; the 7th - // parts-close trips the cap. - for i := 0; i < 7; i++ { - var ch chan codersdk.ChatStreamEvent - select { - case ch = <-opened: - case <-ctx.Done(): - t.Fatalf("timed out waiting for dial %d", i+1) - } - // Closing the parts channel triggers the relayPartsCh - // close branch, which calls retryState.next() and - // schedules the next reconnect. - close(ch) - if i == 6 { - // 7th parts-close should trip the cap; no more - // reconnect timers. - break - } - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - mclk.Advance(call.Duration).MustWait(ctx) - } - - // A terminal error event must arrive on the events channel. - var errEvent *codersdk.ChatStreamEvent - require.Eventually(t, func() bool { - select { - case ev, open := <-events: - if !open { - return errEvent != nil - } - if ev.Type == codersdk.ChatStreamEventTypeError { - errEvent = &ev - return true - } - return false - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast, - "expected a terminal error event after repeated drops hit cap") - require.NotNil(t, errEvent.Error) - require.Contains(t, errEvent.Error.Message, "relay connection failed") - - // We should have observed exactly 7 dials before tear-down. - require.Equal(t, int32(7), call.Load(), - "expected 7 dials (1 initial + 6 reconnect retries) before cap") -} - -// TestRelayStopsAfterIntermittentCap verifies the cap-reached -// tear-down path: after N intermittent failures the merge loop emits -// one error event, closes the events channel, and stops dialing. -func TestRelayStopsAfterIntermittentCap(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - callCount.Add(1) - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: http.StatusBadGateway, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cap") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) - // Advance through N consecutive reconnect timers. Each one - // triggers a dial, which fails and schedules the next timer. - // After the Nth failure the retry state says giveUp=true on - // the next .next() call, so the merge loop tears down. - for i := 0; i < 6; i++ { - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - mclk.Advance(call.Duration).MustWait(ctx) - } - - // Wait for the terminal error event to arrive. mergedEvents - // closes inside the enterprise merge goroutine, but OSS only - // nil-outs relayEvents on close - the outer events channel - // stays open for pubsub/local, so we wait for the error event - // itself rather than channel closure. - var errEvent *codersdk.ChatStreamEvent - require.Eventually(t, func() bool { - select { - case ev, open := <-events: - if !open { - return errEvent != nil - } - if ev.Type == codersdk.ChatStreamEventTypeError { - errEvent = &ev - return true - } - return false - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast, - "expected a terminal error event") - require.NotNil(t, errEvent, "expected a terminal error event") - require.NotNil(t, errEvent.Error) - require.Contains(t, errEvent.Error.Message, "relay connection failed") - require.Contains(t, errEvent.Error.Message, "6") - - // Ensure the cap fires at attempt N+1 - the retry state allows - // relayMaxRetries successful next() calls before flipping - // giveUp. With one initial dial + 6 reconnect-timer fires the - // 7th .next() trips the cap and tears down, so we see 7 dials - // total and nothing further. - totalDials := callCount.Load() - require.Equal(t, int32(7), totalDials, - "expected exactly relayMaxRetries+1 dials before cap; got %d", totalDials) -} - -// chatByIDErrorStore wraps a database.Store and forces GetChatByID -// to return a caller-supplied error once after N successful calls. -// This lets the initial Subscribe call succeed (OSS's initial state -// load needs a real Chat to wire up the relay) while subsequent -// reconnect-branch calls exercise the DB-error retry path. -type chatByIDErrorStore struct { - database.Store - err error - okRemain atomic.Int32 // number of calls allowed to delegate before erroring. -} - -func (s *chatByIDErrorStore) GetChatByID(ctx context.Context, id uuid.UUID) (database.Chat, error) { - if s.okRemain.Add(-1) >= 0 { - return s.Store.GetChatByID(ctx, id) - } - return database.Chat{}, s.err -} - -// TestRelayReconnectStopsAfterDBErrorCap verifies the reconnect-timer -// branch's DB-error path shares the same retry budget as dial -// failures and trips the cap after enough consecutive DB errors. -func TestRelayReconnectStopsAfterDBErrorCap(t *testing.T) { - t.Parallel() - - realDB, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - callCount.Add(1) - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: http.StatusBadGateway, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - // The server sees a DB whose GetChatByID always errors after - // the initial Subscribe snapshot load. Other methods delegate - // to the real DB, so seeding below still works. - failingDB := &chatByIDErrorStore{ - Store: realDB, - err: xerrors.New("mock: GetChatByID always fails"), - } - // Allow one successful GetChatByID (the Subscribe preamble's - // initial state load). All subsequent calls return the mock - // error, exercising the reconnect-branch DB-error path. - failingDB.okRemain.Store(1) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, realDB) - chat := seedWaitingChat(t, realDB, org.ID, user, model, "relay-db-error") - - subscriber := newTestServer(t, failingDB, ps, subscriberID, dialer, mclk) - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Flip to running so the merge loop starts an async dial. The - // dial fails (attempts=1, reconnect scheduled). From there each - // reconnect timer fires, the merge loop calls GetChatByID, the - // failing DB returns an error, and retryState.next() increments. - // - // Budget: 1 dial-failure + 6 DB-failures = 7 next() calls; the - // 7th trips the cap. - setChatRunningAndPublish(ctx, t, realDB, ps, chat.ID, workerID) - for i := 0; i < 6; i++ { - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - mclk.Advance(call.Duration).MustWait(ctx) - } - - var errEvent *codersdk.ChatStreamEvent - require.Eventually(t, func() bool { - select { - case ev, open := <-events: - if !open { - return errEvent != nil - } - if ev.Type == codersdk.ChatStreamEventTypeError { - errEvent = &ev - return true - } - return false - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast, - "expected terminal error event after DB-error cap") - require.NotNil(t, errEvent.Error) - require.Contains(t, errEvent.Error.Message, "relay connection failed") - require.Contains(t, errEvent.Error.Message, "6") - - // Exactly 1 dial fired: the one that triggered the initial - // reconnect schedule. All subsequent next() calls come from the - // DB-error branch without calling the dialer. - require.Equal(t, int32(1), callCount.Load(), - "expected exactly 1 dial; reconnects should short-circuit on DB error") -} - -// TestRelayStopsImmediatelyOnUnauthorized tests the unrecoverable -// branch and its table of status codes. -func TestRelayStopsImmediatelyOnUnauthorized(t *testing.T) { - t.Parallel() - - cases := []struct { - name string - status int - wantUnrecoverable bool - wantMsgContains string - }{ - {"401", http.StatusUnauthorized, true, "401"}, - {"403", http.StatusForbidden, true, "403"}, - {"500_intermittent", http.StatusInternalServerError, false, ""}, - {"zero_intermittent", 0, false, ""}, - } - - for _, tc := range cases { - t.Run(tc.name, func(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - callCount.Add(1) - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: tc.status, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, - "relay-unrec-"+tc.name) - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) - if tc.wantUnrecoverable { - // First dial should tear the relay down. - var errEvent *codersdk.ChatStreamEvent - require.Eventually(t, func() bool { - select { - case ev, open := <-events: - if !open { - return errEvent != nil - } - if ev.Type == codersdk.ChatStreamEventTypeError { - errEvent = &ev - return true - } - return false - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast, - "expected terminal error event") - require.NotNil(t, errEvent) - require.Contains(t, errEvent.Error.Message, "relay authentication failed") - require.Contains(t, errEvent.Error.Message, tc.wantMsgContains) - require.Equal(t, int32(1), callCount.Load(), - "unrecoverable errors must not retry; got %d dials", callCount.Load()) - } else { - // Intermittent: fire one reconnect timer - // and confirm the dialer is called again. - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - mclk.Advance(call.Duration).MustWait(ctx) - require.Eventually(t, func() bool { - return callCount.Load() >= 2 - }, testutil.WaitShort, testutil.IntervalFast, - "intermittent should retry at least once") - } - }) - } -} - -// TestRelayBackoffResetsOnStatusChange checks that closeRelay (driven -// by a status notification) resets the retry counter so subsequent -// dials against a new target start at the floor delay. -func TestRelayBackoffResetsOnStatusChange(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID1 := uuid.New() - workerID2 := uuid.New() - subscriberID := uuid.New() - - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: http.StatusBadGateway, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-reset-on-status") - - _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Drive the async openRelayAsync path with workerID1. - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID1) - - // Drive 3 intermittent failures so attempts=3 and the delay - // has grown past the floor. After each loop iteration the 4th - // reconnect timer is queued - consume it too so our later - // assertion sees the reset's timer, not a stale one. - for i := 0; i < 3; i++ { - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - mclk.Advance(call.Duration).MustWait(ctx) - } - // Grab the next trapped timer (the grown one scheduled after - // the 3rd dial fails) but don't advance it - we want to see it - // replaced by a fresh floor-delay timer after the reset. - grown := trapReconnect.MustWait(ctx) - require.Greater(t, grown.Duration, 500*time.Millisecond, - "sanity: pre-reset delay should have grown past the floor") - grown.MustRelease(ctx) - - // Flip the chat to waiting; closeRelay runs (because the - // status notification no longer points at a running peer) and - // should reset the retry state. - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - }) - require.NoError(t, err) - waitingPayload, err := json.Marshal(coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusWaiting), - }) - require.NoError(t, err) - require.NoError(t, ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), waitingPayload)) - - // Flip back to running on a different worker. This triggers a - // fresh openRelayAsync which fails, arming a reconnect timer. - // That timer's delay must be the floor, proving the reset. - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID2) - - call := trapReconnect.MustWait(ctx) - require.Equal(t, 500*time.Millisecond, call.Duration, - "retry state must reset after status change; got grown delay %v", call.Duration) - call.MustRelease(ctx) -} - -// TestRelayBackoffRespectsContextCancel is a regression guard: the -// reconnect timer must respect ctx cancellation promptly. -func TestRelayBackoffRespectsContextCancel(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - dialer := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - return nil, nil, nil, &entchatd.RelayDialError{ - HTTPStatus: http.StatusBadGateway, - Err: io.EOF, - } - } - - mclk := quartz.NewMock(t) - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, dialer, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-cancel") - - subCtx, subCancel := context.WithCancel(ctx) - _, events, cancel, ok := subscriber.Subscribe(subCtx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - setChatRunningAndPublish(ctx, t, db, ps, chat.ID, workerID) - - // Wait for the first reconnect timer to arm. - call := trapReconnect.MustWait(ctx) - call.MustRelease(ctx) - - // Cancel the subscriber context. The events channel should - // close promptly (the merge goroutine's select exits on - // ctx.Done). - subCancel() - - done := make(chan struct{}) - go func() { - defer close(done) - for { - if _, open := <-events; !open { - return - } - } - }() - select { - case <-done: - case <-time.After(testutil.WaitShort): - t.Fatal("events channel did not close after ctx cancel") - } -} - -// TestDialRelayReal401 exercises the real dialRelay path against an -// httptest server that returns 401 on the stream endpoint. It -// validates that the websocket library's handshake failure -// propagates through as *RelayDialError with HTTPStatus == 401. -// -// This is the one test that uses the real coder/websocket library -// on the failure path - a safety net against library upgrades -// silently breaking status capture. -func TestDialRelayReal401(t *testing.T) { - t.Parallel() - - // An httptest server that 401s every request on the stream - // endpoint. Any other path gets a 404. - srv := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { - if !streamPathRE.MatchString(r.URL.Path) { - http.NotFound(rw, r) - return - } - rw.Header().Set("Content-Type", "application/json") - rw.WriteHeader(http.StatusUnauthorized) - _, _ = rw.Write([]byte(`{"message":"unauthorized"}`)) - })) - t.Cleanup(srv.Close) - - db, _ := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - // Wire real config (no DialerFn override) so dialRelay runs - // end-to-end against the httptest server. Seeding a waiting - // chat (below) keeps Subscribe's initial synchronous dial a - // no-op; we then push a running status notification to the - // merge loop so it invokes dialRelay via the async path, where - // the 401 tear-down logic lives. - cfg := entchatd.MultiReplicaSubscribeConfig{ - ResolveReplicaAddress: func(_ context.Context, _ uuid.UUID) (string, bool) { - return srv.URL, true - }, - ReplicaHTTPClient: srv.Client(), - ReplicaIDFn: func() uuid.UUID { return subscriberID }, - } - subscribeFn := entchatd.NewMultiReplicaSubscribeFn(cfg) - - ctx := testutil.Context(t, testutil.WaitMedium) - user, org, model := seedChatDependencies(t, db) - // Seed a waiting chat - no sync dial - then push a running - // status notification to trigger the async dial via the real - // dialRelay path. - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-real-401") - - statusCh := make(chan osschatd.StatusNotification, 1) - evs := subscribeFn(ctx, osschatd.SubscribeFnParams{ - ChatID: chat.ID, - Chat: chat, - WorkerID: subscriberID, - StatusNotifications: statusCh, - RequestHeader: http.Header{codersdk.SessionTokenHeader: {"test-token"}}, - DB: db, - Logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), - }) - - statusCh <- osschatd.StatusNotification{ - Status: database.ChatStatusRunning, - WorkerID: workerID, - } - - // Wait for a terminal error event. On a real 401 handshake, - // the classifier flags it unrecoverable → one dial, then - // error event, then channel close. - var errEvent *codersdk.ChatStreamEvent - deadline := time.After(testutil.WaitMedium) -waitErr: - for { - select { - case ev, open := <-evs: - if !open { - break waitErr - } - if ev.Type == codersdk.ChatStreamEventTypeError { - errEvent = &ev - } - case <-deadline: - break waitErr - } - } - - require.NotNil(t, errEvent, "expected terminal error event from real 401 dial") - require.NotNil(t, errEvent.Error) - require.Contains(t, errEvent.Error.Message, "relay authentication failed") - require.Contains(t, errEvent.Error.Message, "401") -} - -// streamPathRE matches the chat stream endpoint path built by -// buildRelayURL. Compiled at package scope so the httptest handler -// below doesn't pay regexp.Compile per request. -var streamPathRE = regexp.MustCompile( - `^/api/experimental/chats/[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}/stream$`, -) diff --git a/enterprise/coderd/x/chatd/chatd_test.go b/enterprise/coderd/x/chatd/chatd_test.go index ec43e10928..9dfb2e361f 100644 --- a/enterprise/coderd/x/chatd/chatd_test.go +++ b/enterprise/coderd/x/chatd/chatd_test.go @@ -2,1623 +2,151 @@ package chatd_test import ( "context" - "database/sql" - "encoding/json" - "fmt" - "math" "net/http" "net/http/httptest" - "sync/atomic" "testing" - "time" "github.com/google/uuid" - "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" - "golang.org/x/xerrors" - "cdr.dev/slog/v3/sloggers/slogtest" - "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/dbgen" - "github.com/coder/coder/v2/coderd/database/dbtestutil" - dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" - coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" osschatd "github.com/coder/coder/v2/coderd/x/chatd" - "github.com/coder/coder/v2/coderd/x/chatd/chattest" "github.com/coder/coder/v2/codersdk" entchatd "github.com/coder/coder/v2/enterprise/coderd/x/chatd" - "github.com/coder/coder/v2/testutil" - "github.com/coder/quartz" + "github.com/coder/websocket" ) -func chatLastErrorMessage(raw pqtype.NullRawMessage) string { - if !raw.Valid { - return "" +type fakePartsSession struct { + parts chan osschatd.StreamPart +} + +func newFakePartsSession() *fakePartsSession { + return &fakePartsSession{parts: make(chan osschatd.StreamPart)} +} + +func (*fakePartsSession) SelectEpisode(context.Context, int64, int64) error { return nil } +func (s *fakePartsSession) Parts() <-chan osschatd.StreamPart { return s.parts } +func (s *fakePartsSession) Close() error { + close(s.parts) + return nil +} + +func TestRelayDialErrorIsUnrecoverable(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + status int + want bool + }{ + {"unauthorized", http.StatusUnauthorized, true}, + {"forbidden", http.StatusForbidden, true}, + {"internal_server", http.StatusInternalServerError, false}, + {"bad_gateway", http.StatusBadGateway, false}, + {"pre_response", 0, false}, } - - var payload codersdk.ChatError - if err := json.Unmarshal(raw.RawMessage, &payload); err == nil && payload.Message != "" { - return payload.Message - } - return string(raw.RawMessage) -} - -func newTestServer( - t *testing.T, - db database.Store, - ps dbpubsub.Pubsub, - replicaID uuid.UUID, - dialer func( - ctx context.Context, - chatID uuid.UUID, - workerID uuid.UUID, - requestHeader http.Header, - ) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, - ), - clock quartz.Clock, -) *osschatd.Server { - t.Helper() - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := osschatd.New(osschatd.Config{ - Logger: logger, - Database: db, - ReplicaID: replicaID, - Pubsub: ps, - SubscribeFn: entchatd.NewMultiReplicaSubscribeFn(entchatd.MultiReplicaSubscribeConfig{DialerFn: dialer, Clock: clock}), - PendingChatAcquireInterval: testutil.WaitSuperLong, - }) - server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) - }) - return server -} - -func newActiveWorkerServer( - t *testing.T, - db database.Store, - ps dbpubsub.Pubsub, - replicaID uuid.UUID, -) *osschatd.Server { - t.Helper() - logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - server := osschatd.New(osschatd.Config{ - Logger: logger, - Database: db, - ReplicaID: replicaID, - Pubsub: ps, - PendingChatAcquireInterval: 10 * time.Millisecond, - InFlightChatStaleAfter: testutil.WaitSuperLong, - }) - server.Start() - t.Cleanup(func() { - require.NoError(t, server.Close()) - }) - return server -} - -// seedChatDependencies creates a user, organization, and chat model -// config in the database for use in relay tests. -func seedChatDependencies( - t *testing.T, - db database.Store, -) (database.User, database.Organization, database.ChatModelConfig) { - t.Helper() - - safetyNet := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { - rw.Header().Set("Content-Type", "application/json") - rw.WriteHeader(http.StatusInternalServerError) - _, _ = rw.Write([]byte(`{"error":{"message":"unexpected OpenAI request in chatd relay test safety net"}}`)) - })) - t.Cleanup(safetyNet.Close) - - user := dbgen.User(t, db, database.User{}) - org := dbgen.Organization(t, db, database.Organization{}) - dbgen.OrganizationMember(t, db, database.OrganizationMember{ - UserID: user.ID, - OrganizationID: org.ID, - }) - provider := dbgen.AIProvider(t, db, database.AIProvider{ - Type: database.AiProviderTypeOpenai, - Name: "test-" + uuid.NewString(), - BaseUrl: safetyNet.URL, - }) - dbgen.AIProviderKey(t, db, database.AIProviderKey{ - ProviderID: provider.ID, - }) - model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Provider: "openai", - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - IsDefault: true, - }) - return user, org, model -} - -func seedWaitingChat( - t *testing.T, - db database.Store, - orgID uuid.UUID, - user database.User, - model database.ChatModelConfig, - title string, -) database.Chat { - t.Helper() - - chat := dbgen.Chat(t, db, database.Chat{ - OrganizationID: orgID, - OwnerID: user.ID, - LastModelConfigID: model.ID, - Title: title, - }) - return chat -} - -func seedRemoteRunningChat( - ctx context.Context, - t *testing.T, - db database.Store, - orgID uuid.UUID, - user database.User, - model database.ChatModelConfig, - workerID uuid.UUID, - title string, -) database.Chat { - t.Helper() - - chat := seedWaitingChat(t, db, orgID, user, model, title) - now := time.Now() - chat, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, - StartedAt: sql.NullTime{Time: now, Valid: true}, - HeartbeatAt: sql.NullTime{Time: now, Valid: true}, - }) - require.NoError(t, err) - return chat -} - -func setOpenAIProviderBaseURL( - ctx context.Context, - t *testing.T, - db database.Store, - baseURL string, -) { - t.Helper() - - providers, err := db.GetAIProviders(ctx, database.GetAIProvidersParams{IncludeDisabled: true}) - require.NoError(t, err) - for _, provider := range providers { - if provider.Type != database.AiProviderTypeOpenai { - continue - } - _, err = db.UpdateAIProvider(ctx, database.UpdateAIProviderParams{ - ID: provider.ID, - Type: provider.Type, - DisplayName: provider.DisplayName, - Enabled: provider.Enabled, - BaseUrl: baseURL, - Settings: provider.Settings, - SettingsKeyID: provider.SettingsKeyID, + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + err := &entchatd.RelayDialError{HTTPStatus: tc.status, Err: context.Canceled} + require.Equal(t, tc.want, err.IsUnrecoverable()) }) + } +} + +func TestStreamPartsDialerUsesConfiguredDialer(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + workerID := uuid.New() + headers := http.Header{codersdk.SessionTokenHeader: {"token-value"}} + wantSession := newFakePartsSession() + + var gotInput osschatd.StreamPartsDialInput + dialer := entchatd.NewStreamPartsDialer(entchatd.StreamPartsDialerConfig{ + DialerFn: func(_ context.Context, input osschatd.StreamPartsDialInput) (osschatd.StreamPartsSession, error) { + gotInput = input + return wantSession, nil + }, + }) + + session, err := dialer(context.Background(), osschatd.StreamPartsDialInput{ + ChatID: chatID, + WorkerID: workerID, + RequestHeader: headers, + }) + require.NoError(t, err) + require.Same(t, wantSession, session) + require.Equal(t, chatID, gotInput.ChatID) + require.Equal(t, workerID, gotInput.WorkerID) + require.Equal(t, "token-value", gotInput.RequestHeader.Get(codersdk.SessionTokenHeader)) +} + +func TestStreamPartsDialerDialsPartsEndpoint(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + workerID := uuid.New() + replicaID := uuid.New() + received := make(chan http.Header, 1) + + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + require.Equal(t, "/api/experimental/chats/"+chatID.String()+"/stream/parts", r.URL.Path) + require.Empty(t, r.URL.RawQuery) + received <- r.Header.Clone() + conn, err := websocket.Accept(rw, r, nil) require.NoError(t, err) - return - } - require.Fail(t, "openai provider not found") + _ = conn.Close(websocket.StatusNormalClosure, "") + })) + t.Cleanup(server.Close) + + dialer := entchatd.NewStreamPartsDialer(entchatd.StreamPartsDialerConfig{ + ResolveReplicaAddress: func(_ context.Context, gotWorker uuid.UUID) (string, bool) { + require.Equal(t, workerID, gotWorker) + return server.URL, true + }, + ReplicaHTTPClient: server.Client(), + ReplicaIDFn: func() uuid.UUID { return replicaID }, + }) + + session, err := dialer(context.Background(), osschatd.StreamPartsDialInput{ + ChatID: chatID, + WorkerID: workerID, + RequestHeader: http.Header{ + codersdk.SessionTokenHeader: {"session-token"}, + }, + }) + require.NoError(t, err) + require.NotNil(t, session) + require.NoError(t, session.Close()) + + headers := <-received + require.Equal(t, "session-token", headers.Get(codersdk.SessionTokenHeader)) + require.Equal(t, replicaID.String(), headers.Get(entchatd.RelaySourceHeader)) } -func TestSubscribeRelayReconnectsOnDrop(t *testing.T) { +func TestStreamPartsDialerClassifiesHTTPFailures(t *testing.T) { t.Parallel() - db, ps := dbtestutil.NewDB(t) + chatID := uuid.New() workerID := uuid.New() - subscriberID := uuid.New() + server := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) { + http.Error(rw, "nope", http.StatusUnauthorized) + })) + t.Cleanup(server.Close) - var callCount atomic.Int32 + dialer := entchatd.NewStreamPartsDialer(entchatd.StreamPartsDialerConfig{ + ResolveReplicaAddress: func(context.Context, uuid.UUID) (string, bool) { return server.URL, true }, + ReplicaHTTPClient: server.Client(), + ReplicaIDFn: uuid.New, + }) - provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - ch := make(chan codersdk.ChatStreamEvent, 10) - if call == 1 { - // First relay: send a part then close to simulate a drop. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("first-relay"), - }, - } - close(ch) - } else { - // Second relay: send a different part, keep open. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("second-relay"), - }, - } - // Don't close — keep alive so the subscriber stays connected. - } - return nil, ch, func() {}, nil - } - - mclk := quartz.NewMock(t) - // Trap the reconnect timer so we can fire it deterministically - // instead of waiting real time. - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-reconnect") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Should get the first relay part. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "first-relay" { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // Wait for the reconnect timer to be created after the relay - // drop, then advance the mock clock to fire it immediately. - trapReconnect.MustWait(ctx).MustRelease(ctx) - mclk.Advance(500 * time.Millisecond).MustWait(ctx) - - // After the first relay closes, the reconnection should deliver - // the second relay part. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "second-relay" { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - require.GreaterOrEqual(t, int(callCount.Load()), 2) -} - -func TestSubscribeRelayAsyncDoesNotBlock(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - dialStarted := make(chan struct{}) - dialContinue := make(chan struct{}) - - provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - // Signal that the dial has started, then block until released. - select { - case <-dialStarted: - default: - close(dialStarted) - } - select { - case <-dialContinue: - case <-ctx.Done(): - return nil, nil, nil, ctx.Err() - } - ch := make(chan codersdk.ChatStreamEvent, 10) - return nil, ch, func() {}, nil - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Seed a waiting chat so Subscribe does not trigger a synchronous - // relay. - chat := seedWaitingChat(t, db, org.ID, user, model, "relay-async-nonblock") - - // Subscribe before the chat is marked running so the relay opens - // via pubsub notification (openRelayAsync path). - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Now mark the chat as running on a remote worker. This publishes - // a status notification which triggers openRelayAsync on the - // subscriber. - notify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: workerID.String(), - } - payload, err := json.Marshal(notify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload) - require.NoError(t, err) - - // Wait for the relay dial to actually start (blocking in the - // provider). - select { - case <-dialStarted: - case <-ctx.Done(): - t.Fatal("timed out waiting for relay dial to start") - } - - // While the relay is still dialing (provider is blocked), publish - // another status change. If openRelayAsync blocked the select loop - // this event would never arrive. - statusNotify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusWaiting), - } - statusPayload, err := json.Marshal(statusNotify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), statusPayload) - require.NoError(t, err) - - // The waiting status event should arrive promptly despite the - // relay still dialing. - require.Eventually(t, func() bool { - select { - case event := <-events: - return event.Type == codersdk.ChatStreamEventTypeStatus && - event.Status != nil && - event.Status.Status == codersdk.ChatStatusWaiting - default: - return false - } - }, testutil.WaitShort, testutil.IntervalFast) - - // Unblock the relay dial so the test can clean up. - close(dialContinue) -} - -func TestSubscribeRelaySnapshotDelivered(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - // Return a non-empty snapshot with two parts. - snapshot := []codersdk.ChatStreamEvent{ - { - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("snap-one"), - }, - }, - { - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("snap-two"), - }, - }, - } - ch := make(chan codersdk.ChatStreamEvent, 10) - // Also send a live part after the snapshot. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("live-part"), - }, - } - return snapshot, ch, func() {}, nil - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat := seedRemoteRunningChat(ctx, t, db, org.ID, user, model, workerID, "relay-snapshot") - staleChat := chat - staleChat.Status = database.ChatStatusWaiting - staleChat.WorkerID = uuid.NullUUID{} - staleChat.StartedAt = sql.NullTime{} - staleChat.HeartbeatAt = sql.NullTime{} - - initialSnapshot, events, cancel, ok := subscriber.SubscribeAuthorized(ctx, staleChat, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // The relay snapshot parts are forwarded through the events - // channel by the enterprise SubscribeFn. Collect them along - // with the live part. - var receivedTexts []string - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil { - receivedTexts = append(receivedTexts, event.MessagePart.Part.Text) - } - // We expect snap-one, snap-two, and live-part. - return len(receivedTexts) >= 3 - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - require.Equal(t, []string{"snap-one", "snap-two", "live-part"}, receivedTexts) - - // The initial snapshot should contain the refreshed running status, - // not the stale waiting status passed into SubscribeAuthorized. - var snapshotStatus codersdk.ChatStatus - for _, event := range initialSnapshot { - if event.Type == codersdk.ChatStreamEventTypeStatus && event.Status != nil { - snapshotStatus = event.Status.Status - } - } - require.Equal(t, codersdk.ChatStatusRunning, snapshotStatus) -} - -func TestSubscribeRetryEventAcrossInstances(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var streamCalls atomic.Int32 - firstStreamStarted := make(chan struct{}) - allowFirstFailure := make(chan struct{}) - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("retry-across-instances") - } - if streamCalls.Add(1) == 1 { - select { - case <-firstStreamStarted: - default: - close(firstStreamStarted) - } - <-allowFirstFailure - return chattest.OpenAIRateLimitResponse() - } - return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("retry", " complete")...) - }) - - worker := newActiveWorkerServer(t, db, ps, workerID) - subscriber := newTestServer(t, db, ps, subscriberID, func( - ctx context.Context, - chatID uuid.UUID, - targetWorkerID uuid.UUID, - requestHeader http.Header, - ) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, - ) { - if targetWorkerID != workerID { - return nil, nil, nil, xerrors.Errorf("unexpected relay target %s", targetWorkerID) - } - snapshot, events, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64) - if !ok { - return nil, nil, nil, xerrors.New("worker subscribe failed") - } - return snapshot, events, cancel, nil - }, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - chat, err := worker.CreateChat(ctx, osschatd.CreateOptions{ - OrganizationID: org.ID, - OwnerID: user.ID, - Title: "retry-across-instances", - ModelConfigID: model.ID, - InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - require.Eventually(t, func() bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusRunning && - fromDB.WorkerID.Valid && fromDB.WorkerID.UUID == workerID - }, testutil.WaitMedium, testutil.IntervalFast) - - select { - case <-firstStreamStarted: - case <-ctx.Done(): - t.Fatal("timed out waiting for first streaming attempt") - } - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer cancel() - - close(allowFirstFailure) - - var retryEvent *codersdk.ChatStreamRetry - var waitingSeen bool - var waitingBeforeRetry bool - var assistantMessageBeforeRetry bool - require.Eventually(t, func() bool { - select { - case event, ok := <-events: - if !ok { - return false - } - switch event.Type { - case codersdk.ChatStreamEventTypeRetry: - retryEvent = event.Retry - case codersdk.ChatStreamEventTypeMessage: - if event.Message != nil && event.Message.Role == codersdk.ChatMessageRoleAssistant { - if retryEvent == nil { - assistantMessageBeforeRetry = true - } - } - case codersdk.ChatStreamEventTypeStatus: - if event.Status != nil && event.Status.Status == codersdk.ChatStatusWaiting { - if retryEvent == nil { - waitingBeforeRetry = true - } - waitingSeen = true - } - } - return retryEvent != nil && waitingSeen - default: - return false - } - }, testutil.WaitLong, testutil.IntervalFast) - - require.NotNil(t, retryEvent) - require.Equal(t, 1, retryEvent.Attempt) - require.Greater(t, retryEvent.DelayMs, int64(0)) - require.Equal(t, codersdk.ChatErrorKindRateLimit, retryEvent.Kind) - require.Equal(t, "openai", retryEvent.Provider) - require.Equal(t, 429, retryEvent.StatusCode) - require.Contains(t, retryEvent.Error, "rate limiting requests") - require.False(t, assistantMessageBeforeRetry) - require.False(t, waitingBeforeRetry) - require.GreaterOrEqual(t, streamCalls.Load(), int32(2)) -} - -// TestSubscribeRelayStaleDialDiscardedAfterInterrupt verifies that when a -// user interrupts a streaming chat and sends a new message (which gets -// picked up by a different replica), an in-flight relay dial to the -// OLD replica is canceled/discarded and the relay connects to the -// NEW replica correctly. -func TestSubscribeRelayStaleDialDiscardedAfterInterrupt(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - oldWorkerID := uuid.New() - newWorkerID := uuid.New() - subscriberID := uuid.New() - - // Gate to hold the first dial until we're ready. - firstDialStarted := make(chan struct{}) - releaseFirstDial := make(chan struct{}) - - var callCount atomic.Int32 - - provider := func(ctx context.Context, _ uuid.UUID, workerID uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - ch := make(chan codersdk.ChatStreamEvent, 10) - if call == 1 { - // First dial (to old worker): signal that we started, - // then block until released or context canceled. - close(firstDialStarted) - select { - case <-releaseFirstDial: - case <-ctx.Done(): - return nil, nil, nil, ctx.Err() - } - // If we get here after being released (not canceled), - // return a stale part — this should be discarded. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("stale-part"), - }, - } - close(ch) - return nil, ch, func() {}, nil - } - // Second dial (to new worker): return a valid part. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("new-worker-part"), - }, - } - return nil, ch, func() {}, nil - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Seed the chat in waiting state so Subscribe does not try an initial - // relay. - chat := seedWaitingChat(t, db, org.ID, user, model, "stale-dial-test") - - // Subscribe while chat is in "waiting" state — no relay opened. - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Now simulate the chat being picked up by the OLD worker via pubsub. - // This triggers openRelayAsync in the merge loop. - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: oldWorkerID, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - oldRunningNotify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: oldWorkerID.String(), - } - oldRunningPayload, err := json.Marshal(oldRunningNotify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), oldRunningPayload) - require.NoError(t, err) - - // Wait for the first dial goroutine to start (it's blocked in the provider). - select { - case <-firstDialStarted: - case <-ctx.Done(): - t.Fatal("timed out waiting for first dial to start") - } - - // Simulate interrupt: chat goes to "waiting". - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusWaiting, - }) - require.NoError(t, err) - waitingNotify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusWaiting), - } - waitingPayload, err := json.Marshal(waitingNotify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), waitingPayload) - require.NoError(t, err) - - // Wait for the merge loop to process the waiting notification - // and emit the status event before publishing the new running - // notification. This avoids time.Sleep (banned by project - // policy) and provides a deterministic sync point. - require.Eventually(t, func() bool { - select { - case event := <-events: - return event.Type == codersdk.ChatStreamEventTypeStatus && - event.Status != nil && - event.Status.Status == codersdk.ChatStatusWaiting - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // Now the chat transitions to running on the NEW worker. - _, err = db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: newWorkerID, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - runningNotify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: newWorkerID.String(), - } - runningPayload, err := json.Marshal(runningNotify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), runningPayload) - require.NoError(t, err) - - // Now release the first dial (if it wasn't already canceled). - close(releaseFirstDial) - - // The subscriber should receive parts from the NEW worker, not the stale one. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "new-worker-part" { - return true - } - // If we get the stale part, the bug is present. - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "stale-part" { - t.Fatal("received stale part from old worker — relay did not cancel in-flight dial") - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // Drain the events channel for a while to ensure no late-arriving - // stale part sneaks in after the require.Eventually above returned. - // This closes the timing gap where "stale-part" could arrive after - // "new-worker-part" was already consumed. - require.Never(t, func() bool { - select { - case event := <-events: - return event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "stale-part" - default: - return false - } - }, 2*time.Second, testutil.IntervalFast) -} - -// TestSubscribeCancelDuringInFlightDial verifies that calling the -// subscription's cancel function while a relay dial goroutine is -// still blocking in the provider causes the provider's context to -// be canceled and the goroutine to return cleanly. -func TestSubscribeCancelDuringInFlightDial(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - dialStarted := make(chan struct{}) - dialExited := make(chan struct{}) - - provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - // Signal the dial has started, then block until the context - // is canceled. - close(dialStarted) - <-ctx.Done() - close(dialExited) - return nil, nil, nil, ctx.Err() - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Seed the chat in waiting state so Subscribe does not open a - // synchronous relay. - chat := seedWaitingChat(t, db, org.ID, user, model, "cancel-inflight-dial") - - _, _, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - - // Publish a running notification to trigger openRelayAsync. - notify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: workerID.String(), - } - payload, err := json.Marshal(notify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload) - require.NoError(t, err) - - // Wait for the dial goroutine to block inside the provider. - select { - case <-dialStarted: - case <-ctx.Done(): - t.Fatal("timed out waiting for dial to start") - } - - // Cancel the subscription while the dial is still in-flight. - cancel() - - // The provider context must be canceled, causing the goroutine - // to return cleanly. - require.Eventually(t, func() bool { - select { - case <-dialExited: - return true - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) -} - -// TestSubscribeRelayRunningToRunningSwitch verifies that when a chat -// transitions directly from running(workerA) to running(workerB) -// without an intermediate waiting state, the relay switches to the -// new worker and discards parts from the old one. -func TestSubscribeRelayRunningToRunningSwitch(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerA := uuid.New() - workerB := uuid.New() - subscriberID := uuid.New() - - // Gate to hold workerA's dial until we verify cancellation. - dialAStarted := make(chan struct{}) - dialAExited := make(chan struct{}) - - var callCount atomic.Int32 - - provider := func(ctx context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - if call == 1 { - // First dial (to workerA): signal that we started, - // then block until the context is canceled. - close(dialAStarted) - <-ctx.Done() - close(dialAExited) - return nil, nil, nil, ctx.Err() - } - // Second dial (to workerB): return a valid part. - ch := make(chan codersdk.ChatStreamEvent, 10) - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("worker-b-part"), - }, - } - return nil, ch, func() {}, nil - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Seed the chat in waiting state so Subscribe does not open a relay. - chat := seedWaitingChat(t, db, org.ID, user, model, "running-to-running") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Transition to running on workerA. - notifyA := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: workerA.String(), - } - payloadA, err := json.Marshal(notifyA) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payloadA) - require.NoError(t, err) - - // Wait for the workerA dial goroutine to block inside the - // provider before publishing the workerB notification. - select { - case <-dialAStarted: - case <-ctx.Done(): - t.Fatal("timed out waiting for workerA dial to start") - } - - // Immediately transition to running on workerB (no waiting in - // between). This should cancel workerA's in-flight dial. - notifyB := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: workerB.String(), - } - payloadB, err := json.Marshal(notifyB) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payloadB) - require.NoError(t, err) - - // Verify that the relay canceled workerA's stale dial. - require.Eventually(t, func() bool { - select { - case <-dialAExited: - return true - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // We should receive the part from workerB. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "worker-b-part" { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - require.Equal(t, 2, int(callCount.Load())) -} - -// TestSubscribeRelayFailedDialRetries verifies that when an async relay -// dial fails (returns an error), the merge loop schedules a reconnect -// timer and eventually re-dials successfully. This exercises the -// result.parts == nil path and the scheduleRelayReconnect() logic. -func TestSubscribeRelayFailedDialRetries(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - remoteWorkerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - - provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - if call == 1 { - // First dial: fail with an error to trigger - // scheduleRelayReconnect via the result.parts == nil path. - return nil, nil, nil, xerrors.New("transient dial failure") - } - // Second dial: succeed and return a part. - ch := make(chan codersdk.ChatStreamEvent, 10) - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("retry-success"), - }, - } - return nil, ch, func() {}, nil - } - - mclk := quartz.NewMock(t) - // Trap the reconnect timer so we can fire it deterministically. - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - // Seed the chat in waiting state so Subscribe does not open a - // synchronous relay dial. - chat := seedWaitingChat(t, db, org.ID, user, model, "failed-dial-retry") - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Now mark the chat as running on the remote worker in the DB. - // The reconnect timer calls params.DB.GetChatByID to check if - // the chat is still running on a remote worker, so this must be - // set before we advance the clock. - _, err := db.UpdateChatStatus(ctx, database.UpdateChatStatusParams{ - ID: chat.ID, - Status: database.ChatStatusRunning, - WorkerID: uuid.NullUUID{UUID: remoteWorkerID, Valid: true}, - StartedAt: sql.NullTime{Time: time.Now(), Valid: true}, - HeartbeatAt: sql.NullTime{Time: time.Now(), Valid: true}, - }) - require.NoError(t, err) - - // Publish a running notification with a remote workerID to - // trigger openRelayAsync. The first dial will fail, causing - // scheduleRelayReconnect to be called. - notify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: remoteWorkerID.String(), - } - payload, err := json.Marshal(notify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload) - require.NoError(t, err) - - // Wait for the reconnect timer to be created (after the failed - // dial), then advance the mock clock to fire it. - trapReconnect.MustWait(ctx).MustRelease(ctx) - mclk.Advance(500 * time.Millisecond).MustWait(ctx) - - // The merge loop re-checks the DB, sees the chat is still - // running on the remote worker, and dials again. The second - // dial succeeds. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "retry-success" { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - require.GreaterOrEqual(t, int(callCount.Load()), 2) -} - -// TestSubscribeRunningLocalWorkerClosesRelay verifies that when a chat -// is running on a remote worker and a pubsub notification arrives -// saying the local worker (subscriberID) now owns the chat, the -// existing relay is closed and no new dial is started (the local -// worker serves directly without relaying). -func TestSubscribeRunningLocalWorkerClosesRelay(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - remoteWorkerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - - provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - ch := make(chan codersdk.ChatStreamEvent, 10) - if call == 1 { - // Initial synchronous dial to the remote worker. - ch <- codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessageText("remote-part"), - }, - } - // Keep channel open so the relay stays active. - } - return nil, ch, func() {}, nil - } - - subscriber := newTestServer(t, db, ps, subscriberID, provider, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat := seedRemoteRunningChat( - ctx, - t, - db, - org.ID, - user, - model, - remoteWorkerID, - "local-worker-closes-relay", - ) - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Consume the remote-part from the initial relay. - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == "remote-part" { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // Notify that the LOCAL worker now owns the chat. This should - // close the relay without opening a new one. - notify := coderdpubsub.ChatStreamNotifyMessage{ - Status: string(database.ChatStatusRunning), - WorkerID: subscriberID.String(), - } - payload, err := json.Marshal(notify) - require.NoError(t, err) - err = ps.Publish(coderdpubsub.ChatStreamNotifyChannel(chat.ID), payload) - require.NoError(t, err) - - // Give the system time to process the notification. No additional - // dial should happen — only the initial synchronous one. - require.Never(t, func() bool { - return int(callCount.Load()) > 1 - }, 2*time.Second, testutil.IntervalFast) - - require.Equal(t, 1, int(callCount.Load()), - "only the initial synchronous dial should have happened") -} - -// TestSubscribeRelayMultipleReconnects verifies that the reconnect -// loop handles multiple consecutive relay drops, proving it is -// robust across repeated iterations — not just the single reconnect -// already covered by TestSubscribeRelayReconnectsOnDrop. -func TestSubscribeRelayMultipleReconnects(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var callCount atomic.Int32 - - provider := func(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ http.Header) ( - []codersdk.ChatStreamEvent, <-chan codersdk.ChatStreamEvent, func(), error, - ) { - call := callCount.Add(1) - ch := make(chan codersdk.ChatStreamEvent, 10) - part := codersdk.ChatStreamEvent{ - Type: codersdk.ChatStreamEventTypeMessagePart, - MessagePart: &codersdk.ChatStreamMessagePart{ - Role: "assistant", - Part: codersdk.ChatMessagePart{ - Type: codersdk.ChatMessagePartTypeText, - Text: fmt.Sprintf("relay-%d", call), - }, - }, - } - ch <- part - if call <= 2 { - // First two dials: close channel to simulate relay - // drop. This triggers scheduleRelayReconnect. - close(ch) - } - // Third dial: keep channel open. - return nil, ch, func() {}, nil - } - - mclk := quartz.NewMock(t) - // Trap the reconnect timer so we can fire both reconnects - // deterministically. - trapReconnect := mclk.Trap().NewTimer("reconnect") - defer trapReconnect.Close() - - subscriber := newTestServer(t, db, ps, subscriberID, provider, mclk) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - - chat := seedRemoteRunningChat( - ctx, - t, - db, - org.ID, - user, - model, - workerID, - "multiple-reconnects", - ) - - _, events, cancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - t.Cleanup(cancel) - - // Helper to consume a specific relay part. - consumePart := func(text string) { - t.Helper() - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessagePart && - event.MessagePart != nil && - event.MessagePart.Part.Text == text { - return true - } - return false - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - } - - // First relay: consumed immediately (synchronous dial). - consumePart("relay-1") - - // First relay drops → reconnect timer created. Advance clock - // to fire it. - trapReconnect.MustWait(ctx).MustRelease(ctx) - mclk.Advance(500 * time.Millisecond).MustWait(ctx) - - // Second relay part. - consumePart("relay-2") - - // Second relay drops → another reconnect timer. Advance again. - trapReconnect.MustWait(ctx).MustRelease(ctx) - mclk.Advance(500 * time.Millisecond).MustWait(ctx) - - // Third relay part (channel stays open). - consumePart("relay-3") - require.GreaterOrEqual(t, int(callCount.Load()), 3) -} - -// TestSubscribeRelayDialCanceledOnFastCompletion verifies that a -// subscriber on a remote replica still sees the committed assistant -// response when the worker completes faster than the relay dial. -// -// Scenario: -// 1. Subscriber subscribes to a chat while it's in waiting state (no relay). -// 2. User sends a message → chat becomes pending → worker picks it up. -// 3. Subscriber receives status=running via pubsub → enterprise opens relay async. -// 4. Worker completes quickly → publishes committed message + status=waiting. -// 5. Subscriber receives status=waiting → enterprise cancels the in-progress relay dial. -// 6. Even though the relay never delivered streaming parts, the -// committed assistant message arrives via pubsub so the user -// does not need to refresh to see the response. -// -// Streaming parts for committed turns are intentionally NOT replayed -// via the relay: they would duplicate the durable message on the -// user's screen. The buffer retains in-progress parts only; once an -// assistant turn commits, the parts that built it are claimed by -// the durable message ID and dropped from new buffer snapshots. -func TestSubscribeRelayDialCanceledOnFastCompletion(t *testing.T) { - t.Parallel() - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - var dialAttempted atomic.Bool - - // Gate: closed when the worker finishes processing. - workerDone := make(chan struct{}) - - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("fast-completion-relay-race") - } - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("hello ", "world ", "from ", "the ", "worker")..., - ) - }) - - // Worker server with a 1-hour acquire interval so it only processes - // when explicitly woken by SendMessage's signalWake. - workerLogger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - worker := osschatd.New(osschatd.Config{ - Logger: workerLogger, - Database: db, - ReplicaID: workerID, - Pubsub: ps, - PendingChatAcquireInterval: time.Hour, - InFlightChatStaleAfter: testutil.WaitSuperLong, - }) - worker.Start() - t.Cleanup(func() { - require.NoError(t, worker.Close()) - }) - - // Subscriber's relay dialer blocks until the worker finishes, - // simulating a slow relay dial (network latency between replicas). - // After the worker completes, the dialer connects to the worker - // to retrieve buffered parts from the retained buffer. - subscriber := newTestServer(t, db, ps, subscriberID, func( - ctx context.Context, - chatID uuid.UUID, - targetWorkerID uuid.UUID, - requestHeader http.Header, - ) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, - ) { - dialAttempted.Store(true) - // Block until the worker finishes processing, simulating - // a slow relay dial. - select { - case <-workerDone: - case <-ctx.Done(): - return nil, nil, nil, ctx.Err() - } - // Connect to the worker. The buffer is retained for a - // grace period after processing, so the relay session - // can complete (control events, status updates) even - // though every part has been claimed by its durable - // message and the snapshot is empty. - snapshot, relayEvents, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64) - if !ok { - return nil, nil, nil, xerrors.New("worker subscribe failed") - } - return snapshot, relayEvents, cancel, nil - }, nil) - - ctx := testutil.Context(t, testutil.WaitLong) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - // Create the chat in waiting state so the subscriber sees it - // before the worker picks it up (avoids the synchronous relay - // path in Subscribe). - chat := seedWaitingChat(t, db, org.ID, user, model, "fast-completion-relay-race") - - // Subscribe from the subscriber replica while the chat is idle. - // No relay is opened because the chat is in waiting state. - _, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer subCancel() - - // Send a message via the worker server to transition the chat to - // pending and wake the worker's processing loop. - _, err := worker.SendMessage(ctx, osschatd.SendMessageOptions{ - ChatID: chat.ID, - CreatedBy: user.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // Wait for the worker to fully process the chat. - require.Eventually(t, func() bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusWaiting - }, testutil.WaitMedium, testutil.IntervalFast) - - // Release the relay dial now that the worker is done. - close(workerDone) - - // Collect events that arrived at the subscriber. The committed - // assistant message is guaranteed to arrive via pubsub even when - // the relay dial races worker completion; streaming parts are - // best-effort and are not asserted here because the buffer drops - // already-committed parts to prevent duplicate UI rendering. - var committedAssistantMsgs int - - require.Eventually(t, func() bool { - select { - case event := <-events: - if event.Type == codersdk.ChatStreamEventTypeMessage && - event.Message != nil && - event.Message.Role == codersdk.ChatMessageRoleAssistant { - committedAssistantMsgs++ - } - return committedAssistantMsgs > 0 - default: - return false - } - }, testutil.WaitLong, testutil.IntervalFast) - - // The committed assistant message arrives via pubsub → DB query - // (durable path). - require.Equal(t, 1, committedAssistantMsgs, - "committed assistant message should arrive via pubsub durable path") - - // The relay dial was attempted when status=running arrived. - require.True(t, dialAttempted.Load(), - "relay dial should have been attempted when status changed to running") -} - -// TestSubscribeRelayEstablishedMidStream demonstrates that when the -// relay is established while the worker is still streaming, the -// subscriber receives buffered parts via the relay snapshot and live -// parts through the relay channel. -// -// This is the complementary test to TestSubscribeRelayDialCanceledOnFastCompletion: -// it shows the relay mechanism works correctly when timing is favorable -// (relay connects before the worker finishes), contrasting with the race -// condition where the relay is too slow. -func TestSubscribeRelayEstablishedMidStream(t *testing.T) { - t.Parallel() - // TODO(CODAGT-353): Re-enable this test after the chatd notification flow - // refactor gives workers enough causal information to distinguish stale - // control NOTIFY messages from real interrupts. The current design reuses - // the same status notification shape for wake-only and interrupt intents, - // so a stale NOTIFY can cancel a new processChat run. - t.Skip("skipped until chatd notification flow refactor handles stale control notifications") - - db, ps := dbtestutil.NewDB(t) - workerID := uuid.New() - subscriberID := uuid.New() - - // Gate: worker blocks after first streaming request until we - // release it. This gives the relay time to establish. - firstChunkEmitted := make(chan struct{}) - continueStreaming := make(chan struct{}) - - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { - if !req.Stream { - return chattest.OpenAINonStreamingResponse("mid-stream-relay") - } - // Signal that the first streaming request was received, - // then block until released. - select { - case <-firstChunkEmitted: - default: - close(firstChunkEmitted) - } - <-continueStreaming - return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("continued ", "response")..., - ) - }) - - // Worker with a short fallback poll interval. The primary - // trigger is signalWake() from SendMessage, but under heavy - // CI load the wake goroutine may be delayed. A short poll - // ensures the worker always picks up the pending chat. - workerLogger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) - worker := osschatd.New(osschatd.Config{ - Logger: workerLogger, - Database: db, - ReplicaID: workerID, - Pubsub: ps, - PendingChatAcquireInterval: time.Second, - InFlightChatStaleAfter: testutil.WaitSuperLong, - }) - worker.Start() - t.Cleanup(func() { - require.NoError(t, worker.Close()) - }) - - // Subscriber's dialer connects to the worker with no delay. - // This simulates a relay that succeeds promptly. - subscriber := newTestServer(t, db, ps, subscriberID, func( - ctx context.Context, - chatID uuid.UUID, - targetWorkerID uuid.UUID, - requestHeader http.Header, - ) ( - []codersdk.ChatStreamEvent, - <-chan codersdk.ChatStreamEvent, - func(), - error, - ) { - if targetWorkerID != workerID { - return nil, nil, nil, xerrors.Errorf("unexpected relay target %s", targetWorkerID) - } - snapshot, relayEvents, cancel, ok := worker.Subscribe(ctx, chatID, requestHeader, math.MaxInt64) - if !ok { - return nil, nil, nil, xerrors.New("worker subscribe failed") - } - return snapshot, relayEvents, cancel, nil - }, nil) - - // Use WaitSuperLong so the test survives heavy CI contention. - // The worker pipeline (model resolution, message loading, LLM - // call) involves multiple DB round-trips that can be slow under - // load. - ctx := testutil.Context(t, testutil.WaitSuperLong) - user, org, model := seedChatDependencies(t, db) - setOpenAIProviderBaseURL(ctx, t, db, openAIURL) - - // Create the chat in waiting state. - chat := seedWaitingChat(t, db, org.ID, user, model, "mid-stream-relay") - - // Subscribe from the subscriber replica while the chat is idle. - _, events, subCancel, ok := subscriber.Subscribe(ctx, chat.ID, nil, 0) - require.True(t, ok) - defer subCancel() - - // Send a message to make the chat pending and wake the worker. - _, err := worker.SendMessage(ctx, osschatd.SendMessageOptions{ - ChatID: chat.ID, - CreatedBy: user.ID, - Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, - }) - require.NoError(t, err) - - // Wait for the worker to reach the LLM (first streaming - // request). Also poll the chat status so we fail fast with a - // clear message if the worker errors out instead of timing - // out silently. - ticker := time.NewTicker(250 * time.Millisecond) - defer ticker.Stop() -waitForStream: - for { - select { - case <-firstChunkEmitted: - break waitForStream - case <-ticker.C: - currentChat, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr == nil && currentChat.Status == database.ChatStatusError { - t.Fatalf("worker failed to process chat: status=%s last_error=%s", - currentChat.Status, chatLastErrorMessage(currentChat.LastError)) - } - case <-ctx.Done(): - // Dump the final chat status for debugging. - currentChat, dbErr := db.GetChatByID(context.Background(), chat.ID) - if dbErr == nil { - t.Fatalf("timed out waiting for worker to start streaming (chat status=%s, last_error=%q)", - currentChat.Status, chatLastErrorMessage(currentChat.LastError)) - } - t.Fatal("timed out waiting for worker to start streaming") - } - } - - // Wait for the subscriber to receive the running status, which - // triggers the relay. Because the dialer is non-blocking, the - // relay establishes promptly. - require.Eventually(t, func() bool { - select { - case event := <-events: - return event.Type == codersdk.ChatStreamEventTypeStatus && - event.Status != nil && - event.Status.Status == codersdk.ChatStatusRunning - default: - return false - } - }, testutil.WaitMedium, testutil.IntervalFast) - - // Now release the worker to continue streaming. - close(continueStreaming) - - // Wait for the worker to complete. - require.Eventually(t, func() bool { - fromDB, dbErr := db.GetChatByID(ctx, chat.ID) - if dbErr != nil { - return false - } - return fromDB.Status == database.ChatStatusWaiting - }, testutil.WaitMedium, testutil.IntervalFast) - - // Collect remaining events. - var messageParts []string - var hasCommittedMsg bool - - require.Eventually(t, func() bool { - select { - case event := <-events: - switch event.Type { - case codersdk.ChatStreamEventTypeMessagePart: - if event.MessagePart != nil { - messageParts = append(messageParts, event.MessagePart.Part.Text) - } - case codersdk.ChatStreamEventTypeMessage: - if event.Message != nil && event.Message.Role == codersdk.ChatMessageRoleAssistant { - hasCommittedMsg = true - } - } - return hasCommittedMsg - default: - return false - } - }, testutil.WaitLong, testutil.IntervalFast) - - // The committed message arrives via pubsub. - require.True(t, hasCommittedMsg, - "committed assistant message should arrive") - - // When the relay is established mid-stream, streaming parts - // SHOULD be received through the relay. This contrasts with - // TestSubscribeRelayDialCanceledOnFastCompletion where no parts - // arrive because the relay is never established. - require.NotEmpty(t, messageParts, - "streaming parts should be received when relay establishes while worker is still streaming") + session, err := dialer(context.Background(), osschatd.StreamPartsDialInput{ + ChatID: chatID, + WorkerID: workerID, + }) + require.Nil(t, session) + var dialErr *entchatd.RelayDialError + require.ErrorAs(t, err, &dialErr) + require.Equal(t, http.StatusUnauthorized, dialErr.HTTPStatus) + require.True(t, dialErr.IsUnrecoverable()) } diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index da019143df..76d25ef341 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -250,21 +250,9 @@ coderd_chatd_steps_total{provider="",model=""} 0 # HELP coderd_chatd_stream_buffer_dropped_total Number of chat stream buffer events dropped due to the per-chat buffer cap. # TYPE coderd_chatd_stream_buffer_dropped_total counter coderd_chatd_stream_buffer_dropped_total 0 -# HELP coderd_chatd_stream_buffer_events Sum of current buffer lengths across all chat streams. -# TYPE coderd_chatd_stream_buffer_events gauge -coderd_chatd_stream_buffer_events 0 -# HELP coderd_chatd_stream_buffer_size_max Maximum current buffer length across all chat streams. -# TYPE coderd_chatd_stream_buffer_size_max gauge -coderd_chatd_stream_buffer_size_max 0 # HELP coderd_chatd_stream_retries_total Total LLM stream retries. # TYPE coderd_chatd_stream_retries_total counter coderd_chatd_stream_retries_total{provider="",model="",kind="",chain_broken=""} 0 -# HELP coderd_chatd_stream_subscribers Current number of chat stream subscribers across all chat streams. -# TYPE coderd_chatd_stream_subscribers gauge -coderd_chatd_stream_subscribers 0 -# HELP coderd_chatd_streams_active Current number of chat stream state entries (in-flight plus retained). -# TYPE coderd_chatd_streams_active gauge -coderd_chatd_streams_active 0 # HELP coderd_chatd_tool_errors_total Total tool calls that returned an error result. # TYPE coderd_chatd_tool_errors_total counter coderd_chatd_tool_errors_total{provider="",model="",tool_name=""} 0 diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 34d8c6f18d..7eab4ef02d 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -2701,6 +2701,7 @@ export interface ChatSourcePart { export type ChatStatus = | "completed" | "error" + | "interrupting" | "paused" | "pending" | "requires_action" @@ -2710,6 +2711,7 @@ export type ChatStatus = export const ChatStatuses: ChatStatus[] = [ "completed", "error", + "interrupting", "paused", "pending", "requires_action", @@ -2745,8 +2747,10 @@ export interface ChatStreamEvent { export type ChatStreamEventType = | "action_required" | "error" + | "history_reset" | "message" | "message_part" + | "preview_reset" | "queue_update" | "retry" | "status"; @@ -2754,8 +2758,10 @@ export type ChatStreamEventType = export const ChatStreamEventTypes: ChatStreamEventType[] = [ "action_required", "error", + "history_reset", "message", "message_part", + "preview_reset", "queue_update", "retry", "status", @@ -2768,6 +2774,9 @@ export const ChatStreamEventTypes: ChatStreamEventType[] = [ export interface ChatStreamMessagePart { readonly role?: ChatMessageRole; readonly part: ChatMessagePart; + readonly history_version?: number; + readonly generation_attempt?: number; + readonly seq?: number; } // From codersdk/chats.go diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index a8ba38f4c2..bce82a39e7 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -237,6 +237,19 @@ const makeMessage = ( content: [{ type: "text", text }], }); +const makeMessageWithContent = ( + chatID: string, + id: number, + role: TypesGen.ChatMessageRole, + content: readonly TypesGen.ChatMessagePart[], +): TypesGen.ChatMessage => ({ + id, + chat_id: chatID, + created_at: "2025-01-01T00:00:00.000Z", + role, + content, +}); + const makeQueuedMessage = ( chatID: string, id: number, @@ -344,6 +357,553 @@ describe("useChatStore", () => { }); }); + it("keeps create_workspace durable call without result after preview_reset", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-create-workspace"; + const existingMessage = makeMessage(chatID, 1, "user", "create workspace"); + const assistantMessage = makeMessageWithContent(chatID, 2, "assistant", [ + { + type: "tool-call", + tool_call_id: "create-workspace-1", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + ]); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-call", + tool_call_id: "create-workspace-1", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + }, + }); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect( + result.current.streamState?.toolCalls["create-workspace-1"]?.name, + ).toBe("create_workspace"); + }); + + act(() => { + mockSocket.emitDataBatch([ + { type: "message", chat_id: chatID, message: assistantMessage }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.streamState).toBeNull(); + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.messagesByID.get(2)?.content).toEqual( + assistantMessage.content, + ); + }); + }); + + it("clears stream state when preview_reset arrives after durable tool result", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-tool-result"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const assistantMessage = makeMessageWithContent(chatID, 2, "assistant", [ + { + type: "tool-call", + tool_call_id: "tool-1", + tool_name: "read_template", + args: { template_id: "template-1" }, + }, + ]); + const toolMessage = makeMessageWithContent(chatID, 3, "tool", [ + { + type: "tool-result", + tool_call_id: "tool-1", + tool_name: "read_template", + result: { name: "Template" }, + }, + ]); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-call", + tool_call_id: "tool-1", + tool_name: "read_template", + args: { template_id: "template-1" }, + }, + }, + }); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect(result.current.streamState?.toolCalls["tool-1"]?.name).toBe( + "read_template", + ); + }); + + act(() => { + mockSocket.emitDataBatch([ + { type: "message", chat_id: chatID, message: assistantMessage }, + { type: "preview_reset", chat_id: chatID }, + { + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-result", + tool_call_id: "tool-1", + tool_name: "read_template", + result: { name: "Template" }, + }, + }, + }, + { type: "message", chat_id: chatID, message: toolMessage }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + expect(result.current.streamState).toBeNull(); + }); + }); + + it("preview_reset discards pending buffered parts", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-buffer"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + expect(result.current.streamState).toBeNull(); + }); + + it("keeps only post-reset parts after preview_reset in one batch", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-post-part"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const durableMessage = makeMessage(chatID, 2, "assistant", "done"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "message", chat_id: chatID, message: durableMessage }, + { type: "preview_reset", chat_id: chatID }, + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "fresh" }, + }, + }, + ]); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "fresh" }, + ]); + }); + }); + + it("replaces messages after history_reset", async () => { + const chatID = "chat-history-reset"; + const initialMessages = [ + makeMessage(chatID, 1, "user", "old prompt"), + makeMessage(chatID, 2, "assistant", "old answer"), + makeMessage(chatID, 3, "user", "stale prompt"), + ]; + const replacementMessage = makeMessage(chatID, 1, "user", "new prompt"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: Number.POSITIVE_INFINITY, + refetchOnWindowFocus: false, + networkMode: "offlineFirst", + }, + }, + }); + queryClient.setQueryData(chatMessagesKey(chatID), { + pages: [ + { + messages: [...initialMessages].reverse(), + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "history_reset", chat_id: chatID }, + { type: "message", chat_id: chatID, message: replacementMessage }, + // The server always emits preview_reset after a history + // change in the same sync; it terminates the replacement run. + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1]); + expect(result.current.messagesByID.get(1)?.content).toEqual( + replacementMessage.content, + ); + expect(result.current.streamState).toBeNull(); + }); + + const cached = queryClient.getQueryData<{ + pages: TypesGen.ChatMessagesResponse[]; + pageParams: unknown[]; + }>(chatMessagesKey(chatID)); + expect(cached?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 1, + ]); + }); + + it("buffers a history_reset replacement split across WS frames", async () => { + const chatID = "chat-history-reset-split"; + const initialMessages = [ + makeMessage(chatID, 1, "user", "old prompt"), + makeMessage(chatID, 2, "assistant", "old answer"), + makeMessage(chatID, 3, "user", "stale prompt"), + ]; + const replacementOne = makeMessage(chatID, 1, "user", "new prompt"); + const replacementTwo = makeMessage(chatID, 2, "assistant", "new answer"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: Number.POSITIVE_INFINITY, + refetchOnWindowFocus: false, + networkMode: "offlineFirst", + }, + }, + }); + queryClient.setQueryData(chatMessagesKey(chatID), { + pages: [ + { + messages: [...initialMessages].reverse(), + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + // Frame 1: history_reset plus only part of the replacement + // history. The conversation must not blank or truncate while + // the rest of the run is in flight. + act(() => { + mockSocket.emitDataBatch([ + { type: "history_reset", chat_id: chatID }, + { type: "message", chat_id: chatID, message: replacementOne }, + ]); + }); + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + + // Frame 2: the rest of the replacement, terminated by the + // preview_reset the server emits in the same sync. + act(() => { + mockSocket.emitDataBatch([ + { type: "message", chat_id: chatID, message: replacementTwo }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.messagesByID.get(1)?.content).toEqual( + replacementOne.content, + ); + expect(result.current.messagesByID.get(2)?.content).toEqual( + replacementTwo.content, + ); + }); + + const cached = queryClient.getQueryData<{ + pages: TypesGen.ChatMessagesResponse[]; + pageParams: unknown[]; + }>(chatMessagesKey(chatID)); + expect(cached?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 2, 1, + ]); + }); + it("clears stream state when a new durable message arrives", async () => { immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 8aca1e0bc2..44020f9e10 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts @@ -142,7 +142,8 @@ const reconnectStatesEqual = ( export const isActiveChatStatus = ( status: TypesGen.ChatStatus | null, -): boolean => status === "running" || status === "pending"; +): boolean => + status === "running" || status === "pending" || status === "interrupting"; export type ChatStoreState = { messagesByID: Map; diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts index 80e4a0254f..5e0a259fb4 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts @@ -4,6 +4,7 @@ import { getSubagentDescriptor } from "../ChatElements/tools/subagentDescriptor" import { buildSubagentMaps, getEditableUserMessagePayload, + getPendingToolCallIDs, mergeTools, parseMessageContent, parseMessagesWithMergedTools, @@ -519,6 +520,115 @@ describe("mergeTools", () => { expect(merged).toHaveLength(1); expect(merged[0].status).toBe("completed"); }); + + it("marks unresolved pending calls as running", () => { + const merged = mergeTools([{ id: "1", name: "bash" }], [], { + pendingToolCallIDs: new Set(["1"]), + }); + expect(merged).toHaveLength(1); + expect(merged[0].status).toBe("running"); + }); +}); + +describe("pending durable tool parsing", () => { + const msg = ( + id: number, + role: "assistant" | "tool" | "user", + parts: ChatMessagePart[], + ): ChatMessage => ({ + id, + chat_id: "chat-1", + created_at: new Date(2026, 0, id).toISOString(), + role, + content: parts, + }); + + const toolCall = ( + id: string, + name = "create_workspace", + ): ChatMessagePart => ({ + type: "tool-call", + tool_call_id: id, + tool_name: name, + args: { name: "dev" }, + }); + + const toolResult = ( + id: string, + name = "create_workspace", + isError = false, + ): ChatMessagePart => ({ + type: "tool-result", + tool_call_id: id, + tool_name: name, + result: { workspace_name: "dev", build_id: "build-1" }, + is_error: isError, + }); + + it("marks the latest unresolved assistant tool call as running in an active chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-active")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "running"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("running"); + }); + + it("keeps unresolved historical tool calls completed in an inactive chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-inactive")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "waiting"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("completed"); + }); + + it("uses completed or error status once a matching result exists", () => { + const successMessages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-success")]), + msg(3, "tool", [toolResult("call-success")]), + ]; + const errorMessages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-error")]), + msg(3, "tool", [toolResult("call-error", "create_workspace", true)]), + ]; + + const successParsed = parseMessagesWithMergedTools(successMessages, { + pendingToolCallIDs: getPendingToolCallIDs(successMessages, "running"), + }); + const errorParsed = parseMessagesWithMergedTools(errorMessages, { + pendingToolCallIDs: getPendingToolCallIDs(errorMessages, "running"), + }); + + expect(successParsed[1]?.parsed.tools[0]?.status).toBe("completed"); + expect(errorParsed[1]?.parsed.tools[0]?.status).toBe("error"); + }); + + it("does not mark older unresolved rows running in an active chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "first" }]), + msg(2, "assistant", [toolCall("call-old", "read_file")]), + msg(3, "user", [{ type: "text", text: "second" }]), + msg(4, "assistant", [toolCall("call-latest", "create_workspace")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "running"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("completed"); + expect(parsed[3]?.parsed.tools[0]?.status).toBe("running"); + }); }); describe("parseMessagesWithMergedTools — killedBySignal annotation", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts index 9e118d0656..3f48e6acfa 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts @@ -83,9 +83,65 @@ export const ensureToolBlock = ( return [...blocks, { type: "tool", id }]; }; +const isToolCallPart = ( + part: TypesGen.ChatMessagePart, +): part is TypesGen.ChatToolCallPart => part.type === "tool-call"; + +const isToolResultPart = ( + part: TypesGen.ChatMessagePart, +): part is TypesGen.ChatToolResultPart => part.type === "tool-result"; + +const chatHasActiveToolCalls = (status: TypesGen.ChatStatus | null): boolean => + status === "running" || status === "requires_action"; + +export const getPendingToolCallIDs = ( + messages: readonly TypesGen.ChatMessage[], + chatStatus: TypesGen.ChatStatus | null, +): ReadonlySet | undefined => { + if (!chatHasActiveToolCalls(chatStatus)) { + return undefined; + } + + const resultIDs = new Set(); + for (const message of messages) { + for (const part of message.content ?? []) { + if (isToolResultPart(part) && part.tool_call_id) { + resultIDs.add(part.tool_call_id); + } + } + } + + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index]; + if (!message) { + continue; + } + if (message.role === "user") { + return undefined; + } + if (message.role !== "assistant") { + continue; + } + const pendingToolCallIDs = (message.content ?? []) + .filter(isToolCallPart) + .map((part) => part.tool_call_id) + .filter((id): id is string => Boolean(id && !resultIDs.has(id))); + return pendingToolCallIDs.length > 0 + ? new Set(pendingToolCallIDs) + : undefined; + } + + return undefined; +}; + +type MergeToolsOptions = { + pendingToolCallIDs?: ReadonlySet; +}; + export const mergeTools = ( calls: ParsedToolCall[], results: ParsedToolResult[], + options: MergeToolsOptions = {}, ): MergedTool[] => { const resultById = new Map(results.map((r) => [r.id, r])); const seen = new Set(); @@ -100,13 +156,20 @@ export const mergeTools = ( typeof callArgs?.model_intent === "string" ? callArgs.model_intent : undefined; + const status = result + ? result.isError + ? "error" + : "completed" + : options.pendingToolCallIDs?.has(call.id) + ? "running" + : "completed"; merged.push({ id: call.id, name: call.name, args: call.args, result: result?.result, isError: result?.isError ?? false, - status: result ? (result.isError ? "error" : "completed") : "completed", + status, mcpServerConfigId: call.mcpServerConfigId || result?.mcpServerConfigId, modelIntent, parsedCommands: call.parsedCommands, @@ -272,8 +335,13 @@ export const getEditableUserMessagePayload = ( }; }; +type ParseMessagesWithMergedToolsOptions = { + pendingToolCallIDs?: ReadonlySet; +}; + export const parseMessagesWithMergedTools = ( messages: readonly TypesGen.ChatMessage[], + options: ParseMessagesWithMergedToolsOptions = {}, ): ParsedMessageEntry[] => { const rawParsed = messages.map((message) => ({ message, @@ -303,6 +371,7 @@ export const parseMessagesWithMergedTools = ( parsed.tools = mergeTools( parsed.toolCalls, Array.from(resultById.values()), + { pendingToolCallIDs: options.pendingToolCallIDs }, ); } diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index e2cbe573f1..077a3fd6b6 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -184,6 +184,29 @@ export const useChatStore = ( [chatID, queryClient], ); + const replaceCacheMessages = useCallback( + (messages: readonly TypesGen.ChatMessage[]) => { + if (!chatID) { + return; + } + queryClient.setQueryData< + InfiniteData | undefined + >(chatMessagesKey(chatID), (currentData) => { + if (!currentData?.pages?.length) { + return currentData; + } + const firstPage = currentData.pages[0]; + const updatedMessages = [...messages].sort((a, b) => b.id - a.id); + return { + ...currentData, + pages: [{ ...firstPage, messages: updatedMessages, has_more: false }], + pageParams: currentData.pageParams.slice(0, 1), + }; + }); + }, + [chatID, queryClient], + ); + useEffect(() => { store.batch(() => { // When the active chat changes, clear stale messages @@ -352,6 +375,16 @@ export const useChatStore = ( const partsBuf: TypesGen.ChatMessagePart[] = []; let partsFlushTimer: ReturnType | null = null; + // History replacement state lives at the effect scope because + // the server may split a history_reset and its replacement + // messages across multiple WS frames (the stream handler caps + // frames at a fixed batch size). Replacement messages are + // buffered until a non-message boundary event arrives; the + // server always emits preview_reset after a history change in + // the same sync, so the run is guaranteed to terminate. + let historyResetPending = false; + const historyReplacementBuf: TypesGen.ChatMessage[] = []; + const shouldApplyMessagePart = (): boolean => { const currentStatus = store.getSnapshot().chatStatus; return currentStatus !== "pending" && currentStatus !== "waiting"; @@ -393,9 +426,9 @@ export const useChatStore = ( }; // Discard buffered parts without applying them. Used when - // the stream is no longer active (pending, waiting, retry) + // the preview is reset or the stream is no longer active // so stale buffered parts are not applied after the - // status transition. + // boundary event. const discardBufferedParts = () => { partsBuf.length = 0; if (partsFlushTimer !== null) { @@ -428,6 +461,20 @@ export const useChatStore = ( const pendingMessages: TypesGen.ChatMessage[] = []; let needsStreamReset = false; + // Atomically swap in the buffered replacement history. Called + // when a boundary event signals the replacement run ended, so + // a run split across frames never renders a truncated + // conversation. + const commitHistoryReplacement = () => { + if (!historyResetPending) { + return; + } + historyResetPending = false; + const replacement = historyReplacementBuf.splice(0); + store.replaceMessages(replacement); + replaceCacheMessages(replacement); + }; + // Wrap all store mutations in a batch so subscribers // are notified exactly once at the end, not per event. store.batch(() => { @@ -436,6 +483,7 @@ export const useChatStore = ( if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { continue; } + commitHistoryReplacement(); if (!shouldApplyMessagePart()) { continue; } @@ -447,12 +495,45 @@ export const useChatStore = ( continue; } + if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { + const nextStatus = streamEvent.status?.status; + if (streamEvent.type === "status" && nextStatus) { + store.setSubagentStatusOverride(streamEvent.chat_id, nextStatus); + } + continue; + } + + if (streamEvent.type === "history_reset") { + discardBufferedParts(); + store.clearStreamState(); + // A newer reset supersedes any in-flight replacement + // run, so restart buffering instead of committing. + historyResetPending = true; + historyReplacementBuf.length = 0; + pendingMessages.length = 0; + needsStreamReset = false; + continue; + } + + // Any non-message event for this chat marks the end of + // a replacement run (the server emits replacement + // messages contiguously after history_reset). + if (streamEvent.type !== "message") { + commitHistoryReplacement(); + } + + if (streamEvent.type === "preview_reset") { + discardBufferedParts(); + store.clearStreamState(); + continue; + } + // Only flush buffered parts before events that // need them applied first. `message` events // commit durable state that must include all // stream parts. `error` events should surface // partial output. Other events (status, retry, - // queue_update) must NOT flush — status changes + // queue_update) must not flush. Status changes // need to be visible before parts so the // Thinking indicator can render, and retry // clears stream state which a flush would @@ -467,11 +548,12 @@ export const useChatStore = ( if (!message) { continue; } - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } store.clearRetryState(); - pendingMessages.push(message); + if (historyResetPending) { + historyReplacementBuf.push(message); + } else { + pendingMessages.push(message); + } if ( message.id !== undefined && (lastMessageIdRef.current === undefined || @@ -485,9 +567,6 @@ export const useChatStore = ( continue; } case "queue_update": - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } wsQueueUpdateReceivedRef.current = true; store.applyAuthoritativeQueuedMessages( streamEvent.queued_messages, @@ -500,14 +579,6 @@ export const useChatStore = ( continue; } - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - store.setSubagentStatusOverride( - streamEvent.chat_id, - nextStatus, - ); - continue; - } - wsStatusReceivedRef.current = true; store.clearRetryState(); store.setChatStatus(nextStatus); @@ -529,9 +600,6 @@ export const useChatStore = ( continue; } case "error": { - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } const reason = normalizeChatErrorPayload(streamEvent.error) ?? { kind: "generic", message: "Chat processing failed.", @@ -546,9 +614,6 @@ export const useChatStore = ( continue; } case "retry": { - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } const retry = streamEvent.retry; if (retry) { discardBufferedParts(); @@ -622,6 +687,10 @@ export const useChatStore = ( // apply stale parts from the old connection // into the fresh stream state. discardBufferedParts(); + // Drop any partial replacement run from the old + // socket; the new socket replays a fresh snapshot. + historyResetPending = false; + historyReplacementBuf.length = 0; }, onDisconnect( reconnectState: import("#/utils/reconnectingWebSocket").ReconnectSchedule, @@ -644,7 +713,14 @@ export const useChatStore = ( } activeChatIDRef.current = null; }; - }, [chatID, initialDataLoaded, queryClient, store, upsertCacheMessages]); + }, [ + chatID, + initialDataLoaded, + queryClient, + replaceCacheMessages, + store, + upsertCacheMessages, + ]); return { store, clearStreamError: () => { diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx index 481046dd28..17e3f18340 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx @@ -1,6 +1,7 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { expect, within } from "storybook/test"; import type * as TypesGen from "#/api/typesGenerated"; +import { ChatWorkspaceContext } from "../context/ChatWorkspaceContext"; import { createChatStore } from "./ChatConversation/chatStore"; import { FIXTURE_NOW } from "./ChatConversation/storyFixtures"; import { ChatPageTimeline } from "./ChatPageContent"; @@ -57,6 +58,36 @@ export const SpacerVisibleWhenNotStreaming: Story = { }, }; +export const DurableUnresolvedWorkspaceToolRuns: Story = { + render: () => { + const store = createChatStore(); + store.replaceMessages([ + buildMessage(1, "user", [{ type: "text", text: "Create a workspace" }]), + buildMessage(2, "assistant", [ + { + type: "tool-call", + tool_call_id: "create-workspace-call", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + ]), + ]); + store.setChatStatus("running"); + + return ( + + + + ); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText("Creating workspace…")).toBeInTheDocument(); + expect(canvas.queryByText("Created workspace")).toBeNull(); + expect(canvas.getByText("Loading build logs…")).toBeInTheDocument(); + }, +}; + export const HiddenAssistantPlaceholderDoesNotRender: Story = { render: () => { const store = createChatStore(); diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 84ca4fe606..4c0a88a0a3 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -34,6 +34,7 @@ import { import { LiveStreamTail } from "./ChatConversation/LiveStreamTail"; import { buildSubagentMaps, + getPendingToolCallIDs, parseMessagesWithMergedTools, } from "./ChatConversation/messageParsing"; import { useOnRenderProfiler } from "./ChatConversation/useOnRenderProfiler"; @@ -94,7 +95,10 @@ export const ChatPageTimeline: FC = ({ return message; }) .filter(isChatMessage); - const parsedMessages = parseMessagesWithMergedTools(messages); + const pendingToolCallIDs = getPendingToolCallIDs(messages, chatStatus); + const parsedMessages = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs, + }); const { titles: subagentTitles, variants: subagentVariants } = buildSubagentMaps(parsedMessages); const onRenderProfiler = useOnRenderProfiler(); diff --git a/site/src/pages/AgentsPage/components/ChatsSidebar/tree/statusConfig.ts b/site/src/pages/AgentsPage/components/ChatsSidebar/tree/statusConfig.ts index 534a12ca42..187e7f63ac 100644 --- a/site/src/pages/AgentsPage/components/ChatsSidebar/tree/statusConfig.ts +++ b/site/src/pages/AgentsPage/components/ChatsSidebar/tree/statusConfig.ts @@ -21,6 +21,7 @@ const statusConfig = { pending: { icon: Loader2Icon, className: "text-content-link animate-spin" }, running: { icon: Loader2Icon, className: "text-content-link animate-spin" }, paused: { icon: PauseIcon, className: "text-content-warning" }, + interrupting: { icon: PauseIcon, className: "text-content-warning" }, requires_action: { icon: PauseIcon, className: "text-content-warning" }, error: { icon: AlertTriangleIcon, className: "text-content-destructive" }, completed: { icon: CheckIcon, className: "text-content-secondary" },