diff --git a/cli/testdata/coder_server_--help.golden b/cli/testdata/coder_server_--help.golden index 8abc40867e264..9da5b1df7bd2c 100644 --- a/cli/testdata/coder_server_--help.golden +++ b/cli/testdata/coder_server_--help.golden @@ -281,6 +281,20 @@ Configure the background chat processing daemon. Force chat debug logging on for every chat, bypassing the runtime admin and user opt-in settings. + --chat-hook-enabled bool, $CODER_CHAT_HOOK_ENABLED (default: true) + Whether to dispatch chat agent lifecycle hooks when a hook URL is + configured. Requires the agent-lifecycle-hooks experiment. + + --chat-hook-secret string, $CODER_CHAT_HOOK_SECRET + Shared secret used to sign chat agent lifecycle hook JWTs. + + --chat-hook-timeout duration, $CODER_CHAT_HOOK_TIMEOUT (default: 1.5s) + Maximum time to wait for a chat agent lifecycle hook response. + + --chat-hook-url url, $CODER_CHAT_HOOK_URL + HTTPS URL to receive chat agent lifecycle hook events. Hooks are + disabled when unset. Requires the agent-lifecycle-hooks experiment. + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/cli/testdata/server-config.yaml.golden b/cli/testdata/server-config.yaml.golden index a8f3d90eb32e0..18d3a89cdc177 100644 --- a/cli/testdata/server-config.yaml.golden +++ b/cli/testdata/server-config.yaml.golden @@ -803,6 +803,17 @@ chat: # opt-in settings. # (default: false, type: bool) debugLoggingEnabled: false + # HTTPS URL to receive chat agent lifecycle hook events. Hooks are disabled when + # unset. Requires the agent-lifecycle-hooks experiment. + # (default: , type: url) + hookURL: + # Maximum time to wait for a chat agent lifecycle hook response. + # (default: 1.5s, type: duration) + hookTimeout: 1.5s + # Whether to dispatch chat agent lifecycle hooks when a hook URL is configured. + # Requires the agent-lifecycle-hooks experiment. + # (default: true, type: bool) + hookEnabled: true # Deprecated: AI Gateway routing is now the only routing path. Setting this value # has no effect. This option will be removed in a future release. # (default: true, type: bool) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 2c27365ac0797..d10354f679ea1 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -17078,6 +17078,18 @@ const docTemplate = `{ }, "debug_logging_enabled": { "type": "boolean" + }, + "hook_enabled": { + "type": "boolean" + }, + "hook_secret": { + "type": "string" + }, + "hook_timeout": { + "type": "integer" + }, + "hook_url": { + "$ref": "#/definitions/serpent.URL" } } }, @@ -17328,7 +17340,8 @@ const docTemplate = `{ "usage_limit", "missing_key", "provider_disabled", - "content_filter" + "content_filter", + "hook_dispatch_failed" ], "x-enum-varnames": [ "ChatErrorKindGeneric", @@ -17341,7 +17354,8 @@ const docTemplate = `{ "ChatErrorKindUsageLimit", "ChatErrorKindMissingKey", "ChatErrorKindProviderDisabled", - "ChatErrorKindContentFilter" + "ChatErrorKindContentFilter", + "ChatErrorKindHookDispatchFailed" ] }, "codersdk.ChatFileMetadata": { @@ -18450,9 +18464,19 @@ const docTemplate = `{ "codersdk.CreateChatMessageResponse": { "type": "object", "properties": { + "ended": { + "type": "boolean" + }, "message": { "$ref": "#/definitions/codersdk.ChatMessage" }, + "messages": { + "description": "Messages contains all user-visible messages inserted by an immediate send,\nin insertion order with the user's message last. Clients should upsert the\nfull batch because hooks may prepend notices. Empty for queued or ended sends.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.ChatMessage" + } + }, "queued": { "type": "boolean" }, @@ -19886,9 +19910,19 @@ const docTemplate = `{ "codersdk.EditChatMessageResponse": { "type": "object", "properties": { + "ended": { + "type": "boolean" + }, "message": { "$ref": "#/definitions/codersdk.ChatMessage" }, + "messages": { + "description": "Messages holds every user-visible message the edit inserted, in\ninsertion order with the replacement message last. Lifecycle\nhooks may prepend notices, so clients must upsert all of them\nrather than only Message. Empty for ended edits.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.ChatMessage" + } + }, "warnings": { "type": "array", "items": { @@ -19960,10 +19994,12 @@ const docTemplate = `{ "minimum-implicit-member", "ai-gateway-cost-control", "chat-advisor", - "chat-virtual-desktop" + "chat-virtual-desktop", + "agent-lifecycle-hooks" ], "x-enum-comments": { "ExperimentAIGatewayCostControl": "Enables AI Gateway cost control functionality.", + "ExperimentAgentLifecycleHooks": "Enables chat lifecycle hook webhooks for agent chats.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentChatAdvisor": "Enables the advisor tool for root agent chats.", "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", @@ -19988,7 +20024,8 @@ const docTemplate = `{ "Allows organizations to deviate from the default organization-member roles, in support of Gateway Accounts.", "Enables AI Gateway cost control functionality.", "Enables the advisor tool for root agent chats.", - "Enables virtual desktop and computer use provider for agents." + "Enables virtual desktop and computer use provider for agents.", + "Enables chat lifecycle hook webhooks for agent chats." ], "x-enum-varnames": [ "ExperimentExample", @@ -20002,7 +20039,8 @@ const docTemplate = `{ "ExperimentMinimumImplicitMember", "ExperimentAIGatewayCostControl", "ExperimentChatAdvisor", - "ExperimentChatVirtualDesktop" + "ExperimentChatVirtualDesktop", + "ExperimentAgentLifecycleHooks" ] }, "codersdk.ExternalAPIKeyScopes": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 7f87f2bad29c0..32d7edaa55661 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15347,6 +15347,18 @@ }, "debug_logging_enabled": { "type": "boolean" + }, + "hook_enabled": { + "type": "boolean" + }, + "hook_secret": { + "type": "string" + }, + "hook_timeout": { + "type": "integer" + }, + "hook_url": { + "$ref": "#/definitions/serpent.URL" } } }, @@ -15586,7 +15598,8 @@ "usage_limit", "missing_key", "provider_disabled", - "content_filter" + "content_filter", + "hook_dispatch_failed" ], "x-enum-varnames": [ "ChatErrorKindGeneric", @@ -15599,7 +15612,8 @@ "ChatErrorKindUsageLimit", "ChatErrorKindMissingKey", "ChatErrorKindProviderDisabled", - "ChatErrorKindContentFilter" + "ChatErrorKindContentFilter", + "ChatErrorKindHookDispatchFailed" ] }, "codersdk.ChatFileMetadata": { @@ -16670,9 +16684,19 @@ "codersdk.CreateChatMessageResponse": { "type": "object", "properties": { + "ended": { + "type": "boolean" + }, "message": { "$ref": "#/definitions/codersdk.ChatMessage" }, + "messages": { + "description": "Messages contains all user-visible messages inserted by an immediate send,\nin insertion order with the user's message last. Clients should upsert the\nfull batch because hooks may prepend notices. Empty for queued or ended sends.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.ChatMessage" + } + }, "queued": { "type": "boolean" }, @@ -18056,9 +18080,19 @@ "codersdk.EditChatMessageResponse": { "type": "object", "properties": { + "ended": { + "type": "boolean" + }, "message": { "$ref": "#/definitions/codersdk.ChatMessage" }, + "messages": { + "description": "Messages holds every user-visible message the edit inserted, in\ninsertion order with the replacement message last. Lifecycle\nhooks may prepend notices, so clients must upsert all of them\nrather than only Message. Empty for ended edits.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.ChatMessage" + } + }, "warnings": { "type": "array", "items": { @@ -18126,10 +18160,12 @@ "minimum-implicit-member", "ai-gateway-cost-control", "chat-advisor", - "chat-virtual-desktop" + "chat-virtual-desktop", + "agent-lifecycle-hooks" ], "x-enum-comments": { "ExperimentAIGatewayCostControl": "Enables AI Gateway cost control functionality.", + "ExperimentAgentLifecycleHooks": "Enables chat lifecycle hook webhooks for agent chats.", "ExperimentAutoFillParameters": "This should not be taken out of experiments until we have redesigned the feature.", "ExperimentChatAdvisor": "Enables the advisor tool for root agent chats.", "ExperimentChatVirtualDesktop": "Enables virtual desktop and computer use provider for agents.", @@ -18154,7 +18190,8 @@ "Allows organizations to deviate from the default organization-member roles, in support of Gateway Accounts.", "Enables AI Gateway cost control functionality.", "Enables the advisor tool for root agent chats.", - "Enables virtual desktop and computer use provider for agents." + "Enables virtual desktop and computer use provider for agents.", + "Enables chat lifecycle hook webhooks for agent chats." ], "x-enum-varnames": [ "ExperimentExample", @@ -18168,7 +18205,8 @@ "ExperimentMinimumImplicitMember", "ExperimentAIGatewayCostControl", "ExperimentChatAdvisor", - "ExperimentChatVirtualDesktop" + "ExperimentChatVirtualDesktop", + "ExperimentAgentLifecycleHooks" ] }, "codersdk.ExternalAPIKeyScopes": { diff --git a/coderd/coderd.go b/coderd/coderd.go index ab0332f821ac1..4af4a62bda177 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -100,6 +100,7 @@ import ( "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/mcpclient" + "github.com/coder/coder/v2/coderd/x/chathooks" "github.com/coder/coder/v2/coderd/x/gitsync" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpcsdk" @@ -878,6 +879,28 @@ func New(options *Options) *API { // the chat daemon stays nil and chat HTTP handlers return a // service-unavailable error with a clear remediation message. if options.DeploymentValues.AI.BridgeConfig.Enabled.Value() { + var hookDispatcher *chathooks.Dispatcher + chatConfig := options.DeploymentValues.AI.Chat + hooksConfigured := chatConfig.HookURL.String() != "" && chatConfig.HookEnabled.Value() + hooksExperimentEnabled := experiments.Enabled(codersdk.ExperimentAgentLifecycleHooks) + if hooksConfigured && !hooksExperimentEnabled { + options.Logger.Warn(ctx, "chat lifecycle hooks are configured but inactive; enable the agent-lifecycle-hooks experiment to activate them", + slog.F("experiment", codersdk.ExperimentAgentLifecycleHooks), + ) + } + if hooksConfigured && hooksExperimentEnabled { + hookDispatcher = chathooks.New( + options.Logger, + options.Database, + nil, + chatConfig.HookURL.String(), + chatConfig.HookSecret.Value(), + chatConfig.HookTimeout.Value(), + api.DeploymentID, + buildinfo.Version(), + options.PrometheusRegistry, + ) + } api.chatDaemon = chatd.New(options.Pubsub, chatd.Config{ Logger: options.Logger.Named("chatd"), Database: options.Database, @@ -897,6 +920,7 @@ func New(options *Options) *API { StartWorkspace: api.chatStartWorkspace, StopWorkspace: api.chatStopWorkspace, WebpushDispatcher: options.WebPushDispatcher, + HookDispatcher: hookDispatcher, UsageTracker: options.WorkspaceUsageTracker, PrometheusRegistry: options.PrometheusRegistry, OIDCTokenSource: oidcMCPSrc, diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 3161f3c3e9db0..19f4e63f02034 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2376,6 +2376,13 @@ func (q *querier) DeleteOldChatFiles(ctx context.Context, arg database.DeleteOld return q.db.DeleteOldChatFiles(ctx, arg) } +func (q *querier) DeleteOldChatHookDispatches(ctx context.Context, arg database.DeleteOldChatHookDispatchesParams) (int64, error) { + if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil { + return 0, err + } + return q.db.DeleteOldChatHookDispatches(ctx, arg) +} + func (q *querier) DeleteOldChats(ctx context.Context, arg database.DeleteOldChatsParams) (int64, error) { if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceSystem); err != nil { return 0, err @@ -2739,6 +2746,24 @@ func (q *querier) FetchVolumesResourceMonitorsUpdatedAfter(ctx context.Context, return q.db.FetchVolumesResourceMonitorsUpdatedAfter(ctx, updatedAt) } +func (q *querier) FinalizeChatHookDispatch(ctx context.Context, arg database.FinalizeChatHookDispatchParams) (database.ChatHookDispatch, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if errors.Is(err, sql.ErrNoRows) { + // CreateChat finalizes user_prompt_submit before inserting the chat. + if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).AnyOrganization()); err != nil { + return database.ChatHookDispatch{}, err + } + return q.db.FinalizeChatHookDispatch(ctx, arg) + } + if err != nil { + return database.ChatHookDispatch{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.ChatHookDispatch{}, err + } + return q.db.FinalizeChatHookDispatch(ctx, arg) +} + func (q *querier) FinalizeStaleChatDebugRows(ctx context.Context, updatedBefore database.FinalizeStaleChatDebugRowsParams) (database.FinalizeStaleChatDebugRowsRow, error) { // Background sweep operates across all chats. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { @@ -3187,6 +3212,14 @@ func (q *querier) GetChatDebugStepsByRunID(ctx context.Context, runID uuid.UUID) return q.db.GetChatDebugStepsByRunID(ctx, runID) } +func (q *querier) GetChatDescendantIDsByChatID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + // Callers authorize descendant mutations separately. + if _, err := q.GetChatByID(ctx, id); err != nil { + return nil, err + } + return q.db.GetChatDescendantIDsByChatID(ctx, id) +} + func (q *querier) GetChatDesktopEnabled(ctx context.Context) (bool, error) { // The desktop-enabled flag is a deployment-wide setting read by any // authenticated chat user and by chatd when deciding whether to expose @@ -3361,6 +3394,17 @@ func (q *querier) GetChatHeartbeat(ctx context.Context, arg database.GetChatHear return q.db.GetChatHeartbeat(ctx, arg) } +func (q *querier) GetChatHookDispatchDecision(ctx context.Context, arg database.GetChatHookDispatchDecisionParams) (database.ChatHookDispatch, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if err != nil { + return database.ChatHookDispatch{}, err + } + if err := q.authorizeContext(ctx, policy.ActionRead, chat); err != nil { + return database.ChatHookDispatch{}, err + } + return q.db.GetChatHookDispatchDecision(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 @@ -6012,6 +6056,23 @@ func (q *querier) InsertChatFile(ctx context.Context, arg database.InsertChatFil return insert(q.log, q.auth, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).InOrg(arg.OrganizationID), q.db.InsertChatFile)(ctx, arg) } +func (q *querier) InsertChatHookDispatch(ctx context.Context, arg database.InsertChatHookDispatchParams) (database.ChatHookDispatch, error) { + chat, err := q.db.GetChatByID(ctx, arg.ChatID) + if errors.Is(err, sql.ErrNoRows) { + if err := q.authorizeContext(ctx, policy.ActionCreate, rbac.ResourceChat.WithOwner(arg.OwnerID.String()).AnyOrganization()); err != nil { + return database.ChatHookDispatch{}, err + } + return q.db.InsertChatHookDispatch(ctx, arg) + } + if err != nil { + return database.ChatHookDispatch{}, err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return database.ChatHookDispatch{}, err + } + return q.db.InsertChatHookDispatch(ctx, arg) +} + func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { // Authorize create on the parent chat (using update permission). chat, err := q.db.GetChatByID(ctx, arg.ChatID) @@ -6800,6 +6861,17 @@ func (q *querier) ListChatContextResourcesByChatID(ctx context.Context, chatID u return q.db.ListChatContextResourcesByChatID(ctx, chatID) } +func (q *querier) ListChatHookDispatchesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatHookDispatch, error) { + chat, err := q.db.GetChatByID(ctx, chatID) + if err != nil { + return nil, err + } + if err := q.authorizeContext(ctx, policy.ActionRead, chat); err != nil { + return nil, err + } + return q.db.ListChatHookDispatchesByChatID(ctx, chatID) +} + func (q *querier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil { return nil, err @@ -6918,6 +6990,17 @@ func (q *querier) MarkAllInboxNotificationsAsRead(ctx context.Context, arg datab return q.db.MarkAllInboxNotificationsAsRead(ctx, arg) } +func (q *querier) MarkChatHookDispatchEffectsApplied(ctx context.Context, arg database.MarkChatHookDispatchEffectsAppliedParams) 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 + } + return q.db.MarkChatHookDispatchEffectsApplied(ctx, arg) +} + func (q *querier) MarkChatsContextDirtyByAgent(ctx context.Context, arg database.MarkChatsContextDirtyByAgentParams) ([]database.MarkChatsContextDirtyByAgentRow, error) { // System-level operation: the dirty fan-out runs across every active // chat for the agent in response to a context push. @@ -7335,6 +7418,17 @@ func (q *querier) UpdateChatHeartbeats(ctx context.Context, arg database.UpdateC return q.db.UpdateChatHeartbeats(ctx, arg) } +func (q *querier) UpdateChatHookAllowedTools(ctx context.Context, arg database.UpdateChatHookAllowedToolsParams) error { + chat, err := q.db.GetChatByID(ctx, arg.ID) + if err != nil { + return err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return err + } + return q.db.UpdateChatHookAllowedTools(ctx, arg) +} + func (q *querier) UpdateChatLabelsByID(ctx context.Context, arg database.UpdateChatLabelsByIDParams) (database.Chat, error) { chat, err := q.db.GetChatByID(ctx, arg.ID) if err != nil { @@ -7390,6 +7484,21 @@ func (q *querier) UpdateChatMCPServerIDs(ctx context.Context, arg database.Updat return q.db.UpdateChatMCPServerIDs(ctx, arg) } +func (q *querier) UpdateChatMessageContentByID(ctx context.Context, arg database.UpdateChatMessageContentByIDParams) error { + message, err := q.db.GetChatMessageByID(ctx, arg.ID) + if err != nil { + return err + } + chat, err := q.db.GetChatByID(ctx, message.ChatID) + if err != nil { + return err + } + if err := q.authorizeContext(ctx, policy.ActionUpdate, chat); err != nil { + return err + } + return q.db.UpdateChatMessageContentByID(ctx, arg) +} + func (q *querier) UpdateChatModelConfig(ctx context.Context, arg database.UpdateChatModelConfigParams) (database.ChatModelConfig, error) { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { return database.ChatModelConfig{}, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index c029cd3761146..895eb5820ad8e 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -853,6 +853,13 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().GetChatFamilyIDsByRootID(gomock.Any(), chat.ID).Return(ids, nil).AnyTimes() check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(ids) })) + s.Run("GetChatDescendantIDsByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + ids := []uuid.UUID{uuid.New()} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatDescendantIDsByChatID(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{}) @@ -1013,6 +1020,10 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().DeleteOldChatFiles(gomock.Any(), database.DeleteOldChatFilesParams{}).Return(int64(0), nil).AnyTimes() check.Args(database.DeleteOldChatFilesParams{}).Asserts(rbac.ResourceSystem, policy.ActionDelete) })) + s.Run("DeleteOldChatHookDispatches", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + dbm.EXPECT().DeleteOldChatHookDispatches(gomock.Any(), database.DeleteOldChatHookDispatchesParams{}).Return(int64(0), nil).AnyTimes() + check.Args(database.DeleteOldChatHookDispatchesParams{}).Asserts(rbac.ResourceSystem, policy.ActionDelete) + })) s.Run("DeleteOldChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { dbm.EXPECT().DeleteOldChats(gomock.Any(), database.DeleteOldChatsParams{}).Return(int64(0), nil).AnyTimes() check.Args(database.DeleteOldChatsParams{}).Asserts(rbac.ResourceSystem, policy.ActionDelete) @@ -1279,6 +1290,62 @@ func (s *MethodTestSuite) TestChats() { dbm.EXPECT().InsertChatMessages(gomock.Any(), arg).Return(msgs, nil).AnyTimes() check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(msgs) })) + s.Run("InsertChatHookDispatch", 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.InsertChatHookDispatchParams{ChatID: chat.ID}) + dispatch := testutil.Fake(s.T(), faker, database.ChatHookDispatch{ChatID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().InsertChatHookDispatch(gomock.Any(), arg).Return(dispatch, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(dispatch) + })) + s.Run("FinalizeChatHookDispatch", 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.FinalizeChatHookDispatchParams{ChatID: chat.ID}) + dispatch := testutil.Fake(s.T(), faker, database.ChatHookDispatch{ID: arg.ID, ChatID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().FinalizeChatHookDispatch(gomock.Any(), arg).Return(dispatch, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(dispatch) + })) + s.Run("UpdateChatHookAllowedTools", 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.UpdateChatHookAllowedToolsParams{ID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().UpdateChatHookAllowedTools(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns() + })) + s.Run("MarkChatHookDispatchEffectsApplied", 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.MarkChatHookDispatchEffectsAppliedParams{ChatID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().MarkChatHookDispatchEffectsApplied(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns() + })) + s.Run("ListChatHookDispatchesByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + dispatches := []database.ChatHookDispatch{testutil.Fake(s.T(), faker, database.ChatHookDispatch{ChatID: chat.ID})} + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().ListChatHookDispatchesByChatID(gomock.Any(), chat.ID).Return(dispatches, nil).AnyTimes() + check.Args(chat.ID).Asserts(chat, policy.ActionRead).Returns(dispatches) + })) + s.Run("GetChatHookDispatchDecision", 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.GetChatHookDispatchDecisionParams{ChatID: chat.ID}) + dispatch := testutil.Fake(s.T(), faker, database.ChatHookDispatch{ChatID: chat.ID}) + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().GetChatHookDispatchDecision(gomock.Any(), arg).Return(dispatch, nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionRead).Returns(dispatch) + })) + + s.Run("UpdateChatMessageContentByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + chat := testutil.Fake(s.T(), faker, database.Chat{}) + message := testutil.Fake(s.T(), faker, database.ChatMessage{ChatID: chat.ID}) + arg := testutil.Fake(s.T(), faker, database.UpdateChatMessageContentByIDParams{ID: message.ID}) + dbm.EXPECT().GetChatMessageByID(gomock.Any(), message.ID).Return(message, nil).AnyTimes() + dbm.EXPECT().GetChatByID(gomock.Any(), chat.ID).Return(chat, nil).AnyTimes() + dbm.EXPECT().UpdateChatMessageContentByID(gomock.Any(), arg).Return(nil).AnyTimes() + check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns() + })) + s.Run("InsertChatQueuedMessage", 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.InsertChatQueuedMessageParams{ChatID: chat.ID}) diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 9cdad7e8e82ec..00f6bac35db66 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -93,6 +93,8 @@ func Chat(t testing.TB, db database.Store, seed database.Chat) database.Chat { } chat, err := db.InsertChat(genCtx, database.InsertChatParams{ + ID: uuid.NullUUID{}, + HookAllowedTools: seed.HookAllowedTools, OrganizationID: takeFirst(seed.OrganizationID, uuid.New()), OwnerID: takeFirst(seed.OwnerID, uuid.New()), WorkspaceID: seed.WorkspaceID, @@ -125,6 +127,7 @@ func ChatMessage(t testing.TB, db database.Store, seed database.ChatMessage) dat msgs, err := db.InsertChatMessages(genCtx, database.InsertChatMessagesParams{ ChatID: seed.ChatID, + TurnID: []uuid.UUID{seed.TurnID.UUID}, CreatedBy: []uuid.UUID{seed.CreatedBy.UUID}, ModelConfigID: []uuid.UUID{seed.ModelConfigID.UUID}, ReasoningEffort: []string{string(seed.ReasoningEffort.ChatReasoningEffort)}, diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 57c742457f0e4..c5391a8f5a179 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -745,6 +745,14 @@ func (m queryMetricsStore) DeleteOldChatFiles(ctx context.Context, arg database. return r0, r1 } +func (m queryMetricsStore) DeleteOldChatHookDispatches(ctx context.Context, arg database.DeleteOldChatHookDispatchesParams) (int64, error) { + start := time.Now() + r0, r1 := m.s.DeleteOldChatHookDispatches(ctx, arg) + m.queryLatencies.WithLabelValues("DeleteOldChatHookDispatches").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteOldChatHookDispatches").Inc() + return r0, r1 +} + func (m queryMetricsStore) DeleteOldChats(ctx context.Context, arg database.DeleteOldChatsParams) (int64, error) { start := time.Now() r0, r1 := m.s.DeleteOldChats(ctx, arg) @@ -1057,6 +1065,14 @@ func (m queryMetricsStore) FetchVolumesResourceMonitorsUpdatedAfter(ctx context. return r0, r1 } +func (m queryMetricsStore) FinalizeChatHookDispatch(ctx context.Context, arg database.FinalizeChatHookDispatchParams) (database.ChatHookDispatch, error) { + start := time.Now() + r0, r1 := m.s.FinalizeChatHookDispatch(ctx, arg) + m.queryLatencies.WithLabelValues("FinalizeChatHookDispatch").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "FinalizeChatHookDispatch").Inc() + return r0, r1 +} + func (m queryMetricsStore) FinalizeStaleChatDebugRows(ctx context.Context, updatedBefore database.FinalizeStaleChatDebugRowsParams) (database.FinalizeStaleChatDebugRowsRow, error) { start := time.Now() r0, r1 := m.s.FinalizeStaleChatDebugRows(ctx, updatedBefore) @@ -1521,6 +1537,14 @@ func (m queryMetricsStore) GetChatDebugStepsByRunID(ctx context.Context, runID u return r0, r1 } +func (m queryMetricsStore) GetChatDescendantIDsByChatID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + start := time.Now() + r0, r1 := m.s.GetChatDescendantIDsByChatID(ctx, id) + m.queryLatencies.WithLabelValues("GetChatDescendantIDsByChatID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatDescendantIDsByChatID").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatDesktopEnabled(ctx context.Context) (bool, error) { start := time.Now() r0, r1 := m.s.GetChatDesktopEnabled(ctx) @@ -1625,6 +1649,14 @@ func (m queryMetricsStore) GetChatHeartbeat(ctx context.Context, arg database.Ge return r0, r1 } +func (m queryMetricsStore) GetChatHookDispatchDecision(ctx context.Context, arg database.GetChatHookDispatchDecisionParams) (database.ChatHookDispatch, error) { + start := time.Now() + r0, r1 := m.s.GetChatHookDispatchDecision(ctx, arg) + m.queryLatencies.WithLabelValues("GetChatHookDispatchDecision").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatHookDispatchDecision").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) { start := time.Now() r0, r1 := m.s.GetChatIncludeDefaultSystemPrompt(ctx) @@ -4121,6 +4153,14 @@ func (m queryMetricsStore) InsertChatFile(ctx context.Context, arg database.Inse return r0, r1 } +func (m queryMetricsStore) InsertChatHookDispatch(ctx context.Context, arg database.InsertChatHookDispatchParams) (database.ChatHookDispatch, error) { + start := time.Now() + r0, r1 := m.s.InsertChatHookDispatch(ctx, arg) + m.queryLatencies.WithLabelValues("InsertChatHookDispatch").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "InsertChatHookDispatch").Inc() + return r0, r1 +} + func (m queryMetricsStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { start := time.Now() r0, r1 := m.s.InsertChatMessages(ctx, arg) @@ -4769,6 +4809,14 @@ func (m queryMetricsStore) ListChatContextResourcesByChatID(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) ListChatHookDispatchesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatHookDispatch, error) { + start := time.Now() + r0, r1 := m.s.ListChatHookDispatchesByChatID(ctx, chatID) + m.queryLatencies.WithLabelValues("ListChatHookDispatchesByChatID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatHookDispatchesByChatID").Inc() + return r0, r1 +} + func (m queryMetricsStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) { start := time.Now() r0, r1 := m.s.ListChatUsageLimitGroupOverrides(ctx) @@ -4881,6 +4929,14 @@ func (m queryMetricsStore) MarkAllInboxNotificationsAsRead(ctx context.Context, return r0 } +func (m queryMetricsStore) MarkChatHookDispatchEffectsApplied(ctx context.Context, arg database.MarkChatHookDispatchEffectsAppliedParams) error { + start := time.Now() + r0 := m.s.MarkChatHookDispatchEffectsApplied(ctx, arg) + m.queryLatencies.WithLabelValues("MarkChatHookDispatchEffectsApplied").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "MarkChatHookDispatchEffectsApplied").Inc() + return r0 +} + func (m queryMetricsStore) MarkChatsContextDirtyByAgent(ctx context.Context, arg database.MarkChatsContextDirtyByAgentParams) ([]database.MarkChatsContextDirtyByAgentRow, error) { start := time.Now() r0, r1 := m.s.MarkChatsContextDirtyByAgent(ctx, arg) @@ -5217,6 +5273,14 @@ func (m queryMetricsStore) UpdateChatHeartbeats(ctx context.Context, arg databas return r0, r1 } +func (m queryMetricsStore) UpdateChatHookAllowedTools(ctx context.Context, arg database.UpdateChatHookAllowedToolsParams) error { + start := time.Now() + r0 := m.s.UpdateChatHookAllowedTools(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateChatHookAllowedTools").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatHookAllowedTools").Inc() + return r0 +} + func (m queryMetricsStore) UpdateChatLabelsByID(ctx context.Context, arg database.UpdateChatLabelsByIDParams) (database.Chat, error) { start := time.Now() r0, r1 := m.s.UpdateChatLabelsByID(ctx, arg) @@ -5257,6 +5321,14 @@ func (m queryMetricsStore) UpdateChatMCPServerIDs(ctx context.Context, arg datab return r0, r1 } +func (m queryMetricsStore) UpdateChatMessageContentByID(ctx context.Context, arg database.UpdateChatMessageContentByIDParams) error { + start := time.Now() + r0 := m.s.UpdateChatMessageContentByID(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateChatMessageContentByID").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateChatMessageContentByID").Inc() + return r0 +} + func (m queryMetricsStore) UpdateChatModelConfig(ctx context.Context, arg database.UpdateChatModelConfigParams) (database.ChatModelConfig, error) { start := time.Now() r0, r1 := m.s.UpdateChatModelConfig(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 9103654b6d3e2..b77d32c265859 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1251,6 +1251,21 @@ func (mr *MockStoreMockRecorder) DeleteOldChatFiles(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldChatFiles", reflect.TypeOf((*MockStore)(nil).DeleteOldChatFiles), ctx, arg) } +// DeleteOldChatHookDispatches mocks base method. +func (m *MockStore) DeleteOldChatHookDispatches(ctx context.Context, arg database.DeleteOldChatHookDispatchesParams) (int64, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "DeleteOldChatHookDispatches", ctx, arg) + ret0, _ := ret[0].(int64) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// DeleteOldChatHookDispatches indicates an expected call of DeleteOldChatHookDispatches. +func (mr *MockStoreMockRecorder) DeleteOldChatHookDispatches(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteOldChatHookDispatches", reflect.TypeOf((*MockStore)(nil).DeleteOldChatHookDispatches), ctx, arg) +} + // DeleteOldChats mocks base method. func (m *MockStore) DeleteOldChats(ctx context.Context, arg database.DeleteOldChatsParams) (int64, error) { m.ctrl.T.Helper() @@ -1813,6 +1828,21 @@ func (mr *MockStoreMockRecorder) FetchVolumesResourceMonitorsUpdatedAfter(ctx, u return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FetchVolumesResourceMonitorsUpdatedAfter", reflect.TypeOf((*MockStore)(nil).FetchVolumesResourceMonitorsUpdatedAfter), ctx, updatedAt) } +// FinalizeChatHookDispatch mocks base method. +func (m *MockStore) FinalizeChatHookDispatch(ctx context.Context, arg database.FinalizeChatHookDispatchParams) (database.ChatHookDispatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "FinalizeChatHookDispatch", ctx, arg) + ret0, _ := ret[0].(database.ChatHookDispatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// FinalizeChatHookDispatch indicates an expected call of FinalizeChatHookDispatch. +func (mr *MockStoreMockRecorder) FinalizeChatHookDispatch(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "FinalizeChatHookDispatch", reflect.TypeOf((*MockStore)(nil).FinalizeChatHookDispatch), ctx, arg) +} + // FinalizeStaleChatDebugRows mocks base method. func (m *MockStore) FinalizeStaleChatDebugRows(ctx context.Context, arg database.FinalizeStaleChatDebugRowsParams) (database.FinalizeStaleChatDebugRowsRow, error) { m.ctrl.T.Helper() @@ -2803,6 +2833,21 @@ func (mr *MockStoreMockRecorder) GetChatDebugStepsByRunID(ctx, runID any) *gomoc return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatDebugStepsByRunID", reflect.TypeOf((*MockStore)(nil).GetChatDebugStepsByRunID), ctx, runID) } +// GetChatDescendantIDsByChatID mocks base method. +func (m *MockStore) GetChatDescendantIDsByChatID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatDescendantIDsByChatID", ctx, id) + ret0, _ := ret[0].([]uuid.UUID) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatDescendantIDsByChatID indicates an expected call of GetChatDescendantIDsByChatID. +func (mr *MockStoreMockRecorder) GetChatDescendantIDsByChatID(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatDescendantIDsByChatID", reflect.TypeOf((*MockStore)(nil).GetChatDescendantIDsByChatID), ctx, id) +} + // GetChatDesktopEnabled mocks base method. func (m *MockStore) GetChatDesktopEnabled(ctx context.Context) (bool, error) { m.ctrl.T.Helper() @@ -2998,6 +3043,21 @@ func (mr *MockStoreMockRecorder) GetChatHeartbeat(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatHeartbeat", reflect.TypeOf((*MockStore)(nil).GetChatHeartbeat), ctx, arg) } +// GetChatHookDispatchDecision mocks base method. +func (m *MockStore) GetChatHookDispatchDecision(ctx context.Context, arg database.GetChatHookDispatchDecisionParams) (database.ChatHookDispatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetChatHookDispatchDecision", ctx, arg) + ret0, _ := ret[0].(database.ChatHookDispatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetChatHookDispatchDecision indicates an expected call of GetChatHookDispatchDecision. +func (mr *MockStoreMockRecorder) GetChatHookDispatchDecision(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatHookDispatchDecision", reflect.TypeOf((*MockStore)(nil).GetChatHookDispatchDecision), ctx, arg) +} + // GetChatIncludeDefaultSystemPrompt mocks base method. func (m *MockStore) GetChatIncludeDefaultSystemPrompt(ctx context.Context) (bool, error) { m.ctrl.T.Helper() @@ -7721,6 +7781,21 @@ func (mr *MockStoreMockRecorder) InsertChatFile(ctx, arg any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatFile", reflect.TypeOf((*MockStore)(nil).InsertChatFile), ctx, arg) } +// InsertChatHookDispatch mocks base method. +func (m *MockStore) InsertChatHookDispatch(ctx context.Context, arg database.InsertChatHookDispatchParams) (database.ChatHookDispatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "InsertChatHookDispatch", ctx, arg) + ret0, _ := ret[0].(database.ChatHookDispatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// InsertChatHookDispatch indicates an expected call of InsertChatHookDispatch. +func (mr *MockStoreMockRecorder) InsertChatHookDispatch(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertChatHookDispatch", reflect.TypeOf((*MockStore)(nil).InsertChatHookDispatch), ctx, arg) +} + // InsertChatMessages mocks base method. func (m *MockStore) InsertChatMessages(ctx context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { m.ctrl.T.Helper() @@ -8981,6 +9056,21 @@ func (mr *MockStoreMockRecorder) ListChatContextResourcesByChatID(ctx, chatID an return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatContextResourcesByChatID", reflect.TypeOf((*MockStore)(nil).ListChatContextResourcesByChatID), ctx, chatID) } +// ListChatHookDispatchesByChatID mocks base method. +func (m *MockStore) ListChatHookDispatchesByChatID(ctx context.Context, chatID uuid.UUID) ([]database.ChatHookDispatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListChatHookDispatchesByChatID", ctx, chatID) + ret0, _ := ret[0].([]database.ChatHookDispatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListChatHookDispatchesByChatID indicates an expected call of ListChatHookDispatchesByChatID. +func (mr *MockStoreMockRecorder) ListChatHookDispatchesByChatID(ctx, chatID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatHookDispatchesByChatID", reflect.TypeOf((*MockStore)(nil).ListChatHookDispatchesByChatID), ctx, chatID) +} + // ListChatUsageLimitGroupOverrides mocks base method. func (m *MockStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) { m.ctrl.T.Helper() @@ -9190,6 +9280,20 @@ func (mr *MockStoreMockRecorder) MarkAllInboxNotificationsAsRead(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkAllInboxNotificationsAsRead", reflect.TypeOf((*MockStore)(nil).MarkAllInboxNotificationsAsRead), ctx, arg) } +// MarkChatHookDispatchEffectsApplied mocks base method. +func (m *MockStore) MarkChatHookDispatchEffectsApplied(ctx context.Context, arg database.MarkChatHookDispatchEffectsAppliedParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "MarkChatHookDispatchEffectsApplied", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// MarkChatHookDispatchEffectsApplied indicates an expected call of MarkChatHookDispatchEffectsApplied. +func (mr *MockStoreMockRecorder) MarkChatHookDispatchEffectsApplied(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "MarkChatHookDispatchEffectsApplied", reflect.TypeOf((*MockStore)(nil).MarkChatHookDispatchEffectsApplied), ctx, arg) +} + // MarkChatsContextDirtyByAgent mocks base method. func (m *MockStore) MarkChatsContextDirtyByAgent(ctx context.Context, arg database.MarkChatsContextDirtyByAgentParams) ([]database.MarkChatsContextDirtyByAgentRow, error) { m.ctrl.T.Helper() @@ -9833,6 +9937,20 @@ func (mr *MockStoreMockRecorder) UpdateChatHeartbeats(ctx, arg any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHeartbeats", reflect.TypeOf((*MockStore)(nil).UpdateChatHeartbeats), ctx, arg) } +// UpdateChatHookAllowedTools mocks base method. +func (m *MockStore) UpdateChatHookAllowedTools(ctx context.Context, arg database.UpdateChatHookAllowedToolsParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateChatHookAllowedTools", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpdateChatHookAllowedTools indicates an expected call of UpdateChatHookAllowedTools. +func (mr *MockStoreMockRecorder) UpdateChatHookAllowedTools(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatHookAllowedTools", reflect.TypeOf((*MockStore)(nil).UpdateChatHookAllowedTools), ctx, arg) +} + // UpdateChatLabelsByID mocks base method. func (m *MockStore) UpdateChatLabelsByID(ctx context.Context, arg database.UpdateChatLabelsByIDParams) (database.Chat, error) { m.ctrl.T.Helper() @@ -9907,6 +10025,20 @@ func (mr *MockStoreMockRecorder) UpdateChatMCPServerIDs(ctx, arg any) *gomock.Ca return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatMCPServerIDs", reflect.TypeOf((*MockStore)(nil).UpdateChatMCPServerIDs), ctx, arg) } +// UpdateChatMessageContentByID mocks base method. +func (m *MockStore) UpdateChatMessageContentByID(ctx context.Context, arg database.UpdateChatMessageContentByIDParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateChatMessageContentByID", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpdateChatMessageContentByID indicates an expected call of UpdateChatMessageContentByID. +func (mr *MockStoreMockRecorder) UpdateChatMessageContentByID(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateChatMessageContentByID", reflect.TypeOf((*MockStore)(nil).UpdateChatMessageContentByID), ctx, arg) +} + // UpdateChatModelConfig mocks base method. func (m *MockStore) UpdateChatModelConfig(ctx context.Context, arg database.UpdateChatModelConfigParams) (database.ChatModelConfig, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dbpurge/dbpurge.go b/coderd/database/dbpurge/dbpurge.go index 7c284339d0e45..ac2e02d33153f 100644 --- a/coderd/database/dbpurge/dbpurge.go +++ b/coderd/database/dbpurge/dbpurge.go @@ -50,7 +50,9 @@ const ( chatFilesBatchSize = 1000 // Chat debug run deletions can cascade into steps with large JSONB // payloads, so they use the same conservative batch size. - chatDebugRunsBatchSize = 1000 + chatDebugRunsBatchSize = 1000 + chatHookDispatchRetention = 90 * 24 * time.Hour + chatHookDispatchesBatchSize = 10000 // Chat search tsvector backfill is capped at 5 batches of 10k // rows per tick. Benchmarks on a dogfood-class machine (EPYC 9454P) // with containerized Postgres were measured to take ~800ms per batch. @@ -314,6 +316,15 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. return xerrors.Errorf("failed to delete old workspace build orchestrations: %w", err) } + deleteChatHookDispatchesBefore := start.Add(-chatHookDispatchRetention) + purgedChatHookDispatches, err := tx.DeleteOldChatHookDispatches(ctx, database.DeleteOldChatHookDispatchesParams{ + BeforeTime: deleteChatHookDispatchesBefore, + LimitCount: chatHookDispatchesBatchSize, + }) + if err != nil { + return xerrors.Errorf("failed to delete old chat hook dispatches: %w", err) + } + var purgedChats, purgedChatFiles, purgedChatDebugRuns int64 if purgeChats { purgedChats, purgedChatFiles, err = i.purgeChatsInTx(ctx, tx, start, chatRetentionDays) @@ -338,7 +349,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. // Backfill search_tsv tsvector on chat_messages in batches. Doing this here because it's // potentially too much for a regular migration, especially on larger deployments: // - Each row with search_tsv = NULL is present in idx_chat_messages_search_tsv_pending. - // - Content of chat_messages is not changed after insert. + // - Indexed rows refresh search_tsv during content rewrites. // - Rows that are soft-deleted are no longer part of the index. // NOTE: This should not remain in dbpurge and should be adjusted when the "DBOps" gets // implemented. @@ -363,6 +374,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. slog.F("boundary_logs", purgedBoundaryLogs), slog.F("boundary_sessions", purgedBoundarySessions), slog.F("workspace_build_orchestrations", purgedWorkspaceBuildOrchestrations), + slog.F("chat_hook_dispatches", purgedChatHookDispatches), slog.F("chats", purgedChats), slog.F("chat_files", purgedChatFiles), slog.F("chat_debug_runs", purgedChatDebugRuns), @@ -379,6 +391,7 @@ func (i *instance) purgeTick(ctx context.Context, db database.Store, start time. i.recordsPurged.WithLabelValues("boundary_logs").Add(float64(purgedBoundaryLogs)) i.recordsPurged.WithLabelValues("boundary_sessions").Add(float64(purgedBoundarySessions)) i.recordsPurged.WithLabelValues("workspace_build_orchestrations").Add(float64(purgedWorkspaceBuildOrchestrations)) + i.recordsPurged.WithLabelValues("chat_hook_dispatches").Add(float64(purgedChatHookDispatches)) i.recordsPurged.WithLabelValues("chats").Add(float64(purgedChats)) i.recordsPurged.WithLabelValues("chat_debug_runs").Add(float64(purgedChatDebugRuns)) i.recordsPurged.WithLabelValues("chat_files").Add(float64(purgedChatFiles)) diff --git a/coderd/database/dbpurge/dbpurge_test.go b/coderd/database/dbpurge/dbpurge_test.go index 25e780ec3d2e3..bc88afe576f67 100644 --- a/coderd/database/dbpurge/dbpurge_test.go +++ b/coderd/database/dbpurge/dbpurge_test.go @@ -254,6 +254,7 @@ func TestMetrics(t *testing.T) { mDB.EXPECT().ExpirePrebuildsAPIKeys(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mDB.EXPECT().DeleteOldChatHookDispatches(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() mDB.EXPECT().DeleteOldChatDebugRuns(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatDebugRunsParams{})).Return(int64(0), nil).MinTimes(1) @@ -306,6 +307,7 @@ func TestMetrics(t *testing.T) { mDB.EXPECT().ExpirePrebuildsAPIKeys(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().DeleteOldTelemetryLocks(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().DeleteOldWorkspaceBuildOrchestrations(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() + mDB.EXPECT().DeleteOldChatHookDispatches(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() mDB.EXPECT().DeleteOldAuditLogConnectionEvents(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() mDB.EXPECT().BackfillChatMessagesSearchTsv(gomock.Any(), gomock.Any()).Return(int64(0), nil).AnyTimes() mDB.EXPECT().DeleteOldChats(gomock.Any(), gomock.AssignableToTypeOf(database.DeleteOldChatsParams{})).Return(int64(0), nil).MinTimes(1) @@ -2211,6 +2213,51 @@ func ptr[T any](v T) *T { return &v } +//nolint:paralleltest // It uses LockIDDBPurge. +func TestPurgeChatHookDispatches(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) + clk := quartz.NewMock(t) + clk.Set(now).MustWait(ctx) + db, _ := dbtestutil.NewDB(t) + reg := prometheus.NewRegistry() + chatID := uuid.New() + ownerID := uuid.New() + + insertDispatch := func(startedAt time.Time) database.ChatHookDispatch { + dispatch, err := db.InsertChatHookDispatch(ctx, database.InsertChatHookDispatchParams{ + ID: uuid.New(), + ChatID: chatID, + Event: "stop", + TurnID: uuid.NullUUID{}, + ToolUseID: sql.NullString{}, + OwnerID: ownerID, + WorkspaceID: uuid.NullUUID{}, + StartedAt: startedAt, + }) + require.NoError(t, err) + return dispatch + } + _ = insertDispatch(now.Add(-90*24*time.Hour - time.Second)) + recentDispatch := insertDispatch(now.Add(-89 * 24 * time.Hour)) + + done := awaitDoTick(ctx, t, clk) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, reg, dbpurge.WithClock(clk)) + defer closer.Close() + testutil.TryReceive(ctx, t, done) + + rows, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, recentDispatch.ID, rows[0].ID) + + purged := promhelp.CounterValue(t, reg, "coderd_dbpurge_records_purged_total", prometheus.Labels{ + "record_type": "chat_hook_dispatches", + }) + require.EqualValues(t, 1, purged) +} + //nolint:paralleltest // It uses LockIDDBPurge. func TestPurgeChatDebugRuns(t *testing.T) { now := time.Date(2025, 6, 15, 12, 0, 0, 0, time.UTC) @@ -2531,6 +2578,23 @@ func TestDeleteOldChatFiles(t *testing.T) { // Active chat — should be retained. activeChat := createChat(ctx, t, db, rawDB, deps.user.ID, deps.org.ID, deps.modelConfig.ID, false, now) + // Fresh dispatches prove chat purging is independent of dispatch retention. + insertDispatch := func(chatID uuid.UUID) database.ChatHookDispatch { + dispatch, err := db.InsertChatHookDispatch(ctx, database.InsertChatHookDispatchParams{ + ID: uuid.New(), + ChatID: chatID, + Event: "stop", + OwnerID: deps.user.ID, + StartedAt: now.Add(-time.Hour), + }) + require.NoError(t, err) + return dispatch + } + _ = insertDispatch(oldChat.ID) + recentChatDispatch := insertDispatch(recentChat.ID) + // Denied creates can leave dispatches without a chat; only retention removes them. + orphanDispatch := insertDispatch(uuid.New()) + done := awaitDoTick(ctx, t, clk) closer := dbpurge.New(ctx, logger, db, &codersdk.DeploymentValues{}, prometheus.NewRegistry(), dbpurge.WithClock(clk)) defer closer.Close() @@ -2540,6 +2604,19 @@ func TestDeleteOldChatFiles(t *testing.T) { _, err = db.GetChatByID(ctx, oldChat.ID) require.ErrorIs(t, err, sql.ErrNoRows, "old archived chat should be deleted") + oldChatDispatches, err := db.ListChatHookDispatchesByChatID(ctx, oldChat.ID) + require.NoError(t, err) + require.Empty(t, oldChatDispatches, "dispatches should be deleted with their purged chat") + + recentChatDispatches, err := db.ListChatHookDispatchesByChatID(ctx, recentChat.ID) + require.NoError(t, err) + require.Len(t, recentChatDispatches, 1) + require.Equal(t, recentChatDispatch.ID, recentChatDispatches[0].ID) + + orphanDispatches, err := db.ListChatHookDispatchesByChatID(ctx, orphanDispatch.ChatID) + require.NoError(t, err) + require.Len(t, orphanDispatches, 1, "pre-create dispatch rows are only removed by the time sweep") + // Its messages should be gone too (CASCADE). msgs, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ ChatID: oldChat.ID, diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index b7cd44e1dbf05..59daf98c813b9 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1941,6 +1941,33 @@ CREATE UNLOGGED TABLE chat_heartbeats ( 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_hook_dispatches ( + id uuid NOT NULL, + chat_id uuid NOT NULL, + event text NOT NULL, + turn_id uuid, + tool_use_id text, + owner_id uuid NOT NULL, + workspace_id uuid, + started_at timestamp with time zone NOT NULL, + finished_at timestamp with time zone, + result text DEFAULT 'pending'::text NOT NULL, + http_status integer, + decision text, + input_override jsonb, + original_input jsonb, + model_context text, + user_message text, + allowed_tools jsonb, + end_chat boolean, + error text, + decision_reason text, + effects_applied_at timestamp with time zone, + tool_name text +); + +COMMENT ON TABLE chat_hook_dispatches IS 'Lifecycle hook attempts keyed by dispatch_id (JWT jti).'; + CREATE TABLE chat_messages ( id bigint NOT NULL, chat_id uuid NOT NULL, @@ -1965,7 +1992,8 @@ CREATE TABLE chat_messages ( provider_response_id text, revision bigint NOT NULL, reasoning_effort chat_reasoning_effort, - search_tsv tsvector + search_tsv tsvector, + turn_id uuid ); COMMENT ON COLUMN chat_messages.reasoning_effort IS 'Stores the selected effort for the turn triggered by this message.'; @@ -2017,11 +2045,16 @@ CREATE TABLE chat_queued_messages ( model_config_id uuid, "position" bigint DEFAULT nextval('chat_queued_messages_position_seq'::regclass) NOT NULL, created_by uuid NOT NULL, - reasoning_effort chat_reasoning_effort + reasoning_effort chat_reasoning_effort, + turn_id uuid, + hook_prefix jsonb, + hook_allowed_tools jsonb ); COMMENT ON COLUMN chat_queued_messages.reasoning_effort IS 'Stores the selected effort until the queued row is promoted.'; +COMMENT ON COLUMN chat_queued_messages.hook_allowed_tools IS 'Queued prompt hook policy; NULL means no policy.'; + CREATE SEQUENCE chat_queued_messages_id_seq START WITH 1 INCREMENT BY 1 @@ -2097,6 +2130,7 @@ CREATE TABLE chats ( context_error text DEFAULT ''::text NOT NULL, last_reasoning_effort chat_reasoning_effort, compaction_requested_at timestamp with time zone, + hook_allowed_tools jsonb, 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))), @@ -2122,6 +2156,8 @@ COMMENT ON COLUMN chats.last_reasoning_effort IS 'Stores the most recent message COMMENT ON COLUMN chats.compaction_requested_at IS 'Set when the chat owner manually requests a context compaction. One-shot signal: consumed by the compaction commit and cleared whenever the chat leaves running.'; +COMMENT ON COLUMN chats.hook_allowed_tools IS 'Hook-enforced tool names; NULL means unrestricted. Later policies only narrow.'; + CREATE TABLE users ( id uuid NOT NULL, email text NOT NULL, @@ -2218,7 +2254,8 @@ CREATE VIEW chats_expanded AS c.context_dirty_since, c.context_dirty_resources, c.context_error, - c.compaction_requested_at + c.compaction_requested_at, + c.hook_allowed_tools 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))); @@ -4295,6 +4332,9 @@ ALTER TABLE ONLY chat_files ALTER TABLE ONLY chat_heartbeats ADD CONSTRAINT chat_heartbeats_pkey PRIMARY KEY (chat_id, runner_id); +ALTER TABLE ONLY chat_hook_dispatches + ADD CONSTRAINT chat_hook_dispatches_pkey PRIMARY KEY (id); + ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_pkey PRIMARY KEY (id); @@ -4754,6 +4794,10 @@ CREATE INDEX idx_chat_files_org ON chat_files USING btree (organization_id); CREATE INDEX idx_chat_files_owner ON chat_files USING btree (owner_id); +CREATE INDEX idx_chat_hook_dispatches_chat_id ON chat_hook_dispatches USING btree (chat_id); + +CREATE INDEX idx_chat_hook_dispatches_started_at ON chat_hook_dispatches USING btree (started_at); + CREATE INDEX idx_chat_messages_chat ON chat_messages USING btree (chat_id); CREATE INDEX idx_chat_messages_chat_created ON chat_messages USING btree (chat_id, created_at); diff --git a/coderd/database/migrations/000551_chat_hook_dispatches.down.sql b/coderd/database/migrations/000551_chat_hook_dispatches.down.sql new file mode 100644 index 0000000000000..25deb6ab0c2da --- /dev/null +++ b/coderd/database/migrations/000551_chat_hook_dispatches.down.sql @@ -0,0 +1 @@ +DROP TABLE chat_hook_dispatches; diff --git a/coderd/database/migrations/000551_chat_hook_dispatches.up.sql b/coderd/database/migrations/000551_chat_hook_dispatches.up.sql new file mode 100644 index 0000000000000..d430c38632db5 --- /dev/null +++ b/coderd/database/migrations/000551_chat_hook_dispatches.up.sql @@ -0,0 +1,31 @@ +-- Create-time prompt dispatches precede the chat row, so chat_id has no +-- foreign key. +CREATE TABLE chat_hook_dispatches ( + id uuid PRIMARY KEY, + chat_id uuid NOT NULL, + event text NOT NULL, + turn_id uuid, + tool_use_id text, + owner_id uuid NOT NULL, + workspace_id uuid, + started_at timestamptz NOT NULL, + finished_at timestamptz, + result text NOT NULL DEFAULT 'pending', + http_status integer, + decision text, + input_override jsonb, + original_input jsonb, + model_context text, + user_message text, + allowed_tools jsonb, + end_chat boolean, + error text, + decision_reason text, + effects_applied_at timestamptz, + tool_name text +); + +COMMENT ON TABLE chat_hook_dispatches IS 'Lifecycle hook attempts keyed by dispatch_id (JWT jti).'; + +CREATE INDEX idx_chat_hook_dispatches_chat_id ON chat_hook_dispatches (chat_id); +CREATE INDEX idx_chat_hook_dispatches_started_at ON chat_hook_dispatches (started_at); diff --git a/coderd/database/migrations/000552_chat_hook_columns.down.sql b/coderd/database/migrations/000552_chat_hook_columns.down.sql new file mode 100644 index 0000000000000..e0566def5d18f --- /dev/null +++ b/coderd/database/migrations/000552_chat_hook_columns.down.sql @@ -0,0 +1,58 @@ +DROP VIEW IF EXISTS chats_expanded; + +ALTER TABLE chat_queued_messages + DROP COLUMN hook_allowed_tools, + DROP COLUMN hook_prefix, + DROP COLUMN turn_id; +ALTER TABLE chat_messages DROP COLUMN turn_id; +ALTER TABLE chats DROP COLUMN hook_allowed_tools; + +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.last_reasoning_effort, + 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.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, + c.context_aggregate_hash, + c.context_dirty_since, + c.context_dirty_resources, + c.context_error, + c.compaction_requested_at + 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/000552_chat_hook_columns.up.sql b/coderd/database/migrations/000552_chat_hook_columns.up.sql new file mode 100644 index 0000000000000..16dcea2352f75 --- /dev/null +++ b/coderd/database/migrations/000552_chat_hook_columns.up.sql @@ -0,0 +1,63 @@ +-- Recreate chats_expanded because adding a chats column changes its row type. +DROP VIEW IF EXISTS chats_expanded; + +ALTER TABLE chats ADD COLUMN hook_allowed_tools jsonb; +ALTER TABLE chat_messages ADD COLUMN turn_id uuid; +ALTER TABLE chat_queued_messages + ADD COLUMN turn_id uuid, + ADD COLUMN hook_prefix jsonb, + ADD COLUMN hook_allowed_tools jsonb; + +COMMENT ON COLUMN chats.hook_allowed_tools IS 'Hook-enforced tool names; NULL means unrestricted. Later policies only narrow.'; +COMMENT ON COLUMN chat_queued_messages.hook_allowed_tools IS 'Queued prompt hook policy; NULL means no policy.'; + +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.last_reasoning_effort, + 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.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, + c.context_aggregate_hash, + c.context_dirty_since, + c.context_dirty_resources, + c.context_error, + c.compaction_requested_at, + c.hook_allowed_tools + 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/000551_chat_hook_dispatches.up.sql b/coderd/database/migrations/testdata/fixtures/000551_chat_hook_dispatches.up.sql new file mode 100644 index 0000000000000..38c44804ee28b --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000551_chat_hook_dispatches.up.sql @@ -0,0 +1,28 @@ +INSERT INTO chat_hook_dispatches ( + id, + chat_id, + event, + turn_id, + owner_id, + started_at, + finished_at, + result, + http_status, + decision, + effects_applied_at +) +SELECT + '10000000-0000-0000-0000-00000000d15a'::uuid, + chats.id, + 'pre_tool_use', + '10000000-0000-0000-0000-00000000722d'::uuid, + chats.owner_id, + '2024-01-01 00:00:00+00', + '2024-01-01 00:00:01+00', + 'ok', + 200, + 'allow', + '2024-01-01 00:00:01+00' +FROM chats +ORDER BY created_at, id +LIMIT 1; diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index d0503fc1d9770..ad8240e62d35e 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -843,6 +843,7 @@ func (q *sqlQuerier) GetAuthorizedChats(ctx context.Context, arg GetChatsParams, &i.Chat.ContextDirtyResources, &i.Chat.ContextError, &i.Chat.CompactionRequestedAt, + &i.Chat.HookAllowedTools, &i.HasUnread); err != nil { return nil, err } @@ -922,7 +923,8 @@ func (q *sqlQuerier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID &i.ContextDirtySince, &i.ContextDirtyResources, &i.ContextError, - &i.CompactionRequestedAt); err != nil { + &i.CompactionRequestedAt, + &i.HookAllowedTools); err != nil { return nil, err } items = append(items, i) diff --git a/coderd/database/models.go b/coderd/database/models.go index 7a96121b0214b..5c883c0892617 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4987,6 +4987,7 @@ type Chat struct { ContextDirtyResources pqtype.NullRawMessage `db:"context_dirty_resources" json:"context_dirty_resources"` ContextError string `db:"context_error" json:"context_error"` CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` } // Per-chat pinned copy of the agent context resources a chat is hydrated against. Copied from workspace_agent_context_resources at chat hydration and context refresh; survives agent replacement and workspace rebuilds. @@ -5098,6 +5099,32 @@ type ChatHeartbeat struct { HeartbeatAt time.Time `db:"heartbeat_at" json:"heartbeat_at"` } +// Lifecycle hook attempts keyed by dispatch_id (JWT jti). +type ChatHookDispatch struct { + ID uuid.UUID `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Event string `db:"event" json:"event"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` + ToolUseID sql.NullString `db:"tool_use_id" json:"tool_use_id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + StartedAt time.Time `db:"started_at" json:"started_at"` + FinishedAt sql.NullTime `db:"finished_at" json:"finished_at"` + Result string `db:"result" json:"result"` + HttpStatus sql.NullInt32 `db:"http_status" json:"http_status"` + Decision sql.NullString `db:"decision" json:"decision"` + InputOverride pqtype.NullRawMessage `db:"input_override" json:"input_override"` + OriginalInput pqtype.NullRawMessage `db:"original_input" json:"original_input"` + ModelContext sql.NullString `db:"model_context" json:"model_context"` + UserMessage sql.NullString `db:"user_message" json:"user_message"` + AllowedTools pqtype.NullRawMessage `db:"allowed_tools" json:"allowed_tools"` + EndChat sql.NullBool `db:"end_chat" json:"end_chat"` + Error sql.NullString `db:"error" json:"error"` + DecisionReason sql.NullString `db:"decision_reason" json:"decision_reason"` + EffectsAppliedAt sql.NullTime `db:"effects_applied_at" json:"effects_applied_at"` + ToolName sql.NullString `db:"tool_name" json:"tool_name"` +} + type ChatMessage struct { ID int64 `db:"id" json:"id"` ChatID uuid.UUID `db:"chat_id" json:"chat_id"` @@ -5124,7 +5151,8 @@ type ChatMessage struct { // Stores the selected effort for the turn triggered by this message. ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` // Used for full text search. NULL initially, populated async via background job. - SearchTsv interface{} `db:"search_tsv" json:"search_tsv"` + SearchTsv interface{} `db:"search_tsv" json:"search_tsv"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` } type ChatModelConfig struct { @@ -5155,6 +5183,10 @@ type ChatQueuedMessage struct { CreatedBy uuid.UUID `db:"created_by" json:"created_by"` // Stores the selected effort until the queued row is promoted. ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` + HookPrefix pqtype.NullRawMessage `db:"hook_prefix" json:"hook_prefix"` + // Queued prompt hook policy; NULL means no policy. + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` } type ChatTable struct { @@ -5210,6 +5242,8 @@ type ChatTable struct { LastReasoningEffort NullChatReasoningEffort `db:"last_reasoning_effort" json:"last_reasoning_effort"` // Set when the chat owner manually requests a context compaction. One-shot signal: consumed by the compaction commit and cleared whenever the chat leaves running. CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` + // Hook-enforced tool names; NULL means unrestricted. Later policies only narrow. + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` } type ChatUsageLimitConfig struct { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 5c555cdd86f48..37cdcb6bd6ec8 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -200,9 +200,12 @@ type sqlcQuerier interface { // 2. Files whose every referencing chat has been archived for longer // than the retention period. DeleteOldChatFiles(ctx context.Context, arg DeleteOldChatFilesParams) (int64, error) + DeleteOldChatHookDispatches(ctx context.Context, arg DeleteOldChatHookDispatchesParams) (int64, error) // Deletes chats that have been archived for longer than the given // threshold. Active (non-archived) chats are never deleted. // All chat-scoped child tables are removed via ON DELETE CASCADE. + // Dispatches have no chat FK because they can precede chat creation. + // Delete them explicitly so their payloads respect chat retention. // Parent/root references on child chats are SET NULL. DeleteOldChats(ctx context.Context, arg DeleteOldChatsParams) (int64, error) DeleteOldConnectionLogs(ctx context.Context, arg DeleteOldConnectionLogsParams) (int64, error) @@ -272,6 +275,7 @@ type sqlcQuerier interface { FetchNewMessageMetadata(ctx context.Context, arg FetchNewMessageMetadataParams) (FetchNewMessageMetadataRow, error) FetchVolumesResourceMonitorsByAgentID(ctx context.Context, agentID uuid.UUID) ([]WorkspaceAgentVolumeResourceMonitor, error) FetchVolumesResourceMonitorsUpdatedAfter(ctx context.Context, updatedAt time.Time) ([]WorkspaceAgentVolumeResourceMonitor, error) + FinalizeChatHookDispatch(ctx context.Context, arg FinalizeChatHookDispatchParams) (ChatHookDispatch, error) // Marks orphaned in-progress rows as interrupted so they do not stay // in a non-terminal state forever. The NOT IN list must match the // terminal statuses defined by ChatDebugStatus in codersdk/chats.go. @@ -408,6 +412,7 @@ type sqlcQuerier interface { // Callers must supply an explicit limit to avoid unbounded result sets. GetChatDebugRunsByChatID(ctx context.Context, arg GetChatDebugRunsByChatIDParams) ([]ChatDebugRun, error) GetChatDebugStepsByRunID(ctx context.Context, runID uuid.UUID) ([]ChatDebugStep, error) + GetChatDescendantIDsByChatID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) GetChatDesktopEnabled(ctx context.Context) (bool, error) GetChatDiffStatusByChatID(ctx context.Context, chatID uuid.UUID) (ChatDiffStatus, error) // Returns aggregate PR counts across all agent chats for telemetry. @@ -436,6 +441,7 @@ type sqlcQuerier interface { GetChatGatewayAPIKey(ctx context.Context, arg GetChatGatewayAPIKeyParams) (APIKey, error) GetChatGeneralModelOverride(ctx context.Context) (string, error) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatParams) (ChatHeartbeat, error) + GetChatHookDispatchDecision(ctx context.Context, arg GetChatHookDispatchDecisionParams) (ChatHookDispatch, 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; @@ -1074,6 +1080,7 @@ type sqlcQuerier interface { // with concurrent FinalizeStale under READ COMMITTED isolation. InsertChatDebugStep(ctx context.Context, arg InsertChatDebugStepParams) (ChatDebugStep, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) + InsertChatHookDispatch(ctx context.Context, arg InsertChatHookDispatchParams) (ChatHookDispatch, 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, @@ -1206,6 +1213,7 @@ type sqlcQuerier interface { // Lists a chat's pinned context resources, ordered deterministically by // source. ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatContextResource, error) + ListChatHookDispatchesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatHookDispatch, error) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error) ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) @@ -1229,6 +1237,7 @@ type sqlcQuerier interface { // 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 + MarkChatHookDispatchEffectsApplied(ctx context.Context, arg MarkChatHookDispatchEffectsAppliedParams) error // Flips active, already-hydrated chats for an agent to dirty when the // agent's latest snapshot hash differs from the chat's pinned hash. The // pinned hash is intentionally left untouched; the refresh endpoint @@ -1400,6 +1409,7 @@ type sqlcQuerier interface { // worker. Returns the IDs that were actually updated so the // caller can detect stolen or completed chats via set-difference. UpdateChatHeartbeats(ctx context.Context, arg UpdateChatHeartbeatsParams) ([]uuid.UUID, error) + UpdateChatHookAllowedTools(ctx context.Context, arg UpdateChatHookAllowedToolsParams) error UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLabelsByIDParams) (Chat, error) UpdateChatLastModelConfigByID(ctx context.Context, arg UpdateChatLastModelConfigByIDParams) (Chat, error) // Updates the last read message ID for a chat. This is used to track @@ -1414,6 +1424,9 @@ type sqlcQuerier interface { // 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) + // Preserve NULL as the backfill marker; otherwise refresh search_tsv + // from the new content. + UpdateChatMessageContentByID(ctx context.Context, arg UpdateChatMessageContentByIDParams) error UpdateChatModelConfig(ctx context.Context, arg UpdateChatModelConfigParams) (ChatModelConfig, error) UpdateChatPinOrder(ctx context.Context, arg UpdateChatPinOrderParams) error UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatPlanModeByIDParams) (Chat, error) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 53a477e9e6624..3022b8acce825 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -16551,6 +16551,109 @@ func TestGetChatsSearch(t *testing.T) { } } +func TestUpdateChatMessageContentByIDRefreshesSearchTsv(t *testing.T) { + t.Parallel() + + store, _, sqlDB := dbtestutil.NewDBWithSQLDB(t) + ctx := context.Background() + + org := dbgen.Organization(t, store, database.Organization{}) + user := dbgen.User(t, store, database.User{}) + dbgen.OrganizationMember(t, store, database.OrganizationMember{UserID: user.ID, OrganizationID: org.ID}) + + provider := dbgen.AIProviderWithOptionalKey(t, store, database.AIProvider{ + Type: database.AIProviderTypeOpenai, + }, "test-key") + + modelCfg, err := store.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "test-model-" + uuid.NewString(), + DisplayName: "Test Model", + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Enabled: true, + IsDefault: true, + ContextLimit: 128000, + CompressionThreshold: 80, + Options: json.RawMessage(`{}`), + }) + require.NoError(t, err) + + chat, err := store.InsertChat(ctx, database.InsertChatParams{ + OrganizationID: org.ID, + Status: database.ChatStatusWaiting, + ClientType: database.ChatClientTypeUi, + OwnerID: user.ID, + LastModelConfigID: modelCfg.ID, + Title: "content rewrite", + }) + require.NoError(t, err) + + insertMsg := func(text string) database.ChatMessage { + t.Helper() + msgs, err := store.InsertChatMessages(ctx, database.InsertChatMessagesParams{ + ChatID: chat.ID, + CreatedBy: []uuid.UUID{user.ID}, + ModelConfigID: []uuid.UUID{modelCfg.ID}, + Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, + Content: []string{`[{"type":"text","text":` + strconv.Quote(text) + `}]`}, + ContentVersion: []int16{1}, + 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}, + }) + require.NoError(t, err) + require.Len(t, msgs, 1) + return msgs[0] + } + + searchTsv := func(id int64) sql.NullString { + t.Helper() + var tsv sql.NullString + err := sqlDB.QueryRowContext(ctx, `SELECT search_tsv::text FROM chat_messages WHERE id = $1`, id).Scan(&tsv) + require.NoError(t, err) + return tsv + } + + indexedMsg := insertMsg("original secret phrase") + _, err = store.BackfillChatMessagesSearchTsv(ctx, 1000) + require.NoError(t, err) + require.True(t, searchTsv(indexedMsg.ID).Valid) + + err = store.UpdateChatMessageContentByID(ctx, database.UpdateChatMessageContentByIDParams{ + ID: indexedMsg.ID, + Content: json.RawMessage(`[{"type":"text","text":"replacement override phrase"}]`), + }) + require.NoError(t, err) + + tsv := searchTsv(indexedMsg.ID) + require.True(t, tsv.Valid) + require.Contains(t, tsv.String, "replacement") + require.NotContains(t, tsv.String, "original") + + pendingMsg := insertMsg("pending secret phrase") + err = store.UpdateChatMessageContentByID(ctx, database.UpdateChatMessageContentByIDParams{ + ID: pendingMsg.ID, + Content: json.RawMessage(`[{"type":"text","text":"pending replacement phrase"}]`), + }) + require.NoError(t, err) + require.False(t, searchTsv(pendingMsg.ID).Valid) + + _, err = store.BackfillChatMessagesSearchTsv(ctx, 1000) + require.NoError(t, err) + tsv = searchTsv(pendingMsg.ID) + require.True(t, tsv.Valid) + require.Contains(t, tsv.String, "replacement") +} + func TestChatHasUnread(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 876c952fc5975..84dbf4238bff2 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -5463,6 +5463,357 @@ func (q *sqlQuerier) InsertChatFile(ctx context.Context, arg InsertChatFileParam return i, err } +const deleteOldChatHookDispatches = `-- name: DeleteOldChatHookDispatches :execrows +WITH deletable AS ( + SELECT id + FROM chat_hook_dispatches + WHERE started_at < $1::timestamptz + ORDER BY started_at ASC + LIMIT $2::int +) +DELETE FROM chat_hook_dispatches +USING deletable +WHERE chat_hook_dispatches.id = deletable.id +` + +type DeleteOldChatHookDispatchesParams struct { + BeforeTime time.Time `db:"before_time" json:"before_time"` + LimitCount int32 `db:"limit_count" json:"limit_count"` +} + +func (q *sqlQuerier) DeleteOldChatHookDispatches(ctx context.Context, arg DeleteOldChatHookDispatchesParams) (int64, error) { + result, err := q.db.ExecContext(ctx, deleteOldChatHookDispatches, arg.BeforeTime, arg.LimitCount) + if err != nil { + return 0, err + } + return result.RowsAffected() +} + +const finalizeChatHookDispatch = `-- name: FinalizeChatHookDispatch :one +UPDATE chat_hook_dispatches +SET + finished_at = $1::timestamptz, + result = $2::text, + http_status = $3::integer, + decision = $4::text, + decision_reason = $5::text, + input_override = $6::jsonb, + original_input = $7::jsonb, + model_context = $8::text, + user_message = $9::text, + allowed_tools = $10::jsonb, + end_chat = $11::boolean, + error = $12::text +WHERE id = $13::uuid + AND chat_id = $14::uuid + AND owner_id = $15::uuid +RETURNING id, chat_id, event, turn_id, tool_use_id, owner_id, workspace_id, started_at, finished_at, result, http_status, decision, input_override, original_input, model_context, user_message, allowed_tools, end_chat, error, decision_reason, effects_applied_at, tool_name +` + +type FinalizeChatHookDispatchParams struct { + FinishedAt time.Time `db:"finished_at" json:"finished_at"` + Result string `db:"result" json:"result"` + HttpStatus sql.NullInt32 `db:"http_status" json:"http_status"` + Decision sql.NullString `db:"decision" json:"decision"` + DecisionReason sql.NullString `db:"decision_reason" json:"decision_reason"` + InputOverride pqtype.NullRawMessage `db:"input_override" json:"input_override"` + OriginalInput pqtype.NullRawMessage `db:"original_input" json:"original_input"` + ModelContext sql.NullString `db:"model_context" json:"model_context"` + UserMessage sql.NullString `db:"user_message" json:"user_message"` + AllowedTools pqtype.NullRawMessage `db:"allowed_tools" json:"allowed_tools"` + EndChat sql.NullBool `db:"end_chat" json:"end_chat"` + Error sql.NullString `db:"error" json:"error"` + ID uuid.UUID `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` +} + +func (q *sqlQuerier) FinalizeChatHookDispatch(ctx context.Context, arg FinalizeChatHookDispatchParams) (ChatHookDispatch, error) { + row := q.db.QueryRowContext(ctx, finalizeChatHookDispatch, + arg.FinishedAt, + arg.Result, + arg.HttpStatus, + arg.Decision, + arg.DecisionReason, + arg.InputOverride, + arg.OriginalInput, + arg.ModelContext, + arg.UserMessage, + arg.AllowedTools, + arg.EndChat, + arg.Error, + arg.ID, + arg.ChatID, + arg.OwnerID, + ) + var i ChatHookDispatch + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Event, + &i.TurnID, + &i.ToolUseID, + &i.OwnerID, + &i.WorkspaceID, + &i.StartedAt, + &i.FinishedAt, + &i.Result, + &i.HttpStatus, + &i.Decision, + &i.InputOverride, + &i.OriginalInput, + &i.ModelContext, + &i.UserMessage, + &i.AllowedTools, + &i.EndChat, + &i.Error, + &i.DecisionReason, + &i.EffectsAppliedAt, + &i.ToolName, + ) + return i, err +} + +const getChatHookDispatchDecision = `-- name: GetChatHookDispatchDecision :one +SELECT + id, chat_id, event, turn_id, tool_use_id, owner_id, workspace_id, started_at, finished_at, result, http_status, decision, input_override, original_input, model_context, user_message, allowed_tools, end_chat, error, decision_reason, effects_applied_at, tool_name +FROM + chat_hook_dispatches +WHERE + chat_id = $1::uuid + AND event = 'pre_tool_use' + AND tool_use_id = $2::text + AND tool_name = $3::text + AND ( + original_input = $4::jsonb + OR input_override = $4::jsonb + ) + AND turn_id IS NOT DISTINCT FROM $5::uuid + AND decision IS NOT NULL + AND result IN ('ok', 'denied') +ORDER BY + started_at DESC, + id DESC +LIMIT 1 +` + +type GetChatHookDispatchDecisionParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + ToolUseID string `db:"tool_use_id" json:"tool_use_id"` + ToolName string `db:"tool_name" json:"tool_name"` + ToolInput json.RawMessage `db:"tool_input" json:"tool_input"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` +} + +func (q *sqlQuerier) GetChatHookDispatchDecision(ctx context.Context, arg GetChatHookDispatchDecisionParams) (ChatHookDispatch, error) { + row := q.db.QueryRowContext(ctx, getChatHookDispatchDecision, + arg.ChatID, + arg.ToolUseID, + arg.ToolName, + arg.ToolInput, + arg.TurnID, + ) + var i ChatHookDispatch + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Event, + &i.TurnID, + &i.ToolUseID, + &i.OwnerID, + &i.WorkspaceID, + &i.StartedAt, + &i.FinishedAt, + &i.Result, + &i.HttpStatus, + &i.Decision, + &i.InputOverride, + &i.OriginalInput, + &i.ModelContext, + &i.UserMessage, + &i.AllowedTools, + &i.EndChat, + &i.Error, + &i.DecisionReason, + &i.EffectsAppliedAt, + &i.ToolName, + ) + return i, err +} + +const insertChatHookDispatch = `-- name: InsertChatHookDispatch :one +INSERT INTO chat_hook_dispatches ( + id, + chat_id, + event, + turn_id, + tool_use_id, + tool_name, + owner_id, + workspace_id, + started_at +) VALUES ( + $1::uuid, + $2::uuid, + $3::text, + $4::uuid, + $5::text, + $6::text, + $7::uuid, + $8::uuid, + $9::timestamptz +) +RETURNING id, chat_id, event, turn_id, tool_use_id, owner_id, workspace_id, started_at, finished_at, result, http_status, decision, input_override, original_input, model_context, user_message, allowed_tools, end_chat, error, decision_reason, effects_applied_at, tool_name +` + +type InsertChatHookDispatchParams struct { + ID uuid.UUID `db:"id" json:"id"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + Event string `db:"event" json:"event"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` + ToolUseID sql.NullString `db:"tool_use_id" json:"tool_use_id"` + ToolName sql.NullString `db:"tool_name" json:"tool_name"` + OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` + WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` + StartedAt time.Time `db:"started_at" json:"started_at"` +} + +func (q *sqlQuerier) InsertChatHookDispatch(ctx context.Context, arg InsertChatHookDispatchParams) (ChatHookDispatch, error) { + row := q.db.QueryRowContext(ctx, insertChatHookDispatch, + arg.ID, + arg.ChatID, + arg.Event, + arg.TurnID, + arg.ToolUseID, + arg.ToolName, + arg.OwnerID, + arg.WorkspaceID, + arg.StartedAt, + ) + var i ChatHookDispatch + err := row.Scan( + &i.ID, + &i.ChatID, + &i.Event, + &i.TurnID, + &i.ToolUseID, + &i.OwnerID, + &i.WorkspaceID, + &i.StartedAt, + &i.FinishedAt, + &i.Result, + &i.HttpStatus, + &i.Decision, + &i.InputOverride, + &i.OriginalInput, + &i.ModelContext, + &i.UserMessage, + &i.AllowedTools, + &i.EndChat, + &i.Error, + &i.DecisionReason, + &i.EffectsAppliedAt, + &i.ToolName, + ) + return i, err +} + +const listChatHookDispatchesByChatID = `-- name: ListChatHookDispatchesByChatID :many +SELECT + id, chat_id, event, turn_id, tool_use_id, owner_id, workspace_id, started_at, finished_at, result, http_status, decision, input_override, original_input, model_context, user_message, allowed_tools, end_chat, error, decision_reason, effects_applied_at, tool_name +FROM + chat_hook_dispatches +WHERE + chat_id = $1::uuid +ORDER BY + started_at ASC, + id ASC +` + +func (q *sqlQuerier) ListChatHookDispatchesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatHookDispatch, error) { + rows, err := q.db.QueryContext(ctx, listChatHookDispatchesByChatID, chatID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ChatHookDispatch + for rows.Next() { + var i ChatHookDispatch + if err := rows.Scan( + &i.ID, + &i.ChatID, + &i.Event, + &i.TurnID, + &i.ToolUseID, + &i.OwnerID, + &i.WorkspaceID, + &i.StartedAt, + &i.FinishedAt, + &i.Result, + &i.HttpStatus, + &i.Decision, + &i.InputOverride, + &i.OriginalInput, + &i.ModelContext, + &i.UserMessage, + &i.AllowedTools, + &i.EndChat, + &i.Error, + &i.DecisionReason, + &i.EffectsAppliedAt, + &i.ToolName, + ); 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 markChatHookDispatchEffectsApplied = `-- name: MarkChatHookDispatchEffectsApplied :exec +UPDATE chat_hook_dispatches +SET + effects_applied_at = COALESCE(effects_applied_at, NOW()) +WHERE + chat_id = $1::uuid + AND event = 'pre_tool_use' + AND id = ANY($2::uuid[]) +` + +type MarkChatHookDispatchEffectsAppliedParams struct { + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + DispatchIds []uuid.UUID `db:"dispatch_ids" json:"dispatch_ids"` +} + +func (q *sqlQuerier) MarkChatHookDispatchEffectsApplied(ctx context.Context, arg MarkChatHookDispatchEffectsAppliedParams) error { + _, err := q.db.ExecContext(ctx, markChatHookDispatchEffectsApplied, arg.ChatID, pq.Array(arg.DispatchIds)) + return err +} + +const updateChatHookAllowedTools = `-- name: UpdateChatHookAllowedTools :exec +UPDATE chats +SET + hook_allowed_tools = $1::jsonb, + updated_at = NOW() +WHERE id = $2::uuid +` + +type UpdateChatHookAllowedToolsParams struct { + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` + ID uuid.UUID `db:"id" json:"id"` +} + +func (q *sqlQuerier) UpdateChatHookAllowedTools(ctx context.Context, arg UpdateChatHookAllowedToolsParams) error { + _, err := q.db.ExecContext(ctx, updateChatHookAllowedTools, arg.HookAllowedTools, arg.ID) + return err +} + const deleteChatModelConfigByID = `-- name: DeleteChatModelConfigByID :exec UPDATE chat_model_configs @@ -6005,7 +6356,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -6053,13 +6404,14 @@ chats_expanded AS ( updated_chats.context_dirty_since, updated_chats.context_dirty_resources, updated_chats.context_error, - updated_chats.compaction_requested_at + updated_chats.compaction_requested_at, + updated_chats.hook_allowed_tools FROM updated_chats 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ORDER BY (chats_expanded.id = $1::uuid) DESC, chats_expanded.created_at ASC, chats_expanded.id ASC ` @@ -6119,6 +6471,7 @@ func (q *sqlQuerier) ArchiveChatByID(ctx context.Context, id uuid.UUID) ([]Chat, &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -6169,10 +6522,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.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, c.context_aggregate_hash, c.context_dirty_since, c.context_dirty_resources, c.context_error, c.last_reasoning_effort, c.compaction_requested_at + 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.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, c.context_aggregate_hash, c.context_dirty_since, c.context_dirty_resources, c.context_error, c.last_reasoning_effort, c.compaction_requested_at, c.hook_allowed_tools ) 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.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, a.context_aggregate_hash, a.context_dirty_since, a.context_dirty_resources, a.context_error, a.last_reasoning_effort, a.compaction_requested_at, + 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.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, a.context_aggregate_hash, a.context_dirty_since, a.context_dirty_resources, a.context_error, a.last_reasoning_effort, a.compaction_requested_at, a.hook_allowed_tools, -- Children inherit their root's activity so last_activity_at is never null. COALESCE( t.last_activity_at, @@ -6233,6 +6586,7 @@ type AutoArchiveInactiveChatsRow struct { ContextError string `db:"context_error" json:"context_error"` LastReasoningEffort NullChatReasoningEffort `db:"last_reasoning_effort" json:"last_reasoning_effort"` CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` LastActivityAt time.Time `db:"last_activity_at" json:"last_activity_at"` } @@ -6296,6 +6650,7 @@ func (q *sqlQuerier) AutoArchiveInactiveChats(ctx context.Context, arg AutoArchi &i.ContextError, &i.LastReasoningEffort, &i.CompactionRequestedAt, + &i.HookAllowedTools, &i.LastActivityAt, ); err != nil { return nil, err @@ -6564,6 +6919,12 @@ WITH deletable AS ( AND updated_at < $1::timestamptz ORDER BY updated_at ASC LIMIT $2 + -- Locking keeps the candidate set stable, so hook dispatches are + -- only removed for chats the outer DELETE also removes. + FOR UPDATE +), purged_hook_dispatches AS ( + DELETE FROM chat_hook_dispatches + WHERE chat_id IN (SELECT id FROM deletable) ) DELETE FROM chats USING deletable @@ -6579,6 +6940,8 @@ type DeleteOldChatsParams struct { // Deletes chats that have been archived for longer than the given // threshold. Active (non-archived) chats are never deleted. // All chat-scoped child tables are removed via ON DELETE CASCADE. +// Dispatches have no chat FK because they can precede chat creation. +// Delete them explicitly so their payloads respect chat retention. // Parent/root references on child chats are SET NULL. func (q *sqlQuerier) DeleteOldChats(ctx context.Context, arg DeleteOldChatsParams) (int64, error) { result, err := q.db.ExecContext(ctx, deleteOldChats, arg.BeforeTime, arg.LimitCount) @@ -6602,7 +6965,7 @@ func (q *sqlQuerier) DeleteStaleChatHeartbeats(ctx context.Context, staleSeconds } 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE agent_id = $1::uuid AND archived = false @@ -6667,6 +7030,7 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -6683,7 +7047,7 @@ func (q *sqlQuerier) GetActiveChatsByAgentID(ctx context.Context, agentID uuid.U 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, + 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, chats_expanded.hook_allowed_tools, COALESCE(activity.last_activity_at, chats_expanded.created_at)::timestamptz AS last_activity_at FROM chats_expanded LEFT JOIN LATERAL ( @@ -6759,6 +7123,7 @@ type GetAutoArchiveInactiveChatCandidatesRow struct { ContextDirtyResources pqtype.NullRawMessage `db:"context_dirty_resources" json:"context_dirty_resources"` ContextError string `db:"context_error" json:"context_error"` CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` LastActivityAt time.Time `db:"last_activity_at" json:"last_activity_at"` } @@ -6820,6 +7185,7 @@ func (q *sqlQuerier) GetAutoArchiveInactiveChatCandidates(ctx context.Context, a &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, &i.LastActivityAt, ); err != nil { return nil, err @@ -6858,7 +7224,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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE id = $1::uuid ` @@ -6912,13 +7278,14 @@ func (q *sqlQuerier) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools FROM chats WHERE id = $1::uuid FOR SHARE @@ -6969,13 +7336,14 @@ chats_expanded AS ( shared_chat.context_dirty_since, shared_chat.context_dirty_resources, shared_chat.context_error, - shared_chat.compaction_requested_at + shared_chat.compaction_requested_at, + shared_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -7028,13 +7396,14 @@ func (q *sqlQuerier) GetChatByIDForShare(ctx context.Context, id uuid.UUID) (Cha &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools FROM chats WHERE id = $1::uuid FOR UPDATE @@ -7085,13 +7454,14 @@ chats_expanded AS ( locked_chat.context_dirty_since, locked_chat.context_dirty_resources, locked_chat.context_error, - locked_chat.compaction_requested_at + locked_chat.compaction_requested_at, + locked_chat.hook_allowed_tools FROM locked_chat 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -7144,6 +7514,7 @@ func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Ch &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -7540,6 +7911,44 @@ func (q *sqlQuerier) GetChatCostSummary(ctx context.Context, arg GetChatCostSumm return i, err } +const getChatDescendantIDsByChatID = `-- name: GetChatDescendantIDsByChatID :many +WITH RECURSIVE descendants AS ( + SELECT id, created_at + FROM chats + WHERE parent_chat_id = $1::uuid + UNION ALL + SELECT c.id, c.created_at + FROM chats c + JOIN descendants d ON c.parent_chat_id = d.id +) +SELECT id +FROM descendants +ORDER BY created_at ASC, id ASC +` + +func (q *sqlQuerier) GetChatDescendantIDsByChatID(ctx context.Context, id uuid.UUID) ([]uuid.UUID, error) { + rows, err := q.db.QueryContext(ctx, getChatDescendantIDsByChatID, 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 getChatDiffStatusByChatID = `-- name: GetChatDiffStatusByChatID :one SELECT chat_id, url, pull_request_state, changes_requested, additions, deletions, changed_files, refreshed_at, stale_at, created_at, updated_at, git_branch, git_remote_origin, pull_request_title, pull_request_draft, author_login, author_avatar_url, base_branch, pr_number, commits, approved, reviewer_count, head_branch @@ -7730,7 +8139,7 @@ func (q *sqlQuerier) GetChatHeartbeat(ctx context.Context, arg GetChatHeartbeatP 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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -7766,6 +8175,7 @@ func (q *sqlQuerier) GetChatMessageByID(ctx context.Context, id int64) (ChatMess &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ) return i, err } @@ -7855,7 +8265,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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -7906,6 +8316,7 @@ func (q *sqlQuerier) GetChatMessagesByChatID(ctx context.Context, arg GetChatMes &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -7922,7 +8333,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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -7976,6 +8387,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDAscPaginated(ctx context.Context, ar &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -7992,7 +8404,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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -8059,6 +8471,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -8075,7 +8488,7 @@ func (q *sqlQuerier) GetChatMessagesByChatIDDescPaginated(ctx context.Context, a 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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -8125,6 +8538,7 @@ func (q *sqlQuerier) GetChatMessagesByRevisionForStream(ctx context.Context, arg &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -8157,7 +8571,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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -8232,6 +8646,7 @@ func (q *sqlQuerier) GetChatMessagesForPromptByChatID(ctx context.Context, chatI &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -8295,7 +8710,7 @@ func (q *sqlQuerier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]Get } const getChatQueuedMessageByID = `-- name: GetChatQueuedMessageByID :one -SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools FROM chat_queued_messages WHERE id = $1::bigint AND chat_id = $2::uuid ` @@ -8316,12 +8731,15 @@ func (q *sqlQuerier) GetChatQueuedMessageByID(ctx context.Context, arg GetChatQu &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ) return i, err } const getChatQueuedMessageHead = `-- name: GetChatQueuedMessageHead :one -SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools FROM chat_queued_messages WHERE chat_id = $1::uuid ORDER BY position ASC, id ASC LIMIT 1 @@ -8340,12 +8758,15 @@ func (q *sqlQuerier) GetChatQueuedMessageHead(ctx context.Context, chatID uuid.U &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ) return i, err } const getChatQueuedMessages = `-- name: GetChatQueuedMessages :many -SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools FROM chat_queued_messages WHERE chat_id = $1 ORDER BY created_at ASC, id ASC ` @@ -8368,6 +8789,9 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -8383,7 +8807,7 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID } const getChatQueuedMessagesByPosition = `-- name: GetChatQueuedMessagesByPosition :many -SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort FROM chat_queued_messages +SELECT id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools FROM chat_queued_messages WHERE chat_id = $1::uuid ORDER BY position ASC, id ASC ` @@ -8407,6 +8831,9 @@ func (q *sqlQuerier) GetChatQueuedMessagesByPosition(ctx context.Context, chatID &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -8602,7 +9029,7 @@ func (q *sqlQuerier) GetChatUserPromptsByChatID(ctx context.Context, arg GetChat 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, + 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, chats_expanded.hook_allowed_tools, chat_heartbeats.heartbeat_at AS current_heartbeat_at, NOT EXISTS ( SELECT 1 @@ -8684,6 +9111,7 @@ type GetChatWorkerAcquisitionCandidatesRow struct { ContextDirtyResources pqtype.NullRawMessage `db:"context_dirty_resources" json:"context_dirty_resources"` ContextError string `db:"context_error" json:"context_error"` CompactionRequestedAt sql.NullTime `db:"compaction_requested_at" json:"compaction_requested_at"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` CurrentHeartbeatAt sql.NullTime `db:"current_heartbeat_at" json:"current_heartbeat_at"` HeartbeatStale bool `db:"heartbeat_stale" json:"heartbeat_stale"` } @@ -8754,6 +9182,7 @@ func (q *sqlQuerier) GetChatWorkerAcquisitionCandidates(ctx context.Context, arg &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, &i.CurrentHeartbeatAt, &i.HeartbeatStale, ); err != nil { @@ -8780,7 +9209,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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, + 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, chats_expanded.hook_allowed_tools, EXISTS ( SELECT 1 FROM chat_messages cm WHERE cm.chat_id = chats_expanded.id @@ -9079,6 +9508,7 @@ func (q *sqlQuerier) GetChats(ctx context.Context, arg GetChatsParams) ([]GetCha &i.Chat.ContextDirtyResources, &i.Chat.ContextError, &i.Chat.CompactionRequestedAt, + &i.Chat.HookAllowedTools, &i.HasUnread, ); err != nil { return nil, err @@ -9096,7 +9526,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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at + 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE @@ -9164,6 +9594,7 @@ func (q *sqlQuerier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -9179,7 +9610,7 @@ func (q *sqlQuerier) GetChatsByChatFileID(ctx context.Context, fileID uuid.UUID) } 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE id = ANY($1::uuid[]) ORDER BY id ASC @@ -9240,6 +9671,7 @@ func (q *sqlQuerier) GetChatsByIDsForRunnerSync(ctx context.Context, ids []uuid. &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -9255,7 +9687,7 @@ func (q *sqlQuerier) GetChatsByIDsForRunnerSync(ctx context.Context, ids []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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE archived = false AND workspace_id = ANY($1::uuid[]) @@ -9317,6 +9749,7 @@ func (q *sqlQuerier) GetChatsByWorkspaceIDs(ctx context.Context, ids []uuid.UUID &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -9401,7 +9834,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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, + 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.last_reasoning_effort, 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.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, chats_expanded.context_aggregate_hash, chats_expanded.context_dirty_since, chats_expanded.context_dirty_resources, chats_expanded.context_error, chats_expanded.compaction_requested_at, chats_expanded.hook_allowed_tools, EXISTS ( SELECT 1 FROM chat_messages cm WHERE cm.chat_id = chats_expanded.id @@ -9491,6 +9924,7 @@ func (q *sqlQuerier) GetChildChatsByParentIDs(ctx context.Context, arg GetChildC &i.Chat.ContextDirtyResources, &i.Chat.ContextError, &i.Chat.CompactionRequestedAt, + &i.Chat.HookAllowedTools, &i.HasUnread, ); err != nil { return nil, err @@ -9522,7 +9956,7 @@ func (q *sqlQuerier) GetDatabaseNow(ctx context.Context) (time.Time, error) { 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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id FROM chat_messages WHERE @@ -9568,13 +10002,14 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ) 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at + 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded WHERE @@ -9652,6 +10087,7 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -9846,6 +10282,7 @@ func (q *sqlQuerier) InsertAgentContextResourcesIntoChat(ctx context.Context, ar const insertChat = `-- name: InsertChat :one WITH inserted_chat AS ( INSERT INTO chats ( + id, organization_id, owner_id, workspace_id, @@ -9861,9 +10298,10 @@ INSERT INTO chats ( mcp_server_ids, labels, dynamic_tools, - client_type + client_type, + hook_allowed_tools ) VALUES ( - $1::uuid, + COALESCE($1::uuid, gen_random_uuid()), $2::uuid, $3::uuid, $4::uuid, @@ -9871,16 +10309,18 @@ INSERT INTO chats ( $6::uuid, $7::uuid, $8::uuid, - $9::text, - $10::chat_mode, - $11::chat_plan_mode, - $12::chat_status, - COALESCE($13::uuid[], '{}'::uuid[]), - COALESCE($14::jsonb, '{}'::jsonb), - $15::jsonb, - $16::chat_client_type + $9::uuid, + $10::text, + $11::chat_mode, + $12::chat_plan_mode, + $13::chat_status, + COALESCE($14::uuid[], '{}'::uuid[]), + COALESCE($15::jsonb, '{}'::jsonb), + $16::jsonb, + $17::chat_client_type, + $18::jsonb ) -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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -9928,17 +10368,19 @@ chats_expanded AS ( inserted_chat.context_dirty_since, inserted_chat.context_dirty_resources, inserted_chat.context_error, - inserted_chat.compaction_requested_at + inserted_chat.compaction_requested_at, + inserted_chat.hook_allowed_tools FROM inserted_chat 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` type InsertChatParams struct { + ID uuid.NullUUID `db:"id" json:"id"` OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` OwnerID uuid.UUID `db:"owner_id" json:"owner_id"` WorkspaceID uuid.NullUUID `db:"workspace_id" json:"workspace_id"` @@ -9955,10 +10397,12 @@ type InsertChatParams struct { Labels pqtype.NullRawMessage `db:"labels" json:"labels"` DynamicTools pqtype.NullRawMessage `db:"dynamic_tools" json:"dynamic_tools"` ClientType ChatClientType `db:"client_type" json:"client_type"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` } func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat, error) { row := q.db.QueryRowContext(ctx, insertChat, + arg.ID, arg.OrganizationID, arg.OwnerID, arg.WorkspaceID, @@ -9975,6 +10419,7 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat arg.Labels, arg.DynamicTools, arg.ClientType, + arg.HookAllowedTools, ) var i Chat err := row.Scan( @@ -10023,6 +10468,7 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -10032,7 +10478,7 @@ WITH batch AS ( SELECT ( SELECT val - FROM UNNEST($3::uuid[]) + FROM UNNEST($4::uuid[]) WITH ORDINALITY AS t(val, ord) WHERE val != '00000000-0000-0000-0000-000000000000'::uuid ORDER BY ord DESC @@ -10040,7 +10486,7 @@ WITH batch AS ( ) AS last_model_config_id, ( SELECT NULLIF(val, '')::chat_reasoning_effort - FROM UNNEST($4::text[]) + FROM UNNEST($5::text[]) WITH ORDINALITY AS t(val, ord) WHERE val != '' ORDER BY ord DESC @@ -10063,6 +10509,7 @@ updated_chat AS ( ) INSERT INTO chat_messages ( chat_id, + turn_id, created_by, model_config_id, reasoning_effort, @@ -10085,27 +10532,29 @@ SELECT $1::uuid, NULLIF(UNNEST($2::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), NULLIF(UNNEST($3::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), - NULLIF(UNNEST($4::text[]), '')::chat_reasoning_effort, - UNNEST($5::chat_message_role[]), - UNNEST($6::text[])::jsonb, - UNNEST($7::smallint[]), - UNNEST($8::chat_message_visibility[]), - NULLIF(UNNEST($9::bigint[]), 0), + NULLIF(UNNEST($4::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), + NULLIF(UNNEST($5::text[]), '')::chat_reasoning_effort, + UNNEST($6::chat_message_role[]), + UNNEST($7::text[])::jsonb, + UNNEST($8::smallint[]), + UNNEST($9::chat_message_visibility[]), NULLIF(UNNEST($10::bigint[]), 0), NULLIF(UNNEST($11::bigint[]), 0), NULLIF(UNNEST($12::bigint[]), 0), NULLIF(UNNEST($13::bigint[]), 0), NULLIF(UNNEST($14::bigint[]), 0), NULLIF(UNNEST($15::bigint[]), 0), - UNNEST($16::boolean[]), - NULLIF(UNNEST($17::bigint[]), 0), - NULLIF(UNNEST($18::bigint[]), 0) + NULLIF(UNNEST($16::bigint[]), 0), + UNNEST($17::boolean[]), + NULLIF(UNNEST($18::bigint[]), 0), + NULLIF(UNNEST($19::bigint[]), 0) 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, revision, reasoning_effort, search_tsv + 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, revision, reasoning_effort, search_tsv, turn_id ` type InsertChatMessagesParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + TurnID []uuid.UUID `db:"turn_id" json:"turn_id"` CreatedBy []uuid.UUID `db:"created_by" json:"created_by"` ModelConfigID []uuid.UUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort []string `db:"reasoning_effort" json:"reasoning_effort"` @@ -10128,6 +10577,7 @@ type InsertChatMessagesParams struct { func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, insertChatMessages, arg.ChatID, + pq.Array(arg.TurnID), pq.Array(arg.CreatedBy), pq.Array(arg.ModelConfigID), pq.Array(arg.ReasoningEffort), @@ -10178,6 +10628,7 @@ func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessa &i.Revision, &i.ReasoningEffort, &i.SearchTsv, + &i.TurnID, ); err != nil { return nil, err } @@ -10193,20 +10644,22 @@ 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, reasoning_effort, created_by) +INSERT INTO chat_queued_messages (chat_id, turn_id, content, model_config_id, reasoning_effort, created_by) SELECT $1::uuid, - $2::jsonb, - $3::uuid, - $4::chat_reasoning_effort, + $2::uuid, + $3::jsonb, + $4::uuid, + $5::chat_reasoning_effort, chats.owner_id FROM chats WHERE chats.id = $1::uuid -RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools ` type InsertChatQueuedMessageParams struct { ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` Content json.RawMessage `db:"content" json:"content"` ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` @@ -10218,6 +10671,7 @@ type InsertChatQueuedMessageParams struct { func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChatQueuedMessageParams) (ChatQueuedMessage, error) { row := q.db.QueryRowContext(ctx, insertChatQueuedMessage, arg.ChatID, + arg.TurnID, arg.Content, arg.ModelConfigID, arg.ReasoningEffort, @@ -10232,28 +10686,37 @@ func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChat &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ) return i, err } const insertChatQueuedMessageWithCreator = `-- name: InsertChatQueuedMessageWithCreator :one -INSERT INTO chat_queued_messages (chat_id, content, model_config_id, reasoning_effort, created_by) +INSERT INTO chat_queued_messages (chat_id, turn_id, content, model_config_id, reasoning_effort, created_by, hook_prefix, hook_allowed_tools) VALUES ( $1::uuid, - $2::jsonb, - $3::uuid, - $4::chat_reasoning_effort, - $5::uuid + $2::uuid, + $3::jsonb, + $4::uuid, + $5::chat_reasoning_effort, + $6::uuid, + $7::jsonb, + $8::jsonb ) -RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools ` 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"` - ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` - CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + ChatID uuid.UUID `db:"chat_id" json:"chat_id"` + TurnID uuid.NullUUID `db:"turn_id" json:"turn_id"` + Content json.RawMessage `db:"content" json:"content"` + ModelConfigID uuid.NullUUID `db:"model_config_id" json:"model_config_id"` + ReasoningEffort NullChatReasoningEffort `db:"reasoning_effort" json:"reasoning_effort"` + CreatedBy uuid.UUID `db:"created_by" json:"created_by"` + HookPrefix pqtype.NullRawMessage `db:"hook_prefix" json:"hook_prefix"` + HookAllowedTools pqtype.NullRawMessage `db:"hook_allowed_tools" json:"hook_allowed_tools"` } // Inserts a queued message that carries a position (from the default @@ -10262,10 +10725,13 @@ type InsertChatQueuedMessageWithCreatorParams struct { func (q *sqlQuerier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg InsertChatQueuedMessageWithCreatorParams) (ChatQueuedMessage, error) { row := q.db.QueryRowContext(ctx, insertChatQueuedMessageWithCreator, arg.ChatID, + arg.TurnID, arg.Content, arg.ModelConfigID, arg.ReasoningEffort, arg.CreatedBy, + arg.HookPrefix, + arg.HookAllowedTools, ) var i ChatQueuedMessage err := row.Scan( @@ -10277,6 +10743,9 @@ func (q *sqlQuerier) InsertChatQueuedMessageWithCreator(ctx context.Context, arg &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ) return i, err } @@ -10510,7 +10979,7 @@ WITH bumped_chat AS ( 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -10558,12 +11027,13 @@ chats_expanded AS ( bumped_chat.context_dirty_since, bumped_chat.context_dirty_resources, bumped_chat.context_error, - bumped_chat.compaction_requested_at + bumped_chat.compaction_requested_at, + bumped_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -10620,6 +11090,7 @@ func (q *sqlQuerier) LockChatAndBumpSnapshotVersion(ctx context.Context, id uuid &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -10744,7 +11215,7 @@ WHERE id = ( ORDER BY cqm.created_at ASC, cqm.id ASC LIMIT 1 ) -RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort +RETURNING id, chat_id, content, created_at, model_config_id, position, created_by, reasoning_effort, turn_id, hook_prefix, hook_allowed_tools ` func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) (ChatQueuedMessage, error) { @@ -10759,6 +11230,9 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID) &i.Position, &i.CreatedBy, &i.ReasoningEffort, + &i.TurnID, + &i.HookPrefix, + &i.HookAllowedTools, ) return i, err } @@ -10950,7 +11424,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -10998,13 +11472,14 @@ chats_expanded AS ( updated_chats.context_dirty_since, updated_chats.context_dirty_resources, updated_chats.context_error, - updated_chats.compaction_requested_at + updated_chats.compaction_requested_at, + updated_chats.hook_allowed_tools FROM updated_chats 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ORDER BY (chats_expanded.id = $1::uuid) DESC, chats_expanded.created_at ASC, chats_expanded.id ASC ` @@ -11068,6 +11543,7 @@ func (q *sqlQuerier) UnarchiveChatByID(ctx context.Context, id uuid.UUID) ([]Cha &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ); err != nil { return nil, err } @@ -11170,7 +11646,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11218,13 +11694,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -11283,6 +11760,7 @@ func (q *sqlQuerier) UpdateChatBuildAgentBinding(ctx context.Context, arg Update &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -11296,7 +11774,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11344,13 +11822,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -11408,6 +11887,7 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -11426,7 +11906,7 @@ WITH updated_chat AS ( pin_order = CASE WHEN $2::boolean THEN 0 ELSE pin_order END, updated_at = NOW() WHERE id = $8::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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11474,12 +11954,13 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -11557,6 +12038,7 @@ func (q *sqlQuerier) UpdateChatExecutionState(ctx context.Context, arg UpdateCha &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -11615,7 +12097,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11663,13 +12145,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -11727,6 +12210,7 @@ func (q *sqlQuerier) UpdateChatLabelsByID(ctx context.Context, arg UpdateChatLab &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -11740,7 +12224,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11788,13 +12272,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -11852,6 +12337,7 @@ func (q *sqlQuerier) UpdateChatLastModelConfigByID(ctx context.Context, arg Upda &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -11915,7 +12401,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -11963,13 +12449,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12027,10 +12514,35 @@ func (q *sqlQuerier) UpdateChatMCPServerIDs(ctx context.Context, arg UpdateChatM &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } +const updateChatMessageContentByID = `-- name: UpdateChatMessageContentByID :exec +UPDATE chat_messages +SET content = $1::jsonb, + search_tsv = CASE + WHEN search_tsv IS NULL THEN NULL + ELSE COALESCE( + to_tsvector('simple', chat_message_search_text($1::jsonb)), + ''::tsvector) + END +WHERE id = $2::bigint +` + +type UpdateChatMessageContentByIDParams struct { + Content json.RawMessage `db:"content" json:"content"` + ID int64 `db:"id" json:"id"` +} + +// Preserve NULL as the backfill marker; otherwise refresh search_tsv +// from the new content. +func (q *sqlQuerier) UpdateChatMessageContentByID(ctx context.Context, arg UpdateChatMessageContentByIDParams) error { + _, err := q.db.ExecContext(ctx, updateChatMessageContentByID, arg.Content, arg.ID) + return err +} + const updateChatPinOrder = `-- name: UpdateChatPinOrder :exec WITH target_chat AS ( SELECT @@ -12111,7 +12623,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -12159,13 +12671,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12223,6 +12736,7 @@ func (q *sqlQuerier) UpdateChatPlanModeByID(ctx context.Context, arg UpdateChatP &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -12234,7 +12748,7 @@ WITH updated_chat AS ( 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -12282,12 +12796,13 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12347,6 +12862,7 @@ func (q *sqlQuerier) UpdateChatRetryState(ctx context.Context, arg UpdateChatRet &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -12364,7 +12880,7 @@ SET 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -12412,13 +12928,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12487,6 +13004,7 @@ func (q *sqlQuerier) UpdateChatStatus(ctx context.Context, arg UpdateChatStatusP &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } @@ -12502,7 +13020,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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at +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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), chats_expanded AS ( SELECT @@ -12550,13 +13068,14 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12614,13 +13133,14 @@ func (q *sqlQuerier) UpdateChatTitleByID(ctx context.Context, arg UpdateChatTitl &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } const updateChatWorkspaceBinding = `-- name: UpdateChatWorkspaceBinding :one WITH current_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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools FROM chats WHERE id = $1::uuid ), @@ -12639,13 +13159,13 @@ changed_chat AS ( updated_at = NOW() WHERE id = $1::uuid AND (SELECT changed FROM binding_changed) - 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools ), result_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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools FROM changed_chat UNION ALL - 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at + 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, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, last_reasoning_effort, compaction_requested_at, hook_allowed_tools FROM current_chat WHERE NOT (SELECT changed FROM binding_changed) ), @@ -12695,13 +13215,14 @@ chats_expanded AS ( result_chat.context_dirty_since, result_chat.context_dirty_resources, result_chat.context_error, - result_chat.compaction_requested_at + result_chat.compaction_requested_at, + result_chat.hook_allowed_tools FROM result_chat LEFT JOIN chats root ON root.id = COALESCE(result_chat.root_chat_id, result_chat.parent_chat_id) JOIN visible_users owner ON owner.id = result_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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at +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, last_reasoning_effort, archived, last_error, mode, mcp_server_ids, labels, build_id, agent_id, pin_order, last_read_message_id, 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, context_aggregate_hash, context_dirty_since, context_dirty_resources, context_error, compaction_requested_at, hook_allowed_tools FROM chats_expanded ` @@ -12766,6 +13287,7 @@ func (q *sqlQuerier) UpdateChatWorkspaceBinding(ctx context.Context, arg UpdateC &i.ContextDirtyResources, &i.ContextError, &i.CompactionRequestedAt, + &i.HookAllowedTools, ) return i, err } diff --git a/coderd/database/queries/chathooks.sql b/coderd/database/queries/chathooks.sql new file mode 100644 index 0000000000000..887fab8e06ce8 --- /dev/null +++ b/coderd/database/queries/chathooks.sql @@ -0,0 +1,104 @@ +-- name: InsertChatHookDispatch :one +INSERT INTO chat_hook_dispatches ( + id, + chat_id, + event, + turn_id, + tool_use_id, + tool_name, + owner_id, + workspace_id, + started_at +) VALUES ( + @id::uuid, + @chat_id::uuid, + @event::text, + sqlc.narg('turn_id')::uuid, + sqlc.narg('tool_use_id')::text, + sqlc.narg('tool_name')::text, + @owner_id::uuid, + sqlc.narg('workspace_id')::uuid, + @started_at::timestamptz +) +RETURNING *; + +-- name: FinalizeChatHookDispatch :one +UPDATE chat_hook_dispatches +SET + finished_at = @finished_at::timestamptz, + result = @result::text, + http_status = sqlc.narg('http_status')::integer, + decision = sqlc.narg('decision')::text, + decision_reason = sqlc.narg('decision_reason')::text, + input_override = sqlc.narg('input_override')::jsonb, + original_input = sqlc.narg('original_input')::jsonb, + model_context = sqlc.narg('model_context')::text, + user_message = sqlc.narg('user_message')::text, + allowed_tools = sqlc.narg('allowed_tools')::jsonb, + end_chat = sqlc.narg('end_chat')::boolean, + error = sqlc.narg('error')::text +WHERE id = @id::uuid + AND chat_id = @chat_id::uuid + AND owner_id = @owner_id::uuid +RETURNING *; + +-- name: DeleteOldChatHookDispatches :execrows +WITH deletable AS ( + SELECT id + FROM chat_hook_dispatches + WHERE started_at < @before_time::timestamptz + ORDER BY started_at ASC + LIMIT @limit_count::int +) +DELETE FROM chat_hook_dispatches +USING deletable +WHERE chat_hook_dispatches.id = deletable.id; + +-- name: UpdateChatHookAllowedTools :exec +UPDATE chats +SET + hook_allowed_tools = sqlc.narg('hook_allowed_tools')::jsonb, + updated_at = NOW() +WHERE id = @id::uuid; + +-- name: ListChatHookDispatchesByChatID :many +SELECT + * +FROM + chat_hook_dispatches +WHERE + chat_id = @chat_id::uuid +ORDER BY + started_at ASC, + id ASC; + +-- name: MarkChatHookDispatchEffectsApplied :exec +UPDATE chat_hook_dispatches +SET + effects_applied_at = COALESCE(effects_applied_at, NOW()) +WHERE + chat_id = @chat_id::uuid + AND event = 'pre_tool_use' + AND id = ANY(@dispatch_ids::uuid[]); + +-- name: GetChatHookDispatchDecision :one +SELECT + * +FROM + chat_hook_dispatches +WHERE + chat_id = @chat_id::uuid + AND event = 'pre_tool_use' + AND tool_use_id = @tool_use_id::text + AND tool_name = @tool_name::text + AND ( + original_input = @tool_input::jsonb + OR input_override = @tool_input::jsonb + ) + AND turn_id IS NOT DISTINCT FROM sqlc.narg('turn_id')::uuid + AND decision IS NOT NULL + AND result IN ('ok', 'denied') +ORDER BY + started_at DESC, + id DESC +LIMIT 1; diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index d1a796c54c8f1..97766f03bd077 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -51,7 +51,8 @@ chats_expanded AS ( updated_chats.context_dirty_since, updated_chats.context_dirty_resources, updated_chats.context_error, - updated_chats.compaction_requested_at + updated_chats.compaction_requested_at, + updated_chats.hook_allowed_tools FROM updated_chats LEFT JOIN chats root ON root.id = COALESCE(updated_chats.root_chat_id, updated_chats.parent_chat_id) @@ -119,7 +120,8 @@ chats_expanded AS ( updated_chats.context_dirty_since, updated_chats.context_dirty_resources, updated_chats.context_error, - updated_chats.compaction_requested_at + updated_chats.compaction_requested_at, + updated_chats.hook_allowed_tools FROM updated_chats LEFT JOIN chats root ON root.id = COALESCE(updated_chats.root_chat_id, updated_chats.parent_chat_id) @@ -354,6 +356,20 @@ FROM chats WHERE id = @id::uuid OR root_chat_id = @id::uuid ORDER BY (id = @id::uuid) DESC, created_at ASC, id ASC; +-- name: GetChatDescendantIDsByChatID :many +WITH RECURSIVE descendants AS ( + SELECT id, created_at + FROM chats + WHERE parent_chat_id = @id::uuid + UNION ALL + SELECT c.id, c.created_at + FROM chats c + JOIN descendants d ON c.parent_chat_id = d.id +) +SELECT id +FROM descendants +ORDER BY created_at ASC, id ASC; + -- name: GetChatACLByID :one SELECT user_acl AS users, @@ -381,6 +397,19 @@ WHERE id = @id::bigint AND deleted = false; +-- name: UpdateChatMessageContentByID :exec +-- Preserve NULL as the backfill marker; otherwise refresh search_tsv +-- from the new content. +UPDATE chat_messages +SET content = @content::jsonb, + search_tsv = CASE + WHEN search_tsv IS NULL THEN NULL + ELSE COALESCE( + to_tsvector('simple', chat_message_search_text(@content::jsonb)), + ''::tsvector) + END +WHERE id = @id::bigint; + -- name: GetChatMessagesByChatID :many SELECT * @@ -770,6 +799,7 @@ ORDER BY -- name: InsertChat :one WITH inserted_chat AS ( INSERT INTO chats ( + id, organization_id, owner_id, workspace_id, @@ -785,8 +815,10 @@ INSERT INTO chats ( mcp_server_ids, labels, dynamic_tools, - client_type + client_type, + hook_allowed_tools ) VALUES ( + COALESCE(sqlc.narg('id')::uuid, gen_random_uuid()), @organization_id::uuid, @owner_id::uuid, sqlc.narg('workspace_id')::uuid, @@ -802,7 +834,8 @@ INSERT INTO chats ( COALESCE(@mcp_server_ids::uuid[], '{}'::uuid[]), COALESCE(sqlc.narg('labels')::jsonb, '{}'::jsonb), sqlc.narg('dynamic_tools')::jsonb, - @client_type::chat_client_type + @client_type::chat_client_type, + sqlc.narg('hook_allowed_tools')::jsonb ) RETURNING * ), @@ -852,7 +885,8 @@ chats_expanded AS ( inserted_chat.context_dirty_since, inserted_chat.context_dirty_resources, inserted_chat.context_error, - inserted_chat.compaction_requested_at + inserted_chat.compaction_requested_at, + inserted_chat.hook_allowed_tools FROM inserted_chat LEFT JOIN chats root ON root.id = COALESCE(inserted_chat.root_chat_id, inserted_chat.parent_chat_id) @@ -897,6 +931,7 @@ updated_chat AS ( ) INSERT INTO chat_messages ( chat_id, + turn_id, created_by, model_config_id, reasoning_effort, @@ -917,6 +952,7 @@ INSERT INTO chat_messages ( ) SELECT @chat_id::uuid, + NULLIF(UNNEST(@turn_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), NULLIF(UNNEST(@created_by::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), NULLIF(UNNEST(@model_config_id::uuid[]), '00000000-0000-0000-0000-000000000000'::uuid), NULLIF(UNNEST(@reasoning_effort::text[]), '')::chat_reasoning_effort, @@ -994,7 +1030,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1062,7 +1099,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1128,7 +1166,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1194,7 +1233,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1260,7 +1300,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1346,7 +1387,8 @@ chats_expanded AS ( result_chat.context_dirty_since, result_chat.context_dirty_resources, result_chat.context_error, - result_chat.compaction_requested_at + result_chat.compaction_requested_at, + result_chat.hook_allowed_tools FROM result_chat LEFT JOIN chats root ON root.id = COALESCE(result_chat.root_chat_id, result_chat.parent_chat_id) @@ -1411,7 +1453,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1494,7 +1537,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1702,7 +1746,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools FROM updated_chat LEFT JOIN chats root ON root.id = COALESCE(updated_chat.root_chat_id, updated_chat.parent_chat_id) @@ -1869,9 +1914,10 @@ RETURNING -- 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, reasoning_effort, created_by) +INSERT INTO chat_queued_messages (chat_id, turn_id, content, model_config_id, reasoning_effort, created_by) SELECT @chat_id::uuid, + sqlc.narg('turn_id')::uuid, @content::jsonb, sqlc.narg('model_config_id')::uuid, sqlc.narg('reasoning_effort')::chat_reasoning_effort, @@ -1979,7 +2025,8 @@ chats_expanded AS ( locked_chat.context_dirty_since, locked_chat.context_dirty_resources, locked_chat.context_error, - locked_chat.compaction_requested_at + locked_chat.compaction_requested_at, + locked_chat.hook_allowed_tools FROM locked_chat LEFT JOIN chats root ON root.id = COALESCE(locked_chat.root_chat_id, locked_chat.parent_chat_id) @@ -2041,7 +2088,8 @@ chats_expanded AS ( shared_chat.context_dirty_since, shared_chat.context_dirty_resources, shared_chat.context_error, - shared_chat.compaction_requested_at + shared_chat.compaction_requested_at, + shared_chat.hook_allowed_tools FROM shared_chat LEFT JOIN chats root ON root.id = COALESCE(shared_chat.root_chat_id, shared_chat.parent_chat_id) @@ -2488,6 +2536,8 @@ WHERE id = @id::uuid; -- Deletes chats that have been archived for longer than the given -- threshold. Active (non-archived) chats are never deleted. -- All chat-scoped child tables are removed via ON DELETE CASCADE. +-- Dispatches have no chat FK because they can precede chat creation. +-- Delete them explicitly so their payloads respect chat retention. -- Parent/root references on child chats are SET NULL. WITH deletable AS ( SELECT id @@ -2496,6 +2546,12 @@ WITH deletable AS ( AND updated_at < @before_time::timestamptz ORDER BY updated_at ASC LIMIT @limit_count + -- Locking keeps the candidate set stable, so hook dispatches are + -- only removed for chats the outer DELETE also removes. + FOR UPDATE +), purged_hook_dispatches AS ( + DELETE FROM chat_hook_dispatches + WHERE chat_id IN (SELECT id FROM deletable) ) DELETE FROM chats USING deletable @@ -2716,7 +2772,8 @@ chats_expanded AS ( bumped_chat.context_dirty_since, bumped_chat.context_dirty_resources, bumped_chat.context_error, - bumped_chat.compaction_requested_at + bumped_chat.compaction_requested_at, + bumped_chat.hook_allowed_tools 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 @@ -2791,7 +2848,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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 @@ -2856,7 +2914,8 @@ chats_expanded AS ( updated_chat.context_dirty_since, updated_chat.context_dirty_resources, updated_chat.context_error, - updated_chat.compaction_requested_at + updated_chat.compaction_requested_at, + updated_chat.hook_allowed_tools 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 @@ -2881,13 +2940,16 @@ SELECT NOW()::timestamptz AS now; -- 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, reasoning_effort, created_by) +INSERT INTO chat_queued_messages (chat_id, turn_id, content, model_config_id, reasoning_effort, created_by, hook_prefix, hook_allowed_tools) VALUES ( @chat_id::uuid, + sqlc.narg('turn_id')::uuid, @content::jsonb, sqlc.narg('model_config_id')::uuid, sqlc.narg('reasoning_effort')::chat_reasoning_effort, - @created_by::uuid + @created_by::uuid, + sqlc.narg('hook_prefix')::jsonb, + sqlc.narg('hook_allowed_tools')::jsonb ) RETURNING *; diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 4b1a4376f2db4..06d72a6155246 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -29,6 +29,7 @@ const ( 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); + UniqueChatHookDispatchesPkey UniqueConstraint = "chat_hook_dispatches_pkey" // ALTER TABLE ONLY chat_hook_dispatches ADD CONSTRAINT chat_hook_dispatches_pkey PRIMARY KEY (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 aa95221609c26..7376642207827 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -55,6 +55,7 @@ import ( "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/chathooks" "github.com/coder/coder/v2/coderd/x/gitsync" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/wsjson" @@ -110,6 +111,14 @@ func writeChatUsageLimitExceeded( }) } +// Avoid returning raw dispatch errors, which may expose deployment internals. +func writeChatHookDispatchFailed(ctx context.Context, rw http.ResponseWriter, hookErr *chathooks.DispatchError) { + httpapi.Write(ctx, rw, http.StatusBadGateway, codersdk.Response{ + Message: "Chat lifecycle hook dispatch failed.", + Detail: fmt.Sprintf("Lifecycle hook dispatch %s failed (%s).", hookErr.DispatchID, hookErr.Class), + }) +} + func maybeWriteLimitErr(ctx context.Context, rw http.ResponseWriter, err error) bool { var limitErr *chatd.UsageLimitExceededError if errors.As(err, &limitErr) { @@ -1424,23 +1433,38 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { } chat, err := api.chatDaemon.CreateChat(ctx, chatd.CreateOptions{ - OrganizationID: req.OrganizationID, - OwnerID: apiKey.UserID, - WorkspaceID: workspaceSelection.WorkspaceID, - Title: title, - ModelConfigID: modelConfigID, - ReasoningEffort: reasoningEffort, - PlanMode: planModeToNullChatPlanMode(req.PlanMode), - ClientType: clientType, - SystemPrompt: req.SystemPrompt, - InitialUserContent: contentBlocks, - MCPServerIDs: mcpServerIDs, - Labels: labels, - DynamicTools: dynamicToolsJSON, + OrganizationID: req.OrganizationID, + OwnerID: apiKey.UserID, + WorkspaceID: workspaceSelection.WorkspaceID, + Title: title, + TitleDerivedFromContent: true, + ModelConfigID: modelConfigID, + ReasoningEffort: reasoningEffort, + PlanMode: planModeToNullChatPlanMode(req.PlanMode), + ClientType: clientType, + SystemPrompt: req.SystemPrompt, + InitialUserContent: contentBlocks, + MCPServerIDs: mcpServerIDs, + Labels: labels, + DynamicTools: dynamicToolsJSON, // IMPORTANT: users can only create root chats at the time of writing. ParentChatID: uuid.NullUUID{}, }) if err != nil { + var denied *chatd.UserPromptDeniedError + if errors.As(err, &denied) { + message := denied.UserMessage + if message == "" { + message = "Chat creation denied by lifecycle hook." + } + httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{Message: message}) + return + } + var hookErr *chathooks.DispatchError + if errors.As(err, &hookErr) { + writeChatHookDispatchFailed(ctx, rw, hookErr) + return + } if maybeWriteLimitErr(ctx, rw, err) { return } @@ -1483,10 +1507,22 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { return } - // Link any user-uploaded files referenced in the initial - // message to this newly created chat (best-effort; cap - // enforced in SQL). - unlinked, capExceeded := api.linkFilesToChat(ctx, chat.ID, fileIDs) + linkFileIDs := fileIDs + if len(fileIDs) > 0 { + initialUser, err := api.Database.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + }) + if err != nil { + api.Logger.Warn(ctx, "load initial message for file linking", + slog.F("chat_id", chat.ID), + slog.Error(err), + ) + } else { + linkFileIDs = api.linkedFileIDsFromContent(ctx, initialUser, fileIDs) + } + } + unlinked, capExceeded := api.linkFilesToChat(ctx, chat.ID, linkFileIDs) // Re-read the chat so the response reflects the authoritative // database state (file links are deduped in the join table). @@ -3381,9 +3417,29 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { }, ) if sendErr != nil { + var denied *chatd.UserPromptDeniedError + if errors.As(sendErr, &denied) { + message := denied.UserMessage + if message == "" { + message = "Chat message denied by lifecycle hook." + } + httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{Message: message}) + return + } + var hookErr *chathooks.DispatchError + if errors.As(sendErr, &hookErr) { + writeChatHookDispatchFailed(ctx, rw, hookErr) + return + } if maybeWriteLimitErr(ctx, rw, sendErr) { return } + if errors.Is(sendErr, chatstate.ErrChatNotRoot) { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Hook end_chat can only archive a root chat.", + }) + return + } if xerrors.Is(sendErr, chatd.ErrChatArchived) { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Cannot send messages to an archived chat.", @@ -3435,9 +3491,24 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { return } - // Link any user-uploaded files referenced in this message - // to the chat (best-effort; cap enforced in SQL). - unlinked, capExceeded := api.linkFilesToChat(ctx, chatID, fileIDs) + if sendResult.Ended { + httpapi.Write(ctx, rw, http.StatusOK, codersdk.CreateChatMessageResponse{Ended: true}) + return + } + + linkFileIDs := fileIDs + if sendResult.Queued { + if sendResult.QueuedMessage != nil { + linkFileIDs = api.linkedFileIDsFromContent(ctx, database.ChatMessage{ + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, + Content: pqtype.NullRawMessage{RawMessage: sendResult.QueuedMessage.Content, Valid: true}, + }, fileIDs) + } + } else { + linkFileIDs = api.linkedFileIDsFromContent(ctx, sendResult.Message, fileIDs) + } + unlinked, capExceeded := api.linkFilesToChat(ctx, chatID, linkFileIDs) response := codersdk.CreateChatMessageResponse{Queued: sendResult.Queued} if sendResult.Queued { if sendResult.QueuedMessage != nil { @@ -3446,6 +3517,15 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { } else { message := convertChatMessage(sendResult.Message) response.Message = &message + // Return the full inserted batch because hook notices can have lower IDs. + // A client that caches only the user message could skip them when + // reconnecting with after_id. + for _, inserted := range sendResult.InsertedMessages { + if inserted.Visibility == database.ChatMessageVisibilityModel { + continue + } + response.Messages = append(response.Messages, convertChatMessage(inserted)) + } } if len(unlinked) > 0 { if capExceeded { @@ -3550,11 +3630,29 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { ReasoningEffort: editReasoningEffort, }) if editErr != nil { + var denied *chatd.UserPromptDeniedError + if errors.As(editErr, &denied) { + message := denied.UserMessage + if message == "" { + message = "Chat message denied by lifecycle hook." + } + httpapi.Write(ctx, rw, http.StatusForbidden, codersdk.Response{Message: message}) + return + } + var hookErr *chathooks.DispatchError + if errors.As(editErr, &hookErr) { + writeChatHookDispatchFailed(ctx, rw, hookErr) + return + } if maybeWriteLimitErr(ctx, rw, editErr) { return } switch { + case errors.Is(editErr, chatstate.ErrChatNotRoot): + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Hook end_chat can only archive a root chat.", + }) case xerrors.Is(editErr, chatd.ErrChatArchived): httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Cannot edit messages in an archived chat.", @@ -3594,11 +3692,23 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { return } - // Link any user-uploaded files referenced in the edited - // message to the chat (best-effort; cap enforced in SQL). - unlinked, capExceeded := api.linkFilesToChat(ctx, chat.ID, fileIDs) - response := codersdk.EditChatMessageResponse{ - Message: convertChatMessage(editResult.Message), + if editResult.Ended { + httpapi.Write(ctx, rw, http.StatusOK, codersdk.EditChatMessageResponse{Ended: true}) + return + } + + unlinked, capExceeded := api.linkFilesToChat(ctx, chat.ID, api.linkedFileIDsFromContent(ctx, editResult.Message, fileIDs)) + message := convertChatMessage(editResult.Message) + response := codersdk.EditChatMessageResponse{Message: &message} + // Hook notices and synthetic cancellations are inserted before the + // replacement with lower IDs; clients that seed their transcript + // cache from this response need all of them, or a stream reconnect + // with after_id set to the replacement would skip the earlier rows. + for _, inserted := range editResult.InsertedMessages { + if inserted.Visibility == database.ChatMessageVisibilityModel { + continue + } + response.Messages = append(response.Messages, convertChatMessage(inserted)) } if len(unlinked) > 0 { if capExceeded { @@ -6818,14 +6928,30 @@ func createChatInputFromParts( return content, pasteData, fileIDs, nil } -// linkFilesToChat inserts file-link rows into the chat_file_links -// join table. Cap enforcement and dedup are handled atomically in -// SQL. On success returns (nil, false). On failure returns the full -// input fileIDs slice — linking is all-or-nothing because the -// SQL operates on the batch atomically. capExceeded indicates -// whether the failure was due to the cap being exceeded (true) -// or a database error (false). -// Failures are logged but never block the caller. +// A prompt override may remove file parts, so derive links from persisted +// content. Fall back to request IDs if parsing fails. +func (api *API) linkedFileIDsFromContent(ctx context.Context, msg database.ChatMessage, requestFileIDs []uuid.UUID) []uuid.UUID { + if len(requestFileIDs) == 0 { + return nil + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + api.Logger.Warn(ctx, "parse persisted message for file linking", + slog.F("message_id", msg.ID), + slog.Error(err), + ) + return requestFileIDs + } + var ids []uuid.UUID + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeFile && part.FileID.Valid { + ids = append(ids, part.FileID.UUID) + } + } + return ids +} + +// Link each batch atomically; report the entire batch when linking fails. func (api *API) linkFilesToChat(ctx context.Context, chatID uuid.UUID, fileIDs []uuid.UUID) (unlinked []uuid.UUID, capExceeded bool) { if len(fileIDs) == 0 { return nil, false @@ -8157,7 +8283,10 @@ func (api *API) postChatToolResults(rw http.ResponseWriter, r *http.Request) { if err != nil { var validationErr *chatd.ToolResultValidationError var conflictErr *chatd.ToolResultStatusConflictError + var hookErr *chathooks.DispatchError switch { + case errors.As(err, &hookErr): + writeChatHookDispatchFailed(ctx, rw, hookErr) case xerrors.Is(err, chatd.ErrChatArchived): httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ Message: "Cannot submit tool results to an archived chat.", diff --git a/coderd/exp_chats_hooks_test.go b/coderd/exp_chats_hooks_test.go new file mode 100644 index 0000000000000..5ee9f4cc1db7b --- /dev/null +++ b/coderd/exp_chats_hooks_test.go @@ -0,0 +1,415 @@ +package coderd_test + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "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/x/chatd/chattest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" + "github.com/coder/serpent" +) + +func TestPostChatsInitialPromptHookErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + response string + wantStatus int + wantMessage string + }{ + { + name: "deny", + statusCode: http.StatusOK, + response: `{"permission":{"decision":"deny"},"user_message":"blocked by policy"}`, + wantStatus: http.StatusForbidden, + wantMessage: "blocked by policy", + }, + { + name: "dispatch failure", + statusCode: http.StatusInternalServerError, + wantStatus: http.StatusBadGateway, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + requests := make(chan agenthooks.Request, 2) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + requests <- request + w.WriteHeader(test.statusCode) + if test.response != "" { + _, err := w.Write([]byte(test.response)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + + client, db := newChatClientWithDatabase(t, func(opts *coderdtest.Options) { + opts.ChatWorkerDisabled = true + require.NoError(t, opts.DeploymentValues.AI.Chat.HookURL.Set(consumer.URL)) + opts.DeploymentValues.AI.Chat.HookSecret = serpent.String("test-hook-secret-32-bytes-minimum!!") + opts.DeploymentValues.AI.Chat.HookTimeout = serpent.Duration(time.Second) + opts.DeploymentValues.AI.Chat.HookEnabled = serpent.Bool(true) + }) + user := coderdtest.CreateFirstUser(t, client.Client) + model := createAdditionalChatModelConfig(t, client, "openai", "gpt-4.1") + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, + ModelConfigID: &model.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "blocked prompt", + }}, + }) + sdkErr := coderdtest.SDKError(t, err) + require.Equal(t, test.wantStatus, sdkErr.StatusCode()) + if test.wantMessage != "" { + require.Equal(t, test.wantMessage, sdkErr.Message) + } + request := testutil.RequireReceive(ctx, t, requests) + require.Equal(t, agenthooks.EventUserPromptSubmit, request.Type) + require.NotEqual(t, uuid.Nil, request.Meta.ChatID) + _, err = db.GetChatByID(dbauthz.AsSystemRestricted(ctx), request.Meta.ChatID) + require.ErrorIs(t, err, sql.ErrNoRows) + }) + } +} + +func TestChatLifecycleHooksExperimentDisabled(t *testing.T) { + t.Parallel() + + var hookRequests atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hookRequests.Add(1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(consumer.Close) + + client, db := newChatClientWithDatabase(t, func(opts *coderdtest.Options) { + opts.ChatWorkerDisabled = true + opts.DeploymentValues.Experiments = serpent.StringArray{ + string(codersdk.ExperimentChatAdvisor), + string(codersdk.ExperimentChatVirtualDesktop), + } + require.NoError(t, opts.DeploymentValues.AI.Chat.HookURL.Set(consumer.URL)) + opts.DeploymentValues.AI.Chat.HookSecret = serpent.String("test-hook-secret-32-bytes-minimum!!") + opts.DeploymentValues.AI.Chat.HookTimeout = serpent.Duration(time.Second) + opts.DeploymentValues.AI.Chat.HookEnabled = serpent.Bool(true) + }) + user := coderdtest.CreateFirstUser(t, client.Client) + model := createAdditionalChatModelConfig(t, client, "openai", "gpt-4.1") + ctx := testutil.Context(t, testutil.WaitLong) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, + ModelConfigID: &model.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "prompt with hooks disabled", + }}, + }) + require.NoError(t, err) + + require.Zero(t, hookRequests.Load()) + rows, err := db.ListChatHookDispatchesByChatID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Empty(t, rows) +} + +func TestChatLifecycleHooksWorkedExample(t *testing.T) { + t.Parallel() + + const ( + secret = "test-hook-secret-32-bytes-minimum!!" + deniedToolCallID = "call_denied" + allowedToolCallID = "call_allowed" + ) + ctx := testutil.Context(t, testutil.WaitLong) + var modelCalls atomic.Int32 + secondModelRequest := make(chan []byte, 1) + thirdModelRequest := make(chan []byte, 1) + modelURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("Lifecycle hooks") + } + switch modelCalls.Add(1) { + case 1: + chunk := chattest.OpenAIToolCallChunk("read_secret", `{"path":"/tmp/secret"}`) + chunk.Choices[0].ToolCalls[0].ID = deniedToolCallID + return chattest.OpenAIStreamingResponse(chunk) + case 2: + secondModelRequest <- bytes.Clone(req.RawBody) + chunk := chattest.OpenAIToolCallChunk("search_docs", `{"query":"customer secret"}`) + chunk.Choices[0].ToolCalls[0].ID = allowedToolCallID + return chattest.OpenAIStreamingResponse(chunk) + case 3: + thirdModelRequest <- bytes.Clone(req.RawBody) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + default: + return chattest.OpenAIErrorResponse(http.StatusInternalServerError, "unexpected_call", "unexpected model call") + } + }) + + hookEvents := make(chan agenthooks.EventType, 16) + recordHook := func(event agenthooks.EventType) { + hookEvents <- event + } + consumer := httptest.NewServer(agenthooks.NewHTTPHandler([]byte(secret), agenthooks.Hooks{ + SessionStart: func(context.Context, agenthooks.Meta, agenthooks.SessionStartData) (agenthooks.Response, error) { + recordHook(agenthooks.EventSessionStart) + return agenthooks.Response{}, nil + }, + UserPromptSubmit: func(context.Context, agenthooks.Meta, agenthooks.UserPromptSubmitData) (agenthooks.Response, error) { + recordHook(agenthooks.EventUserPromptSubmit) + return agenthooks.Response{}, nil + }, + PreToolUse: func(_ context.Context, _ agenthooks.Meta, tool agenthooks.PreToolUseData) (agenthooks.Response, error) { + recordHook(agenthooks.EventPreToolUse) + switch tool.ToolUseID { + case deniedToolCallID: + return agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionDeny, + Reason: "secret reads are blocked", + }}, nil + case allowedToolCallID: + return agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAllow, + InputOverride: json.RawMessage(`{"query":"public documentation"}`), + }}, nil + default: + return agenthooks.Response{}, nil + } + }, + PostToolUse: func(context.Context, agenthooks.Meta, agenthooks.PostToolUseData) (agenthooks.Response, error) { + recordHook(agenthooks.EventPostToolUse) + return agenthooks.Response{ + ModelContext: "The approved search result is safe to use.", + UserMessage: "Search result approved by policy.", + }, nil + }, + Stop: func(context.Context, agenthooks.Meta, agenthooks.StopData) (agenthooks.Response, error) { + recordHook(agenthooks.EventStop) + return agenthooks.Response{}, nil + }, + })) + t.Cleanup(consumer.Close) + + client, db := newChatClientWithDatabase(t, func(opts *coderdtest.Options) { + require.NoError(t, opts.DeploymentValues.AI.Chat.HookURL.Set(consumer.URL)) + opts.DeploymentValues.AI.Chat.HookSecret = serpent.String(secret) + opts.DeploymentValues.AI.Chat.HookTimeout = serpent.Duration(time.Second) + opts.DeploymentValues.AI.Chat.HookEnabled = serpent.Bool(true) + }) + user := coderdtest.CreateFirstUser(t, client.Client) + model := createChatModelConfigWithBaseURL(t, client, modelURL) + + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, + ModelConfigID: &model.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "Find the deployment documentation.", + }}, + UnsafeDynamicTools: []codersdk.DynamicTool{ + { + Name: "read_secret", + Description: "Read a secret file.", + InputSchema: json.RawMessage(`{"type":"object"}`), + }, + { + Name: "search_docs", + Description: "Search public documentation.", + InputSchema: json.RawMessage(`{"type":"object"}`), + }, + }, + }) + require.NoError(t, err) + + var stored database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + stored, err = db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + return err == nil && stored.Status == database.ChatStatusRequiresAction + }, testutil.IntervalFast) + require.Equal(t, int32(2), modelCalls.Load()) + require.Contains(t, string(testutil.RequireReceive(ctx, t, secondModelRequest)), "DENIED: secret reads are blocked") + + messages, err := client.GetChatMessages(ctx, chat.ID, nil) + require.NoError(t, err) + var allowedCall *codersdk.ChatMessagePart + for _, message := range messages.Messages { + for i := range message.Content { + part := &message.Content[i] + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolCallID == allowedToolCallID { + allowedCall = part + } + } + } + require.NotNil(t, allowedCall) + require.JSONEq(t, `{"query":"public documentation"}`, string(allowedCall.Args)) + + err = client.SubmitToolResults(ctx, chat.ID, codersdk.SubmitToolResultsRequest{ + Results: []codersdk.ToolResult{{ + ToolCallID: allowedToolCallID, + Output: json.RawMessage(`{"matches":["agent hooks"]}`), + }}, + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + stored, err = db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + return err == nil && stored.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.Contains(t, string(testutil.RequireReceive(ctx, t, thirdModelRequest)), "The approved search result is safe to use.") + require.Equal(t, int32(3), modelCalls.Load()) + + rows, err := db.ListChatHookDispatchesByChatID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Len(t, rows, 6) + assertDispatch := func(event agenthooks.EventType, toolUseID, result, decision string) database.ChatHookDispatch { + t.Helper() + for _, row := range rows { + if row.Event != string(event) || row.ToolUseID.String != toolUseID { + continue + } + require.Equal(t, result, row.Result) + require.True(t, row.FinishedAt.Valid) + require.Equal(t, int32(http.StatusOK), row.HttpStatus.Int32) + require.Equal(t, decision != "", row.Decision.Valid) + if decision != "" { + require.Equal(t, decision, row.Decision.String) + } + return row + } + require.FailNow(t, "hook dispatch not found", "event=%s tool_use_id=%s", event, toolUseID) + return database.ChatHookDispatch{} + } + assertDispatch(agenthooks.EventUserPromptSubmit, "", "ok", "") + assertDispatch(agenthooks.EventSessionStart, "", "ok", "") + denied := assertDispatch(agenthooks.EventPreToolUse, deniedToolCallID, "denied", "deny") + require.Equal(t, "secret reads are blocked", denied.DecisionReason.String) + allowed := assertDispatch(agenthooks.EventPreToolUse, allowedToolCallID, "ok", "allow") + require.JSONEq(t, `{"query":"public documentation"}`, string(allowed.InputOverride.RawMessage)) + post := assertDispatch(agenthooks.EventPostToolUse, allowedToolCallID, "ok", "") + require.Equal(t, "The approved search result is safe to use.", post.ModelContext.String) + assertDispatch(agenthooks.EventStop, "", "ok", "") + + seenEvents := make([]agenthooks.EventType, 0, len(rows)) + for range rows { + seenEvents = append(seenEvents, testutil.RequireReceive(ctx, t, hookEvents)) + } + require.ElementsMatch(t, []agenthooks.EventType{ + agenthooks.EventUserPromptSubmit, + agenthooks.EventSessionStart, + agenthooks.EventPreToolUse, + agenthooks.EventPreToolUse, + agenthooks.EventPostToolUse, + agenthooks.EventStop, + }, seenEvents) +} + +func TestChatHooksFileLinksAfterPromptOverride(t *testing.T) { + t.Parallel() + + const secret = "test-hook-secret-32-bytes-minimum!!" + ctx := testutil.Context(t, testutil.WaitLong) + modelURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + consumer := httptest.NewServer(agenthooks.NewHTTPHandler([]byte(secret), agenthooks.Hooks{ + UserPromptSubmit: func(_ context.Context, _ agenthooks.Meta, data agenthooks.UserPromptSubmitData) (agenthooks.Response, error) { + if strings.Contains(data.Prompt, "REDACTME") { + return agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAllow, + InputOverride: json.RawMessage(`{"prompt":"redacted"}`), + }}, nil + } + return agenthooks.Response{}, nil + }, + })) + t.Cleanup(consumer.Close) + + client, api := newChatClientWithAPI(t, func(opts *coderdtest.Options) { + require.NoError(t, opts.DeploymentValues.AI.Chat.HookURL.Set(consumer.URL)) + opts.DeploymentValues.AI.Chat.HookSecret = serpent.String(secret) + opts.DeploymentValues.AI.Chat.HookTimeout = serpent.Duration(time.Second) + opts.DeploymentValues.AI.Chat.HookEnabled = serpent.Bool(true) + }) + user := coderdtest.CreateFirstUser(t, client.Client) + model := createChatModelConfigWithBaseURL(t, client, modelURL) + + uploadFile := func(name string) uuid.UUID { + pngData := append([]byte{0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A}, make([]byte, 16)...) + resp, err := client.UploadChatFile(ctx, user.OrganizationID, "image/png", name, bytes.NewReader(pngData)) + require.NoError(t, err) + return resp.ID + } + + redactedFile := uploadFile("redacted.png") + chat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: user.OrganizationID, + ModelConfigID: &model.ID, + Content: []codersdk.ChatInputPart{ + {Type: codersdk.ChatInputPartTypeText, Text: "REDACTME create"}, + {Type: codersdk.ChatInputPartTypeFile, FileID: redactedFile}, + }, + }) + require.NoError(t, err) + created, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Empty(t, created.Files, "overridden create must not link dropped attachments") + + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) + + keptFile := uploadFile("kept.png") + sendResp, err := client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{ + {Type: codersdk.ChatInputPartTypeText, Text: "keep this"}, + {Type: codersdk.ChatInputPartTypeFile, FileID: keptFile}, + }, + }) + require.NoError(t, err) + require.False(t, sendResp.Queued) + afterSend, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Len(t, afterSend.Files, 1) + require.Equal(t, keptFile, afterSend.Files[0].ID) + + coderdtest.WaitForChatSettled(ctx, t, api, chat.ID) + + droppedFile := uploadFile("dropped.png") + _, err = client.CreateChatMessage(ctx, chat.ID, codersdk.CreateChatMessageRequest{ + Content: []codersdk.ChatInputPart{ + {Type: codersdk.ChatInputPartTypeText, Text: "REDACTME send"}, + {Type: codersdk.ChatInputPartTypeFile, FileID: droppedFile}, + }, + }) + require.NoError(t, err) + afterOverride, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.Len(t, afterOverride.Files, 1, "overridden send must not link dropped attachments") + require.Equal(t, keptFile, afterOverride.Files[0].ID) +} diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 09d934b341a3b..8ae88600afb5f 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -75,6 +75,7 @@ func newChatTestOptions( values.Experiments = serpent.StringArray{ string(codersdk.ExperimentChatAdvisor), string(codersdk.ExperimentChatVirtualDesktop), + string(codersdk.ExperimentAgentLifecycleHooks), } } diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index fa64cd8bd44f5..b6b3cffdb8686 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -110,6 +110,7 @@ I don't recommend reading the rest of section thoroughly if this is your first t - `Create(initialMessages)` creates a new chat, initializes `snapshot_version` to 1, inserts its initial history, and lands in `running`. The inserted initial history sets `history_version` to 1. Since the queue has not changed, `queue_version` remains 0. This transition is a special case: since the chat does not exist at the time it's run, the chat row cannot be locked before the transition is applied. - `SetArchived(archived)` sets or clears the archived marker for one chat. +- `EndChat(prefixMessages)` inserts hook response messages, clears the queue and ownership state, then lands a root chat in archived `waiting`. - `SendMessage(m, busy_behavior)` inserts a user message directly when the chat is idle, or queues it when the chat is busy. `busy_behavior` must be either `queue` or `interrupt`. With `busy_behavior=interrupt`, it also requests interruption or cancels a pending dynamic-tool action as needed. - `EditMessage(k, replacement)` clears queued messages, cancels or obsoletes active work, marks the truncated active-history suffix as deleted, inserts the replacement turn, and lands in `running`. - `DeleteQueuedMessage(qid)` removes one queued message without changing the active history. @@ -117,6 +118,7 @@ I don't recommend reading the rest of section thoroughly if this is your first t - `Interrupt(reason)` requests cancellation of an active generation or closes pending dynamic-tool action. It preserves queued backlog. - `CompleteRequiresAction(results)` inserts submitted tool-result messages, clears `requires_action_deadline_at`, and lands in `running`. It preserves queued messages. - `RequestCompaction` records a manual compaction request on an idle chat by setting `compaction_requested_at` and landing in `running` without inserting any message. The chat worker picks the chat up like any other running chat and consumes the request. See [Manual compaction](#manual-compaction). +- `FailIdle(err)` moves a waiting chat to `error` and persists `last_error = err` without requiring runner ownership. Used when a blocking lifecycle hook dispatch fails while handling `SendMessage` or `EditMessage`, before any generation starts. ### Transitions used by the chat worker @@ -149,11 +151,14 @@ stateDiagram-v2 W --> R0: SendMessage W --> R0: EditMessage W --> R0: RequestCompaction + W --> E0: FailIdle W --> XW: SetArchived(true) + W --> XW: EndChat E0 --> R0: SendMessage E0 --> R0: EditMessage E0 --> XE0: SetArchived(true) + E0 --> XW: EndChat E1 --> R1: SendMessage E1 --> R0: EditMessage @@ -162,6 +167,7 @@ stateDiagram-v2 E1 --> R0: PromoteQueuedMessage / promoted last queued E1 --> R1: PromoteQueuedMessage / queue still non-empty E1 --> XE1: SetArchived(true) + E1 --> XW: EndChat R0 --> R0: RecordGenerationAttempt R0 --> R0: RecordRetryState @@ -173,6 +179,7 @@ stateDiagram-v2 R0 --> W: FinishTurn / queue empty R0 --> E0: FinishError R0 --> R1: SendMessage(queue) + R0 --> XW: EndChat R1 --> R1: RecordGenerationAttempt R1 --> R1: RecordRetryState @@ -183,6 +190,7 @@ stateDiagram-v2 R1 --> R0: EditMessage R1 --> E1: FinishError R1 --> R1: SendMessage(queue) + R1 --> XW: EndChat R1 --> R0: DeleteQueuedMessage / removed last queued R1 --> R1: DeleteQueuedMessage / queue still non-empty R1 --> I1: PromoteQueuedMessage @@ -192,18 +200,21 @@ stateDiagram-v2 I0 --> I1: SendMessage I0 --> R0: EditMessage I0 --> W: FinishInterruption + I0 --> XW: EndChat I1 --> I1: SendMessage I1 --> R0: EditMessage I1 --> I0: DeleteQueuedMessage / removed last queued I1 --> I1: DeleteQueuedMessage / queue still non-empty I1 --> I1: PromoteQueuedMessage + I1 --> XW: EndChat I1 --> R0: FinishInterruption / promoted last queued I1 --> R1: FinishInterruption / queue still non-empty after promoting head A0 --> R0: CompleteRequiresAction A0 --> R0: Interrupt A0 --> R0: CancelRequiresAction + A0 --> XW: EndChat A0 --> A1: SendMessage(queue) A0 --> R1: SendMessage(interrupt) A0 --> R0: EditMessage @@ -211,6 +222,7 @@ stateDiagram-v2 A1 --> R1: CompleteRequiresAction A1 --> R1: Interrupt A1 --> R1: CancelRequiresAction + A1 --> XW: EndChat A1 --> A1: SendMessage(queue) A1 --> R1: SendMessage(interrupt) A1 --> R0: EditMessage diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index d594af63e0cf5..8f6d872319902 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -51,8 +51,10 @@ import ( "github.com/coder/coder/v2/coderd/x/chatd/chattool" "github.com/coder/coder/v2/coderd/x/chatd/mcpclient" "github.com/coder/coder/v2/coderd/x/chatd/messagepartbuffer" + "github.com/coder/coder/v2/coderd/x/chathooks" skillspkg "github.com/coder/coder/v2/coderd/x/skills" "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" "github.com/coder/coder/v2/codersdk/workspacesdk" "github.com/coder/quartz" ) @@ -176,6 +178,7 @@ type Server struct { stopWorkspaceFn chattool.StopWorkspaceFn pubsub pubsub.Pubsub webpushDispatcher webpush.Dispatcher + hookDispatcher *chathooks.Dispatcher providerAPIKeys chatprovider.ProviderAPIKeys allowBYOK bool oidcTokenSource mcpclient.UserOIDCTokenSource @@ -1164,24 +1167,25 @@ func (e *UsageLimitExceededError) Error() string { // CreateOptions controls chat creation in the shared chat mutation path. type CreateOptions struct { - OrganizationID uuid.UUID - OwnerID uuid.UUID - WorkspaceID uuid.NullUUID - BuildID uuid.NullUUID - AgentID uuid.NullUUID - ParentChatID uuid.NullUUID - RootChatID uuid.NullUUID - Title string - ModelConfigID uuid.UUID - ReasoningEffort *string - ChatMode database.NullChatMode - PlanMode database.NullChatPlanMode - ClientType database.ChatClientType - SystemPrompt string - InitialUserContent []codersdk.ChatMessagePart - MCPServerIDs []uuid.UUID - Labels database.StringMap - DynamicTools json.RawMessage + OrganizationID uuid.UUID + OwnerID uuid.UUID + WorkspaceID uuid.NullUUID + BuildID uuid.NullUUID + AgentID uuid.NullUUID + ParentChatID uuid.NullUUID + RootChatID uuid.NullUUID + Title string + TitleDerivedFromContent bool + ModelConfigID uuid.UUID + ReasoningEffort *string + ChatMode database.NullChatMode + PlanMode database.NullChatPlanMode + ClientType database.ChatClientType + SystemPrompt string + InitialUserContent []codersdk.ChatMessagePart + MCPServerIDs []uuid.UUID + Labels database.StringMap + DynamicTools json.RawMessage } // SendMessageBusyBehavior controls what happens when a chat is already active. @@ -1212,9 +1216,13 @@ type SendMessageOptions struct { // SendMessageResult contains the outcome of user message processing. type SendMessageResult struct { Queued bool + Ended bool QueuedMessage *database.ChatQueuedMessage Message database.ChatMessage - Chat database.Chat + // InsertedMessages holds non-queued send messages in insertion order. + // It is empty for queued sends. + InsertedMessages []database.ChatMessage + Chat database.Chat } // EditMessageOptions controls user message edits via soft-delete and re-insert. @@ -1232,8 +1240,13 @@ type EditMessageOptions struct { // EditMessageResult contains the replacement user message and chat status. type EditMessageResult struct { + Ended bool Message database.ChatMessage - Chat database.Chat + // InsertedMessages holds every message the edit inserted, in + // insertion order: synthetic tool cancellations, hook prefix + // messages, then the replacement user message. + InsertedMessages []database.ChatMessage + Chat database.Chat } // PromoteQueuedOptions controls queued-message promotion. @@ -1301,6 +1314,41 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C return database.Chat{}, xerrors.Errorf("marshal labels: %w", err) } + chatID := uuid.New() + var turnID *uuid.UUID + contentParts := opts.InitialUserContent + var hookResponse agenthooks.Response + if p.hookDispatcher != nil && p.hookDispatcher.Enabled() { + // Hook dispatch needs the same model admission that the insert enforces. + if err := validateCreateModelConfigID(ctx, p.db, opts.ModelConfigID); err != nil { + return database.Chat{}, err + } + mintedTurnID := uuid.New() + turnID = &mintedTurnID + hookChat := database.Chat{} + hookChat.ID = chatID + hookChat.OwnerID = opts.OwnerID + hookChat.WorkspaceID = opts.WorkspaceID + hookResponse, err = p.dispatchUserPromptSubmit(ctx, hookChat, *turnID, contentParts) + if err != nil { + return database.Chat{}, err + } + if hookResponse.EndChat { + return database.Chat{}, &UserPromptDeniedError{UserMessage: hookResponse.UserMessage} + } + override, overridden, err := userPromptOverride(hookResponse) + if err != nil { + return database.Chat{}, err + } + if overridden { + contentParts = []codersdk.ChatMessagePart{codersdk.ChatMessageText(override)} + // Avoid deriving titles from the prompt that policy replaced. + if opts.TitleDerivedFromContent { + opts.Title = chatprompt.FallbackTitle(chatprompt.TitleText(contentParts, nil)) + } + } + } + userPrompt := SanitizePromptText(opts.SystemPrompt) workspaceAwareness := workspaceDetachedAwareness if opts.WorkspaceID.Valid { @@ -1312,7 +1360,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C if err != nil { return database.Chat{}, xerrors.Errorf("marshal workspace awareness: %w", err) } - userContent, err := chatprompt.MarshalParts(opts.InitialUserContent) + userContent, err := chatprompt.MarshalParts(contentParts) if err != nil { return database.Chat{}, xerrors.Errorf("marshal initial user content: %w", err) } @@ -1337,9 +1385,22 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C initialMessages = append(initialMessages, systemMessage(userPromptContent, opts.ModelConfigID)) } initialMessages = append(initialMessages, systemMessage(workspaceAwarenessContent, opts.ModelConfigID)) - initialMessages = append(initialMessages, userMessage(userContent, opts.ModelConfigID, opts.OwnerID, opts.ReasoningEffort)) + prefixMessages, err := hookPrefixMessages(hookResponse, opts.ModelConfigID, turnID) + if err != nil { + return database.Chat{}, err + } + initialMessages = append(initialMessages, prefixMessages...) + initialUserMessage := userMessage(userContent, opts.ModelConfigID, opts.OwnerID, opts.ReasoningEffort) + if turnID != nil { + initialUserMessage.TurnID = uuid.NullUUID{UUID: *turnID, Valid: true} + } + initialMessages = append(initialMessages, initialUserMessage) - result, err := chatstate.CreateChat(ctx, p.db, p.pubsub, chatstate.CreateChatInput{ + hookAllowedTools, err := hookAllowedTools(hookResponse) + if err != nil { + return database.Chat{}, err + } + result, err := chatstate.CreateChatWithID(ctx, p.db, p.pubsub, chatID, hookAllowedTools, chatstate.CreateChatInput{ OrganizationID: opts.OrganizationID, OwnerID: opts.OwnerID, WorkspaceID: opts.WorkspaceID, @@ -1409,7 +1470,53 @@ func (p *Server) SendMessage( return SendMessageResult{}, xerrors.Errorf("invalid busy behavior %q", opts.BusyBehavior) } - content, err := chatprompt.MarshalParts(opts.Content) + turnID := uuid.New() + contentParts := opts.Content + var hookResponse agenthooks.Response + if p.hookDispatcher != nil && p.hookDispatcher.Enabled() { + chat, err := p.db.GetChatByID(ctx, opts.ChatID) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("load chat for user_prompt_submit: %w", err) + } + // Hook dispatch needs preflight admission; the transaction repeats it under lock. + if chat.Archived { + return SendMessageResult{}, ErrChatArchived + } + if err := p.checkUsageLimit(ctx, p.db, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); err != nil { + return SendMessageResult{}, err + } + if _, err := resolveSendMessageModelConfigID(ctx, p.db, chat, opts.ModelConfigID); err != nil { + return SendMessageResult{}, err + } + // A non-empty queue forces the capacity-checked path under the transaction lock. + queuedCount, err := p.db.CountChatQueuedMessages(ctx, opts.ChatID) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("count queued messages: %w", err) + } + if queuedCount >= chatstate.MaxQueueSize { + return SendMessageResult{}, &chatstate.MessageQueueFullError{Max: chatstate.MaxQueueSize} + } + hookResponse, err = p.dispatchUserPromptSubmit(ctx, chat, turnID, contentParts) + if err != nil { + var denied *UserPromptDeniedError + if errors.As(err, &denied) && hookResponse.EndChat { + if endErr := p.endChatAfterPromptDenial(ctx, opts.ChatID, nil); endErr != nil { + return SendMessageResult{}, errors.Join(err, endErr) + } + return SendMessageResult{}, err + } + return SendMessageResult{}, p.handleUserPromptDispatchError(ctx, opts.ChatID, err) + } + override, overridden, err := userPromptOverride(hookResponse) + if err != nil { + return SendMessageResult{}, err + } + if overridden { + contentParts = []codersdk.ChatMessagePart{codersdk.ChatMessageText(override)} + } + } + + content, err := chatprompt.MarshalParts(contentParts) if err != nil { return SendMessageResult{}, xerrors.Errorf("marshal message content: %w", err) } @@ -1418,6 +1525,7 @@ func (p *Server) SendMessage( requestedMCPServerIDs := opts.MCPServerIDs var result SendMessageResult + var endedDescendants []database.Chat machine := p.newChatMachine(opts.ChatID) updateErr := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { lockedChat, err := store.GetChatByID(ctx, opts.ChatID) @@ -1454,6 +1562,30 @@ func (p *Server) SendMessage( return err } + prefixMessages, err := hookPrefixMessages(hookResponse, modelConfigID, &turnID) + if err != nil { + return err + } + // Keep the tool policy with queued prompts until promotion. + hookToolPolicy, err := hookAllowedTools(hookResponse) + if err != nil { + return err + } + if hookResponse.EndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return err + } + endedDescendants = endResult.EndedDescendants + result.Ended = true + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after end: %w", err) + } + result.Chat = refreshed + return nil + } + // Update MCP server IDs on the chat when explicitly provided. // Explore child chats keep the spawn-time snapshot immutable. if requestedMCPServerIDs != nil { @@ -1480,9 +1612,13 @@ func (p *Server) SendMessage( // Queue capacity is enforced inside tx.SendMessage; this // wrapper only propagates the typed error. + message := userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort) + message.TurnID = uuid.NullUUID{UUID: turnID, Valid: true} sendResult, err := tx.SendMessage(chatstate.SendMessageInput{ - Message: userMessage(content, modelConfigID, messageCreatedBy, opts.ReasoningEffort), - BusyBehavior: busyBehaviorToChatState(busyBehavior), + Message: message, + PrefixMessages: prefixMessages, + HookAllowedTools: hookToolPolicy, + BusyBehavior: busyBehaviorToChatState(busyBehavior), }) if err != nil { return err @@ -1496,6 +1632,7 @@ func (p *Server) SendMessage( // cancellation messages; the user message is always // last in the inserted slice. result.Message = sendResult.InsertedMessages[len(sendResult.InsertedMessages)-1] + result.InsertedMessages = sendResult.InsertedMessages } // Capture the post-transition chat inside the same // transaction so the returned chat and the watch event @@ -1514,7 +1651,11 @@ func (p *Server) SendMessage( // 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) + if result.Ended { + p.publishEndChatSideEffects(ctx, result.Chat, endedDescendants) + } else { + p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) + } return result, nil } @@ -1590,6 +1731,20 @@ func requireEnabledChatModelConfig( return nil } +func validateCreateModelConfigID(ctx context.Context, store database.Store, modelConfigID uuid.UUID) error { + if modelConfigID == uuid.Nil { + return xerrors.Errorf("%w: %s", ErrInvalidModelConfigID, modelConfigID) + } + chatdCtx := chatdModelConfigLookupContext(ctx) + if _, err := store.GetChatModelConfigByID(chatdCtx, modelConfigID); err != nil { + if errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("%w: %s", ErrInvalidModelConfigID, modelConfigID) + } + return xerrors.Errorf("get requested model config %s: %w", modelConfigID, err) + } + return nil +} + func resolveFallbackModelConfigID( ctx context.Context, store database.Store, @@ -1633,6 +1788,37 @@ func resolveFallbackModelConfigID( return defaultConfig.ID, nil } +func validateModelConfigOverride( + ctx context.Context, + store database.Store, + requested uuid.UUID, +) (uuid.NullUUID, error) { + if requested == uuid.Nil { + return uuid.NullUUID{}, nil + } + if err := requireEnabledChatModelConfig(ctx, store, requested); err != nil { + return uuid.NullUUID{}, err + } + return uuid.NullUUID{UUID: requested, Valid: true}, nil +} + +func validateEditTarget(ctx context.Context, store database.Store, chatID uuid.UUID, messageID int64) error { + target, err := store.GetChatMessageByID(ctx, messageID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return ErrEditedMessageNotFound + } + return xerrors.Errorf("get edited message: %w", err) + } + if target.ChatID != chatID || target.Deleted { + return ErrEditedMessageNotFound + } + if target.Role != database.ChatMessageRoleUser { + return ErrEditedMessageNotUser + } + return nil +} + // 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 @@ -1651,14 +1837,68 @@ func (p *Server) EditMessage( return EditMessageResult{}, xerrors.New("content is required") } - content, err := chatprompt.MarshalParts(opts.Content) + turnID := uuid.New() + contentParts := opts.Content + var sessionStartResponse, hookResponse agenthooks.Response + if p.hookDispatcher != nil && p.hookDispatcher.Enabled() { + chat, err := p.db.GetChatByID(ctx, opts.ChatID) + if err != nil { + return EditMessageResult{}, xerrors.Errorf("load chat for edit hooks: %w", err) + } + // Hook dispatch needs preflight admission; the transaction repeats it under lock. + if chat.Archived { + return EditMessageResult{}, ErrChatArchived + } + if err := p.checkUsageLimit(ctx, p.db, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); err != nil { + return EditMessageResult{}, err + } + if err := validateEditTarget(ctx, p.db, opts.ChatID, opts.EditedMessageID); err != nil { + return EditMessageResult{}, err + } + if _, err := validateModelConfigOverride(ctx, p.db, opts.ModelConfigID); err != nil { + return EditMessageResult{}, err + } + sessionStartResponse, err = p.dispatchLifecycleHook(ctx, chat, &turnID, agenthooks.EventSessionStart, agenthooks.SessionStartData{Source: sessionStartSourceClear}) + if err != nil { + return EditMessageResult{}, p.handleAPIDispatchError(ctx, opts.ChatID, agenthooks.EventSessionStart, err) + } + if sessionStartResponse.EndChat { + return p.endChatFromEditSessionStart(ctx, chat, &turnID, sessionStartResponse) + } + hookResponse, err = p.dispatchUserPromptSubmit(ctx, chat, turnID, contentParts) + if err != nil { + // Accepted end_chat effects take precedence over prompt denial. + var denied *UserPromptDeniedError + if errors.As(err, &denied) && hookResponse.EndChat { + sessionMessages, prefixErr := hookPrefixMessages(sessionStartResponse, chat.LastModelConfigID, &turnID) + if prefixErr != nil { + return EditMessageResult{}, errors.Join(err, prefixErr) + } + if endErr := p.endChatAfterPromptDenial(ctx, opts.ChatID, sessionMessages); endErr != nil { + return EditMessageResult{}, errors.Join(err, endErr) + } + return EditMessageResult{}, err + } + return EditMessageResult{}, p.handleUserPromptDispatchError(ctx, opts.ChatID, err) + } + override, overridden, err := userPromptOverride(hookResponse) + if err != nil { + return EditMessageResult{}, err + } + if overridden { + contentParts = []codersdk.ChatMessagePart{codersdk.ChatMessageText(override)} + } + } + + content, err := chatprompt.MarshalParts(contentParts) if err != nil { return EditMessageResult{}, xerrors.Errorf("marshal message content: %w", err) } var ( - result EditMessageResult - editedMsg database.ChatMessage - editedCutoffT time.Time + result EditMessageResult + editedMsg database.ChatMessage + editedCutoffT time.Time + endedDescendants []database.Chat ) machine := p.newChatMachine(opts.ChatID) err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { @@ -1686,18 +1926,19 @@ func (p *Server) EditMessage( if target.ChatID != opts.ChatID { return ErrEditedMessageNotFound } + if target.Deleted { + return ErrEditedMessageNotFound + } + if target.Role != database.ChatMessageRoleUser { + return ErrEditedMessageNotUser + } editedMsg = target - // 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 := requireEnabledChatModelConfig(ctx, store, opts.ModelConfigID); err != nil { - return err - } - modelOverride = uuid.NullUUID{UUID: opts.ModelConfigID, Valid: true} - } else { + modelOverride, err := validateModelConfigOverride(ctx, store, opts.ModelConfigID) + if err != nil { + return err + } + if !modelOverride.Valid { // Without an explicit override the transition preserves // the edited message's original model, which may have been // disabled since; resolve it like a normal message send. @@ -1714,6 +1955,34 @@ func (p *Server) EditMessage( } } + modelConfigID := target.ModelConfigID.UUID + if modelOverride.Valid { + modelConfigID = modelOverride.UUID + } + responses := []agenthooks.Response{sessionStartResponse, hookResponse} + prefixMessages, hookEndChat, err := hookResponseMessages(responses, modelConfigID, &turnID) + if err != nil { + return err + } + if err := applyHookAllowedToolsResponses(ctx, store, opts.ChatID, responses); err != nil { + return err + } + if hookEndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return err + } + endedDescendants = endResult.EndedDescendants + result.Ended = true + refreshed, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after end: %w", err) + } + result.Chat = refreshed + editedCutoffT = refreshed.UpdatedAt + return nil + } + var reasoningEffortOverride database.NullChatReasoningEffort if opts.ReasoningEffort != nil && *opts.ReasoningEffort != "" { reasoningEffortOverride = database.NullChatReasoningEffort{ChatReasoningEffort: database.ChatReasoningEffort(*opts.ReasoningEffort), Valid: true} @@ -1721,6 +1990,8 @@ func (p *Server) EditMessage( editResult, err := tx.EditMessage(chatstate.EditMessageInput{ MessageID: opts.EditedMessageID, + TurnID: turnID, + PrefixMessages: prefixMessages, CreatedBy: opts.CreatedBy, Content: content, ModelConfigIDOverride: modelOverride, @@ -1733,6 +2004,10 @@ func (p *Server) EditMessage( return err } result.Message = editResult.ReplacementMessage + inserted := make([]database.ChatMessage, 0, len(editResult.CancellationMessages)+1) + inserted = append(inserted, editResult.CancellationMessages...) + inserted = append(inserted, editResult.ReplacementMessage) + result.InsertedMessages = inserted // 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. @@ -1750,6 +2025,10 @@ func (p *Server) EditMessage( // Sidebar watch event keeps the chat list responsive. Stream // side effects are handled by chat:update consumers. + if result.Ended { + p.publishEndChatSideEffects(ctx, result.Chat, endedDescendants) + return result, nil + } p.publishChatPubsubEvent(result.Chat, codersdk.ChatWatchEventKindStatusChange, nil) // Editing can race with an interrupted worker still flushing its @@ -1977,21 +2256,152 @@ func (e *ToolResultStatusConflictError) Error() string { ) } -// SubmitToolResults validates and persists client-provided tool -// 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. +type dynamicPostToolUseState struct { + chat database.Chat + turnID *uuid.UUID + modelConfigID uuid.UUID + toolNames map[string]string +} + +func loadDynamicPostToolUseState( + ctx context.Context, + machine *chatstate.ChatMachine, + opts SubmitToolResultsOptions, +) (dynamicPostToolUseState, error) { + var state dynamicPostToolUseState + err := machine.ReadLock(ctx, func(store database.Store) error { + chat, err := store.GetChatByID(ctx, opts.ChatID) + if err != nil { + return xerrors.Errorf("load chat: %w", err) + } + if chat.Archived { + return ErrChatArchived + } + if chat.Status != database.ChatStatusRequiresAction { + return &ToolResultStatusConflictError{ActualStatus: chat.Status} + } + messages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: opts.ChatID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("load chat messages: %w", err) + } + _, pending, err := unresolvedToolCallsFromHistory(messages, dynamicToolNamesFromChat(chat)) + if err != nil { + return xerrors.Errorf("load pending dynamic tool calls: %w", err) + } + toolNames := make(map[string]string, len(pending)) + for _, call := range pending { + toolNames[call.ToolCallID] = call.ToolName + } + if err := validateSubmittedToolResults(opts.Results, toolNames); err != nil { + return err + } + modelConfigID := opts.ModelConfigID + if modelConfigID == uuid.Nil { + modelConfigID = chat.LastModelConfigID + } + state = dynamicPostToolUseState{ + chat: chat, + turnID: activeTurnID(messages), + modelConfigID: modelConfigID, + toolNames: toolNames, + } + return nil + }) + return state, err +} + +func validateSubmittedToolResults(results []codersdk.ToolResult, toolNames map[string]string) error { + submitted := make(map[string]struct{}, len(results)) + for _, result := range results { + if _, ok := submitted[result.ToolCallID]; ok { + return &ToolResultValidationError{ + Message: "Duplicate tool_call_id in results.", + Detail: fmt.Sprintf("Duplicate tool call ID %q.", result.ToolCallID), + } + } + 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), + } + } + if _, ok := toolNames[result.ToolCallID]; !ok { + return &ToolResultValidationError{ + Message: "Unexpected tool result.", + Detail: fmt.Sprintf("No pending tool call with ID %q.", result.ToolCallID), + } + } + submitted[result.ToolCallID] = struct{}{} + } + for toolCallID := range toolNames { + if _, ok := submitted[toolCallID]; !ok { + return &ToolResultValidationError{ + Message: "Missing tool result.", + Detail: fmt.Sprintf("Missing result for tool call %q.", toolCallID), + } + } + } + return nil +} + +func dynamicPostToolUseData(result codersdk.ToolResult, toolName string) agenthooks.PostToolUseData { + data := agenthooks.PostToolUseData{ + ToolUseID: result.ToolCallID, + ToolName: toolName, + } + if result.IsError { + if err := json.Unmarshal(result.Output, &data.ToolError); err != nil { + data.ToolError = string(result.Output) + } + } else { + data.ToolResponse = append(json.RawMessage(nil), result.Output...) + } + return data +} + +// SubmitToolResults dispatches hooks before completing the requires_action transition. func (p *Server) SubmitToolResults( ctx context.Context, opts SubmitToolResultsOptions, ) error { + machine := p.newChatMachine(opts.ChatID) + var hookResponses []agenthooks.Response + var hookSuffix []chatstate.Message + hookEndChat := false + if p.hookDispatcher != nil && p.hookDispatcher.Enabled() { + state, err := loadDynamicPostToolUseState(ctx, machine, opts) + if err != nil { + return err + } + for _, result := range opts.Results { + response, err := p.dispatchPostToolUseData(ctx, state.chat, state.turnID, dynamicPostToolUseData(result, state.toolNames[result.ToolCallID])) + if err != nil { + // Accepted end_chat effects take precedence over later dispatch failures. + if hookEndChat { + return p.endChatAfterToolHookFailure(ctx, machine, opts.ChatID, hookSuffix) + } + // Leave pending calls intact so the client can resubmit after recovery. + return generationHookDispatchError(agenthooks.EventPostToolUse, err) + } + hookResponses = append(hookResponses, response) + responseMessages, err := hookPrefixMessages(response, state.modelConfigID, state.turnID) + if err != nil { + return err + } + hookSuffix = append(hookSuffix, responseMessages...) + hookEndChat = hookEndChat || response.EndChat + } + } + var ( - statusConflict *ToolResultStatusConflictError - refreshChat database.Chat - refreshedOK bool + statusConflict *ToolResultStatusConflictError + refreshChat database.Chat + refreshedOK bool + endedDescendants []database.Chat ) - 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 { @@ -2002,21 +2412,25 @@ func (p *Server) SubmitToolResults( } toolResults := make([]chatstate.ToolResultInput, 0, len(opts.Results)) - for _, r := range opts.Results { + for _, result := range opts.Results { toolResults = append(toolResults, chatstate.ToolResultInput{ - ToolCallID: r.ToolCallID, - Output: r.Output, - IsError: r.IsError, + ToolCallID: result.ToolCallID, + Output: result.Output, + IsError: result.IsError, }) } modelConfigID := opts.ModelConfigID if modelConfigID == uuid.Nil { modelConfigID = locked.LastModelConfigID } + if err := applyHookAllowedToolsResponses(ctx, store, opts.ChatID, hookResponses); err != nil { + return err + } if _, err := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{ - CreatedBy: opts.UserID, - ModelConfigID: modelConfigID, - Results: toolResults, + CreatedBy: opts.UserID, + ModelConfigID: modelConfigID, + Results: toolResults, + SuffixMessages: hookSuffix, }); err != nil { if !errors.Is(err, chatstate.ErrInvalidState) && locked.Status != database.ChatStatusRequiresAction && @@ -2028,9 +2442,13 @@ func (p *Server) SubmitToolResults( } 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. + if hookEndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{}) + if err != nil { + return xerrors.Errorf("end chat from post_tool_use: %w", err) + } + endedDescendants = endResult.EndedDescendants + } refreshed, err := store.GetChatByID(ctx, opts.ChatID) if err != nil { return xerrors.Errorf("reload chat after tool results: %w", err) @@ -2047,7 +2465,11 @@ func (p *Server) SubmitToolResults( } if refreshedOK { - p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) + if refreshChat.Archived { + p.publishEndChatSideEffects(ctx, refreshChat, endedDescendants) + } else { + p.publishChatPubsubEvent(refreshChat, codersdk.ChatWatchEventKindStatusChange, nil) + } } return nil } @@ -2831,6 +3253,7 @@ func appendMessageFields( params *database.InsertChatMessagesParams, msg chatMessage, ) { + params.TurnID = append(params.TurnID, uuid.Nil) params.CreatedBy = append(params.CreatedBy, msg.createdBy) params.ModelConfigID = append(params.ModelConfigID, msg.modelConfigID) params.ReasoningEffort = append(params.ReasoningEffort, "") @@ -2894,6 +3317,7 @@ type Config struct { AllowBYOKSet bool AlwaysEnableDebugLogs bool WebpushDispatcher webpush.Dispatcher + HookDispatcher *chathooks.Dispatcher UsageTracker *workspacestats.UsageTracker Clock quartz.Clock AIBridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] @@ -2961,6 +3385,13 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { if cfg.AllowBYOKSet { allowBYOK = cfg.AllowBYOK } + + // Require the experiment even for injected dispatchers to preserve explicit opt-in. + hookDispatcher := cfg.HookDispatcher + if hookDispatcher != nil && !cfg.Experiments.Enabled(codersdk.ExperimentAgentLifecycleHooks) { + cfg.Logger.Warn(ctx, "ignoring chat lifecycle hook dispatcher; the agent-lifecycle-hooks experiment is not enabled") + hookDispatcher = nil + } p := &Server{ cancel: cancel, db: cfg.Database, @@ -2975,6 +3406,7 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { stopWorkspaceFn: cfg.StopWorkspace, pubsub: ps, webpushDispatcher: cfg.WebpushDispatcher, + hookDispatcher: hookDispatcher, providerAPIKeys: cfg.ProviderAPIKeys, allowBYOK: allowBYOK, oidcTokenSource: cfg.OIDCTokenSource, @@ -3128,6 +3560,12 @@ func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.Ch } } +func (p *Server) publishEndChatSideEffects(ctx context.Context, ended database.Chat, descendants []database.Chat) { + p.scheduleArchiveDebugCleanup(ctx, append([]database.Chat{ended}, descendants...)) + p.publishChatPubsubEvents(descendants, codersdk.ChatWatchEventKindDeleted) + p.publishChatPubsubEvent(ended, codersdk.ChatWatchEventKindDeleted, nil) +} + // chatWatchEventSDKChat builds the chat embedded in ChatWatchEvent // notifications. These payloads travel through PostgreSQL NOTIFY, so // omit fields that can grow large and that watch consumers already read diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go index ff659f1021a68..e9e84687f6268 100644 --- a/coderd/x/chatd/chatloop/chatloop.go +++ b/coderd/x/chatd/chatloop/chatloop.go @@ -274,6 +274,7 @@ type GenerateCompactionOptions struct { ContextLimit int64 ContextLimitFallback int64 SummaryPrompt string + SummaryHint string SystemSummaryPrefix string StepUsage fantasy.Usage StepMetadata fantasy.ProviderMetadata diff --git a/coderd/x/chatd/chatloop/compaction.go b/coderd/x/chatd/chatloop/compaction.go index e0bcef5bb9c63..2304cadd5ee30 100644 --- a/coderd/x/chatd/chatloop/compaction.go +++ b/coderd/x/chatd/chatloop/compaction.go @@ -87,6 +87,7 @@ type CompactionOptions struct { ThresholdPercent int32 ContextLimit int64 SummaryPrompt string + SummaryHint string SystemSummaryPrefix string Persist func(context.Context, CompactionResult) error DebugSvc *chatdebug.Service @@ -213,6 +214,7 @@ func normalizedCompactionGenerateConfig(opts GenerateCompactionOptions) (Compact ThresholdPercent: opts.ThresholdPercent, ContextLimit: opts.ContextLimit, SummaryPrompt: opts.SummaryPrompt, + SummaryHint: opts.SummaryHint, SystemSummaryPrefix: opts.SystemSummaryPrefix, DebugSvc: opts.DebugSvc, ChatID: opts.ChatID, @@ -416,11 +418,13 @@ func generateCompactionSummary( ) (summary string, err error) { summaryPrompt := make([]fantasy.Message, 0, len(messages)+1) summaryPrompt = append(summaryPrompt, messages...) + summaryParts := []fantasy.MessagePart{fantasy.TextPart{Text: options.SummaryPrompt}} + if strings.TrimSpace(options.SummaryHint) != "" { + summaryParts = append(summaryParts, fantasy.TextPart{Text: options.SummaryHint}) + } summaryPrompt = append(summaryPrompt, fantasy.Message{ - Role: fantasy.MessageRoleUser, - Content: []fantasy.MessagePart{ - fantasy.TextPart{Text: options.SummaryPrompt}, - }, + Role: fantasy.MessageRoleUser, + Content: summaryParts, }) toolChoice := fantasy.ToolChoiceNone diff --git a/coderd/x/chatd/chatstate/family_test.go b/coderd/x/chatd/chatstate/family_test.go index 7fcbd18310f63..4e84687647274 100644 --- a/coderd/x/chatd/chatstate/family_test.go +++ b/coderd/x/chatd/chatstate/family_test.go @@ -9,6 +9,7 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + 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" @@ -197,6 +198,155 @@ func TestSetFamilyArchivedAcceptsAlreadyDesiredMembers(t *testing.T) { require.True(t, childAfter.Archived) } +func TestEndChatArchivesChildren(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.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + grandchild := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "grandchild", + Status: database.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: child.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + 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() + machine := chatstate.NewChatMachine(db, pub, root.ID) + var endResult chatstate.EndChatResult + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + var err error + endResult, err = tx.EndChatFamily(chatstate.EndChatInput{}) + return err + })) + + for _, chatID := range []uuid.UUID{root.ID, child.ID, grandchild.ID} { + after, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.True(t, after.Archived, "family member must be archived") + require.Equal(t, database.ChatStatusWaiting, after.Status) + require.False(t, after.WorkerID.Valid, "worker ownership must be cleared") + require.False(t, after.RunnerID.Valid, "runner ownership must be cleared") + } + count, err := db.CountChatQueuedMessages(ctx, child.ID) + require.NoError(t, err) + require.Zero(t, count, "child queue must be cleared") + for _, chatID := range []uuid.UUID{child.ID, grandchild.ID} { + require.Contains(t, pub.channels, coderdpubsub.ChatStateUpdateChannel(chatID), + "ended child must publish a chat:update") + } + endedIDs := make([]uuid.UUID, 0, len(endResult.EndedDescendants)) + for _, desc := range endResult.EndedDescendants { + require.True(t, desc.Archived, "returned descendant must carry the post-transition row") + endedIDs = append(endedIDs, desc.ID) + } + require.ElementsMatch(t, []uuid.UUID{child.ID, grandchild.ID}, endedIDs, + "cascade must surface newly ended descendants for caller side effects") +} + +func TestEndChatOnChildArchivesSubtreeOnly(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.ChatStatusRunning, + }) + 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}, + }) + grandchild := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "grandchild", + Status: database.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: child.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + sibling := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Title: "sibling", + Status: database.ChatStatusRunning, + ParentChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: root.ID, Valid: true}, + }) + + pub := newRecordingPubsub() + machine := chatstate.NewChatMachine(db, pub, child.ID) + var endResult chatstate.EndChatResult + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + var err error + endResult, err = tx.EndChatFamily(chatstate.EndChatInput{}) + return err + })) + + for _, chatID := range []uuid.UUID{child.ID, grandchild.ID} { + after, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.True(t, after.Archived, "subtree member must be archived") + require.Equal(t, database.ChatStatusWaiting, after.Status) + require.False(t, after.WorkerID.Valid) + require.False(t, after.RunnerID.Valid) + } + for _, chatID := range []uuid.UUID{root.ID, sibling.ID} { + after, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.False(t, after.Archived, "root and sibling must stay active") + require.Equal(t, database.ChatStatusRunning, after.Status) + } + require.Contains(t, pub.channels, coderdpubsub.ChatStateUpdateChannel(grandchild.ID), + "ended grandchild must publish a chat:update") + endedIDs := make([]uuid.UUID, 0, len(endResult.EndedDescendants)) + for _, desc := range endResult.EndedDescendants { + endedIDs = append(endedIDs, desc.ID) + } + require.ElementsMatch(t, []uuid.UUID{grandchild.ID}, endedIDs, + "subtree cascade must surface only newly ended descendants") +} + func seedFamilyDeps(t *testing.T, db database.Store) (database.User, database.Organization, database.ChatModelConfig) { t.Helper() user := dbgen.User(t, db, database.User{}) diff --git a/coderd/x/chatd/chatstate/hook_prefix_test.go b/coderd/x/chatd/chatstate/hook_prefix_test.go new file mode 100644 index 0000000000000..ceac1d9eac6aa --- /dev/null +++ b/coderd/x/chatd/chatstate/hook_prefix_test.go @@ -0,0 +1,233 @@ +package chatstate_test + +import ( + "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" +) + +func TestQueuedHookPrefixDeferredUntilPromotion(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) + baseHistory := activeHistoryIDs(ctx, t, f, created.Chat.ID) + + turnID := uuid.New() + prompt := userTextMessage("queued prompt", f.User.ID, f.Model.ID) + prompt.TurnID = uuid.NullUUID{UUID: turnID, Valid: true} + prefixContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("hook context")}) + require.NoError(t, err) + prefix := chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: prefixContent, + Visibility: database.ChatMessageVisibilityModel, + ModelConfigID: uuid.NullUUID{UUID: f.Model.ID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + } + + 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: prompt, + PrefixMessages: []chatstate.Message{prefix}, + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + require.NotNil(t, send.QueuedMessage) + require.True(t, send.QueuedMessage.HookPrefix.Valid) + require.Equal(t, baseHistory, activeHistoryIDs(ctx, t, f, created.Chat.ID), + "queueing a prompt must not change active history") + + 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) + + history := activeHistoryIDs(ctx, t, f, created.Chat.ID) + require.Len(t, history, len(baseHistory)+1, "promotion inserts the prompt") + promptRow := requireChatMessageByID(ctx, t, f, history[len(history)-1]) + require.Equal(t, finish.PromotedMessage.ID, promptRow.ID) + assertChatMessageText(t, promptRow, "queued prompt") + + promptMessages, err := f.DB.GetChatMessagesForPromptByChatID(ctx, created.Chat.ID) + require.NoError(t, err) + prefixRow := promptMessages[len(promptMessages)-2] + require.Equal(t, database.ChatMessageVisibilityModel, prefixRow.Visibility) + assertChatMessageText(t, prefixRow, "hook context") + require.Equal(t, uuid.NullUUID{UUID: turnID, Valid: true}, prefixRow.TurnID, + "prefix adopts the queued prompt's turn ID") + require.Equal(t, promptRow.ID, promptMessages[len(promptMessages)-1].ID, + "prefix lands immediately before the promoted prompt") +} + +func TestQueuedHookAllowedToolsDeferredUntilPromotion(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) + + policy := pqtype.NullRawMessage{RawMessage: []byte(`["read_file"]`), Valid: true} + prompt := userTextMessage("queued prompt", f.User.ID, f.Model.ID) + prompt.TurnID = uuid.NullUUID{UUID: uuid.New(), Valid: true} + + 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: prompt, + HookAllowedTools: policy, + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + require.NotNil(t, send.QueuedMessage) + require.True(t, send.QueuedMessage.HookAllowedTools.Valid) + + chat, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.False(t, chat.HookAllowedTools.Valid, + "queueing a prompt must not change the chat tool policy") + + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + + chat, err = f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.True(t, chat.HookAllowedTools.Valid, "promotion applies the queued prompt's tool policy") + require.JSONEq(t, `["read_file"]`, string(chat.HookAllowedTools.RawMessage)) +} + +func TestNarrowHookAllowedTools(t *testing.T) { + t.Parallel() + + valid := func(tools string) pqtype.NullRawMessage { + return pqtype.NullRawMessage{RawMessage: []byte(tools), Valid: true} + } + + tests := []struct { + name string + current pqtype.NullRawMessage + incoming pqtype.NullRawMessage + want string + wantNull bool + }{ + { + name: "NoIncomingKeepsCurrent", + current: valid(`["a"]`), + incoming: pqtype.NullRawMessage{}, + want: `["a"]`, + }, + { + name: "NoPolicyAdoptsIncoming", + current: pqtype.NullRawMessage{}, + incoming: valid(`["a","b"]`), + want: `["a","b"]`, + }, + { + name: "BothNullStaysNull", + current: pqtype.NullRawMessage{}, + incoming: pqtype.NullRawMessage{}, + wantNull: true, + }, + { + name: "IncomingNarrows", + current: valid(`["a","b","c"]`), + incoming: valid(`["b","c","d"]`), + want: `["b","c"]`, + }, + { + name: "WideningIsIgnored", + current: valid(`["a"]`), + incoming: valid(`["a","b","c"]`), + want: `["a"]`, + }, + { + name: "EmptyPolicyStaysEmpty", + current: valid(`[]`), + incoming: valid(`["a","b"]`), + want: `[]`, + }, + { + name: "IncomingEmptyRestrictsAll", + current: valid(`["a","b"]`), + incoming: valid(`[]`), + want: `[]`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, err := chatstate.NarrowHookAllowedTools(tt.current, tt.incoming) + require.NoError(t, err) + if tt.wantNull { + require.False(t, got.Valid) + return + } + require.True(t, got.Valid) + require.JSONEq(t, tt.want, string(got.RawMessage)) + }) + } +} + +func TestQueuedHookPolicyCannotWidenOnPromotion(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) + + prompt := userTextMessage("queued prompt", f.User.ID, f.Model.ID) + prompt.TurnID = uuid.NullUUID{UUID: uuid.New(), Valid: true} + 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: prompt, + HookAllowedTools: pqtype.NullRawMessage{RawMessage: []byte(`["read_file","execute"]`), Valid: true}, + BusyBehavior: chatstate.BusyBehaviorQueue, + }) + return err + })) + require.NotNil(t, send.QueuedMessage) + + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + return store.UpdateChatHookAllowedTools(ctx, database.UpdateChatHookAllowedToolsParams{ + HookAllowedTools: pqtype.NullRawMessage{RawMessage: []byte(`["read_file"]`), Valid: true}, + ID: created.Chat.ID, + }) + })) + + require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + _, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + return err + })) + + chat, err := f.DB.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.True(t, chat.HookAllowedTools.Valid) + require.JSONEq(t, `["read_file"]`, string(chat.HookAllowedTools.RawMessage), + "promotion must intersect with the narrower live policy, not replace it") +} diff --git a/coderd/x/chatd/chatstate/machine.go b/coderd/x/chatd/chatstate/machine.go index afe85ae1dac74..8a326eea16ef1 100644 --- a/coderd/x/chatd/chatstate/machine.go +++ b/coderd/x/chatd/chatstate/machine.go @@ -63,6 +63,10 @@ type Tx struct { ctx context.Context store database.Store chatID uuid.UUID + // publisher holds the active [PublishBuffer], not the live pubsub. + // Nested machines wrap it so their publications flush only after the + // outer transaction commits. + publisher Publisher } // Ctx returns the context the surrounding [ChatMachine.Update] call @@ -169,9 +173,10 @@ func (m *ChatMachine) Update( return xerrors.Errorf("lock chat and bump snapshot: %w", err) } tx := &Tx{ - ctx: ctx, - store: store, - chatID: m.chatID, + ctx: ctx, + store: store, + chatID: m.chatID, + publisher: buffer, } if err := fn(tx, store); err != nil { return err diff --git a/coderd/x/chatd/chatstate/messages.go b/coderd/x/chatd/chatstate/messages.go index b867c1f05eac8..21500b42981d6 100644 --- a/coderd/x/chatd/chatstate/messages.go +++ b/coderd/x/chatd/chatstate/messages.go @@ -21,6 +21,7 @@ type Message struct { Role database.ChatMessageRole Content pqtype.NullRawMessage Visibility database.ChatMessageVisibility + TurnID uuid.NullUUID ModelConfigID uuid.NullUUID ReasoningEffort database.NullChatReasoningEffort CreatedBy uuid.NullUUID @@ -47,6 +48,7 @@ func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMes n := len(messages) params := database.InsertChatMessagesParams{ ChatID: chatID, + TurnID: make([]uuid.UUID, n), CreatedBy: make([]uuid.UUID, n), ModelConfigID: make([]uuid.UUID, n), ReasoningEffort: make([]string, n), @@ -66,6 +68,7 @@ func toInsertParams(chatID uuid.UUID, messages []Message) database.InsertChatMes RuntimeMs: make([]int64, n), } for i, m := range messages { + params.TurnID[i] = nullUUIDOrNil(m.TurnID) params.CreatedBy[i] = nullUUIDOrNil(m.CreatedBy) params.ModelConfigID[i] = nullUUIDOrNil(m.ModelConfigID) if m.ReasoningEffort.Valid { diff --git a/coderd/x/chatd/chatstate/synthetics.go b/coderd/x/chatd/chatstate/synthetics.go index d442843d17c3f..9702478243afa 100644 --- a/coderd/x/chatd/chatstate/synthetics.go +++ b/coderd/x/chatd/chatstate/synthetics.go @@ -31,6 +31,8 @@ import ( // The synthetic results use the supplied chat's last_model_config_id. // Returns (nil, nil) when there is nothing to synthesize. // +// pendingInserts count as handled to avoid canceling calls resolved in the same transaction. +// //nolint:revive // dynamicOnly is a domain flag, not a control flag. func synthesizePendingToolCancellations( ctx context.Context, @@ -38,6 +40,7 @@ func synthesizePendingToolCancellations( chat database.Chat, reason string, dynamicOnly bool, + pendingInserts ...Message, ) ([]Message, error) { var dynamicToolNames map[string]bool if dynamicOnly { @@ -99,6 +102,24 @@ func synthesizePendingToolCancellations( } } } + for _, msg := range pendingInserts { + if msg.Role != database.ChatMessageRoleTool { + continue + } + parts, err := chatprompt.ParseContent(database.ChatMessage{ + Role: msg.Role, + ContentVersion: msg.ContentVersion, + Content: msg.Content, + }) + if err != nil { + 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 { @@ -142,6 +163,62 @@ func synthesizePendingToolCancellations( return out, nil } +// synthesizeBatchToolCancellations returns cancellations that must be +// inserted after the batch so every result follows its call. +func synthesizeBatchToolCancellations(chat database.Chat, reason string, batch []Message) ([]Message, error) { + handled := make(map[string]bool) + var calls []codersdk.ChatMessagePart + for _, msg := range batch { + parts, err := chatprompt.ParseContent(database.ChatMessage{ + Role: msg.Role, + ContentVersion: msg.ContentVersion, + Content: msg.Content, + }) + if err != nil { + continue + } + for _, part := range parts { + switch part.Type { + case codersdk.ChatMessagePartTypeToolCall: + if msg.Role == database.ChatMessageRoleAssistant && !part.ProviderExecuted { + calls = append(calls, part) + } + case codersdk.ChatMessagePartTypeToolResult: + handled[part.ToolCallID] = true + default: + } + } + } + out := make([]Message, 0) + for _, part := range calls { + 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 diff --git a/coderd/x/chatd/chatstate/transition.go b/coderd/x/chatd/chatstate/transition.go index f7b6c3634c5c2..565a23ae19cfb 100644 --- a/coderd/x/chatd/chatstate/transition.go +++ b/coderd/x/chatd/chatstate/transition.go @@ -12,6 +12,7 @@ type Transition string const ( TransitionCreateChat Transition = "CreateChat" TransitionSetArchived Transition = "SetArchived" + TransitionEndChat Transition = "EndChat" TransitionSendMessage Transition = "SendMessage" TransitionEditMessage Transition = "EditMessage" TransitionRequestCompaction Transition = "RequestCompaction" @@ -28,6 +29,7 @@ const ( TransitionFinishInterruption Transition = "FinishInterruption" TransitionFinishTurn Transition = "FinishTurn" TransitionFinishError Transition = "FinishError" + TransitionFailIdle Transition = "FailIdle" TransitionCancelRequiresAction Transition = "CancelRequiresAction" TransitionReconcileInvalidState Transition = "ReconcileInvalidState" ) @@ -43,6 +45,7 @@ func (t Transition) String() string { return string(t) } var AllExecutionTransitions = []Transition{ TransitionCreateChat, TransitionSetArchived, + TransitionEndChat, TransitionSendMessage, TransitionEditMessage, TransitionRequestCompaction, @@ -57,6 +60,7 @@ var AllExecutionTransitions = []Transition{ TransitionFinishInterruption, TransitionFinishTurn, TransitionFinishError, + TransitionFailIdle, TransitionCancelRequiresAction, TransitionReconcileInvalidState, } @@ -77,17 +81,21 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionCreateChat: {StateR0}, }, StateW: { + TransitionEndChat: {StateXW}, TransitionSetArchived: {StateXW}, TransitionSendMessage: {StateR0}, TransitionEditMessage: {StateR0}, TransitionRequestCompaction: {StateR0}, + TransitionFailIdle: {StateE0}, }, StateE0: { + TransitionEndChat: {StateXW}, TransitionSetArchived: {StateXE0}, TransitionSendMessage: {StateR0}, TransitionEditMessage: {StateR0}, }, StateE1: { + TransitionEndChat: {StateXW}, TransitionSetArchived: {StateXE1}, TransitionSendMessage: {StateR1}, TransitionEditMessage: {StateR0}, @@ -95,6 +103,7 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionPromoteQueuedMessage: {StateR0, StateR1}, }, StateR0: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateR1, StateI1}, TransitionEditMessage: {StateR0}, TransitionInterrupt: {StateI0}, @@ -106,6 +115,7 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionFinishError: {StateE0}, }, StateR1: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateR1, StateI1}, TransitionEditMessage: {StateR0}, TransitionDeleteQueuedMessage: {StateR0, StateR1}, @@ -119,11 +129,13 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionFinishError: {StateE1}, }, StateI0: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateI1}, TransitionEditMessage: {StateR0}, TransitionFinishInterruption: {StateW}, }, StateI1: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateI1}, TransitionEditMessage: {StateR0}, TransitionDeleteQueuedMessage: {StateI0, StateI1}, @@ -131,6 +143,7 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionFinishInterruption: {StateR0, StateR1}, }, StateA0: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateA1, StateR1}, TransitionEditMessage: {StateR0}, TransitionInterrupt: {StateR0}, @@ -138,6 +151,7 @@ var transitionMatrix = map[ExecutionState]map[Transition][]ExecutionState{ TransitionCancelRequiresAction: {StateR0}, }, StateA1: { + TransitionEndChat: {StateXW}, TransitionSendMessage: {StateA1, StateR1}, TransitionEditMessage: {StateR0}, TransitionDeleteQueuedMessage: {StateA0, StateA1}, diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go index 6b8593eff8322..b119eac864594 100644 --- a/coderd/x/chatd/chatstate/transitions.go +++ b/coderd/x/chatd/chatstate/transitions.go @@ -60,6 +60,32 @@ func CreateChat( store database.Store, publisher Publisher, input CreateChatInput, +) (CreateChatResult, error) { + return insertChat(ctx, store, publisher, uuid.NullUUID{}, pqtype.NullRawMessage{}, input) +} + +// CreateChatWithID creates a chat using a caller-minted ID and initial hook tool policy. +func CreateChatWithID( + ctx context.Context, + store database.Store, + publisher Publisher, + chatID uuid.UUID, + hookAllowedTools pqtype.NullRawMessage, + input CreateChatInput, +) (CreateChatResult, error) { + if chatID == uuid.Nil { + return CreateChatResult{}, xerrors.New("chatstate: CreateChatWithID called with nil chat ID") + } + return insertChat(ctx, store, publisher, uuid.NullUUID{UUID: chatID, Valid: true}, hookAllowedTools, input) +} + +func insertChat( + ctx context.Context, + store database.Store, + publisher Publisher, + chatID uuid.NullUUID, + hookAllowedTools pqtype.NullRawMessage, + input CreateChatInput, ) (CreateChatResult, error) { if store == nil { return CreateChatResult{}, xerrors.New("chatstate: CreateChat called with nil store") @@ -78,6 +104,7 @@ func CreateChat( defer buffer.Discard() err := store.InTx(func(store database.Store) error { chat, err := store.InsertChat(ctx, database.InsertChatParams{ + ID: chatID, OrganizationID: input.OrganizationID, OwnerID: input.OwnerID, WorkspaceID: input.WorkspaceID, @@ -94,6 +121,7 @@ func CreateChat( Labels: input.Labels, DynamicTools: input.DynamicTools, ClientType: input.ClientType, + HookAllowedTools: hookAllowedTools, }) if err != nil { return xerrors.Errorf("insert chat: %w", err) @@ -227,7 +255,8 @@ func (tx *Tx) requireQueueCapacity() error { // 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) { +// Hook prefix messages remain queued until their prompt is promoted. +func (tx *Tx) insertQueuedMessage(ownerFallback uuid.UUID, m Message, hookPrefix []Message, hookAllowedTools pqtype.NullRawMessage) (database.ChatQueuedMessage, error) { createdBy := ownerFallback if m.CreatedBy.Valid { createdBy = m.CreatedBy.UUID @@ -236,18 +265,150 @@ func (tx *Tx) insertQueuedMessage(ownerFallback uuid.UUID, m Message) (database. if !m.Content.Valid || len(rawContent) == 0 { rawContent = json.RawMessage("null") } + encodedPrefix, err := encodeQueuedHookPrefix(hookPrefix) + if err != nil { + return database.ChatQueuedMessage{}, err + } 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, - ReasoningEffort: m.ReasoningEffort, - CreatedBy: createdBy, + ChatID: tx.chatID, + TurnID: m.TurnID, + Content: rawContent, + ModelConfigID: m.ModelConfigID, + ReasoningEffort: m.ReasoningEffort, + CreatedBy: createdBy, + HookPrefix: encodedPrefix, + HookAllowedTools: hookAllowedTools, }) } +// applyHookAllowedTools intersects prompt policy with current policy so replay cannot widen restrictions. +func (tx *Tx) applyHookAllowedTools(allowedTools pqtype.NullRawMessage) error { + if !allowedTools.Valid { + return nil + } + chat, err := tx.store.GetChatByID(tx.ctx, tx.chatID) + if err != nil { + return xerrors.Errorf("load chat for hook allowed tools: %w", err) + } + narrowed, err := NarrowHookAllowedTools(chat.HookAllowedTools, allowedTools) + if err != nil { + return err + } + if err := tx.store.UpdateChatHookAllowedTools(tx.ctx, database.UpdateChatHookAllowedToolsParams{ + HookAllowedTools: narrowed, + ID: tx.chatID, + }); err != nil { + return xerrors.Errorf("apply hook allowed tools: %w", err) + } + return nil +} + +// NarrowHookAllowedTools intersects an incoming policy with the current policy. +// A missing current policy adopts the incoming value; an existing policy only shrinks. +func NarrowHookAllowedTools(current, incoming pqtype.NullRawMessage) (pqtype.NullRawMessage, error) { + if !incoming.Valid { + return current, nil + } + var incomingTools []string + if err := json.Unmarshal(incoming.RawMessage, &incomingTools); err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("decode incoming hook allowed tools: %w", err) + } + if !current.Valid { + return incoming, nil + } + var currentTools []string + if err := json.Unmarshal(current.RawMessage, ¤tTools); err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("decode current hook allowed tools: %w", err) + } + allowed := make(map[string]struct{}, len(currentTools)) + for _, tool := range currentTools { + allowed[tool] = struct{}{} + } + intersection := make([]string, 0, len(incomingTools)) + for _, tool := range incomingTools { + if _, ok := allowed[tool]; ok { + intersection = append(intersection, tool) + } + } + encoded, err := json.Marshal(intersection) + if err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("marshal narrowed hook allowed tools: %w", err) + } + return pqtype.NullRawMessage{RawMessage: encoded, Valid: true}, nil +} + +// queuedHookPrefixMessage adopts the queued prompt's turn ID at promotion. +type queuedHookPrefixMessage struct { + Role database.ChatMessageRole `json:"role"` + Content json.RawMessage `json:"content"` + Visibility database.ChatMessageVisibility `json:"visibility"` + ModelConfigID *uuid.UUID `json:"model_config_id,omitempty"` + CreatedBy *uuid.UUID `json:"created_by,omitempty"` + ContentVersion int16 `json:"content_version"` +} + +func encodeQueuedHookPrefix(prefix []Message) (pqtype.NullRawMessage, error) { + if len(prefix) == 0 { + return pqtype.NullRawMessage{}, nil + } + encoded := make([]queuedHookPrefixMessage, 0, len(prefix)) + for _, m := range prefix { + e := queuedHookPrefixMessage{ + Role: m.Role, + Visibility: m.Visibility, + ContentVersion: m.ContentVersion, + } + if m.Content.Valid { + e.Content = m.Content.RawMessage + } + if m.ModelConfigID.Valid { + id := m.ModelConfigID.UUID + e.ModelConfigID = &id + } + if m.CreatedBy.Valid { + id := m.CreatedBy.UUID + e.CreatedBy = &id + } + encoded = append(encoded, e) + } + raw, err := json.Marshal(encoded) + if err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("marshal queued hook prefix: %w", err) + } + return pqtype.NullRawMessage{RawMessage: raw, Valid: true}, nil +} + +func decodeQueuedHookPrefix(q database.ChatQueuedMessage) ([]Message, error) { + if !q.HookPrefix.Valid { + return nil, nil + } + var encoded []queuedHookPrefixMessage + if err := json.Unmarshal(q.HookPrefix.RawMessage, &encoded); err != nil { + return nil, xerrors.Errorf("unmarshal queued hook prefix: %w", err) + } + messages := make([]Message, 0, len(encoded)) + for _, e := range encoded { + m := Message{ + Role: e.Role, + Content: pqtype.NullRawMessage{RawMessage: e.Content, Valid: len(e.Content) > 0}, + Visibility: e.Visibility, + TurnID: q.TurnID, + ContentVersion: e.ContentVersion, + } + if e.ModelConfigID != nil { + m.ModelConfigID = uuid.NullUUID{UUID: *e.ModelConfigID, Valid: true} + } + if e.CreatedBy != nil { + m.CreatedBy = uuid.NullUUID{UUID: *e.CreatedBy, Valid: true} + } + messages = append(messages, m) + } + return messages, nil +} + // messageFromQueuedRow synthesizes a Message from a stored queued row, // suitable for promoting into active history. func messageFromQueuedRow(q database.ChatQueuedMessage) Message { @@ -255,6 +416,7 @@ func messageFromQueuedRow(q database.ChatQueuedMessage) Message { Role: database.ChatMessageRoleUser, Content: pqtype.NullRawMessage{RawMessage: q.Content, Valid: q.Content != nil}, Visibility: database.ChatMessageVisibilityBoth, + TurnID: q.TurnID, ModelConfigID: q.ModelConfigID, ReasoningEffort: q.ReasoningEffort, CreatedBy: uuid.NullUUID{UUID: q.CreatedBy, Valid: true}, @@ -301,6 +463,125 @@ func (tx *Tx) SetArchived(input SetArchivedInput) (SetArchivedResult, error) { return SetArchivedResult{}, nil } +// EndChatInput configures [Tx.EndChatFamily]. +type EndChatInput struct { + PrefixMessages []Message +} + +// EndChatResult is returned by [Tx.EndChatFamily]. +type EndChatResult struct { + Chat database.Chat + InsertedMessages []database.ChatMessage + DeletedQueuedMessageIDs []int64 + // EndedDescendants contains newly archived descendants for post-commit side effects. + EndedDescendants []database.Chat +} + +// EndChatFamily orchestrates the single-chat EndChat transition across the +// addressed chat and each of its descendants, archiving them and clearing +// active execution state. +func (tx *Tx) EndChatFamily(input EndChatInput) (EndChatResult, error) { + chat, _, err := tx.requireFromAllowed(TransitionEndChat) + if err != nil { + return EndChatResult{}, err + } + result, err := tx.applyEndChat(chat, input.PrefixMessages) + if err != nil { + return EndChatResult{}, err + } + descendants, err := tx.endDescendantChats(chat) + if err != nil { + return EndChatResult{}, err + } + result.EndedDescendants = descendants + return result, nil +} + +func (tx *Tx) applyEndChat(chat database.Chat, prefixMessages []Message) (EndChatResult, error) { + const endChatCancelReason = "Tool execution interrupted because the chat was ended" + cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, endChatCancelReason, false, prefixMessages...) + if err != nil { + return EndChatResult{}, err + } + // Batch cancellations follow the calls they resolve. + batchCancels, err := synthesizeBatchToolCancellations(chat, endChatCancelReason, prefixMessages) + if err != nil { + return EndChatResult{}, err + } + inserted, err := tx.insertMessages(append(append(cancels, prefixMessages...), batchCancels...)) + if err != nil { + return EndChatResult{}, xerrors.Errorf("insert end chat prefix messages: %w", err) + } + deletedQueuedIDs, err := tx.clearQueue() + if err != nil { + return EndChatResult{}, err + } + updated, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusWaiting, + Archived: true, + WorkerID: uuid.NullUUID{}, + RunnerID: uuid.NullUUID{}, + LastError: pqtype.NullRawMessage{}, + RequiresActionDeadlineAt: sql.NullTime{}, + }) + if err != nil { + return EndChatResult{}, xerrors.Errorf("end chat: %w", err) + } + return EndChatResult{ + Chat: updated, + InsertedMessages: inserted, + DeletedQueuedMessageIDs: deletedQueuedIDs, + }, nil +} + +// endDescendantChats uses nested machines so descendant snapshots and +// publications share the outer transaction. IDs follow the family lock order. +func (tx *Tx) endDescendantChats(chat database.Chat) ([]database.Chat, error) { + var ids []uuid.UUID + var err error + if chat.ParentChatID.Valid { + ids, err = tx.store.GetChatDescendantIDsByChatID(tx.ctx, chat.ID) + } else { + ids, err = tx.store.GetChatFamilyIDsByRootID(tx.ctx, chat.ID) + } + if err != nil { + return nil, xerrors.Errorf("get chat descendants: %w", err) + } + var ended []database.Chat + for _, id := range ids { + if id == chat.ID { + continue + } + machine := NewChatMachine(tx.store, tx.publisher, id) + err := machine.Update(tx.ctx, func(child *Tx, _ database.Store) error { + chat, from, err := child.loadState() + if err != nil { + return err + } + // Preserve invalid-state detection for archived descendants. + if from == StateInvalid { + return ErrInvalidState + } + if chat.Archived { + return nil + } + if err := requireExecutionTransition(TransitionEndChat, from); err != nil { + return err + } + result, err := child.applyEndChat(chat, nil) + if err != nil { + return err + } + ended = append(ended, result.Chat) + return nil + }) + if err != nil { + return nil, xerrors.Errorf("end child chat %s: %w", id, err) + } + } + return ended, nil +} + // BusyBehavior controls how SendMessage behaves when the chat is // currently busy (R*/I*/A*). From idle/error states the two behaviors // are equivalent. @@ -313,8 +594,11 @@ const ( // SendMessageInput configures [Tx.SendMessage]. type SendMessageInput struct { - Message Message - BusyBehavior BusyBehavior + Message Message + PrefixMessages []Message + // HookAllowedTools applies when the prompt enters history. + HookAllowedTools pqtype.NullRawMessage + BusyBehavior BusyBehavior } // SendMessageResult is returned by [Tx.SendMessage]. @@ -355,51 +639,54 @@ func (tx *Tx) SendMessage(input SendMessageInput) (SendMessageResult, error) { // Idle / empty-queue error: insert directly into history, clear // last_error, leave queue alone. case StateW, StateE0: - return tx.sendMessageDirect(chat, input.Message) + return tx.sendMessageDirect(chat, input) // Error-with-queue: append to tail, promote previous head into // history, clear last_error. case StateE1: - return tx.sendMessageE1(chat, input.Message) + return tx.sendMessageE1(chat, input) // 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, database.ChatStatusInterrupting, chat.LastError, chat.RequiresActionDeadlineAt) } - return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + return tx.sendMessageQueueAndSetStatus(chat, input, 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, database.ChatStatusInterrupting, chat.LastError, chat.RequiresActionDeadlineAt) } - return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + return tx.sendMessageQueueAndSetStatus(chat, input, 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) + return tx.sendMessageQueueAndSetStatus(chat, input, 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.sendMessageInterruptRequiresAction(chat, input) } - return tx.sendMessageQueueAndSetStatus(chat, input.Message, chat.Status, chat.LastError, chat.RequiresActionDeadlineAt) + return tx.sendMessageQueueAndSetStatus(chat, input, 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) { +func (tx *Tx) sendMessageDirect(chat database.Chat, input SendMessageInput) (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)) + inserted, err := tx.insertMessages(append(append(cancels, input.PrefixMessages...), input.Message)) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert direct user message: %w", err) } + if err := tx.applyHookAllowedTools(input.HookAllowedTools); err != nil { + return SendMessageResult{}, err + } if _, err := tx.applyExecutionState(executionStateUpdate{ Status: database.ChatStatusRunning, Archived: false, @@ -415,8 +702,8 @@ func (tx *Tx) sendMessageDirect(chat database.Chat, m Message) (SendMessageResul }, nil } -func (tx *Tx) sendMessageE1(chat database.Chat, m Message) (SendMessageResult, error) { - queued, err := tx.insertQueuedMessage(chat.OwnerID, m) +func (tx *Tx) sendMessageE1(chat database.Chat, input SendMessageInput) (SendMessageResult, error) { + queued, err := tx.insertQueuedMessage(chat.OwnerID, input.Message, input.PrefixMessages, input.HookAllowedTools) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert queued: %w", err) } @@ -428,11 +715,18 @@ func (tx *Tx) sendMessageE1(chat database.Chat, m Message) (SendMessageResult, e if err != nil { return SendMessageResult{}, err } + headPrefix, err := decodeQueuedHookPrefix(head) + if err != nil { + return SendMessageResult{}, err + } promoted := messageFromQueuedRow(head) - inserted, err := tx.insertMessages(append(cancels, promoted)) + inserted, err := tx.insertMessages(append(append(cancels, headPrefix...), promoted)) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert promoted queued head: %w", err) } + if err := tx.applyHookAllowedTools(head.HookAllowedTools); err != nil { + return SendMessageResult{}, err + } if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ ID: head.ID, ChatID: tx.chatID, @@ -457,12 +751,12 @@ func (tx *Tx) sendMessageE1(chat database.Chat, m Message) (SendMessageResult, e func (tx *Tx) sendMessageQueueAndSetStatus( chat database.Chat, - m Message, + input SendMessageInput, status database.ChatStatus, lastError pqtype.NullRawMessage, deadline sql.NullTime, ) (SendMessageResult, error) { - queued, err := tx.insertQueuedMessage(chat.OwnerID, m) + queued, err := tx.insertQueuedMessage(chat.OwnerID, input.Message, input.PrefixMessages, input.HookAllowedTools) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert queued: %w", err) } @@ -485,7 +779,7 @@ func (tx *Tx) sendMessageQueueAndSetStatus( }, nil } -func (tx *Tx) sendMessageInterruptRequiresAction(chat database.Chat, m Message) (SendMessageResult, error) { +func (tx *Tx) sendMessageInterruptRequiresAction(chat database.Chat, input SendMessageInput) (SendMessageResult, error) { cancels, err := synthesizePendingToolCancellations(tx.ctx, tx.store, chat, "Tool execution interrupted by user message", true) if err != nil { return SendMessageResult{}, err @@ -493,12 +787,14 @@ func (tx *Tx) sendMessageInterruptRequiresAction(chat database.Chat, m Message) 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{}) + return tx.sendMessageQueueAndSetStatus(chat, input, database.ChatStatusRunning, chat.LastError, sql.NullTime{}) } // EditMessageInput configures [Tx.EditMessage]. type EditMessageInput struct { MessageID int64 + TurnID uuid.UUID + PrefixMessages []Message CreatedBy uuid.UUID Content pqtype.NullRawMessage ModelConfigIDOverride uuid.NullUUID @@ -569,7 +865,7 @@ func (tx *Tx) EditMessage(input EditMessageInput) (EditMessageResult, error) { if err != nil { return EditMessageResult{}, err } - cancellationMessages, err := tx.insertMessages(cancels) + cancellationMessages, err := tx.insertMessages(append(cancels, input.PrefixMessages...)) if err != nil { return EditMessageResult{}, xerrors.Errorf("insert message edit cancellations: %w", err) } @@ -586,6 +882,7 @@ func (tx *Tx) EditMessage(input EditMessageInput) (EditMessageResult, error) { Role: database.ChatMessageRoleUser, Content: input.Content, Visibility: target.Visibility, + TurnID: uuid.NullUUID{UUID: input.TurnID, Valid: input.TurnID != uuid.Nil}, ModelConfigID: modelConfig, ReasoningEffort: reasoningEffort, CreatedBy: uuid.NullUUID{UUID: input.CreatedBy, Valid: true}, @@ -770,15 +1067,22 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu if err != nil { return PromoteQueuedMessageResult{}, err } + targetPrefix, err := decodeQueuedHookPrefix(target) + if err != nil { + return PromoteQueuedMessageResult{}, err + } promotedMsg := messageFromQueuedRow(target) - inserted, err := tx.insertMessages(append(cancels, promotedMsg)) + inserted, err := tx.insertMessages(append(append(cancels, targetPrefix...), promotedMsg)) if err != nil { return PromoteQueuedMessageResult{}, xerrors.Errorf("insert promoted queued message: %w", err) } - if len(inserted) != len(cancels)+1 { + if err := tx.applyHookAllowedTools(target.HookAllowedTools); err != nil { + return PromoteQueuedMessageResult{}, err + } + if len(inserted) != len(cancels)+len(targetPrefix)+1 { return PromoteQueuedMessageResult{}, xerrors.Errorf( "insert promoted queued message: expected %d rows, got %d", - len(cancels)+1, len(inserted), + len(cancels)+len(targetPrefix)+1, len(inserted), ) } if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ @@ -797,7 +1101,7 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu }); err != nil { return PromoteQueuedMessageResult{}, xerrors.Errorf("set running: %w", err) } - cancellations := inserted[:len(inserted)-1] + cancellations := inserted[:len(cancels)] insertedUserMsg := inserted[len(inserted)-1] return PromoteQueuedMessageResult{ QueuedMessage: target, @@ -877,9 +1181,10 @@ type ToolResultInput struct { // CompleteRequiresActionInput configures [Tx.CompleteRequiresAction]. type CompleteRequiresActionInput struct { - CreatedBy uuid.UUID - ModelConfigID uuid.UUID - Results []ToolResultInput + CreatedBy uuid.UUID + ModelConfigID uuid.UUID + Results []ToolResultInput + SuffixMessages []Message } // CompleteRequiresActionResult is returned by [Tx.CompleteRequiresAction]. @@ -957,7 +1262,7 @@ func (tx *Tx) CompleteRequiresAction(input CompleteRequiresActionInput) (Complet ContentVersion: chatprompt.CurrentContentVersion, }) } - inserted, err := tx.insertMessages(messages) + inserted, err := tx.insertMessages(append(messages, input.SuffixMessages...)) if err != nil { return CompleteRequiresActionResult{}, xerrors.Errorf("insert tool results: %w", err) } @@ -1275,11 +1580,18 @@ func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterrupt if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("get queue head: %w", err) } + headPrefix, err := decodeQueuedHookPrefix(head) + if err != nil { + return FinishInterruptionResult{}, err + } promotedMsg := messageFromQueuedRow(head) - insertedHead, err := tx.insertMessages([]Message{promotedMsg}) + insertedHead, err := tx.insertMessages(append(headPrefix, promotedMsg)) if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("insert promoted queue head: %w", err) } + if err := tx.applyHookAllowedTools(head.HookAllowedTools); err != nil { + return FinishInterruptionResult{}, err + } if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ ID: head.ID, ChatID: tx.chatID, @@ -1298,8 +1610,8 @@ func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterrupt } insertedPartial = append(insertedPartial, insertedHead...) var promoted *database.ChatMessage - if len(insertedHead) == 1 { - promoted = &insertedHead[0] + if len(insertedHead) > 0 { + promoted = &insertedHead[len(insertedHead)-1] } return FinishInterruptionResult{ InsertedMessages: insertedPartial, @@ -1345,11 +1657,18 @@ func (tx *Tx) FinishTurn(_ FinishTurnInput) (FinishTurnResult, error) { if err != nil { return FinishTurnResult{}, err } + headPrefix, err := decodeQueuedHookPrefix(head) + if err != nil { + return FinishTurnResult{}, err + } promotedMsg := messageFromQueuedRow(head) - inserted, err := tx.insertMessages(append(cancels, promotedMsg)) + inserted, err := tx.insertMessages(append(append(cancels, headPrefix...), promotedMsg)) if err != nil { return FinishTurnResult{}, xerrors.Errorf("insert promoted queue head: %w", err) } + if err := tx.applyHookAllowedTools(head.HookAllowedTools); err != nil { + return FinishTurnResult{}, err + } if _, err := tx.store.DeleteChatQueuedMessageReturningCount(tx.ctx, database.DeleteChatQueuedMessageReturningCountParams{ ID: head.ID, ChatID: tx.chatID, @@ -1404,6 +1723,46 @@ func (tx *Tx) FinishError(input FinishErrorInput) (FinishErrorResult, error) { return FinishErrorResult{}, nil } +// FailIdleInput configures [Tx.FailIdle]. +type FailIdleInput struct { + LastError string + // Kind classifies the persisted error; empty means generic. + Kind codersdk.ChatErrorKind +} + +// FailIdleResult is returned by [Tx.FailIdle]. +type FailIdleResult struct{} + +// FailIdle moves a waiting chat to error without requiring runner ownership. +func (tx *Tx) FailIdle(input FailIdleInput) (FailIdleResult, error) { + chat, _, err := tx.requireFromAllowed(TransitionFailIdle) + if err != nil { + return FailIdleResult{}, err + } + kind := input.Kind + if kind == "" { + kind = codersdk.ChatErrorKindGeneric + } + lastError, err := json.Marshal(codersdk.ChatError{ + Message: input.LastError, + Kind: kind, + }) + if err != nil { + return FailIdleResult{}, xerrors.Errorf("encode last error: %w", err) + } + if _, err := tx.applyExecutionState(executionStateUpdate{ + Status: database.ChatStatusError, + Archived: false, + WorkerID: chat.WorkerID, + RunnerID: chat.RunnerID, + LastError: pqtype.NullRawMessage{RawMessage: lastError, Valid: true}, + RequiresActionDeadlineAt: sql.NullTime{}, + }); err != nil { + return FailIdleResult{}, xerrors.Errorf("set error: %w", err) + } + return FailIdleResult{}, nil +} + // CancelRequiresActionInput configures [Tx.CancelRequiresAction]. type CancelRequiresActionInput struct { Reason string diff --git a/coderd/x/chatd/chatstate/transitions_matrix_test.go b/coderd/x/chatd/chatstate/transitions_matrix_test.go index 1f1597f498bf0..948615b7d36ed 100644 --- a/coderd/x/chatd/chatstate/transitions_matrix_test.go +++ b/coderd/x/chatd/chatstate/transitions_matrix_test.go @@ -108,6 +108,13 @@ func applySetArchived(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededCh return err } +func applyEndChat(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.endChat, err = tx.EndChatFamily(chatstate.EndChatInput{}) + return err +} + func applySendMessageQueue(t *testing.T, f *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { t.Helper() var err error @@ -262,6 +269,13 @@ func applyFinishError(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededCh return err } +func applyFailIdle(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { + t.Helper() + var err error + result.failIdle, err = tx.FailIdle(chatstate.FailIdleInput{LastError: "hook dispatch failed"}) + return err +} + func applyCancelRequiresAction(t *testing.T, _ *testFixture, tx *chatstate.Tx, _ seededChat, _ chatstate.ExecutionState, result *transitionCaseResult) error { t.Helper() var err error @@ -285,6 +299,8 @@ func defaultApplier(tr chatstate.Transition) applierFn { switch tr { case chatstate.TransitionSetArchived: return applySetArchived + case chatstate.TransitionEndChat: + return applyEndChat case chatstate.TransitionSendMessage: return applySendMessageQueue case chatstate.TransitionEditMessage: @@ -313,6 +329,8 @@ func defaultApplier(tr chatstate.Transition) applierFn { return applyFinishTurn case chatstate.TransitionFinishError: return applyFinishError + case chatstate.TransitionFailIdle: + return applyFailIdle case chatstate.TransitionCancelRequiresAction: return applyCancelRequiresAction case chatstate.TransitionReconcileInvalidState: @@ -340,6 +358,7 @@ func mustMarshalParts(t *testing.T, parts []codersdk.ChatMessagePart) pqtype.Nul // AllowedExecutionTransitionsFrom and AllowedExecutionTransitionOutputs. type transitionCaseResult struct { + endChat chatstate.EndChatResult sendMessage chatstate.SendMessageResult editMessage chatstate.EditMessageResult requestCompaction chatstate.RequestCompactionResult @@ -354,6 +373,7 @@ type transitionCaseResult struct { finishInterruption chatstate.FinishInterruptionResult finishTurn chatstate.FinishTurnResult finishError chatstate.FinishErrorResult + failIdle chatstate.FailIdleResult cancelRequiresAction chatstate.CancelRequiresActionResult reconcileInvalidState chatstate.ReconcileInvalidStateResult } @@ -741,6 +761,16 @@ func matrixCases() []transitionCaseSpec { setArchivedCase(chatstate.StateXE0, chatstate.StateE0, database.ChatStatusError), setArchivedCase(chatstate.StateXE1, chatstate.StateE1, database.ChatStatusError), + endChatCase(chatstate.StateW), + endChatCase(chatstate.StateE0), + endChatCase(chatstate.StateE1), + endChatCase(chatstate.StateR0), + endChatCase(chatstate.StateR1), + endChatCase(chatstate.StateI0), + endChatCase(chatstate.StateI1), + endChatCase(chatstate.StateA0), + endChatCase(chatstate.StateA1), + // SendMessage(queue) cases: idle states insert directly, // busy states append to the queue tail. sendMessageQueueCase(chatstate.StateW, chatstate.StateR0, true, 0), @@ -869,6 +899,8 @@ func matrixCases() []transitionCaseSpec { finishErrorCase(chatstate.StateR0, chatstate.StateE0), finishErrorCase(chatstate.StateR1, chatstate.StateE1), + failIdleCase(), + // ReconcileInvalidState cases: Invalid with empty queue // lands in E0; Invalid with non-empty queue lands in E1. reconcileInvalidStateCase(chatstate.StateE0, queueShapeDefault), @@ -909,6 +941,46 @@ func setArchivedCase(from, want chatstate.ExecutionState, wantStatus database.Ch } } +func endChatCase(from chatstate.ExecutionState) transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionEndChat, + from: from, + want: chatstate.StateXW, + apply: applyEndChat, + 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.Archived) + require.Equal(t, database.ChatStatusWaiting, after.Status) + require.False(t, after.WorkerID.Valid) + require.False(t, after.RunnerID.Valid) + require.False(t, after.LastError.Valid) + require.False(t, after.RequiresActionDeadlineAt.Valid) + switch from { + case chatstate.StateA0, chatstate.StateA1: + require.Len(t, result.endChat.InsertedMessages, 1, + "EndChat from A* synthesizes one tool-result cancellation") + cancel := requireChatMessageByID(ctx, t, f, + result.endChat.InsertedMessages[0].ID) + assertToolResultForCall(t, cancel, seeded.pendingToolCallID) + expectedHistory := append([]int64{}, base.historyIDs...) + expectedHistory = append(expectedHistory, cancel.ID) + require.Equal(t, expectedHistory, activeHistoryIDs(ctx, t, f, seeded.chatID), + "EndChat from A* archives history with the cancellation appended") + default: + require.Empty(t, result.endChat.InsertedMessages) + require.Equal(t, base.historyIDs, activeHistoryIDs(ctx, t, f, seeded.chatID)) + } + if len(base.queueIDs) == 0 { + require.Empty(t, result.endChat.DeletedQueuedMessageIDs) + } else { + require.Equal(t, base.queueIDs, result.endChat.DeletedQueuedMessageIDs) + } + require.Empty(t, queuedIDsByPosition(ctx, t, f, seeded.chatID)) + }, + } +} + func sendMessageQueueCase(from, want chatstate.ExecutionState, directInsert bool, queueDelta int64) transitionCaseSpec { return transitionCaseSpec{ transition: chatstate.TransitionSendMessage, @@ -1846,6 +1918,27 @@ func finishErrorCase(from, want chatstate.ExecutionState) transitionCaseSpec { } } +func failIdleCase() transitionCaseSpec { + return transitionCaseSpec{ + transition: chatstate.TransitionFailIdle, + from: chatstate.StateW, + want: chatstate.StateE0, + apply: applyFailIdle, + 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) + require.True(t, after.LastError.Valid) + require.JSONEq(t, `{"message":"hook dispatch failed","kind":"generic","retryable":false}`, string(after.LastError.RawMessage)) + require.Equal(t, base.chat.WorkerID, after.WorkerID) + require.Equal(t, base.chat.RunnerID, after.RunnerID) + require.Equal(t, base.historyVersion, after.HistoryVersion) + require.Equal(t, base.queueVersion, after.QueueVersion) + }, + } +} + func reconcileInvalidStateCase(want chatstate.ExecutionState, shape queueShape) transitionCaseSpec { spec := transitionCaseSpec{ transition: chatstate.TransitionReconcileInvalidState, diff --git a/coderd/x/chatd/compaction_hooks_test.go b/coderd/x/chatd/compaction_hooks_test.go new file mode 100644 index 0000000000000..efdc360aac758 --- /dev/null +++ b/coderd/x/chatd/compaction_hooks_test.go @@ -0,0 +1,264 @@ +package chatd_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "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/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/codersdk/agenthooks" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" + "github.com/coder/coder/v2/testutil" +) + +func TestCompactionHooksHintAndPostCommitResponses(t *testing.T) { + t.Parallel() + + var postSawCommitted atomic.Bool + fixture := startCompactionHookChat(t, + func(t *testing.T, db database.Store, request agenthooks.Request) (int, string) { + switch request.Type { + case agenthooks.EventPreCompact: + return http.StatusOK, `{"model_context":"preserve deployment constraints","user_message":"compaction starting","allowed_tools":["read_file"],"end_chat":true}` + case agenthooks.EventPostCompact: + postSawCommitted.Store(hasCompactionRows(t, db, request.Meta.ChatID)) + return http.StatusOK, `{"model_context":"post compact context","user_message":"compaction complete"}` + default: + return http.StatusOK, `{}` + } + }, + func(t *testing.T, body string) { + require.Contains(t, body, "preserve deployment constraints") + }, + ) + + waitCtx := testutil.Context(t, testutil.WaitLong) + testutil.Eventually(waitCtx, t, func(context.Context) bool { + updated, err := fixture.db.GetChatByID(waitCtx, fixture.chat.ID) + return err == nil && updated.Archived + }, testutil.IntervalFast) + // post_compact dispatches before the commit so its effects land + // atomically with the compaction step; the hook observes the + // pre-compaction history. + require.False(t, postSawCommitted.Load()) + require.Equal(t, int32(1), fixture.compactionCalls.Load()) + + updated, err := fixture.db.GetChatByID(fixture.ctx, fixture.chat.ID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(updated.HookAllowedTools.RawMessage)) + userMessages := chatMessages(fixture.ctx, t, fixture.db, fixture.chat.ID) + promptMessages, err := fixture.db.GetChatMessagesForPromptByChatID(fixture.ctx, fixture.chat.ID) + require.NoError(t, err) + require.True(t, hasMessageText(t, userMessages, "compaction starting", database.ChatMessageVisibilityUser)) + require.True(t, hasMessageText(t, userMessages, "compaction complete", database.ChatMessageVisibilityUser)) + require.True(t, hasMessageText(t, promptMessages, "post compact context", database.ChatMessageVisibilityModel)) + require.False(t, hasMessageText(t, promptMessages, "preserve deployment constraints", database.ChatMessageVisibilityModel)) +} + +func TestPreCompactHookFailureAbortsCompaction(t *testing.T) { + t.Parallel() + + fixture := startCompactionHookChat(t, + func(_ *testing.T, _ database.Store, request agenthooks.Request) (int, string) { + if request.Type == agenthooks.EventPreCompact { + return http.StatusInternalServerError, "" + } + return http.StatusOK, `{}` + }, + func(t *testing.T, _ string) { + require.FailNow(t, "compaction model called after pre_compact failure") + }, + ) + waitCtx := testutil.Context(t, testutil.WaitLong) + failed := waitForChatStatus(waitCtx, t, fixture.db, fixture.chat.ID, database.ChatStatusError) + require.Equal(t, int32(0), fixture.compactionCalls.Load()) + require.False(t, hasCompactionRows(t, fixture.db, fixture.chat.ID)) + dispatch := lifecycleDispatch(t, fixture.db, fixture.chat.ID, agenthooks.EventPreCompact) + require.Equal(t, "http_error", dispatch.Result) + lastError := chatLastErrorMessage(failed.LastError) + require.Contains(t, lastError, "hook dispatch failed: pre_compact: http_error") + require.Contains(t, lastError, dispatch.ID.String()) +} + +func TestPostCompactHookFailureKeepsCompaction(t *testing.T) { + t.Parallel() + + var postSawCommitted atomic.Bool + fixture := startCompactionHookChat(t, + func(t *testing.T, db database.Store, request agenthooks.Request) (int, string) { + if request.Type == agenthooks.EventPostCompact { + postSawCommitted.Store(hasCompactionRows(t, db, request.Meta.ChatID)) + return http.StatusInternalServerError, "" + } + return http.StatusOK, `{}` + }, + func(*testing.T, string) {}, + ) + waitCtx := testutil.Context(t, testutil.WaitLong) + failed := waitForChatStatus(waitCtx, t, fixture.db, fixture.chat.ID, database.ChatStatusError) + // The failure commits atomically with the compaction step, so the + // hook fires before the rows exist while the rows still persist. + require.False(t, postSawCommitted.Load()) + require.Equal(t, int32(1), fixture.compactionCalls.Load()) + require.True(t, hasCompactionRows(t, fixture.db, fixture.chat.ID)) + dispatch := lifecycleDispatch(t, fixture.db, fixture.chat.ID, agenthooks.EventPostCompact) + require.Equal(t, "http_error", dispatch.Result) + lastError := chatLastErrorMessage(failed.LastError) + require.Contains(t, lastError, "hook dispatch failed: post_compact: http_error") + require.Contains(t, lastError, dispatch.ID.String()) +} + +func TestPostCompactHookFailureAppliesPreCompactEndChat(t *testing.T) { + t.Parallel() + + fixture := startCompactionHookChat(t, + func(_ *testing.T, _ database.Store, request agenthooks.Request) (int, string) { + switch request.Type { + case agenthooks.EventPreCompact: + return http.StatusOK, `{"end_chat":true}` + case agenthooks.EventPostCompact: + return http.StatusInternalServerError, "" + default: + return http.StatusOK, `{}` + } + }, + func(*testing.T, string) {}, + ) + waitCtx := testutil.Context(t, testutil.WaitLong) + var archived database.Chat + testutil.Eventually(waitCtx, t, func(ctx context.Context) bool { + updated, err := fixture.db.GetChatByID(ctx, fixture.chat.ID) + if err != nil { + return false + } + archived = updated + return updated.Archived && updated.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.False(t, archived.LastError.Valid) + require.Equal(t, int32(1), fixture.compactionCalls.Load()) + require.True(t, hasCompactionRows(t, fixture.db, fixture.chat.ID)) + dispatch := lifecycleDispatch(t, fixture.db, fixture.chat.ID, agenthooks.EventPostCompact) + require.Equal(t, "http_error", dispatch.Result) +} + +type compactionHookFixture struct { + ctx context.Context + db database.Store + chat database.Chat + compactionCalls *atomic.Int32 +} + +func startCompactionHookChat( + t *testing.T, + hookResponse func(*testing.T, database.Store, agenthooks.Request) (int, string), + inspectCompaction func(*testing.T, string), +) compactionHookFixture { + t.Helper() + + const ( + contextLimit = int64(100) + thresholdPercent = int32(70) + ) + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var compactionCalls atomic.Int32 + var streamCalls 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") { + compactionCalls.Add(1) + inspectCompaction(t, body) + return anthropicCompactionResponse("hook compaction summary") + } + return chattest.AnthropicNonStreamingResponse("title") + } + if streamCalls.Add(1) == 1 { + return highUsageReadFileResponse("/tmp/hook.txt") + } + 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) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + status, body := hookResponse(t, db, request) + w.WriteHeader(status) + if body != "" { + _, err := w.Write([]byte(body)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/hook.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.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, anthropicURL, chattest.WithPreservePath())) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "compaction-hooks", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("trigger compaction hooks"), + }, + }) + require.NoError(t, err) + return compactionHookFixture{ctx: ctx, db: db, chat: chat, compactionCalls: &compactionCalls} +} + +func hasCompactionRows(t *testing.T, db database.Store, chatID uuid.UUID) bool { + t.Helper() + userMessages := chatMessages(t.Context(), t, db, chatID) + promptMessages, err := db.GetChatMessagesForPromptByChatID(t.Context(), chatID) + require.NoError(t, err) + compressed := compressedChatSummarizedMessages(t, append(promptMessages, userMessages...)) + return len(compressed.summaries) > 0 && len(compressed.calls) > 0 && len(compressed.results) > 0 +} + +func hasMessageText(t *testing.T, messages []database.ChatMessage, text string, visibility database.ChatMessageVisibility) bool { + t.Helper() + for _, message := range messages { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if message.Visibility == visibility && len(parts) == 1 && parts[0].Text == text { + return true + } + } + return false +} diff --git a/coderd/x/chatd/create_hooks_test.go b/coderd/x/chatd/create_hooks_test.go new file mode 100644 index 0000000000000..5c559f20db989 --- /dev/null +++ b/coderd/x/chatd/create_hooks_test.go @@ -0,0 +1,279 @@ +package chatd_test + +import ( + "context" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/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/chathooks" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" +) + +func TestCreateChatUserPromptSubmitHook(t *testing.T) { + t.Parallel() + + t.Run("passthrough", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusOK, `{}`) + + chat, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "passthrough")) + require.NoError(t, err) + request := testutil.RequireReceive(ctx, t, requests) + require.Equal(t, agenthooks.EventUserPromptSubmit, request.Type) + require.Equal(t, chat.ID, request.Meta.ChatID) + require.Equal(t, user.ID, request.Meta.OwnerID) + require.NotNil(t, request.Meta.TurnID) + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.UserPromptSubmitData) + require.True(t, ok) + require.Equal(t, "passthrough", data.Prompt) + var hookParts []codersdk.ChatMessagePart + require.NoError(t, json.Unmarshal(data.Parts, &hookParts)) + require.Equal(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("passthrough")}, hookParts) + + messages := chatMessages(ctx, t, db, chat.ID) + initialUser := messages[len(messages)-1] + require.Equal(t, database.ChatMessageRoleUser, initialUser.Role) + require.Equal(t, database.ChatMessageVisibilityBoth, initialUser.Visibility) + require.Equal(t, "passthrough", hookMessageText(t, initialUser)) + require.Equal(t, uuid.NullUUID{UUID: *request.Meta.TurnID, Valid: true}, initialUser.TurnID) + }) + + t.Run("override", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"permission":{"decision":"allow","input_override":{"prompt":"redacted"}}}`) + + opts := createHookOptions(t, db, user.ID, org.ID, model.ID, "secret") + opts.Title = chatprompt.FallbackTitle(chatprompt.TitleText(opts.InitialUserContent, nil)) + opts.TitleDerivedFromContent = true + chat, err := server.CreateChat(ctx, opts) + require.NoError(t, err) + request := testutil.RequireReceive(ctx, t, requests) + require.NotNil(t, request.Meta.TurnID) + messages := chatMessages(ctx, t, db, chat.ID) + initialUser := messages[len(messages)-1] + require.Equal(t, "redacted", hookMessageText(t, initialUser)) + require.Equal(t, uuid.NullUUID{UUID: *request.Meta.TurnID, Valid: true}, initialUser.TurnID) + require.Equal(t, "redacted", chat.Title, "prompt-derived title must be recomputed from the override") + }) + + t.Run("override keeps explicit title", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, _ := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"permission":{"decision":"allow","input_override":{"prompt":"redacted"}}}`) + + chat, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "secret")) + require.NoError(t, err) + require.Equal(t, "create hook test", chat.Title) + }) + + t.Run("invalid model config rejected before dispatch", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusOK, `{}`) + + opts := createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt") + opts.ModelConfigID = uuid.New() + _, err := server.CreateChat(ctx, opts) + require.ErrorIs(t, err, chatd.ErrInvalidModelConfigID) + select { + case request := <-requests: + t.Fatalf("unexpected hook dispatch %s for rejected create", request.Type) + default: + } + }) + + t.Run("override recomputes paste-derived title", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, _ := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"permission":{"decision":"allow","input_override":{"prompt":"redacted"}}}`) + + opts := createHookOptions(t, db, user.ID, org.ID, model.ID, " ") + opts.Title = chatprompt.FallbackTitle("secret paste content") + opts.TitleDerivedFromContent = true + chat, err := server.CreateChat(ctx, opts) + require.NoError(t, err) + require.Equal(t, "redacted", chat.Title, + "paste-derived title must be recomputed from the override") + }) + + t.Run("response messages", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, _ := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"model_context":"model only","user_message":"user only"}`) + + chat, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt")) + require.NoError(t, err) + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + require.GreaterOrEqual(t, len(promptMessages), 2) + modelContext := promptMessages[len(promptMessages)-2] + initialUser := promptMessages[len(promptMessages)-1] + require.Equal(t, database.ChatMessageRoleUser, modelContext.Role) + require.Equal(t, database.ChatMessageVisibilityModel, modelContext.Visibility) + require.Equal(t, "model only", hookMessageText(t, modelContext)) + require.Equal(t, database.ChatMessageRoleUser, initialUser.Role) + require.Equal(t, database.ChatMessageVisibilityBoth, initialUser.Visibility) + require.Equal(t, "prompt", hookMessageText(t, initialUser)) + require.Equal(t, initialUser.TurnID, modelContext.TurnID) + + userMessages := chatMessages(ctx, t, db, chat.ID) + require.GreaterOrEqual(t, len(userMessages), 2) + userNotice := userMessages[len(userMessages)-2] + require.Equal(t, database.ChatMessageRoleSystem, userNotice.Role) + require.Equal(t, database.ChatMessageVisibilityUser, userNotice.Visibility) + require.Equal(t, "user only", hookMessageText(t, userNotice)) + require.Equal(t, initialUser.TurnID, userNotice.TurnID) + }) + + t.Run("allowed tools", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, _ := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"allowed_tools":["read_file","write_file"]}`) + + chat, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt")) + require.NoError(t, err) + require.True(t, chat.HookAllowedTools.Valid) + var allowedTools []string + require.NoError(t, json.Unmarshal(chat.HookAllowedTools.RawMessage, &allowedTools)) + require.Equal(t, []string{"read_file", "write_file"}, allowedTools) + }) + + t.Run("deny", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"permission":{"decision":"deny"},"user_message":"blocked"}`) + + _, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt")) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + require.Equal(t, "blocked", denied.UserMessage) + request := testutil.RequireReceive(ctx, t, requests) + requireCreateHookChatMissing(ctx, t, db, request.Meta.ChatID) + }) + + t.Run("dispatch failure", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusInternalServerError, "") + + _, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt")) + var dispatchErr *chathooks.DispatchError + require.ErrorAs(t, err, &dispatchErr) + require.Equal(t, chathooks.ResultHTTPError, dispatchErr.Class) + request := testutil.RequireReceive(ctx, t, requests) + requireCreateHookChatMissing(ctx, t, db, request.Meta.ChatID) + }) + + t.Run("end chat", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server, requests := newCreateHookTestServer(t, db, ps, http.StatusOK, `{"user_message":"ended","end_chat":true}`) + + _, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "prompt")) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + require.Equal(t, "ended", denied.UserMessage) + request := testutil.RequireReceive(ctx, t, requests) + requireCreateHookChatMissing(ctx, t, db, request.Meta.ChatID) + }) +} + +func TestCreateChatHooksDisabledUnchanged(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + server := newTestServer(t, db, ps, uuid.New()) + + chat, err := server.CreateChat(ctx, createHookOptions(t, db, user.ID, org.ID, model.ID, "unchanged")) + require.NoError(t, err) + messages := chatMessages(ctx, t, db, chat.ID) + initialUser := messages[len(messages)-1] + require.False(t, initialUser.TurnID.Valid) + require.Equal(t, "unchanged", hookMessageText(t, initialUser)) + require.False(t, chat.HookAllowedTools.Valid) +} + +func newCreateHookTestServer( + t *testing.T, + db database.Store, + ps dbpubsub.Pubsub, + statusCode int, + response string, +) (*chatd.Server, <-chan agenthooks.Request) { + t.Helper() + requests := make(chan agenthooks.Request, 2) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + requests <- request + w.WriteHeader(statusCode) + if response != "" { + _, err := w.Write([]byte(response)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + return newHookTestServer(t, db, ps, consumer), requests +} + +func createHookOptions( + t *testing.T, + db database.Store, + userID uuid.UUID, + organizationID uuid.UUID, + modelConfigID uuid.UUID, + prompt string, +) chatd.CreateOptions { + t.Helper() + return chatd.CreateOptions{ + OrganizationID: organizationID, + OwnerID: userID, + Title: "create hook test", + ModelConfigID: modelConfigID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)}, + } +} + +func requireCreateHookChatMissing(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID) { + t.Helper() + _, err := db.GetChatByID(ctx, chatID) + require.ErrorIs(t, err, sql.ErrNoRows) +} diff --git a/coderd/x/chatd/generation.go b/coderd/x/chatd/generation.go index 443909157df89..6ed0be12e2811 100644 --- a/coderd/x/chatd/generation.go +++ b/coderd/x/chatd/generation.go @@ -22,6 +22,7 @@ import ( "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/codersdk/agenthooks" ) // generationPrepareInput contains the committed state used to prepare one @@ -317,6 +318,37 @@ func unresolvedToolCallsFromHistory( return localCalls, dynamicCalls, nil } +// priorToolCallIDsInTurn returns tool call IDs from assistant steps +// before the latest one in the current turn. Their recorded +// pre_tool_use decisions must not be replayed for a repeated call. +func priorToolCallIDsInTurn(messages []database.ChatMessage) (map[string]bool, error) { + assistantIndex := lastMessageIndex(messages, func(msg database.ChatMessage) bool { + return msg.Role == database.ChatMessageRoleAssistant + }) + prior := make(map[string]bool) + // Assistant steps lack turn IDs, so user-visible prompts bound the turn. + // Including earlier turns would duplicate hook effects on retry. + start := currentTurnStartIndex(messages) + if assistantIndex < start { + return prior, nil + } + for _, msg := range messages[start:assistantIndex] { + if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleAssistant { + continue + } + parts, err := chatprompt.ParseContent(msg) + if err != nil { + return nil, xerrors.Errorf("parse assistant message: %w", err) + } + for _, part := range parts { + if part.Type == codersdk.ChatMessagePartTypeToolCall && part.ToolCallID != "" { + prior[part.ToolCallID] = true + } + } + } + return prior, nil +} + func hasExclusiveToolCall(toolCalls []fantasy.ToolCallContent, exclusiveToolNames map[string]bool) bool { if len(exclusiveToolNames) == 0 { return false @@ -329,13 +361,64 @@ func hasExclusiveToolCall(toolCalls []fantasy.ToolCallContent, exclusiveToolName return false } +func (s *taskStarter) startGenerationSession( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + chat database.Chat, + messages []database.ChatMessage, +) (result sessionStartResult, dispatched bool, err error) { + dispatched, complete, err := input.SessionStart.claim(ctx) + if err != nil { + return sessionStartResult{}, false, errors.Join(errTaskExpectedExit, xerrors.Errorf("claim session_start: %w", err)) + } + if !dispatched { + return sessionStartResult{Chat: chat}, false, nil + } + + completed := false + // Re-arm the claim until its response is applied so a replacement task + // can replay session_start effects. + defer func() { complete(completed) }() + turnID := activeTurnID(messages) + response, err := s.server.dispatchLifecycleHook(ctx, chat, turnID, agenthooks.EventSessionStart, agenthooks.SessionStartData{Source: sessionStartSource(messages)}) + if err != nil { + return sessionStartResult{}, true, sessionStartDispatchError(err) + } + result, err = applySessionStartResponse(ctx, machine, input, chat, turnID, response) + if err != nil { + return sessionStartResult{}, true, err + } + completed = true + return result, true, nil +} + func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskStartInput) error { + if input.StopNudges == nil { + input.StopNudges = &stopNudgeTracker{} + } machine := chatstate.NewChatMachine(s.opts.Store, s.opts.Pubsub, input.ChatID) for { chat, messages, err := loadGenerationState(ctx, machine, input) if err != nil { return xerrors.Errorf("load generation state: %w", err) } + if s.server.hookDispatcher != nil && s.server.hookDispatcher.Enabled() { + result, dispatched, err := s.startGenerationSession(ctx, machine, input, chat, messages) + if err != nil { + if errors.Is(err, errTaskExpectedExit) { + return err + } + return s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired) + } + if dispatched { + if result.Ended { + return s.finishSessionStartEnd(ctx, input, result) + } + input.HistoryVersion = result.Chat.HistoryVersion + continue + } + } prepareInput := generationPrepareInput{ Chat: chat, Messages: messages, @@ -350,20 +433,25 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS return s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired) } cleanup := prepared.Cleanup - decision, err := retryGenerationPhase(ctx, s, "decide", 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, - compactionThresholdPercent: generationCompactionThreshold(prepared.Compaction), - compactionContextLimit: generationCompactionContextLimit(prepared.Compaction), + var decision generationDecision + if input.StopNudges.consume(activeTurnID(prepared.Messages)) { + decision = generationDecision{kind: generationActionGenerateAssistant} + } else { + decision, err = retryGenerationPhase(ctx, s, "decide", 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, + compactionThresholdPercent: generationCompactionThreshold(prepared.Compaction), + compactionContextLimit: generationCompactionContextLimit(prepared.Compaction), + }) }) - }) + } if err != nil { cleanup() if errors.Is(err, errTaskExpectedExit) || errors.Is(err, errTaskRetryable) { @@ -384,8 +472,29 @@ func (s *taskStarter) StartGeneration(ctx context.Context, input chatWorkerTaskS var actionErr error switch decision.kind { case generationActionEnterRequiresAction: - cleanup() - return s.enterRequiresAction(ctx, machine, input) + toolCalls := make([]fantasy.ToolCallContent, 0, len(decision.pendingDynamicToolCalls)) + for _, toolCall := range decision.pendingDynamicToolCalls { + toolCalls = append(toolCalls, fantasy.ToolCallContent{ + ToolCallID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + Input: toolCall.Args, + }) + } + priorToolCallIDs, err := priorToolCallIDsInTurn(prepared.Messages) + if err != nil { + cleanup() + return s.finishGenerationError(ctx, machine, input, err, generationAttemptNotRequired) + } + preflight, err := s.server.preflightPendingToolCalls(ctx, prepared.Chat, activeTurnID(prepared.Messages), toolCalls, priorToolCallIDs) + if err != nil { + cleanup() + return s.finishGenerationError(ctx, machine, input, generationHookDispatchError(agenthooks.EventPreToolUse, err), generationAttemptNotRequired) + } + if len(preflight.Denied) == 0 { + cleanup() + return s.enterRequiresAction(ctx, machine, input, prepared, preflight) + } + actionErr = s.commitPreToolUseDeniedResults(ctx, machine, input, prepared, preflight) case generationActionFinishTurn: cleanup() return s.finishGenerationTurn(ctx, machine, input, decision, generationAttemptNotRequired) @@ -636,10 +745,53 @@ func (s *taskStarter) generateAssistant( if len(outcome.Step.Content) == 0 { return s.finishGenerationTurn(ctx, machine, input, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, requireGenerationAttempt(attempt.number)) } + turnID := activeTurnID(prepared.Messages) + preflight, err := s.server.preflightToolCalls(ctx, prepared.Chat, turnID, outcome.Step, outcome.ToolCalls) + if err != nil { + return generationHookDispatchError(agenthooks.EventPreToolUse, err) + } messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ modelConfigID: prepared.ModelConfigID, modelCallConfig: prepared.ModelConfig, - step: stepDataFromPersisted(outcome.Step), + step: stepDataFromPersisted(preflight.Step), + toolNameToConfigID: prepared.ToolNameToConfigID, + logger: s.opts.Logger, + contentVersion: chatprompt.CurrentContentVersion, + }) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + messages, endChat, err := applyHookResponseMessages(messages, preflight.Responses, prepared.ModelConfigID, turnID) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionGenerateAssistant, messages, generationCommitHooks{ + Responses: preflight.Responses, + EndChat: endChat, + EffectDispatchIDs: preflight.EffectDispatchIDs, + }) +} + +func (s *taskStarter) commitPreToolUseDeniedResults( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + prepared generationPrepared, + preflight preToolUseExecutionResult, +) error { + attempt, err := s.beginGenerationAttempt(ctx, machine, input) + if err != nil { + return xerrors.Errorf("begin generation attempt: %w", err) + } + defer attempt.closeEpisode() + content := make([]fantasy.Content, 0, len(preflight.Denied)) + for _, denied := range preflight.Denied { + content = append(content, denied) + } + messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ + modelConfigID: prepared.ModelConfigID, + modelCallConfig: prepared.ModelConfig, + step: stepDataFromPersisted(chatloop.PersistedStep{Content: content}), toolNameToConfigID: prepared.ToolNameToConfigID, logger: s.opts.Logger, contentVersion: chatprompt.CurrentContentVersion, @@ -647,7 +799,17 @@ func (s *taskStarter) generateAssistant( if err != nil { return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) } - return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionGenerateAssistant, messages) + turnID := activeTurnID(prepared.Messages) + messages, endChat, err := applyHookResponseMessages(messages, preflight.Responses, prepared.ModelConfigID, turnID) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionExecuteLocalTools, messages, generationCommitHooks{ + Responses: preflight.Responses, + Overrides: preflight.Overrides, + EndChat: endChat, + EffectDispatchIDs: preflight.EffectDispatchIDs, + }) } func (s *taskStarter) executeLocalTools( @@ -657,6 +819,25 @@ func (s *taskStarter) executeLocalTools( prepared generationPrepared, decision generationDecision, ) error { + turnID := activeTurnID(prepared.Messages) + priorToolCallIDs, err := priorToolCallIDsInTurn(prepared.Messages) + if err != nil { + return err + } + preflight, err := s.server.preflightPendingToolCalls(ctx, prepared.Chat, turnID, decision.localToolCalls, priorToolCallIDs) + if err != nil { + return generationHookDispatchError(agenthooks.EventPreToolUse, err) + } + for _, response := range preflight.Responses { + if response.EndChat { + // end_chat commits hook effects before any allowed tool executes. + return s.commitPreToolUseDeniedResults(ctx, machine, input, prepared, preflight) + } + } + pendingPolicy, hasPendingPolicy, err := narrowHookAllowedToolsResponses(prepared.Chat.HookAllowedTools, preflight.Responses) + if err != nil { + return xerrors.Errorf("narrow pending hook allowed tools: %w", err) + } attempt, err := s.beginGenerationAttempt(ctx, machine, input) if err != nil { return xerrors.Errorf("beginGenerationAttempt: %w", err) @@ -668,25 +849,37 @@ func (s *taskStarter) executeLocalTools( provider = prepared.Model.Provider() modelName = prepared.Model.Model() } - outcome, err := chatloop.ExecuteLocalTools(ctx, chatloop.ExecuteLocalToolsOptions{ - Tools: prepared.Tools, - ActiveTools: prepared.ActiveTools, - ProviderTools: prepared.ProviderTools, - ToolCalls: decision.localToolCalls, - ExclusiveToolNames: prepared.ExclusiveToolNames, - BuiltinToolNames: prepared.BuiltinToolNames, - ModelProvider: provider, - ModelName: modelName, - ContextLimit: prepared.ContextLimitFallback, - ToolNameAliases: subagentToolNameAliases, - PublishMessagePart: attempt.publish, - Logger: s.opts.Logger, - Metrics: s.server.metrics, - Clock: s.opts.Clock, - }) - if err != nil { - return xerrors.Errorf("execute local tools: %w", err) + toolCtx := ctx + if hasPendingPolicy { + toolCtx = withPendingHookAllowedTools(toolCtx, prepared.Chat.ID, pendingPolicy) + } + var outcome chatloop.ToolExecutionOutcome + if len(preflight.Allowed) > 0 { + outcome, err = chatloop.ExecuteLocalTools(toolCtx, chatloop.ExecuteLocalToolsOptions{ + Tools: prepared.Tools, + ActiveTools: prepared.ActiveTools, + ProviderTools: prepared.ProviderTools, + ToolCalls: preflight.Allowed, + ExclusiveToolNames: prepared.ExclusiveToolNames, + BuiltinToolNames: prepared.BuiltinToolNames, + ModelProvider: provider, + ModelName: modelName, + ContextLimit: prepared.ContextLimitFallback, + ToolNameAliases: subagentToolNameAliases, + PublishMessagePart: attempt.publish, + Logger: s.opts.Logger, + Metrics: s.server.metrics, + Clock: s.opts.Clock, + }) + if err != nil { + return xerrors.Errorf("execute local tools: %w", err) + } } + postResponses, postDispatchErr := s.server.dispatchPostToolUseResults(ctx, prepared.Chat, turnID, outcome.Step.Content) + for _, denied := range preflight.Denied { + outcome.Step.Content = append(outcome.Step.Content, denied) + } + restoreToolCallOrder(outcome.Step.Content, decision.localToolCalls) messages, err := buildCommitStepMessages(buildCommitStepMessagesInput{ modelConfigID: prepared.ModelConfigID, modelCallConfig: prepared.ModelConfig, @@ -698,7 +891,31 @@ func (s *taskStarter) executeLocalTools( if err != nil { return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) } - return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionExecuteLocalTools, messages) + messages, preEndChat, err := applyHookResponseMessages(messages, preflight.Responses, prepared.ModelConfigID, turnID) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + messages, postEndChat, err := appendHookResponseMessages(messages, postResponses, prepared.ModelConfigID, turnID) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + responses := make([]agenthooks.Response, 0, len(preflight.Responses)+len(postResponses)) + responses = append(responses, preflight.Responses...) + responses = append(responses, postResponses...) + endChat := preEndChat || postEndChat + // A post-hook failure prevents a later end_chat from archiving an unreviewed turn. + var postCommitErr error + if postDispatchErr != nil { + postCommitErr = generationHookDispatchError(agenthooks.EventPostToolUse, postDispatchErr) + endChat = preEndChat + } + return s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionExecuteLocalTools, messages, generationCommitHooks{ + Responses: responses, + Overrides: preflight.Overrides, + EndChat: endChat, + PostCommitError: postCommitErr, + EffectDispatchIDs: preflight.EffectDispatchIDs, + }) } // compactionSourceForDecision maps a compact decision to the @@ -751,6 +968,12 @@ func (s *taskStarter) generateCompaction( overrideModel.modelConfig, ) } + turnID := activeTurnID(prepared.Messages) + preResponse, err := s.server.dispatchLifecycleHook(ctx, prepared.Chat, turnID, agenthooks.EventPreCompact, agenthooks.PreCompactData{}) + if err != nil { + return generationHookDispatchError(agenthooks.EventPreCompact, err) + } + compactionOpts.SummaryHint = preResponse.ModelContext compactionOpts.PublishMessagePart = attempt.publish compactionOpts.Source = source compactionOpts.Force = source == chatloop.CompactionSourceManual @@ -779,14 +1002,50 @@ func (s *taskStarter) generateCompaction( s.server.metrics.RecordCompaction(metricProvider, metricModel, false, err) return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) } - err = s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionCompact, stepMessagesForCommit{ + persistedPreResponse := preResponse + persistedPreResponse.ModelContext = "" + commitMessages, preEndChat, err := applyHookResponseMessages(stepMessagesForCommit{ Messages: messages.Messages, VisibleIndexes: visibleMessageIndexes(messages.Messages), ConsumeCompactionRequest: true, + }, []agenthooks.Response{persistedPreResponse}, prepared.ModelConfigID, turnID) + if err != nil { + s.server.metrics.RecordCompaction(metricProvider, metricModel, false, err) + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + // post_compact dispatches before the commit so its effects, or the + // fail-closed dispatch error, land atomically with the compaction + // step. Committing first would publish a runnable running state + // whose new history lets the runner start another model step before + // the response applies, and a crash between the commits would drop + // the response entirely. + postResponse, postDispatchErr := s.server.dispatchLifecycleHook(ctx, prepared.Chat, turnID, agenthooks.EventPostCompact, agenthooks.PostCompactData{}) + responses := []agenthooks.Response{preResponse} + endChat := preEndChat + var postCommitErr error + if postDispatchErr != nil { + // An accepted pre_compact end_chat outranks a post_compact failure. + if !preEndChat { + postCommitErr = generationHookDispatchError(agenthooks.EventPostCompact, postDispatchErr) + } + } else { + var postEndChat bool + commitMessages, postEndChat, err = appendHookResponseMessages(commitMessages, []agenthooks.Response{postResponse}, prepared.ModelConfigID, turnID) + if err != nil { + s.server.metrics.RecordCompaction(metricProvider, metricModel, false, err) + return s.finishGenerationError(ctx, machine, input, err, requireGenerationAttempt(attempt.number)) + } + responses = append(responses, postResponse) + endChat = preEndChat || postEndChat + } + err = s.commitGenerationStep(ctx, machine, input, attempt.number, generationActionCompact, commitMessages, generationCommitHooks{ + Responses: responses, + EndChat: endChat, + PostCommitError: postCommitErr, }) s.server.metrics.RecordCompaction(metricProvider, metricModel, err == nil, err) if err != nil { - return xerrors.Errorf("commit generation step: %w", err) + return xerrors.Errorf("commit compaction step: %w", err) } return nil } @@ -871,6 +1130,14 @@ func (s *taskStarter) beginGenerationAttempt( }, nil } +type generationCommitHooks struct { + Responses []agenthooks.Response + Overrides map[string]json.RawMessage + EndChat bool + PostCommitError error + EffectDispatchIDs []uuid.UUID +} + func (s *taskStarter) commitGenerationStep( ctx context.Context, machine *chatstate.ChatMachine, @@ -878,36 +1145,110 @@ func (s *taskStarter) commitGenerationStep( attempt int64, kind generationActionKind, messages stepMessagesForCommit, + hooks generationCommitHooks, ) error { - if len(messages.Messages) == 0 { + // end_chat must commit even when the step carries no messages. + if len(messages.Messages) == 0 && !hooks.EndChat { + if hooks.PostCommitError != nil { + return s.finishGenerationError(ctx, machine, input, hooks.PostCommitError, requireGenerationAttempt(attempt)) + } return s.finishGenerationTurn(ctx, machine, input, generationDecision{kind: generationActionFinishTurn, finishReason: generationFinishReasonComplete}, requireGenerationAttempt(attempt)) } + // An accepted end_chat outranks a post-hook failure: the archived + // chat is terminal, so nothing can continue unreviewed. + failClosed := hooks.PostCommitError != nil && !hooks.EndChat + var postCommitLastError pqtype.NullRawMessage + var postCommitMessage string + if hooks.PostCommitError != nil { + classified := chaterror.Classify(hooks.PostCommitError) + s.opts.Logger.Warn(ctx, "chat generation failed", + slog.F("chat_id", input.ChatID), + slog.F("worker_id", input.WorkerID), + slog.F("generation_attempt", input.GenerationAttempt), + slog.F("error_kind", classified.Kind), + slog.F("provider", classified.Provider), + slog.F("status_code", classified.StatusCode), + slog.F("retryable", classified.Retryable), + slog.Error(hooks.PostCommitError), + ) + postCommitLastError, postCommitMessage = generationLastError(hooks.PostCommitError) + } var committed database.Chat + var endedDescendants []database.Chat insertedMessages := []runnerActionMessage{} err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { if _, err := loadChatForGeneration(ctx, store, input, requireGenerationAttempt(attempt)); err != nil { return xerrors.Errorf("load chat for generation: %w", err) } - commitResult, err := tx.CommitStep(chatstate.CommitStepInput{ - Messages: messages.Messages, - ConsumeCompactionRequest: messages.ConsumeCompactionRequest, - }) - if err != nil { - return xerrors.Errorf("tx.CommitStep: %w", err) + if err := replacePersistedToolCallInputs(ctx, store, input.ChatID, hooks.Overrides); err != nil { + return err + } + if err := applyHookAllowedToolsResponses(ctx, store, input.ChatID, hooks.Responses); err != nil { + return err + } + if err := markHookDispatchEffectsApplied(ctx, store, input.ChatID, hooks.EffectDispatchIDs); err != nil { + return err } - insertedMessages = make([]runnerActionMessage, 0, len(commitResult.InsertedMessages)) - for _, msg := range commitResult.InsertedMessages { + var inserted []database.ChatMessage + if hooks.EndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: messages.Messages}) + if err != nil { + return xerrors.Errorf("tx.EndChatFamily: %w", err) + } + inserted = endResult.InsertedMessages + endedDescendants = endResult.EndedDescendants + } else { + commitResult, err := tx.CommitStep(chatstate.CommitStepInput{ + Messages: messages.Messages, + ConsumeCompactionRequest: messages.ConsumeCompactionRequest, + }) + if err != nil { + return xerrors.Errorf("tx.CommitStep: %w", err) + } + inserted = commitResult.InsertedMessages + // The fail-closed hook error must land in the same + // transaction as the step: committing them separately + // publishes a runnable running state whose new history + // lets the runner start another model step, and a crash + // between the commits drops the error entirely. + if failClosed { + if _, err := tx.FinishError(chatstate.FinishErrorInput{LastError: postCommitLastError}); err != nil { + return xerrors.Errorf("tx.FinishError: %w", err) + } + } + } + insertedMessages = make([]runnerActionMessage, 0, len(inserted)) + for _, msg := range inserted { insertedMessages = append(insertedMessages, runnerActionMessage{ID: msg.ID, Role: codersdk.ChatMessageRole(msg.Role)}) } - committed, err = store.GetChatByID(ctx, input.ChatID) + loadedChat, err := store.GetChatByID(ctx, input.ChatID) if err != nil { return xerrors.Errorf("load committed chat: %w", err) } + committed = loadedChat return nil }) if err != nil { return normalizeTaskTransitionError(err, "commit generation step") } + if hooks.EndChat { + input.StopNudges.reset() + return s.finishEndedChat(ctx, input, committed, endedDescendants) + } + if failClosed { + input.DebugTurn.RecordOutcome(chatdebug.StatusError) + postCommitCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postCommitWatchPublishTimeout) + defer cancel() + if err := s.publishWatchAndRoute(postCommitCtx, committed, codersdk.ChatWatchEventKindStatusChange); err != nil { + return xerrors.Errorf("publish watch and route: %w", err) + } + return s.afterGenerationOutcome(postCommitCtx, generationOutcome{ + Chat: committed, + Kind: runnerActionKindFinishError, + WatchEventKind: codersdk.ChatWatchEventKindStatusChange, + LastError: postCommitMessage, + }) + } s.routeStateHint(ctx, stateUpdateFromChat(committed)) return s.afterGenerationOutcome(ctx, generationOutcome{ Chat: committed, @@ -920,14 +1261,52 @@ func (s *taskStarter) enterRequiresAction( ctx context.Context, machine *chatstate.ChatMachine, input chatWorkerTaskStartInput, + prepared generationPrepared, + preflight preToolUseExecutionResult, ) error { + turnID := activeTurnID(prepared.Messages) + messages, endChat, err := applyHookResponseMessages(stepMessagesForCommit{}, preflight.Responses, prepared.ModelConfigID, turnID) + if err != nil { + return err + } var committed database.Chat - err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + var endedDescendants []database.Chat + insertedMessages := []runnerActionMessage{} + err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { if _, err := loadChatForTask(ctx, store, input, database.ChatStatusRunning, taskFenceOptions{requireHistory: true}); err != nil { return xerrors.Errorf("load chat for task: %w", err) } - if _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}); err != nil { - return xerrors.Errorf("tx.EnterRequiresAction: %w", err) + if err := replacePersistedToolCallInputs(ctx, store, input.ChatID, preflight.Overrides); err != nil { + return err + } + if err := applyHookAllowedToolsResponses(ctx, store, input.ChatID, preflight.Responses); err != nil { + return err + } + if err := markHookDispatchEffectsApplied(ctx, store, input.ChatID, preflight.EffectDispatchIDs); err != nil { + return err + } + var inserted []database.ChatMessage + if endChat { + result, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: messages.Messages}) + if err != nil { + return xerrors.Errorf("tx.EndChatFamily: %w", err) + } + inserted = result.InsertedMessages + endedDescendants = result.EndedDescendants + } else { + if len(messages.Messages) > 0 { + result, err := tx.CommitStep(chatstate.CommitStepInput{Messages: messages.Messages}) + if err != nil { + return xerrors.Errorf("tx.CommitStep: %w", err) + } + inserted = result.InsertedMessages + } + if _, err := tx.EnterRequiresAction(chatstate.EnterRequiresActionInput{}); err != nil { + return xerrors.Errorf("tx.EnterRequiresAction: %w", err) + } + } + for _, message := range inserted { + insertedMessages = append(insertedMessages, runnerActionMessage{ID: message.ID, Role: codersdk.ChatMessageRole(message.Role)}) } chat, err := store.GetChatByID(ctx, input.ChatID) if err != nil { @@ -939,13 +1318,17 @@ func (s *taskStarter) enterRequiresAction( if err != nil { return normalizeTaskTransitionError(err, "enter requires action") } + if endChat { + return s.finishEndedChat(ctx, input, committed, endedDescendants) + } if err := s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindActionRequired); err != nil { return xerrors.Errorf("publish watch and route: %w", err) } return s.afterGenerationOutcome(ctx, generationOutcome{ - Chat: committed, - Kind: runnerActionKindEnterRequiresAction, - WatchEventKind: codersdk.ChatWatchEventKindActionRequired, + Chat: committed, + Kind: runnerActionKindEnterRequiresAction, + WatchEventKind: codersdk.ChatWatchEventKindActionRequired, + InsertedMessages: insertedMessages, }) } @@ -997,7 +1380,14 @@ func recordGenerationFinishFailure(turn *runnerDebugTurn, err error) { turn.RecordOutcome(chatdebug.StatusError) } -func (s *taskStarter) finishGenerationTurn( +func (s *taskStarter) finishEndedChat(ctx context.Context, input chatWorkerTaskStartInput, committed database.Chat, descendants []database.Chat) error { + input.DebugTurn.RecordOutcome(chatdebug.StatusCompleted) + s.server.scheduleArchiveDebugCleanup(ctx, append([]database.Chat{committed}, descendants...)) + s.server.publishChatPubsubEvents(descendants, codersdk.ChatWatchEventKindDeleted) + return s.publishWatchAndRoute(ctx, committed, codersdk.ChatWatchEventKindDeleted) +} + +func (s *taskStarter) finishGenerationTurnWithoutHook( ctx context.Context, machine *chatstate.ChatMachine, input chatWorkerTaskStartInput, @@ -1024,6 +1414,129 @@ func (s *taskStarter) finishGenerationTurn( recordGenerationFinishFailure(input.DebugTurn, err) return err } + input.StopNudges.reset() + 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 xerrors.Errorf("publish watch and route: %w", err) + } + if err := s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKindFinishTurn, + WatchEventKind: codersdk.ChatWatchEventKindStatusChange, + PromotedMessageID: decision.promotedMessageID, + }); err != nil { + return xerrors.Errorf("after generation outcome: %w", err) + } + s.routeStateHint(ctx, stateUpdateFromChat(committed)) + return nil +} + +func (s *taskStarter) finishGenerationTurn( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + decision generationDecision, + fence generationAttemptFence, +) error { + if s.server.hookDispatcher == nil || !s.server.hookDispatcher.Enabled() { + return s.finishGenerationTurnWithoutHook(ctx, machine, input, decision, fence) + } + var chat database.Chat + var messages []database.ChatMessage + err := machine.ReadLock(ctx, func(store database.Store) error { + loadedChat, err := loadChatForGeneration(ctx, store, input, fence) + if err != nil { + return xerrors.Errorf("load chat for stop hook: %w", err) + } + loadedMessages, err := store.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ + ChatID: input.ChatID, + AfterID: 0, + }) + if err != nil { + return xerrors.Errorf("load messages for stop hook: %w", err) + } + chat = loadedChat + messages = loadedMessages + return nil + }) + if err != nil { + return normalizeTaskTransitionError(err, "load stop hook state") + } + turnID := activeTurnID(messages) + response, err := s.server.dispatchLifecycleHook(ctx, chat, turnID, agenthooks.EventStop, agenthooks.StopData{}) + if err != nil { + return s.finishGenerationError(ctx, machine, input, generationHookDispatchError(agenthooks.EventStop, err), fence) + } + prefixMessages, err := hookPrefixMessages(response, chat.LastModelConfigID, turnID) + if err != nil { + return s.finishGenerationError(ctx, machine, input, err, fence) + } + continueTurn := response.ModelContext != "" && !response.EndChat && input.StopNudges.claim(turnID) + + var committed database.Chat + var endedDescendants []database.Chat + ended := false + err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := loadChatForGeneration(ctx, store, input, fence); err != nil { + return xerrors.Errorf("load chat for generation: %w", err) + } + if err := applyHookAllowedTools(ctx, store, input.ChatID, response); err != nil { + return err + } + if response.EndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return xerrors.Errorf("tx.EndChatFamily: %w", err) + } + endedDescendants = endResult.EndedDescendants + ended = true + } else { + if len(prefixMessages) > 0 { + if _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: prefixMessages}); err != nil { + return xerrors.Errorf("commit stop hook messages: %w", err) + } + } + if !continueTurn { + finishResult, err := tx.FinishTurn(chatstate.FinishTurnInput{}) + if err != nil { + return xerrors.Errorf("tx.FinishTurn: %w", err) + } + if finishResult.PromotedMessage != nil { + decision.promotedMessageID = finishResult.PromotedMessage.ID + } + committed = finishResult.Chat + return nil + } + } + loadedChat, err := store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("load committed chat: %w", err) + } + committed = loadedChat + return nil + }) + if err != nil { + if continueTurn { + input.StopNudges.cancel(turnID) + } + err := normalizeTaskTransitionError(err, "finish generation turn") + recordGenerationFinishFailure(input.DebugTurn, err) + return err + } + if ended { + input.StopNudges.reset() + return s.finishEndedChat(ctx, input, committed, endedDescendants) + } + if continueTurn { + s.routeStateHint(ctx, stateUpdateFromChat(committed)) + return s.afterGenerationOutcome(ctx, generationOutcome{ + Chat: committed, + Kind: runnerActionKind(generationActionGenerateAssistant), + }) + } + input.StopNudges.reset() input.DebugTurn.RecordOutcome(chatdebug.StatusCompleted) watchCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), postCommitWatchPublishTimeout) defer cancel() diff --git a/coderd/x/chatd/generation_internal_test.go b/coderd/x/chatd/generation_internal_test.go index aa8e93a3d9609..c84ea68de4df1 100644 --- a/coderd/x/chatd/generation_internal_test.go +++ b/coderd/x/chatd/generation_internal_test.go @@ -1,15 +1,18 @@ package chatd //nolint:testpackage // Exercises unexported generation helpers. import ( + "encoding/json" "testing" "github.com/stretchr/testify/require" "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/chatloop" "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" ) @@ -87,3 +90,66 @@ func TestRecordGenerationFinishFailure(t *testing.T) { }) } } + +func TestPriorToolCallIDsInTurn(t *testing.T) { + t.Parallel() + + t.Run("ExcludesEarlierTurns", func(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("first prompt")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("reused-1", "run_command", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("reused-1", "run_command", json.RawMessage(`{}`), false, false)), + dbMessage(t, 4, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("second prompt")), + dbMessage(t, 5, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("reused-1", "run_command", json.RawMessage(`{}`))), + } + prior, err := priorToolCallIDsInTurn(messages) + require.NoError(t, err) + require.Empty(t, prior) + }) + + t.Run("IncludesEarlierStepsInTurn", func(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("prompt")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "run_command", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("call-1", "run_command", json.RawMessage(`{}`), false, false)), + dbMessage(t, 4, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "run_command", json.RawMessage(`{}`))), + } + prior, err := priorToolCallIDsInTurn(messages) + require.NoError(t, err) + require.Equal(t, map[string]bool{"call-1": true}, prior) + }) + + t.Run("HookContextDoesNotSplitTurn", func(t *testing.T) { + t.Parallel() + + hookContext := dbMessage(t, 4, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hook context")) + hookContext.Visibility = database.ChatMessageVisibilityModel + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("prompt")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "run_command", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("call-1", "run_command", json.RawMessage(`{}`), false, false)), + hookContext, + dbMessage(t, 5, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("call-1", "run_command", json.RawMessage(`{}`))), + } + prior, err := priorToolCallIDsInTurn(messages) + require.NoError(t, err) + require.Equal(t, map[string]bool{"call-1": true}, prior) + }) + + t.Run("NoAssistantInCurrentTurn", func(t *testing.T) { + t.Parallel() + + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("first prompt")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageToolCall("old-1", "run_command", json.RawMessage(`{}`))), + dbMessage(t, 3, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("second prompt")), + } + prior, err := priorToolCallIDsInTurn(messages) + require.NoError(t, err) + require.Empty(t, prior) + }) +} diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 06b47af4e4218..6f25c8869e3aa 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -9,6 +9,7 @@ import ( "charm.land/fantasy" "github.com/google/uuid" + "github.com/sqlc-dev/pqtype" "golang.org/x/sync/errgroup" "golang.org/x/xerrors" @@ -25,6 +26,56 @@ import ( "github.com/coder/coder/v2/codersdk" ) +func filterHookAllowedTools( + hookAllowedTools pqtype.NullRawMessage, + tools []fantasy.AgentTool, + providerTools []chatloop.ProviderTool, +) ([]fantasy.AgentTool, []chatloop.ProviderTool, map[string]struct{}, error) { + if !hookAllowedTools.Valid { + return tools, providerTools, nil, nil + } + var names []string + if err := json.Unmarshal(hookAllowedTools.RawMessage, &names); err != nil { + return nil, nil, nil, xerrors.Errorf("decode hook allowed tools: %w", err) + } + allowed := make(map[string]struct{}, len(names)) + for _, name := range names { + allowed[name] = struct{}{} + } + tools = slices.DeleteFunc(tools, func(tool fantasy.AgentTool) bool { + _, ok := allowed[tool.Info().Name] + return !ok + }) + providerTools = slices.DeleteFunc(providerTools, func(tool chatloop.ProviderTool) bool { + _, ok := allowed[tool.Definition.GetName()] + return !ok + }) + return tools, providerTools, allowed, nil +} + +func filterToolNameMap(names map[string]bool, allowed map[string]struct{}) { + for name := range names { + if _, ok := allowed[name]; !ok { + delete(names, name) + } + } +} + +// hookAllowedToolsNotice prevents the model from fabricating tools removed +// from the request schema by hook policy. +func hookAllowedToolsNotice(toolCount int) string { + if toolCount == 0 { + return "An external policy has removed all tool access for this chat. " + + "Respond in plain text only. Do not write tool-call syntax in your " + + "reply and do not claim any tool action completed." + } + return "Tool access for this chat is restricted by an external policy. " + + "Use only the tools available in this request. If a tool you need is " + + "unavailable, tell the user that policy prevents the action. Do not " + + "write tool-call syntax in your reply for unavailable tools and do " + + "not claim an unexecuted tool action completed." +} + func (server *Server) prepareGeneration( ctx context.Context, input generationPrepareInput, @@ -534,6 +585,19 @@ func (server *Server) prepareGeneration( } } + tools, providerTools, hookAllowedNames, err := filterHookAllowedTools(chat.HookAllowedTools, tools, providerTools) + if err != nil { + cleanup() + return generationPrepared{}, err + } + if chat.HookAllowedTools.Valid { + filterToolNameMap(builtinToolNames, hookAllowedNames) + filterToolNameMap(exclusiveToolNames, hookAllowedNames) + // Excluded dynamic tools must resolve as inactive instead of requiring action. + filterToolNameMap(dynamicToolNames, hookAllowedNames) + prompt = chatprompt.InsertSystem(prompt, hookAllowedToolsNotice(len(tools)+len(providerTools))) + } + var requestedEffort *string if chat.LastReasoningEffort.Valid { requestedEffort = new(string(chat.LastReasoningEffort.ChatReasoningEffort)) diff --git a/coderd/x/chatd/hooks.go b/coderd/x/chatd/hooks.go new file mode 100644 index 0000000000000..d1ff3455d23ab --- /dev/null +++ b/coderd/x/chatd/hooks.go @@ -0,0 +1,944 @@ +package chatd + +import ( + "bytes" + "context" + "database/sql" + "encoding/json" + "errors" + "fmt" + "io" + "slices" + "strings" + + "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/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/chatstate" + "github.com/coder/coder/v2/coderd/x/chathooks" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" +) + +const ( + sessionStartSourceStartup = "startup" + sessionStartSourceResume = "resume" + sessionStartSourceClear = "clear" +) + +func lifecycleHookEvent( + chat database.Chat, + turnID *uuid.UUID, + eventType agenthooks.EventType, + data any, +) chathooks.Event { + var workspaceID *uuid.UUID + if chat.WorkspaceID.Valid { + workspaceID = &chat.WorkspaceID.UUID + } + var parentChatID *uuid.UUID + if chat.ParentChatID.Valid { + parentChatID = &chat.ParentChatID.UUID + } + var rootChatID *uuid.UUID + if chat.RootChatID.Valid { + rootChatID = &chat.RootChatID.UUID + } + return chathooks.Event{ + Type: eventType, + ChatRef: agenthooks.ChatRef{ + ChatID: chat.ID, + OwnerID: chat.OwnerID, + WorkspaceID: workspaceID, + TurnID: turnID, + ParentChatID: parentChatID, + RootChatID: rootChatID, + }, + Data: data, + } +} + +func (p *Server) dispatchLifecycleHook( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + eventType agenthooks.EventType, + data any, +) (agenthooks.Response, error) { + if p.hookDispatcher == nil || !p.hookDispatcher.Enabled() { + return agenthooks.Response{}, nil + } + resp, _, err := p.hookDispatcher.Dispatch(ctx, lifecycleHookEvent(chat, turnID, eventType, data)) + return resp, err +} + +type preToolUseResult struct { + Step chatloop.PersistedStep + Responses []agenthooks.Response + EffectDispatchIDs []uuid.UUID +} + +func (p *Server) dispatchPreToolUse( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + toolCall fantasy.ToolCallContent, +) (agenthooks.Response, uuid.UUID, error) { + event := lifecycleHookEvent(chat, turnID, agenthooks.EventPreToolUse, agenthooks.PreToolUseData{ + ToolUseID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + ToolInput: json.RawMessage(toolCall.Input), + }) + return p.hookDispatcher.Dispatch(ctx, event) +} + +func (p *Server) dispatchPostToolUseData( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + data agenthooks.PostToolUseData, +) (agenthooks.Response, error) { + event := lifecycleHookEvent(chat, turnID, agenthooks.EventPostToolUse, data) + resp, _, err := p.hookDispatcher.Dispatch(ctx, event) + return resp, err +} + +func (p *Server) dispatchPostToolUse( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + toolResult fantasy.ToolResultContent, +) (agenthooks.Response, error) { + data := agenthooks.PostToolUseData{ + ToolUseID: toolResult.ToolCallID, + ToolName: toolResult.ToolName, + } + switch output := toolResult.Result.(type) { + case fantasy.ToolResultOutputContentError: + if output.Error != nil { + data.ToolError = output.Error.Error() + } + case *fantasy.ToolResultOutputContentError: + if output != nil && output.Error != nil { + data.ToolError = output.Error.Error() + } + default: + encoded, err := json.Marshal(toolResult.Result) + if err != nil { + return agenthooks.Response{}, xerrors.Errorf("marshal post_tool_use response: %w", err) + } + data.ToolResponse = encoded + } + return p.dispatchPostToolUseData(ctx, chat, turnID, data) +} + +func (p *Server) dispatchPostToolUseResults( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + content []fantasy.Content, +) ([]agenthooks.Response, error) { + if p.hookDispatcher == nil || !p.hookDispatcher.Enabled() { + return nil, nil + } + responses := make([]agenthooks.Response, 0, len(content)) + // Dispatch every completed result so each outcome is recorded in + // chat_hook_dispatches. Preserve only the first failure before an + // accepted end_chat. + var firstErr error + endChatSeen := false + for _, block := range content { + toolResult, ok := asToolResultContent(block) + if !ok || toolResult.ProviderExecuted { + continue + } + response, err := p.dispatchPostToolUse(ctx, chat, turnID, toolResult) + if err != nil { + if firstErr == nil && !endChatSeen { + firstErr = err + } + continue + } + endChatSeen = endChatSeen || response.EndChat + responses = append(responses, response) + } + return responses, firstErr +} + +func (p *Server) preflightToolCalls( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + step chatloop.PersistedStep, + toolCalls []fantasy.ToolCallContent, +) (preToolUseResult, error) { + result := preToolUseResult{Step: step} + if p.hookDispatcher == nil || !p.hookDispatcher.Enabled() { + return result, nil + } + if err := rejectDuplicateToolUseIDs(toolCalls); err != nil { + return preToolUseResult{}, err + } + + for _, toolCall := range toolCalls { + if toolCall.ProviderExecuted { + continue + } + response, dispatchID, err := p.dispatchPreToolUse(ctx, chat, turnID, toolCall) + if err != nil { + return preToolUseResult{}, err + } + if err := applyPreToolUsePermission(&result.Step, toolCall, response); err != nil { + return preToolUseResult{}, err + } + result.Responses = append(result.Responses, response) + result.EffectDispatchIDs = append(result.EffectDispatchIDs, dispatchID) + // Accepted end_chat effects take precedence over later dispatch failures. + if response.EndChat { + break + } + } + return result, nil +} + +// restoreToolCallOrder reorders tool results to match the assistant's +// call order because providers pair results with calls positionally. +// Entries that are not tool results for the given calls keep their slots. +func restoreToolCallOrder(content []fantasy.Content, calls []fantasy.ToolCallContent) { + position := make(map[string]int, len(calls)) + for index, call := range calls { + position[call.ToolCallID] = index + } + slots := make([]int, 0, len(content)) + results := make([]fantasy.ToolResultContent, 0, len(content)) + for index, entry := range content { + result, ok := entry.(fantasy.ToolResultContent) + if !ok { + continue + } + if _, known := position[result.ToolCallID]; !known { + continue + } + slots = append(slots, index) + results = append(results, result) + } + slices.SortStableFunc(results, func(a, b fantasy.ToolResultContent) int { + return position[a.ToolCallID] - position[b.ToolCallID] + }) + for index, slot := range slots { + content[slot] = results[index] + } +} + +// rejectDuplicateToolUseIDs fails closed because persisted decisions are +// keyed by tool-use ID within a turn. +func rejectDuplicateToolUseIDs(toolCalls []fantasy.ToolCallContent) error { + seen := make(map[string]struct{}, len(toolCalls)) + for _, toolCall := range toolCalls { + if toolCall.ProviderExecuted { + continue + } + if _, ok := seen[toolCall.ToolCallID]; ok { + return xerrors.Errorf("duplicate tool use ID %q in one step; lifecycle hook decisions cannot be attributed unambiguously", toolCall.ToolCallID) + } + seen[toolCall.ToolCallID] = struct{}{} + } + return nil +} + +type preToolUseExecutionResult struct { + Allowed []fantasy.ToolCallContent + Denied []fantasy.ToolResultContent + Responses []agenthooks.Response + Overrides map[string]json.RawMessage + EffectDispatchIDs []uuid.UUID +} + +func (p *Server) preflightPendingToolCalls( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + toolCalls []fantasy.ToolCallContent, + priorToolCallIDs map[string]bool, +) (preToolUseExecutionResult, error) { + result := preToolUseExecutionResult{ + Allowed: make([]fantasy.ToolCallContent, 0, len(toolCalls)), + Overrides: make(map[string]json.RawMessage), + } + if p.hookDispatcher == nil || !p.hookDispatcher.Enabled() { + result.Allowed = append(result.Allowed, toolCalls...) + return result, nil + } + if err := rejectDuplicateToolUseIDs(toolCalls); err != nil { + return preToolUseExecutionResult{}, err + } + + for _, toolCall := range toolCalls { + // Re-dispatch invalid JSON or IDs reused earlier in the turn because no + // persisted decision can safely authorize those calls. + row, err := database.ChatHookDispatch{}, sql.ErrNoRows + if json.Valid([]byte(toolCall.Input)) && !priorToolCallIDs[toolCall.ToolCallID] { + row, err = p.db.GetChatHookDispatchDecision(ctx, database.GetChatHookDispatchDecisionParams{ + ChatID: chat.ID, + ToolUseID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + ToolInput: json.RawMessage(toolCall.Input), + TurnID: hookTurnID(turnID), + }) + } + if err == nil { + // Replay unapplied effects from a finalized dispatch without re-dispatching. + endChat := false + if !row.EffectsAppliedAt.Valid { + response := dispatchRowResponse(row) + endChat = response.EndChat + result.Responses = append(result.Responses, response) + result.EffectDispatchIDs = append(result.EffectDispatchIDs, row.ID) + } + switch agenthooks.PermissionDecision(row.Decision.String) { + case agenthooks.PermissionAllow: + if row.InputOverride.Valid { + toolCall.Input = string(row.InputOverride.RawMessage) + result.Overrides[toolCall.ToolCallID] = row.InputOverride.RawMessage + } + result.Allowed = append(result.Allowed, toolCall) + case agenthooks.PermissionDeny: + result.Denied = append(result.Denied, deniedToolResult(toolCall, row.DecisionReason.String)) + } + // Accepted end_chat effects take precedence over later dispatch failures. + if endChat { + break + } + continue + } + if !errors.Is(err, sql.ErrNoRows) { + return preToolUseExecutionResult{}, xerrors.Errorf("get pre_tool_use decision: %w", err) + } + + response, dispatchID, err := p.dispatchPreToolUse(ctx, chat, turnID, toolCall) + if err != nil { + return preToolUseExecutionResult{}, err + } + result.Responses = append(result.Responses, response) + result.EffectDispatchIDs = append(result.EffectDispatchIDs, dispatchID) + if response.Permission == nil { + result.Allowed = append(result.Allowed, toolCall) + } else { + switch response.Permission.Decision { + case agenthooks.PermissionAllow: + toolCall.Input = string(response.Permission.InputOverride) + result.Overrides[toolCall.ToolCallID] = response.Permission.InputOverride + result.Allowed = append(result.Allowed, toolCall) + case agenthooks.PermissionDeny: + result.Denied = append(result.Denied, deniedToolResult(toolCall, response.Permission.Reason)) + } + } + // Accepted end_chat effects take precedence over later dispatch failures. + if response.EndChat { + break + } + } + return result, nil +} + +func dispatchRowResponse(row database.ChatHookDispatch) agenthooks.Response { + response := agenthooks.Response{ + EndChat: row.EndChat.Valid && row.EndChat.Bool, + } + if row.ModelContext.Valid { + response.ModelContext = row.ModelContext.String + } + if row.UserMessage.Valid { + response.UserMessage = row.UserMessage.String + } + if row.AllowedTools.Valid { + var allowed []string + if err := json.Unmarshal(row.AllowedTools.RawMessage, &allowed); err == nil { + response.AllowedTools = &allowed + } + } + return response +} + +// Mark effects in the same transaction so rollback leaves them replayable. +func markHookDispatchEffectsApplied( + ctx context.Context, + store database.Store, + chatID uuid.UUID, + dispatchIDs []uuid.UUID, +) error { + if len(dispatchIDs) == 0 { + return nil + } + if err := store.MarkChatHookDispatchEffectsApplied(ctx, database.MarkChatHookDispatchEffectsAppliedParams{ + ChatID: chatID, + DispatchIds: dispatchIDs, + }); err != nil { + return xerrors.Errorf("mark hook dispatch effects applied: %w", err) + } + return nil +} + +func replacePersistedToolCallInputs( + ctx context.Context, + store database.Store, + chatID uuid.UUID, + overrides map[string]json.RawMessage, +) error { + if len(overrides) == 0 { + return nil + } + assistant, err := store.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chatID, + Role: database.ChatMessageRoleAssistant, + }) + if err != nil { + return xerrors.Errorf("get assistant message for tool override: %w", err) + } + parts, err := chatprompt.ParseContent(assistant) + if err != nil { + return xerrors.Errorf("parse assistant message for tool override: %w", err) + } + for i := range parts { + if override, ok := overrides[parts[i].ToolCallID]; ok && parts[i].Type == codersdk.ChatMessagePartTypeToolCall { + parts[i].Args = override + } + } + content, err := chatprompt.MarshalParts(parts) + if err != nil { + return xerrors.Errorf("marshal assistant message with tool override: %w", err) + } + if err := store.UpdateChatMessageContentByID(ctx, database.UpdateChatMessageContentByIDParams{ + Content: content.RawMessage, + ID: assistant.ID, + }); err != nil { + return xerrors.Errorf("update assistant message with tool override: %w", err) + } + return nil +} + +func applyPreToolUsePermission(step *chatloop.PersistedStep, toolCall fantasy.ToolCallContent, response agenthooks.Response) error { + if response.Permission == nil { + return nil + } + switch response.Permission.Decision { + case agenthooks.PermissionAllow: + if !replaceToolCallInput(step.Content, toolCall.ToolCallID, string(response.Permission.InputOverride)) { + return xerrors.Errorf("tool call %q is missing from generated step", toolCall.ToolCallID) + } + case agenthooks.PermissionDeny: + step.Content = append(step.Content, deniedToolResult(toolCall, response.Permission.Reason)) + } + return nil +} + +func deniedToolResult(toolCall fantasy.ToolCallContent, reason string) fantasy.ToolResultContent { + reason = strings.TrimSpace(reason) + if reason == "" { + reason = "denied by lifecycle hook" + } + return fantasy.ToolResultContent{ + ToolCallID: toolCall.ToolCallID, + ToolName: toolCall.ToolName, + Result: fantasy.ToolResultOutputContentError{ + Error: xerrors.New("DENIED: " + reason), + }, + } +} + +func replaceToolCallInput(content []fantasy.Content, toolCallID, input string) bool { + for i, block := range content { + if toolCall, ok := fantasy.AsContentType[fantasy.ToolCallContent](block); ok && toolCall.ToolCallID == toolCallID { + toolCall.Input = input + content[i] = toolCall + return true + } + if toolCall, ok := fantasy.AsContentType[*fantasy.ToolCallContent](block); ok && toolCall != nil && toolCall.ToolCallID == toolCallID { + updated := *toolCall + updated.Input = input + content[i] = updated + return true + } + } + return false +} + +func sessionStartSource(messages []database.ChatMessage) string { + for _, message := range messages { + if message.Role == database.ChatMessageRoleAssistant { + return sessionStartSourceResume + } + } + return sessionStartSourceStartup +} + +func activeTurnID(messages []database.ChatMessage) *uuid.UUID { + for i := len(messages) - 1; i >= 0; i-- { + if messages[i].TurnID.Valid { + turnID := messages[i].TurnID.UUID + return &turnID + } + } + return nil +} + +// UserPromptDeniedError reports that a lifecycle hook rejected a prompt. +type UserPromptDeniedError struct { + UserMessage string +} + +func (*UserPromptDeniedError) Error() string { + return "user prompt denied by lifecycle hook" +} + +func (p *Server) dispatchUserPromptSubmit( + ctx context.Context, + chat database.Chat, + turnID uuid.UUID, + parts []codersdk.ChatMessagePart, +) (agenthooks.Response, error) { + encodedParts, err := chatprompt.MarshalParts(parts) + if err != nil { + return agenthooks.Response{}, xerrors.Errorf("marshal prompt parts for hook: %w", err) + } + response, err := p.dispatchLifecycleHook(ctx, chat, &turnID, agenthooks.EventUserPromptSubmit, agenthooks.UserPromptSubmitData{ + Prompt: textFromParts(parts), + Parts: encodedParts.RawMessage, + }) + if err != nil { + return agenthooks.Response{}, err + } + if response.Permission != nil && response.Permission.Decision == agenthooks.PermissionDeny { + // Preserve end_chat from a denied response. + return response, &UserPromptDeniedError{UserMessage: response.UserMessage} + } + return response, nil +} + +// Treat an already archived or deleted chat as satisfying accepted end_chat. +func (p *Server) endChatAfterPromptDenial(ctx context.Context, chatID uuid.UUID, prefixMessages []chatstate.Message) error { + var ( + ended database.Chat + descendants []database.Chat + ) + machine := p.newChatMachine(chatID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return err + } + descendants = endResult.EndedDescendants + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return xerrors.Errorf("reload ended chat: %w", err) + } + ended = chat + return nil + }) + if errors.Is(err, chatstate.ErrTransitionNotAllowed) || errors.Is(err, chatstate.ErrChatNotFound) { + return nil + } + if err != nil { + return xerrors.Errorf("end chat after prompt denial: %w", err) + } + p.publishEndChatSideEffects(ctx, ended, descendants) + return nil +} + +func (p *Server) endChatFromEditSessionStart( + ctx context.Context, + chat database.Chat, + turnID *uuid.UUID, + response agenthooks.Response, +) (EditMessageResult, error) { + prefixMessages, err := hookPrefixMessages(response, chat.LastModelConfigID, turnID) + if err != nil { + return EditMessageResult{}, err + } + var ( + result EditMessageResult + descendants []database.Chat + ) + machine := p.newChatMachine(chat.ID) + err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return xerrors.Errorf("end chat from session_start: %w", err) + } + descendants = endResult.EndedDescendants + refreshed, err := store.GetChatByID(ctx, chat.ID) + if err != nil { + return xerrors.Errorf("reload ended chat: %w", err) + } + result.Chat = refreshed + result.Ended = true + return nil + }) + if err != nil { + return EditMessageResult{}, err + } + p.publishEndChatSideEffects(ctx, result.Chat, descendants) + return result, nil +} + +func (p *Server) endChatAfterToolHookFailure( + ctx context.Context, + machine *chatstate.ChatMachine, + chatID uuid.UUID, + suffixMessages []chatstate.Message, +) error { + var ( + ended database.Chat + descendants []database.Chat + ) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: suffixMessages}) + if err != nil { + return err + } + descendants = endResult.EndedDescendants + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return xerrors.Errorf("reload ended chat: %w", err) + } + ended = chat + return nil + }) + if errors.Is(err, chatstate.ErrTransitionNotAllowed) || errors.Is(err, chatstate.ErrChatNotFound) { + return nil + } + if err != nil { + return xerrors.Errorf("end chat after tool hook failure: %w", err) + } + p.publishEndChatSideEffects(ctx, ended, descendants) + return nil +} + +func (p *Server) handleUserPromptDispatchError(ctx context.Context, chatID uuid.UUID, dispatchErr error) error { + return p.handleAPIDispatchError(ctx, chatID, agenthooks.EventUserPromptSubmit, dispatchErr) +} + +func (p *Server) handleAPIDispatchError(ctx context.Context, chatID uuid.UUID, eventType agenthooks.EventType, dispatchErr error) error { + lastError, ok := hookDispatchErrorMessage(eventType, dispatchErr) + if !ok { + return dispatchErr + } + var failedChat database.Chat + machine := p.newChatMachine(chatID) + err := machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := tx.FailIdle(chatstate.FailIdleInput{ + LastError: lastError, + Kind: codersdk.ChatErrorKindHookDispatchFailed, + }); err != nil { + return err + } + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return xerrors.Errorf("reload chat after hook failure: %w", err) + } + failedChat = chat + return nil + }) + if errors.Is(err, chatstate.ErrTransitionNotAllowed) { + return dispatchErr + } + if err != nil { + return errors.Join(dispatchErr, xerrors.Errorf("fail idle chat after hook dispatch: %w", err)) + } + p.publishChatPubsubEvent(failedChat, codersdk.ChatWatchEventKindStatusChange, nil) + return dispatchErr +} + +func hookDispatchErrorMessage(eventType agenthooks.EventType, dispatchErr error) (string, bool) { + var structured *chathooks.DispatchError + if !errors.As(dispatchErr, &structured) { + return "", false + } + return fmt.Sprintf( + "hook dispatch failed: %s: %s (dispatch %s)", + eventType, + structured.Class, + structured.DispatchID, + ), true +} + +func sessionStartDispatchError(dispatchErr error) error { + return generationHookDispatchError(agenthooks.EventSessionStart, dispatchErr) +} + +func generationHookDispatchError(eventType agenthooks.EventType, dispatchErr error) error { + message, ok := hookDispatchErrorMessage(eventType, dispatchErr) + if !ok { + message = dispatchErr.Error() + } + return chaterror.WithClassification(dispatchErr, chaterror.ClassifiedError{ + Message: message, + Kind: codersdk.ChatErrorKindHookDispatchFailed, + }) +} + +type sessionStartResult struct { + Chat database.Chat + Ended bool + EndedDescendants []database.Chat +} + +func applySessionStartResponse( + ctx context.Context, + machine *chatstate.ChatMachine, + input chatWorkerTaskStartInput, + chat database.Chat, + turnID *uuid.UUID, + response agenthooks.Response, +) (sessionStartResult, error) { + if response.ModelContext == "" && response.UserMessage == "" && response.AllowedTools == nil && !response.EndChat { + return sessionStartResult{Chat: chat}, nil + } + + // Pre-turn-ID histories persist hook messages with a NULL turn_id. + prefixMessages, err := hookPrefixMessages(response, chat.LastModelConfigID, turnID) + if err != nil { + return sessionStartResult{}, err + } + + var result sessionStartResult + err = machine.Update(ctx, func(tx *chatstate.Tx, store database.Store) error { + if _, err := loadChatForGeneration(ctx, store, input, generationAttemptNotRequired); err != nil { + return xerrors.Errorf("load chat for session_start response: %w", err) + } + if err := applyHookAllowedTools(ctx, store, input.ChatID, response); err != nil { + return err + } + if response.EndChat { + endResult, err := tx.EndChatFamily(chatstate.EndChatInput{PrefixMessages: prefixMessages}) + if err != nil { + return xerrors.Errorf("end chat from session_start: %w", err) + } + result.Ended = true + result.EndedDescendants = endResult.EndedDescendants + } else if len(prefixMessages) > 0 { + if _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: prefixMessages}); err != nil { + return xerrors.Errorf("insert session_start response messages: %w", err) + } + } + result.Chat, err = store.GetChatByID(ctx, input.ChatID) + if err != nil { + return xerrors.Errorf("reload chat after session_start response: %w", err) + } + return nil + }) + if err != nil { + return sessionStartResult{}, normalizeTaskTransitionError(err, "apply session_start response") + } + return result, nil +} + +func (s *taskStarter) finishSessionStartEnd(ctx context.Context, input chatWorkerTaskStartInput, result sessionStartResult) error { + return s.finishEndedChat(ctx, input, result.Chat, result.EndedDescendants) +} + +func hookPrefixMessages(response agenthooks.Response, modelConfigID uuid.UUID, turnID *uuid.UUID) ([]chatstate.Message, error) { + messages := make([]chatstate.Message, 0, 2) + if response.ModelContext != "" { + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(response.ModelContext)}) + if err != nil { + return nil, xerrors.Errorf("marshal hook model context: %w", err) + } + messages = append(messages, chatstate.Message{ + Role: database.ChatMessageRoleUser, + Content: content, + Visibility: database.ChatMessageVisibilityModel, + TurnID: hookTurnID(turnID), + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + }) + } + if response.UserMessage != "" { + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(response.UserMessage)}) + if err != nil { + return nil, xerrors.Errorf("marshal hook user message: %w", err) + } + messages = append(messages, chatstate.Message{ + Role: database.ChatMessageRoleSystem, + Content: content, + Visibility: database.ChatMessageVisibilityUser, + TurnID: hookTurnID(turnID), + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil}, + ContentVersion: chatprompt.CurrentContentVersion, + }) + } + return messages, nil +} + +func hookTurnID(turnID *uuid.UUID) uuid.NullUUID { + if turnID == nil { + return uuid.NullUUID{} + } + return uuid.NullUUID{UUID: *turnID, Valid: true} +} + +func hookResponseMessages( + responses []agenthooks.Response, + modelConfigID uuid.UUID, + turnID *uuid.UUID, +) ([]chatstate.Message, bool, error) { + var messages []chatstate.Message + endChat := false + for _, response := range responses { + responseMessages, err := hookPrefixMessages(response, modelConfigID, turnID) + if err != nil { + return nil, false, err + } + messages = append(messages, responseMessages...) + endChat = endChat || response.EndChat + } + return messages, endChat, nil +} + +func applyHookResponseMessages( + messages stepMessagesForCommit, + responses []agenthooks.Response, + modelConfigID uuid.UUID, + turnID *uuid.UUID, +) (stepMessagesForCommit, bool, error) { + prefix, endChat, err := hookResponseMessages(responses, modelConfigID, turnID) + if err != nil { + return stepMessagesForCommit{}, false, err + } + if len(prefix) > 0 { + messages.Messages = append(prefix, messages.Messages...) + messages.VisibleIndexes = visibleMessageIndexes(messages.Messages) + } + return messages, endChat, nil +} + +func appendHookResponseMessages( + messages stepMessagesForCommit, + responses []agenthooks.Response, + modelConfigID uuid.UUID, + turnID *uuid.UUID, +) (stepMessagesForCommit, bool, error) { + suffix, endChat, err := hookResponseMessages(responses, modelConfigID, turnID) + if err != nil { + return stepMessagesForCommit{}, false, err + } + if len(suffix) > 0 { + messages.Messages = append(messages.Messages, suffix...) + messages.VisibleIndexes = visibleMessageIndexes(messages.Messages) + } + return messages, endChat, nil +} + +func hookAllowedTools(response agenthooks.Response) (pqtype.NullRawMessage, error) { + if response.AllowedTools == nil { + return pqtype.NullRawMessage{}, nil + } + encoded, err := json.Marshal(response.AllowedTools) + if err != nil { + return pqtype.NullRawMessage{}, xerrors.Errorf("marshal hook allowed tools: %w", err) + } + return pqtype.NullRawMessage{RawMessage: encoded, Valid: true}, nil +} + +type pendingHookAllowedToolsContextKey struct{} + +type pendingHookAllowedToolsContextValue struct { + chatID uuid.UUID + policy pqtype.NullRawMessage +} + +func withPendingHookAllowedTools(ctx context.Context, chatID uuid.UUID, policy pqtype.NullRawMessage) context.Context { + return context.WithValue(ctx, pendingHookAllowedToolsContextKey{}, pendingHookAllowedToolsContextValue{ + chatID: chatID, + policy: policy, + }) +} + +func pendingHookAllowedToolsFromContext(ctx context.Context, chatID uuid.UUID) (pqtype.NullRawMessage, bool) { + value, ok := ctx.Value(pendingHookAllowedToolsContextKey{}).(pendingHookAllowedToolsContextValue) + if !ok || value.chatID != chatID { + return pqtype.NullRawMessage{}, false + } + return value.policy, true +} + +func narrowHookAllowedToolsResponses(current pqtype.NullRawMessage, responses []agenthooks.Response) (pqtype.NullRawMessage, bool, error) { + narrowed := current + applied := false + for _, response := range responses { + incoming, err := hookAllowedTools(response) + if err != nil { + return pqtype.NullRawMessage{}, false, err + } + if !incoming.Valid { + continue + } + narrowed, err = chatstate.NarrowHookAllowedTools(narrowed, incoming) + if err != nil { + return pqtype.NullRawMessage{}, false, err + } + applied = true + } + return narrowed, applied, nil +} + +func applyHookAllowedTools(ctx context.Context, store database.Store, chatID uuid.UUID, response agenthooks.Response) error { + allowedTools, err := hookAllowedTools(response) + if err != nil { + return err + } + if !allowedTools.Valid { + return nil + } + chat, err := store.GetChatByID(ctx, chatID) + if err != nil { + return xerrors.Errorf("load chat for hook allowed tools: %w", err) + } + narrowed, err := chatstate.NarrowHookAllowedTools(chat.HookAllowedTools, allowedTools) + if err != nil { + return err + } + if err := store.UpdateChatHookAllowedTools(ctx, database.UpdateChatHookAllowedToolsParams{ + HookAllowedTools: narrowed, + ID: chatID, + }); err != nil { + return xerrors.Errorf("update hook allowed tools: %w", err) + } + return nil +} + +func applyHookAllowedToolsResponses(ctx context.Context, store database.Store, chatID uuid.UUID, responses []agenthooks.Response) error { + for _, response := range responses { + if err := applyHookAllowedTools(ctx, store, chatID, response); err != nil { + return err + } + } + return nil +} + +func userPromptOverride(response agenthooks.Response) (string, bool, error) { + if response.Permission == nil || response.Permission.Decision != agenthooks.PermissionAllow { + return "", false, nil + } + var override struct { + Prompt *string `json:"prompt"` + } + decoder := json.NewDecoder(bytes.NewReader(response.Permission.InputOverride)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&override); err != nil { + return "", false, xerrors.Errorf("decode user prompt input override: %w", err) + } + if override.Prompt == nil { + return "", false, xerrors.New("decode user prompt input override: prompt is required") + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return "", false, xerrors.New("decode user prompt input override: trailing JSON value") + } + return *override.Prompt, true, nil +} diff --git a/coderd/x/chatd/hooks_internal_test.go b/coderd/x/chatd/hooks_internal_test.go new file mode 100644 index 0000000000000..80faecbea9323 --- /dev/null +++ b/coderd/x/chatd/hooks_internal_test.go @@ -0,0 +1,479 @@ +package chatd + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + "time" + + "charm.land/fantasy" + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "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/chatloop" + "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/chathooks" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" +) + +func TestSessionStartDispatchSources(t *testing.T) { + t.Parallel() + + const secret = "test-hook-secret-32-bytes-minimum!!" + type received struct { + request agenthooks.Request + claims agenthooks.Claims + data *agenthooks.SessionStartData + } + receivedCh := make(chan received, 2) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + claims, err := agenthooks.Verify(r.Header.Get("Authorization"), []byte(secret)) + require.NoError(t, err) + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.SessionStartData) + require.True(t, ok) + receivedCh <- received{request: request, claims: claims, data: data} + _, err = w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + db, _ := dbtestutil.NewDB(t) + dispatcher := chathooks.New( + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + db, + consumer.Client(), + consumer.URL, + secret, + time.Second, + "test-deployment", + "test-version", + prometheus.NewRegistry(), + ) + server := &Server{hookDispatcher: dispatcher} + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{OwnerID: user.ID, OrganizationID: org.ID, LastModelConfigID: model.ID}) + turnID := uuid.New() + ctx := testutil.Context(t, testutil.WaitLong) + + _, err := server.dispatchLifecycleHook(ctx, chat, &turnID, agenthooks.EventSessionStart, agenthooks.SessionStartData{Source: sessionStartSource(nil)}) + require.NoError(t, err) + _, err = server.dispatchLifecycleHook(ctx, chat, &turnID, agenthooks.EventSessionStart, agenthooks.SessionStartData{Source: sessionStartSource([]database.ChatMessage{{Role: database.ChatMessageRoleAssistant}})}) + require.NoError(t, err) + + startup := <-receivedCh + resume := <-receivedCh + require.Equal(t, agenthooks.EventSessionStart, startup.request.Type) + require.Equal(t, sessionStartSourceStartup, startup.data.Source) + require.Equal(t, startup.request.Meta.DispatchID, startup.claims.JTI) + require.Equal(t, agenthooks.EventSessionStart, resume.request.Type) + require.Equal(t, sessionStartSourceResume, resume.data.Source) + require.Equal(t, resume.request.Meta.DispatchID, resume.claims.JTI) + require.NotEqual(t, startup.claims.JTI, resume.claims.JTI) +} + +func TestSessionStartTrackerRetriesIncompleteDispatch(t *testing.T) { + t.Parallel() + tracker := &sessionStartTracker{} + claimed, complete, err := tracker.claim(t.Context()) + require.NoError(t, err) + require.True(t, claimed) + + canceled, cancel := context.WithCancel(t.Context()) + cancel() + _, _, err = tracker.claim(canceled) + require.ErrorIs(t, err, context.Canceled) + complete(false) + + claimed, complete, err = tracker.claim(t.Context()) + require.NoError(t, err) + require.True(t, claimed) + complete(true) + claimed, _, err = tracker.claim(t.Context()) + require.NoError(t, err) + require.False(t, claimed) +} + +func TestSessionStartDispatchFailureFinishesGeneration(t *testing.T) { + t.Parallel() + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + chat = f.acquireChat(t, chat.ID, workerID, runnerID) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(consumer.Close) + dispatcher := chathooks.New( + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + f.db, + consumer.Client(), + consumer.URL, + "test-hook-secret-32-bytes-minimum!!", + time.Second, + "test-deployment", + "test-version", + prometheus.NewRegistry(), + ) + starter := newTestTaskStarter(t, f, newTaskSideEffectRecorder()) + starter.server.hookDispatcher = dispatcher + ctx := testutil.Context(t, testutil.WaitLong) + debugTurn := newRunnerDebugTurn(ctx, starter.opts.Logger) + defer debugTurn.Finalize(ctx) + err := starter.StartGeneration(ctx, chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: chat.HistoryVersion, + GenerationAttempt: chat.GenerationAttempt, + Status: database.ChatStatusRunning, + DebugTurn: debugTurn, + SessionStart: &sessionStartTracker{}, + }) + require.NoError(t, err) + updated, err := f.db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusError, updated.Status) + var chatErr codersdk.ChatError + require.NoError(t, json.Unmarshal(updated.LastError.RawMessage, &chatErr)) + require.Equal(t, codersdk.ChatErrorKindHookDispatchFailed, chatErr.Kind) + require.Contains(t, chatErr.Message, "hook dispatch failed: session_start: http_error (dispatch ") + require.False(t, chatErr.Retryable) +} + +func TestApplySessionStartResponse(t *testing.T) { + t.Parallel() + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + chat = f.acquireChat(t, chat.ID, workerID, runnerID) + ctx := testutil.Context(t, testutil.WaitLong) + messages, err := f.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.NotEmpty(t, messages) + turnID := messages[len(messages)-1].TurnID.UUID + input := chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: chat.HistoryVersion, + Status: database.ChatStatusRunning, + } + result, err := applySessionStartResponse( + ctx, + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + input, + chat, + &turnID, + agenthooks.Response{ + ModelContext: "model context", + UserMessage: "user notice", + AllowedTools: &[]string{"read_file"}, + }, + ) + require.NoError(t, err) + require.False(t, result.Ended) + + promptRows, err := f.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, "model context", hookMessageTextInternal(t, promptRows[len(promptRows)-1])) + require.Equal(t, database.ChatMessageVisibilityModel, promptRows[len(promptRows)-1].Visibility) + allRows, err := f.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + userNotice := allRows[len(allRows)-1] + require.Equal(t, database.ChatMessageRoleSystem, userNotice.Role) + require.Equal(t, database.ChatMessageVisibilityUser, userNotice.Visibility) + require.Equal(t, "user notice", hookMessageTextInternal(t, userNotice)) + require.JSONEq(t, `["read_file"]`, string(result.Chat.HookAllowedTools.RawMessage)) +} + +func TestApplySessionStartResponseNilTurnID(t *testing.T) { + t.Parallel() + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + chat = f.acquireChat(t, chat.ID, workerID, runnerID) + ctx := testutil.Context(t, testutil.WaitLong) + input := chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: chat.HistoryVersion, + Status: database.ChatStatusRunning, + } + result, err := applySessionStartResponse( + ctx, + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + input, + chat, + nil, + agenthooks.Response{ + ModelContext: "legacy context", + UserMessage: "legacy notice", + }, + ) + require.NoError(t, err) + require.False(t, result.Ended) + + promptRows, err := f.db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + modelContext := promptRows[len(promptRows)-1] + require.Equal(t, "legacy context", hookMessageTextInternal(t, modelContext)) + require.False(t, modelContext.TurnID.Valid) + allRows, err := f.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + userNotice := allRows[len(allRows)-1] + require.Equal(t, "legacy notice", hookMessageTextInternal(t, userNotice)) + require.False(t, userNotice.TurnID.Valid) +} + +func TestApplySessionStartResponseNoOp(t *testing.T) { + t.Parallel() + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + f.pubsub.clear() + result, err := applySessionStartResponse( + testutil.Context(t, testutil.WaitLong), + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + chatWorkerTaskStartInput{}, + chat, + nil, + agenthooks.Response{}, + ) + require.NoError(t, err) + require.Equal(t, chat.SnapshotVersion, result.Chat.SnapshotVersion) + require.Empty(t, f.pubsub.events()) +} + +func TestApplySessionStartResponseEndChat(t *testing.T) { + t.Parallel() + f := newTaskTestFixture(t) + chat := f.createRunningChat(t) + workerID := uuid.New() + runnerID := uuid.New() + chat = f.acquireChat(t, chat.ID, workerID, runnerID) + ctx := testutil.Context(t, testutil.WaitLong) + messages, err := f.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + turnID := messages[len(messages)-1].TurnID.UUID + result, err := applySessionStartResponse( + ctx, + chatstate.NewChatMachine(f.db, f.pubsub, chat.ID), + chatWorkerTaskStartInput{ + ChatID: chat.ID, + WorkerID: workerID, + RunnerID: runnerID, + HistoryVersion: chat.HistoryVersion, + Status: database.ChatStatusRunning, + }, + chat, + &turnID, + agenthooks.Response{UserMessage: "ended by hook", EndChat: true}, + ) + require.NoError(t, err) + require.True(t, result.Ended) + require.True(t, result.Chat.Archived) + require.Equal(t, database.ChatStatusWaiting, result.Chat.Status) + require.False(t, result.Chat.WorkerID.Valid) + require.False(t, result.Chat.RunnerID.Valid) + rows, err := f.db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Equal(t, "ended by hook", hookMessageTextInternal(t, rows[len(rows)-1])) +} + +func hookMessageTextInternal(t *testing.T, message database.ChatMessage) string { + t.Helper() + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + require.Len(t, parts, 1) + return parts[0].Text +} + +func TestRejectDuplicateToolUseIDs(t *testing.T) { + t.Parallel() + + require.NoError(t, rejectDuplicateToolUseIDs([]fantasy.ToolCallContent{ + {ToolCallID: "first", ToolName: "read_file", Input: `{}`}, + {ToolCallID: "second", ToolName: "execute", Input: `{}`}, + })) + require.ErrorContains(t, rejectDuplicateToolUseIDs([]fantasy.ToolCallContent{ + {ToolCallID: "duplicate", ToolName: "read_file", Input: `{}`}, + {ToolCallID: "duplicate", ToolName: "execute", Input: `{}`}, + }), "duplicate tool use ID") +} + +func TestNarrowHookAllowedToolsResponses(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + current pqtype.NullRawMessage + responses []agenthooks.Response + want string + wantValid bool + wantUsed bool + }{ + { + name: "None", + }, + { + name: "Empty", + responses: []agenthooks.Response{ + {AllowedTools: &[]string{}}, + }, + want: `[]`, + wantValid: true, + wantUsed: true, + }, + { + name: "Intersection", + current: pqtype.NullRawMessage{RawMessage: []byte(`["read_file","spawn_agent","execute"]`), Valid: true}, + responses: []agenthooks.Response{ + {AllowedTools: &[]string{"read_file", "spawn_agent"}}, + {}, + {AllowedTools: &[]string{"read_file", "execute"}}, + }, + want: `["read_file"]`, + wantValid: true, + wantUsed: true, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + got, used, err := narrowHookAllowedToolsResponses(tt.current, tt.responses) + require.NoError(t, err) + require.Equal(t, tt.wantUsed, used) + require.Equal(t, tt.wantValid, got.Valid) + if tt.wantValid { + require.JSONEq(t, tt.want, string(got.RawMessage)) + } + }) + } +} + +func TestPendingHookAllowedToolsContextChatScope(t *testing.T) { + t.Parallel() + + chatID := uuid.New() + policy := pqtype.NullRawMessage{RawMessage: []byte(`[]`), Valid: true} + ctx := withPendingHookAllowedTools(t.Context(), chatID, policy) + + got, ok := pendingHookAllowedToolsFromContext(ctx, chatID) + require.True(t, ok) + require.Equal(t, policy, got) + _, ok = pendingHookAllowedToolsFromContext(ctx, uuid.New()) + require.False(t, ok) +} + +func TestRestoreToolCallOrder(t *testing.T) { + t.Parallel() + + calls := []fantasy.ToolCallContent{ + {ToolCallID: "call_a", ToolName: "write_file"}, + {ToolCallID: "call_b", ToolName: "read_file"}, + {ToolCallID: "call_c", ToolName: "execute"}, + } + content := []fantasy.Content{ + fantasy.ToolResultContent{ToolCallID: "call_c", ToolName: "execute"}, + fantasy.ToolResultContent{ToolCallID: "call_b", ToolName: "read_file"}, + fantasy.ToolResultContent{ToolCallID: "call_a", ToolName: "write_file"}, + } + restoreToolCallOrder(content, calls) + gotIDs := make([]string, 0, len(content)) + for _, entry := range content { + result, ok := entry.(fantasy.ToolResultContent) + require.True(t, ok) + gotIDs = append(gotIDs, result.ToolCallID) + } + require.Equal(t, []string{"call_a", "call_b", "call_c"}, gotIDs) + + mixed := []fantasy.Content{ + fantasy.ToolResultContent{ToolCallID: "call_b", ToolName: "read_file"}, + fantasy.TextContent{Text: "note"}, + fantasy.ToolResultContent{ToolCallID: "unknown", ToolName: "other"}, + fantasy.ToolResultContent{ToolCallID: "call_a", ToolName: "write_file"}, + } + restoreToolCallOrder(mixed, calls) + first, ok := mixed[0].(fantasy.ToolResultContent) + require.True(t, ok) + require.Equal(t, "call_a", first.ToolCallID) + _, ok = mixed[1].(fantasy.TextContent) + require.True(t, ok) + unknown, ok := mixed[2].(fantasy.ToolResultContent) + require.True(t, ok) + require.Equal(t, "unknown", unknown.ToolCallID) + last, ok := mixed[3].(fantasy.ToolResultContent) + require.True(t, ok) + require.Equal(t, "call_b", last.ToolCallID) +} + +func TestFilterHookAllowedTools(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + allowed pqtype.NullRawMessage + wantTools []string + wantProviders []string + }{ + { + name: "null", + wantTools: []string{"read_file", "dynamic_tool"}, + wantProviders: []string{"web_search"}, + }, + { + name: "empty", + allowed: pqtype.NullRawMessage{RawMessage: []byte(`[]`), Valid: true}, + wantTools: []string{}, + wantProviders: []string{}, + }, + { + name: "subset", + allowed: pqtype.NullRawMessage{RawMessage: []byte(`["dynamic_tool","web_search"]`), Valid: true}, + wantTools: []string{"dynamic_tool"}, + wantProviders: []string{"web_search"}, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + tools := []fantasy.AgentTool{newTestAgentTool("read_file"), newTestAgentTool("dynamic_tool")} + providerTools := []chatloop.ProviderTool{{ + Definition: fantasy.ProviderDefinedTool{ID: "web_search", Name: "web_search"}, + }} + tools, providerTools, _, err := filterHookAllowedTools(test.allowed, tools, providerTools) + require.NoError(t, err) + toolNames := make([]string, 0, len(tools)) + for _, tool := range tools { + toolNames = append(toolNames, tool.Info().Name) + } + providerNames := make([]string, 0, len(providerTools)) + for _, tool := range providerTools { + providerNames = append(providerNames, tool.Definition.GetName()) + } + require.Equal(t, test.wantTools, toolNames) + require.Equal(t, test.wantProviders, providerNames) + }) + } +} diff --git a/coderd/x/chatd/hooks_test.go b/coderd/x/chatd/hooks_test.go new file mode 100644 index 0000000000000..274c24869f91d --- /dev/null +++ b/coderd/x/chatd/hooks_test.go @@ -0,0 +1,978 @@ +package chatd_test + +import ( + "context" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "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" + dbpubsub "github.com/coder/coder/v2/coderd/database/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/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/coderd/x/chathooks" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" +) + +func TestSendMessageUserPromptSubmitHook(t *testing.T) { + t.Parallel() + + t.Run("override", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + + submitted := []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("before"), + codersdk.ChatMessageFileReference("main.go", 1, 3, "package main"), + } + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.UserPromptSubmitData) + require.True(t, ok) + require.Equal(t, "before", data.Prompt) + var hookParts []codersdk.ChatMessagePart + require.NoError(t, json.Unmarshal(data.Parts, &hookParts)) + require.Equal(t, submitted, hookParts, "hook payload must carry non-text parts") + require.NotNil(t, request.Meta.TurnID) + _, err = w.Write([]byte(`{"permission":{"decision":"allow","input_override":{"prompt":"after"}},"model_context":"model only","user_message":"user only","allowed_tools":["read","write"]}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newHookTestServer(t, db, ps, consumer) + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: submitted, + }) + require.NoError(t, err) + require.True(t, result.Message.TurnID.Valid) + parts, err := chatprompt.ParseContent(result.Message) + require.NoError(t, err) + require.Equal(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("after")}, parts) + // Prefix messages let clients seed their transcript cache. + require.Len(t, result.InsertedMessages, 3) + require.Equal(t, database.ChatMessageRoleUser, result.InsertedMessages[0].Role) + require.Equal(t, database.ChatMessageVisibilityModel, result.InsertedMessages[0].Visibility) + require.Equal(t, database.ChatMessageRoleSystem, result.InsertedMessages[1].Role) + require.Equal(t, database.ChatMessageVisibilityUser, result.InsertedMessages[1].Visibility) + require.Equal(t, result.Message.ID, result.InsertedMessages[2].ID) + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, updated.HookAllowedTools.Valid) + require.JSONEq(t, `["read","write"]`, string(updated.HookAllowedTools.RawMessage)) + }) + + t.Run("deny", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write([]byte(`{"permission":{"decision":"deny"},"user_message":"blocked"}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newHookTestServer(t, db, ps, consumer) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("blocked prompt")}, + }) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + require.Equal(t, "blocked", denied.UserMessage) + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Empty(t, messages) + }) +} + +func TestUserPromptSubmitDenyEndChat(t *testing.T) { + t.Parallel() + + consumerResponse := `{"permission":{"decision":"deny"},"user_message":"blocked","end_chat":true}` + + t.Run("send", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + server := newHookTestServer(t, db, ps, hookConsumer(t, consumerResponse)) + + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("blocked prompt")}, + }) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + require.Equal(t, "blocked", denied.UserMessage) + + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, updated.Archived, "denied end_chat must archive the chat") + require.Equal(t, database.ChatStatusWaiting, updated.Status) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Empty(t, messages, "denied prompt must not persist") + }) + + t.Run("edit", func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) + require.NoError(t, err) + inserted, err := db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( + chat.ID, database.ChatMessageRoleUser, content, database.ChatMessageVisibilityBoth, model.ID, chatprompt.CurrentContentVersion, user.ID, + )) + require.NoError(t, err) + require.Len(t, inserted, 1) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + response := `{}` + if request.Type == agenthooks.EventUserPromptSubmit { + response = consumerResponse + } + _, err := w.Write([]byte(response)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited")}, + }) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, updated.Archived, "denied end_chat must archive the chat") + lastUser, err := db.GetLastChatMessageByRole(ctx, database.GetLastChatMessageByRoleParams{ + ChatID: chat.ID, + Role: database.ChatMessageRoleUser, + }) + require.NoError(t, err) + require.Equal(t, "original", hookMessageText(t, lastUser), + "denied edit must keep the original message") + }) +} + +func TestHookAllowedToolsPolicyNotice(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + hookResponse string + wantTools bool + wantNotice string + }{ + { + name: "no policy", + hookResponse: `{}`, + wantTools: true, + }, + { + name: "subset", + hookResponse: `{"allowed_tools":["read_file"]}`, + wantTools: true, + wantNotice: "restricted by an external policy", + }, + { + name: "all tools removed", + hookResponse: `{"allowed_tools":["no_such_tool"]}`, + wantNotice: "removed all tool access", + }, + } + 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) + var ( + requestMu sync.Mutex + requestTools []string + requestMsgs []chattest.OpenAIMessage + ) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + requestMu.Lock() + requestTools = requestTools[:0] + for _, tool := range req.Tools { + requestTools = append(requestTools, tool.Function.Name) + } + requestMsgs = append([]chattest.OpenAIMessage(nil), req.Messages...) + requestMu.Unlock() + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("ok")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + + consumer := hookConsumer(t, tt.hookResponse) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "allowed tools notice", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + requestMu.Lock() + defer requestMu.Unlock() + if tt.wantTools { + require.NotEmpty(t, requestTools) + } else { + require.Empty(t, requestTools) + } + var notice string + for _, msg := range requestMsgs { + if msg.Role == "system" && strings.Contains(msg.Content, "external policy") { + notice = msg.Content + } + } + if tt.wantNotice == "" { + require.Empty(t, notice, "chat without a tool policy must not receive the policy notice") + } else { + require.Contains(t, notice, tt.wantNotice) + } + }) + } +} + +func newHookDispatcher(t *testing.T, db database.Store, consumer *httptest.Server) *chathooks.Dispatcher { + t.Helper() + return chathooks.New( + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + db, + consumer.Client(), + consumer.URL, + "test-hook-secret-32-bytes-minimum!!", + time.Second, + "test-deployment", + "test-version", + prometheus.NewRegistry(), + ) +} + +func newHookTestServer(t *testing.T, db database.Store, ps dbpubsub.Pubsub, consumer *httptest.Server) *chatd.Server { + t.Helper() + return newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) +} + +func TestHookDispatcherRequiresExperiment(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + + var hookRequests atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hookRequests.Add(1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(consumer.Close) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.Experiments = slices.DeleteFunc( + slices.Clone(codersdk.ExperimentsKnown), + func(e codersdk.Experiment) bool { return e == codersdk.ExperimentAgentLifecycleHooks }, + ) + }) + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, + }) + require.NoError(t, err) + parts, err := chatprompt.ParseContent(result.Message) + require.NoError(t, err) + require.Equal(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")}, parts) + + require.Zero(t, hookRequests.Load()) + rows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + require.Empty(t, rows) +} + +func lifecycleDispatch(t *testing.T, db database.Store, chatID uuid.UUID, event agenthooks.EventType) database.ChatHookDispatch { + t.Helper() + rows, err := db.ListChatHookDispatchesByChatID(t.Context(), chatID) + require.NoError(t, err) + for _, row := range rows { + if row.Event == string(event) { + return row + } + } + require.FailNow(t, "lifecycle dispatch not found", "event: %s", event) + return database.ChatHookDispatch{} +} + +func TestSendMessageUserPromptSubmitPassthrough(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + server := newHookTestServer(t, db, ps, hookConsumer(t, `{}`)) + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("passthrough")}, + }) + require.NoError(t, err) + require.Equal(t, "passthrough", hookMessageText(t, result.Message)) +} + +func TestSendMessageUserPromptSubmitEndChat(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + consumer := hookConsumer(t, `{"end_chat":true}`) + server := newHookTestServer(t, db, ps, consumer) + + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do not persist")}, + }) + require.NoError(t, err) + require.True(t, result.Ended) + require.True(t, result.Chat.Archived) + require.Equal(t, database.ChatStatusWaiting, result.Chat.Status) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Empty(t, messages) +} + +func TestSendMessageUserPromptSubmitQueue(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat, err := newTestServer(t, db, ps, uuid.New()).CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "queued hook", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("running")}, + }) + require.NoError(t, err) + consumer := hookConsumer(t, `{"permission":{"decision":"allow","input_override":{"prompt":"queued override"}},"model_context":"queued context","allowed_tools":["read_file"]}`) + server := newHookTestServer(t, db, ps, consumer) + + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued original")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + require.True(t, result.Queued) + require.NotNil(t, result.QueuedMessage) + require.True(t, result.QueuedMessage.TurnID.Valid) + queuedParts, err := chatprompt.ParseContent(database.ChatMessage{ + Role: database.ChatMessageRoleUser, + Content: pqtype.NullRawMessage{RawMessage: result.QueuedMessage.Content, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + }) + require.NoError(t, err) + require.Equal(t, []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued override")}, queuedParts) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + for i := range messages { + require.NotEqual(t, messages[i].TurnID, result.QueuedMessage.TurnID, + "queued prompt's turn must have no history rows before promotion") + } + require.True(t, result.QueuedMessage.HookPrefix.Valid) + require.Contains(t, string(result.QueuedMessage.HookPrefix.RawMessage), "queued context") + require.True(t, result.QueuedMessage.HookAllowedTools.Valid) + require.JSONEq(t, `["read_file"]`, string(result.QueuedMessage.HookAllowedTools.RawMessage)) + queuedChat, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.False(t, queuedChat.HookAllowedTools.Valid, + "queued prompt's tool policy must not apply before promotion") +} + +func TestSendMessageUserPromptSubmitQueuedRejections(t *testing.T) { + t.Parallel() + tests := []struct { + name string + statusCode int + response string + assertErr func(*testing.T, error) + }{ + { + name: "deny", + statusCode: http.StatusOK, + response: `{"permission":{"decision":"deny"},"user_message":"blocked"}`, + assertErr: func(t *testing.T, err error) { + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + }, + }, + { + name: "dispatch failure", + statusCode: http.StatusInternalServerError, + assertErr: func(t *testing.T, err error) { + var dispatchErr *chathooks.DispatchError + require.ErrorAs(t, err, &dispatchErr) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat, err := newTestServer(t, db, ps, uuid.New()).CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "queued rejection", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("running")}, + }) + require.NoError(t, err) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(test.statusCode) + if test.response != "" { + _, err := w.Write([]byte(test.response)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + test.assertErr(t, err) + queued, err := db.GetChatQueuedMessages(ctx, chat.ID) + require.NoError(t, err) + require.Empty(t, queued) + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRunning, updated.Status) + require.False(t, updated.LastError.Valid) + }) + } +} + +func TestSendMessageUserPromptSubmitDispatchFailure(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("fails")}, + }) + var dispatchErr *chathooks.DispatchError + require.ErrorAs(t, err, &dispatchErr) + require.Equal(t, chathooks.ResultHTTPError, dispatchErr.Class) + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusError, updated.Status) + var chatErr codersdk.ChatError + require.NoError(t, json.Unmarshal(updated.LastError.RawMessage, &chatErr)) + require.Equal(t, "hook dispatch failed: user_prompt_submit: http_error (dispatch "+dispatchErr.DispatchID.String()+")", chatErr.Message) +} + +func TestEditMessageUserPromptSubmitHook(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) + require.NoError(t, err) + inserted, err := db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( + chat.ID, database.ChatMessageRoleUser, content, database.ChatMessageVisibilityBoth, model.ID, chatprompt.CurrentContentVersion, user.ID, + )) + require.NoError(t, err) + require.Len(t, inserted, 1) + type receivedHook struct { + request agenthooks.Request + claims agenthooks.Claims + } + received := make([]receivedHook, 0, 2) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + claims, err := agenthooks.Verify(r.Header.Get("Authorization"), []byte("test-hook-secret-32-bytes-minimum!!")) + require.NoError(t, err) + received = append(received, receivedHook{request: request, claims: claims}) + response := `{"model_context":"clear context","user_message":"clear notice","allowed_tools":["read_file"]}` + if request.Type == agenthooks.EventUserPromptSubmit { + response = `{"permission":{"decision":"allow","input_override":{"prompt":"edited override"}}}` + } + _, err = w.Write([]byte(response)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + result, err := server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited original")}, + }) + require.NoError(t, err) + require.True(t, result.Message.TurnID.Valid) + require.NotEqual(t, inserted[0].TurnID, result.Message.TurnID) + require.Equal(t, "edited override", hookMessageText(t, result.Message)) + require.Len(t, received, 2) + require.Equal(t, agenthooks.EventSessionStart, received[0].request.Type) + data, err := received[0].request.Decode() + require.NoError(t, err) + require.Equal(t, &agenthooks.SessionStartData{Source: "clear"}, data) + require.Equal(t, received[0].request.Meta.DispatchID, received[0].claims.JTI) + require.Equal(t, agenthooks.EventUserPromptSubmit, received[1].request.Type) + require.Equal(t, received[1].request.Meta.DispatchID, received[1].claims.JTI) + require.NotEqual(t, received[0].claims.JTI, received[1].claims.JTI) + rows, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + var foundNotice bool + for _, row := range rows { + if row.Role == database.ChatMessageRoleSystem && row.Visibility == database.ChatMessageVisibilityUser && hookMessageText(t, row) == "clear notice" { + foundNotice = true + } + } + require.True(t, foundNotice) + promptRows, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + var foundContext bool + for _, row := range promptRows { + if row.Visibility == database.ChatMessageVisibilityModel && hookMessageText(t, row) == "clear context" { + foundContext = true + } + } + require.True(t, foundContext) + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(updated.HookAllowedTools.RawMessage)) +} + +func TestEditMessageSessionStartEndChatSkipsPromptDispatch(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) + require.NoError(t, err) + inserted, err := db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( + chat.ID, database.ChatMessageRoleUser, content, database.ChatMessageVisibilityBoth, model.ID, chatprompt.CurrentContentVersion, user.ID, + )) + require.NoError(t, err) + require.Len(t, inserted, 1) + var promptDispatches atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type == agenthooks.EventUserPromptSubmit { + promptDispatches.Add(1) + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + _, err := w.Write([]byte(`{"end_chat":true,"user_message":"ended by session policy","model_context":"session context note"}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + result, err := server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited original")}, + }) + require.NoError(t, err) + require.True(t, result.Ended, "session end_chat must return an ended edit result") + require.True(t, result.Chat.Archived) + require.Zero(t, promptDispatches.Load(), "edited prompt must not be dispatched after session end_chat") + + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, updated.Archived, "accepted session end_chat must archive the chat") + require.Equal(t, database.ChatStatusWaiting, updated.Status) + require.False(t, updated.LastError.Valid) + original, err := db.GetChatMessageByID(ctx, inserted[0].ID) + require.NoError(t, err) + require.Equal(t, "original", hookMessageText(t, original)) + require.False(t, original.Deleted) + + assertSessionEndMessagesPersisted(ctx, t, db, chat.ID) +} + +func TestEditMessagePromptDenialEndChatPersistsSessionMessages(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) + require.NoError(t, err) + inserted, err := db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( + chat.ID, database.ChatMessageRoleUser, content, database.ChatMessageVisibilityBoth, model.ID, chatprompt.CurrentContentVersion, user.ID, + )) + require.NoError(t, err) + require.Len(t, inserted, 1) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type == agenthooks.EventUserPromptSubmit { + _, err := w.Write([]byte(`{"permission":{"decision":"deny"},"end_chat":true}`)) + require.NoError(t, err) + return + } + _, err := w.Write([]byte(`{"user_message":"ended by session policy","model_context":"session context note"}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edited original")}, + }) + var denied *chatd.UserPromptDeniedError + require.ErrorAs(t, err, &denied) + + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, updated.Archived, "accepted denial end_chat must archive the chat") + require.Equal(t, database.ChatStatusWaiting, updated.Status) + require.False(t, updated.LastError.Valid) + original, err := db.GetChatMessageByID(ctx, inserted[0].ID) + require.NoError(t, err) + require.Equal(t, "original", hookMessageText(t, original)) + require.False(t, original.Deleted) + + assertSessionEndMessagesPersisted(ctx, t, db, chat.ID) +} + +func assertSessionEndMessagesPersisted(ctx context.Context, t *testing.T, db database.Store, chatID uuid.UUID) { + t.Helper() + rows, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chatID}) + require.NoError(t, err) + noticePersisted := false + for _, row := range rows { + if row.Role == database.ChatMessageRoleSystem && row.Visibility == database.ChatMessageVisibilityUser && hookMessageText(t, row) == "ended by session policy" { + noticePersisted = true + } + } + require.True(t, noticePersisted, "accepted session user_message must persist with the archive") + promptRows, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + contextPersisted := false + for _, row := range promptRows { + if row.Visibility == database.ChatMessageVisibilityModel && hookMessageText(t, row) == "session context note" { + contextPersisted = true + } + } + require.True(t, contextPersisted, "accepted session model_context must persist with the archive") +} + +func TestEditMessageInvalidTargetSkipsHooks(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + var dispatched atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + dispatched.Add(1) + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + _, err := server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: 999999, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit of nothing")}, + }) + require.ErrorIs(t, err, chatd.ErrEditedMessageNotFound) + + // dbgen.Chat ignores seed.Archived; archive explicitly. + archived := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + _, err = db.ArchiveChatByID(ctx, archived.ID) + require.NoError(t, err) + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: archived.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("send to archived")}, + }) + require.ErrorIs(t, err, chatd.ErrChatArchived) + _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: archived.ID, + CreatedBy: user.ID, + EditedMessageID: 1, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("edit archived")}, + }) + require.ErrorIs(t, err, chatd.ErrChatArchived) + + require.Zero(t, dispatched.Load(), "invalid targets must not dispatch hooks") + for _, chatID := range []uuid.UUID{chat.ID, archived.ID} { + dispatches, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + require.Empty(t, dispatches) + } +} + +func TestPromptHooksAdmissionPreflight(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + var dispatched atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + dispatched.Add(1) + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + server := newHookTestServer(t, db, ps, consumer) + + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + _, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("bad model")}, + ModelConfigID: uuid.New(), + }) + require.ErrorIs(t, err, chatd.ErrInvalidModelConfigID) + + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("original")}) + require.NoError(t, err) + inserted, err := db.InsertChatMessages(ctx, chatd.BuildSingleChatMessageInsertParams( + chat.ID, database.ChatMessageRoleUser, content, database.ChatMessageVisibilityBoth, model.ID, chatprompt.CurrentContentVersion, user.ID, + )) + require.NoError(t, err) + require.Len(t, inserted, 1) + _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("bad model edit")}, + ModelConfigID: uuid.New(), + }) + require.ErrorIs(t, err, chatd.ErrInvalidModelConfigID) + + busy := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + Status: database.ChatStatusRunning, + }) + queuedContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("queued")}) + require.NoError(t, err) + for range chatstate.MaxQueueSize { + _, err = db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: busy.ID, + Content: queuedContent.RawMessage, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + CreatedBy: user.ID, + }) + require.NoError(t, err) + } + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: busy.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("queue full")}, + }) + require.ErrorIs(t, err, chatstate.ErrMessageQueueFull) + + // Usage limits are deployment-wide, so these cases run last. + _, err = db.UpsertChatUsageLimitConfig(ctx, database.UpsertChatUsageLimitConfigParams{ + Enabled: true, + DefaultLimitMicros: 100, + Period: string(codersdk.ChatUsageLimitPeriodDay), + }) + require.NoError(t, err) + 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}, + }) + var limitErr *chatd.UsageLimitExceededError + _, err = server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("over limit")}, + }) + require.ErrorAs(t, err, &limitErr) + _, err = server.EditMessage(ctx, chatd.EditMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + EditedMessageID: inserted[0].ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("over limit edit")}, + }) + require.ErrorAs(t, err, &limitErr) + + require.Zero(t, dispatched.Load(), "admission-rejected prompts must not dispatch hooks") + for _, chatID := range []uuid.UUID{chat.ID, busy.ID} { + dispatches, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + require.Empty(t, dispatches) + } +} + +func TestSendMessageHooksDisabled(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, model := seedChatDependencies(t, db) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + server := newTestServer(t, db, ps, uuid.New()) + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("unchanged")}, + }) + require.NoError(t, err) + require.True(t, result.Message.TurnID.Valid) + require.Equal(t, "unchanged", hookMessageText(t, result.Message)) +} + +func hookConsumer(t *testing.T, response string) *httptest.Server { + t.Helper() + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write([]byte(response)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + return consumer +} + +func hookMessageText(t *testing.T, message database.ChatMessage) string { + t.Helper() + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + require.Len(t, parts, 1) + return parts[0].Text +} diff --git a/coderd/x/chatd/message_conversion.go b/coderd/x/chatd/message_conversion.go index c70f1efcd8cf8..5de5d5c65b24b 100644 --- a/coderd/x/chatd/message_conversion.go +++ b/coderd/x/chatd/message_conversion.go @@ -336,19 +336,23 @@ func buildCompactionMessages(input buildCompactionMessagesInput) (compactionMess return compactionMessagesForCommit{Messages: messages, HiddenCount: 1}, nil } -func currentTurnStepCount(messages []database.ChatMessage) int { - latestUser := -1 +// Hook model-context messages use the user role but must not reset per-turn guards. +func currentTurnStartIndex(messages []database.ChatMessage) int { + start := 0 for i, msg := range messages { if msg.Deleted || msg.Compressed { continue } - if msg.Role == database.ChatMessageRoleUser { - latestUser = i + if msg.Role == database.ChatMessageRoleUser && msg.Visibility != database.ChatMessageVisibilityModel { + start = i + 1 } } + return start +} + +func currentTurnStepCount(messages []database.ChatMessage) int { count := 0 - for i := latestUser + 1; i < len(messages); i++ { - msg := messages[i] + for _, msg := range messages[currentTurnStartIndex(messages):] { if msg.Deleted || msg.Compressed { continue } @@ -477,16 +481,7 @@ func historyHasStopAfterToolResult(messages []database.ChatMessage, stopAfterToo 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:] { + for _, msg := range messages[currentTurnStartIndex(messages):] { if msg.Deleted || msg.Compressed || msg.Role != database.ChatMessageRoleTool { continue } diff --git a/coderd/x/chatd/message_conversion_test.go b/coderd/x/chatd/message_conversion_test.go index 4a14d3cd34ae5..40eb2fa291d22 100644 --- a/coderd/x/chatd/message_conversion_test.go +++ b/coderd/x/chatd/message_conversion_test.go @@ -271,6 +271,22 @@ func TestCurrentTurnStepCount_CountsAssistantMessagesAfterLatestUser(t *testing. require.Equal(t, 2, got) } +func TestCurrentTurnStepCount_IgnoresHookModelContext(t *testing.T) { + t.Parallel() + + hookContext := dbMessage(t, 4, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hook context")) + hookContext.Visibility = database.ChatMessageVisibilityModel + messages := []database.ChatMessage{ + dbMessage(t, 1, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("prompt")), + dbMessage(t, 2, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("one")), + dbMessage(t, 3, database.ChatMessageRoleTool, false, codersdk.ChatMessageToolResult("call", "tool", json.RawMessage(`{}`), false, false)), + hookContext, + dbMessage(t, 5, database.ChatMessageRoleAssistant, false, codersdk.ChatMessageText("two")), + } + got := currentTurnStepCount(messages) + require.Equal(t, 2, got) +} + func TestDecisionCompactsAgainAfterPostCompactionTurn(t *testing.T) { t.Parallel() @@ -545,6 +561,22 @@ func TestDecisionDetectsStopAfterToolFromCommittedHistory(t *testing.T) { require.False(t, got) } +func TestDecisionDetectsStopAfterToolAcrossHookContext(t *testing.T) { + t.Parallel() + + hookContext := dbMessage(t, 4, database.ChatMessageRoleUser, false, codersdk.ChatMessageText("hook context")) + hookContext.Visibility = database.ChatMessageVisibilityModel + 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)), + hookContext, + } + got, err := historyHasStopAfterToolResult(messages, map[string]struct{}{"propose_plan": {}}) + require.NoError(t, err) + require.True(t, got) +} + func TestDecisionDetectsCurrentHistoryCompletion(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/options.go b/coderd/x/chatd/options.go index ff3dbdd3d9a30..356653aad1fce 100644 --- a/coderd/x/chatd/options.go +++ b/coderd/x/chatd/options.go @@ -3,6 +3,7 @@ package chatd import ( "context" "database/sql" + "sync" "sync/atomic" "time" @@ -62,6 +63,102 @@ type chatWorkerTaskStartInput struct { Status database.ChatStatus RequiresActionDeadlineAt sql.NullTime DebugTurn *runnerDebugTurn + SessionStart *sessionStartTracker + StopNudges *stopNudgeTracker +} + +type stopNudgeTracker struct { + mu sync.Mutex + turnID uuid.UUID + claimed bool + pending bool +} + +func stopNudgeTurnID(turnID *uuid.UUID) uuid.UUID { + if turnID == nil { + return uuid.Nil + } + return *turnID +} + +func (t *stopNudgeTracker) claim(turnID *uuid.UUID) bool { + currentTurnID := stopNudgeTurnID(turnID) + t.mu.Lock() + defer t.mu.Unlock() + if t.turnID != currentTurnID { + t.turnID = currentTurnID + t.claimed = false + } + if t.claimed { + return false + } + t.claimed = true + t.pending = true + return true +} + +func (t *stopNudgeTracker) consume(turnID *uuid.UUID) bool { + currentTurnID := stopNudgeTurnID(turnID) + t.mu.Lock() + defer t.mu.Unlock() + if t.turnID != currentTurnID || !t.pending { + return false + } + t.pending = false + return true +} + +func (t *stopNudgeTracker) cancel(turnID *uuid.UUID) { + currentTurnID := stopNudgeTurnID(turnID) + t.mu.Lock() + defer t.mu.Unlock() + if t.turnID != currentTurnID || !t.pending { + return + } + t.pending = false + t.claimed = false +} + +func (t *stopNudgeTracker) reset() { + t.mu.Lock() + t.turnID = uuid.Nil + t.claimed = false + t.pending = false + t.mu.Unlock() +} + +type sessionStartTracker struct { + mu sync.Mutex + completed bool + inFlight chan struct{} +} + +func (t *sessionStartTracker) claim(ctx context.Context) (bool, func(bool), error) { + for { + t.mu.Lock() + if t.completed { + t.mu.Unlock() + return false, nil, nil + } + if t.inFlight == nil { + t.inFlight = make(chan struct{}) + t.mu.Unlock() + return true, func(completed bool) { + t.mu.Lock() + t.completed = completed + close(t.inFlight) + t.inFlight = nil + t.mu.Unlock() + }, nil + } + inFlight := t.inFlight + t.mu.Unlock() + select { + case <-inFlight: + case <-ctx.Done(): + return false, nil, ctx.Err() + } + } } // chatWorkerOptions configures a chatWorker. diff --git a/coderd/x/chatd/post_tool_use_test.go b/coderd/x/chatd/post_tool_use_test.go new file mode 100644 index 0000000000000..205fce3e94768 --- /dev/null +++ b/coderd/x/chatd/post_tool_use_test.go @@ -0,0 +1,694 @@ +package chatd_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "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/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/coderd/x/chathooks" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" + "github.com/coder/coder/v2/testutil" +) + +func TestPostToolUseHookResponsesCommitWithResults(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + first := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/first.txt"}`) + first.Choices[0].ToolCalls[0].ID = "call_first" + second := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/second.txt"}`).Choices[0].ToolCalls[0] + second.ID = "call_second" + second.Index = 1 + first.Choices[0].ToolCalls = append(first.Choices[0].ToolCalls, second) + return chattest.OpenAIStreamingResponse(first) + } + toolResultIndex := -1 + contextIndex := -1 + for i, message := range req.Messages { + if message.Role == "tool" && strings.Contains(message.Content, "data") { + toolResultIndex = i + } + if strings.Contains(message.Content, "lint feedback") { + contextIndex = i + } + } + require.NotEqual(t, -1, toolResultIndex) + require.NotEqual(t, -1, contextIndex) + require.Less(t, toolResultIndex, contextIndex) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + var mu sync.Mutex + var received []agenthooks.PostToolUseData + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPostToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PostToolUseData) + mu.Lock() + received = append(received, *data) + index := len(received) + mu.Unlock() + + messages := chatMessages(ctx, t, db, request.Meta.ChatID) + for _, message := range messages { + require.NotEqual(t, database.ChatMessageRoleTool, message.Role) + } + if index == 1 { + _, err = w.Write([]byte(`{"model_context":"lint feedback","user_message":"tool notice","allowed_tools":["read_file"]}`)) + } else { + _, err = w.Write([]byte(`{}`)) + } + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(2) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "post-tool-use-responses", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read both files"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + mu.Lock() + receivedSnapshot := append([]agenthooks.PostToolUseData(nil), received...) + mu.Unlock() + require.Len(t, receivedSnapshot, 2) + require.Equal(t, "call_first", receivedSnapshot[0].ToolUseID) + require.Equal(t, "call_second", receivedSnapshot[1].ToolUseID) + require.Equal(t, "read_file", receivedSnapshot[0].ToolName) + require.Empty(t, receivedSnapshot[0].ToolError) + require.Contains(t, string(receivedSnapshot[0].ToolResponse), "data") + + var toolResults, userMessages int + for _, message := range chatMessages(ctx, t, db, chat.ID) { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if message.Role == database.ChatMessageRoleTool { + toolResults++ + } + if len(parts) == 1 && parts[0].Text == "tool notice" { + userMessages++ + require.Equal(t, database.ChatMessageVisibilityUser, message.Visibility) + } + } + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + var modelContexts int + for _, message := range promptMessages { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "lint feedback" { + modelContexts++ + require.Equal(t, database.ChatMessageVisibilityModel, message.Visibility) + } + } + require.Equal(t, 2, toolResults) + require.Equal(t, 1, modelContexts) + require.Equal(t, 1, userMessages) + updated, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(updated.HookAllowedTools.RawMessage)) +} + +func TestPostToolUseHookDynamicResult(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"query":"value"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_dynamic_result" + return chattest.OpenAIStreamingResponse(chunk) + } + resultIndex := -1 + contextIndex := -1 + for i, message := range req.Messages { + if message.Role == "tool" && strings.Contains(message.Content, "answer") { + resultIndex = i + } + if strings.Contains(message.Content, "dynamic feedback") { + contextIndex = i + } + } + require.NotEqual(t, -1, resultIndex) + require.NotEqual(t, -1, contextIndex) + require.Less(t, resultIndex, contextIndex) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + var postCalls atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPostToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + postCalls.Add(1) + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PostToolUseData) + require.Equal(t, "call_dynamic_result", data.ToolUseID) + require.Equal(t, "my_dynamic_tool", data.ToolName) + require.JSONEq(t, `{"answer":42}`, string(data.ToolResponse)) + _, err = w.Write([]byte(`{"model_context":"dynamic feedback"}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "post-tool-use-dynamic", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(context.Context) bool { + updated, err := db.GetChatByID(ctx, chat.ID) + return err == nil && updated.Status == database.ChatStatusRequiresAction + }, testutil.IntervalFast) + + err = server.SubmitToolResults(ctx, chatd.SubmitToolResultsOptions{ + ChatID: chat.ID, + UserID: user.ID, + ModelConfigID: model.ID, + Results: []codersdk.ToolResult{{ + ToolCallID: "call_dynamic_result", + Output: json.RawMessage(`{"answer":42}`), + }}, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(1), postCalls.Load()) +} + +func TestPostToolUseHookDynamicFailureRejectsSubmission(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{}`) + chunk.Choices[0].ToolCalls[0].ID = "call_dynamic_failure" + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + var failPostToolUse atomic.Bool + failPostToolUse.Store(true) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type == agenthooks.EventPostToolUse && failPostToolUse.Load() { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "post-tool-use-dynamic-failure", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(context.Context) bool { + updated, err := db.GetChatByID(ctx, chat.ID) + return err == nil && updated.Status == database.ChatStatusRequiresAction + }, testutil.IntervalFast) + + results := []codersdk.ToolResult{{ + ToolCallID: "call_dynamic_failure", + Output: json.RawMessage(`{"answer":42}`), + }} + err = server.SubmitToolResults(ctx, chatd.SubmitToolResultsOptions{ + ChatID: chat.ID, + UserID: user.ID, + ModelConfigID: model.ID, + Results: results, + }) + var dispatchErr *chathooks.DispatchError + require.ErrorAs(t, err, &dispatchErr) + + unchanged, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, database.ChatStatusRequiresAction, unchanged.Status) + require.False(t, unchanged.LastError.Valid) + for _, part := range chatToolParts(ctx, t, db, chat.ID) { + require.NotEqual(t, codersdk.ChatMessagePartTypeToolResult, part.Type, + "rejected submission must not commit tool results") + } + dispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventPostToolUse) + require.Equal(t, "http_error", dispatch.Result) + + failPostToolUse.Store(false) + require.NoError(t, server.SubmitToolResults(ctx, chatd.SubmitToolResultsOptions{ + ChatID: chat.ID, + UserID: user.ID, + ModelConfigID: model.ID, + Results: results, + })) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), "my_dynamic_tool") + require.JSONEq(t, `{"answer":42}`, string(result.Result)) +} + +func TestPostToolUseHookDynamicFailurePreservesAcceptedEndChat(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{}`) + chunk.Choices[0].ToolCalls[0].ID = "call_end_first" + second := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{}`) + second.Choices[0].ToolCalls[0].ID = "call_fail_second" + second.Choices[0].ToolCalls[0].Index = 1 + chunk.Choices[0].ToolCalls = append(chunk.Choices[0].ToolCalls, second.Choices[0].ToolCalls[0]) + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPostToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.PostToolUseData) + require.True(t, ok) + if data.ToolUseID == "call_fail_second" { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err = w.Write([]byte(`{"end_chat":true}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "post-tool-use-end-chat-precedence", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool twice"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(context.Context) bool { + updated, err := db.GetChatByID(ctx, chat.ID) + return err == nil && updated.Status == database.ChatStatusRequiresAction + }, testutil.IntervalFast) + + require.NoError(t, server.SubmitToolResults(ctx, chatd.SubmitToolResultsOptions{ + ChatID: chat.ID, + UserID: user.ID, + ModelConfigID: model.ID, + Results: []codersdk.ToolResult{ + {ToolCallID: "call_end_first", Output: json.RawMessage(`{"answer":42}`)}, + {ToolCallID: "call_fail_second", Output: json.RawMessage(`{"answer":43}`)}, + }, + })) + + ended, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.True(t, ended.Archived, "accepted end_chat must archive the chat") + require.Equal(t, database.ChatStatusWaiting, ended.Status) + require.False(t, ended.LastError.Valid) +} + +func TestPostToolUseHookFailureCommitsResultThenErrors(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunk := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/file.txt"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_failure" + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type == agenthooks.EventPostToolUse { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/file.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "post-tool-use-failure", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file"), + }, + }) + require.NoError(t, err) + failed := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), "read_file") + require.Contains(t, string(result.Result), "data") + postDispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventPostToolUse) + require.Equal(t, "http_error", postDispatch.Result) + require.Equal(t, "call_failure", postDispatch.ToolUseID.String) + lastError := chatLastErrorMessage(failed.LastError) + require.Contains(t, lastError, "hook dispatch failed: post_tool_use: http_error") + require.Contains(t, lastError, postDispatch.ID.String()) +} + +func TestPostToolUseHookFailureDispatchesRemainingResults(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + first := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/first.txt"}`) + first.Choices[0].ToolCalls[0].ID = "call_first" + second := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/second.txt"}`).Choices[0].ToolCalls[0] + second.ID = "call_second" + second.Index = 1 + first.Choices[0].ToolCalls = append(first.Choices[0].ToolCalls, second) + return chattest.OpenAIStreamingResponse(first) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPostToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PostToolUseData) + if data.ToolUseID == "call_first" { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err = w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(2) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "post-tool-use-continue", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read both files"), + }, + }) + require.NoError(t, err) + failed := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + + rows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + results := map[string]string{} + for _, row := range rows { + if row.Event == string(agenthooks.EventPostToolUse) { + results[row.ToolUseID.String] = row.Result + } + } + require.Equal(t, map[string]string{ + "call_first": "http_error", + "call_second": "ok", + }, results, "every executed result must have a dispatch row") + require.Contains(t, chatLastErrorMessage(failed.LastError), "hook dispatch failed: post_tool_use: http_error") +} + +func TestPostToolUseHookEndChatOrderingWithFailure(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + endChatID string + failID string + wantEnded bool + wantErrors bool + }{ + { + name: "failure before end_chat fails closed", + failID: "call_first", + endChatID: "call_second", + wantErrors: true, + }, + { + name: "failure after end_chat archives", + endChatID: "call_first", + failID: "call_second", + wantEnded: true, + }, + } + 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) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + first := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/first.txt"}`) + first.Choices[0].ToolCalls[0].ID = "call_first" + second := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/second.txt"}`).Choices[0].ToolCalls[0] + second.ID = "call_second" + second.Index = 1 + first.Choices[0].ToolCalls = append(first.Choices[0].ToolCalls, second) + return chattest.OpenAIStreamingResponse(first) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPostToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PostToolUseData) + switch data.ToolUseID { + case tt.failID: + w.WriteHeader(http.StatusInternalServerError) + case tt.endChatID: + _, err = w.Write([]byte(`{"end_chat":true}`)) + require.NoError(t, err) + default: + _, err = w.Write([]byte(`{}`)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(2) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "post-tool-use-ordering", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read both files"), + }, + }) + require.NoError(t, err) + + if tt.wantErrors { + failed := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + require.False(t, failed.Archived) + require.Contains(t, chatLastErrorMessage(failed.LastError), "hook dispatch failed: post_tool_use: http_error") + } + if tt.wantEnded { + var ended database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + ended, err = db.GetChatByID(ctx, chat.ID) + return err == nil && ended.Archived && ended.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.False(t, ended.LastError.Valid) + } + + rows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + results := map[string]string{} + for _, row := range rows { + if row.Event == string(agenthooks.EventPostToolUse) { + results[row.ToolUseID.String] = row.Result + } + } + require.Equal(t, map[string]string{ + tt.failID: "http_error", + tt.endChatID: "ok", + }, results, "both executed results must have dispatch rows") + }) + } +} diff --git a/coderd/x/chatd/pre_tool_use_test.go b/coderd/x/chatd/pre_tool_use_test.go new file mode 100644 index 0000000000000..9366f12d178ef --- /dev/null +++ b/coderd/x/chatd/pre_tool_use_test.go @@ -0,0 +1,1819 @@ +package chatd_test + +import ( + "context" + "database/sql" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "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/dbtestutil" + dbpubsub "github.com/coder/coder/v2/coderd/database/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/coderd/x/chatd/chattest" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/agenthooks" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock" + "github.com/coder/coder/v2/testutil" +) + +func TestPreToolUseHookAllow(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + response string + expectedPath string + decision string + }{ + { + name: "passthrough", + response: `{}`, + expectedPath: "/tmp/before.txt", + decision: "allow", + }, + { + name: "override", + response: `{"permission":{"decision":"allow","input_override":{"path":"/tmp/after.txt"}},"user_message":"tool approved","allowed_tools":["read_file"]}`, + expectedPath: "/tmp/after.txt", + decision: "allow", + }, + } + 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) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/before.txt"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_non_uuid" + return chattest.OpenAIStreamingResponse(chunk) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + hookCalls.Add(1) + require.Equal(t, "call_non_uuid", data.ToolUseID) + require.Equal(t, "read_file", data.ToolName) + require.JSONEq(t, `{"path":"/tmp/before.txt"}`, string(data.ToolInput)) + return tt.response + }) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), tt.expectedPath, int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(1) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "pre-tool-use-allow", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the file"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + require.Equal(t, int32(1), hookCalls.Load()) + call := requireToolCallPart(t, chatToolParts(ctx, t, db, chat.ID), "read_file") + require.JSONEq(t, `{"path":"`+tt.expectedPath+`"}`, string(call.Args)) + dispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventPreToolUse) + require.Equal(t, "call_non_uuid", dispatch.ToolUseID.String) + require.Equal(t, tt.decision != "", dispatch.Decision.Valid) + if tt.decision != "" { + require.Equal(t, tt.decision, dispatch.Decision.String) + } + if tt.name == "override" { + chatResult, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(chatResult.HookAllowedTools.RawMessage)) + var foundUserMessage bool + for _, message := range chatMessages(ctx, t, db, chat.ID) { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "tool approved" { + foundUserMessage = true + } + } + require.True(t, foundUserMessage) + } + require.JSONEq(t, `{"path":"/tmp/before.txt"}`, string(dispatch.OriginalInput.RawMessage)) + }) + } +} + +func TestPreToolUseHookDeny(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + var secondMessages []chattest.OpenAIMessage + var messagesMu sync.Mutex + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("read_file", `{"path":"/tmp/secret.txt"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_denied" + return chattest.OpenAIStreamingResponse(chunk) + } + messagesMu.Lock() + secondMessages = append([]chattest.OpenAIMessage(nil), req.Messages...) + messagesMu.Unlock() + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("used another approach")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + require.Equal(t, "call_denied", data.ToolUseID) + return `{"permission":{"decision":"deny","reason":"blocked by policy"},"model_context":"Do not read secrets."}` + }) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + 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, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + Title: "pre-tool-use-deny", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("read the secret"), + }, + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + parts := chatToolParts(ctx, t, db, chat.ID) + result := requireToolResultPart(t, parts, "read_file") + require.True(t, result.IsError) + require.Contains(t, string(result.Result), "DENIED: blocked by policy") + + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + var foundModelContext bool + for _, message := range messages { + if message.Visibility != database.ChatMessageVisibilityModel { + continue + } + parsed, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parsed) == 1 && parsed[0].Text == "Do not read secrets." { + foundModelContext = true + } + } + require.True(t, foundModelContext) + + messagesMu.Lock() + modelMessages := append([]chattest.OpenAIMessage(nil), secondMessages...) + messagesMu.Unlock() + require.True(t, openAIMessagesContain(modelMessages, "DENIED: blocked by policy")) + require.True(t, openAIMessagesContain(modelMessages, "Do not read secrets.")) + dispatchRows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + for _, row := range dispatchRows { + require.NotEqual(t, string(agenthooks.EventPostToolUse), row.Event) + } + dispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventPreToolUse) + require.Equal(t, "deny", dispatch.Decision.String) +} + +func TestPreToolUseSkipsProviderExecutedTools(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( + anthropicWebSearchPairChunks("ws-hook-skip", `{"query":"coder"}`, "search done", "end_turn")..., + ) + }) + user, org, model := seedAnthropicChatDependencies(t, db, anthropicURL) + model = enableAnthropicWebSearchForTest(t, db, model) + var preToolCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(agenthooks.PreToolUseData) string { + preToolCalls.Add(1) + return `{}` + }) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, anthropicURL, chattest.WithPreservePath())) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, model.ID, "search for coder") + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + require.Zero(t, preToolCalls.Load()) + call := requireToolCallPart(t, chatToolParts(ctx, t, db, chat.ID), "web_search") + require.True(t, call.ProviderExecuted) +} + +func TestPreToolUseHookDynamicAllowResponse(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"query":"original"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_dynamic_allow" + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + require.Equal(t, "call_dynamic_allow", data.ToolUseID) + return `{"permission":{"decision":"allow","input_override":{"query":"redacted"}},"model_context":"dynamic context","user_message":"dynamic notice","allowed_tools":["my_dynamic_tool"]}` + }) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-dynamic-allow", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + + var action database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + action, err = db.GetChatByID(ctx, chat.ID) + return err == nil && action.Status == database.ChatStatusRequiresAction + }, testutil.IntervalFast) + require.JSONEq(t, `["my_dynamic_tool"]`, string(action.HookAllowedTools.RawMessage)) + call := requireToolCallPart(t, chatToolParts(ctx, t, db, chat.ID), "my_dynamic_tool") + require.JSONEq(t, `{"query":"redacted"}`, string(call.Args)) + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + var foundContext bool + for _, message := range promptMessages { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "dynamic context" { + foundContext = true + } + } + require.True(t, foundContext) + var foundNotice bool + for _, message := range chatMessages(ctx, t, db, chat.ID) { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "dynamic notice" { + foundNotice = true + } + } + require.True(t, foundNotice) +} + +func TestPreToolUseHookEndChat(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + modelCalls.Add(1) + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{}`) + chunk.Choices[0].ToolCalls[0].ID = "call_end_chat" + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + require.Equal(t, "call_end_chat", data.ToolUseID) + return `{"end_chat":true}` + }) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-end-chat", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("stop before the tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + + var archived database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + archived, err = db.GetChatByID(ctx, chat.ID) + return err == nil && archived.Archived && archived.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.Equal(t, int32(1), modelCalls.Load()) + require.False(t, archived.RequiresActionDeadlineAt.Valid) + parts := chatToolParts(ctx, t, db, chat.ID) + call := requireToolCallPart(t, parts, "my_dynamic_tool") + require.Equal(t, "call_end_chat", call.ToolCallID) + result := requireToolResultPart(t, parts, "my_dynamic_tool") + require.Equal(t, "call_end_chat", result.ToolCallID) + require.True(t, result.IsError) + require.Contains(t, string(result.Result), "chat was ended") +} + +func TestPreToolUseHookEndChatShortCircuitsGeneratedStep(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + first := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"arg":"a"}`) + first.Choices[0].ToolCalls[0].ID = "call_end_first" + second := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"arg":"b"}`).Choices[0].ToolCalls[0] + second.ID = "call_fail_second" + second.Index = 1 + first.Choices[0].ToolCalls = append(first.Choices[0].ToolCalls, second) + return chattest.OpenAIStreamingResponse(first) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + + var hookCalls atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + hookCalls.Add(1) + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PreToolUseData) + if data.ToolUseID != "call_end_first" { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err = w.Write([]byte(`{"end_chat":true}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-end-chat-generated", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("stop before the tools"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + + var archived database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + archived, err = db.GetChatByID(ctx, chat.ID) + return err == nil && archived.Archived && archived.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.False(t, archived.LastError.Valid) + require.Equal(t, int32(1), hookCalls.Load(), "the second call must not dispatch after end_chat") + + rows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + for _, row := range rows { + if row.Event == string(agenthooks.EventPreToolUse) { + require.Equal(t, "call_end_first", row.ToolUseID.String) + } + } +} + +func TestPreToolUseHookEndChatSkipsLocalToolExecution(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_local_end_chat", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/secret.txt"}`, + }) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + hookCalls.Add(1) + require.Equal(t, "call_local_end_chat", data.ToolUseID) + return `{"end_chat":true}` + }) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + + var archived database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + var err error + archived, err = db.GetChatByID(ctx, chatID) + return err == nil && archived.Archived && archived.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.Equal(t, int32(1), hookCalls.Load()) + require.False(t, archived.LastError.Valid) + parts := chatToolParts(ctx, t, db, chatID) + call := requireToolCallPart(t, parts, "read_file") + require.Equal(t, "call_local_end_chat", call.ToolCallID) + result := requireToolResultPart(t, parts, "read_file") + require.Equal(t, "call_local_end_chat", result.ToolCallID) + require.True(t, result.IsError, "never-executed call must persist as a synthetic cancellation") +} + +func TestPreToolUseHookEndChatShortCircuitsLaterDispatches(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("resume")}) + require.NoError(t, err) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + LastModelConfigID: model.ID, + Title: "pending-tool-calls", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: userContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + chatID := created.Chat.ID + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "call_end_chat_first", + ToolName: "read_file", + Args: json.RawMessage(`{"path":"/tmp/first.txt"}`), + }, + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: "call_never_dispatched", + ToolName: "read_file", + Args: json.RawMessage(`{"path":"/tmp/second.txt"}`), + }, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(db, ps, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: created.InitialMessages[0].TurnID, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }}) + return err + })) + + var hookCalls atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + hookCalls.Add(1) + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.PreToolUseData) + require.True(t, ok) + if data.ToolUseID != "call_end_chat_first" { + w.WriteHeader(http.StatusInternalServerError) + return + } + _, err = w.Write([]byte(`{"end_chat":true}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + + var archived database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + var err error + archived, err = db.GetChatByID(ctx, chatID) + return err == nil && archived.Archived && archived.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.False(t, archived.LastError.Valid) + require.Equal(t, int32(1), hookCalls.Load()) + + rows, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + preToolUse := 0 + for _, row := range rows { + if row.Event == string(agenthooks.EventPreToolUse) { + preToolUse++ + require.Equal(t, "call_end_chat_first", row.ToolUseID.String) + } + } + require.Equal(t, 1, preToolUse) + + results := map[string]bool{} + for _, part := range chatToolParts(ctx, t, db, chatID) { + if part.Type == codersdk.ChatMessagePartTypeToolResult { + results[part.ToolCallID] = part.IsError + } + } + require.Equal(t, map[string]bool{ + "call_end_chat_first": true, + "call_never_dispatched": true, + }, results, "both never-executed calls must persist as synthetic cancellations") +} + +func TestPreToolUseHookRepeatedToolCallIDDispatchesFresh(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + + const toolUseID = "call_repeated" + toolInput := `{"path":"/tmp/dup.txt"}` + userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("resume")}) + require.NoError(t, err) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true}, + LastModelConfigID: model.ID, + Title: "repeated-tool-call-id", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: userContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + chatID := created.Chat.ID + turnID := created.InitialMessages[0].TurnID + toolCallContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: toolUseID, + ToolName: "read_file", + Args: json.RawMessage(toolInput), + }}) + require.NoError(t, err) + toolResultContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeToolResult, + ToolCallID: toolUseID, + ToolName: "read_file", + Result: json.RawMessage(`{"output":"data"}`), + }}) + require.NoError(t, err) + machine := chatstate.NewChatMachine(db, ps, chatID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: toolCallContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: turnID, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + { + Role: database.ChatMessageRoleTool, + Content: toolResultContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: turnID, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + { + Role: database.ChatMessageRoleAssistant, + Content: toolCallContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: turnID, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: model.ID, Valid: true}, + }, + }}) + return err + })) + + recordPreToolUseDecision(ctx, t, db, chatID, user.ID, turnID.UUID, toolUseID, "read_file", json.RawMessage(toolInput), agenthooks.PermissionAllow, "", nil) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + hookCalls.Add(1) + require.Equal(t, toolUseID, data.ToolUseID) + return `{"permission":{"decision":"deny","reason":"repeated call"}}` + }) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Equal(t, int32(1), hookCalls.Load(), "the repeated occurrence must dispatch fresh") + denied := 0 + for _, part := range chatToolParts(ctx, t, db, chatID) { + if part.Type == codersdk.ChatMessagePartTypeToolResult && part.IsError && + strings.Contains(string(part.Result), "repeated call") { + denied++ + } + } + require.Equal(t, 1, denied, "the repeated call must persist the fresh deny result") +} + +func TestPreToolUseHookRecordedEndChatSkipsExecution(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_recorded_end_chat", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/end.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + turnID := messages[0].TurnID.UUID + recordPreToolUseResponse(ctx, t, db, chatID, user.ID, turnID, "call_recorded_end_chat", "read_file", json.RawMessage(`{"path":"/tmp/end.txt"}`), agenthooks.PermissionAllow, "", nil, agenthooks.Response{ + EndChat: true, + }) + + var preToolUseCalls, postToolUseCalls atomic.Int32 + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + switch request.Type { + case agenthooks.EventPreToolUse: + preToolUseCalls.Add(1) + case agenthooks.EventPostToolUse: + postToolUseCalls.Add(1) + } + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + + var archived database.Chat + testutil.Eventually(ctx, t, func(ctx context.Context) bool { + var err error + archived, err = db.GetChatByID(ctx, chatID) + return err == nil && archived.Archived && archived.Status == database.ChatStatusWaiting + }, testutil.IntervalFast) + require.False(t, archived.LastError.Valid) + require.Zero(t, preToolUseCalls.Load(), "recorded decision must not re-dispatch") + require.Zero(t, postToolUseCalls.Load(), "never-executed call must not dispatch post_tool_use") + + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.Equal(t, "call_recorded_end_chat", result.ToolCallID) + require.True(t, result.IsError, "never-executed call must persist as a synthetic cancellation") + row, err := db.GetChatHookDispatchDecision(ctx, database.GetChatHookDispatchDecisionParams{ + ChatID: chatID, + ToolUseID: "call_recorded_end_chat", + ToolName: "read_file", + ToolInput: json.RawMessage(`{"path":"/tmp/end.txt"}`), + TurnID: uuid.NullUUID{UUID: turnID, Valid: true}, + }) + require.NoError(t, err) + require.True(t, row.EffectsAppliedAt.Valid, "replayed end_chat effects must be marked applied") +} + +func TestPreToolUseHookDispatchFailure(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + statusCode int + response string + result string + }{ + { + name: "http error", + statusCode: http.StatusInternalServerError, + result: "http_error", + }, + { + name: "ask protocol error", + response: `{"permission":{"decision":"ask"}}`, + result: "protocol_error", + }, + } + 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) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{}`) + chunk.Choices[0].ToolCalls[0].ID = "call_failure" + return chattest.OpenAIStreamingResponse(chunk) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + if tt.statusCode != 0 { + w.WriteHeader(tt.statusCode) + return + } + _, err := w.Write([]byte(tt.response)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-failure", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("fail before commit"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + + failed, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + dispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventPreToolUse) + require.Equal(t, tt.result, dispatch.Result) + lastError := chatLastErrorMessage(failed.LastError) + require.Contains(t, lastError, "hook dispatch failed: pre_tool_use: "+tt.result) + require.Contains(t, lastError, dispatch.ID.String()) + messages := chatMessages(ctx, t, db, chat.ID) + require.Len(t, messages, 1) + require.Equal(t, database.ChatMessageRoleUser, messages[0].Role) + }) + } +} + +func TestPreToolUseHookFailureAbortsWholeStep(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + first := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"value":1}`) + first.Choices[0].ToolCalls[0].ID = "call_first" + second := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"value":2}`).Choices[0].ToolCalls[0] + second.ID = "call_second" + second.Index = 1 + first.Choices[0].ToolCalls = append(first.Choices[0].ToolCalls, second) + return chattest.OpenAIStreamingResponse(first) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PreToolUseData) + if data.ToolUseID == "call_first" { + _, err = w.Write([]byte(`{"permission":{"decision":"allow","input_override":{"value":1}}}`)) + require.NoError(t, err) + return + } + require.Equal(t, "call_second", data.ToolUseID) + w.WriteHeader(http.StatusInternalServerError) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-atomic-failure", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call both tools"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + + messages := chatMessages(ctx, t, db, chat.ID) + require.Len(t, messages, 1) + rows, err := db.ListChatHookDispatchesByChatID(ctx, chat.ID) + require.NoError(t, err) + var toolRows []database.ChatHookDispatch + for _, row := range rows { + if row.Event == string(agenthooks.EventPreToolUse) { + toolRows = append(toolRows, row) + } + } + require.Len(t, toolRows, 2) + require.Equal(t, "call_first", toolRows[0].ToolUseID.String) + require.Equal(t, "allow", toolRows[0].Decision.String) + require.Equal(t, "call_second", toolRows[1].ToolUseID.String) + require.Equal(t, "http_error", toolRows[1].Result) +} + +func TestPreToolUseHookResumeRecordedDeny(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + modelCalls.Add(1) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("replanned")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_recorded_deny", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/secret.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + recordPreToolUseDecision(ctx, t, db, chatID, user.ID, uuid.New(), "call_recorded_deny", "read_file", json.RawMessage(`{"path":"/tmp/secret.txt"}`), agenthooks.PermissionAllow, "", nil) + recordPreToolUseDecision(ctx, t, db, chatID, user.ID, messages[0].TurnID.UUID, "call_recorded_deny", "read_file", json.RawMessage(`{"path":"/tmp/secret.txt"}`), agenthooks.PermissionDeny, "recorded policy reason", nil) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(agenthooks.PreToolUseData) string { + hookCalls.Add(1) + return `{}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Zero(t, hookCalls.Load()) + require.Equal(t, int32(1), modelCalls.Load()) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.True(t, result.IsError) + require.Contains(t, string(result.Result), "DENIED: recorded policy reason") +} + +func TestPreToolUseHookRecordedDecisionRequiresMatchingToolCall(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + recordedToolName string + recordedToolInput json.RawMessage + }{ + { + name: "InputMismatch", + recordedToolName: "read_file", + recordedToolInput: json.RawMessage(`{"path":"/tmp/secret.txt"}`), + }, + { + name: "ToolNameMismatch", + recordedToolName: "execute", + recordedToolInput: json.RawMessage(`{"path":"/tmp/harmless.txt"}`), + }, + { + name: "MissingToolName", + recordedToolInput: json.RawMessage(`{"path":"/tmp/harmless.txt"}`), + }, + } + 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) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + toolCallID := "call_binding_mismatch" + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: toolCallID, + ToolName: "read_file", + ToolInput: `{"path":"/tmp/harmless.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + staleDispatchID := recordPreToolUseDecision(ctx, t, db, chatID, user.ID, messages[0].TurnID.UUID, toolCallID, tt.recordedToolName, tt.recordedToolInput, agenthooks.PermissionDeny, "stale decision", nil) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + hookCalls.Add(1) + require.Equal(t, "read_file", data.ToolName) + require.JSONEq(t, `{"path":"/tmp/harmless.txt"}`, string(data.ToolInput)) + return `{}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/harmless.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(1) + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Positive(t, hookCalls.Load()) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.False(t, result.IsError) + + dispatches, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + var freshDispatches int + for _, dispatch := range dispatches { + if dispatch.Event != string(agenthooks.EventPreToolUse) { + continue + } + if dispatch.ID == staleDispatchID { + require.False(t, dispatch.EffectsAppliedAt.Valid) + continue + } + freshDispatches++ + require.True(t, dispatch.EffectsAppliedAt.Valid) + } + require.Equal(t, 1, freshDispatches) + }) + } +} + +func recordPreToolUseDecision( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + ownerID uuid.UUID, + turnID uuid.UUID, + toolUseID string, + toolName string, + toolInput json.RawMessage, + decision agenthooks.PermissionDecision, + reason string, + override json.RawMessage, +) uuid.UUID { + t.Helper() + return recordPreToolUseResponse(ctx, t, db, chatID, ownerID, turnID, toolUseID, toolName, toolInput, decision, reason, override, agenthooks.Response{}) +} + +func recordPreToolUseResponse( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + ownerID uuid.UUID, + turnID uuid.UUID, + toolUseID string, + toolName string, + toolInput json.RawMessage, + decision agenthooks.PermissionDecision, + reason string, + override json.RawMessage, + response agenthooks.Response, +) uuid.UUID { + t.Helper() + var allowedTools json.RawMessage + if response.AllowedTools != nil { + encoded, err := json.Marshal(response.AllowedTools) + require.NoError(t, err) + allowedTools = encoded + } + dispatchID := uuid.New() + _, err := db.InsertChatHookDispatch(ctx, database.InsertChatHookDispatchParams{ + ID: dispatchID, + ChatID: chatID, + Event: string(agenthooks.EventPreToolUse), + TurnID: uuid.NullUUID{UUID: turnID, Valid: true}, + ToolUseID: sql.NullString{String: toolUseID, Valid: true}, + ToolName: sql.NullString{String: toolName, Valid: toolName != ""}, + OwnerID: ownerID, + StartedAt: time.Now(), + }) + require.NoError(t, err) + _, err = db.FinalizeChatHookDispatch(ctx, database.FinalizeChatHookDispatchParams{ + FinishedAt: time.Now(), + Result: "ok", + Decision: sql.NullString{String: string(decision), Valid: true}, + DecisionReason: sql.NullString{String: reason, Valid: reason != ""}, + InputOverride: nullRawMessage(override), + OriginalInput: nullRawMessage(toolInput), + ModelContext: sql.NullString{String: response.ModelContext, Valid: response.ModelContext != ""}, + UserMessage: sql.NullString{String: response.UserMessage, Valid: response.UserMessage != ""}, + AllowedTools: nullRawMessage(allowedTools), + EndChat: sql.NullBool{Bool: response.EndChat, Valid: response.EndChat}, + ID: dispatchID, + ChatID: chatID, + OwnerID: ownerID, + }) + require.NoError(t, err) + return dispatchID +} + +func TestPreToolUseHookReplaysOverrideForOriginalInput(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_recorded_override", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/original.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + dispatchID := recordPreToolUseDecision( + ctx, + t, + db, + chatID, + user.ID, + messages[0].TurnID.UUID, + "call_recorded_override", + "read_file", + json.RawMessage(`{"path":"/tmp/original.txt"}`), + agenthooks.PermissionAllow, + "", + json.RawMessage(`{"path":"/tmp/reviewed.txt"}`), + ) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(agenthooks.PreToolUseData) string { + hookCalls.Add(1) + return `{"permission":{"decision":"allow","input_override":{"path":"/tmp/reviewed.txt"}}}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/reviewed.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(1) + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Zero(t, hookCalls.Load()) + call := requireToolCallPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.JSONEq(t, `{"path":"/tmp/reviewed.txt"}`, string(call.Args)) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.False(t, result.IsError) + + dispatches, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + var preToolDispatches int + var recordedDispatchFound bool + for _, dispatch := range dispatches { + if dispatch.Event != string(agenthooks.EventPreToolUse) { + continue + } + preToolDispatches++ + if dispatch.ID == dispatchID { + recordedDispatchFound = true + require.True(t, dispatch.EffectsAppliedAt.Valid) + } + } + require.True(t, recordedDispatchFound) + require.Equal(t, 1, preToolDispatches) +} + +func TestPreToolUseHookPendingPolicyNarrowsSpawnedChild(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + ModelConfigID: model.ID, + ToolCallID: "call_spawn_with_policy", + ToolName: "spawn_agent", + ToolInput: `{"type":"general","prompt":"inspect the workspace","title":"child"}`, + }) + require.NoError(t, db.UpdateChatHookAllowedTools(ctx, database.UpdateChatHookAllowedToolsParams{ + ID: chatID, + HookAllowedTools: nullRawMessage([]byte(`["read_file","spawn_agent"]`)), + })) + + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + response := `{}` + if request.Type == agenthooks.EventPreToolUse { + decoded, err := request.Decode() + require.NoError(t, err) + data := decoded.(*agenthooks.PreToolUseData) + require.Equal(t, "call_spawn_with_policy", data.ToolUseID) + require.Equal(t, "spawn_agent", data.ToolName) + response = `{"allowed_tools":["read_file"]}` + } + _, err := w.Write([]byte(response)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + parent, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(parent.HookAllowedTools.RawMessage)) + children, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{ + ParentIds: []uuid.UUID{chatID}, + }) + require.NoError(t, err) + require.Len(t, children, 1) + child := children[0].Chat + require.JSONEq(t, `["read_file"]`, string(child.HookAllowedTools.RawMessage)) + + dispatch := lifecycleDispatch(t, db, chatID, agenthooks.EventPreToolUse) + require.Equal(t, "call_spawn_with_policy", dispatch.ToolUseID.String) + require.True(t, dispatch.EffectsAppliedAt.Valid) +} + +func TestPreToolUseHookReplaysUnappliedEffects(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_replay_effects", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/a.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + allowed := []string{"read_file"} + replayedDispatchID := recordPreToolUseResponse(ctx, t, db, chatID, user.ID, messages[0].TurnID.UUID, "call_replay_effects", "read_file", json.RawMessage(`{"path":"/tmp/a.txt"}`), agenthooks.PermissionAllow, "", nil, agenthooks.Response{ + ModelContext: "recorded context", + UserMessage: "recorded notice", + AllowedTools: &allowed, + }) + staleDispatchID := recordPreToolUseResponse(ctx, t, db, chatID, user.ID, messages[0].TurnID.UUID, "call_replay_effects", "execute", json.RawMessage(`{"path":"/tmp/a.txt"}`), agenthooks.PermissionDeny, "stale decision", nil, agenthooks.Response{ + ModelContext: "stale context", + }) + + var preToolUseCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(agenthooks.PreToolUseData) string { + preToolUseCalls.Add(1) + return `{}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true}, nil).AnyTimes() + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Zero(t, preToolUseCalls.Load(), "recorded decision must not re-dispatch") + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.False(t, result.IsError) + + promptRows, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + var contextCount, staleContextCount int + for _, row := range promptRows { + switch hookMessageText(t, row) { + case "recorded context": + contextCount++ + case "stale context": + staleContextCount++ + } + } + var noticeCount int + for _, row := range chatMessages(ctx, t, db, chatID) { + if hookMessageText(t, row) == "recorded notice" { + noticeCount++ + } + } + require.Equal(t, 1, contextCount, "replayed model context must commit exactly once") + require.Zero(t, staleContextCount) + require.Equal(t, 1, noticeCount, "replayed user notice must commit exactly once") + updated, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(updated.HookAllowedTools.RawMessage)) + dispatches, err := db.ListChatHookDispatchesByChatID(ctx, chatID) + require.NoError(t, err) + var foundReplayed, foundStale bool + for _, dispatch := range dispatches { + switch dispatch.ID { + case replayedDispatchID: + foundReplayed = true + require.True(t, dispatch.EffectsAppliedAt.Valid) + case staleDispatchID: + foundStale = true + require.False(t, dispatch.EffectsAppliedAt.Valid) + } + } + require.True(t, foundReplayed) + require.True(t, foundStale) +} + +func TestPreToolUseHookSkipsAppliedEffects(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_applied_effects", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/a.txt"}`, + }) + messages, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + require.True(t, messages[0].TurnID.Valid) + allowed := []string{"read_file"} + dispatchID := recordPreToolUseResponse(ctx, t, db, chatID, user.ID, messages[0].TurnID.UUID, "call_applied_effects", "read_file", json.RawMessage(`{"path":"/tmp/a.txt"}`), agenthooks.PermissionAllow, "", nil, agenthooks.Response{ + ModelContext: "recorded context", + AllowedTools: &allowed, + }) + require.NoError(t, db.MarkChatHookDispatchEffectsApplied(ctx, database.MarkChatHookDispatchEffectsAppliedParams{ + ChatID: chatID, + DispatchIds: []uuid.UUID{dispatchID}, + })) + + var preToolUseCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(agenthooks.PreToolUseData) string { + preToolUseCalls.Add(1) + return `{}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{Success: true}, nil).AnyTimes() + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Zero(t, preToolUseCalls.Load(), "recorded decision must not re-dispatch") + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.False(t, result.IsError) + + promptRows, err := db.GetChatMessagesForPromptByChatID(ctx, chatID) + require.NoError(t, err) + for _, row := range promptRows { + require.NotEqual(t, "recorded context", hookMessageText(t, row), "applied effects must not replay") + } + updated, err := db.GetChatByID(ctx, chatID) + require.NoError(t, err) + require.False(t, updated.HookAllowedTools.Valid, "applied allowed_tools must not replay") +} + +func TestPreToolUseHookResumeFallback(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + modelCalls.Add(1) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID) + chatID := seedPendingToolCall(ctx, t, db, ps, pendingToolCallSeed{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: ws.ID, + AgentID: dbAgent.ID, + ModelConfigID: model.ID, + ToolCallID: "call_resume_fallback", + ToolName: "read_file", + ToolInput: `{"path":"/tmp/original.txt"}`, + }) + + var hookCalls atomic.Int32 + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + hookCalls.Add(1) + require.Equal(t, "call_resume_fallback", data.ToolUseID) + return `{"permission":{"decision":"allow","input_override":{"path":"/tmp/resume.txt"}}}` + }) + ctrl := gomock.NewController(t) + mockConn := agentconnmock.NewMockAgentConn(ctrl) + setupToolExecutionAgentConn(t, mockConn) + mockConn.EXPECT().ReadFileLines(gomock.Any(), "/tmp/resume.txt", int64(1), int64(0), gomock.Any()). + Return(workspacesdk.ReadFileLinesResponse{ + Success: true, FileSize: 4, TotalLines: 1, LinesRead: 1, Content: "data", + }, nil). + Times(1) + + server := newTestServer(t, db, ps, uuid.New(), func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + cfg.AgentConn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) { + require.Equal(t, dbAgent.ID, agentID) + return mockConn, func() {}, nil + } + }) + server.Start() + waitForChatStatus(ctx, t, db, chatID, database.ChatStatusWaiting) + + require.Equal(t, int32(1), hookCalls.Load()) + require.Equal(t, int32(1), modelCalls.Load()) + dispatch := lifecycleDispatch(t, db, chatID, agenthooks.EventPreToolUse) + require.Equal(t, "allow", dispatch.Decision.String) + call := requireToolCallPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.JSONEq(t, `{"path":"/tmp/resume.txt"}`, string(call.Args)) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chatID), "read_file") + require.False(t, result.IsError) +} + +type pendingToolCallSeed struct { + OrganizationID uuid.UUID + OwnerID uuid.UUID + WorkspaceID uuid.UUID + AgentID uuid.UUID + ModelConfigID uuid.UUID + ToolCallID string + ToolName string + ToolInput string + DynamicTools json.RawMessage +} + +func seedPendingToolCall( + ctx context.Context, + t *testing.T, + db database.Store, + ps dbpubsub.Pubsub, + seed pendingToolCallSeed, +) uuid.UUID { + t.Helper() + userContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText("resume")}) + require.NoError(t, err) + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: seed.OrganizationID, + OwnerID: seed.OwnerID, + WorkspaceID: uuid.NullUUID{UUID: seed.WorkspaceID, Valid: seed.WorkspaceID != uuid.Nil}, + AgentID: uuid.NullUUID{UUID: seed.AgentID, Valid: seed.AgentID != uuid.Nil}, + LastModelConfigID: seed.ModelConfigID, + Title: "pending-tool-call", + DynamicTools: nullRawMessage(seed.DynamicTools), + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: userContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: uuid.NullUUID{UUID: uuid.New(), Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + CreatedBy: uuid.NullUUID{UUID: seed.OwnerID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: seed.ModelConfigID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + assistantContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{ + { + Type: codersdk.ChatMessagePartTypeToolCall, + ToolCallID: seed.ToolCallID, + ToolName: seed.ToolName, + Args: json.RawMessage(seed.ToolInput), + }, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(db, ps, created.Chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.CommitStep(chatstate.CommitStepInput{Messages: []chatstate.Message{ + { + Role: database.ChatMessageRoleAssistant, + Content: assistantContent, + Visibility: database.ChatMessageVisibilityBoth, + TurnID: created.InitialMessages[0].TurnID, + ContentVersion: chatprompt.CurrentContentVersion, + ModelConfigID: uuid.NullUUID{UUID: seed.ModelConfigID, Valid: true}, + }, + }}) + return err + })) + return created.Chat.ID +} + +func TestPreToolUseHookDynamicDeny(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"query":"test"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_dynamic_denied" + return chattest.OpenAIStreamingResponse(chunk) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("replanned")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := preToolUseConsumer(t, func(data agenthooks.PreToolUseData) string { + require.Equal(t, "call_dynamic_denied", data.ToolUseID) + return `{"permission":{"decision":"deny","reason":"dynamic denied"}}` + }) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "pre-tool-use-dynamic-deny", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + chatResult, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.False(t, chatResult.RequiresActionDeadlineAt.Valid) + require.Equal(t, int32(2), modelCalls.Load()) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), "my_dynamic_tool") + require.True(t, result.IsError) + require.Contains(t, string(result.Result), "DENIED: dynamic denied") +} + +func TestHookAllowedToolsExcludesDynamicTool(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if modelCalls.Add(1) == 1 { + chunk := chattest.OpenAIToolCallChunk("my_dynamic_tool", `{"query":"test"}`) + chunk.Choices[0].ToolCalls[0].ID = "call_excluded_dynamic" + return chattest.OpenAIStreamingResponse(chunk) + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("replanned")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + body := `{}` + if request.Type == agenthooks.EventUserPromptSubmit { + body = `{"allowed_tools":["read_file"]}` + } + _, err := w.Write([]byte(body)) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "allowed-tools-excludes-dynamic", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("call the dynamic tool"), + }, + DynamicTools: dynamicToolJSON(t, "my_dynamic_tool"), + }) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + chatResult, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.JSONEq(t, `["read_file"]`, string(chatResult.HookAllowedTools.RawMessage)) + require.False(t, chatResult.RequiresActionDeadlineAt.Valid) + require.Equal(t, int32(2), modelCalls.Load()) + result := requireToolResultPart(t, chatToolParts(ctx, t, db, chat.ID), "my_dynamic_tool") + require.True(t, result.IsError) + require.Contains(t, string(result.Result), "Tool not active in this turn") +} + +func preToolUseConsumer(t *testing.T, response func(agenthooks.PreToolUseData) string) *httptest.Server { + t.Helper() + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventPreToolUse { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + decoded, err := request.Decode() + require.NoError(t, err) + data, ok := decoded.(*agenthooks.PreToolUseData) + require.True(t, ok) + _, err = w.Write([]byte(response(*data))) + require.NoError(t, err) + })) + t.Cleanup(consumer.Close) + return consumer +} diff --git a/coderd/x/chatd/runner.go b/coderd/x/chatd/runner.go index e36364aab8629..72b916a38a6ea 100644 --- a/coderd/x/chatd/runner.go +++ b/coderd/x/chatd/runner.go @@ -57,6 +57,8 @@ type runner struct { tasksByIndex map[taskIndexKey]taskInstanceID localLocks *localLockSet debugTurn *runnerDebugTurn + sessionStart sessionStartTracker + stopNudges stopNudgeTracker } func newRunner(ctx context.Context, mgr *runnerManager, rec *runnerRecord, opts chatWorkerOptions) *runner { @@ -227,6 +229,8 @@ func (r *runner) spawnTaskIfNeeded(kind taskKind, state runnerStateUpdate) { Status: state.Status, RequiresActionDeadlineAt: state.RequiresActionDeadlineAt, DebugTurn: r.debugTurn, + SessionStart: &r.sessionStart, + StopNudges: &r.stopNudges, } go r.runTask(taskCtx, kind, key, input, done) } diff --git a/coderd/x/chatd/runner_test.go b/coderd/x/chatd/runner_test.go index ef95069518632..0438127892443 100644 --- a/coderd/x/chatd/runner_test.go +++ b/coderd/x/chatd/runner_test.go @@ -42,6 +42,7 @@ func TestRunner_CancelsActiveTaskWhenHistoryChanges(t *testing.T) { require.NotErrorIs(t, context.Cause(first.ctx), errTaskTimeout) second := starter.waitCall(t, taskKindGeneration, chat.ID) require.Equal(t, updated.HistoryVersion, second.input.HistoryVersion) + require.Same(t, first.input.SessionStart, second.input.SessionStart) } func TestRunner_CancelsActiveTaskWhenStatusChanges(t *testing.T) { diff --git a/coderd/x/chatd/stop_test.go b/coderd/x/chatd/stop_test.go new file mode 100644 index 0000000000000..37667c460825a --- /dev/null +++ b/coderd/x/chatd/stop_test.go @@ -0,0 +1,221 @@ +package chatd_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/database" + "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/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" +) + +func TestStopHookNoOpFinishesTurn(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + var stopCalls atomic.Int32 + consumer := stopConsumer(t, func() (int, string) { + stopCalls.Add(1) + return http.StatusOK, `{}` + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "stop-noop", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("finish normally"), + }, + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(context.Context) bool { + return stopCalls.Load() == 1 + }, testutil.IntervalFast) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(1), stopCalls.Load()) + require.Equal(t, "ok", lifecycleDispatch(t, db, chat.ID, agenthooks.EventStop).Result) +} + +func TestStopHookNudgeContinuesOnce(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + var modelCalls atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + switch modelCalls.Add(1) { + case 1: + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("first answer")...) + case 2: + var found bool + for _, message := range req.Messages { + found = found || strings.Contains(message.Content, "continue please") + } + require.True(t, found) + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("second answer")...) + default: + require.FailNow(t, "stop nudge exceeded continuation cap") + return chattest.OpenAIStreamingResponse() + } + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + var stopCalls atomic.Int32 + consumer := stopConsumer(t, func() (int, string) { + stopCalls.Add(1) + return http.StatusOK, `{"model_context":"continue please"}` + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "stop-nudge", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("continue once"), + }, + }) + require.NoError(t, err) + testutil.Eventually(ctx, t, func(context.Context) bool { + return stopCalls.Load() == 2 + }, testutil.IntervalFast) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + require.Equal(t, int32(2), modelCalls.Load()) + require.Equal(t, int32(2), stopCalls.Load()) + + promptMessages, err := db.GetChatMessagesForPromptByChatID(ctx, chat.ID) + require.NoError(t, err) + var contextRows int + for _, message := range promptMessages { + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "continue please" { + contextRows++ + require.Equal(t, database.ChatMessageVisibilityModel, message.Visibility) + } + } + require.Equal(t, 2, contextRows) +} + +func TestStopHookDispatchFailureErrorsChat(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := stopConsumer(t, func() (int, string) { + return http.StatusInternalServerError, "" + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "stop-failure", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("fail on stop"), + }, + }) + require.NoError(t, err) + failed := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) + dispatch := lifecycleDispatch(t, db, chat.ID, agenthooks.EventStop) + require.Equal(t, "http_error", dispatch.Result) + lastError := chatLastErrorMessage(failed.LastError) + require.Contains(t, lastError, "hook dispatch failed: stop: http_error") + require.Contains(t, lastError, dispatch.ID.String()) +} + +func TestStopHookEndChatArchives(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + db, ps := dbtestutil.NewDB(t) + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("done")...) + }) + user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + consumer := stopConsumer(t, func() (int, string) { + return http.StatusOK, `{"user_message":"chat ended","end_chat":true}` + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + cfg.HookDispatcher = newHookDispatcher(t, db, consumer) + }) + chat, err := server.CreateChat(ctx, chatd.CreateOptions{ + OrganizationID: org.ID, + OwnerID: user.ID, + Title: "stop-end-chat", + ModelConfigID: model.ID, + InitialUserContent: []codersdk.ChatMessagePart{ + codersdk.ChatMessageText("archive on stop"), + }, + }) + require.NoError(t, err) + + testutil.Eventually(ctx, t, func(context.Context) bool { + updated, err := db.GetChatByID(ctx, chat.ID) + return err == nil && updated.Archived + }, testutil.IntervalFast) +} + +func stopConsumer(t *testing.T, response func() (int, string)) *httptest.Server { + t.Helper() + consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var request agenthooks.Request + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + if request.Type != agenthooks.EventStop { + _, err := w.Write([]byte(`{}`)) + require.NoError(t, err) + return + } + status, body := response() + w.WriteHeader(status) + if body != "" { + _, err := w.Write([]byte(body)) + require.NoError(t, err) + } + })) + t.Cleanup(consumer.Close) + return consumer +} diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 48c5b9bf487ab..53e3224d887ec 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -25,6 +25,7 @@ import ( "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/agenthooks" "github.com/coder/coder/v2/codersdk/workspacesdk" ) @@ -608,6 +609,11 @@ func (p *Server) subagentTools( options, ) if err != nil { + // Surface the consumer reason so the model can adjust its prompt. + var denied *UserPromptDeniedError + if errors.As(err, &denied) && denied.UserMessage != "" { + return fantasy.NewTextErrorResponse(err.Error() + ": " + denied.UserMessage), nil + } return fantasy.NewTextErrorResponse(err.Error()), nil } @@ -948,6 +954,13 @@ func (p *Server) loadSubagentSpawnParentChat( if err := validateSubagentSpawnParent(parent); err != nil { return database.Chat{}, err } + if pendingPolicy, ok := pendingHookAllowedToolsFromContext(ctx, parent.ID); ok { + narrowed, err := chatstate.NarrowHookAllowedTools(parent.HookAllowedTools, pendingPolicy) + if err != nil { + return database.Chat{}, xerrors.Errorf("narrow pending parent hook allowed tools: %w", err) + } + parent.HookAllowedTools = narrowed + } return parent, nil } @@ -1040,9 +1053,6 @@ func (p *Server) createChildSubagentChatWithOptions( } title = strings.TrimSpace(title) - if title == "" { - title = subagentFallbackChatTitle(prompt) - } rootChatID := parent.ID if parent.RootChatID.Valid { @@ -1086,6 +1096,40 @@ func (p *Server) createChildSubagentChatWithOptions( return database.Chat{}, limitErr } + // Review before persistence so spawned chats cannot bypass prompt policy. + childChatID := uuid.New() + var childTurnID *uuid.UUID + var hookResponse agenthooks.Response + if p.hookDispatcher != nil && p.hookDispatcher.Enabled() { + mintedTurnID := uuid.New() + childTurnID = &mintedTurnID + hookChat := database.Chat{} + hookChat.ID = childChatID + hookChat.OwnerID = parent.OwnerID + hookChat.WorkspaceID = parent.WorkspaceID + hookChat.ParentChatID = uuid.NullUUID{UUID: parent.ID, Valid: true} + hookChat.RootChatID = uuid.NullUUID{UUID: rootChatID, Valid: true} + hookResponse, err = p.dispatchUserPromptSubmit(ctx, hookChat, mintedTurnID, []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)}) + if err != nil { + return database.Chat{}, err + } + // The child chat does not exist yet, so an end_chat response + // denies the spawn instead of archiving. + if hookResponse.EndChat { + return database.Chat{}, &UserPromptDeniedError{UserMessage: hookResponse.UserMessage} + } + override, overridden, overrideErr := userPromptOverride(hookResponse) + if overrideErr != nil { + return database.Chat{}, overrideErr + } + if overridden { + prompt = override + } + } + if title == "" { + title = subagentFallbackChatTitle(prompt) + } + workspaceAwareness := workspaceDetachedNoCreateAwareness if parent.WorkspaceID.Valid { workspaceAwareness = workspaceAttachedAwareness @@ -1122,17 +1166,37 @@ func (p *Server) createChildSubagentChatWithOptions( } initialMessages = append(initialMessages, systemMessage(workspaceAwarenessContent, modelConfigID)) + prefixMessages, err := hookPrefixMessages(hookResponse, modelConfigID, childTurnID) + if err != nil { + return database.Chat{}, err + } + initialMessages = append(initialMessages, prefixMessages...) + // The child shares the parent's workspace and agent, so it inherits // workspace context the same way a top-level chat does: pinned from the // agent's latest snapshot (see hydrateChatContextOnCreate below). The // parent's context is not copied into child history. - initialMessages = append(initialMessages, userMessage(userContent, modelConfigID, parent.OwnerID, opts.reasoningEffortOverride)) + childUserMessage := userMessage(userContent, modelConfigID, parent.OwnerID, opts.reasoningEffortOverride) + if childTurnID != nil { + childUserMessage.TurnID = uuid.NullUUID{UUID: *childTurnID, Valid: true} + } + initialMessages = append(initialMessages, childUserMessage) + + // Child policy may narrow the parent policy but cannot widen it. + responseTools, err := hookAllowedTools(hookResponse) + if err != nil { + return database.Chat{}, err + } + childPolicy, err := chatstate.NarrowHookAllowedTools(parent.HookAllowedTools, responseTools) + if err != nil { + return database.Chat{}, err + } publisher := p.pubsub if publisher == nil { publisher = dbpubsub.NewInMemory() } - result, err := chatstate.CreateChat(ctx, p.db, publisher, chatstate.CreateChatInput{ + result, err := chatstate.CreateChatWithID(ctx, p.db, publisher, childChatID, childPolicy, chatstate.CreateChatInput{ OrganizationID: parent.OrganizationID, OwnerID: parent.OwnerID, WorkspaceID: parent.WorkspaceID, diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 7859d91302eb1..10d4a0595c8db 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -15,6 +15,7 @@ import ( "charm.land/fantasy" "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -34,6 +35,7 @@ import ( "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/chattool" + "github.com/coder/coder/v2/coderd/x/chathooks" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" "github.com/coder/quartz" @@ -267,6 +269,144 @@ func insertInternalAIProvider( }) } +func TestCreateChildSubagentChatDispatchesUserPromptSubmit(t *testing.T) { + t.Parallel() + + newFixture := func(t *testing.T, handler http.HandlerFunc) (context.Context, database.Store, database.Chat, *Server) { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + db, _ := 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}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + parent := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + HookAllowedTools: pqtype.NullRawMessage{RawMessage: []byte(`["read_file","spawn_agent"]`), Valid: true}, + }) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + ctx = aibridge.WithDelegatedAPIKeyID(ctx, apiKey.ID) + + consumer := httptest.NewServer(handler) + t.Cleanup(consumer.Close) + server := &Server{ + db: db, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + hookDispatcher: chathooks.New( + slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + db, + consumer.Client(), + consumer.URL, + "test-hook-secret-32-bytes-minimum!!", + time.Second, + "test-deployment", + "test-version", + prometheus.NewRegistry(), + ), + } + return ctx, db, parent, server + } + + t.Run("RewriteAndPolicyIntersection", func(t *testing.T) { + t.Parallel() + + var meta struct { + sync.Mutex + parentChatID string + prompt string + } + ctx, db, parent, server := newFixture(t, func(rw http.ResponseWriter, r *http.Request) { + var request struct { + Type string `json:"type"` + Meta struct { + ParentChatID *uuid.UUID `json:"parent_chat_id"` + } `json:"meta"` + Data struct { + Prompt string `json:"prompt"` + } `json:"data"` + } + require.NoError(t, json.NewDecoder(r.Body).Decode(&request)) + require.Equal(t, "user_prompt_submit", request.Type) + meta.Lock() + if request.Meta.ParentChatID != nil { + meta.parentChatID = request.Meta.ParentChatID.String() + } + meta.prompt = request.Data.Prompt + meta.Unlock() + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write([]byte(`{ + "permission": {"decision": "allow", "input_override": {"prompt": "REVIEWED: inspect"}}, + "allowed_tools": ["read_file", "execute"] + }`)) + }) + + child, err := server.createChildSubagentChatWithOptions(ctx, parent, "inspect the workspace", "", childSubagentChatOptions{}) + require.NoError(t, err) + + meta.Lock() + require.Equal(t, parent.ID.String(), meta.parentChatID, "spawn dispatch must identify the parent chat") + require.Equal(t, "inspect the workspace", meta.prompt) + meta.Unlock() + + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: child.ID}) + require.NoError(t, err) + var childUserMessage database.ChatMessage + for _, message := range messages { + if message.Role == database.ChatMessageRoleUser { + childUserMessage = message + break + } + } + require.NotZero(t, childUserMessage.ID) + require.True(t, childUserMessage.Content.Valid) + require.Contains(t, string(childUserMessage.Content.RawMessage), "REVIEWED: inspect", + "the hook rewrite must land as the child's initial prompt") + require.NotContains(t, string(childUserMessage.Content.RawMessage), "inspect the workspace") + require.True(t, childUserMessage.TurnID.Valid, "the gated spawn prompt starts the child's first turn") + + require.True(t, child.HookAllowedTools.Valid) + require.JSONEq(t, `["read_file"]`, string(child.HookAllowedTools.RawMessage), + "child policy is the parent policy narrowed by the response") + }) + + t.Run("DenyRefusesSpawn", func(t *testing.T) { + t.Parallel() + + ctx, db, parent, server := newFixture(t, func(rw http.ResponseWriter, _ *http.Request) { + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write([]byte(`{"permission": {"decision": "deny", "reason": "spawn blocked"}, "user_message": "not allowed"}`)) + }) + + _, err := server.createChildSubagentChatWithOptions(ctx, parent, "exfiltrate secrets", "", childSubagentChatOptions{}) + var denied *UserPromptDeniedError + require.ErrorAs(t, err, &denied) + require.Equal(t, "not allowed", denied.UserMessage) + + chats, err := db.GetChildChatsByParentIDs(ctx, database.GetChildChatsByParentIDsParams{ + ParentIds: []uuid.UUID{parent.ID}, + }) + require.NoError(t, err) + require.Empty(t, chats, "a denied spawn must not create a child chat") + }) + + t.Run("ParentPolicyInheritedWithoutResponsePolicy", func(t *testing.T) { + t.Parallel() + + ctx, _, parent, server := newFixture(t, func(rw http.ResponseWriter, _ *http.Request) { + rw.Header().Set("Content-Type", "application/json") + _, _ = rw.Write([]byte(`{}`)) + }) + + child, err := server.createChildSubagentChatWithOptions(ctx, parent, "inspect the workspace", "", childSubagentChatOptions{}) + require.NoError(t, err) + require.True(t, child.HookAllowedTools.Valid) + require.JSONEq(t, `["read_file","spawn_agent"]`, string(child.HookAllowedTools.RawMessage), + "a restricted parent must not spawn an unrestricted child") + }) +} + func TestResolveUserProviderAPIKeys_AIProvider(t *testing.T) { t.Parallel() diff --git a/coderd/x/chathooks/dispatcher.go b/coderd/x/chathooks/dispatcher.go new file mode 100644 index 0000000000000..a7b0cd73da353 --- /dev/null +++ b/coderd/x/chathooks/dispatcher.go @@ -0,0 +1,717 @@ +// Package chathooks dispatches chat lifecycle events to an external webhook. +package chathooks + +import ( + "bytes" + "context" + "crypto/sha256" + "database/sql" + "encoding/hex" + "encoding/json" + "errors" + "io" + "math" + "net" + "net/http" + "syscall" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + "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/codersdk/agenthooks" +) + +const ( + maxConcurrentDispatches = 256 + maxResponseBodyBytes = 1_048_576 + maxModelContextBytes = 16_384 + capacityWaitLimit = 250 * time.Millisecond + retryBackoff = 250 * time.Millisecond + finalizeTimeout = 2 * time.Second + // clockSkewLeeway tolerates small clock differences with hook consumers. + clockSkewLeeway = 30 * time.Second +) + +// DispatchResult classifies the terminal outcome of a dispatch attempt. +type DispatchResult string + +const ( + ResultOK DispatchResult = "ok" + ResultDenied DispatchResult = "denied" + ResultHTTPError DispatchResult = "http_error" + ResultProtocolError DispatchResult = "protocol_error" + ResultTimeout DispatchResult = "timeout" + ResultConnectionError DispatchResult = "connection_error" + ResultOverCapacity DispatchResult = "over_capacity" + ResultInternalError DispatchResult = "internal_error" +) + +type store interface { + InsertChatHookDispatch(context.Context, database.InsertChatHookDispatchParams) (database.ChatHookDispatch, error) + FinalizeChatHookDispatch(context.Context, database.FinalizeChatHookDispatchParams) (database.ChatHookDispatch, error) +} + +// Event carries the identities persisted with each delivery attempt. +type Event struct { + Type agenthooks.EventType + agenthooks.ChatRef + Data any +} + +func (e Event) toolMetadata() (toolUseID, toolName *string) { + switch e.Type { + case agenthooks.EventPreToolUse: + if data, ok := dataValue[agenthooks.PreToolUseData](e.Data); ok { + return &data.ToolUseID, &data.ToolName + } + case agenthooks.EventPostToolUse: + if data, ok := dataValue[agenthooks.PostToolUseData](e.Data); ok { + return &data.ToolUseID, &data.ToolName + } + } + return nil, nil +} + +// DispatchError preserves the attempt ID and failure class. +type DispatchError struct { + Class DispatchResult + DispatchID uuid.UUID + Err error +} + +func (e *DispatchError) Error() string { + return e.Err.Error() +} + +func (e *DispatchError) Unwrap() error { + return e.Err +} + +func newDispatchError(class DispatchResult, dispatchID uuid.UUID, err error) error { + if err == nil { + return nil + } + return &DispatchError{Class: class, DispatchID: dispatchID, Err: err} +} + +// Dispatcher persists and delivers lifecycle hook attempts. +type Dispatcher struct { + logger slog.Logger + db store + client *http.Client + hookURL string + secret []byte + timeout time.Duration + deploymentID string + userAgent string + semaphore chan struct{} + metrics *metrics +} + +// New copies (or creates) the HTTP client and disables redirects for signed +// requests. +func New( + logger slog.Logger, + db database.Store, + client *http.Client, + hookURL string, + secret string, + timeout time.Duration, + deploymentID string, + coderVersion string, + reg prometheus.Registerer, +) *Dispatcher { + if client == nil { + client = &http.Client{} + } else { + clientCopy := *client + client = &clientCopy + } + client.CheckRedirect = func(_ *http.Request, _ []*http.Request) error { + return http.ErrUseLastResponse + } + return &Dispatcher{ + logger: logger.Named("chat_hook_dispatcher"), + db: db, + client: client, + hookURL: hookURL, + secret: []byte(secret), + timeout: timeout, + deploymentID: deploymentID, + userAgent: "coderd/" + coderVersion, + semaphore: make(chan struct{}, maxConcurrentDispatches), + metrics: newMetrics(reg), + } +} + +func (d *Dispatcher) Enabled() bool { + return d != nil && d.hookURL != "" +} + +// Dispatch persists and delivers one event. The returned ID identifies the +// dispatch so effects can bind to a single attempt. +func (d *Dispatcher) Dispatch(ctx context.Context, event Event) (agenthooks.Response, uuid.UUID, error) { + if !d.Enabled() { + return agenthooks.Response{}, uuid.Nil, xerrors.New("chat hook dispatcher is not enabled") + } + + startedAt := time.Now() + dispatchID := uuid.New() + wait := min(d.timeout, capacityWaitLimit) + if wait < 0 { + wait = 0 + } + capacityTimer := time.NewTimer(wait) + defer capacityTimer.Stop() + + select { + case d.semaphore <- struct{}{}: + defer func() { <-d.semaphore }() + case <-ctx.Done(): + response, err := d.finishWithoutPost(ctx, event, dispatchID, startedAt, ResultTimeout, ctx.Err()) + return response, dispatchID, err + case <-capacityTimer.C: + response, err := d.finishWithoutPost(ctx, event, dispatchID, startedAt, ResultOverCapacity, context.DeadlineExceeded) + return response, dispatchID, err + } + + // Insert with a detached context so a caller canceled right after + // acquiring capacity still leaves an audit row, matching finishWithoutPost. + insertCtx, cancelInsert := context.WithTimeout(context.WithoutCancel(ctx), finalizeTimeout) + insertErr := d.insert(insertCtx, event, dispatchID, startedAt) + cancelInsert() + if insertErr != nil { + d.metrics.observe(event.Type, ResultInternalError, agenthooks.Response{}, time.Since(startedAt)) + return agenthooks.Response{}, dispatchID, newDispatchError(ResultInternalError, dispatchID, xerrors.Errorf("insert chat hook dispatch: %w", insertErr)) + } + + response, outcome := d.prepareAndPost(ctx, event, dispatchID) + finalizeErr := d.finalize(ctx, event, dispatchID, outcome) + d.metrics.observe(event.Type, outcome.result, outcome.response, time.Since(startedAt)) + if finalizeErr != nil { + d.logger.Error(context.WithoutCancel(ctx), "failed to finalize chat hook dispatch", slog.Error(finalizeErr)) + if outcome.err != nil { + return agenthooks.Response{}, dispatchID, newDispatchError(outcome.result, dispatchID, errors.Join(outcome.err, finalizeErr)) + } + return agenthooks.Response{}, dispatchID, newDispatchError(ResultInternalError, dispatchID, finalizeErr) + } + return response, dispatchID, newDispatchError(outcome.result, dispatchID, outcome.err) +} + +func (d *Dispatcher) finishWithoutPost( + ctx context.Context, + event Event, + dispatchID uuid.UUID, + startedAt time.Time, + result DispatchResult, + dispatchErr error, +) (agenthooks.Response, error) { + persistCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), finalizeTimeout) + defer cancel() + if err := d.insert(persistCtx, event, dispatchID, startedAt); err != nil { + return agenthooks.Response{}, newDispatchError(result, dispatchID, errors.Join(dispatchErr, xerrors.Errorf("insert chat hook dispatch: %w", err))) + } + outcome := dispatchOutcome{result: result, err: dispatchErr} + finalizeErr := d.finalize(ctx, event, dispatchID, outcome) + d.metrics.observe(event.Type, result, agenthooks.Response{}, time.Since(startedAt)) + if finalizeErr != nil { + return agenthooks.Response{}, newDispatchError(result, dispatchID, errors.Join(dispatchErr, finalizeErr)) + } + return agenthooks.Response{}, newDispatchError(result, dispatchID, dispatchErr) +} + +func (d *Dispatcher) insert(ctx context.Context, event Event, dispatchID uuid.UUID, startedAt time.Time) error { + toolUseID, toolName := event.toolMetadata() + _, err := d.db.InsertChatHookDispatch(ctx, database.InsertChatHookDispatchParams{ + ID: dispatchID, + ChatID: event.ChatID, + Event: string(event.Type), + TurnID: nullUUID(event.TurnID), + ToolUseID: nullStringPtr(toolUseID), + ToolName: nullStringPtr(toolName), + OwnerID: event.OwnerID, + WorkspaceID: nullUUID(event.WorkspaceID), + StartedAt: startedAt, + }) + return err +} + +type dispatchOutcome struct { + result DispatchResult + httpStatus sql.NullInt32 + response agenthooks.Response + original pqtype.NullRawMessage + err error +} + +func (d *Dispatcher) prepareAndPost(ctx context.Context, event Event, dispatchID uuid.UUID) (agenthooks.Response, dispatchOutcome) { + data, original, err := marshalEventData(event) + if err != nil { + return agenthooks.Response{}, dispatchOutcome{result: ResultProtocolError, original: original, err: err} + } + request := agenthooks.Request{ + Type: event.Type, + Meta: agenthooks.Meta{ + DispatchID: dispatchID, + SchemaVersion: agenthooks.SchemaVersion, + ChatRef: event.ChatRef, + }, + Data: data, + } + body, err := json.Marshal(request) + if err != nil { + return agenthooks.Response{}, dispatchOutcome{result: ResultProtocolError, original: original, err: xerrors.Errorf("marshal request: %w", err)} + } + + digest := sha256.Sum256(body) + now := time.Now() + token, err := agenthooks.SignClaims(d.secret, agenthooks.Claims{ + Issuer: d.deploymentID, + Subject: "coder:chat:" + event.ChatID.String(), + Audience: d.hookURL, + IssuedAt: now.Unix(), + // Backdate nbf to tolerate consumer clock skew. + NotBefore: now.Add(-clockSkewLeeway).Unix(), + Expires: now.Add(d.timeout + clockSkewLeeway).Unix(), + JTI: dispatchID, + Type: event.Type, + BodySHA256: hex.EncodeToString(digest[:]), + }) + if err != nil { + return agenthooks.Response{}, dispatchOutcome{result: ResultProtocolError, original: original, err: xerrors.Errorf("sign request: %w", err)} + } + + response, status, result, err := d.post(ctx, body, token) + outcome := dispatchOutcome{ + result: result, + httpStatus: status, + response: response, + original: original, + err: err, + } + if err != nil { + return agenthooks.Response{}, outcome + } + if err := validateResponse(event.Type, response); err != nil { + outcome.result = ResultProtocolError + outcome.err = err + return agenthooks.Response{}, outcome + } + if response.Permission != nil && response.Permission.Decision == agenthooks.PermissionDeny { + outcome.result = ResultDenied + } + return response, outcome +} + +func (d *Dispatcher) post( + ctx context.Context, + body []byte, + token string, +) (response agenthooks.Response, status sql.NullInt32, result DispatchResult, err error) { + for attempt := range 2 { + attemptCtx, cancel := context.WithTimeout(ctx, d.timeout) + req, reqErr := http.NewRequestWithContext(attemptCtx, http.MethodPost, d.hookURL, bytes.NewReader(body)) + if reqErr != nil { + cancel() + return agenthooks.Response{}, sql.NullInt32{}, ResultProtocolError, xerrors.Errorf("create request: %w", reqErr) + } + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("User-Agent", d.userAgent) + + httpResponse, requestErr := d.client.Do(req) + if requestErr != nil { + attemptErr := attemptCtx.Err() + cancel() + if isTimeoutError(attemptErr) || isTimeoutError(requestErr) || errors.Is(requestErr, context.Canceled) { + return agenthooks.Response{}, sql.NullInt32{}, ResultTimeout, xerrors.Errorf("post lifecycle hook: %w", requestErr) + } + if !isConnectionError(requestErr) { + return agenthooks.Response{}, sql.NullInt32{}, ResultProtocolError, xerrors.Errorf("post lifecycle hook: %w", requestErr) + } + if attempt == 1 { + return agenthooks.Response{}, sql.NullInt32{}, ResultConnectionError, xerrors.Errorf("post lifecycle hook: %w", requestErr) + } + backoff := time.NewTimer(retryBackoff) + select { + case <-ctx.Done(): + backoff.Stop() + return agenthooks.Response{}, sql.NullInt32{}, ResultTimeout, xerrors.Errorf("post lifecycle hook: %w", ctx.Err()) + case <-backoff.C: + } + continue + } + + statusCode := int64(httpResponse.StatusCode) + if statusCode < math.MinInt32 || statusCode > math.MaxInt32 { + _ = httpResponse.Body.Close() + cancel() + return agenthooks.Response{}, sql.NullInt32{}, ResultProtocolError, xerrors.Errorf("lifecycle hook returned invalid HTTP status %d", httpResponse.StatusCode) + } + status = sql.NullInt32{Int32: int32(statusCode), Valid: true} + if httpResponse.StatusCode < http.StatusOK || httpResponse.StatusCode >= http.StatusMultipleChoices { + _ = httpResponse.Body.Close() + cancel() + return agenthooks.Response{}, status, ResultHTTPError, xerrors.Errorf("lifecycle hook returned HTTP status %d", httpResponse.StatusCode) + } + + responseBody, readErr := io.ReadAll(io.LimitReader(httpResponse.Body, maxResponseBodyBytes+1)) + attemptErr := attemptCtx.Err() + _ = httpResponse.Body.Close() + cancel() + if readErr != nil { + switch { + case isTimeoutError(attemptErr), isTimeoutError(readErr), errors.Is(readErr, context.Canceled): + return agenthooks.Response{}, status, ResultTimeout, xerrors.Errorf("read lifecycle hook response: %w", readErr) + case isConnectionError(readErr): + if attempt == 1 { + return agenthooks.Response{}, status, ResultConnectionError, xerrors.Errorf("read lifecycle hook response: %w", readErr) + } + default: + return agenthooks.Response{}, status, ResultProtocolError, xerrors.Errorf("read lifecycle hook response: %w", readErr) + } + // Mid-body connection drops get the same single retry as dial + // failures, reusing the dispatch ID. + backoff := time.NewTimer(retryBackoff) + select { + case <-ctx.Done(): + backoff.Stop() + return agenthooks.Response{}, status, ResultTimeout, xerrors.Errorf("read lifecycle hook response: %w", ctx.Err()) + case <-backoff.C: + } + continue + } + if len(responseBody) > maxResponseBodyBytes { + return agenthooks.Response{}, status, ResultProtocolError, xerrors.New("lifecycle hook response exceeds 1 MiB") + } + trimmed := bytes.TrimSpace(responseBody) + if len(trimmed) == 0 { + return agenthooks.Response{}, status, ResultOK, nil + } + if bytes.Equal(trimmed, []byte("null")) { + return agenthooks.Response{}, status, ResultProtocolError, xerrors.New("lifecycle hook response must be a JSON object") + } + if err := json.Unmarshal(trimmed, &response); err != nil { + return agenthooks.Response{}, status, ResultProtocolError, xerrors.Errorf("decode lifecycle hook response: %w", err) + } + return response, status, ResultOK, nil + } + panic("unreachable") +} + +func validateResponse(eventType agenthooks.EventType, response agenthooks.Response) error { + if len(response.ModelContext) > maxModelContextBytes { + return xerrors.New("model_context exceeds 16 KiB") + } + if response.Permission == nil { + return nil + } + if eventType != agenthooks.EventUserPromptSubmit && eventType != agenthooks.EventPreToolUse { + return xerrors.Errorf("permission is not valid for event %q", eventType) + } + + switch response.Permission.Decision { + case agenthooks.PermissionAllow: + inputOverride := bytes.TrimSpace(response.Permission.InputOverride) + if len(inputOverride) == 0 || bytes.Equal(inputOverride, []byte("null")) { + return xerrors.New("allow decision requires input_override") + } + if eventType == agenthooks.EventUserPromptSubmit { + if err := validateUserPromptSubmitOverride(inputOverride); err != nil { + return err + } + } + case agenthooks.PermissionDeny: + // A persisted deny override would poison decision reuse, which + // matches future inputs against original_input OR input_override. + inputOverride := bytes.TrimSpace(response.Permission.InputOverride) + if len(inputOverride) > 0 && !bytes.Equal(inputOverride, []byte("null")) { + return xerrors.New("deny decision must not include input_override") + } + case agenthooks.PermissionAsk: + return xerrors.New("ask decision is not supported") + default: + return xerrors.Errorf("invalid permission decision %q", response.Permission.Decision) + } + return nil +} + +func validateUserPromptSubmitOverride(input json.RawMessage) error { + var override struct { + Prompt *string `json:"prompt"` + } + decoder := json.NewDecoder(bytes.NewReader(input)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&override); err != nil { + return xerrors.Errorf("user_prompt_submit input_override must be {\"prompt\": string}: %w", err) + } + if override.Prompt == nil { + return xerrors.New("user_prompt_submit input_override must be {\"prompt\": string}") + } + if err := decoder.Decode(&struct{}{}); err != io.EOF { + return xerrors.New("user_prompt_submit input_override must contain one JSON object") + } + return nil +} + +func marshalEventData(event Event) (data json.RawMessage, original pqtype.NullRawMessage, err error) { + switch event.Type { + case agenthooks.EventSessionStart: + if !isData[agenthooks.SessionStartData](event.Data) { + return nil, original, xerrors.New("session_start data has the wrong type") + } + case agenthooks.EventUserPromptSubmit: + value, ok := dataValue[agenthooks.UserPromptSubmitData](event.Data) + if !ok { + return nil, original, xerrors.New("user_prompt_submit data has the wrong type") + } + encoded, marshalErr := json.Marshal(value.Prompt) + if marshalErr != nil { + return nil, original, xerrors.Errorf("marshal original prompt: %w", marshalErr) + } + original = pqtype.NullRawMessage{RawMessage: encoded, Valid: true} + case agenthooks.EventPreToolUse: + value, ok := dataValue[agenthooks.PreToolUseData](event.Data) + if !ok { + return nil, original, xerrors.New("pre_tool_use data has the wrong type") + } + if value.ToolUseID == "" || value.ToolName == "" { + return nil, original, xerrors.New("pre_tool_use data requires tool_use_id and tool_name") + } + // original_input is a jsonb column, so malformed model output + // cannot be persisted there. Leave it NULL so the dispatch can + // still finalize as a protocol error instead of staying pending. + if json.Valid(value.ToolInput) { + original = pqtype.NullRawMessage{RawMessage: bytes.Clone(value.ToolInput), Valid: true} + } + case agenthooks.EventPostToolUse: + value, ok := dataValue[agenthooks.PostToolUseData](event.Data) + if !ok { + return nil, original, xerrors.New("post_tool_use data has the wrong type") + } + if value.ToolUseID == "" || value.ToolName == "" { + return nil, original, xerrors.New("post_tool_use data requires tool_use_id and tool_name") + } + case agenthooks.EventPreCompact: + if !isData[agenthooks.PreCompactData](event.Data) { + return nil, original, xerrors.New("pre_compact data has the wrong type") + } + case agenthooks.EventPostCompact: + if !isData[agenthooks.PostCompactData](event.Data) { + return nil, original, xerrors.New("post_compact data has the wrong type") + } + case agenthooks.EventStop: + if !isData[agenthooks.StopData](event.Data) { + return nil, original, xerrors.New("stop data has the wrong type") + } + default: + return nil, original, xerrors.Errorf("unknown event type %q", event.Type) + } + + encoded, marshalErr := json.Marshal(event.Data) + if marshalErr != nil { + return nil, original, xerrors.Errorf("marshal event data: %w", marshalErr) + } + return encoded, original, nil +} + +func isData[T any](value any) bool { + _, ok := dataValue[T](value) + return ok +} + +func dataValue[T any](value any) (T, bool) { + if typed, ok := value.(T); ok { + return typed, true + } + if typed, ok := value.(*T); ok && typed != nil { + return *typed, true + } + var zero T + return zero, false +} + +func (d *Dispatcher) finalize(ctx context.Context, event Event, dispatchID uuid.UUID, outcome dispatchOutcome) error { + response := outcome.response + if outcome.err != nil { + d.logger.Warn(context.WithoutCancel(ctx), "chat hook dispatch failed", + slog.F("dispatch_id", dispatchID), + slog.F("event", event.Type), + slog.F("result", outcome.result), + slog.Error(outcome.err), + ) + } + // Only accepted outcomes may persist reusable decisions. + accepted := outcome.result == ResultOK || outcome.result == ResultDenied + var decision sql.NullString + var inputOverride pqtype.NullRawMessage + if response.Permission != nil && accepted { + decision = sql.NullString{String: string(response.Permission.Decision), Valid: true} + // A deny override is at most JSON null (validateResponse rejects the + // rest), and persisting that null would let decision reuse match + // future null tool inputs against it. Persist allow overrides only. + if response.Permission.Decision == agenthooks.PermissionAllow && response.Permission.InputOverride != nil { + inputOverride = pqtype.NullRawMessage{RawMessage: bytes.Clone(response.Permission.InputOverride), Valid: true} + } + } else if event.Type == agenthooks.EventPreToolUse && outcome.result == ResultOK { + decision = sql.NullString{String: string(agenthooks.PermissionAllow), Valid: true} + } + allowedTools := pqtype.NullRawMessage{} + if response.AllowedTools != nil { + encoded, err := json.Marshal(response.AllowedTools) + if err != nil { + return xerrors.Errorf("marshal allowed tools: %w", err) + } + allowedTools = pqtype.NullRawMessage{RawMessage: encoded, Valid: true} + } + + finalizeCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), finalizeTimeout) + defer cancel() + _, err := d.db.FinalizeChatHookDispatch(finalizeCtx, database.FinalizeChatHookDispatchParams{ + FinishedAt: time.Now(), + Result: string(outcome.result), + HttpStatus: outcome.httpStatus, + Decision: decision, + DecisionReason: nullPermissionReason(response.Permission), + InputOverride: inputOverride, + OriginalInput: outcome.original, + ModelContext: nullString(response.ModelContext), + UserMessage: nullString(response.UserMessage), + AllowedTools: allowedTools, + EndChat: sql.NullBool{Bool: response.EndChat, Valid: response.EndChat}, + Error: nullError(outcome.err), + ID: dispatchID, + ChatID: event.ChatID, + OwnerID: event.OwnerID, + }) + if err != nil { + return xerrors.Errorf("finalize chat hook dispatch: %w", err) + } + return nil +} + +func nullPermissionReason(permission *agenthooks.Permission) sql.NullString { + if permission == nil || permission.Reason == "" { + return sql.NullString{} + } + return sql.NullString{String: permission.Reason, Valid: true} +} + +func nullUUID(value *uuid.UUID) uuid.NullUUID { + if value == nil { + return uuid.NullUUID{} + } + return uuid.NullUUID{UUID: *value, Valid: true} +} + +func nullStringPtr(value *string) sql.NullString { + if value == nil { + return sql.NullString{} + } + return sql.NullString{String: *value, Valid: true} +} + +func nullString(value string) sql.NullString { + return sql.NullString{String: value, Valid: value != ""} +} + +func nullError(err error) sql.NullString { + if err == nil { + return sql.NullString{} + } + return sql.NullString{String: err.Error(), Valid: true} +} + +func isTimeoutError(err error) bool { + if err == nil { + return false + } + if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, syscall.ETIMEDOUT) { + return true + } + var netErr net.Error + return errors.As(err, &netErr) && netErr.Timeout() +} + +func isConnectionError(err error) bool { + if errors.Is(err, io.EOF) || errors.Is(err, io.ErrUnexpectedEOF) || + errors.Is(err, net.ErrClosed) || errors.Is(err, syscall.ECONNRESET) || + errors.Is(err, syscall.ECONNREFUSED) || errors.Is(err, syscall.EPIPE) { + return true + } + var opErr *net.OpError + return errors.As(err, &opErr) +} + +type metrics struct { + dispatches *prometheus.CounterVec + duration *prometheus.HistogramVec + decisions *prometheus.CounterVec + contextSize *prometheus.HistogramVec + inputOverrides *prometheus.CounterVec +} + +func newMetrics(reg prometheus.Registerer) *metrics { + if reg == nil { + reg = prometheus.NewRegistry() + } + factory := promauto.With(reg) + return &metrics{ + dispatches: factory.NewCounterVec(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "hook_dispatches_total", + Help: "Total lifecycle hook dispatches by event and result.", + }, []string{"event", "result"}), + duration: factory.NewHistogramVec(prometheus.HistogramOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "hook_dispatch_seconds", + Help: "Lifecycle hook dispatch duration in seconds.", + }, []string{"event"}), + decisions: factory.NewCounterVec(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "hook_decisions_total", + Help: "Total lifecycle hook permission decisions by event and decision.", + }, []string{"event", "decision"}), + contextSize: factory.NewHistogramVec(prometheus.HistogramOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "hook_context_size_bytes", + Help: "Lifecycle hook model context response size in bytes.", + Buckets: prometheus.ExponentialBuckets(64, 2, 10), + }, []string{"event"}), + inputOverrides: factory.NewCounterVec(prometheus.CounterOpts{ + Namespace: "coderd", + Subsystem: "chatd", + Name: "hook_input_overrides_total", + Help: "Total lifecycle hook input overrides by event.", + }, []string{"event"}), + } +} + +func (m *metrics) observe(eventType agenthooks.EventType, result DispatchResult, response agenthooks.Response, duration time.Duration) { + event := string(eventType) + m.dispatches.WithLabelValues(event, string(result)).Inc() + m.duration.WithLabelValues(event).Observe(duration.Seconds()) + if response.ModelContext != "" { + m.contextSize.WithLabelValues(event).Observe(float64(len(response.ModelContext))) + } + if response.Permission == nil { + return + } + switch response.Permission.Decision { + case agenthooks.PermissionAllow, agenthooks.PermissionDeny, agenthooks.PermissionAsk: + m.decisions.WithLabelValues(event, string(response.Permission.Decision)).Inc() + } + if response.Permission.Decision == agenthooks.PermissionAllow && response.Permission.InputOverride != nil { + m.inputOverrides.WithLabelValues(event).Inc() + } +} diff --git a/coderd/x/chathooks/dispatcher_internal_test.go b/coderd/x/chathooks/dispatcher_internal_test.go new file mode 100644 index 0000000000000..c1d595aa52321 --- /dev/null +++ b/coderd/x/chathooks/dispatcher_internal_test.go @@ -0,0 +1,600 @@ +package chathooks + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" + + "github.com/google/uuid" + "github.com/prometheus/client_golang/prometheus" + "github.com/stretchr/testify/assert" + "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/codersdk/agenthooks" + "github.com/coder/coder/v2/testutil" +) + +const ( + testSecret = "test-hook-secret-32-bytes-minimum!!" + testDeploymentID = "test-deployment" + testVersion = "test-version" +) + +func TestDispatcherSuccess(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventSessionStart, agenthooks.SessionStartData{Source: "new"}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + assert.NoError(t, err) + assert.Equal(t, "application/json", r.Header.Get("Content-Type")) + assert.Equal(t, "coderd/"+testVersion, r.Header.Get("User-Agent")) + + claims, err := agenthooks.Verify(r.Header.Get("Authorization"), []byte(testSecret)) + assert.NoError(t, err) + assert.Equal(t, testDeploymentID, claims.Issuer) + assert.Equal(t, serverURL(r), claims.Audience) + assert.Equal(t, event.Type, claims.Type) + chatID, err := claims.ChatID() + assert.NoError(t, err) + assert.Equal(t, event.ChatID, chatID) + digest := sha256.Sum256(body) + assert.Equal(t, hex.EncodeToString(digest[:]), claims.BodySHA256) + assert.Equal(t, claims.IssuedAt-int64(clockSkewLeeway/time.Second), claims.NotBefore) + + var request agenthooks.Request + assert.NoError(t, json.Unmarshal(body, &request)) + assert.Equal(t, claims.JTI, request.Meta.DispatchID) + assert.Equal(t, agenthooks.SchemaVersion, request.Meta.SchemaVersion) + decoded, err := request.Decode() + assert.NoError(t, err) + assert.Equal(t, &agenthooks.SessionStartData{Source: "new"}, decoded) + + rows, err := db.ListChatHookDispatchesByChatID(r.Context(), event.ChatID) + assert.NoError(t, err) + if assert.Len(t, rows, 1) { + assert.Equal(t, "pending", rows[0].Result) + assert.False(t, rows[0].FinishedAt.Valid) + } + _, err = w.Write([]byte(`{}`)) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + dispatcher := newTestDispatcher(t, db, server.Client(), server.URL, 2*time.Second) + response, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event) + require.NoError(t, err) + require.Equal(t, agenthooks.Response{}, response) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultOK), row.Result) + require.True(t, row.FinishedAt.Valid) + require.Equal(t, int32(http.StatusOK), row.HttpStatus.Int32) + require.False(t, row.Error.Valid) +} + +func TestDispatcherDeny(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventUserPromptSubmit, agenthooks.UserPromptSubmitData{Prompt: "delete everything"}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write([]byte(`{"permission":{"decision":"deny","reason":"blocked"},"user_message":"not allowed"}`)) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + response, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.NoError(t, err) + require.NotNil(t, response.Permission) + require.Equal(t, agenthooks.PermissionDeny, response.Permission.Decision) + require.Equal(t, "not allowed", response.UserMessage) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultDenied), row.Result) + require.Equal(t, string(agenthooks.PermissionDeny), row.Decision.String) + require.Equal(t, "blocked", row.DecisionReason.String) + require.Equal(t, "not allowed", row.UserMessage.String) + require.JSONEq(t, `"delete everything"`, string(row.OriginalInput.RawMessage)) +} + +func TestDispatcherDenyNullOverrideNotPersisted(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventPreToolUse, agenthooks.PreToolUseData{ + ToolUseID: "tool-use-1", + ToolName: "execute", + ToolInput: json.RawMessage(`{"cmd":"rm"}`), + }) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write([]byte(`{"permission":{"decision":"deny","input_override":null}}`)) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + response, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.NoError(t, err) + require.NotNil(t, response.Permission) + require.Equal(t, agenthooks.PermissionDeny, response.Permission.Decision) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultDenied), row.Result) + require.Equal(t, string(agenthooks.PermissionDeny), row.Decision.String) + require.False(t, row.InputOverride.Valid) +} + +func TestDispatcherAllowInputOverride(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + toolInput := json.RawMessage(`{"path":"before"}`) + toolUseID := "call_" + uuid.NewString() + event := newTestEvent(t, db, agenthooks.EventPreToolUse, agenthooks.PreToolUseData{ + ToolUseID: toolUseID, + ToolName: "edit", + ToolInput: toolInput, + }) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write([]byte(`{"permission":{"decision":"allow","input_override":{"path":"after"}}}`)) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + response, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.NoError(t, err) + require.NotNil(t, response.Permission) + require.Equal(t, agenthooks.PermissionAllow, response.Permission.Decision) + require.JSONEq(t, `{"path":"after"}`, string(response.Permission.InputOverride)) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultOK), row.Result) + require.Equal(t, toolUseID, row.ToolUseID.String) + require.JSONEq(t, `{"path":"after"}`, string(row.InputOverride.RawMessage)) + require.JSONEq(t, `{"path":"before"}`, string(row.OriginalInput.RawMessage)) +} + +func TestDispatcherTimeoutNoRetry(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventStop, agenthooks.StopData{}) + var requests atomic.Int32 + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusOK) + assert.NoError(t, http.NewResponseController(w).Flush()) + <-release + })) + t.Cleanup(server.Close) + + _, _, err := newTestDispatcher(t, db, server.Client(), server.URL, 50*time.Millisecond).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + close(release) + require.Error(t, err) + require.Equal(t, int32(1), requests.Load()) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultTimeout), row.Result) + require.Equal(t, int32(http.StatusOK), row.HttpStatus.Int32) + require.True(t, row.Error.Valid) +} + +func TestDispatcherRetriesConnectionErrorWithSameJTI(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventPostCompact, agenthooks.PostCompactData{}) + claimsCh := make(chan agenthooks.Claims, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + claims, err := agenthooks.Verify(r.Header.Get("Authorization"), []byte(testSecret)) + assert.NoError(t, err) + claimsCh <- claims + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(server.Close) + + var attempts atomic.Int32 + baseTransport := server.Client().Transport + client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + claims, err := agenthooks.Verify(req.Header.Get("Authorization"), []byte(testSecret)) + if err != nil { + return nil, err + } + if attempts.Add(1) == 1 { + claimsCh <- claims + _, err = io.Copy(io.Discard, req.Body) + if err != nil { + return nil, err + } + return nil, io.EOF + } + return baseTransport.RoundTrip(req) + })} + + _, _, err := newTestDispatcher(t, db, client, server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.NoError(t, err) + require.Equal(t, int32(2), attempts.Load()) + first := <-claimsCh + second := <-claimsCh + require.Equal(t, first.JTI, second.JTI) + chatID, err := first.ChatID() + require.NoError(t, err) + require.Equal(t, event.ChatID, chatID) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultOK), row.Result) + require.Equal(t, first.JTI, row.ID) +} + +func TestDispatcherRetriesMidBodyConnectionError(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventPostCompact, agenthooks.PostCompactData{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusNoContent) + })) + t.Cleanup(server.Close) + + var attempts atomic.Int32 + baseTransport := server.Client().Transport + client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + if attempts.Add(1) == 1 { + _, err := io.Copy(io.Discard, req.Body) + if err != nil { + return nil, err + } + // A response whose body errors mid-read simulates the + // connection dropping after headers were received. + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(errReader{err: io.ErrUnexpectedEOF}), + Header: http.Header{}, + }, nil + } + return baseTransport.RoundTrip(req) + })} + + _, _, err := newTestDispatcher(t, db, client, server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.NoError(t, err) + require.Equal(t, int32(2), attempts.Load()) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultOK), row.Result) +} + +type errReader struct{ err error } + +func (r errReader) Read([]byte) (int, error) { return 0, r.err } + +func TestDispatcherTLSFailureNoRetry(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventStop, agenthooks.StopData{}) + server := httptest.NewTLSServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + t.Fatal("request with an untrusted certificate reached the handler") + })) + t.Cleanup(server.Close) + + transport := http.DefaultTransport.(*http.Transport).Clone() + t.Cleanup(transport.CloseIdleConnections) + var attempts atomic.Int32 + client := &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + attempts.Add(1) + return transport.RoundTrip(req) + })} + + _, _, err := newTestDispatcher(t, db, client, server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.Error(t, err) + require.Equal(t, int32(1), attempts.Load()) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultProtocolError), row.Result) +} + +func TestDispatcherNon2xxNoRetry(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventPreCompact, agenthooks.PreCompactData{}) + var requests atomic.Int32 + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + requests.Add(1) + w.WriteHeader(http.StatusServiceUnavailable) + assert.NoError(t, http.NewResponseController(w).Flush()) + <-release + })) + t.Cleanup(server.Close) + + _, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + close(release) + require.Error(t, err) + require.Equal(t, int32(1), requests.Load()) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultHTTPError), row.Result) + require.Equal(t, int32(http.StatusServiceUnavailable), row.HttpStatus.Int32) +} + +func TestDispatcherProtocolErrors(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + eventType agenthooks.EventType + data any + responseBody []byte + assertRow func(*testing.T, database.ChatHookDispatch) + }{ + { + name: "malformed JSON", + eventType: agenthooks.EventStop, + data: agenthooks.StopData{}, + responseBody: []byte(`{"end_chat":`), + }, + { + name: "oversized model context", + eventType: agenthooks.EventStop, + data: agenthooks.StopData{}, + responseBody: mustJSON(t, agenthooks.Response{ModelContext: string(bytes.Repeat([]byte("x"), maxModelContextBytes+1))}), + assertRow: func(t *testing.T, row database.ChatHookDispatch) { + require.True(t, row.ModelContext.Valid) + require.Len(t, row.ModelContext.String, maxModelContextBytes+1) + }, + }, + { + name: "invalid user prompt override shape", + eventType: agenthooks.EventUserPromptSubmit, + data: agenthooks.UserPromptSubmitData{Prompt: "question"}, + responseBody: mustJSON(t, agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAllow, + InputOverride: json.RawMessage(`{"unexpected":"value"}`), + }}), + }, + { + name: "deny with input override", + eventType: agenthooks.EventPreToolUse, + data: agenthooks.PreToolUseData{ + ToolUseID: "call_deny_override", + ToolName: "edit", + ToolInput: json.RawMessage(`{"path":"a"}`), + }, + responseBody: mustJSON(t, agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionDeny, + InputOverride: json.RawMessage(`{"path":"b"}`), + }}), + }, + { + name: "ask decision", + eventType: agenthooks.EventUserPromptSubmit, + data: agenthooks.UserPromptSubmitData{Prompt: "question"}, + responseBody: mustJSON(t, agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAsk, + }}), + assertRow: func(t *testing.T, row database.ChatHookDispatch) { + require.False(t, row.Decision.Valid, "rejected responses must not persist a reusable decision") + }, + }, + { + name: "pre_tool_use allow without input_override", + eventType: agenthooks.EventPreToolUse, + data: agenthooks.PreToolUseData{ToolUseID: "call_no_override", ToolName: "run_command", ToolInput: json.RawMessage(`{"cmd":"ls"}`)}, + responseBody: mustJSON(t, agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAllow, + }}), + assertRow: func(t *testing.T, row database.ChatHookDispatch) { + require.False(t, row.Decision.Valid, "rejected responses must not persist a reusable decision") + }, + }, + { + name: "pre_tool_use without tool_use_id", + eventType: agenthooks.EventPreToolUse, + data: agenthooks.PreToolUseData{ToolName: "run_command", ToolInput: json.RawMessage(`{"cmd":"ls"}`)}, + responseBody: []byte(`{}`), + }, + { + name: "post_tool_use without tool_name", + eventType: agenthooks.EventPostToolUse, + data: agenthooks.PostToolUseData{ToolUseID: "call_no_name"}, + responseBody: []byte(`{}`), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, test.eventType, test.data) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, err := w.Write(test.responseBody) + assert.NoError(t, err) + })) + t.Cleanup(server.Close) + + _, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.Error(t, err) + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultProtocolError), row.Result) + require.True(t, row.Error.Valid) + if test.assertRow != nil { + test.assertRow(t, row) + } + }) + } +} + +func TestDispatcherOverCapacity(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventStop, agenthooks.StopData{}) + dispatcher := newTestDispatcher(t, db, nil, "http://unused.test", 10*time.Millisecond) + for range maxConcurrentDispatches { + dispatcher.semaphore <- struct{}{} + } + defer func() { + for range maxConcurrentDispatches { + <-dispatcher.semaphore + } + }() + + _, _, err := dispatcher.Dispatch(testutil.Context(t, testutil.WaitLong), event) + require.ErrorIs(t, err, context.DeadlineExceeded) + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultOverCapacity), row.Result) + require.False(t, row.HttpStatus.Valid) + require.True(t, row.FinishedAt.Valid) +} + +func TestDispatcherCanceledContextPersistsTimeout(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + event := newTestEvent(t, db, agenthooks.EventStop, agenthooks.StopData{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(server.Close) + + // Whichever select branch wins with an already-canceled context, a + // timeout row must be persisted. + ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitLong)) + cancel() + _, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch(ctx, event) + require.Error(t, err) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultTimeout), row.Result) + require.True(t, row.FinishedAt.Valid) +} + +func TestDispatcherInvalidToolInputFinalizesProtocolError(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + toolUseID := "call_" + uuid.NewString() + event := newTestEvent(t, db, agenthooks.EventPreToolUse, agenthooks.PreToolUseData{ + ToolUseID: toolUseID, + ToolName: "edit", + ToolInput: json.RawMessage(`{"path":`), + }) + + var hookRequests atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + hookRequests.Add(1) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(server.Close) + + _, _, err := newTestDispatcher(t, db, server.Client(), server.URL, time.Second).Dispatch( + testutil.Context(t, testutil.WaitLong), event, + ) + require.Error(t, err) + require.Zero(t, hookRequests.Load()) + + row := singleDispatch(t, db, event.ChatID) + require.Equal(t, string(ResultProtocolError), row.Result) + require.True(t, row.FinishedAt.Valid, "dispatch must finalize instead of staying pending") + require.False(t, row.OriginalInput.Valid, "malformed input must not persist as jsonb") +} + +func newTestDispatcher( + t *testing.T, + db database.Store, + client *http.Client, + hookURL string, + timeout time.Duration, +) *Dispatcher { + t.Helper() + return New( + testutil.Logger(t), + db, + client, + hookURL, + testSecret, + timeout, + testDeploymentID, + testVersion, + prometheus.NewRegistry(), + ) +} + +func newTestEvent(t *testing.T, db database.Store, eventType agenthooks.EventType, data any) Event { + t.Helper() + user := dbgen.User(t, db, database.User{}) + organization := dbgen.Organization(t, db, database.Organization{}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: organization.ID, + OwnerID: user.ID, + LastModelConfigID: model.ID, + }) + return Event{ + Type: eventType, + ChatRef: agenthooks.ChatRef{ + ChatID: chat.ID, + OwnerID: user.ID, + }, + Data: data, + } +} + +func singleDispatch(t *testing.T, db database.Store, chatID uuid.UUID) database.ChatHookDispatch { + t.Helper() + rows, err := db.ListChatHookDispatchesByChatID(testutil.Context(t, testutil.WaitLong), chatID) + require.NoError(t, err) + require.Len(t, rows, 1) + return rows[0] +} + +func mustJSON(t *testing.T, value any) []byte { + t.Helper() + encoded, err := json.Marshal(value) + require.NoError(t, err) + return encoded +} + +func serverURL(r *http.Request) string { + return "http://" + r.Host +} + +type roundTripFunc func(*http.Request) (*http.Response, error) + +func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} diff --git a/codersdk/agenthooks/agenthooks_test.go b/codersdk/agenthooks/agenthooks_test.go new file mode 100644 index 0000000000000..a16be407d6b77 --- /dev/null +++ b/codersdk/agenthooks/agenthooks_test.go @@ -0,0 +1,422 @@ +package agenthooks_test + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/codersdk/agenthooks" +) + +var testSecret = []byte("0123456789abcdef0123456789abcdef") + +func TestSignClaimsVerify(t *testing.T) { + t.Parallel() + + claims := validClaims(t, "https://hooks.example.com/coder", agenthooks.EventPreToolUse, nil) + token, err := agenthooks.SignClaims(testSecret, claims) + require.NoError(t, err) + + got, err := agenthooks.Verify("Bearer "+token, testSecret) + require.NoError(t, err) + require.Equal(t, claims, got) +} + +func TestVerifyRejectsAlgorithmConfusion(t *testing.T) { + t.Parallel() + + claims := validClaims(t, "https://hooks.example.com/coder", agenthooks.EventStop, nil) + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.HS512, Key: bytes.Repeat([]byte{1}, 64)}, + new(jose.SignerOptions).WithType("JWT"), + ) + require.NoError(t, err) + payload, err := json.Marshal(claims) + require.NoError(t, err) + signed, err := signer.Sign(payload) + require.NoError(t, err) + token, err := signed.CompactSerialize() + require.NoError(t, err) + + _, err = agenthooks.Verify("Bearer "+token, testSecret) + require.Error(t, err) +} + +func TestVerifyTimeBounds(t *testing.T) { + t.Parallel() + + now := time.Now() + tests := []struct { + name string + update func(*agenthooks.Claims) + }{ + { + name: "expired", + update: func(claims *agenthooks.Claims) { + claims.Expires = now.Add(-time.Minute).Unix() + }, + }, + { + name: "not before", + update: func(claims *agenthooks.Claims) { + claims.NotBefore = now.Add(time.Minute).Unix() + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + claims := validClaims(t, "https://hooks.example.com/coder", agenthooks.EventStop, nil) + test.update(&claims) + token, err := agenthooks.SignClaims(testSecret, claims) + require.NoError(t, err) + + _, err = agenthooks.Verify("Bearer "+token, testSecret) + require.Error(t, err) + }) + } +} + +func TestHTTPHandlerRoutesEvents(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + event agenthooks.EventType + data any + install func(*agenthooks.Hooks, *bool) + }{ + { + name: "session start", + event: agenthooks.EventSessionStart, + data: agenthooks.SessionStartData{Source: "startup"}, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.SessionStart = func(_ context.Context, _ agenthooks.Meta, data agenthooks.SessionStartData) (agenthooks.Response, error) { + *called = true + require.Equal(t, "startup", data.Source) + return agenthooks.Response{UserMessage: "session start"}, nil + } + }, + }, + { + name: "user prompt submit", + event: agenthooks.EventUserPromptSubmit, + data: agenthooks.UserPromptSubmitData{Prompt: "hello"}, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.UserPromptSubmit = func(_ context.Context, _ agenthooks.Meta, data agenthooks.UserPromptSubmitData) (agenthooks.Response, error) { + *called = true + require.Equal(t, "hello", data.Prompt) + return agenthooks.Response{UserMessage: "user prompt submit"}, nil + } + }, + }, + { + name: "pre tool use", + event: agenthooks.EventPreToolUse, + data: agenthooks.PreToolUseData{ + ToolUseID: "call_" + uuid.NewString(), + ToolName: "execute", + ToolInput: json.RawMessage(`{"command":"pwd"}`), + }, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.PreToolUse = func(_ context.Context, _ agenthooks.Meta, data agenthooks.PreToolUseData) (agenthooks.Response, error) { + *called = true + require.Equal(t, "execute", data.ToolName) + return agenthooks.Response{UserMessage: "pre tool use"}, nil + } + }, + }, + { + name: "post tool use", + event: agenthooks.EventPostToolUse, + data: agenthooks.PostToolUseData{ + ToolUseID: "call_" + uuid.NewString(), + ToolName: "execute", + ToolResponse: json.RawMessage(`{"output":"ok"}`), + }, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.PostToolUse = func(_ context.Context, _ agenthooks.Meta, data agenthooks.PostToolUseData) (agenthooks.Response, error) { + *called = true + require.Equal(t, "execute", data.ToolName) + return agenthooks.Response{UserMessage: "post tool use"}, nil + } + }, + }, + { + name: "pre compact", + event: agenthooks.EventPreCompact, + data: agenthooks.PreCompactData{}, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.PreCompact = func(context.Context, agenthooks.Meta, agenthooks.PreCompactData) (agenthooks.Response, error) { + *called = true + return agenthooks.Response{UserMessage: "pre compact"}, nil + } + }, + }, + { + name: "post compact", + event: agenthooks.EventPostCompact, + data: agenthooks.PostCompactData{}, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.PostCompact = func(context.Context, agenthooks.Meta, agenthooks.PostCompactData) (agenthooks.Response, error) { + *called = true + return agenthooks.Response{UserMessage: "post compact"}, nil + } + }, + }, + { + name: "stop", + event: agenthooks.EventStop, + data: agenthooks.StopData{}, + install: func(hooks *agenthooks.Hooks, called *bool) { + hooks.Stop = func(context.Context, agenthooks.Meta, agenthooks.StopData) (agenthooks.Response, error) { + *called = true + return agenthooks.Response{UserMessage: "stop"}, nil + } + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + called := false + var hooks agenthooks.Hooks + test.install(&hooks, &called) + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, hooks)) + t.Cleanup(server.Close) + + response := postEvent(t, server.URL, test.event, test.data, nil, nil) + defer response.Body.Close() + require.Equal(t, http.StatusOK, response.StatusCode) + var got agenthooks.Response + require.NoError(t, json.NewDecoder(response.Body).Decode(&got)) + require.Equal(t, test.name, got.UserMessage) + require.True(t, called) + }) + } +} + +func TestHTTPHandlerNoOpHookDoesNotDecodeData(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, agenthooks.Hooks{})) + t.Cleanup(server.Close) + response := postEvent(t, server.URL, agenthooks.EventStop, "unused", nil, nil) + defer response.Body.Close() + require.Equal(t, http.StatusOK, response.StatusCode) + var got agenthooks.Response + require.NoError(t, json.NewDecoder(response.Body).Decode(&got)) + require.Equal(t, agenthooks.Response{}, got) +} + +func TestHTTPHandlerRejectsMismatches(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + updateRequest func(*agenthooks.Request) + updateClaims func(*agenthooks.Claims) + }{ + { + name: "dispatch ID", + updateRequest: func(request *agenthooks.Request) { + request.Meta.DispatchID = uuid.New() + }, + }, + { + name: "event type", + updateRequest: func(request *agenthooks.Request) { + request.Type = agenthooks.EventPreCompact + }, + }, + { + name: "chat ID", + updateRequest: func(request *agenthooks.Request) { + request.Meta.ChatID = uuid.New() + }, + }, + { + name: "audience", + updateClaims: func(claims *agenthooks.Claims) { + claims.Audience = "https://hooks.example.com/other" + }, + }, + { + name: "body SHA-256", + updateClaims: func(claims *agenthooks.Claims) { + claims.BodySHA256 = hex.EncodeToString(bytes.Repeat([]byte{1}, sha256.Size)) + }, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, agenthooks.Hooks{})) + t.Cleanup(server.Close) + response := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, test.updateRequest, test.updateClaims) + defer response.Body.Close() + require.Equal(t, http.StatusBadRequest, response.StatusCode) + }) + } +} + +func TestHTTPHandlerExpectedIssuer(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler( + testSecret, + agenthooks.Hooks{}, + agenthooks.WithExpectedIssuer("deployment-a"), + )) + t.Cleanup(server.Close) + + matching := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, nil, func(claims *agenthooks.Claims) { + claims.Issuer = "deployment-a" + }) + defer matching.Body.Close() + require.Equal(t, http.StatusOK, matching.StatusCode) + + mismatched := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, nil, func(claims *agenthooks.Claims) { + claims.Issuer = "deployment-b" + }) + defer mismatched.Body.Close() + require.Equal(t, http.StatusUnauthorized, mismatched.StatusCode) +} + +func TestResponseAllowedToolsWireShape(t *testing.T) { + t.Parallel() + + unset, err := json.Marshal(agenthooks.Response{}) + require.NoError(t, err) + require.NotContains(t, string(unset), "allowed_tools") + + restrictAll, err := json.Marshal(agenthooks.Response{AllowedTools: &[]string{}}) + require.NoError(t, err) + require.Contains(t, string(restrictAll), `"allowed_tools":[]`) + + var decoded agenthooks.Response + require.NoError(t, json.Unmarshal([]byte(`{"allowed_tools":null}`), &decoded)) + require.Nil(t, decoded.AllowedTools) + require.NoError(t, json.Unmarshal([]byte(`{"allowed_tools":[]}`), &decoded)) + require.NotNil(t, decoded.AllowedTools) + require.Empty(t, *decoded.AllowedTools) +} + +func TestHTTPHandlerAcceptsTrailingSlashAudience(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, agenthooks.Hooks{})) + t.Cleanup(server.Close) + response := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, nil, func(claims *agenthooks.Claims) { + claims.Audience = server.URL + "/" + }) + defer response.Body.Close() + require.Equal(t, http.StatusOK, response.StatusCode) +} + +func TestHTTPHandlerHonorsForwardedProto(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, agenthooks.Hooks{})) + t.Cleanup(server.Close) + httpsAudience := "https" + strings.TrimPrefix(server.URL, "http") + response := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, nil, func(claims *agenthooks.Claims) { + claims.Audience = httpsAudience + }, func(r *http.Request) { + r.Header.Set("X-Forwarded-Proto", "https, http") + }) + defer response.Body.Close() + require.Equal(t, http.StatusOK, response.StatusCode) +} + +func TestHTTPHandlerHonorsForwardedHost(t *testing.T) { + t.Parallel() + + server := httptest.NewServer(agenthooks.NewHTTPHandler(testSecret, agenthooks.Hooks{})) + t.Cleanup(server.Close) + response := postEvent(t, server.URL, agenthooks.EventStop, agenthooks.StopData{}, nil, func(claims *agenthooks.Claims) { + claims.Audience = "https://hooks.example.com" + }, func(r *http.Request) { + r.Header.Set("X-Forwarded-Proto", "https") + r.Header.Set("X-Forwarded-Host", "hooks.example.com, internal-lb") + }) + defer response.Body.Close() + require.Equal(t, http.StatusOK, response.StatusCode) +} + +func postEvent(t *testing.T, target string, eventType agenthooks.EventType, data any, updateRequest func(*agenthooks.Request), updateClaims func(*agenthooks.Claims), updateHTTPRequest ...func(*http.Request)) *http.Response { + t.Helper() + + dataJSON, err := json.Marshal(data) + require.NoError(t, err) + request := agenthooks.Request{ + Type: eventType, + Meta: agenthooks.Meta{ + DispatchID: uuid.New(), + SchemaVersion: agenthooks.SchemaVersion, + ChatRef: agenthooks.ChatRef{ + ChatID: uuid.New(), + OwnerID: uuid.New(), + }, + }, + Data: dataJSON, + } + claims := validClaims(t, target, eventType, &request) + if updateRequest != nil { + updateRequest(&request) + } + body, err := json.Marshal(request) + require.NoError(t, err) + digest := sha256.Sum256(body) + claims.BodySHA256 = hex.EncodeToString(digest[:]) + if updateClaims != nil { + updateClaims(&claims) + } + token, err := agenthooks.SignClaims(testSecret, claims) + require.NoError(t, err) + httpRequest, err := http.NewRequestWithContext(t.Context(), http.MethodPost, target, bytes.NewReader(body)) + require.NoError(t, err) + httpRequest.Header.Set("Authorization", "Bearer "+token) + for _, update := range updateHTTPRequest { + update(httpRequest) + } + response, err := http.DefaultClient.Do(httpRequest) + require.NoError(t, err) + return response +} + +func validClaims(t *testing.T, audience string, eventType agenthooks.EventType, request *agenthooks.Request) agenthooks.Claims { + t.Helper() + + now := time.Now() + claims := agenthooks.Claims{ + Issuer: uuid.NewString(), + Subject: "coder:chat:" + uuid.NewString(), + Audience: audience, + IssuedAt: now.Unix(), + NotBefore: now.Add(-time.Second).Unix(), + Expires: now.Add(time.Minute).Unix(), + JTI: uuid.New(), + Type: eventType, + BodySHA256: hex.EncodeToString(make([]byte, sha256.Size)), + } + if request != nil { + claims.Subject = "coder:chat:" + request.Meta.ChatID.String() + claims.JTI = request.Meta.DispatchID + } + return claims +} diff --git a/codersdk/agenthooks/http.go b/codersdk/agenthooks/http.go new file mode 100644 index 0000000000000..3be6ee9f2e4e0 --- /dev/null +++ b/codersdk/agenthooks/http.go @@ -0,0 +1,193 @@ +package agenthooks + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "io" + "net/http" + "net/url" + "strings" + + "golang.org/x/xerrors" +) + +// Hooks lets a consumer implement only the lifecycle events it uses. +type Hooks struct { + SessionStart func(context.Context, Meta, SessionStartData) (Response, error) + UserPromptSubmit func(context.Context, Meta, UserPromptSubmitData) (Response, error) + PreToolUse func(context.Context, Meta, PreToolUseData) (Response, error) + PostToolUse func(context.Context, Meta, PostToolUseData) (Response, error) + PreCompact func(context.Context, Meta, PreCompactData) (Response, error) + PostCompact func(context.Context, Meta, PostCompactData) (Response, error) + Stop func(context.Context, Meta, StopData) (Response, error) +} + +// HandlerOption configures NewHTTPHandler. +type HandlerOption func(*handlerOptions) + +type handlerOptions struct { + expectedIssuer string +} + +// WithExpectedIssuer requires the verified iss claim to match issuer. +// If omitted, NewHTTPHandler accepts any non-empty issuer signed with the secret. +func WithExpectedIssuer(issuer string) HandlerOption { + return func(options *handlerOptions) { + options.expectedIssuer = issuer + } +} + +func NewHTTPHandler(secret []byte, hooks Hooks, opts ...HandlerOption) http.Handler { + var options handlerOptions + for _, opt := range opts { + opt(&options) + } + return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + rw.Header().Set("Allow", http.MethodPost) + http.Error(rw, "method not allowed", http.StatusMethodNotAllowed) + return + } + + claims, err := Verify(r.Header.Get("Authorization"), secret) + if err != nil { + http.Error(rw, err.Error(), http.StatusUnauthorized) + return + } + if options.expectedIssuer != "" && claims.Issuer != options.expectedIssuer { + http.Error(rw, "unexpected issuer", http.StatusUnauthorized) + return + } + body, err := io.ReadAll(r.Body) + if err != nil { + http.Error(rw, "read request body", http.StatusBadRequest) + return + } + var request Request + if err := json.Unmarshal(body, &request); err != nil { + http.Error(rw, "decode request body", http.StatusBadRequest) + return + } + if err := verifyBody(r, body, claims, request); err != nil { + http.Error(rw, err.Error(), http.StatusBadRequest) + return + } + response, err := dispatch(r.Context(), hooks, request) + if err != nil { + http.Error(rw, err.Error(), http.StatusInternalServerError) + return + } + + rw.Header().Set("Content-Type", "application/json") + if err := json.NewEncoder(rw).Encode(response); err != nil { + return + } + }) +} + +func verifyBody(r *http.Request, body []byte, claims Claims, request Request) error { + digest := sha256.Sum256(body) + if claims.BodySHA256 != hex.EncodeToString(digest[:]) { + return xerrors.New("request body does not match body_sha256 claim") + } + if canonicalAudience(claims.Audience) != requestAudience(r) { + return xerrors.New("request URL does not match audience claim") + } + if request.Meta.SchemaVersion != SchemaVersion { + return xerrors.New("unsupported schema version") + } + if request.Meta.DispatchID != claims.JTI { + return xerrors.New("dispatch ID does not match JWT ID") + } + if request.Type != claims.Type { + return xerrors.New("request type does not match type claim") + } + chatID, err := claims.ChatID() + if err != nil { + return err + } + if request.Meta.ChatID != chatID { + return xerrors.New("chat ID does not match subject claim") + } + return nil +} + +func requestAudience(r *http.Request) string { + requestURL := *r.URL + if requestURL.Scheme == "" { + requestURL.Scheme = "http" + if r.TLS != nil { + requestURL.Scheme = "https" + } + // Forwarded values reconstruct the signed client-facing audience. + if proto := forwardedProto(r); proto != "" { + requestURL.Scheme = proto + } + } + if requestURL.Host == "" { + requestURL.Host = r.Host + if host := forwardedHost(r); host != "" { + requestURL.Host = host + } + } + return canonicalAudience(requestURL.String()) +} + +func forwardedProto(r *http.Request) string { + proto := r.Header.Get("X-Forwarded-Proto") + // Proxies append values; the first is client-facing. + proto, _, _ = strings.Cut(proto, ",") + return strings.ToLower(strings.TrimSpace(proto)) +} + +func forwardedHost(r *http.Request) string { + host := r.Header.Get("X-Forwarded-Host") + host, _, _ = strings.Cut(host, ",") + return strings.TrimSpace(host) +} + +// canonicalAudience treats root URLs with and without a trailing slash as equivalent. +func canonicalAudience(audience string) string { + parsed, err := url.Parse(audience) + if err != nil { + return audience + } + if parsed.Path == "/" && parsed.RawPath == "" { + parsed.Path = "" + } + return parsed.String() +} + +func dispatch(ctx context.Context, hooks Hooks, request Request) (Response, error) { + switch request.Type { + case EventSessionStart: + return dispatchHook(ctx, request, hooks.SessionStart) + case EventUserPromptSubmit: + return dispatchHook(ctx, request, hooks.UserPromptSubmit) + case EventPreToolUse: + return dispatchHook(ctx, request, hooks.PreToolUse) + case EventPostToolUse: + return dispatchHook(ctx, request, hooks.PostToolUse) + case EventPreCompact: + return dispatchHook(ctx, request, hooks.PreCompact) + case EventPostCompact: + return dispatchHook(ctx, request, hooks.PostCompact) + case EventStop: + return dispatchHook(ctx, request, hooks.Stop) + default: + return Response{}, xerrors.Errorf("unknown event type %q", request.Type) + } +} + +func dispatchHook[T any](ctx context.Context, request Request, hook func(context.Context, Meta, T) (Response, error)) (Response, error) { + if hook == nil { + return Response{}, nil + } + var data T + if err := json.Unmarshal(request.Data, &data); err != nil { + return Response{}, xerrors.Errorf("decode %q event data: %w", request.Type, err) + } + return hook(ctx, request.Meta, data) +} diff --git a/codersdk/agenthooks/jwt.go b/codersdk/agenthooks/jwt.go new file mode 100644 index 0000000000000..bd8b2147ef355 --- /dev/null +++ b/codersdk/agenthooks/jwt.go @@ -0,0 +1,126 @@ +package agenthooks + +import ( + "encoding/hex" + "encoding/json" + "strings" + "time" + + "github.com/go-jose/go-jose/v4" + "github.com/google/uuid" + "golang.org/x/xerrors" +) + +const jwtType = "JWT" + +// SignClaims signs claims with the shared secret using HS256. +func SignClaims(secret []byte, claims Claims) (string, error) { + signer, err := jose.NewSigner( + jose.SigningKey{Algorithm: jose.HS256, Key: secret}, + new(jose.SignerOptions).WithType(jwtType), + ) + if err != nil { + return "", xerrors.Errorf("create signer: %w", err) + } + + payload, err := json.Marshal(claims) + if err != nil { + return "", xerrors.Errorf("marshal claims: %w", err) + } + signed, err := signer.Sign(payload) + if err != nil { + return "", xerrors.Errorf("sign claims: %w", err) + } + token, err := signed.CompactSerialize() + if err != nil { + return "", xerrors.Errorf("serialize token: %w", err) + } + return token, nil +} + +func Verify(authzHeader string, secret []byte) (Claims, error) { + const bearerPrefix = "Bearer " + token, ok := strings.CutPrefix(authzHeader, bearerPrefix) + if !ok || token == "" || strings.ContainsAny(token, " \t\r\n") { + return Claims{}, xerrors.New("authorization header must contain one Bearer token") + } + + object, err := jose.ParseSigned(token, []jose.SignatureAlgorithm{jose.HS256}) + if err != nil { + return Claims{}, xerrors.Errorf("parse token: %w", err) + } + if len(object.Signatures) != 1 { + return Claims{}, xerrors.New("token must contain one signature") + } + header := object.Signatures[0].Header + if header.Algorithm != string(jose.HS256) { + return Claims{}, xerrors.Errorf("token algorithm must be %q", jose.HS256) + } + typ, ok := header.ExtraHeaders[jose.HeaderType].(string) + if !ok || typ != jwtType { + return Claims{}, xerrors.Errorf("token type must be %q", jwtType) + } + + payload, err := object.Verify(secret) + if err != nil { + return Claims{}, xerrors.Errorf("verify token: %w", err) + } + var claims Claims + if err := json.Unmarshal(payload, &claims); err != nil { + return Claims{}, xerrors.Errorf("decode claims: %w", err) + } + if err := validateClaims(claims, time.Now()); err != nil { + return Claims{}, err + } + return claims, nil +} + +func validateClaims(claims Claims, now time.Time) error { + switch { + case claims.Issuer == "": + return xerrors.New("issuer is required") + case claims.Subject == "": + return xerrors.New("subject is required") + case claims.Audience == "": + return xerrors.New("audience is required") + case claims.IssuedAt == 0: + return xerrors.New("issued at is required") + case claims.NotBefore == 0: + return xerrors.New("not before is required") + case claims.Expires == 0: + return xerrors.New("expiry is required") + case claims.JTI == uuid.Nil: + return xerrors.New("JWT ID is required") + case !validEventType(claims.Type): + return xerrors.Errorf("invalid event type %q", claims.Type) + case !validSHA256(claims.BodySHA256): + return xerrors.New("body SHA-256 must be a hexadecimal SHA-256 digest") + case claims.NotBefore > claims.Expires: + return xerrors.New("not before must not be after expiry") + case claims.IssuedAt > claims.Expires: + return xerrors.New("issued at must not be after expiry") + case now.Unix() < claims.NotBefore: + return xerrors.New("token is not valid yet") + case now.Unix() >= claims.Expires: + return xerrors.New("token has expired") + } + if _, err := claims.ChatID(); err != nil { + return err + } + return nil +} + +func validEventType(eventType EventType) bool { + switch eventType { + case EventSessionStart, EventUserPromptSubmit, EventPreToolUse, + EventPostToolUse, EventPreCompact, EventPostCompact, EventStop: + return true + default: + return false + } +} + +func validSHA256(value string) bool { + decoded, err := hex.DecodeString(value) + return err == nil && len(decoded) == 32 +} diff --git a/codersdk/agenthooks/types.go b/codersdk/agenthooks/types.go new file mode 100644 index 0000000000000..11fa581e5cb2c --- /dev/null +++ b/codersdk/agenthooks/types.go @@ -0,0 +1,159 @@ +// Package agenthooks defines the experimental wire protocol for Coder agent +// lifecycle hooks. It requires the agent-lifecycle-hooks experiment and has no +// backward-compatibility guarantee, including for SchemaVersion 1. +package agenthooks + +import ( + "encoding/json" + "strings" + + "github.com/google/uuid" + "golang.org/x/xerrors" +) + +// SchemaVersion is the current lifecycle hook request schema version. +const SchemaVersion = 1 + +type EventType string + +const ( + EventSessionStart EventType = "session_start" + EventUserPromptSubmit EventType = "user_prompt_submit" + EventPreToolUse EventType = "pre_tool_use" + EventPostToolUse EventType = "post_tool_use" + EventPreCompact EventType = "pre_compact" + EventPostCompact EventType = "post_compact" + EventStop EventType = "stop" +) + +// Request is the body coderd posts to the configured lifecycle hook URL. +type Request struct { + Type EventType `json:"type"` + Meta Meta `json:"meta"` + Data json.RawMessage `json:"data"` +} + +func (r Request) Decode() (any, error) { + var data any + switch r.Type { + case EventSessionStart: + data = &SessionStartData{} + case EventUserPromptSubmit: + data = &UserPromptSubmitData{} + case EventPreToolUse: + data = &PreToolUseData{} + case EventPostToolUse: + data = &PostToolUseData{} + case EventPreCompact: + data = &PreCompactData{} + case EventPostCompact: + data = &PostCompactData{} + case EventStop: + data = &StopData{} + default: + return nil, xerrors.Errorf("unknown event type %q", r.Type) + } + + if err := json.Unmarshal(r.Data, data); err != nil { + return nil, xerrors.Errorf("decode %q event data: %w", r.Type, err) + } + return data, nil +} + +type Meta struct { + DispatchID uuid.UUID `json:"dispatch_id"` + SchemaVersion int `json:"schema_version"` + ChatRef +} + +// ChatRef identifies the chat a lifecycle hook event refers to. Embedded +// structs flatten in JSON, so it adds no nesting on the wire. +type ChatRef struct { + ChatID uuid.UUID `json:"chat_id"` + OwnerID uuid.UUID `json:"owner_id"` + WorkspaceID *uuid.UUID `json:"workspace_id,omitempty"` + TurnID *uuid.UUID `json:"turn_id,omitempty"` + ParentChatID *uuid.UUID `json:"parent_chat_id,omitempty"` + // RootChatID groups a subagent subtree with its user-facing conversation. + // Unset for top-level chats. + RootChatID *uuid.UUID `json:"root_chat_id,omitempty"` +} + +type SessionStartData struct { + Source string `json:"source"` +} + +// UserPromptSubmitData includes concatenated text and persisted parts. +// Inspect Parts when structure matters. +type UserPromptSubmitData struct { + Prompt string `json:"prompt"` + Parts json.RawMessage `json:"parts,omitempty"` +} + +type PreToolUseData struct { + ToolUseID string `json:"tool_use_id"` + ToolName string `json:"tool_name"` + ToolInput json.RawMessage `json:"tool_input"` +} + +type PostToolUseData struct { + ToolUseID string `json:"tool_use_id"` + ToolName string `json:"tool_name"` + ToolResponse json.RawMessage `json:"tool_response,omitempty"` + ToolError string `json:"tool_error,omitempty"` +} + +type PreCompactData struct{} + +type PostCompactData struct{} + +type StopData struct{} + +type Response struct { + Permission *Permission `json:"permission,omitempty"` + ModelContext string `json:"model_context,omitempty"` + UserMessage string `json:"user_message,omitempty"` + // AllowedTools distinguishes unchanged (nil), no tools (empty), and named tools. + AllowedTools *[]string `json:"allowed_tools,omitempty"` + EndChat bool `json:"end_chat,omitempty"` +} + +// Permission controls whether mutable hook input may proceed. +type Permission struct { + Decision PermissionDecision `json:"decision"` + Reason string `json:"reason,omitempty"` + InputOverride json.RawMessage `json:"input_override,omitempty"` +} + +type PermissionDecision string + +const ( + PermissionAllow PermissionDecision = "allow" + PermissionDeny PermissionDecision = "deny" + PermissionAsk PermissionDecision = "ask" +) + +// Claims describes the JWT minted by coderd for a lifecycle hook dispatch. +type Claims struct { + Issuer string `json:"iss"` + Subject string `json:"sub"` + Audience string `json:"aud"` + IssuedAt int64 `json:"iat"` + NotBefore int64 `json:"nbf"` + Expires int64 `json:"exp"` + JTI uuid.UUID `json:"jti"` + Type EventType `json:"type"` + BodySHA256 string `json:"body_sha256"` +} + +func (c Claims) ChatID() (uuid.UUID, error) { + value, ok := strings.CutPrefix(c.Subject, "coder:chat:") + if !ok { + return uuid.Nil, xerrors.Errorf("invalid subject %q", c.Subject) + } + chatID, err := uuid.Parse(value) + if err != nil { + return uuid.Nil, xerrors.Errorf("parse chat ID: %w", err) + } + return chatID, nil +} diff --git a/codersdk/chats.go b/codersdk/chats.go index c6990e8a776f8..0f4149068f71a 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -631,18 +631,27 @@ type EditChatMessageRequest struct { // CreateChatMessageResponse is the response from adding a message to a chat. type CreateChatMessageResponse struct { - Message *ChatMessage `json:"message,omitempty"` + Message *ChatMessage `json:"message,omitempty"` + // Messages contains all user-visible messages inserted by an immediate send, + // in insertion order with the user's message last. Clients should upsert the + // full batch because hooks may prepend notices. Empty for queued or ended sends. + Messages []ChatMessage `json:"messages,omitempty"` QueuedMessage *ChatQueuedMessage `json:"queued_message,omitempty"` Queued bool `json:"queued"` + Ended bool `json:"ended"` Warnings []string `json:"warnings,omitempty"` } // EditChatMessageResponse is the response from editing a message in a chat. -// Edits are always synchronous (no queueing), so the message is returned -// directly. type EditChatMessageResponse struct { - Message ChatMessage `json:"message"` - Warnings []string `json:"warnings,omitempty"` + Message *ChatMessage `json:"message,omitempty"` + // Messages holds every user-visible message the edit inserted, in + // insertion order with the replacement message last. Lifecycle + // hooks may prepend notices, so clients must upsert all of them + // rather than only Message. Empty for ended edits. + Messages []ChatMessage `json:"messages,omitempty"` + Ended bool `json:"ended"` + Warnings []string `json:"warnings,omitempty"` } // UploadChatFileResponse is the response from uploading a chat file. @@ -1718,6 +1727,7 @@ const ( ChatErrorKindMissingKey ChatErrorKind = "missing_key" ChatErrorKindProviderDisabled ChatErrorKind = "provider_disabled" ChatErrorKindContentFilter ChatErrorKind = "content_filter" + ChatErrorKindHookDispatchFailed ChatErrorKind = "hook_dispatch_failed" ) // AllChatErrorKinds contains every ChatErrorKind value. @@ -1734,6 +1744,7 @@ var AllChatErrorKinds = []ChatErrorKind{ ChatErrorKindMissingKey, ChatErrorKindProviderDisabled, ChatErrorKindContentFilter, + ChatErrorKindHookDispatchFailed, } // ChatError represents a terminal chat error in persisted chat state or the diff --git a/codersdk/deployment.go b/codersdk/deployment.go index 2c4a0bffa9d24..dc10160b49882 100644 --- a/codersdk/deployment.go +++ b/codersdk/deployment.go @@ -4293,6 +4293,47 @@ Write out the current server config as YAML to stdout.`, Group: &deploymentGroupChat, YAML: "debugLoggingEnabled", }, + { + Name: "Chat: Hook URL", + Description: "HTTPS URL to receive chat agent lifecycle hook events. Hooks are disabled when unset. Requires the agent-lifecycle-hooks experiment.", + Flag: "chat-hook-url", + Env: "CODER_CHAT_HOOK_URL", + Value: &c.AI.Chat.HookURL, + Default: "", + Group: &deploymentGroupChat, + YAML: "hookURL", + }, + { + Name: "Chat: Hook Secret", + Description: "Shared secret used to sign chat agent lifecycle hook JWTs.", + Flag: "chat-hook-secret", + Env: "CODER_CHAT_HOOK_SECRET", + Value: &c.AI.Chat.HookSecret, + Default: "", + Group: &deploymentGroupChat, + Annotations: serpent.Annotations{}.Mark(annotationSecretKey, "true"), + }, + { + Name: "Chat: Hook Timeout", + Description: "Maximum time to wait for a chat agent lifecycle hook response.", + Flag: "chat-hook-timeout", + Env: "CODER_CHAT_HOOK_TIMEOUT", + Value: &c.AI.Chat.HookTimeout, + Default: (1500 * time.Millisecond).String(), + Group: &deploymentGroupChat, + YAML: "hookTimeout", + Annotations: serpent.Annotations{}.Mark(annotationFormatDuration, "true"), + }, + { + Name: "Chat: Hook Enabled", + Description: "Whether to dispatch chat agent lifecycle hooks when a hook URL is configured. Requires the agent-lifecycle-hooks experiment.", + Flag: "chat-hook-enabled", + Env: "CODER_CHAT_HOOK_ENABLED", + Value: &c.AI.Chat.HookEnabled, + Default: "true", + Group: &deploymentGroupChat, + YAML: "hookEnabled", + }, { Name: "Chat: AI Gateway Routing Enabled", Description: "Deprecated: AI Gateway routing is now the only routing path. Setting this value has no effect. This option will be removed in a future release.", @@ -4975,8 +5016,12 @@ type AIBridgeProxyConfig struct { } type ChatConfig struct { - AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` - DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` + AcquireBatchSize serpent.Int64 `json:"acquire_batch_size" typescript:",notnull"` + DebugLoggingEnabled serpent.Bool `json:"debug_logging_enabled" typescript:",notnull"` + HookURL serpent.URL `json:"hook_url" typescript:",notnull"` + HookSecret serpent.String `json:"hook_secret" typescript:",notnull"` + HookTimeout serpent.Duration `json:"hook_timeout" typescript:",notnull"` + HookEnabled serpent.Bool `json:"hook_enabled" typescript:",notnull"` // Deprecated: AI Gateway routing is now the only routing path. Setting this // value has no effect. This option will be removed in a future release. AIGatewayRoutingEnabled serpent.Bool `json:"ai_gateway_routing_enabled" typescript:",notnull" swaggerignore:"true"` @@ -5026,6 +5071,38 @@ func (c *DeploymentValues) Validate() error { refresh, access, ) } + + // Disabled hooks must not validate inert settings. + if c.AI.Chat.HookEnabled.Value() { + if c.AI.Chat.HookURL.String() != "" { + hookURL := c.AI.Chat.HookURL.Value() + if hookURL.Scheme != "https" { + return xerrors.New("chat hook URL must use HTTPS; set --chat-hook-url to an HTTPS URL") + } + if hookURL.Host == "" { + return xerrors.New("chat hook URL must include a host; set --chat-hook-url to a complete HTTPS URL") + } + // The configured string is signed verbatim as the JWT audience, + // but fragments and userinfo never reach the consumer, so its + // reconstructed audience would mismatch on every dispatch. + if hookURL.Fragment != "" || hookURL.RawFragment != "" || hookURL.User != nil { + return xerrors.New("chat hook URL must not contain a fragment or userinfo; set --chat-hook-url to a plain HTTPS URL") + } + if c.AI.Chat.HookSecret.Value() == "" { + return xerrors.New("chat hook secret is required when chat hook URL is set; set --chat-hook-secret") + } + // go-jose requires HS256 secrets to be at least 32 bytes. + if len(c.AI.Chat.HookSecret.Value()) < 32 { + return xerrors.New("chat hook secret must be at least 32 bytes of cryptographically random data; set --chat-hook-secret to a longer value") + } + } + + hookTimeout := c.AI.Chat.HookTimeout.Value() + if hookTimeout <= 0 || hookTimeout > 5*time.Second { + return xerrors.Errorf("chat hook timeout (%s) must be greater than zero and no more than 5s; set --chat-hook-timeout to a valid duration", hookTimeout) + } + } + return nil } @@ -5246,6 +5323,7 @@ const ( ExperimentAIGatewayCostControl Experiment = "ai-gateway-cost-control" // Enables AI Gateway cost control functionality. ExperimentChatAdvisor Experiment = "chat-advisor" // Enables the advisor tool for root agent chats. ExperimentChatVirtualDesktop Experiment = "chat-virtual-desktop" // Enables virtual desktop and computer use provider for agents. + ExperimentAgentLifecycleHooks Experiment = "agent-lifecycle-hooks" // Enables chat lifecycle hook webhooks for agent chats. ) func (e Experiment) DisplayName() string { @@ -5274,6 +5352,8 @@ func (e Experiment) DisplayName() string { return "Chat Advisor" case ExperimentChatVirtualDesktop: return "Chat Virtual Desktop" + case ExperimentAgentLifecycleHooks: + return "Agent Lifecycle Hooks" default: // Split on hyphen and convert to title case // e.g. "mcp-server-http" -> "Mcp Server Http" @@ -5296,6 +5376,7 @@ var ExperimentsKnown = Experiments{ ExperimentAIGatewayCostControl, ExperimentChatAdvisor, ExperimentChatVirtualDesktop, + ExperimentAgentLifecycleHooks, } // ExperimentsSafe should include all experiments that are safe for diff --git a/codersdk/deployment_test.go b/codersdk/deployment_test.go index a70fef938a64c..901ba9f3578ae 100644 --- a/codersdk/deployment_test.go +++ b/codersdk/deployment_test.go @@ -82,6 +82,9 @@ func TestDeploymentValues_HighlyConfigurable(t *testing.T) { "Email Auth: Password": { yaml: true, }, + "Chat: Hook Secret": { + yaml: true, + }, "Notifications: Email Auth: Password": { yaml: true, }, @@ -726,6 +729,7 @@ func TestDeploymentValues_Validate_RefreshLifetime(t *testing.T) { dv := &codersdk.DeploymentValues{} dv.Sessions.DefaultDuration = serpent.Duration(access) dv.Sessions.RefreshDefaultDuration = serpent.Duration(refresh) + dv.AI.Chat.HookTimeout = serpent.Duration(1500 * time.Millisecond) return dv } @@ -770,6 +774,115 @@ func TestDeploymentValues_Validate_RefreshLifetime(t *testing.T) { }) } +func TestDeploymentValues_Validate_ChatHooks(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + disabled bool + url string + secret string + timeout time.Duration + wantErr string + }{ + { + name: "NoURL", + timeout: 1500 * time.Millisecond, + }, + { + name: "DisabledSkipsValidation", + disabled: true, + url: "http://hooks.example.com/agent", + timeout: 0, + }, + { + name: "Valid", + url: "https://hooks.example.com/agent", + secret: "0123456789abcdef0123456789abcdef", + timeout: 5 * time.Second, + }, + { + name: "HTTPURL", + url: "http://hooks.example.com/agent", + secret: "0123456789abcdef0123456789abcdef", + timeout: 1500 * time.Millisecond, + wantErr: "chat hook URL must use HTTPS", + }, + { + name: "HostlessURL", + url: "https:///hook", + secret: "0123456789abcdef0123456789abcdef", + timeout: 1500 * time.Millisecond, + wantErr: "must include a host", + }, + { + name: "FragmentURL", + url: "https://hooks.example.com/agent#frag", + secret: "0123456789abcdef0123456789abcdef", + timeout: 1500 * time.Millisecond, + wantErr: "must not contain a fragment or userinfo", + }, + { + name: "UserinfoURL", + url: "https://user:pass@hooks.example.com/agent", + secret: "0123456789abcdef0123456789abcdef", + timeout: 1500 * time.Millisecond, + wantErr: "must not contain a fragment or userinfo", + }, + { + name: "MissingSecret", + url: "https://hooks.example.com/agent", + timeout: 1500 * time.Millisecond, + wantErr: "chat hook secret is required", + }, + { + name: "ShortSecret", + url: "https://hooks.example.com/agent", + secret: "0123456789abcdef0123456789abcde", + timeout: 1500 * time.Millisecond, + wantErr: "chat hook secret must be at least 32 bytes", + }, + { + name: "ZeroTimeout", + timeout: 0, + wantErr: "chat hook timeout", + }, + { + name: "NegativeTimeout", + timeout: -time.Millisecond, + wantErr: "chat hook timeout", + }, + { + name: "TimeoutAboveMaximum", + timeout: 5*time.Second + time.Millisecond, + wantErr: "chat hook timeout", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() + + dv := &codersdk.DeploymentValues{} + dv.Sessions.DefaultDuration = serpent.Duration(time.Hour) + dv.Sessions.RefreshDefaultDuration = serpent.Duration(48 * time.Hour) + dv.AI.Chat.HookEnabled = serpent.Bool(!tt.disabled) + dv.AI.Chat.HookSecret = serpent.String(tt.secret) + dv.AI.Chat.HookTimeout = serpent.Duration(tt.timeout) + if tt.url != "" { + require.NoError(t, dv.AI.Chat.HookURL.Set(tt.url)) + } + + err := dv.Validate() + if tt.wantErr == "" { + require.NoError(t, err) + return + } + require.ErrorContains(t, err, tt.wantErr) + }) + } +} + func TestDeploymentValues_DurationFormatNanoseconds(t *testing.T) { t.Parallel() diff --git a/docs/admin/integrations/prometheus.md b/docs/admin/integrations/prometheus.md index 26f71eca7aafe..d69dbdd4c1cc4 100644 --- a/docs/admin/integrations/prometheus.md +++ b/docs/admin/integrations/prometheus.md @@ -225,6 +225,11 @@ deployment. They will always be available from the agent. | `coderd_chat_auto_archive_records_archived_total` | counter | Total number of chats archived by the auto-archive job (counting both roots and cascaded children). | | | `coderd_chatd_chats` | gauge | Number of chats being processed, by state. | `state` | | `coderd_chatd_compaction_total` | counter | Total compaction outcomes (only recorded when compaction was triggered or failed). | `model` `provider` `result` | +| `coderd_chatd_hook_context_size_bytes` | histogram | Lifecycle hook model context response size in bytes. | `event` | +| `coderd_chatd_hook_decisions_total` | counter | Total lifecycle hook permission decisions by event and decision. | `decision` `event` | +| `coderd_chatd_hook_dispatch_seconds` | histogram | Lifecycle hook dispatch duration in seconds. | `event` | +| `coderd_chatd_hook_dispatches_total` | counter | Total lifecycle hook dispatches by event and result. | `event` `result` | +| `coderd_chatd_hook_input_overrides_total` | counter | Total lifecycle hook input overrides by event. | `event` | | `coderd_chatd_message_count` | histogram | Number of messages in the prompt per LLM request. | `model` `provider` | | `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` | diff --git a/docs/admin/security/audit-logs.md b/docs/admin/security/audit-logs.md index af5967e908a92..d14fcf327624c 100644 --- a/docs/admin/security/audit-logs.md +++ b/docs/admin/security/audit-logs.md @@ -13,41 +13,41 @@ We track the following resources: -| Resource | | | -|-----------------------------------------------------------------|----------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| AIGatewayKey
create, delete | |
FieldTracked
created_atfalse
hashed_secrettrue
idtrue
last_heartbeat_atfalse
nametrue
secret_prefixtrue
| -| AIProvider
create, write, delete | |
FieldTracked
base_urltrue
created_atfalse
deletedtrue
display_nametrue
enabledtrue
icontrue
idtrue
nametrue
settingstrue
settings_key_idfalse
typetrue
updated_atfalse
| -| AIProviderKey
create, delete | |
FieldTracked
api_keytrue
api_key_key_idfalse
created_atfalse
idtrue
provider_idtrue
updated_atfalse
| -| AISeatState
create | |
FieldTracked
first_used_attrue
last_event_descriptiontrue
last_event_typetrue
last_used_atfalse
updated_atfalse
user_idtrue
| -| APIKey
login, logout, register, create, write, delete | |
FieldTracked
allow_listfalse
created_attrue
expires_attrue
hashed_secretfalse
idfalse
ip_addressfalse
last_usedtrue
lifetime_secondsfalse
login_typefalse
scopesfalse
token_namefalse
updated_atfalse
user_idtrue
| -| AuditOAuthConvertState
| |
FieldTracked
created_attrue
expires_attrue
from_login_typetrue
to_login_typetrue
user_idtrue
| -| Group
create, write, delete | |
FieldTracked
avatar_urltrue
chat_spend_limit_microstrue
display_nametrue
idtrue
memberstrue
nametrue
organization_idfalse
quota_allowancetrue
sourcefalse
| -| 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
compaction_requested_atfalse
context_aggregate_hashfalse
context_dirty_resourcesfalse
context_dirty_sincefalse
context_errorfalse
created_atfalse
dynamic_toolsfalse
generation_attemptfalse
group_acltrue
heartbeat_atfalse
history_versionfalse
idtrue
labelstrue
last_errorfalse
last_model_config_idfalse
last_read_message_idfalse
last_reasoning_effortfalse
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
| -| HealthSettings
| |
FieldTracked
dismissed_healthcheckstrue
idfalse
| -| License
create, delete | |
FieldTracked
exptrue
idfalse
jwtfalse
uploaded_attrue
uuidtrue
| -| NotificationTemplate
| |
FieldTracked
actionstrue
body_templatetrue
enabled_by_defaulttrue
grouptrue
idfalse
kindtrue
methodtrue
nametrue
title_templatetrue
| -| NotificationsSettings
| |
FieldTracked
idfalse
notifier_pausedtrue
| -| OAuth2ProviderApp
| |
FieldTracked
callback_urltrue
client_id_issued_atfalse
client_secret_expires_attrue
client_typetrue
client_uritrue
contactstrue
created_atfalse
dynamically_registeredtrue
grant_typestrue
icontrue
idfalse
jwkstrue
jwks_uritrue
logo_uritrue
nametrue
policy_uritrue
redirect_uristrue
registration_access_tokentrue
registration_client_uritrue
response_typestrue
scopetrue
software_idtrue
software_versiontrue
token_endpoint_auth_methodtrue
tos_uritrue
updated_atfalse
| -| OAuth2ProviderAppSecret
| |
FieldTracked
app_idfalse
created_atfalse
display_secretfalse
hashed_secretfalse
idfalse
last_used_atfalse
secret_prefixfalse
| -| Organization
| |
FieldTracked
created_atfalse
default_org_member_rolestrue
deletedtrue
descriptiontrue
display_nametrue
icontrue
idfalse
is_defaulttrue
nametrue
shareable_workspace_ownerstrue
updated_attrue
| -| OrganizationSyncSettings
| |
FieldTracked
assign_defaulttrue
fieldtrue
mappingtrue
| -| PrebuildsSettings
| |
FieldTracked
idfalse
reconciliation_pausedtrue
| -| RoleSyncSettings
| |
FieldTracked
fieldtrue
mappingtrue
| -| TaskTable
| |
FieldTracked
created_atfalse
deleted_atfalse
display_nametrue
idtrue
nametrue
organization_idfalse
owner_idtrue
prompttrue
template_parameterstrue
template_version_idtrue
workspace_idtrue
| -| Template
write, delete | |
FieldTracked
active_version_idtrue
activity_bumptrue
allow_user_autostarttrue
allow_user_autostoptrue
allow_user_cancel_workspace_jobstrue
autostart_block_days_of_weektrue
autostop_requirement_days_of_weektrue
autostop_requirement_weekstrue
cors_behaviortrue
created_atfalse
created_bytrue
created_by_avatar_urlfalse
created_by_namefalse
created_by_usernamefalse
default_ttltrue
deletedfalse
deprecatedtrue
descriptiontrue
disable_module_cachetrue
display_nametrue
failure_ttltrue
group_acltrue
icontrue
idtrue
max_port_sharing_leveltrue
nametrue
organization_display_namefalse
organization_iconfalse
organization_idfalse
organization_namefalse
provisionertrue
require_active_versiontrue
time_til_autostop_notifytrue
time_til_dormanttrue
time_til_dormant_autodeletetrue
updated_atfalse
use_classic_parameter_flowtrue
user_acltrue
| -| TemplateVersion
create, write | |
FieldTracked
archivedtrue
created_atfalse
created_bytrue
created_by_avatar_urlfalse
created_by_namefalse
created_by_usernamefalse
external_auth_providersfalse
has_ai_taskfalse
has_external_agentfalse
idtrue
job_idfalse
messagefalse
nametrue
organization_idfalse
readmetrue
source_example_idfalse
template_idtrue
updated_atfalse
| -| User
create, write, delete | |
FieldTracked
avatar_urlfalse
chat_spend_limit_microstrue
created_atfalse
deletedtrue
emailtrue
github_com_user_idfalse
hashed_one_time_passcodefalse
hashed_passwordtrue
idtrue
is_service_accounttrue
is_systemtrue
last_seen_atfalse
login_typetrue
nametrue
one_time_passcode_expires_attrue
quiet_hours_scheduletrue
rbac_rolestrue
statustrue
updated_atfalse
usernametrue
| -| UserSecret
create, write, delete | |
FieldTracked
created_atfalse
descriptiontrue
env_nametrue
file_pathtrue
idtrue
nametrue
updated_atfalse
user_idtrue
valuetrue
value_key_idfalse
| -| UserSkill
create, write, delete | |
FieldTracked
contenttrue
created_atfalse
descriptiontrue
idtrue
nametrue
updated_atfalse
user_idtrue
| -| WorkspaceBuild
start, stop | |
FieldTracked
build_numberfalse
created_atfalse
daily_costfalse
deadlinefalse
has_ai_taskfalse
has_external_agentfalse
idfalse
initiator_by_avatar_urlfalse
initiator_by_namefalse
initiator_by_usernamefalse
initiator_idfalse
job_idfalse
max_deadlinefalse
notified_autostop_deadlinefalse
reasonfalse
template_version_idtrue
template_version_preset_idfalse
transitionfalse
updated_atfalse
workspace_idfalse
| -| WorkspaceProxy
| |
FieldTracked
created_attrue
deletedfalse
derp_enabledtrue
derp_onlytrue
display_nametrue
icontrue
idtrue
nametrue
region_idtrue
token_hashed_secrettrue
updated_atfalse
urltrue
versiontrue
wildcard_hostnametrue
| -| WorkspaceTable
| |
FieldTracked
automatic_updatestrue
autostart_scheduletrue
created_atfalse
deletedfalse
deleting_attrue
dormant_attrue
favoritetrue
group_acltrue
idtrue
last_used_atfalse
nametrue
next_start_attrue
organization_idfalse
owner_idtrue
template_idtrue
ttltrue
updated_atfalse
user_acltrue
| +| Resource | | | +|-----------------------------------------------------------------|----------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| AIGatewayKey
create, delete | |
FieldTracked
created_atfalse
hashed_secrettrue
idtrue
last_heartbeat_atfalse
nametrue
secret_prefixtrue
| +| AIProvider
create, write, delete | |
FieldTracked
base_urltrue
created_atfalse
deletedtrue
display_nametrue
enabledtrue
icontrue
idtrue
nametrue
settingstrue
settings_key_idfalse
typetrue
updated_atfalse
| +| AIProviderKey
create, delete | |
FieldTracked
api_keytrue
api_key_key_idfalse
created_atfalse
idtrue
provider_idtrue
updated_atfalse
| +| AISeatState
create | |
FieldTracked
first_used_attrue
last_event_descriptiontrue
last_event_typetrue
last_used_atfalse
updated_atfalse
user_idtrue
| +| APIKey
login, logout, register, create, write, delete | |
FieldTracked
allow_listfalse
created_attrue
expires_attrue
hashed_secretfalse
idfalse
ip_addressfalse
last_usedtrue
lifetime_secondsfalse
login_typefalse
scopesfalse
token_namefalse
updated_atfalse
user_idtrue
| +| AuditOAuthConvertState
| |
FieldTracked
created_attrue
expires_attrue
from_login_typetrue
to_login_typetrue
user_idtrue
| +| Group
create, write, delete | |
FieldTracked
avatar_urltrue
chat_spend_limit_microstrue
display_nametrue
idtrue
memberstrue
nametrue
organization_idfalse
quota_allowancetrue
sourcefalse
| +| 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
compaction_requested_atfalse
context_aggregate_hashfalse
context_dirty_resourcesfalse
context_dirty_sincefalse
context_errorfalse
created_atfalse
dynamic_toolsfalse
generation_attemptfalse
group_acltrue
heartbeat_atfalse
history_versionfalse
hook_allowed_toolsfalse
idtrue
labelstrue
last_errorfalse
last_model_config_idfalse
last_read_message_idfalse
last_reasoning_effortfalse
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
| +| HealthSettings
| |
FieldTracked
dismissed_healthcheckstrue
idfalse
| +| License
create, delete | |
FieldTracked
exptrue
idfalse
jwtfalse
uploaded_attrue
uuidtrue
| +| NotificationTemplate
| |
FieldTracked
actionstrue
body_templatetrue
enabled_by_defaulttrue
grouptrue
idfalse
kindtrue
methodtrue
nametrue
title_templatetrue
| +| NotificationsSettings
| |
FieldTracked
idfalse
notifier_pausedtrue
| +| OAuth2ProviderApp
| |
FieldTracked
callback_urltrue
client_id_issued_atfalse
client_secret_expires_attrue
client_typetrue
client_uritrue
contactstrue
created_atfalse
dynamically_registeredtrue
grant_typestrue
icontrue
idfalse
jwkstrue
jwks_uritrue
logo_uritrue
nametrue
policy_uritrue
redirect_uristrue
registration_access_tokentrue
registration_client_uritrue
response_typestrue
scopetrue
software_idtrue
software_versiontrue
token_endpoint_auth_methodtrue
tos_uritrue
updated_atfalse
| +| OAuth2ProviderAppSecret
| |
FieldTracked
app_idfalse
created_atfalse
display_secretfalse
hashed_secretfalse
idfalse
last_used_atfalse
secret_prefixfalse
| +| Organization
| |
FieldTracked
created_atfalse
default_org_member_rolestrue
deletedtrue
descriptiontrue
display_nametrue
icontrue
idfalse
is_defaulttrue
nametrue
shareable_workspace_ownerstrue
updated_attrue
| +| OrganizationSyncSettings
| |
FieldTracked
assign_defaulttrue
fieldtrue
mappingtrue
| +| PrebuildsSettings
| |
FieldTracked
idfalse
reconciliation_pausedtrue
| +| RoleSyncSettings
| |
FieldTracked
fieldtrue
mappingtrue
| +| TaskTable
| |
FieldTracked
created_atfalse
deleted_atfalse
display_nametrue
idtrue
nametrue
organization_idfalse
owner_idtrue
prompttrue
template_parameterstrue
template_version_idtrue
workspace_idtrue
| +| Template
write, delete | |
FieldTracked
active_version_idtrue
activity_bumptrue
allow_user_autostarttrue
allow_user_autostoptrue
allow_user_cancel_workspace_jobstrue
autostart_block_days_of_weektrue
autostop_requirement_days_of_weektrue
autostop_requirement_weekstrue
cors_behaviortrue
created_atfalse
created_bytrue
created_by_avatar_urlfalse
created_by_namefalse
created_by_usernamefalse
default_ttltrue
deletedfalse
deprecatedtrue
descriptiontrue
disable_module_cachetrue
display_nametrue
failure_ttltrue
group_acltrue
icontrue
idtrue
max_port_sharing_leveltrue
nametrue
organization_display_namefalse
organization_iconfalse
organization_idfalse
organization_namefalse
provisionertrue
require_active_versiontrue
time_til_autostop_notifytrue
time_til_dormanttrue
time_til_dormant_autodeletetrue
updated_atfalse
use_classic_parameter_flowtrue
user_acltrue
| +| TemplateVersion
create, write | |
FieldTracked
archivedtrue
created_atfalse
created_bytrue
created_by_avatar_urlfalse
created_by_namefalse
created_by_usernamefalse
external_auth_providersfalse
has_ai_taskfalse
has_external_agentfalse
idtrue
job_idfalse
messagefalse
nametrue
organization_idfalse
readmetrue
source_example_idfalse
template_idtrue
updated_atfalse
| +| User
create, write, delete | |
FieldTracked
avatar_urlfalse
chat_spend_limit_microstrue
created_atfalse
deletedtrue
emailtrue
github_com_user_idfalse
hashed_one_time_passcodefalse
hashed_passwordtrue
idtrue
is_service_accounttrue
is_systemtrue
last_seen_atfalse
login_typetrue
nametrue
one_time_passcode_expires_attrue
quiet_hours_scheduletrue
rbac_rolestrue
statustrue
updated_atfalse
usernametrue
| +| UserSecret
create, write, delete | |
FieldTracked
created_atfalse
descriptiontrue
env_nametrue
file_pathtrue
idtrue
nametrue
updated_atfalse
user_idtrue
valuetrue
value_key_idfalse
| +| UserSkill
create, write, delete | |
FieldTracked
contenttrue
created_atfalse
descriptiontrue
idtrue
nametrue
updated_atfalse
user_idtrue
| +| WorkspaceBuild
start, stop | |
FieldTracked
build_numberfalse
created_atfalse
daily_costfalse
deadlinefalse
has_ai_taskfalse
has_external_agentfalse
idfalse
initiator_by_avatar_urlfalse
initiator_by_namefalse
initiator_by_usernamefalse
initiator_idfalse
job_idfalse
max_deadlinefalse
notified_autostop_deadlinefalse
reasonfalse
template_version_idtrue
template_version_preset_idfalse
transitionfalse
updated_atfalse
workspace_idfalse
| +| WorkspaceProxy
| |
FieldTracked
created_attrue
deletedfalse
derp_enabledtrue
derp_onlytrue
display_nametrue
icontrue
idtrue
nametrue
region_idtrue
token_hashed_secrettrue
updated_atfalse
urltrue
versiontrue
wildcard_hostnametrue
| +| WorkspaceTable
| |
FieldTracked
automatic_updatestrue
autostart_scheduletrue
created_atfalse
deletedfalse
deleting_attrue
dormant_attrue
favoritetrue
group_acltrue
idtrue
last_used_atfalse
nametrue
next_start_attrue
organization_idfalse
owner_idtrue
template_idtrue
ttltrue
updated_atfalse
user_acltrue
| diff --git a/docs/admin/setup/chat-lifecycle-hooks.md b/docs/admin/setup/chat-lifecycle-hooks.md new file mode 100644 index 0000000000000..c002928ac87c6 --- /dev/null +++ b/docs/admin/setup/chat-lifecycle-hooks.md @@ -0,0 +1,187 @@ +# Configure chat lifecycle hooks + +> [!NOTE] +> Chat lifecycle hooks are an experimental feature. +> The feature requires the `agent-lifecycle-hooks` experiment, and the consumer contract (including the request schema and JWT claims) may change or be removed in any release without a compatibility guarantee. + +This reference is for Coder deployment administrators who need to apply an external policy service to the agent loop. +It covers deployment configuration, the consumer contract, failure behavior, rollout, and dispatch auditing. + +Chat lifecycle hooks send events from the agent loop to 1 deployment-wide webhook endpoint. +The configured consumer can observe all 7 lifecycle events, add model or user context, restrict tools, replace mutable input, deny selected actions, or end a chat. + +> [!IMPORTANT] +> A consumer can block agent activity across the deployment. +> Start with an observe-only consumer and test failure recovery before enforcing policy. + +## Configure the deployment + +Enable the experiment first: + +```env +CODER_EXPERIMENTS=agent-lifecycle-hooks +``` + +Without the experiment, hook configuration is accepted but inactive: `coder server` logs a warning at startup and dispatches no hook events. +The experiment list is read at startup, so enabling or disabling it requires a `coder server` restart. + +Set the following deployment options on `coder server`. + +| Environment variable | CLI flag | Default | Requirement | +|---------------------------|-----------------------|---------|------------------------------------------------------------------------------------------------------------------------------------------| +| `CODER_CHAT_HOOK_URL` | `--chat-hook-url` | Empty | Use an `https` URL. Hooks are inactive when this value is empty. | +| `CODER_CHAT_HOOK_SECRET` | `--chat-hook-secret` | Empty | Required when the hook URL is set, at least 32 bytes of cryptographically random data. Coder uses this shared secret to sign HS256 JWTs. | +| `CODER_CHAT_HOOK_TIMEOUT` | `--chat-hook-timeout` | `1.5s` | Must be greater than `0` and no more than `5s`. The timeout applies to each request. | +| `CODER_CHAT_HOOK_ENABLED` | `--chat-hook-enabled` | `true` | Set to `false` to stop dispatching without removing the URL or secret. | + +Treat `CODER_CHAT_HOOK_ENABLED=false` as the break-glass control. +Changing deployment options requires the normal `coder server` configuration rollout for your installation. + +Use a dedicated secret and rotate it through your existing secret-management process. +Rotation is a hard cutover: Coder signs with exactly one secret, so dispatches fail until the consumer accepts the new value. +Rotate during a maintenance window, or temporarily set `CODER_CHAT_HOOK_ENABLED=false` for the cutover if blocked chats are worse than unreviewed ones for your deployment. +Coder requires the configured URL to use HTTPS. +A TLS terminator can forward the request to a consumer over plain HTTP on a trusted local network. +It must set `X-Forwarded-Proto: https` and either preserve the original `Host` header or carry it in `X-Forwarded-Host` for the SDK handler's audience check. +The SDK trusts those forwarded headers, so the audience check is only as strong as that proxy boundary: the proxy must strip or overwrite client-supplied forwarded headers, and the consumer must not be reachable except through the proxy. + +## Handle lifecycle events + +Coder sends an HTTP `POST` request for each event. +The JSON body contains `type`, `meta`, and event-specific `data`. +The `meta` object includes `dispatch_id`, `schema_version`, `chat_id`, `owner_id`, and optional workspace and turn IDs. +Events from subagent chats also carry `parent_chat_id` and `root_chat_id` so a consumer can correlate a subagent subtree with the user-facing conversation and apply the parent's policy context. +The current `schema_version` is `1`. + +| Event | When Coder sends it | Decision-relevant data | +|----------------------|---------------------------------------------------------------------|------------------------------------------------------------------------| +| `session_start` | A chat session starts, resumes, or clears | `source` (`startup`, `resume`, or `clear`) | +| `user_prompt_submit` | A user submits a prompt, or `spawn_agent` submits a subagent prompt | `prompt` and `parts` | +| `pre_tool_use` | Before a non-provider-executed tool runs | `tool_use_id`, `tool_name`, and `tool_input` | +| `post_tool_use` | After a non-provider-executed tool returns | `tool_use_id`, `tool_name`, and either `tool_response` or `tool_error` | +| `pre_compact` | Before Coder compacts chat context | No event-specific fields | +| `post_compact` | After Coder compacts chat context | No event-specific fields | +| `stop` | The model stops a turn | No event-specific fields | + +Provider-executed tools don't produce `pre_tool_use` or `post_tool_use` events because the provider executes them outside Coder's tool runtime. + +For `user_prompt_submit`, `prompt` concatenates the text parts of the message, and `parts` carries the full structured message exactly as Coder stores it and sends it to the model, including non-text parts such as file references. +A consumer that gates prompt content must inspect `parts`. + +### Verify each request + +Coder sends the JWT in the `Authorization: Bearer ` header. +A consumer must apply all of the following checks before it uses the body: + +- Accept only the `HS256` algorithm and verify the signature with `CODER_CHAT_HOOK_SECRET`. +- Check that `iss` is the Coder deployment ID associated with the secret. +- Check that `aud` exactly matches `CODER_CHAT_HOOK_URL`. +- Check `nbf` and reject expired tokens using `exp`. +- Check that `jti` equals the request `meta.dispatch_id`. +- Check that the JWT event `type` equals the body event `type`. +- Compute SHA-256 over the exact request body bytes and compare it with `body_sha256`. +- Check that the chat ID in `sub` matches `meta.chat_id`. + +The Go consumer SDK in `codersdk/agenthooks` implements the wire types, JWT verification, body binding, audience checks, and event routing. +Use `agenthooks.NewHTTPHandler` to build an `http.Handler` from callbacks for the events your consumer handles. +Pass `agenthooks.WithExpectedIssuer` with the deployment ID associated with the secret to enforce the `iss` check. +Without it, `NewHTTPHandler` accepts any non-empty `iss` signed with the shared secret, so use a secret dedicated to one deployment or always set the expected issuer. + +### Return a response + +Return any `2xx` status with an empty body for a no-op response. +An empty JSON object has the same effect. +If the response has a body, return a JSON object with these optional fields. + +| Field | Effect | +|-----------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `permission` | Allows or denies mutable input for `user_prompt_submit` and `pre_tool_use` only. | +| `model_context` | Adds text visible to the model. The value is limited to 16 KiB. The text never appears in the user's transcript, so users cannot tell that the consumer steered the model; the only record is the dispatch audit row. | +| `user_message` | Adds a system message visible to the user. | +| `allowed_tools` | Narrows the chat's hook tool policy. Omit the field to preserve the policy. The first non-omitted value initializes the policy; later values intersect with it, so an established policy only ever shrinks and `[]` permanently restricts all tools. To restore wider access, end the chat and start a new one. The policy applies from the next generation step onward: excluded tools are removed from the model request, a system notice tells the model that tool access is restricted by policy, and a pending call to an excluded dynamic tool resolves as inactive. The policy also covers provider tools that Coder requests, such as web search, but tools injected outside Coder's request, for example by an AI gateway, are beyond its reach. A subagent chat starts with its parent's policy intersected with any policy returned for the spawn prompt. | +| `end_chat` | Ends the chat when `true`. Ending a chat archives the chat identified by `meta.chat_id` together with any subagent chats it spawned. For an event from a subagent chat, this ends that subagent subtree and its parent chat continues. A `pre_tool_use` `end_chat` ends the chat before the tool runs, and the pending calls become synthetic cancellations. | + +The `permission.decision` value supports `allow` and `deny`. +The `ask` value isn't supported and causes the dispatch to fail closed. + +Permission rules depend on the event: + +- For `user_prompt_submit`, `allow` requires `input_override` in the exact form `{"prompt":"replacement text"}`. + Coder stores and sends the replacement prompt instead of the original prompt. +- For `pre_tool_use`, `allow` requires `input_override` containing the replacement tool input. + Coder persists the replacement with the tool call and executes the tool with it. + Nothing marks the call as rewritten in the chat, so the model may misattribute the changed behavior; a consumer that rewrites input should also return `user_message` explaining the change. +- For either event, `deny` blocks the input and must not include `input_override`. + A denied prompt isn't persisted, and a denied tool call becomes a synthetic error result so the model can choose another action. + A `user_prompt_submit` denial that also sets `end_chat` rejects the prompt and ends the existing chat. During chat creation there is no chat to end. + When the `clear` `session_start` emitted for a message edit sets `end_chat`, Coder ends the chat immediately and never dispatches `user_prompt_submit` for the edited content. +- For all other events, omit `permission`. + +## Plan failure recovery + +Lifecycle hooks are fail closed. +Coder treats a timeout, connection failure, non-`2xx` response, malformed response, or unsupported response field combination as a hook dispatch failure. +An in-progress chat enters the error state and records the dispatch ID in its error details. +If the first `user_prompt_submit` dispatch fails during chat creation, Coder rejects the request and doesn't create the chat. +If `post_tool_use` fails for a client-submitted tool result, Coder rejects the submission without committing the results, and the client can resubmit them after the consumer recovers. +If `post_tool_use` fails for a tool that Coder already executed, Coder commits the tool result first so the transcript reflects the completed side effect, then moves the chat to the error state. +An `end_chat` instruction that Coder already accepted from a successful dispatch in the same step takes precedence over the error state: if `post_tool_use` or `post_compact` fails after an accepted `end_chat`, Coder still ends the chat and records the failed dispatch. + +Dispatch precedes persistence, so a delivered event doesn't guarantee that the operation commits. +Coder checks admission before dispatching, but concurrent requests can still fail admission afterward, for example two sends racing for the last queue slot or duplicate submissions of the same tool results. +The consumer then observes an event for a request that Coder rejects, and the rejected request doesn't persist a prompt or tool result. +Treat events as attempt notifications rather than proof of a committed operation, and key idempotent tool-event processing on `tool_use_id`. + +Delivery is at least once. +Coder retries one connection failure per dispatch with the same JWT, so the consumer can receive the same `dispatch_id` more than once, and `session_start` response effects can repeat when a runner or process is replaced mid-turn. +A side-effectful consumer must deduplicate durably by `dispatch_id` and replay its previous response for a duplicate; rejecting duplicates breaks Coder's own retries. + +After the consumer is healthy, send another message to an existing errored chat to resume it. +Coder emits `session_start` with `source` set to `resume` when the agent loop starts again. +If the consumer continues blocking chat activity, set `CODER_CHAT_HOOK_ENABLED=false` and roll out the Coder deployment configuration before users retry. + +## Roll out enforcement in stages + +Use the following rollout sequence: + +1. Deploy a consumer that verifies every request, logs the event and identifiers, and always returns a `2xx` status with an empty response. +2. Configure the hook URL, secret, and timeout on a test deployment. +3. Exercise normal chats, tool calls, compaction, consumer timeouts, and consumer restarts. +4. Review `chat_hook_dispatches` for event coverage and unexpected failures. +5. Add policy responses for a narrow event or tool set. +6. Expand enforcement after the dispatch audit shows the expected decisions and failure rate. + +Keep the break-glass procedure available throughout the rollout. + +## Start from the reference consumer + +The reference consumer at `scripts/agenthooks-server` uses `agenthooks.NewHTTPHandler` and logs 1 JSON object for each event. +Log-only mode returns an empty response for every verified event. +With log-only mode disabled, the optional example flags can deny tool names by regular expression or replace matching prompt text before the agent loop uses it. +Coder retains the original prompt in the dispatch audit row. + +Run the consumer from a Coder source checkout: + +```sh +CODER_AGENTHOOKS_SECRET='' \ + go run ./scripts/agenthooks-server \ + --listen 127.0.0.1:8081 \ + --log-only=true +``` + +The reference server accepts optional TLS certificate and key paths. +For local testing with plain HTTP, place an HTTPS reverse proxy in front of it because `CODER_CHAT_HOOK_URL` accepts only HTTPS URLs. +Run `go run ./scripts/agenthooks-server --help` for all flags and environment variable names. + +## Audit dispatches + +Coder records every attempted dispatch in the `chat_hook_dispatches` database table. +Each row includes the event, chat and turn identifiers, tool-use ID when present, timestamps, HTTP status, result class, permission decision, input override, response context, and error details. +Use the `dispatch_id` from the request or chat error as the row ID when correlating consumer logs with Coder state. + +The table can contain prompts, tool input, response context, and user messages. +When a response rewrites a prompt or tool input, the row keeps the original value alongside the override, so rewriting doesn't remove the original content from the deployment. +Rows have no foreign key to chats: they can outlive deleted chats, and denied create attempts leave rows for chats that never existed. +Apply the same access controls and database protection that you use for other sensitive chat data. +The `dbpurge` service removes dispatch rows after 90 days to bound table growth; the retention period isn't configurable. +Dispatch rows are also removed together with their chat when chat retention purges it, so a shorter chat retention bounds dispatch payloads too. diff --git a/docs/manifest.json b/docs/manifest.json index ffa3db636851a..a0b8bfcb17d19 100644 --- a/docs/manifest.json +++ b/docs/manifest.json @@ -429,6 +429,12 @@ "description": "Learn what usage telemetry Coder collects", "path": "./admin/setup/telemetry.md" }, + { + "title": "Chat lifecycle hooks", + "description": "Configure an external policy service for chat agent lifecycle events", + "path": "./admin/setup/chat-lifecycle-hooks.md", + "state": ["early access"] + }, { "title": "Data Retention", "description": "Configure data retention policies for database tables", diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 1b89aeb459dfc..bafca4ffaf533 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -225,12 +225,12 @@ Status Code **200** #### Enumerated Values -| Property | Value(s) | -|---------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `client_type` | `api`, `ui` | -| `kind` | `auth`, `config`, `content_filter`, `generic`, `instruction_file`, `mcp_config`, `mcp_server`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `skill`, `stream_silence_timeout`, `timeout`, `usage_limit` | -| `status` | `error`, `excluded`, `interrupting`, `invalid`, `ok`, `oversize`, `requires_action`, `running`, `unreadable`, `waiting` | -| `plan_mode` | `plan` | +| Property | Value(s) | +|---------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `client_type` | `api`, `ui` | +| `kind` | `auth`, `config`, `content_filter`, `generic`, `hook_dispatch_failed`, `instruction_file`, `mcp_config`, `mcp_server`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `skill`, `stream_silence_timeout`, `timeout`, `usage_limit` | +| `status` | `error`, `excluded`, `interrupting`, `invalid`, `ok`, `oversize`, `requires_action`, `running`, `unreadable`, `waiting` | +| `plan_mode` | `plan` | To perform this operation, you must be authenticated. [Learn more](authentication.md). @@ -1800,6 +1800,7 @@ Experimental: this endpoint is subject to change. ```json { + "ended": true, "message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", "content": [ @@ -1881,6 +1882,89 @@ Experimental: this endpoint is subject to change. "total_tokens": 0 } }, + "messages": [ + { + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "content": [ + { + "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" + } + ], + "created_at": "2019-08-24T14:15:22Z", + "created_by": "ee824cad-d7a6-4f48-87dc-e8461a9201c4", + "id": 0, + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", + "role": "system", + "usage": { + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "context_limit": 0, + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + "total_tokens": 0 + } + } + ], "queued": true, "queued_message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", @@ -2016,6 +2100,7 @@ Experimental: this endpoint is subject to change. ```json { + "ended": true, "message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", "content": [ @@ -2097,6 +2182,89 @@ Experimental: this endpoint is subject to change. "total_tokens": 0 } }, + "messages": [ + { + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "content": [ + { + "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" + } + ], + "created_at": "2019-08-24T14:15:22Z", + "created_by": "ee824cad-d7a6-4f48-87dc-e8461a9201c4", + "id": 0, + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", + "role": "system", + "usage": { + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "context_limit": 0, + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + "total_tokens": 0 + } + } + ], "warnings": [ "string" ] diff --git a/docs/reference/api/general.md b/docs/reference/api/general.md index 4715ff82458cb..3a8abd793cccd 100644 --- a/docs/reference/api/general.md +++ b/docs/reference/api/general.md @@ -232,7 +232,23 @@ curl -X GET http://coder-server:8080/api/v2/deployment/config \ }, "chat": { "acquire_batch_size": 0, - "debug_logging_enabled": true + "debug_logging_enabled": true, + "hook_enabled": true, + "hook_secret": "string", + "hook_timeout": 0, + "hook_url": { + "forceQuery": true, + "fragment": "string", + "host": "string", + "omitHost": true, + "opaque": "string", + "path": "string", + "rawFragment": "string", + "rawPath": "string", + "rawQuery": "string", + "scheme": "string", + "user": {} + } } }, "allow_workspace_renames": true, diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index ddba40ad34a87..efffec87dd4da 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -1020,7 +1020,23 @@ }, "chat": { "acquire_batch_size": 0, - "debug_logging_enabled": true + "debug_logging_enabled": true, + "hook_enabled": true, + "hook_secret": "string", + "hook_timeout": 0, + "hook_url": { + "forceQuery": true, + "fragment": "string", + "host": "string", + "omitHost": true, + "opaque": "string", + "path": "string", + "rawFragment": "string", + "rawPath": "string", + "rawQuery": "string", + "scheme": "string", + "user": {} + } } } ``` @@ -2390,16 +2406,36 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in ```json { "acquire_batch_size": 0, - "debug_logging_enabled": true + "debug_logging_enabled": true, + "hook_enabled": true, + "hook_secret": "string", + "hook_timeout": 0, + "hook_url": { + "forceQuery": true, + "fragment": "string", + "host": "string", + "omitHost": true, + "opaque": "string", + "path": "string", + "rawFragment": "string", + "rawPath": "string", + "rawQuery": "string", + "scheme": "string", + "user": {} + } } ``` ### Properties -| Name | Type | Required | Restrictions | Description | -|-------------------------|---------|----------|--------------|-------------| -| `acquire_batch_size` | integer | false | | | -| `debug_logging_enabled` | boolean | false | | | +| Name | Type | Required | Restrictions | Description | +|-------------------------|----------------------------|----------|--------------|-------------| +| `acquire_batch_size` | integer | false | | | +| `debug_logging_enabled` | boolean | false | | | +| `hook_enabled` | boolean | false | | | +| `hook_secret` | string | false | | | +| `hook_timeout` | integer | false | | | +| `hook_url` | [serpent.URL](#serpenturl) | false | | | ## codersdk.ChatContext @@ -2622,9 +2658,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in #### Enumerated Values -| Value(s) | -|-------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `auth`, `config`, `content_filter`, `generic`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `stream_silence_timeout`, `timeout`, `usage_limit` | +| Value(s) | +|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `auth`, `config`, `content_filter`, `generic`, `hook_dispatch_failed`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `stream_silence_timeout`, `timeout`, `usage_limit` | ## codersdk.ChatFileMetadata @@ -4465,6 +4501,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in ```json { + "ended": true, "message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", "content": [ @@ -4546,6 +4583,89 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in "total_tokens": 0 } }, + "messages": [ + { + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "content": [ + { + "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" + } + ], + "created_at": "2019-08-24T14:15:22Z", + "created_by": "ee824cad-d7a6-4f48-87dc-e8461a9201c4", + "id": 0, + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", + "role": "system", + "usage": { + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "context_limit": 0, + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + "total_tokens": 0 + } + } + ], "queued": true, "queued_message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", @@ -4625,12 +4745,14 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in ### Properties -| Name | Type | Required | Restrictions | Description | -|------------------|----------------------------------------------------------|----------|--------------|-------------| -| `message` | [codersdk.ChatMessage](#codersdkchatmessage) | false | | | -| `queued` | boolean | false | | | -| `queued_message` | [codersdk.ChatQueuedMessage](#codersdkchatqueuedmessage) | false | | | -| `warnings` | array of string | false | | | +| Name | Type | Required | Restrictions | Description | +|------------------|----------------------------------------------------------|----------|--------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `ended` | boolean | false | | | +| `message` | [codersdk.ChatMessage](#codersdkchatmessage) | false | | | +| `messages` | array of [codersdk.ChatMessage](#codersdkchatmessage) | false | | Messages contains all user-visible messages inserted by an immediate send, in insertion order with the user's message last. Clients should upsert the full batch because hooks may prepend notices. Empty for queued or ended sends. | +| `queued` | boolean | false | | | +| `queued_message` | [codersdk.ChatQueuedMessage](#codersdkchatqueuedmessage) | false | | | +| `warnings` | array of string | false | | | ## codersdk.CreateChatRequest @@ -5729,7 +5851,23 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o }, "chat": { "acquire_batch_size": 0, - "debug_logging_enabled": true + "debug_logging_enabled": true, + "hook_enabled": true, + "hook_secret": "string", + "hook_timeout": 0, + "hook_url": { + "forceQuery": true, + "fragment": "string", + "host": "string", + "omitHost": true, + "opaque": "string", + "path": "string", + "rawFragment": "string", + "rawPath": "string", + "rawQuery": "string", + "scheme": "string", + "user": {} + } } }, "allow_workspace_renames": true, @@ -6339,7 +6477,23 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o }, "chat": { "acquire_batch_size": 0, - "debug_logging_enabled": true + "debug_logging_enabled": true, + "hook_enabled": true, + "hook_secret": "string", + "hook_timeout": 0, + "hook_url": { + "forceQuery": true, + "fragment": "string", + "host": "string", + "omitHost": true, + "opaque": "string", + "path": "string", + "rawFragment": "string", + "rawPath": "string", + "rawQuery": "string", + "scheme": "string", + "user": {} + } } }, "allow_workspace_renames": true, @@ -7046,6 +7200,7 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o ```json { + "ended": true, "message": { "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", "content": [ @@ -7127,6 +7282,89 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o "total_tokens": 0 } }, + "messages": [ + { + "chat_id": "efc9fe20-a1e5-4a8c-9c48-f1b30c1e4f86", + "content": [ + { + "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" + } + ], + "created_at": "2019-08-24T14:15:22Z", + "created_by": "ee824cad-d7a6-4f48-87dc-e8461a9201c4", + "id": 0, + "model_config_id": "f5fb4d91-62ca-4377-9ee6-5d43ba00d205", + "role": "system", + "usage": { + "cache_creation_tokens": 0, + "cache_read_tokens": 0, + "context_limit": 0, + "input_tokens": 0, + "output_tokens": 0, + "reasoning_tokens": 0, + "total_tokens": 0 + } + } + ], "warnings": [ "string" ] @@ -7135,10 +7373,12 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o ### Properties -| Name | Type | Required | Restrictions | Description | -|------------|----------------------------------------------|----------|--------------|-------------| -| `message` | [codersdk.ChatMessage](#codersdkchatmessage) | false | | | -| `warnings` | array of string | false | | | +| Name | Type | Required | Restrictions | Description | +|------------|-------------------------------------------------------|----------|--------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `ended` | boolean | false | | | +| `message` | [codersdk.ChatMessage](#codersdkchatmessage) | false | | | +| `messages` | array of [codersdk.ChatMessage](#codersdkchatmessage) | false | | Messages holds every user-visible message the edit inserted, in insertion order with the replacement message last. Lifecycle hooks may prepend notices, so clients must upsert all of them rather than only Message. Empty for ended edits. | +| `warnings` | array of string | false | | | ## codersdk.Entitlement @@ -7218,9 +7458,9 @@ CreateWorkspaceRequest provides options for creating a new workspace. Only one o #### Enumerated Values -| Value(s) | -|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `ai-gateway-cost-control`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `minimum-implicit-member`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-usage` | +| Value(s) | +|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `agent-lifecycle-hooks`, `ai-gateway-cost-control`, `auto-fill-parameters`, `chat-advisor`, `chat-virtual-desktop`, `example`, `mcp-server-http`, `minimum-implicit-member`, `nats_pubsub`, `notifications`, `oauth2`, `workspace-build-updates`, `workspace-usage` | ## codersdk.ExternalAPIKeyScopes diff --git a/docs/reference/cli/server.md b/docs/reference/cli/server.md index f4ae058d8781f..bd7f5215028de 100644 --- a/docs/reference/cli/server.md +++ b/docs/reference/cli/server.md @@ -1745,6 +1745,47 @@ Hide AI tasks from the dashboard. Force chat debug logging on for every chat, bypassing the runtime admin and user opt-in settings. +### --chat-hook-url + +| | | +|-------------|-----------------------------------| +| Type | url | +| Environment | $CODER_CHAT_HOOK_URL | +| YAML | chat.hookURL | + +HTTPS URL to receive chat agent lifecycle hook events. Hooks are disabled when unset. Requires the agent-lifecycle-hooks experiment. + +### --chat-hook-secret + +| | | +|-------------|--------------------------------------| +| Type | string | +| Environment | $CODER_CHAT_HOOK_SECRET | + +Shared secret used to sign chat agent lifecycle hook JWTs. + +### --chat-hook-timeout + +| | | +|-------------|---------------------------------------| +| Type | duration | +| Environment | $CODER_CHAT_HOOK_TIMEOUT | +| YAML | chat.hookTimeout | +| Default | 1.5s | + +Maximum time to wait for a chat agent lifecycle hook response. + +### --chat-hook-enabled + +| | | +|-------------|---------------------------------------| +| Type | bool | +| Environment | $CODER_CHAT_HOOK_ENABLED | +| YAML | chat.hookEnabled | +| Default | true | + +Whether to dispatch chat agent lifecycle hooks when a hook URL is configured. Requires the agent-lifecycle-hooks experiment. + ### --ai-gateway-enabled | | | diff --git a/enterprise/audit/table.go b/enterprise/audit/table.go index 50532c6b72876..1703716a74a32 100644 --- a/enterprise/audit/table.go +++ b/enterprise/audit/table.go @@ -472,6 +472,7 @@ var auditableResourcesTypes = map[any]map[string]Action{ "context_dirty_since": ActionIgnore, // Agent-pushed context snapshot state. "context_dirty_resources": ActionIgnore, // Agent-pushed context snapshot state. "context_error": ActionIgnore, // Agent-pushed context snapshot state. + "hook_allowed_tools": ActionIgnore, // System-managed hook output. "dynamic_tools": ActionIgnore, // Internal lifecycle. "plan_mode": ActionIgnore, // Can flip back and forth during a session. "client_type": ActionIgnore, // Set at creation. diff --git a/enterprise/cli/testdata/coder_server_--help.golden b/enterprise/cli/testdata/coder_server_--help.golden index 369b2fe72c805..e1962f4046499 100644 --- a/enterprise/cli/testdata/coder_server_--help.golden +++ b/enterprise/cli/testdata/coder_server_--help.golden @@ -282,6 +282,20 @@ Configure the background chat processing daemon. Force chat debug logging on for every chat, bypassing the runtime admin and user opt-in settings. + --chat-hook-enabled bool, $CODER_CHAT_HOOK_ENABLED (default: true) + Whether to dispatch chat agent lifecycle hooks when a hook URL is + configured. Requires the agent-lifecycle-hooks experiment. + + --chat-hook-secret string, $CODER_CHAT_HOOK_SECRET + Shared secret used to sign chat agent lifecycle hook JWTs. + + --chat-hook-timeout duration, $CODER_CHAT_HOOK_TIMEOUT (default: 1.5s) + Maximum time to wait for a chat agent lifecycle hook response. + + --chat-hook-url url, $CODER_CHAT_HOOK_URL + HTTPS URL to receive chat agent lifecycle hook events. Hooks are + disabled when unset. Requires the agent-lifecycle-hooks experiment. + CLIENT OPTIONS: These options change the behavior of how clients interact with the Coder. Clients include the Coder CLI, Coder Desktop, IDE extensions, and the web UI. diff --git a/scripts/agenthooks-server/main.go b/scripts/agenthooks-server/main.go new file mode 100644 index 0000000000000..796b0d98fdb3b --- /dev/null +++ b/scripts/agenthooks-server/main.go @@ -0,0 +1,236 @@ +// agenthooks-server is a reference consumer that logs verified lifecycle +// events as JSON. +package main + +import ( + "context" + "encoding/json" + "errors" + "flag" + "fmt" + "net/http" + "os" + "os/signal" + "regexp" + "strconv" + "sync" + "syscall" + "time" + + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/codersdk/agenthooks" +) + +type config struct { + listen string + secret string + issuer string + tlsCert string + tlsKey string + logOnly bool + denyToolPattern string + redactPrompt string +} + +type eventLog struct { + Event agenthooks.EventType `json:"event"` + DispatchID string `json:"dispatch_id"` + ChatID string `json:"chat_id"` + TurnID string `json:"turn_id,omitempty"` + ToolUseID string `json:"tool_use_id,omitempty"` + ToolName string `json:"tool_name,omitempty"` + Source string `json:"source,omitempty"` + Prompt string `json:"prompt,omitempty"` + ToolInput json.RawMessage `json:"tool_input,omitempty"` + ToolOutput json.RawMessage `json:"tool_output,omitempty"` + ToolError string `json:"tool_error,omitempty"` +} + +func main() { + if err := run(); err != nil { + _, _ = fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } +} + +func run() error { + cfg, err := parseFlags() + if err != nil { + return err + } + if cfg.secret == "" { + return xerrors.New("secret is required through --secret or CODER_AGENTHOOKS_SECRET") + } + if (cfg.tlsCert == "") != (cfg.tlsKey == "") { + return xerrors.New("TLS certificate and key must be configured together") + } + + var denyTool *regexp.Regexp + if cfg.denyToolPattern != "" { + denyTool, err = regexp.Compile(cfg.denyToolPattern) + if err != nil { + return xerrors.Errorf("compile deny tool pattern: %w", err) + } + } + var redactPrompt *regexp.Regexp + if cfg.redactPrompt != "" { + redactPrompt, err = regexp.Compile(cfg.redactPrompt) + if err != nil { + return xerrors.Errorf("compile redact prompt pattern: %w", err) + } + } + + var logMu sync.Mutex + encoder := json.NewEncoder(os.Stdout) + logEvent := func(event eventLog) error { + logMu.Lock() + defer logMu.Unlock() + if err := encoder.Encode(event); err != nil { + return xerrors.Errorf("encode event: %w", err) + } + return nil + } + baseEvent := func(event agenthooks.EventType, meta agenthooks.Meta) eventLog { + entry := eventLog{ + Event: event, + DispatchID: meta.DispatchID.String(), + ChatID: meta.ChatID.String(), + } + if meta.TurnID != nil { + entry.TurnID = meta.TurnID.String() + } + return entry + } + + hooks := agenthooks.Hooks{ + SessionStart: func(_ context.Context, meta agenthooks.Meta, data agenthooks.SessionStartData) (agenthooks.Response, error) { + entry := baseEvent(agenthooks.EventSessionStart, meta) + entry.Source = data.Source + return agenthooks.Response{}, logEvent(entry) + }, + UserPromptSubmit: func(_ context.Context, meta agenthooks.Meta, data agenthooks.UserPromptSubmitData) (agenthooks.Response, error) { + entry := baseEvent(agenthooks.EventUserPromptSubmit, meta) + entry.Prompt = data.Prompt + matches := redactPrompt != nil && redactPrompt.MatchString(data.Prompt) + if matches { + entry.Prompt = redactPrompt.ReplaceAllString(data.Prompt, "[REDACTED]") + } + if err := logEvent(entry); err != nil { + return agenthooks.Response{}, err + } + // Log-only mode still redacts the log entry above; it only + // suppresses the prompt override response. + if cfg.logOnly || !matches { + return agenthooks.Response{}, nil + } + override, err := json.Marshal(map[string]string{"prompt": entry.Prompt}) + if err != nil { + return agenthooks.Response{}, xerrors.Errorf("marshal prompt override: %w", err) + } + return agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionAllow, + InputOverride: override, + }}, nil + }, + PreToolUse: func(_ context.Context, meta agenthooks.Meta, data agenthooks.PreToolUseData) (agenthooks.Response, error) { + entry := baseEvent(agenthooks.EventPreToolUse, meta) + entry.ToolUseID = data.ToolUseID + entry.ToolName = data.ToolName + entry.ToolInput = data.ToolInput + if err := logEvent(entry); err != nil { + return agenthooks.Response{}, err + } + if cfg.logOnly || denyTool == nil || !denyTool.MatchString(data.ToolName) { + return agenthooks.Response{}, nil + } + return agenthooks.Response{Permission: &agenthooks.Permission{ + Decision: agenthooks.PermissionDeny, + Reason: "tool name matched the configured deny pattern", + }}, nil + }, + PostToolUse: func(_ context.Context, meta agenthooks.Meta, data agenthooks.PostToolUseData) (agenthooks.Response, error) { + entry := baseEvent(agenthooks.EventPostToolUse, meta) + entry.ToolUseID = data.ToolUseID + entry.ToolName = data.ToolName + entry.ToolOutput = data.ToolResponse + entry.ToolError = data.ToolError + return agenthooks.Response{}, logEvent(entry) + }, + PreCompact: func(_ context.Context, meta agenthooks.Meta, _ agenthooks.PreCompactData) (agenthooks.Response, error) { + return agenthooks.Response{}, logEvent(baseEvent(agenthooks.EventPreCompact, meta)) + }, + PostCompact: func(_ context.Context, meta agenthooks.Meta, _ agenthooks.PostCompactData) (agenthooks.Response, error) { + return agenthooks.Response{}, logEvent(baseEvent(agenthooks.EventPostCompact, meta)) + }, + Stop: func(_ context.Context, meta agenthooks.Meta, _ agenthooks.StopData) (agenthooks.Response, error) { + return agenthooks.Response{}, logEvent(baseEvent(agenthooks.EventStop, meta)) + }, + } + + var handlerOpts []agenthooks.HandlerOption + if cfg.issuer != "" { + handlerOpts = append(handlerOpts, agenthooks.WithExpectedIssuer(cfg.issuer)) + } + handler := agenthooks.NewHTTPHandler([]byte(cfg.secret), hooks, handlerOpts...) + server := &http.Server{ + Addr: cfg.listen, + Handler: handler, + ReadHeaderTimeout: 10 * time.Second, + } + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + go func() { + <-ctx.Done() + _ = server.Close() + }() + + _, _ = fmt.Fprintf(os.Stderr, "Agent hooks server listening on %s\n", cfg.listen) + if cfg.tlsCert != "" { + err = server.ListenAndServeTLS(cfg.tlsCert, cfg.tlsKey) + } else { + err = server.ListenAndServe() + } + if err != nil && !errors.Is(err, http.ErrServerClosed) { + return xerrors.Errorf("serve lifecycle hooks: %w", err) + } + return nil +} + +func parseFlags() (config, error) { + logOnly, err := envBool("CODER_AGENTHOOKS_LOG_ONLY", true) + if err != nil { + return config{}, err + } + var cfg config + cfg.logOnly = logOnly + flag.StringVar(&cfg.listen, "listen", envOrDefault("CODER_AGENTHOOKS_LISTEN", "127.0.0.1:8081"), "Listen address (CODER_AGENTHOOKS_LISTEN)") + flag.StringVar(&cfg.secret, "secret", os.Getenv("CODER_AGENTHOOKS_SECRET"), "Shared HS256 secret, required (CODER_AGENTHOOKS_SECRET)") + flag.StringVar(&cfg.issuer, "issuer", os.Getenv("CODER_AGENTHOOKS_ISSUER"), "Expected iss claim, normally the Coder deployment ID (CODER_AGENTHOOKS_ISSUER)") + flag.StringVar(&cfg.tlsCert, "tls-cert", os.Getenv("CODER_AGENTHOOKS_TLS_CERT"), "TLS certificate path (CODER_AGENTHOOKS_TLS_CERT)") + flag.StringVar(&cfg.tlsKey, "tls-key", os.Getenv("CODER_AGENTHOOKS_TLS_KEY"), "TLS private key path (CODER_AGENTHOOKS_TLS_KEY)") + flag.BoolVar(&cfg.logOnly, "log-only", cfg.logOnly, "Return an empty response for every event (CODER_AGENTHOOKS_LOG_ONLY)") + flag.StringVar(&cfg.denyToolPattern, "deny-tool-pattern", os.Getenv("CODER_AGENTHOOKS_DENY_TOOL_PATTERN"), "Example regexp for denied tool names (CODER_AGENTHOOKS_DENY_TOOL_PATTERN)") + flag.StringVar(&cfg.redactPrompt, "redact-prompt-pattern", os.Getenv("CODER_AGENTHOOKS_REDACT_PROMPT_PATTERN"), "Example regexp to redact in prompts (CODER_AGENTHOOKS_REDACT_PROMPT_PATTERN)") + flag.Parse() + return cfg, nil +} + +func envOrDefault(name, fallback string) string { + if value := os.Getenv(name); value != "" { + return value + } + return fallback +} + +func envBool(name string, fallback bool) (bool, error) { + value := os.Getenv(name) + if value == "" { + return fallback, nil + } + parsed, err := strconv.ParseBool(value) + if err != nil { + return false, xerrors.Errorf("parse %s: %w", name, err) + } + return parsed, nil +} diff --git a/scripts/apitypings/main.go b/scripts/apitypings/main.go index 77c648a050b3c..eca00aaf9df57 100644 --- a/scripts/apitypings/main.go +++ b/scripts/apitypings/main.go @@ -26,6 +26,7 @@ func main() { generateDirectories := map[string]string{ "github.com/coder/coder/v2/codersdk": "", "github.com/coder/coder/v2/coderd/healthcheck/health": "Health", + "github.com/coder/coder/v2/codersdk/agenthooks": "AgentHook", "github.com/coder/coder/v2/codersdk/healthsdk": "", } for dir, prefix := range generateDirectories { @@ -78,6 +79,7 @@ func TSMutations(ts *guts.Typescript) { config.NotNullMaps, FixSerpentStruct, DiscriminatedChatMessagePart, + AgentHookRawMessages, // Prefer enums as types config.EnumAsTypes, // Enum list generator @@ -146,6 +148,43 @@ func TypeMappings(gen *guts.GoParser) error { return nil } +// AgentHookRawMessages maps agent-hook raw JSON fields to unknown instead of +// the global object type. +func AgentHookRawMessages(ts *guts.Typescript) { + if _, ok := ts.Node("AgentHookRequest"); !ok { + return + } + unknown := bindings.KeywordUnknown + fields := map[string]string{ + "AgentHookRequest": "data", + "AgentHookUserPromptSubmitData": "parts", + "AgentHookPreToolUseData": "tool_input", + "AgentHookPostToolUseData": "tool_response", + "AgentHookPermission": "input_override", + } + for typeName, fieldName := range fields { + node, ok := ts.Node(typeName) + if !ok { + panic(fmt.Sprintf("agent hook type %q was not generated", typeName)) + } + iface, ok := node.(*bindings.Interface) + if !ok { + panic(fmt.Sprintf("agent hook type %q is not an interface", typeName)) + } + found := false + for _, field := range iface.Fields { + if field.Name == fieldName { + field.Type = &unknown + found = true + break + } + } + if !found { + panic(fmt.Sprintf("agent hook field %q.%s was not generated", typeName, fieldName)) + } + } +} + // DiscriminatedChatMessagePart splits the flat ChatMessagePart // interface into a discriminated union of per-type sub-interfaces. // Each sub-interface narrows the `type` field to a string literal diff --git a/scripts/metricsdocgen/generated_metrics b/scripts/metricsdocgen/generated_metrics index 65e613adf6b8e..a4b45ef0f2d5d 100644 --- a/scripts/metricsdocgen/generated_metrics +++ b/scripts/metricsdocgen/generated_metrics @@ -277,6 +277,21 @@ coderd_chatd_chats{state=""} 0 # HELP coderd_chatd_compaction_total Total compaction outcomes (only recorded when compaction was triggered or failed). # TYPE coderd_chatd_compaction_total counter coderd_chatd_compaction_total{provider="",model="",result=""} 0 +# HELP coderd_chatd_hook_context_size_bytes Lifecycle hook model context response size in bytes. +# TYPE coderd_chatd_hook_context_size_bytes histogram +coderd_chatd_hook_context_size_bytes{event=""} 0 +# HELP coderd_chatd_hook_decisions_total Total lifecycle hook permission decisions by event and decision. +# TYPE coderd_chatd_hook_decisions_total counter +coderd_chatd_hook_decisions_total{event="",decision=""} 0 +# HELP coderd_chatd_hook_dispatch_seconds Lifecycle hook dispatch duration in seconds. +# TYPE coderd_chatd_hook_dispatch_seconds histogram +coderd_chatd_hook_dispatch_seconds{event=""} 0 +# HELP coderd_chatd_hook_dispatches_total Total lifecycle hook dispatches by event and result. +# TYPE coderd_chatd_hook_dispatches_total counter +coderd_chatd_hook_dispatches_total{event="",result=""} 0 +# HELP coderd_chatd_hook_input_overrides_total Total lifecycle hook input overrides by event. +# TYPE coderd_chatd_hook_input_overrides_total counter +coderd_chatd_hook_input_overrides_total{event=""} 0 # HELP coderd_chatd_message_count Number of messages in the prompt per LLM request. # TYPE coderd_chatd_message_count histogram coderd_chatd_message_count{provider="",model=""} 0 diff --git a/site/src/api/queries/chatMessageEdits.ts b/site/src/api/queries/chatMessageEdits.ts index 2fbefa12741f1..66824fd0871bf 100644 --- a/site/src/api/queries/chatMessageEdits.ts +++ b/site/src/api/queries/chatMessageEdits.ts @@ -117,28 +117,34 @@ export const projectEditedConversationIntoCache = ({ export const reconcileEditedMessageInCache = ({ currentData, optimisticMessageId, - responseMessage, + responseMessages, }: { currentData: InfiniteData | undefined; optimisticMessageId: number; - responseMessage: TypesGen.ChatMessage; + // Every message the edit inserted, replacement last. Hook notices + // precede the replacement with lower IDs and must land in the + // cache too, or a stream reconnect keyed on the highest cached ID + // would skip them. + responseMessages: readonly TypesGen.ChatMessage[]; }): InfiniteData | undefined => { - if (!currentData?.pages?.length) { + if (!currentData?.pages?.length || responseMessages.length === 0) { return currentData; } + const responseIDs = new Set(responseMessages.map((message) => message.id)); const replacedPages = currentData.pages.map((page, pageIndex) => { const preservedMessages = page.messages.filter( (message) => - message.id !== optimisticMessageId && message.id !== responseMessage.id, + message.id !== optimisticMessageId && !responseIDs.has(message.id), ); if (pageIndex !== 0) { return { ...page, messages: preservedMessages }; } - return { - ...page, - messages: upsertFirstPageMessage(preservedMessages, responseMessage), - }; + let messages = preservedMessages; + for (const responseMessage of responseMessages) { + messages = upsertFirstPageMessage(messages, responseMessage); + } + return { ...page, messages }; }); return { diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 19e3f92fc23f6..750bf83797a42 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -1074,6 +1074,51 @@ describe("mutation invalidation scope", () => { ).toBe(true); }); + it("editChatMessage onSuccess without a message restores the snapshot and invalidates", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const messages = [3, 2, 1].map((id) => makeMsg(chatId, id)); + const previousData: InfMessages = { + pages: [{ messages, queued_messages: [], has_more: false }], + pageParams: [undefined], + }; + + queryClient.setQueryData(chatMessagesKey(chatId), { + pages: [ + { + messages: [makeMsg(chatId, 2), makeMsg(chatId, 1)], + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + queryClient.setQueryData(chatKey(chatId), makeChat(chatId)); + + const mutation = editChatMessage(queryClient, chatId); + mutation.onSuccess( + { ended: true }, + { messageId: 2, req: editReq }, + { previousData }, + ); + + await new Promise((r) => setTimeout(r, 0)); + + expect( + queryClient.getQueryData(chatMessagesKey(chatId)), + "hook-ended edit should restore the pre-edit snapshot", + ).toEqual(previousData); + const messagesState = queryClient.getQueryState(chatMessagesKey(chatId)); + expect( + messagesState?.isInvalidated, + "chatMessagesKey should be invalidated when the edit did not commit", + ).toBe(true); + expect( + queryClient.getQueryState(chatKey(chatId))?.isInvalidated, + "chatKey should be invalidated so an archived chat cannot stay stale", + ).toBe(true); + }); + // Shared type for the infinite messages cache shape used by // editChatMessage tests below. type InfMessages = { @@ -1265,8 +1310,9 @@ describe("mutation invalidation scope", () => { }, ); mutation.onSuccess( - { message: responseMessage }, + { message: responseMessage, ended: false }, { messageId: 3, optimisticMessage, req: editReq }, + undefined, ); const data = queryClient.getQueryData(chatMessagesKey(chatId)); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index fdff89970c08f..da421d0aadc74 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -1406,18 +1406,47 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ queryKey: chatMessagesKey(chatId), exact: true, }); + // Hook denial may archive the chat even though the edit fails. + // Refresh lists in case the chat-watch delete event is missed. + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ + queryKey: chatsByWorkspaceKeyPrefix, + }); }, onSuccess: ( response: TypesGen.EditChatMessageResponse, variables: EditChatMessageMutationArgs, + context: EditChatMessageMutationContext | undefined, ) => { + // A hook-ended edit has no committed message, so restore and refresh the transcript. + const responseMessage = response.message; + if (!responseMessage) { + if (context?.previousData) { + queryClient.setQueryData(chatMessagesKey(chatId), context.previousData); + } + void queryClient.invalidateQueries({ + queryKey: chatMessagesKey(chatId), + exact: true, + }); + // Refresh the chat detail and lists in case the WebSocket + // delete event was missed. + void queryClient.invalidateQueries({ + queryKey: chatKey(chatId), + exact: true, + }); + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ + queryKey: chatsByWorkspaceKeyPrefix, + }); + return; + } queryClient.setQueryData< InfiniteData | undefined >(chatMessagesKey(chatId), (current) => reconcileEditedMessageInCache({ currentData: current, optimisticMessageId: variables.messageId, - responseMessage: response.message, + responseMessages: response.messages ?? [responseMessage], }), ); }, diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 75a8d7c8d445d..f0777f4bbc634 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1217,6 +1217,173 @@ export interface AgentFirewallSessionLogsResponse { readonly results: readonly AgentFirewallLog[]; } +// From agenthooks/types.go +/** + * ChatRef identifies the chat a lifecycle hook event refers to. Embedded + * structs flatten in JSON, so it adds no nesting on the wire. + */ +export interface AgentHookChatRef { + readonly chat_id: string; + readonly owner_id: string; + readonly workspace_id?: string; + readonly turn_id?: string; + readonly parent_chat_id?: string; + /** + * RootChatID groups a subagent subtree with its user-facing conversation. + * Unset for top-level chats. + */ + readonly root_chat_id?: string; +} + +// From agenthooks/types.go +/** + * Claims describes the JWT minted by coderd for a lifecycle hook dispatch. + */ +export interface AgentHookClaims { + readonly iss: string; + readonly sub: string; + readonly aud: string; + readonly iat: number; + readonly nbf: number; + readonly exp: number; + readonly jti: string; + readonly type: AgentHookEventType; + readonly body_sha256: string; +} + +// From agenthooks/types.go +export type AgentHookEventType = + | "post_compact" + | "post_tool_use" + | "pre_compact" + | "pre_tool_use" + | "session_start" + | "stop" + | "user_prompt_submit"; + +export const AgentHookEventTypes: AgentHookEventType[] = [ + "post_compact", + "post_tool_use", + "pre_compact", + "pre_tool_use", + "session_start", + "stop", + "user_prompt_submit", +]; + +// From agenthooks/http.go +/** + * Hooks lets a consumer implement only the lifecycle events it uses. + */ +export interface AgentHookHooks { + // Function type detected, and unsupported. Leaving the type as unknown + readonly SessionStart: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly UserPromptSubmit: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly PreToolUse: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly PostToolUse: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly PreCompact: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly PostCompact: unknown; + // Function type detected, and unsupported. Leaving the type as unknown + readonly Stop: unknown; +} + +// From agenthooks/types.go +export interface AgentHookMeta extends AgentHookChatRef { + readonly dispatch_id: string; + readonly schema_version: number; +} + +// From agenthooks/types.go +/** + * Permission controls whether mutable hook input may proceed. + */ +export interface AgentHookPermission { + readonly decision: AgentHookPermissionDecision; + readonly reason?: string; + readonly input_override?: unknown; +} + +// From agenthooks/types.go +export type AgentHookPermissionDecision = "allow" | "ask" | "deny"; + +export const AgentHookPermissionDecisions: AgentHookPermissionDecision[] = [ + "allow", + "ask", + "deny", +]; + +// From agenthooks/types.go +export interface AgentHookPostCompactData {} + +// From agenthooks/types.go +export interface AgentHookPostToolUseData { + readonly tool_use_id: string; + readonly tool_name: string; + readonly tool_response?: unknown; + readonly tool_error?: string; +} + +// From agenthooks/types.go +export interface AgentHookPreCompactData {} + +// From agenthooks/types.go +export interface AgentHookPreToolUseData { + readonly tool_use_id: string; + readonly tool_name: string; + readonly tool_input: unknown; +} + +// From agenthooks/types.go +/** + * Request is the body coderd posts to the configured lifecycle hook URL. + */ +export interface AgentHookRequest { + readonly type: AgentHookEventType; + readonly meta: AgentHookMeta; + readonly data: unknown; +} + +// From agenthooks/types.go +export interface AgentHookResponse { + readonly permission?: AgentHookPermission; + readonly model_context?: string; + readonly user_message?: string; + /** + * AllowedTools distinguishes unchanged (nil), no tools (empty), and named tools. + */ + readonly allowed_tools?: string[]; + readonly end_chat?: boolean; +} + +// From agenthooks/types.go +/** + * SchemaVersion is the current lifecycle hook request schema version. + */ +export const AgentHookSchemaVersion = 1; + +// From agenthooks/types.go +export interface AgentHookSessionStartData { + readonly source: string; +} + +// From agenthooks/types.go +export interface AgentHookStopData {} + +// From agenthooks/types.go +/** + * UserPromptSubmitData includes concatenated text and persisted parts. + * Inspect Parts when structure matters. + */ +export interface AgentHookUserPromptSubmitData { + readonly prompt: string; + readonly parts?: unknown; +} + // From codersdk/workspacebuilds.go export interface AgentScriptTiming { readonly started_at: string; @@ -1792,6 +1959,10 @@ export const ChatComputerUseProviders: ChatComputerUseProvider[] = [ export interface ChatConfig { readonly acquire_batch_size: number; readonly debug_logging_enabled: boolean; + readonly hook_url: string; + readonly hook_secret: string; + readonly hook_timeout: number; + readonly hook_enabled: boolean; /** * @deprecated AI Gateway routing is now the only routing path. Setting this * value has no effect. This option will be removed in a future release. @@ -2252,6 +2423,7 @@ export type ChatErrorKind = | "config" | "content_filter" | "generic" + | "hook_dispatch_failed" | "missing_key" | "overloaded" | "provider_disabled" @@ -2265,6 +2437,7 @@ export const ChatErrorKinds: ChatErrorKind[] = [ "config", "content_filter", "generic", + "hook_dispatch_failed", "missing_key", "overloaded", "provider_disabled", @@ -3682,8 +3855,15 @@ export interface CreateChatMessageRequest { */ export interface CreateChatMessageResponse { readonly message?: ChatMessage; + /** + * Messages contains all user-visible messages inserted by an immediate send, + * in insertion order with the user's message last. Clients should upsert the + * full batch because hooks may prepend notices. Empty for queued or ended sends. + */ + readonly messages?: readonly ChatMessage[]; readonly queued_message?: ChatQueuedMessage; readonly queued: boolean; + readonly ended: boolean; readonly warnings?: readonly string[]; } @@ -4756,11 +4936,17 @@ export interface EditChatMessageRequest { // From codersdk/chats.go /** * EditChatMessageResponse is the response from editing a message in a chat. - * Edits are always synchronous (no queueing), so the message is returned - * directly. */ export interface EditChatMessageResponse { - readonly message: ChatMessage; + readonly message?: ChatMessage; + /** + * Messages holds every user-visible message the edit inserted, in + * insertion order with the replacement message last. Lifecycle + * hooks may prepend notices, so clients must upsert all of them + * rather than only Message. Empty for ended edits. + */ + readonly messages?: readonly ChatMessage[]; + readonly ended: boolean; readonly warnings?: readonly string[]; } @@ -4811,6 +4997,7 @@ export const EntitlementsWarningHeader = "X-Coder-Entitlements-Warning"; // From codersdk/deployment.go export type Experiment = | "ai-gateway-cost-control" + | "agent-lifecycle-hooks" | "auto-fill-parameters" | "chat-advisor" | "chat-virtual-desktop" @@ -4825,6 +5012,7 @@ export type Experiment = export const Experiments: Experiment[] = [ "ai-gateway-cost-control", + "agent-lifecycle-hooks", "auto-fill-parameters", "chat-advisor", "chat-virtual-desktop", diff --git a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx index 0307fe89c3e88..40ccc41113ede 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.stories.tsx @@ -1210,6 +1210,7 @@ export const WithMessageHistory: Story = { id: 5, created_at: "2026-02-18T00:03:00.000Z", }, + ended: false, }); }, play: async ({ canvasElement }) => { @@ -2744,6 +2745,7 @@ export const SlashCompactQueuedEditSaves: Story = { "createChatMessage", ).mockResolvedValue({ queued: true, + ended: false, queued_message: { ...MockChatQueuedMessage, id: 4, @@ -2816,6 +2818,7 @@ export const SlashCompactYieldsToPersonalSkill: Story = { "createChatMessage", ).mockResolvedValue({ queued: false, + ended: false, message: { id: 3, chat_id: CHAT_ID, diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 945a5724e6e4f..2f40f3ec94c04 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -967,6 +967,7 @@ describe("submitEditAndScroll", () => { const callOrder: string[] = []; const editMessage = vi.fn(async () => { callOrder.push("editMessage"); + return { ended: false }; }); const scrollToBottom = vi.fn(() => { callOrder.push("scrollToBottom"); @@ -982,6 +983,19 @@ describe("submitEditAndScroll", () => { expect(callOrder).toEqual(["editMessage", "scrollToBottom"]); }); + it("returns the editMessage response so callers can handle ended chats", async () => { + const editMessage = vi.fn().mockResolvedValue({ ended: true }); + + const response = await submitEditAndScroll({ + editMessage, + editArgs: dummyArgs, + scrollToBottom: vi.fn(), + onError: vi.fn(), + }); + + expect(response).toEqual({ ended: true }); + }); + it("does not call scrollToBottom when editMessage throws", async () => { const scrollToBottom = vi.fn(); const onError = vi.fn(); @@ -1003,7 +1017,7 @@ describe("submitEditAndScroll", () => { }); it("tolerates null scrollToBottom", async () => { - const editMessage = vi.fn().mockResolvedValue(undefined); + const editMessage = vi.fn().mockResolvedValue({ ended: false }); await submitEditAndScroll({ editMessage, diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 46f9afd01856e..3c2ea436b698e 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -28,14 +28,17 @@ import { chat, chatKey, chatMessagesForInfiniteScroll, + chatMessagesKey, chatModelConfigs, chatModels, chatProviderConfigs, + chatsByWorkspaceKeyPrefix, compactChat, createChatMessage, deleteChatQueuedMessage, editChatMessage, interruptChat, + invalidateChatListQueries, mcpServerConfigs, promoteChatQueuedMessage, updateChatPlanMode, @@ -242,7 +245,7 @@ export async function submitEditAndScroll({ messageId: number; optimisticMessage?: TypesGen.ChatMessage; req: TypesGen.EditChatMessageRequest; - }) => Promise; + }) => Promise; editArgs: { messageId: number; optimisticMessage?: TypesGen.ChatMessage; @@ -250,9 +253,10 @@ export async function submitEditAndScroll({ }; scrollToBottom: (() => void) | null | undefined; onError: (error: unknown) => void; -}): Promise { +}): Promise { + let response: TypesGen.EditChatMessageResponse; try { - await editMessage(editArgs); + response = await editMessage(editArgs); } catch (error) { onError(error); throw error; @@ -264,6 +268,7 @@ export async function submitEditAndScroll({ // as the IntersectionObserver reacts to rapid layout // shifts between the old and truncated content. scrollToBottom?.(); + return response; } /** @internal Exported for testing. */ @@ -1585,7 +1590,7 @@ const AgentChatPage: FC = () => { store.setChatStatus("running"); store.clearStreamState(); }); - await submitEditAndScroll({ + const editResponse = await submitEditAndScroll({ editMessage, editArgs: { messageId: editedMessageID, @@ -1598,6 +1603,11 @@ const AgentChatPage: FC = () => { handleUsageLimitError(error); }, }); + if (editResponse.ended) { + restoreOptimisticRequestSnapshot(store, previousSnapshot); + store.clearStreamState(); + return; + } if (editSelectedModelConfigID) { localStorage.setItem( lastModelConfigIDStorageKey, @@ -1636,8 +1646,33 @@ const AgentChatPage: FC = () => { response = await sendMessage(request); } catch (error) { handleUsageLimitError(error); + // A failed hook dispatch can still archive or error the chat. + // Refresh the transcript and lists in case the chat-watch event is missed. + void queryClient.invalidateQueries({ queryKey: chatKey(agentId) }); + void queryClient.invalidateQueries({ + queryKey: chatMessagesKey(agentId), + exact: true, + }); + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ + queryKey: chatsByWorkspaceKeyPrefix, + }); throw error; } + // A hook-ended send may add notices and archive the chat, so refresh both views. + if (response.ended) { + store.clearStreamState(); + void queryClient.invalidateQueries({ queryKey: chatKey(agentId) }); + void queryClient.invalidateQueries({ + queryKey: chatMessagesKey(agentId), + exact: true, + }); + void invalidateChatListQueries(queryClient); + void queryClient.invalidateQueries({ + queryKey: chatsByWorkspaceKeyPrefix, + }); + return; + } // When the server accepts the message immediately (not // queued), clear the stream and insert the user's message // so it appears in the timeline without waiting for the @@ -1653,9 +1688,14 @@ const AgentChatPage: FC = () => { // to error/pending instead, the WebSocket event // overrides this optimistic value. store.setChatStatus("running"); - if (response.message) { - store.upsertDurableMessage(response.message); - upsertCacheMessages([response.message]); + // Prefer the full inserted batch: hooks may prepend notices + // with lower IDs than the user message, and a stream + // reconnect keyed on the highest cached ID would skip them. + const insertedMessages = + response.messages ?? (response.message ? [response.message] : []); + if (insertedMessages.length > 0) { + store.upsertDurableMessages(insertedMessages); + upsertCacheMessages(insertedMessages); } } if (selectedModelConfigID) { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx index 4c2f707272296..68faafefee61e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.stories.tsx @@ -404,6 +404,74 @@ const meta: Meta = { export default meta; type Story = StoryObj; +export const LifecycleHookNotice: Story = { + args: { + ...defaultArgs, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "system", + content: [ + { + type: "text", + text: "Your organization requires an approval before deployment.", + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const notice = canvas.getByRole("alert"); + expect(notice).toBeVisible(); + expect(within(notice).getByText("Lifecycle hook")).toBeVisible(); + expect( + within(notice).getByText( + "Your organization requires an approval before deployment.", + ), + ).toBeVisible(); + expect( + canvas.queryByRole("button", { name: "Copy message" }), + ).not.toBeInTheDocument(); + }, +}; + +export const LifecycleHookNoticeDimmedWhileEditing: Story = { + args: { + ...defaultArgs, + editingMessageId: 1, + parsedMessages: buildMessages([ + { + ...baseMessage, + id: 1, + role: "user", + content: [{ type: "text", text: "original prompt" }], + }, + { + ...baseMessage, + id: 2, + role: "system", + content: [ + { + type: "text", + text: "See the [policy](https://example.com/policy) for details.", + }, + ], + }, + ]), + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + const notice = canvas.getByRole("alert"); + const link = notice.querySelector("a"); + expect(link).not.toBeNull(); + link?.focus(); + expect(document.activeElement).not.toBe(link); + expect(notice.closest("[inert]")).not.toBeNull(); + }, +}; + export const DurableListTemplatesToolLifecycle: Story = { args: { ...defaultArgs, diff --git a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx index fa74ec3b7e75a..f786b5a9def00 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/ConversationTimeline.tsx @@ -14,6 +14,7 @@ import { preferenceSettings } from "#/api/queries/users"; import type * as TypesGen from "#/api/typesGenerated"; import type { ThinkingDisplayMode } from "#/api/typesGenerated"; +import { Alert, AlertTitle } from "#/components/Alert/Alert"; import { Button } from "#/components/Button/Button"; import { CopyButton } from "#/components/CopyButton/CopyButton"; import { @@ -591,6 +592,25 @@ const ChatMessageItem = memo<{ if (displayState.shouldHide) { return null; } + if (message.role === "system") { + return ( +
+ +
+ Lifecycle hook + {parsed.markdown} +
+
+
+ ); + } const conversationItemProps: { role: "user" | "assistant" } = { role: isUser ? "user" : "assistant", @@ -1091,6 +1111,10 @@ function computeLastInChainFlags( let nextVisibleIsUser = true; for (let i = displayMessages.length - 1; i >= 0; i--) { const entry = displayMessages[i]; + if (entry.message.role === "system") { + nextVisibleIsUser = true; + continue; + } if (entry.message.role !== "user") { flags[i] = nextVisibleIsUser; }