mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Michael Suchacz
91c7232d97
Adds generic chat state and query capabilities that the lifecycle hooks integration (#27429) builds on. Part of the lifecycle hooks stack (#27401, #27429, #27430). - `chatstate`: `EditMessage` accepts caller-provided suffix messages inserted after the replacement in the same transaction, transitions can carry a typed error kind, and `FinishError` is also allowed from waiting chats so admission-time failures can park an idle chat in error. - `chatstate`: `ValidateToolResults` holds the submitted-tool-result rules (duplicate, invalid JSON, missing, unexpected) in one place, so `CompleteRequiresAction` and API-level prechecks reject the same payloads with the same typed causes. - `database`: `InsertChat` accepts an optional caller-provided ID. No hook-specific state or behavior is introduced here; these primitives are usable by any caller. An earlier revision added a message-content rewrite primitive so a `pre_tool_use` override could update an already-committed tool call. Message content is immutable by design, and @hugodutka pushed back on changing that. The rewrite is gone: #27429 now dispatches the hook before the assistant message is stored, so the stored input is the one that runs and nothing needs updating. > This PR was written by Mux, an AI coding agent, on Mike's behalf.
850 lines
32 KiB
Go
850 lines
32 KiB
Go
package chatstate_test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/sqlc-dev/pqtype"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/coder/coder/v2/coderd/database"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
|
"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
|
|
"github.com/coder/coder/v2/codersdk"
|
|
"github.com/coder/coder/v2/testutil"
|
|
)
|
|
|
|
// CreateChat tests.
|
|
//
|
|
// CreateChat is the only transition that originates from StateN and it
|
|
// is not exercised through ChatMachine.Update, so it lives outside
|
|
// TestTransitionMatrix_AllCombinations.
|
|
|
|
// TestTransitionCreate_NToR0 verifies that CreateChat lands a fresh
|
|
// chat in R0 with snapshot_version 1, the initial user message
|
|
// recorded at revision 1, queue_version still 0, and the post-commit
|
|
// publish requesting an ownership hint plus a chat:update.
|
|
func TestTransitionCreate_NToR0(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
res := createTestChat(t, f)
|
|
|
|
require.Equal(t, database.ChatStatusRunning, res.Chat.Status)
|
|
require.False(t, res.Chat.Archived)
|
|
require.Equal(t, int64(1), res.Chat.SnapshotVersion, "snapshot_version starts at 1")
|
|
require.Equal(t, int64(1), res.Chat.HistoryVersion, "history_version set by trigger after initial insert")
|
|
require.Equal(t, int64(0), res.Chat.QueueVersion, "queue_version stays 0 when no queue rows")
|
|
require.Equal(t, int64(0), res.Chat.GenerationAttempt)
|
|
require.NotEmpty(t, res.InitialMessages)
|
|
require.Equal(t, int64(1), res.InitialMessages[0].Revision)
|
|
require.Equal(t, chatstate.StateR0, f.classify(ctx, t, res.Chat.ID))
|
|
require.True(t, f.Pub.hasOwnership(), "newly created chat is runnable and unowned")
|
|
f.Pub.expectChatUpdate(t, res.Chat.ID, 1)
|
|
}
|
|
|
|
// TestCreateChat_RejectsEmptyInitialMessages verifies that CreateChat
|
|
// rejects an empty InitialMessages slice with ErrTransitionNotAllowed
|
|
// and does not publish anything.
|
|
func TestCreateChat_RejectsEmptyInitialMessages(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
_, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{
|
|
OrganizationID: f.Org.ID,
|
|
OwnerID: f.User.ID,
|
|
LastModelConfigID: f.Model.ID,
|
|
ClientType: database.ChatClientTypeApi,
|
|
Title: "t",
|
|
InitialMessages: nil,
|
|
})
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed)
|
|
require.Empty(t, f.Pub.channels, "rejected create must not publish")
|
|
}
|
|
|
|
func TestCreateChat_AllowsNoUserMessages(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
assistant := userTextMessage("oops", f.User.ID, f.Model.ID)
|
|
assistant.Role = database.ChatMessageRoleAssistant
|
|
res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{
|
|
OrganizationID: f.Org.ID,
|
|
OwnerID: f.User.ID,
|
|
LastModelConfigID: f.Model.ID,
|
|
Title: "t",
|
|
ClientType: database.ChatClientTypeApi,
|
|
InitialMessages: []chatstate.Message{assistant},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, res.InitialMessages, 1)
|
|
}
|
|
|
|
func TestCreateChat_AllowsNonFinalUserMessage(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
res, err := chatstate.CreateChat(ctx, f.DB, f.Pub, chatstate.CreateChatInput{
|
|
OrganizationID: f.Org.ID,
|
|
OwnerID: f.User.ID,
|
|
LastModelConfigID: f.Model.ID,
|
|
Title: "t",
|
|
ClientType: database.ChatClientTypeApi,
|
|
InitialMessages: []chatstate.Message{
|
|
userTextMessage("context user", f.User.ID, f.Model.ID),
|
|
userTextMessage("final user", f.User.ID, f.Model.ID),
|
|
},
|
|
})
|
|
require.NoError(t, err)
|
|
require.Len(t, res.InitialMessages, 2)
|
|
}
|
|
|
|
// Input-specific rejection cases.
|
|
//
|
|
// These tests cover the same matrix rows as TestTransitionMatrix_AllCombinations
|
|
// but exercise legal source states with invalid transition inputs. They are
|
|
// intentionally outside the matrix entry point so the matrix focus stays on
|
|
// positive cases and generated disallowed cases.
|
|
|
|
type setArchivedWrongDirectionCase struct {
|
|
from chatstate.ExecutionState
|
|
wantArchive bool
|
|
label string
|
|
}
|
|
|
|
func setArchivedWrongDirectionCases() []setArchivedWrongDirectionCase {
|
|
return []setArchivedWrongDirectionCase{
|
|
// Non-archived states with archived=false: no-op.
|
|
{from: chatstate.StateW, wantArchive: false, label: "W_to_W"},
|
|
{from: chatstate.StateE0, wantArchive: false, label: "E0_to_E0"},
|
|
{from: chatstate.StateE1, wantArchive: false, label: "E1_to_E1"},
|
|
// Archived states with archived=true: no-op.
|
|
{from: chatstate.StateXW, wantArchive: true, label: "XW_to_XW"},
|
|
{from: chatstate.StateXE0, wantArchive: true, label: "XE0_to_XE0"},
|
|
{from: chatstate.StateXE1, wantArchive: true, label: "XE1_to_XE1"},
|
|
}
|
|
}
|
|
|
|
var invalidBusyBehaviors = []chatstate.BusyBehavior{
|
|
chatstate.BusyBehavior(""),
|
|
chatstate.BusyBehavior("not-a-real-mode"),
|
|
}
|
|
|
|
func runSetArchivedWrongDirectionCase(t *testing.T, tc setArchivedWrongDirectionCase) {
|
|
t.Helper()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
seeded := seedState(t, f, tc.from)
|
|
require.Equal(t, tc.from, f.classify(ctx, t, seeded.chatID))
|
|
|
|
base := captureBaseline(ctx, t, f, seeded)
|
|
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID)
|
|
err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, serr := tx.SetArchived(chatstate.SetArchivedInput{Archived: tc.wantArchive})
|
|
return serr
|
|
})
|
|
require.Error(t, err, "SetArchived must reject when Archived matches the current value")
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed,
|
|
"SetArchived must wrap ErrTransitionNotAllowed")
|
|
var te *chatstate.TransitionError
|
|
require.ErrorAs(t, err, &te,
|
|
"SetArchived must return a typed TransitionError")
|
|
require.Equal(t, chatstate.TransitionSetArchived, te.Transition)
|
|
require.Equal(t, tc.from, te.From, "TransitionError records the loaded from-state")
|
|
|
|
require.Equal(t, tc.from, f.classify(ctx, t, seeded.chatID),
|
|
"rejected SetArchived must leave the chat in the same state")
|
|
assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base)
|
|
}
|
|
|
|
func runInvalidBusyBehaviorCase(t *testing.T, from chatstate.ExecutionState, bb chatstate.BusyBehavior) {
|
|
t.Helper()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
seeded := seedState(t, f, from)
|
|
require.Equal(t, from, f.classify(ctx, t, seeded.chatID),
|
|
"seed must land in %s", from)
|
|
base := captureBaseline(ctx, t, f, seeded)
|
|
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID)
|
|
err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, serr := tx.SendMessage(chatstate.SendMessageInput{
|
|
Message: userTextMessage("invalid-bb", f.User.ID, f.Model.ID),
|
|
BusyBehavior: bb,
|
|
})
|
|
return serr
|
|
})
|
|
require.Error(t, err, "SendMessage must reject invalid BusyBehavior")
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed,
|
|
"SendMessage rejection must wrap ErrTransitionNotAllowed")
|
|
var te *chatstate.TransitionError
|
|
require.ErrorAs(t, err, &te,
|
|
"SendMessage must return a typed TransitionError")
|
|
require.Equal(t, chatstate.TransitionSendMessage, te.Transition)
|
|
require.Equal(t, from, te.From,
|
|
"TransitionError records the source state")
|
|
|
|
require.Equal(t, from, f.classify(ctx, t, seeded.chatID),
|
|
"rejected SendMessage must leave the chat in the same state")
|
|
assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base)
|
|
}
|
|
|
|
type completeRequiresActionRejectCase struct {
|
|
name string
|
|
results func(seeded seededChat) []chatstate.ToolResultInput
|
|
}
|
|
|
|
type recordRetryStateRejectCase struct {
|
|
name string
|
|
retryState pqtype.NullRawMessage
|
|
}
|
|
|
|
func completeRequiresActionRejectCases() []completeRequiresActionRejectCase {
|
|
valid := func(id string) chatstate.ToolResultInput {
|
|
return chatstate.ToolResultInput{
|
|
ToolCallID: id,
|
|
Output: json.RawMessage(`{"ok":true}`),
|
|
}
|
|
}
|
|
return []completeRequiresActionRejectCase{
|
|
{
|
|
name: "missing_required_tool_result",
|
|
results: func(seeded seededChat) []chatstate.ToolResultInput { return nil },
|
|
},
|
|
{
|
|
name: "extra_tool_result",
|
|
results: func(seeded seededChat) []chatstate.ToolResultInput {
|
|
return []chatstate.ToolResultInput{valid(seeded.pendingToolCallID), valid("call_extra")}
|
|
},
|
|
},
|
|
{
|
|
name: "duplicate_tool_call_id",
|
|
results: func(seeded seededChat) []chatstate.ToolResultInput {
|
|
return []chatstate.ToolResultInput{valid(seeded.pendingToolCallID), valid(seeded.pendingToolCallID)}
|
|
},
|
|
},
|
|
{
|
|
name: "mismatched_tool_call_id",
|
|
results: func(seeded seededChat) []chatstate.ToolResultInput {
|
|
return []chatstate.ToolResultInput{valid("call_mismatch")}
|
|
},
|
|
},
|
|
{
|
|
name: "invalid_json_output",
|
|
results: func(seeded seededChat) []chatstate.ToolResultInput {
|
|
return []chatstate.ToolResultInput{{ToolCallID: seeded.pendingToolCallID, Output: json.RawMessage(`{`)}}
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func recordRetryStateRejectCases() []recordRetryStateRejectCase {
|
|
return []recordRetryStateRejectCase{
|
|
{
|
|
name: "sql_null_payload",
|
|
},
|
|
{
|
|
name: "empty_payload",
|
|
retryState: pqtype.NullRawMessage{RawMessage: json.RawMessage(``), Valid: true},
|
|
},
|
|
{
|
|
name: "invalid_json_payload",
|
|
retryState: pqtype.NullRawMessage{RawMessage: json.RawMessage(`{`), Valid: true},
|
|
},
|
|
}
|
|
}
|
|
|
|
func runCompleteRequiresActionRejectCase(t *testing.T, tc completeRequiresActionRejectCase) {
|
|
t.Helper()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
seeded := seedAOrA1(t, f, 0, "reject_complete_requires_action")
|
|
require.Equal(t, chatstate.StateA0, f.classify(ctx, t, seeded.chatID))
|
|
base := captureBaseline(ctx, t, f, seeded)
|
|
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID)
|
|
err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, cerr := tx.CompleteRequiresAction(chatstate.CompleteRequiresActionInput{
|
|
CreatedBy: f.User.ID,
|
|
ModelConfigID: f.Model.ID,
|
|
Results: tc.results(seeded),
|
|
})
|
|
return cerr
|
|
})
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed)
|
|
var te *chatstate.TransitionError
|
|
require.ErrorAs(t, err, &te)
|
|
require.Equal(t, chatstate.TransitionCompleteRequiresAction, te.Transition)
|
|
require.Equal(t, chatstate.StateA0, te.From)
|
|
assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base)
|
|
}
|
|
|
|
func runRecordRetryStateRejectCase(t *testing.T, tc recordRetryStateRejectCase) {
|
|
t.Helper()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
seeded := seedState(t, f, chatstate.StateR0)
|
|
require.Equal(t, chatstate.StateR0, f.classify(ctx, t, seeded.chatID))
|
|
base := captureBaseline(ctx, t, f, seeded)
|
|
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID)
|
|
err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, rerr := tx.RecordRetryState(chatstate.RecordRetryStateInput{
|
|
RetryState: tc.retryState,
|
|
})
|
|
return rerr
|
|
})
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed)
|
|
var te *chatstate.TransitionError
|
|
require.ErrorAs(t, err, &te)
|
|
require.Equal(t, chatstate.TransitionRecordRetryState, te.Transition)
|
|
require.Equal(t, chatstate.StateR0, te.From)
|
|
assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base)
|
|
}
|
|
|
|
// TestTransitionInputValidation groups every input-specific rejection
|
|
// test. The matrix coverage entry point in
|
|
// TestTransitionMatrix_AllCombinations intentionally focuses on
|
|
// positive cases and generated disallowed cases; rejection cases that
|
|
// exercise legal matrix rows with invalid inputs live here so the
|
|
// matrix entry point stays focused.
|
|
func TestTransitionInputValidation(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
t.Run("SetArchived_wrong_direction", func(t *testing.T) {
|
|
t.Parallel()
|
|
for _, tc := range setArchivedWrongDirectionCases() {
|
|
t.Run(tc.label, func(t *testing.T) {
|
|
t.Parallel()
|
|
runSetArchivedWrongDirectionCase(t, tc)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("SendMessage_invalid_busy_behavior", func(t *testing.T) {
|
|
t.Parallel()
|
|
for _, from := range chatstate.AllowedInputStates(chatstate.TransitionSendMessage) {
|
|
for _, bb := range invalidBusyBehaviors {
|
|
label := from.String() + "/" + string(bb)
|
|
if bb == "" {
|
|
label = from.String() + "/empty"
|
|
}
|
|
t.Run(label, func(t *testing.T) {
|
|
t.Parallel()
|
|
runInvalidBusyBehaviorCase(t, from, bb)
|
|
})
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("CompleteRequiresAction_invalid_results", func(t *testing.T) {
|
|
t.Parallel()
|
|
for _, tc := range completeRequiresActionRejectCases() {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
runCompleteRequiresActionRejectCase(t, tc)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("RecordRetryState_invalid_payload", func(t *testing.T) {
|
|
t.Parallel()
|
|
for _, tc := range recordRetryStateRejectCases() {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
runRecordRetryStateRejectCase(t, tc)
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestSendMessageQueueCapRejectsQueueAppend seeds a chat with the
|
|
// maximum queued messages and asserts that the next SendMessage in
|
|
// a queue-appending state returns chatstate.ErrMessageQueueFull and
|
|
// rolls back without persisting another queued row.
|
|
func TestSendMessageQueueCapRejectsQueueAppend(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
created := createTestChat(t, f)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
|
|
// createTestChat lands the chat in R0; SendMessage in R0 with
|
|
// BusyBehaviorQueue queues. Fill the queue to MaxQueueSize.
|
|
for i := 0; i < chatstate.MaxQueueSize; i++ {
|
|
sendQueuedMessage(t, f, m, "filler")
|
|
}
|
|
count, err := f.DB.CountChatQueuedMessages(ctx, created.Chat.ID)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, chatstate.MaxQueueSize, count)
|
|
chatBefore := f.readChat(ctx, t, created.Chat.ID)
|
|
|
|
// The next queue append must fail with ErrMessageQueueFull and a
|
|
// typed wrapper that exposes the cap.
|
|
err = m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, serr := tx.SendMessage(chatstate.SendMessageInput{
|
|
Message: userTextMessage("overflow", f.User.ID, f.Model.ID),
|
|
BusyBehavior: chatstate.BusyBehaviorQueue,
|
|
})
|
|
return serr
|
|
})
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, chatstate.ErrMessageQueueFull,
|
|
"queue-append over the cap returns ErrMessageQueueFull")
|
|
var typed *chatstate.MessageQueueFullError
|
|
require.ErrorAs(t, err, &typed, "ErrMessageQueueFull is carried as a typed error")
|
|
require.EqualValues(t, chatstate.MaxQueueSize, typed.Max)
|
|
|
|
// The transaction rolled back: queue size, snapshot version,
|
|
// and queue version are unchanged.
|
|
countAfter, err := f.DB.CountChatQueuedMessages(ctx, created.Chat.ID)
|
|
require.NoError(t, err)
|
|
require.EqualValues(t, chatstate.MaxQueueSize, countAfter,
|
|
"queue size must not change when the cap rejects the append")
|
|
chatAfter := f.readChat(ctx, t, created.Chat.ID)
|
|
require.Equal(t, chatBefore.SnapshotVersion, chatAfter.SnapshotVersion,
|
|
"failed queue append must not bump snapshot_version")
|
|
require.Equal(t, chatBefore.QueueVersion, chatAfter.QueueVersion,
|
|
"failed queue append must not bump queue_version")
|
|
}
|
|
|
|
func TestSendMessageInterruptRequiresActionReturnsCancellations(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
seeded := seedAOrA1(t, f, 0, "interrupt_cancels")
|
|
require.Equal(t, chatstate.StateA0, f.classify(ctx, t, seeded.chatID))
|
|
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, seeded.chatID)
|
|
var send chatstate.SendMessageResult
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
var err error
|
|
send, err = tx.SendMessage(chatstate.SendMessageInput{
|
|
Message: userTextMessage("interrupt", f.User.ID, f.Model.ID),
|
|
BusyBehavior: chatstate.BusyBehaviorInterrupt,
|
|
})
|
|
return err
|
|
}))
|
|
|
|
require.NotNil(t, send.QueuedMessage, "interrupt from A0 queues the user message")
|
|
require.Len(t, send.InsertedMessages, 1,
|
|
"the synthetic tool cancellation must be reported to callers")
|
|
cancel := send.InsertedMessages[0]
|
|
require.Equal(t, database.ChatMessageRoleTool, cancel.Role)
|
|
require.Equal(t, database.ChatMessageVisibilityBoth, cancel.Visibility)
|
|
parts, err := chatprompt.ParseContent(cancel)
|
|
require.NoError(t, err)
|
|
require.Len(t, parts, 1)
|
|
require.Equal(t, seeded.pendingToolCallID, parts[0].ToolCallID)
|
|
}
|
|
|
|
// TestEditMessageNonUserReturnsSentinel asserts that editing a
|
|
// non-user message returns chatstate.ErrEditedMessageNotUser via
|
|
// the TransitionError cause chain, and still matches the generic
|
|
// transition sentinel.
|
|
func TestEditMessageNonUserReturnsSentinel(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
created := createTestChat(t, f)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
|
|
// Insert an assistant message via CommitStep so we have a
|
|
// non-user message to target.
|
|
var assistantID int64
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
assistant := userTextMessage("assistant", f.User.ID, f.Model.ID)
|
|
assistant.Role = database.ChatMessageRoleAssistant
|
|
step, err := tx.CommitStep(chatstate.CommitStepInput{
|
|
Messages: []chatstate.Message{assistant},
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
require.Len(t, step.InsertedMessages, 1)
|
|
assistantID = step.InsertedMessages[0].ID
|
|
return nil
|
|
}))
|
|
|
|
rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("new content"),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
editErr := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, eerr := tx.EditMessage(chatstate.EditMessageInput{
|
|
MessageID: assistantID,
|
|
CreatedBy: f.User.ID,
|
|
Content: rawContent,
|
|
})
|
|
return eerr
|
|
})
|
|
require.Error(t, editErr)
|
|
require.ErrorIs(t, editErr, chatstate.ErrEditedMessageNotUser,
|
|
"non-user edit returns ErrEditedMessageNotUser via TransitionError cause")
|
|
require.ErrorIs(t, editErr, chatstate.ErrTransitionNotAllowed,
|
|
"ErrEditedMessageNotUser still matches the generic transition sentinel")
|
|
}
|
|
|
|
func TestEditMessageInsertsSuffixMessages(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)
|
|
|
|
target := userTextMessage("original prompt", f.User.ID, f.Model.ID)
|
|
|
|
var targetID int64
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
step, err := tx.CommitStep(chatstate.CommitStepInput{
|
|
Messages: []chatstate.Message{target},
|
|
})
|
|
if err != nil {
|
|
return err
|
|
}
|
|
require.Len(t, step.InsertedMessages, 1)
|
|
targetID = step.InsertedMessages[0].ID
|
|
return nil
|
|
}))
|
|
|
|
firstSuffix := userTextMessage("first notice", f.User.ID, f.Model.ID)
|
|
firstSuffix.Role = database.ChatMessageRoleSystem
|
|
secondSuffix := userTextMessage("second notice", f.User.ID, f.Model.ID)
|
|
secondSuffix.Role = database.ChatMessageRoleSystem
|
|
rawContent, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{
|
|
codersdk.ChatMessageText("edited prompt"),
|
|
})
|
|
require.NoError(t, err)
|
|
|
|
var result chatstate.EditMessageResult
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
result, err = tx.EditMessage(chatstate.EditMessageInput{
|
|
MessageID: targetID,
|
|
SuffixMessages: []chatstate.Message{firstSuffix, secondSuffix},
|
|
CreatedBy: f.User.ID,
|
|
Content: rawContent,
|
|
})
|
|
return err
|
|
}))
|
|
|
|
require.Contains(t, result.DeletedMessageIDs, targetID)
|
|
require.Len(t, result.SuffixMessages, 2)
|
|
assertChatMessageText(t, result.SuffixMessages[0], "first notice")
|
|
assertChatMessageText(t, result.SuffixMessages[1], "second notice")
|
|
require.Greater(t, result.SuffixMessages[0].ID, result.ReplacementMessage.ID,
|
|
"the suffix messages must follow the replacement")
|
|
require.Greater(t, result.SuffixMessages[1].ID, result.SuffixMessages[0].ID,
|
|
"suffix message IDs must ascend with the input array")
|
|
_, err = f.DB.GetChatMessageByID(ctx, result.ReplacementMessage.ID)
|
|
require.NoError(t, err, "the replacement message must stay active")
|
|
|
|
// Reloading is what the generation path does, so the input order has to
|
|
// survive the round trip and not just the returned slice.
|
|
history, err := f.DB.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{
|
|
ChatID: created.Chat.ID,
|
|
})
|
|
require.NoError(t, err)
|
|
require.GreaterOrEqual(t, len(history), 3)
|
|
tail := history[len(history)-3:]
|
|
require.Equal(t, []int64{
|
|
result.ReplacementMessage.ID,
|
|
result.SuffixMessages[0].ID,
|
|
result.SuffixMessages[1].ID,
|
|
}, []int64{tail[0].ID, tail[1].ID, tail[2].ID})
|
|
assertChatMessageText(t, tail[0], "edited prompt")
|
|
assertChatMessageText(t, tail[1], "first notice")
|
|
assertChatMessageText(t, tail[2], "second notice")
|
|
|
|
for _, message := range history {
|
|
require.NotEqual(t, targetID, message.ID,
|
|
"the edited message must not stay in active history")
|
|
}
|
|
}
|
|
|
|
// TestTransitionAbandon_RejectsUnowned verifies that calling Abandon
|
|
// on a chat the runner does not own returns ErrTransitionNotAllowed
|
|
// wrapped in a TransitionError that records the loaded from-state,
|
|
// without mutating chat state or publishing anything.
|
|
func TestTransitionAbandon_RejectsUnowned(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
created := createTestChat(t, f)
|
|
seeded := seededChat{chatID: created.Chat.ID, exists: true}
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
base := captureBaseline(ctx, t, f, seeded)
|
|
|
|
err := m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, aerr := tx.Abandon(chatstate.AbandonInput{})
|
|
return aerr
|
|
})
|
|
require.Error(t, err)
|
|
require.ErrorIs(t, err, chatstate.ErrTransitionNotAllowed)
|
|
var te *chatstate.TransitionError
|
|
require.ErrorAs(t, err, &te)
|
|
require.Equal(t, chatstate.TransitionAbandon, te.Transition)
|
|
// createTestChat lands the chat in R0; Abandon's precondition
|
|
// rejects an unowned chat there.
|
|
require.Equal(t, chatstate.StateR0, te.From)
|
|
assertNoMutationOrPublish(ctx, t, f, seeded.chatID, base)
|
|
}
|
|
|
|
// TestTransitionAbandon_ClearsOwnership verifies the Acquire/Abandon
|
|
// round-trip: after Acquire the chat carries a worker+runner and a
|
|
// fresh heartbeat row exists, and after Abandon both ownership fields
|
|
// are cleared. The heartbeat row is not deleted by Abandon; heartbeat
|
|
// cleanup is a separate concern.
|
|
func TestTransitionAbandon_ClearsOwnership(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
created := createTestChat(t, f)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
|
|
worker := uuid.New()
|
|
runner := uuid.New()
|
|
|
|
// Acquire writes ownership and a fresh heartbeat row.
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner})
|
|
return err
|
|
}))
|
|
owned := f.readChat(ctx, t, created.Chat.ID)
|
|
require.Equal(t, worker, owned.WorkerID.UUID)
|
|
require.Equal(t, runner, owned.RunnerID.UUID)
|
|
hb, err := f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{
|
|
ChatID: created.Chat.ID,
|
|
RunnerID: runner,
|
|
})
|
|
require.NoError(t, err, "Acquire writes a fresh heartbeat row")
|
|
require.Equal(t, runner, hb.RunnerID)
|
|
|
|
// Abandon clears ownership but leaves the heartbeat row intact.
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Abandon(chatstate.AbandonInput{})
|
|
return err
|
|
}))
|
|
hb, err = f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{
|
|
ChatID: created.Chat.ID,
|
|
RunnerID: runner,
|
|
})
|
|
require.NoError(t, err, "Abandon does not delete the heartbeat row")
|
|
abandoned := f.readChat(ctx, t, created.Chat.ID)
|
|
require.False(t, abandoned.WorkerID.Valid, "Abandon clears worker_id")
|
|
require.False(t, abandoned.RunnerID.Valid, "Abandon clears runner_id")
|
|
}
|
|
|
|
// TestTransitionAcquire_OverwritesFreshOwnership verifies that Acquire
|
|
// is an unconditional ownership handoff: a second worker calling
|
|
// Acquire on a chat that was *just* acquired by another worker
|
|
// successfully replaces ownership without inspecting heartbeat
|
|
// freshness. It also asserts that Acquire itself does not request an
|
|
// ownership hint, so the post-commit publish stays quiet on
|
|
// `chat:ownership` when the resulting heartbeat is fresh.
|
|
func TestTransitionAcquire_OverwritesFreshOwnership(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
created := createTestChat(t, f)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
|
|
firstWorker := uuid.New()
|
|
firstRunner := uuid.New()
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Acquire(chatstate.AcquireInput{WorkerID: firstWorker, RunnerID: firstRunner})
|
|
return err
|
|
}))
|
|
|
|
// The chat is now owned with a fresh (chat_id, firstRunner)
|
|
// heartbeat written by the first Acquire.
|
|
firstChat := f.readChat(ctx, t, created.Chat.ID)
|
|
require.Equal(t, firstWorker, firstChat.WorkerID.UUID)
|
|
require.Equal(t, firstRunner, firstChat.RunnerID.UUID)
|
|
_, err := f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{
|
|
ChatID: created.Chat.ID,
|
|
RunnerID: firstRunner,
|
|
})
|
|
require.NoError(t, err, "first Acquire wrote a fresh heartbeat")
|
|
// Sanity check: heartbeat is not stale by the same threshold the
|
|
// machine uses for ownership-hint decisions.
|
|
stale, err := f.DB.IsChatHeartbeatStale(ctx, database.IsChatHeartbeatStaleParams{
|
|
ChatID: created.Chat.ID,
|
|
RunnerID: firstRunner,
|
|
StaleSeconds: chatstate.HeartbeatStaleSeconds,
|
|
})
|
|
require.NoError(t, err)
|
|
require.False(t, stale, "first runner's heartbeat is fresh before the second Acquire")
|
|
|
|
// Snapshot publish counts before the takeover so we can assert
|
|
// Acquire does not publish an ownership hint itself.
|
|
ownershipBefore := f.Pub.ownershipPublishCount()
|
|
beforeChat := f.readChat(ctx, t, created.Chat.ID)
|
|
|
|
secondWorker := uuid.New()
|
|
secondRunner := uuid.New()
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Acquire(chatstate.AcquireInput{WorkerID: secondWorker, RunnerID: secondRunner})
|
|
return err
|
|
}))
|
|
|
|
after := f.readChat(ctx, t, created.Chat.ID)
|
|
require.Equal(t, secondWorker, after.WorkerID.UUID, "ownership replaced")
|
|
require.Equal(t, secondRunner, after.RunnerID.UUID, "runner replaced")
|
|
require.Equal(t, beforeChat.SnapshotVersion+1, after.SnapshotVersion, "snapshot bumps exactly once")
|
|
f.Pub.expectChatUpdate(t, created.Chat.ID, after.SnapshotVersion)
|
|
|
|
// The new (chat_id, secondRunner) heartbeat exists. The old
|
|
// (chat_id, firstRunner) row may or may not exist; Acquire is not
|
|
// responsible for cleaning it up.
|
|
_, err = f.DB.GetChatHeartbeat(ctx, database.GetChatHeartbeatParams{
|
|
ChatID: created.Chat.ID,
|
|
RunnerID: secondRunner,
|
|
})
|
|
require.NoError(t, err, "second Acquire wrote a heartbeat for the new runner")
|
|
|
|
// Acquire does not publish an ownership hint when it writes a fresh
|
|
// heartbeat. The post-commit ownership-hint logic in Update stays
|
|
// quiet because the new heartbeat is fresh, so no `chat:ownership`
|
|
// notification fires.
|
|
require.Equal(t, ownershipBefore, f.Pub.ownershipPublishCount(),
|
|
"Acquire must not publish an ownership hint when the resulting heartbeat is fresh")
|
|
}
|
|
|
|
// TestTransitionAcquire_ExecutionStateOrthogonal verifies that Acquire
|
|
// preserves every execution-state field on the chat across
|
|
// representative valid execution states, including idle, runnable, and
|
|
// archived states. The transition only mutates ownership.
|
|
func TestTransitionAcquire_ExecutionStateOrthogonal(t *testing.T) {
|
|
t.Parallel()
|
|
|
|
// Each setup leaves the chat in the named state and returns the
|
|
// chat ID for downstream assertions.
|
|
cases := []struct {
|
|
name string
|
|
state chatstate.ExecutionState
|
|
setup func(t *testing.T, f *testFixture) uuid.UUID
|
|
}{
|
|
{
|
|
name: "R0",
|
|
state: chatstate.StateR0,
|
|
setup: func(t *testing.T, f *testFixture) uuid.UUID {
|
|
return createTestChat(t, f).Chat.ID
|
|
},
|
|
},
|
|
{
|
|
name: "W",
|
|
state: chatstate.StateW,
|
|
setup: func(t *testing.T, f *testFixture) uuid.UUID {
|
|
created := createTestChat(t, f)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.FinishTurn(chatstate.FinishTurnInput{})
|
|
return err
|
|
}))
|
|
return created.Chat.ID
|
|
},
|
|
},
|
|
{
|
|
name: "E0",
|
|
state: chatstate.StateE0,
|
|
setup: func(t *testing.T, f *testFixture) uuid.UUID {
|
|
created := createTestChat(t, f)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.FinishError(chatstate.FinishErrorInput{
|
|
LastError: pqtype.NullRawMessage{
|
|
RawMessage: json.RawMessage(`{"message":"boom"}`),
|
|
Valid: true,
|
|
},
|
|
})
|
|
return err
|
|
}))
|
|
return created.Chat.ID
|
|
},
|
|
},
|
|
{
|
|
name: "I0",
|
|
state: chatstate.StateI0,
|
|
setup: func(t *testing.T, f *testFixture) uuid.UUID {
|
|
created := createTestChat(t, f)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Interrupt(chatstate.InterruptInput{Reason: "test"})
|
|
return err
|
|
}))
|
|
return created.Chat.ID
|
|
},
|
|
},
|
|
{
|
|
name: "XW",
|
|
state: chatstate.StateXW,
|
|
setup: func(t *testing.T, f *testFixture) uuid.UUID {
|
|
created := createTestChat(t, f)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID)
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.FinishTurn(chatstate.FinishTurnInput{})
|
|
return err
|
|
}))
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.SetArchived(chatstate.SetArchivedInput{Archived: true})
|
|
return err
|
|
}))
|
|
return created.Chat.ID
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range cases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
t.Parallel()
|
|
f := newTestFixture(t)
|
|
ctx := testutil.Context(t, testutil.WaitShort)
|
|
chatID := tc.setup(t, f)
|
|
require.Equal(t, tc.state, f.classify(ctx, t, chatID), "test setup must leave chat in %s", tc.state)
|
|
|
|
before := f.readChat(ctx, t, chatID)
|
|
queueBefore, err := f.DB.CountChatQueuedMessages(ctx, chatID)
|
|
require.NoError(t, err)
|
|
historyBefore := historyMessageIDs(ctx, t, f, chatID)
|
|
|
|
worker := uuid.New()
|
|
runner := uuid.New()
|
|
m := chatstate.NewChatMachine(f.DB, f.Pub, chatID)
|
|
require.NoError(t, m.Update(ctx, func(tx *chatstate.Tx, store database.Store) error {
|
|
_, err := tx.Acquire(chatstate.AcquireInput{WorkerID: worker, RunnerID: runner})
|
|
return err
|
|
}))
|
|
|
|
after := f.readChat(ctx, t, chatID)
|
|
// Ownership updated.
|
|
require.Equal(t, worker, after.WorkerID.UUID)
|
|
require.Equal(t, runner, after.RunnerID.UUID)
|
|
// Execution state preserved.
|
|
require.Equal(t, before.Status, after.Status, "status preserved")
|
|
require.Equal(t, before.Archived, after.Archived, "archived flag preserved")
|
|
require.Equal(t, before.RequiresActionDeadlineAt, after.RequiresActionDeadlineAt, "requires-action deadline preserved")
|
|
require.Equal(t, before.LastError, after.LastError, "last_error preserved")
|
|
require.Equal(t, before.HistoryVersion, after.HistoryVersion, "history_version preserved")
|
|
require.Equal(t, before.QueueVersion, after.QueueVersion, "queue_version preserved")
|
|
require.Equal(t, before.GenerationAttempt, after.GenerationAttempt, "generation_attempt preserved")
|
|
// Classified state unchanged.
|
|
require.Equal(t, tc.state, f.classify(ctx, t, chatID), "execution state preserved by Acquire")
|
|
// Queue and history rows untouched.
|
|
queueAfter, err := f.DB.CountChatQueuedMessages(ctx, chatID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, queueBefore, queueAfter, "queue cardinality preserved")
|
|
require.Equal(t, historyBefore, historyMessageIDs(ctx, t, f, chatID), "history preserved")
|
|
})
|
|
}
|
|
}
|