diff --git a/server/internal/handler/chat_input_ownership_test.go b/server/internal/handler/chat_input_ownership_test.go index 93cfe4a3442..304b0ded3b1 100644 --- a/server/internal/handler/chat_input_ownership_test.go +++ b/server/internal/handler/chat_input_ownership_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "testing" + "github.com/jackc/pgx/v5/pgtype" db "github.com/multica-ai/multica/server/pkg/db/generated" "github.com/multica-ai/multica/server/pkg/protocol" ) @@ -138,6 +139,244 @@ func TestDirectChat_RunningTaskDoesNotAbsorbNewMessage(t *testing.T) { } } +// TestChannelChat_QueuedSuccessorKeepsMessageAcrossPredecessorReply reproduces +// the Feishu race where U2 is persisted and T2 queued while T1 runs, then T1 +// writes an assistant row before T2 is claimed. Channel enqueue must seal U2 +// onto T2 before that reply changes the legacy trailing-message cursor. +func TestChannelChat_QueuedSuccessorKeepsMessageAcrossPredecessorReply(t *testing.T) { + if testHandler == nil { + t.Skip("database not available") + } + ctx := context.Background() + _, sessionID, runtimeID, daemonID := setupDirectChatSession(t, ctx, "channel ownership chat") + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load chat session: %v", err) + } + appendUser := func(content string) pgtype.UUID { + t.Helper() + message, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: content, + }) + if err != nil { + t.Fatalf("append channel user message: %v", err) + } + return message.ID + } + + firstMessage := appendUser("今天北京天气如何?") + first, err := testHandler.TaskService.EnqueueChannelChatTask(ctx, session, parseUUID(testUserID), false, []pgtype.UUID{firstMessage}) + if err != nil { + t.Fatalf("enqueue first channel task: %v", err) + } + claimed := claimTaskForRuntimeGuard(t, runtimeID, daemonID) + if claimed.ChatMessage != "今天北京天气如何?" { + t.Fatalf("first claim input = %q", claimed.ChatMessage) + } + markTaskRunning(t, ctx, uuidToString(first.ID)) + + secondMessage := appendUser("今天上海天气怎么样") + second, err := testHandler.TaskService.EnqueueChannelChatTask(ctx, session, parseUUID(testUserID), false, []pgtype.UUID{secondMessage}) + if err != nil { + t.Fatalf("enqueue second channel task: %v", err) + } + if _, err := testHandler.TaskService.CompleteTask(ctx, first.ID, completeResult(t, "北京天气回复"), "", ""); err != nil { + t.Fatalf("complete first channel task: %v", err) + } + + claimed2 := claimTaskForRuntimeGuard(t, runtimeID, daemonID) + if claimed2.ChatMessage != "今天上海天气怎么样" { + t.Fatalf("second claim lost its owned input; got %q", claimed2.ChatMessage) + } + assertTaskInputOwner(t, ctx, uuidToString(second.ID), uuidToString(second.ID)) +} + +func TestChannelChat_DuplicateMessageIDsAreClaimedOnce(t *testing.T) { + if testHandler == nil { + t.Skip("database not available") + } + ctx := context.Background() + _, sessionID, _, _ := setupDirectChatSession(t, ctx, "channel duplicate input") + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load chat session: %v", err) + } + message, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: "only once", + }) + if err != nil { + t.Fatalf("append channel user message: %v", err) + } + + task, err := testHandler.TaskService.EnqueueChannelChatTask( + ctx, session, parseUUID(testUserID), false, []pgtype.UUID{message.ID, message.ID}, + ) + if err != nil { + t.Fatalf("enqueue duplicate channel input: %v", err) + } + owned, err := testHandler.Queries.ListChatInputMessages(ctx, task.ID) + if err != nil { + t.Fatalf("list owned channel input: %v", err) + } + if len(owned) != 1 || owned[0].ID != message.ID { + t.Fatalf("owned input = %+v, want the message exactly once", owned) + } +} + +func TestChannelChat_RecoversUnownedInputAfterRouterRestart(t *testing.T) { + if testHandler == nil { + t.Skip("database not available") + } + ctx := context.Background() + _, sessionID, _, _ := setupDirectChatSession(t, ctx, "channel restart recovery") + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load chat session: %v", err) + } + appendUser := func(content string) pgtype.UUID { + t.Helper() + message, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: content, + }) + if err != nil { + t.Fatalf("append channel input: %v", err) + } + return message.ID + } + + appendUser("U1 before restart") + u2 := appendUser("U2 after restart") + task, err := testHandler.TaskService.EnqueueChannelChatTask( + ctx, session, parseUUID(testUserID), false, []pgtype.UUID{u2}, + ) + if err != nil { + t.Fatalf("enqueue post-restart channel task: %v", err) + } + owned, err := testHandler.Queries.ListChatInputMessages(ctx, task.ID) + if err != nil { + t.Fatalf("list recovered input: %v", err) + } + if got := msgContents(owned); len(got) != 2 || got[0] != "U1 before restart" || got[1] != "U2 after restart" { + t.Fatalf("recovered input = %v, want U1 then U2", got) + } +} + +func TestChannelChat_RecoverySkipsAnsweredLegacyHistory(t *testing.T) { + if testHandler == nil { + t.Skip("database not available") + } + ctx := context.Background() + agentID, sessionID, runtimeID, _ := setupDirectChatSession(t, ctx, "channel rolling upgrade recovery") + legacyUser, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: parseUUID(sessionID), + Role: "user", + Content: "answered before upgrade", + }) + if err != nil { + t.Fatalf("create legacy input: %v", err) + } + legacyTask := insertChannelChatTask(t, ctx, agentID, runtimeID, sessionID) + if _, err := testHandler.TaskService.CompleteTask(ctx, parseUUID(legacyTask), completeResult(t, "legacy answer"), "", ""); err != nil { + t.Fatalf("complete legacy channel task: %v", err) + } + var legacyOwner pgtype.UUID + if err := testPool.QueryRow(ctx, `SELECT task_id FROM chat_message WHERE id = $1`, legacyUser).Scan(&legacyOwner); err == nil && legacyOwner.Valid { + t.Fatalf("legacy channel input unexpectedly owned by %s", uuidToString(legacyOwner)) + } + + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load chat session: %v", err) + } + recovered, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: "unflushed after upgrade", + }) + if err != nil { + t.Fatalf("create recoverable input: %v", err) + } + current, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: "current after restart", + }) + if err != nil { + t.Fatalf("create current input: %v", err) + } + task, err := testHandler.TaskService.EnqueueChannelChatTask(ctx, session, parseUUID(testUserID), false, []pgtype.UUID{current.ID}) + if err != nil { + t.Fatalf("enqueue post-upgrade channel task: %v", err) + } + owned, err := testHandler.Queries.ListChatInputMessages(ctx, task.ID) + if err != nil { + t.Fatalf("list recovered input: %v", err) + } + if got := msgContents(owned); len(got) != 2 || got[0] != "unflushed after upgrade" || got[1] != "current after restart" { + t.Fatalf("recovered input = %v, want only post-assistant U1 then U2", got) + } + var recoveredOwner string + if err := testPool.QueryRow(ctx, `SELECT task_id FROM chat_message WHERE id = $1`, recovered.ID).Scan(&recoveredOwner); err != nil { + t.Fatalf("read recovered owner: %v", err) + } + if recoveredOwner != uuidToString(task.ID) { + t.Fatalf("recovered input owner = %s, want %s", recoveredOwner, uuidToString(task.ID)) + } +} + +func TestChannelChat_RecoversRolledBackBatchAfterRouterRestart(t *testing.T) { + if testHandler == nil { + t.Skip("database not available") + } + ctx := context.Background() + _, sessionID, _, _ := setupDirectChatSession(t, ctx, "channel failed flush recovery") + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load chat session: %v", err) + } + appendUser := func(content string) pgtype.UUID { + t.Helper() + message, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: session.ID, + Role: "user", + Content: content, + }) + if err != nil { + t.Fatalf("append channel input: %v", err) + } + return message.ID + } + + u1 := appendUser("failed U1") + missing := parseUUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa") + if _, err := testHandler.TaskService.EnqueueChannelChatTask( + ctx, session, parseUUID(testUserID), false, []pgtype.UUID{u1, missing}, + ); err == nil { + t.Fatal("enqueue with a missing explicit message must roll back") + } + + u2 := appendUser("later U2") + task, err := testHandler.TaskService.EnqueueChannelChatTask( + ctx, session, parseUUID(testUserID), false, []pgtype.UUID{u2}, + ) + if err != nil { + t.Fatalf("enqueue after simulated restart: %v", err) + } + owned, err := testHandler.Queries.ListChatInputMessages(ctx, task.ID) + if err != nil { + t.Fatalf("list recovered failed batch: %v", err) + } + if got := msgContents(owned); len(got) != 2 || got[0] != "failed U1" || got[1] != "later U2" { + t.Fatalf("recovered failed input = %v, want U1 then U2", got) + } +} + // TestCompleteTask_ChatEmptyOutputWritesNoResponse: an empty final output is a // visible, terminal no_response outcome — exactly one assistant row with // message_kind='no_response' and a non-empty fallback body, task completed, and @@ -383,4 +622,46 @@ func TestCompleteTask_ChannelEmptyOutputWritesNoRow(t *testing.T) { if rows[0].MessageKind != protocol.ChatMessageKindMessage || rows[0].Content != "channel reply" { t.Fatalf("channel message = kind %q content %q, want message/'channel reply'", rows[0].MessageKind, rows[0].Content) } + + // A task-owned channel batch must keep the same silent-empty semantics. + var installationID string + if err := testPool.QueryRow(ctx, ` + INSERT INTO channel_installation (workspace_id, agent_id, channel_type, config, installer_user_id) + VALUES ($1, $2, 'feishu', '{}', $3) + RETURNING id + `, testWorkspaceID, agentID, testUserID).Scan(&installationID); err != nil { + t.Fatalf("create channel installation: %v", err) + } + if _, err := testPool.Exec(ctx, ` + INSERT INTO channel_chat_session_binding + (chat_session_id, installation_id, channel_type, channel_chat_id, chat_type) + VALUES ($1, $2, 'feishu', $3, 'p2p') + `, sessionID, installationID, "owned-channel-"+sessionID); err != nil { + t.Fatalf("create channel binding: %v", err) + } + message, err := testHandler.Queries.CreateChatMessage(ctx, db.CreateChatMessageParams{ + ChatSessionID: parseUUID(sessionID), + Role: "user", + Content: "tool-only channel turn", + }) + if err != nil { + t.Fatalf("create owned channel input: %v", err) + } + session, err := testHandler.Queries.GetChatSession(ctx, parseUUID(sessionID)) + if err != nil { + t.Fatalf("load channel session: %v", err) + } + owned, err := testHandler.TaskService.EnqueueChannelChatTask( + ctx, session, parseUUID(testUserID), false, []pgtype.UUID{message.ID}, + ) + if err != nil { + t.Fatalf("enqueue owned channel task: %v", err) + } + markTaskRunning(t, ctx, uuidToString(owned.ID)) + if _, err := testHandler.TaskService.CompleteTask(ctx, owned.ID, completeResult(t, " "), "", ""); err != nil { + t.Fatalf("complete owned channel task: %v", err) + } + if rows := assistantRows(t, ctx, sessionID); len(rows) != 1 { + t.Fatalf("owned channel empty completion must not add an assistant row, got %d total", len(rows)) + } } diff --git a/server/internal/integrations/channel/engine/batcher.go b/server/internal/integrations/channel/engine/batcher.go index 6cdb156de54..ffe6abf8a0e 100644 --- a/server/internal/integrations/channel/engine/batcher.go +++ b/server/internal/integrations/channel/engine/batcher.go @@ -103,6 +103,30 @@ func (b *pendingBatcher) Schedule(key string, flush func()) { b.mu.Unlock() } +// ScheduleIfAbsent arms a retry without replacing a newer inbound message's +// pending flush closure. It returns whether a timer was added. +func (b *pendingBatcher) ScheduleIfAbsent(key string, flush func()) bool { + b.mu.Lock() + if b.stopped { + b.mu.Unlock() + flush() + return true + } + if _, ok := b.pending[key]; ok { + b.mu.Unlock() + return false + } + b.seq++ + gen := b.seq + b.pending[key] = &pendingEntry{ + flush: flush, + gen: gen, + timer: b.afterFunc(b.window, func() { b.onFire(key, gen) }), + } + b.mu.Unlock() + return true +} + // onFire runs the flush for key if it is still the live, armed generation. It // is the timer callback; in production it runs on time.AfterFunc's goroutine, // so the flush is naturally detached from the inbound path. diff --git a/server/internal/integrations/channel/engine/batcher_test.go b/server/internal/integrations/channel/engine/batcher_test.go index ff0d0b8e912..3dc2ba58aa2 100644 --- a/server/internal/integrations/channel/engine/batcher_test.go +++ b/server/internal/integrations/channel/engine/batcher_test.go @@ -138,6 +138,29 @@ func TestPendingBatcher_StaleTimerFireIsNoop(t *testing.T) { } } +func TestPendingBatcher_ScheduleIfAbsentPreservesNewerFlush(t *testing.T) { + f := &fakeTimerFactory{} + b := newTestBatcher(f) + var calls []string + + b.Schedule("s", func() { calls = append(calls, "inbound") }) + if b.ScheduleIfAbsent("s", func() { calls = append(calls, "retry") }) { + t.Fatal("retry must not replace an already armed inbound flush") + } + f.fireArmed() + if len(calls) != 1 || calls[0] != "inbound" { + t.Fatalf("flush calls = %v, want only the newer inbound closure", calls) + } + + if !b.ScheduleIfAbsent("s", func() { calls = append(calls, "retry") }) { + t.Fatal("retry must arm when the session has no pending flush") + } + f.fireArmed() + if len(calls) != 2 || calls[1] != "retry" { + t.Fatalf("flush calls = %v, want retry after inbound", calls) + } +} + func TestPendingBatcher_FlushAllDrainsPending(t *testing.T) { f := &fakeTimerFactory{} b := newTestBatcher(f) diff --git a/server/internal/integrations/channel/engine/resolvers.go b/server/internal/integrations/channel/engine/resolvers.go index 80b1d76a6cc..41edd19591c 100644 --- a/server/internal/integrations/channel/engine/resolvers.go +++ b/server/internal/integrations/channel/engine/resolvers.go @@ -99,6 +99,8 @@ type AppendParams struct { // AppendResult reports what AppendMessage decided. type AppendResult struct { + // MessageID is the persisted chat_message row for this inbound message. + MessageID pgtype.UUID // IssueCommand is non-nil when the message was an /issue command. IssueCommand *IssueCommand // DedupMarked is true when AppendMessage finalized the dedup claim in its @@ -215,7 +217,7 @@ type IssueCreator interface { // TaskEnqueuer is the narrow subset of service.TaskService the Router needs to // trigger a chat run. Shared across platforms. type TaskEnqueuer interface { - EnqueueChatTask(ctx context.Context, session db.ChatSession, initiatorUserID pgtype.UUID, forceFreshSession bool) (db.AgentTaskQueue, error) + EnqueueChannelChatTask(ctx context.Context, session db.ChatSession, initiatorUserID pgtype.UUID, forceFreshSession bool, messageIDs []pgtype.UUID) (db.AgentTaskQueue, error) } // SessionReader reads the rows the debounced flush + /issue identifier need. diff --git a/server/internal/integrations/channel/engine/router.go b/server/internal/integrations/channel/engine/router.go index 1335ea54875..6386fec3a1d 100644 --- a/server/internal/integrations/channel/engine/router.go +++ b/server/internal/integrations/channel/engine/router.go @@ -43,10 +43,18 @@ type Router struct { logger *slog.Logger - pendingFreshMu sync.Mutex + flushMu sync.Mutex + flushLocks map[string]*flushLock + pendingInputMu sync.Mutex + pendingInput map[string][]pgtype.UUID pendingFresh map[string]bool } +type flushLock struct { + mu sync.Mutex + refCount int +} + // Config tunes the Router. Zero values default. type RouterConfig struct { // ReplyTimeout caps a single detached OutboundReplier.Reply / typing @@ -74,7 +82,9 @@ func NewRouter(issues IssueCreator, tasks TaskEnqueuer, reader SessionReader, cf reader: reader, replyTimeout: cfg.ReplyTimeout, logger: cfg.Logger, + flushLocks: make(map[string]*flushLock), pendingFresh: make(map[string]bool), + pendingInput: make(map[string][]pgtype.UUID), } } @@ -311,24 +321,30 @@ func (r *Router) processClaimed(ctx context.Context, set ResolverSet, msg channe // THIS message's sender (the task initiator), deliberately not the // session creator (group sessions are creator=installer). Latest sender // in a window wins (MUL-2645). - r.scheduleRun(set, inst, msg, sessionID, identity.UserID) + r.scheduleRun(set, inst, msg, sessionID, identity.UserID, appendRes.MessageID) return res, postAppendFinalize, nil } // scheduleRun hands the per-session run trigger to the debouncer (or fires it // inline when batching is disabled). -func (r *Router) scheduleRun(set ResolverSet, inst ResolvedInstallation, msg channel.InboundMessage, sessionID, initiatorUserID pgtype.UUID) { +func (r *Router) scheduleRun(set ResolverSet, inst ResolvedInstallation, msg channel.InboundMessage, sessionID, initiatorUserID, messageID pgtype.UUID) { key := keyForSession(sessionID) fresh := msg.ForceFresh - if r.batcher == nil { - r.flushChatRun(set, inst, msg, sessionID, initiatorUserID, fresh) - return - } + r.appendPendingInput(key, messageID) if fresh { r.markPendingFresh(key) } + if r.batcher == nil { + messageIDs, forceFresh := r.takePendingBatch(key, fresh) + r.flushChatRun(set, inst, msg, sessionID, initiatorUserID, forceFresh, messageIDs) + return + } flush := func() { - r.flushChatRun(set, inst, msg, sessionID, initiatorUserID, r.takePendingFresh(key, fresh)) + messageIDs, forceFresh := r.takePendingBatch(key, fresh) + if len(messageIDs) == 0 { + return + } + r.flushChatRun(set, inst, msg, sessionID, initiatorUserID, forceFresh, messageIDs) } r.batcher.Schedule(key, flush) } @@ -340,18 +356,22 @@ const chatRunFlushTimeout = 10 * time.Second // flushChatRun is the debounced run-trigger: reload session, enqueue exactly // one chat task for the window, and emit the offline/archived notice (only // known here now) via the replier. Errors are logged, not returned. -func (r *Router) flushChatRun(set ResolverSet, inst ResolvedInstallation, msg channel.InboundMessage, sessionID, initiatorUserID pgtype.UUID, forceFresh bool) { +func (r *Router) flushChatRun(set ResolverSet, inst ResolvedInstallation, msg channel.InboundMessage, sessionID, initiatorUserID pgtype.UUID, forceFresh bool, messageIDs []pgtype.UUID) { + unlock := r.lockFlush(keyForSession(sessionID)) + defer unlock() + ctx, cancel := context.WithTimeout(context.Background(), chatRunFlushTimeout) defer cancel() session, err := r.reader.GetChatSession(ctx, sessionID) if err != nil { + r.restoreAndScheduleRun(set, inst, msg, sessionID, initiatorUserID, messageIDs, forceFresh) r.logger.Error("channel router: flush reload chat session failed", "chat_session_id", uuidString(sessionID), "err", err.Error()) r.clearTyping(ctx, set, sessionID) return } - if _, err := r.tasks.EnqueueChatTask(ctx, session, initiatorUserID, forceFresh); err != nil { + if _, err := r.tasks.EnqueueChannelChatTask(ctx, session, initiatorUserID, forceFresh, messageIDs); err != nil { // No task was enqueued, so no task lifecycle event will ever publish and // the platform's bus-driven typing clear can never fire. Clear the // indicator here (before any notice) so the "processing" reaction does @@ -363,12 +383,67 @@ func (r *Router) flushChatRun(set ResolverSet, inst ResolvedInstallation, msg ch case errors.Is(err, service.ErrChatTaskAgentArchived): r.emitFlushReply(ctx, set, inst, msg, sessionID, OutcomeAgentArchived) default: + r.restoreAndScheduleRun(set, inst, msg, sessionID, initiatorUserID, messageIDs, forceFresh) r.logger.Error("channel router: flush enqueue chat task failed", "chat_session_id", uuidString(sessionID), "err", err.Error()) } } } +func (r *Router) lockFlush(key string) func() { + r.flushMu.Lock() + lock := r.flushLocks[key] + if lock == nil { + lock = &flushLock{} + r.flushLocks[key] = lock + } + lock.refCount++ + r.flushMu.Unlock() + + lock.mu.Lock() + return func() { + lock.mu.Unlock() + r.flushMu.Lock() + lock.refCount-- + if lock.refCount == 0 { + delete(r.flushLocks, key) + } + r.flushMu.Unlock() + } +} + +func (r *Router) restoreAndScheduleRun(set ResolverSet, inst ResolvedInstallation, msg channel.InboundMessage, sessionID, initiatorUserID pgtype.UUID, ids []pgtype.UUID, forceFresh bool) { + key := keyForSession(sessionID) + r.restorePendingBatch(key, ids, forceFresh) + if r.batcher == nil { + return + } + r.batcher.ScheduleIfAbsent(key, func() { + messageIDs, retryFresh := r.takePendingBatch(key, forceFresh) + if len(messageIDs) == 0 { + return + } + r.flushChatRun(set, inst, msg, sessionID, initiatorUserID, retryFresh, messageIDs) + }) +} + +func (r *Router) appendPendingInput(key string, id pgtype.UUID) { + r.pendingInputMu.Lock() + defer r.pendingInputMu.Unlock() + r.pendingInput[key] = append(r.pendingInput[key], id) +} + +func (r *Router) restorePendingBatch(key string, ids []pgtype.UUID, forceFresh bool) { + r.pendingInputMu.Lock() + defer r.pendingInputMu.Unlock() + // Preserve arrival order when messages enter a new debounce window while + // the previous batch is being flushed. + r.pendingInput[key] = append(append([]pgtype.UUID(nil), ids...), r.pendingInput[key]...) + if forceFresh { + r.pendingFresh[key] = true + } +} + // clearTyping asks the platform to drop the "processing" indicator for a session // whose flush produced no task run. A nil TypingNotifier (platform without the // feature) is a no-op. @@ -379,17 +454,19 @@ func (r *Router) clearTyping(ctx context.Context, set ResolverSet, sessionID pgt } func (r *Router) markPendingFresh(key string) { - r.pendingFreshMu.Lock() - defer r.pendingFreshMu.Unlock() + r.pendingInputMu.Lock() + defer r.pendingInputMu.Unlock() r.pendingFresh[key] = true } -func (r *Router) takePendingFresh(key string, fallback bool) bool { - r.pendingFreshMu.Lock() - defer r.pendingFreshMu.Unlock() - fresh := fallback || r.pendingFresh[key] +func (r *Router) takePendingBatch(key string, fallbackFresh bool) ([]pgtype.UUID, bool) { + r.pendingInputMu.Lock() + defer r.pendingInputMu.Unlock() + ids := r.pendingInput[key] + delete(r.pendingInput, key) + fresh := fallbackFresh || r.pendingFresh[key] delete(r.pendingFresh, key) - return fresh + return ids, fresh } // emitFlushReply delivers an offline/archived notice for a flushed run. diff --git a/server/internal/integrations/channel/engine/router_test.go b/server/internal/integrations/channel/engine/router_test.go index 5afe8b8a7d5..93dc7971995 100644 --- a/server/internal/integrations/channel/engine/router_test.go +++ b/server/internal/integrations/channel/engine/router_test.go @@ -157,20 +157,46 @@ type fakeTasks struct { mu sync.Mutex called bool forceFresh bool + freshArgs []bool initiator pgtype.UUID + messageIDs [][]pgtype.UUID err error + gate chan struct{} + entered chan struct{} } -func (f *fakeTasks) EnqueueChatTask(_ context.Context, _ db.ChatSession, initiator pgtype.UUID, forceFresh bool) (db.AgentTaskQueue, error) { +func (f *fakeTasks) EnqueueChannelChatTask(_ context.Context, _ db.ChatSession, initiator pgtype.UUID, forceFresh bool, messageIDs []pgtype.UUID) (db.AgentTaskQueue, error) { + if f.entered != nil { + f.entered <- struct{}{} + } + if f.gate != nil { + <-f.gate + } f.mu.Lock() defer f.mu.Unlock() f.called = true f.forceFresh = forceFresh + f.freshArgs = append(f.freshArgs, forceFresh) f.initiator = initiator + f.messageIDs = append(f.messageIDs, append([]pgtype.UUID(nil), messageIDs...)) return db.AgentTaskQueue{}, f.err } func (f *fakeTasks) wasCalled() bool { f.mu.Lock(); defer f.mu.Unlock(); return f.called } func (f *fakeTasks) freshArg() bool { f.mu.Lock(); defer f.mu.Unlock(); return f.forceFresh } +func (f *fakeTasks) freshCalls() []bool { + f.mu.Lock() + defer f.mu.Unlock() + return append([]bool(nil), f.freshArgs...) +} +func (f *fakeTasks) inputCalls() [][]pgtype.UUID { + f.mu.Lock() + defer f.mu.Unlock() + calls := make([][]pgtype.UUID, len(f.messageIDs)) + for i := range f.messageIDs { + calls[i] = append([]pgtype.UUID(nil), f.messageIDs[i]...) + } + return calls +} type fakeReader struct { session db.ChatSession @@ -498,6 +524,104 @@ func TestRouter_FlushSuccess_DoesNotClearTyping(t *testing.T) { } } +func TestRouter_FlushFailureRestoresInputForNextFlush(t *testing.T) { + h := newHarness(t) + h.tasks.err = errors.New("transient enqueue failure") + firstID := uuidFromString(t, "66666666-6666-6666-6666-666666666666") + h.binder.appendResult.MessageID = firstID + first := p2pMessage(t) + first.ForceFresh = true + if err := h.router.Handle(context.Background(), first); err != nil { + t.Fatalf("first handle: %v", err) + } + + h.tasks.err = nil + secondID := uuidFromString(t, "77777777-7777-7777-7777-777777777777") + second := p2pMessage(t) + second.EventID = "evt-2" + second.MessageID = "om-2" + h.binder.appendResult.MessageID = secondID + if err := h.router.Handle(context.Background(), second); err != nil { + t.Fatalf("second handle: %v", err) + } + + calls := h.tasks.inputCalls() + if len(calls) != 2 { + t.Fatalf("enqueue calls = %d, want 2", len(calls)) + } + if len(calls[1]) != 2 || calls[1][0] != firstID || calls[1][1] != secondID { + t.Fatalf("retry input = %v, want failed batch followed by new input", calls[1]) + } + freshCalls := h.tasks.freshCalls() + if len(freshCalls) != 2 || !freshCalls[0] || !freshCalls[1] { + t.Fatalf("forceFresh calls = %v, want failed and retried batches to stay fresh", freshCalls) + } +} + +func TestRouter_FlushFailureRearmsRetryWithoutAnotherMessage(t *testing.T) { + h := newHarness(t) + timers := &fakeTimerFactory{} + h.router.batcher = newTestBatcher(timers) + h.tasks.err = errors.New("transient enqueue failure") + messageID := uuidFromString(t, "66666666-6666-6666-6666-666666666666") + h.binder.appendResult.MessageID = messageID + if err := h.router.Handle(context.Background(), p2pMessage(t)); err != nil { + t.Fatalf("handle: %v", err) + } + + timers.fireArmed() + if got := h.router.batcher.pendingCount(); got != 1 { + t.Fatalf("failed flush must rearm one retry timer, pending=%d", got) + } + h.tasks.mu.Lock() + h.tasks.err = nil + h.tasks.mu.Unlock() + timers.fireArmed() + + calls := h.tasks.inputCalls() + if len(calls) != 2 { + t.Fatalf("enqueue calls = %d, want failed attempt plus timer retry", len(calls)) + } + if len(calls[1]) != 1 || calls[1][0] != messageID { + t.Fatalf("retry input = %v, want original failed batch", calls[1]) + } +} + +func TestRouter_FlushesSameSessionSerially(t *testing.T) { + h := newHarness(t) + h.tasks.gate = make(chan struct{}) + h.tasks.entered = make(chan struct{}, 1) + firstID := uuidFromString(t, "66666666-6666-6666-6666-666666666666") + secondID := uuidFromString(t, "77777777-7777-7777-7777-777777777777") + + firstDone := make(chan struct{}) + go func() { + h.router.flushChatRun(ResolverSet{}, h.inst.inst, p2pMessage(t), h.binder.ensureID, h.ident.id.UserID, false, []pgtype.UUID{firstID}) + close(firstDone) + }() + <-h.tasks.entered + + secondDone := make(chan struct{}) + go func() { + h.router.flushChatRun(ResolverSet{}, h.inst.inst, p2pMessage(t), h.binder.ensureID, h.ident.id.UserID, false, []pgtype.UUID{secondID}) + close(secondDone) + }() + + select { + case <-h.tasks.entered: + t.Fatal("second same-session flush reached enqueue before the first completed") + case <-time.After(50 * time.Millisecond): + } + + close(h.tasks.gate) + <-firstDone + <-secondDone + calls := h.tasks.inputCalls() + if len(calls) != 2 || len(calls[0]) != 1 || calls[0][0] != firstID || len(calls[1]) != 1 || calls[1][0] != secondID { + t.Fatalf("flush input calls = %v, want first then second", calls) + } +} + func TestRouter_ForceFresh_Propagates(t *testing.T) { h := newHarness(t) msg := p2pMessage(t) diff --git a/server/internal/integrations/channel/engine/session.go b/server/internal/integrations/channel/engine/session.go index cba3e0a835b..934ed546d31 100644 --- a/server/internal/integrations/channel/engine/session.go +++ b/server/internal/integrations/channel/engine/session.go @@ -280,11 +280,12 @@ func (s *ChatSession) AppendUserMessage(ctx context.Context, in AppendInput) (Ap } } - if _, err := qtx.CreateChatMessage(ctx, db.CreateChatMessageParams{ + message, err := qtx.CreateChatMessage(ctx, db.CreateChatMessageParams{ ChatSessionID: in.SessionID, Role: "user", Content: in.Body, - }); err != nil { + }) + if err != nil { return AppendResult{}, fmt.Errorf("create chat message: %w", err) } if err := qtx.TouchChatSession(ctx, in.SessionID); err != nil { @@ -324,7 +325,7 @@ func (s *ChatSession) AppendUserMessage(ctx context.Context, in AppendInput) (Ap if err := tx.Commit(ctx); err != nil { return AppendResult{}, fmt.Errorf("commit: %w", err) } - return AppendResult{IssueCommand: cmd, DedupMarked: markedInTx}, nil + return AppendResult{MessageID: message.ID, IssueCommand: cmd, DedupMarked: markedInTx}, nil } func isUniqueViolation(err error) bool { diff --git a/server/internal/service/task.go b/server/internal/service/task.go index de320b2bbac..184dd400509 100644 --- a/server/internal/service/task.go +++ b/server/internal/service/task.go @@ -1449,6 +1449,112 @@ func (s *TaskService) EnqueueChatTask(ctx context.Context, chatSession db.ChatSe return task, nil } +// EnqueueChannelChatTask creates a channel task and atomically seals the +// currently pending channel messages as its immutable input batch. Unlike the +// legacy claim-time trailing-message scan, this ownership survives a predecessor +// writing its assistant reply before the queued successor is claimed. +func (s *TaskService) EnqueueChannelChatTask(ctx context.Context, chatSession db.ChatSession, initiatorUserID pgtype.UUID, forceFreshSession bool, messageIDs []pgtype.UUID) (db.AgentTaskQueue, error) { + messageIDs = uniqueUUIDs(messageIDs) + if len(messageIDs) == 0 { + return db.AgentTaskQueue{}, errors.New("channel chat task: empty input batch") + } + agent, err := s.Queries.GetAgent(ctx, chatSession.AgentID) + if err != nil { + return db.AgentTaskQueue{}, fmt.Errorf("load agent: %w", err) + } + if agent.ArchivedAt.Valid { + return db.AgentTaskQueue{}, ErrChatTaskAgentArchived + } + if !agent.RuntimeID.Valid { + return db.AgentTaskQueue{}, ErrChatTaskAgentNoRuntime + } + + overlay := s.buildRuntimeMCPOverlay(ctx, initiatorUserID, agent) + attr := attribution.DirectHumanRun(initiatorUserID, attribution.EvidenceChat, chatSession.ID) + attr, err = s.applyAttributionFallback(ctx, attr, agent) + if err != nil { + return db.AgentTaskQueue{}, err + } + attrSource, _, attrEvidenceKind, attrEvidenceRef := attributionCreateParams(attr) + + var task db.AgentTaskQueue + if err := s.runInTx(ctx, func(qtx *db.Queries) error { + if _, err := qtx.LockChatSessionForChannelInput(ctx, chatSession.ID); err != nil { + return fmt.Errorf("lock channel chat input: %w", err) + } + task, err = qtx.CreateChatTask(ctx, db.CreateChatTaskParams{ + AgentID: chatSession.AgentID, + RuntimeID: agent.RuntimeID, + Priority: 2, + ChatSessionID: chatSession.ID, + InitiatorUserID: initiatorUserID, + OriginatorUserID: initiatorUserID, + AccountableUserID: attr.AccountableUserID, + ForceFreshSession: pgtype.Bool{Bool: forceFreshSession, Valid: true}, + RuntimeMcpOverlay: overlay.Overlay, + RuntimeConnectedApps: overlay.ConnectedApps, + OriginatorSource: attrSource, + TriggerEvidenceKind: attrEvidenceKind, + TriggerEvidenceRefID: attrEvidenceRef, + }) + if err != nil { + return fmt.Errorf("create channel chat task: %w", err) + } + task, err = qtx.SetChatTaskInputOwnerSelf(ctx, task.ID) + if err != nil { + return fmt.Errorf("stamp channel chat input owner: %w", err) + } + claimed, err := qtx.ClaimChannelChatInputMessages(ctx, db.ClaimChannelChatInputMessagesParams{ + ChatSessionID: chatSession.ID, + TaskID: task.ID, + MessageIds: messageIDs, + }) + if err != nil { + return fmt.Errorf("claim channel chat input: %w", err) + } + if !containsAllUUIDs(claimed, messageIDs) { + return fmt.Errorf("claim channel chat input: one or more explicit messages were unavailable") + } + return nil + }); err != nil { + return db.AgentTaskQueue{}, err + } + + slog.Info("channel chat task enqueued", + "task_id", util.UUIDToString(task.ID), + "chat_session_id", util.UUIDToString(chatSession.ID), + "agent_id", util.UUIDToString(chatSession.AgentID)) + s.broadcastTaskEvent(ctx, protocol.EventTaskQueued, task) + s.NotifyTaskEnqueued(ctx, task) + return task, nil +} + +func uniqueUUIDs(ids []pgtype.UUID) []pgtype.UUID { + unique := make([]pgtype.UUID, 0, len(ids)) + seen := make(map[pgtype.UUID]struct{}, len(ids)) + for _, id := range ids { + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + unique = append(unique, id) + } + return unique +} + +func containsAllUUIDs(haystack, needles []pgtype.UUID) bool { + found := make(map[pgtype.UUID]struct{}, len(haystack)) + for _, id := range haystack { + found[id] = struct{}{} + } + for _, id := range needles { + if _, ok := found[id]; !ok { + return false + } + } + return true +} + // DirectChatSendResult carries the rows a transactional direct-chat send // persisted, so the handler can broadcast the user message and shape its // response without re-reading them. @@ -2766,8 +2872,14 @@ func (s *TaskService) writeChatCompletionOutcome(ctx context.Context, qtx *db.Qu // row, only an empty chat:done for typing/lifecycle. Keeps the Slack/Lark // silent-drop path. Attachments still force a row — the agent produced a // deliverable the user must see. - if isEmpty && pendingAttachments == 0 && !task.ChatInputTaskID.Valid { - return nil, nil + if isEmpty && pendingAttachments == 0 { + isChannel, err := qtx.IsChannelChatSession(ctx, task.ChatSessionID) + if err != nil { + return nil, fmt.Errorf("detect channel chat session: %w", err) + } + if isChannel || !task.ChatInputTaskID.Valid { + return nil, nil + } } params := db.CreateChatMessageParams{ diff --git a/server/pkg/db/generated/chat.sql.go b/server/pkg/db/generated/chat.sql.go index dbaf81d44b7..4581015e62f 100644 --- a/server/pkg/db/generated/chat.sql.go +++ b/server/pkg/db/generated/chat.sql.go @@ -31,6 +31,67 @@ func (q *Queries) ChatSessionHasUserMessage(ctx context.Context, chatSessionID p return has_user_message, err } +const claimChannelChatInputMessages = `-- name: ClaimChannelChatInputMessages :many +WITH boundary AS ( + SELECT created_at, id + FROM chat_message + WHERE chat_session_id = $2 + AND id = ANY($3::uuid[]) + AND role = 'user' + ORDER BY created_at DESC, id DESC + LIMIT 1 +), last_assistant AS ( + SELECT created_at, id + FROM chat_message + WHERE chat_session_id = $2 + AND role = 'assistant' + ORDER BY created_at DESC, id DESC + LIMIT 1 +) +UPDATE chat_message m +SET task_id = $1 +FROM boundary b +LEFT JOIN last_assistant a ON TRUE +WHERE m.chat_session_id = $2 + AND m.role = 'user' + AND m.task_id IS NULL + AND (m.created_at, m.id) <= (b.created_at, b.id) + AND (a.id IS NULL OR (m.created_at, m.id) > (a.created_at, a.id)) +RETURNING m.id +` + +type ClaimChannelChatInputMessagesParams struct { + TaskID pgtype.UUID `json:"task_id"` + ChatSessionID pgtype.UUID `json:"chat_session_id"` + MessageIds []pgtype.UUID `json:"message_ids"` +} + +// Seals the explicit debounce batch plus older unowned channel user messages. +// The newest explicit row is the recovery boundary: this rediscovers durable +// messages whose in-memory trigger was lost across a process restart without +// absorbing messages appended after this flush began. The lower bound is the +// last assistant reply so a rolling deploy cannot replay already-answered +// legacy channel history whose user rows still have task_id NULL. +func (q *Queries) ClaimChannelChatInputMessages(ctx context.Context, arg ClaimChannelChatInputMessagesParams) ([]pgtype.UUID, error) { + rows, err := q.db.Query(ctx, claimChannelChatInputMessages, arg.TaskID, arg.ChatSessionID, arg.MessageIds) + if err != nil { + return nil, err + } + defer rows.Close() + items := []pgtype.UUID{} + for rows.Next() { + var id pgtype.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const createChatDraftRestore = `-- name: CreateChatDraftRestore :one INSERT INTO chat_draft_restore (id, chat_session_id, task_id, content, attachment_ids) VALUES ($1, $2, $3, $4, $5) @@ -585,6 +646,19 @@ func (q *Queries) HasPendingChatTasksByCreator(ctx context.Context, arg HasPendi return has_pending, err } +const isChannelChatSession = `-- name: IsChannelChatSession :one +SELECT EXISTS ( + SELECT 1 FROM channel_chat_session_binding WHERE chat_session_id = $1 +) +` + +func (q *Queries) IsChannelChatSession(ctx context.Context, chatSessionID pgtype.UUID) (bool, error) { + row := q.db.QueryRow(ctx, isChannelChatSession, chatSessionID) + var exists bool + err := row.Scan(&exists) + return exists, err +} + const linkChatMessageToTask = `-- name: LinkChatMessageToTask :exec UPDATE chat_message SET task_id = $2 @@ -1025,6 +1099,22 @@ func (q *Queries) ListPendingChatTasksByCreator(ctx context.Context, arg ListPen return items, nil } +const lockChatSessionForChannelInput = `-- name: LockChatSessionForChannelInput :one +SELECT id FROM chat_session +WHERE id = $1 +FOR UPDATE +` + +// Serializes channel flushes for one session. The lock also fences concurrent +// chat_message inserts through their FK key-share lock, so the recovery claim +// below has a stable upper boundary and cannot absorb a newer inbound message. +func (q *Queries) LockChatSessionForChannelInput(ctx context.Context, id pgtype.UUID) (pgtype.UUID, error) { + row := q.db.QueryRow(ctx, lockChatSessionForChannelInput, id) + var id_2 pgtype.UUID + err := row.Scan(&id_2) + return id_2, err +} + const lockChatSessionForDelete = `-- name: LockChatSessionForDelete :one SELECT id FROM chat_session WHERE id = $1 diff --git a/server/pkg/db/queries/chat.sql b/server/pkg/db/queries/chat.sql index fef86178a8a..cf9d9eabee2 100644 --- a/server/pkg/db/queries/chat.sql +++ b/server/pkg/db/queries/chat.sql @@ -132,6 +132,14 @@ SELECT id FROM chat_session WHERE id = $1 FOR UPDATE; +-- name: LockChatSessionForChannelInput :one +-- Serializes channel flushes for one session. The lock also fences concurrent +-- chat_message inserts through their FK key-share lock, so the recovery claim +-- below has a stable upper boundary and cannot absorb a newer inbound message. +SELECT id FROM chat_session +WHERE id = $1 +FOR UPDATE; + -- name: DeleteChatSession :exec -- Hard delete. chat_message rows cascade via FK ON DELETE CASCADE; the -- chat_session_id on agent_task_queue is set NULL by FK so completed/failed @@ -230,6 +238,45 @@ SET chat_input_task_id = id WHERE id = $1 RETURNING *; +-- name: ClaimChannelChatInputMessages :many +-- Seals the explicit debounce batch plus older unowned channel user messages. +-- The newest explicit row is the recovery boundary: this rediscovers durable +-- messages whose in-memory trigger was lost across a process restart without +-- absorbing messages appended after this flush began. The lower bound is the +-- last assistant reply so a rolling deploy cannot replay already-answered +-- legacy channel history whose user rows still have task_id NULL. +WITH boundary AS ( + SELECT created_at, id + FROM chat_message + WHERE chat_session_id = sqlc.arg(chat_session_id) + AND id = ANY(sqlc.arg(message_ids)::uuid[]) + AND role = 'user' + ORDER BY created_at DESC, id DESC + LIMIT 1 +), last_assistant AS ( + SELECT created_at, id + FROM chat_message + WHERE chat_session_id = sqlc.arg(chat_session_id) + AND role = 'assistant' + ORDER BY created_at DESC, id DESC + LIMIT 1 +) +UPDATE chat_message m +SET task_id = sqlc.arg(task_id) +FROM boundary b +LEFT JOIN last_assistant a ON TRUE +WHERE m.chat_session_id = sqlc.arg(chat_session_id) + AND m.role = 'user' + AND m.task_id IS NULL + AND (m.created_at, m.id) <= (b.created_at, b.id) + AND (a.id IS NULL OR (m.created_at, m.id) > (a.created_at, a.id)) +RETURNING m.id; + +-- name: IsChannelChatSession :one +SELECT EXISTS ( + SELECT 1 FROM channel_chat_session_binding WHERE chat_session_id = $1 +); + -- name: GetLastChatTaskSession :one -- Returns the most recent task in this chat session that managed to record a -- session_id. Includes both completed and failed tasks: even a failed task