Files
coder/coderd/x/chatd/hooks_test.go
T
Michael Suchacz 4b9880afa6 feat: add --chat-hook-allow-insecure to allow plain HTTP chat hook URLs (#27896)
Adds a hidden `--chat-hook-allow-insecure` /
`CODER_CHAT_HOOK_ALLOW_INSECURE` deployment option (default `false`)
that allows the chat lifecycle hook URL to use plain HTTP for any host.

The HTTPS requirement is enforced at two points, and the flag relaxes
both: `DeploymentValues.Validate()` rejects `http` hook URLs at startup,
and the hook dispatcher's `validateHookURL` allows `http` only for
loopback hosts. With the flag set, any-host `http` is accepted; the
host, fragment/userinfo, secret, and timeout checks are unchanged, and
non-http(s) schemes still fail. This removes the need for an HTTPS
reverse proxy when testing a hook consumer on a trusted network.

Following security review feedback, the flag description and docs state
that plain HTTP lets an on-path attacker forge hook responses (which
control agent execution), and `coder server` logs a startup warning
(with a redacted hook URL) when hooks run over plain HTTP.

Docs, generated API types, and the server config golden are updated
accordingly.

> Mux acted on Mike's behalf to create this PR.
2026-08-05 22:41:17 +02:00

694 lines
26 KiB
Go

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)
upload := codersdk.ChatMessageFile(uuid.New(), "image/png", "edited.png")
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)
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
}