package chatd_test import ( "encoding/json" "net/http" "net/http/httptest" "slices" "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/agenthooks/dispatch" "github.com/coder/coder/v2/coderd/x/chatd" "github.com/coder/coder/v2/coderd/x/chatd/chathooks" "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/x/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)) data := decodeHookData[agenthooks.UserPromptSubmitData](t, request) 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"}`)) 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) parts, err := chatprompt.ParseContent(result.Message) require.NoError(t, err) require.Equal(t, []codersdk.ChatMessagePart{ codersdk.ChatMessageText("after"), codersdk.ChatMessageFileReference("main.go", 1, 3, "package main"), {Type: codersdk.ChatMessagePartTypeHookContext, Text: "model only"}, {Type: codersdk.ChatMessagePartTypeHookNotice, Text: "user only"}, }, parts) require.Len(t, result.InsertedMessages, 1) require.Equal(t, result.Message.ID, result.InsertedMessages[0].ID) }) 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 *chathooks.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 newHookDispatcher(t *testing.T, _ database.Store, consumer *httptest.Server) *dispatch.Dispatcher { t.Helper() return dispatch.New( slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), consumer.Client(), consumer.URL, false, "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()) } 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, }) var received agenthooks.Request consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, json.NewDecoder(r.Body).Decode(&received)) _, err := w.Write([]byte(`{}`)) require.NoError(t, err) })) t.Cleanup(consumer.Close) server := newHookTestServer(t, db, ps, consumer) 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)) require.Equal(t, agenthooks.EventUserPromptSubmit, received.Type) promptData := decodeHookData[agenthooks.UserPromptSubmitData](t, received) require.Equal(t, "passthrough", promptData.Prompt) // The persisted content is jsonb-normalized, so compare JSON // semantics rather than raw bytes. require.JSONEq(t, string(result.Message.Content.RawMessage), string(promptData.Parts)) } 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) var received agenthooks.Request consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, json.NewDecoder(r.Body).Decode(&received)) _, err := w.Write([]byte(`{"permission":{"decision":"allow","input_override":{"prompt":"queued override"}},"model_context":"queued context","user_message":"queued notice"}`)) require.NoError(t, err) })) t.Cleanup(consumer.Close) 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) 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) wantQueuedParts := []codersdk.ChatMessagePart{ codersdk.ChatMessageText("queued override"), {Type: codersdk.ChatMessagePartTypeHookContext, Text: "queued context"}, {Type: codersdk.ChatMessagePartTypeHookNotice, Text: "queued notice"}, } require.Equal(t, wantQueuedParts, queuedParts) queued, err := db.GetChatQueuedMessages(ctx, chat.ID) require.NoError(t, err) require.Len(t, queued, 1) persistedParts, err := chatprompt.ParseContent(database.ChatMessage{ Role: database.ChatMessageRoleUser, Content: pqtype.NullRawMessage{RawMessage: queued[0].Content, Valid: true}, ContentVersion: chatprompt.CurrentContentVersion, }) require.NoError(t, err) require.Equal(t, wantQueuedParts, persistedParts) require.Equal(t, agenthooks.EventUserPromptSubmit, received.Type) require.Equal(t, "queued original", decodeHookData[agenthooks.UserPromptSubmitData](t, received).Prompt) } 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 *chathooks.UserPromptDeniedError require.ErrorAs(t, err, &denied) }, }, { name: "dispatch failure", statusCode: http.StatusInternalServerError, assertErr: func(t *testing.T, err error) { var dispatchErr *dispatch.Error 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 TestSubagentSpawnHookDispatchFailureFailsTurn(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("spawn_agent", `{"type":"general","prompt":"child admission prompt"}`) chunk.Choices[0].ToolCalls[0].ID = "call_spawn" 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.EventUserPromptSubmit { data := decodeHookData[agenthooks.UserPromptSubmitData](t, request) if data.Prompt == "child admission prompt" { 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: "spawn-hook-failure", ModelConfigID: model.ID, InitialUserContent: []codersdk.ChatMessagePart{ codersdk.ChatMessageText("spawn a child"), }, }) require.NoError(t, err) failed := waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusError) require.Contains(t, chatLastErrorMessage(failed.LastError), "hook dispatch failed: user_prompt_submit: http_error") messages := chatMessages(ctx, t, db, chat.ID) require.Len(t, messages, 3) require.Equal(t, database.ChatMessageRoleUser, messages[0].Role) require.Equal(t, database.ChatMessageRoleAssistant, messages[1].Role) require.Equal(t, database.ChatMessageRoleTool, messages[2].Role) require.Contains(t, string(messages[2].Content.RawMessage), "lifecycle hook returned HTTP status 500") chats, err := db.GetChats(ctx, database.GetChatsParams{ OwnedOnly: true, ViewerID: user.ID, AfterID: uuid.Nil, OffsetOpt: 0, LimitOpt: 100, }) require.NoError(t, err) require.Len(t, chats, 1) } 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, }) var received agenthooks.Request consumer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { require.NoError(t, json.NewDecoder(r.Body).Decode(&received)) 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 *dispatch.Error require.ErrorAs(t, err, &dispatchErr) require.Equal(t, dispatch.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) require.Equal(t, agenthooks.EventUserPromptSubmit, received.Type) prompt := decodeHookData[agenthooks.UserPromptSubmitData](t, received) require.Equal(t, "fails", prompt.Prompt) messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) require.NoError(t, err) require.Empty(t, messages) queued, err := db.GetChatQueuedMessages(ctx, chat.ID) require.NoError(t, err) require.Empty(t, queued) } 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 } var receivedMu sync.Mutex 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) receivedMu.Lock() received = append(received, receivedHook{request: request, claims: claims}) receivedMu.Unlock() response := `{"model_context":"clear context","user_message":"clear notice"}` if request.Type == agenthooks.EventUserPromptSubmit { response = `{"permission":{"decision":"allow","input_override":{"prompt":"edited override"}},"model_context":"edit context","user_message":"edit notice"}` } _, err = w.Write([]byte(response)) require.NoError(t, err) })) t.Cleanup(consumer.Close) server := newHookTestServer(t, db, ps, consumer) chatFile, err := db.InsertChatFile(ctx, database.InsertChatFileParams{ OwnerID: user.ID, OrganizationID: org.ID, Name: "edited.png", Mimetype: "image/png", Data: []byte("png-bytes"), }) require.NoError(t, err) upload := codersdk.ChatMessageFile(chatFile.ID, chatFile.Mimetype, chatFile.Name) reference := codersdk.ChatMessageFileReference("main.go", 1, 3, "package main") result, err := server.EditMessage(ctx, chatd.EditMessageOptions{ ChatID: chat.ID, CreatedBy: user.ID, EditedMessageID: inserted[0].ID, Content: []codersdk.ChatMessagePart{ reference, codersdk.ChatMessageText("edited original"), upload, }, }) require.NoError(t, err) linkedFiles, err := db.GetChatFileMetadataByChatID(ctx, chat.ID) require.NoError(t, err) require.Len(t, linkedFiles, 1) require.Equal(t, chatFile.ID, linkedFiles[0].ID) parts, err := chatprompt.ParseContent(result.Message) require.NoError(t, err) require.Equal(t, []codersdk.ChatMessagePart{ reference, codersdk.ChatMessageText("edited override"), upload, {Type: codersdk.ChatMessagePartTypeHookContext, Text: "edit context"}, {Type: codersdk.ChatMessagePartTypeHookNotice, Text: "edit notice"}, }, parts) receivedMu.Lock() received = slices.Clone(received) receivedMu.Unlock() require.Len(t, received, 2) require.Equal(t, agenthooks.EventSessionStart, received[0].request.Type) require.Equal(t, agenthooks.SessionStartData{Source: "clear"}, decodeHookData[agenthooks.SessionStartData](t, received[0].request)) require.Equal(t, received[0].request.Meta.DispatchID, received[0].claims.JTI) require.Equal(t, agenthooks.EventUserPromptSubmit, received[1].request.Type) prompt := decodeHookData[agenthooks.UserPromptSubmitData](t, received[1].request) require.Equal(t, "edited original", prompt.Prompt) require.NotNil(t, received[0].request.Meta.TurnID) require.Equal(t, received[0].request.Meta.TurnID, received[1].request.Meta.TurnID) 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) } 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") } func TestPromptHooksAdmissionPreflight(t *testing.T) { t.Parallel() db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) user, org, model := seedChatDependencies(t, db) received := make(chan agenthooks.Request, 8) 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)) received <- request _, 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) select { case request := <-received: t.Fatalf("admission-rejected prompt dispatched %s", request.Type) default: } } 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.Equal(t, "unchanged", hookMessageText(t, result.Message)) } 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 } func decodeHookData[T any](t *testing.T, request agenthooks.Request) T { t.Helper() var data T require.NoError(t, json.Unmarshal(request.Data, &data)) return data }