From dea04b2bed88f82c35f0fd7d2b0ab44a9efef983 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Sun, 31 May 2026 12:36:13 +0100 Subject: [PATCH] fix(backend): close clarification overlay after agent MCP timeout and dedup retried questions (#1185) --- apps/backend/cmd/kandev/helpers.go | 8 +- .../internal/clarification/canceller.go | 75 +++++++--- .../internal/clarification/canceller_test.go | 72 +++++++++- .../internal/clarification/handlers.go | 6 +- .../internal/clarification/handlers_test.go | 8 +- apps/backend/internal/clarification/store.go | 89 +++++++++--- .../internal/clarification/store_test.go | 115 ++++++++++++--- apps/backend/internal/clarification/types.go | 12 +- .../integration/mcp_integration_test.go | 2 +- .../backend/internal/mcp/handlers/handlers.go | 28 ++-- .../internal/mcp/handlers/handlers_test.go | 133 ++++++++++++++++++ .../task/handlers/process_handlers_test.go | 3 + .../internal/task/repository/interface.go | 1 + .../task/repository/sqlite/message.go | 21 +++ 14 files changed, 490 insertions(+), 83 deletions(-) diff --git a/apps/backend/cmd/kandev/helpers.go b/apps/backend/cmd/kandev/helpers.go index 58eeec1f5e..e6cdcfa5ec 100644 --- a/apps/backend/cmd/kandev/helpers.go +++ b/apps/backend/cmd/kandev/helpers.go @@ -506,7 +506,7 @@ func registerRoutes(p routeParams) { p.gateway.SetupRoutes(p.router) registerTaskRoutes(p, planService, handoffSvc) - registerSecondaryRoutes(p, workflowCtrl, clarificationStore, planService, handoffSvc) + registerSecondaryRoutes(p, workflowCtrl, clarificationStore, clarificationCanceller, planService, handoffSvc) // /health is a readiness probe, not a liveness probe. It only // returns 200 after main has flipped the package-level `ready` @@ -662,6 +662,7 @@ func registerSecondaryRoutes( p routeParams, workflowCtrl *workflowcontroller.Controller, clarificationStore *clarification.Store, + clarificationCanceller *clarification.Canceller, planService *taskservice.PlanService, handoffSvc *taskservice.HandoffService, ) { @@ -775,7 +776,7 @@ func registerSecondaryRoutes( p.log.Debug("Registered Improve Kandev handlers (HTTP)") } - registerMCPAndDebugRoutes(p, workflowCtrl, clarificationStore, planService, handoffSvc) + registerMCPAndDebugRoutes(p, workflowCtrl, clarificationStore, clarificationCanceller, planService, handoffSvc) var automationSvc *automation.Service if p.services.Automation != nil { @@ -882,12 +883,13 @@ func registerMCPAndDebugRoutes( p routeParams, wfCtrl *workflowcontroller.Controller, clarificationStore *clarification.Store, + clarificationCanceller *clarification.Canceller, planService *taskservice.PlanService, handoffSvc *taskservice.HandoffService, ) { mcpHandlers := mcphandlers.NewHandlers( p.taskSvc, wfCtrl, - clarificationStore, p.msgCreator, p.taskRepo, p.taskRepo, p.eventBus, planService, p.orchestratorSvc, p.orchestratorSvc.GetMessageQueue(), p.log, + clarificationStore, clarificationCanceller, p.msgCreator, p.taskRepo, p.taskRepo, p.eventBus, planService, p.orchestratorSvc, p.orchestratorSvc.GetMessageQueue(), p.log, ) // Wire config-mode dependencies for agent-native configuration mcpHandlers.SetConfigDeps(p.services.Workflow, p.agentSettingsController, p.mcpConfigSvc) diff --git a/apps/backend/internal/clarification/canceller.go b/apps/backend/internal/clarification/canceller.go index ecf5199ddf..0cfceaa479 100644 --- a/apps/backend/internal/clarification/canceller.go +++ b/apps/backend/internal/clarification/canceller.go @@ -32,6 +32,39 @@ func NewCanceller(store *Store, repo messageStore, eventBus EventBus, log *logge } } +// isTerminalStatus returns true for statuses that should not be overwritten. +func isTerminalStatus(status string) bool { + switch status { + case string(StatusAnswered), string(StatusExpired), string(StatusRejected), string(StatusCancelled): + return true + } + return false +} + +// markMessagesExpired updates the given messages to status=expired and publishes +// a message.updated event for each one. It is idempotent: already-terminal +// messages are skipped. +func (c *Canceller) markMessagesExpired(ctx context.Context, msgs []*taskmodels.Message, pendingID string) { + for _, msg := range msgs { + if msg.Metadata == nil { + msg.Metadata = map[string]any{} + } + if current, _ := msg.Metadata["status"].(string); isTerminalStatus(current) { + continue + } + msg.Metadata["agent_disconnected"] = true + msg.Metadata["status"] = string(StatusExpired) + if err := c.repo.UpdateMessage(ctx, msg); err != nil { + c.logger.Warn("failed to update message with expired status", + zap.String("pending_id", pendingID), + zap.String("message_id", msg.ID), + zap.Error(err)) + continue + } + c.publishMessageUpdated(ctx, msg) + } +} + // CancelSessionAndNotify cancels all pending clarifications for a session, // unblocking WaitForResponse callers, and marks the database messages // as expired (with agent_disconnected=true for context) so the frontend @@ -44,10 +77,8 @@ func NewCanceller(store *Store, repo messageStore, eventBus EventBus, log *logge // forwarded to the orchestrator as "User declined to answer"). func (c *Canceller) CancelSessionAndNotify(ctx context.Context, sessionID string) int { pendingIDs := c.store.CancelSession(sessionID) - if len(pendingIDs) == 0 { - return 0 - } + handled := make(map[string]bool, len(pendingIDs)) for _, id := range pendingIDs { msgs, err := c.repo.FindMessagesByPendingID(ctx, id) if err != nil || len(msgs) == 0 { @@ -56,24 +87,30 @@ func (c *Canceller) CancelSessionAndNotify(ctx context.Context, sessionID string zap.Error(err)) continue } - for _, msg := range msgs { - if msg.Metadata == nil { - msg.Metadata = map[string]any{} - } - msg.Metadata["agent_disconnected"] = true - msg.Metadata["status"] = string(StatusExpired) - if err := c.repo.UpdateMessage(ctx, msg); err != nil { - c.logger.Warn("failed to update message with expired status", - zap.String("pending_id", id), - zap.String("message_id", msg.ID), - zap.Error(err)) - continue - } - c.publishMessageUpdated(ctx, msg) - } + c.markMessagesExpired(ctx, msgs, id) + handled[id] = true } - return len(pendingIDs) + // Defense-in-depth: also find any still-pending clarification messages in + // the DB whose in-memory store entry was already removed (e.g. by a racing + // timeout path). Group by pending_id and mark each bundle expired. + msgs, err := c.repo.FindPendingClarificationMessagesBySessionID(ctx, sessionID) + if err != nil || len(msgs) == 0 { + return len(pendingIDs) + } + + byPendingID := make(map[string][]*taskmodels.Message) + for _, msg := range msgs { + pid := stringFromMetadata(msg.Metadata, "pending_id") + if pid == "" || handled[pid] { + continue + } + byPendingID[pid] = append(byPendingID[pid], msg) + } + for pid, bundle := range byPendingID { + c.markMessagesExpired(ctx, bundle, pid) + } + return len(pendingIDs) + len(byPendingID) } // publishMessageUpdated publishes a message.updated event to the event bus. diff --git a/apps/backend/internal/clarification/canceller_test.go b/apps/backend/internal/clarification/canceller_test.go index 81ec97fa17..b814c88195 100644 --- a/apps/backend/internal/clarification/canceller_test.go +++ b/apps/backend/internal/clarification/canceller_test.go @@ -36,6 +36,20 @@ func (s *stubMessageStore) FindMessagesByPendingID(_ context.Context, pendingID return msgs, nil } +func (s *stubMessageStore) FindPendingClarificationMessagesBySessionID(_ context.Context, sessionID string) ([]*taskmodels.Message, error) { + var out []*taskmodels.Message + for _, msgs := range s.messages { + for _, m := range msgs { + if m.TaskSessionID == sessionID { + if status, _ := m.Metadata["status"].(string); status == "pending" { + out = append(out, m) + } + } + } + } + return out, nil +} + func (s *stubMessageStore) UpdateMessage(_ context.Context, m *taskmodels.Message) error { s.updated = append(s.updated, m) return nil @@ -67,7 +81,7 @@ func TestCanceller_MarksStatusExpiredOnDisconnect(t *testing.T) { msgs := map[string][]*taskmodels.Message{} c, repo, _ := newTestCanceller(t, msgs) - pendingID := c.store.CreateRequest(&Request{SessionID: "s1"}) + pendingID, _ := c.store.CreateRequest(&Request{SessionID: "s1"}) msgs[pendingID] = []*taskmodels.Message{{ ID: "m1", TaskSessionID: "s1", @@ -112,7 +126,7 @@ func TestCanceller_PublishesMessageUpdatedEvent(t *testing.T) { msgs := map[string][]*taskmodels.Message{} c, _, eventBus := newTestCanceller(t, msgs) - pendingID := c.store.CreateRequest(&Request{SessionID: "s1"}) + pendingID, _ := c.store.CreateRequest(&Request{SessionID: "s1"}) msgs[pendingID] = []*taskmodels.Message{{ ID: "m1", TaskSessionID: "s1", @@ -133,7 +147,7 @@ func TestCanceller_MultiQuestion_MarksAllMessagesExpired(t *testing.T) { msgs := map[string][]*taskmodels.Message{} c, _, eventBus := newTestCanceller(t, msgs) - pendingID := c.store.CreateRequest(&Request{SessionID: "s1"}) + pendingID, _ := c.store.CreateRequest(&Request{SessionID: "s1"}) msgs[pendingID] = []*taskmodels.Message{ {ID: "m1", TaskSessionID: "s1", Metadata: map[string]any{"status": "pending", "question_id": "q1"}}, {ID: "m2", TaskSessionID: "s1", Metadata: map[string]any{"status": "pending", "question_id": "q2"}}, @@ -148,3 +162,55 @@ func TestCanceller_MultiQuestion_MarksAllMessagesExpired(t *testing.T) { t.Fatalf("expected 3 message.updated events, got %d", len(eventBus.events)) } } + +// TestCanceller_MarksExpiredWhenStoreAlreadyDrained verifies that even when +// the in-memory store entry has already been removed (racing MCP-timeout path), +// CancelSessionAndNotify still finds and marks the DB messages expired. +func TestCanceller_MarksExpiredWhenStoreAlreadyDrained(t *testing.T) { + msgs := map[string][]*taskmodels.Message{} + c, repo, _ := newTestCanceller(t, msgs) + + pendingID := "orphan-pending-id" + msgs[pendingID] = []*taskmodels.Message{{ + ID: "m1", + TaskSessionID: "s1", + Metadata: map[string]any{ + "status": "pending", + "pending_id": pendingID, + }, + }} + + cancelled := c.CancelSessionAndNotify(context.Background(), "s1") + if cancelled != 1 { + t.Fatalf("expected 1 cancelled bundle from DB fallback, got %d", cancelled) + } + if len(repo.updated) != 1 { + t.Fatalf("expected 1 message update, got %d", len(repo.updated)) + } + updated := repo.updated[0] + if got, _ := updated.Metadata["status"].(string); got != "expired" { + t.Errorf("expected status=expired, got %q", got) + } +} + +// TestCanceller_Idempotent_SkipsAnsweredMessages verifies that an already-answered +// message is not clobbered back to expired. +func TestCanceller_Idempotent_SkipsAnsweredMessages(t *testing.T) { + msgs := map[string][]*taskmodels.Message{} + c, repo, _ := newTestCanceller(t, msgs) + + pendingID, _ := c.store.CreateRequest(&Request{SessionID: "s1"}) + msgs[pendingID] = []*taskmodels.Message{ + {ID: "m1", TaskSessionID: "s1", Metadata: map[string]any{"status": "answered"}}, + {ID: "m2", TaskSessionID: "s1", Metadata: map[string]any{"status": "pending"}}, + } + + c.CancelSessionAndNotify(context.Background(), "s1") + + if len(repo.updated) != 1 { + t.Fatalf("expected 1 message update (only the pending one), got %d", len(repo.updated)) + } + if repo.updated[0].ID != "m2" { + t.Errorf("expected m2 to be updated, got %s", repo.updated[0].ID) + } +} diff --git a/apps/backend/internal/clarification/handlers.go b/apps/backend/internal/clarification/handlers.go index 0719ece1c9..3b6cefc1fa 100644 --- a/apps/backend/internal/clarification/handlers.go +++ b/apps/backend/internal/clarification/handlers.go @@ -31,6 +31,7 @@ type messageStore interface { GetTaskSession(ctx context.Context, id string) (*taskmodels.TaskSession, error) FindMessageByPendingID(ctx context.Context, pendingID string) (*taskmodels.Message, error) FindMessagesByPendingID(ctx context.Context, pendingID string) ([]*taskmodels.Message, error) + FindPendingClarificationMessagesBySessionID(ctx context.Context, sessionID string) ([]*taskmodels.Message, error) UpdateMessage(ctx context.Context, message *taskmodels.Message) error } @@ -139,14 +140,15 @@ func (h *Handlers) httpCreateRequest(c *gin.Context) { Context: body.Context, } - pendingID := h.store.CreateRequest(req) + pendingID, isNew := h.store.CreateRequest(req) // Create one message per question in the database; all share the same // pending_id and are rendered as a stacked group on the frontend. The // session.message.added WebSocket event fires per message. On failure we // also cancel the in-store pending entry so any blocking WaitForResponse // caller unblocks immediately rather than waiting for the MCP timeout. - if h.messageCreator != nil { + // When dedup fires (isNew=false) the messages already exist, so skip creation. + if isNew && h.messageCreator != nil { _, err := h.messageCreator.CreateClarificationRequestMessages( c.Request.Context(), taskID, diff --git a/apps/backend/internal/clarification/handlers_test.go b/apps/backend/internal/clarification/handlers_test.go index 7dee9fb5d9..21ce93bd35 100644 --- a/apps/backend/internal/clarification/handlers_test.go +++ b/apps/backend/internal/clarification/handlers_test.go @@ -155,7 +155,7 @@ func TestHttpRespond_AnsweredAfterTimeout_PublishesEvent(t *testing.T) { // receive a phantom answer for the question that was actually skipped. func TestHttpRespond_DuplicateQuestionID_Rejected400(t *testing.T) { h, _, _, _ := setupTestHandler(t, map[string][]*taskmodels.Message{}) - pendingID := h.store.CreateRequest(&Request{ + pendingID, _ := h.store.CreateRequest(&Request{ SessionID: "s1", TaskID: "t1", Questions: []Question{ @@ -179,7 +179,7 @@ func TestHttpRespond_DuplicateQuestionID_Rejected400(t *testing.T) { // are rejected even with the right cardinality. func TestHttpRespond_UnknownQuestionID_Rejected400(t *testing.T) { h, _, _, _ := setupTestHandler(t, map[string][]*taskmodels.Message{}) - pendingID := h.store.CreateRequest(&Request{ + pendingID, _ := h.store.CreateRequest(&Request{ SessionID: "s1", TaskID: "t1", Questions: []Question{ @@ -201,7 +201,7 @@ func TestHttpRespond_UnknownQuestionID_Rejected400(t *testing.T) { func TestHttpRespond_PartialAnswers_Rejected400(t *testing.T) { h, _, _, _ := setupTestHandler(t, map[string][]*taskmodels.Message{}) - pendingID := h.store.CreateRequest(&Request{ + pendingID, _ := h.store.CreateRequest(&Request{ SessionID: "s1", TaskID: "t1", Questions: []Question{ @@ -230,7 +230,7 @@ func TestHttpRespond_PartialAnswers_Rejected400(t *testing.T) { func TestHttpRespond_AllAnswers_PrimaryPath_Success(t *testing.T) { h, _, _, msgCreator := setupTestHandler(t, map[string][]*taskmodels.Message{}) - pendingID := h.store.CreateRequest(&Request{ + pendingID, _ := h.store.CreateRequest(&Request{ SessionID: "s1", TaskID: "t1", Questions: []Question{ diff --git a/apps/backend/internal/clarification/store.go b/apps/backend/internal/clarification/store.go index ccf5a47452..bed74d302b 100644 --- a/apps/backend/internal/clarification/store.go +++ b/apps/backend/internal/clarification/store.go @@ -36,25 +36,40 @@ func NewStore(timeout time.Duration) *Store { } } -// CreateRequest creates a new clarification request and returns its pending ID. -// The request will be stored until a response is received or it times out. -func (s *Store) CreateRequest(req *Request) string { +// CreateRequest creates a new clarification request and returns its pending ID +// plus a boolean indicating whether a new entry was created (true) or an +// existing one was reused (false). If a pending entry for the same session +// with identical normalised questions already exists, the existing pending ID +// is returned and isNew is false. +func (s *Store) CreateRequest(req *Request) (string, bool) { s.mu.Lock() defer s.mu.Unlock() + // Normalise in-place so dedup keys are stable even when the caller + // hasn't assigned IDs yet. + _ = NormalizeAndValidateQuestions(req.Questions) + + // Deduplicate: if a pending entry for the same session with identical + // normalised questions already exists, return the existing pending ID. + for _, existing := range s.pending { + if existing.Request.SessionID == req.SessionID && questionsEqual(existing.Request.Questions, req.Questions) { + return existing.Request.PendingID, false + } + } + if req.PendingID == "" { req.PendingID = uuid.New().String() } req.CreatedAt = time.Now() s.pending[req.PendingID] = &PendingClarification{ - Request: req, - ResponseCh: make(chan *Response, 1), - CancelCh: make(chan struct{}), - CreatedAt: time.Now(), + Request: req, + done: make(chan struct{}), + CancelCh: make(chan struct{}), + CreatedAt: time.Now(), } - return req.PendingID + return req.PendingID, true } // GetRequest returns a pending clarification request by ID. @@ -85,12 +100,12 @@ func (s *Store) WaitForResponse(ctx context.Context, pendingID string) (*Respons defer cancel() select { - case resp := <-pending.ResponseCh: + case <-pending.done: // Clean up after receiving response s.mu.Lock() delete(s.pending, pendingID) s.mu.Unlock() - return resp, nil + return pending.resp, nil case <-pending.CancelCh: // Agent's turn completed — cancel the blocking wait s.mu.Lock() @@ -98,13 +113,15 @@ func (s *Store) WaitForResponse(ctx context.Context, pendingID string) (*Respons s.mu.Unlock() return nil, fmt.Errorf("clarification cancelled (agent moved on): %s", pendingID) case <-timeoutCtx.Done(): - // Clean up on timeout + if ctx.Err() != nil { + // Parent context cancelled — do not delete the shared entry + // because another waiter may still be blocked on it. + return nil, ctx.Err() + } + // Store-level timeout — safe to clean up. s.mu.Lock() delete(s.pending, pendingID) s.mu.Unlock() - if ctx.Err() != nil { - return nil, ctx.Err() - } return nil, fmt.Errorf("clarification request timed out: %s", pendingID) } } @@ -120,16 +137,19 @@ func (s *Store) Respond(pendingID string, resp *Response) error { return fmt.Errorf("%w: %s", ErrNotFound, pendingID) } - resp.PendingID = pendingID - resp.RespondedAt = time.Now() + pending.mu.Lock() + defer pending.mu.Unlock() - // Non-blocking send (channel has buffer of 1) - select { - case pending.ResponseCh <- resp: - return nil - default: + if pending.resolved { return fmt.Errorf("%w: %s", ErrAlreadyResponded, pendingID) } + + resp.PendingID = pendingID + resp.RespondedAt = time.Now() + pending.resp = resp + pending.resolved = true + close(pending.done) + return nil } // CancelRequest cancels a single pending clarification by id, unblocking any @@ -179,3 +199,30 @@ func (s *Store) CancelSession(sessionID string) []string { } return cancelled } + +func questionsEqual(a, b []Question) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].Prompt != b[i].Prompt { + return false + } + if !optionsEqual(a[i].Options, b[i].Options) { + return false + } + } + return true +} + +func optionsEqual(a, b []Option) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i].ID != b[i].ID || a[i].Label != b[i].Label || a[i].Description != b[i].Description { + return false + } + } + return true +} diff --git a/apps/backend/internal/clarification/store_test.go b/apps/backend/internal/clarification/store_test.go index 5305982fbf..8728853cf0 100644 --- a/apps/backend/internal/clarification/store_test.go +++ b/apps/backend/internal/clarification/store_test.go @@ -26,7 +26,7 @@ func TestCreateRequest_GeneratesID(t *testing.T) { s := NewStore(time.Minute) req := &Request{SessionID: "s1", Questions: []Question{{Prompt: "test?"}}} - id := s.CreateRequest(req) + id, _ := s.CreateRequest(req) if id == "" { t.Fatal("expected non-empty pending ID") @@ -40,7 +40,7 @@ func TestCreateRequest_PreservesExistingID(t *testing.T) { s := NewStore(time.Minute) req := &Request{PendingID: "custom-id", SessionID: "s1"} - id := s.CreateRequest(req) + id, _ := s.CreateRequest(req) if id != "custom-id" { t.Errorf("expected preserved ID %q, got %q", "custom-id", id) @@ -49,7 +49,7 @@ func TestCreateRequest_PreservesExistingID(t *testing.T) { func TestGetRequest_Found(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "test?"}}}) + id, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "test?"}}}) req, ok := s.GetRequest(id) @@ -73,7 +73,7 @@ func TestGetRequest_NotFound(t *testing.T) { func TestWaitForResponse_Success(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) // Respond in a goroutine go func() { @@ -107,7 +107,7 @@ func TestWaitForResponse_NotFound(t *testing.T) { func TestWaitForResponse_ContextCancelled(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) ctx, cancel := context.WithCancel(context.Background()) go func() { @@ -123,7 +123,7 @@ func TestWaitForResponse_ContextCancelled(t *testing.T) { func TestWaitForResponse_CancelCh(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) // Cancel via CancelSession in a goroutine go func() { @@ -142,7 +142,7 @@ func TestWaitForResponse_CancelCh(t *testing.T) { func TestWaitForResponse_StoreTimeout(t *testing.T) { s := NewStore(50 * time.Millisecond) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) _, err := s.WaitForResponse(context.Background(), id) if err == nil { @@ -155,7 +155,7 @@ func TestWaitForResponse_StoreTimeout(t *testing.T) { func TestRespond_Success(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) err := s.Respond(id, &Response{Answers: []Answer{{CustomText: "yes"}}}) if err != nil { @@ -174,7 +174,7 @@ func TestRespond_NotFound(t *testing.T) { func TestRespond_Duplicate(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) // First respond succeeds if err := s.Respond(id, &Response{}); err != nil { @@ -188,9 +188,9 @@ func TestRespond_Duplicate(t *testing.T) { func TestCancelSession_CancelsMatchingRequests(t *testing.T) { s := NewStore(time.Minute) - id1 := s.CreateRequest(&Request{SessionID: "s1"}) - id2 := s.CreateRequest(&Request{SessionID: "s1"}) - id3 := s.CreateRequest(&Request{SessionID: "s2"}) + id1, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "q1?", Options: []Option{{ID: "o1", Label: "A"}}}}}) + id2, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "q2?", Options: []Option{{ID: "o1", Label: "B"}}}}}) + id3, _ := s.CreateRequest(&Request{SessionID: "s2", Questions: []Question{{Prompt: "q3?", Options: []Option{{ID: "o1", Label: "C"}}}}}) cancelled := s.CancelSession("s1") @@ -222,7 +222,7 @@ func TestCancelSession_CancelsMatchingRequests(t *testing.T) { // paths return an error from WaitForResponse — that's the contract under test. func TestCancelRequest(t *testing.T) { s := NewStore(time.Minute) - id := s.CreateRequest(&Request{SessionID: "s1"}) + id, _ := s.CreateRequest(&Request{SessionID: "s1"}) started := make(chan struct{}) done := make(chan error, 1) @@ -251,7 +251,7 @@ func TestCancelRequest(t *testing.T) { func TestCancelSession_NoMatch(t *testing.T) { s := NewStore(time.Minute) - s.CreateRequest(&Request{SessionID: "s1"}) + _, _ = s.CreateRequest(&Request{SessionID: "s1"}) cancelled := s.CancelSession("other") @@ -274,13 +274,13 @@ func TestListPendingPermissions_Empty(t *testing.T) { func TestListPendingPermissions_ReturnsPendingRequests(t *testing.T) { s := NewStore(time.Minute) - s.CreateRequest(&Request{ + _, _ = s.CreateRequest(&Request{ SessionID: "session-1", TaskID: "task-1", Questions: []Question{{Prompt: "Allow bash execution?"}}, Context: "tool permission", }) - s.CreateRequest(&Request{ + _, _ = s.CreateRequest(&Request{ SessionID: "session-2", TaskID: "task-2", Questions: []Question{{Prompt: "Write to /tmp?"}}, @@ -321,8 +321,8 @@ func TestListPendingPermissions_ReturnsPendingRequests(t *testing.T) { func TestListPendingPermissions_ExcludesCancelled(t *testing.T) { s := NewStore(time.Minute) - s.CreateRequest(&Request{SessionID: "s1"}) - s.CreateRequest(&Request{SessionID: "s2"}) + _, _ = s.CreateRequest(&Request{SessionID: "s1"}) + _, _ = s.CreateRequest(&Request{SessionID: "s2"}) // Cancel session s1 s.CancelSession("s1") @@ -344,3 +344,82 @@ func TestListPendingPermissions_ImplementsInterface(t *testing.T) { ListPendingPermissions() []shared.PendingPermission } = s } + +func TestCreateRequest_Dedup_SameSessionAndQuestions(t *testing.T) { + s := NewStore(time.Minute) + q := []Question{{Prompt: "What colour?", Options: []Option{{ID: "o1", Label: "Red"}, {ID: "o2", Label: "Blue"}}}} + + id1, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: q}) + id2, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: q}) + + if id1 != id2 { + t.Fatalf("expected duplicate request to reuse pending ID %q, got %q", id1, id2) + } + + pending := s.ListPending() + if len(pending) != 1 { + t.Fatalf("expected 1 pending request after dedup, got %d", len(pending)) + } +} + +func TestCreateRequest_NoDedup_DifferentQuestions(t *testing.T) { + s := NewStore(time.Minute) + id1, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "Q1?", Options: []Option{{ID: "o1", Label: "A"}}}}}) + id2, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "Q2?", Options: []Option{{ID: "o1", Label: "A"}}}}}) + + if id1 == id2 { + t.Fatal("expected different pending IDs for different questions") + } +} + +// testStore wraps Store so WaitForResponse can signal that it has passed the +// initial lookup and is about to block. This eliminates the race where +// Respond fires before a late goroutine even enters WaitForResponse. +type testStore struct { + *Store + entered chan string +} + +func (ts *testStore) WaitForResponse(ctx context.Context, id string) (*Response, error) { + ts.entered <- id + return ts.Store.WaitForResponse(ctx, id) +} + +func TestWaitForResponse_Broadcast_MultipleWaiters(t *testing.T) { + s := NewStore(time.Minute) + ts := &testStore{Store: s, entered: make(chan string, 2)} + id, _ := s.CreateRequest(&Request{SessionID: "s1", Questions: []Question{{Prompt: "test?", Options: []Option{{ID: "o1", Label: "A"}}}}}) + + done := make(chan *Response, 2) + for i := 0; i < 2; i++ { + go func() { + resp, err := ts.WaitForResponse(context.Background(), id) + if err != nil { + done <- nil + return + } + done <- resp + }() + } + <-ts.entered + <-ts.entered + + if err := s.Respond(id, &Response{Answers: []Answer{{CustomText: "hello"}}}); err != nil { + t.Fatalf("unexpected respond error: %v", err) + } + + var got int + for i := 0; i < 2; i++ { + select { + case resp := <-done: + if resp != nil && len(resp.Answers) == 1 && resp.Answers[0].CustomText == "hello" { + got++ + } + case <-time.After(time.Second): + t.Fatal("WaitForResponse did not return") + } + } + if got != 2 { + t.Fatalf("expected both waiters to receive response, got %d", got) + } +} diff --git a/apps/backend/internal/clarification/types.go b/apps/backend/internal/clarification/types.go index 436d2e2eb3..9d7c52f826 100644 --- a/apps/backend/internal/clarification/types.go +++ b/apps/backend/internal/clarification/types.go @@ -3,6 +3,7 @@ package clarification import ( + "sync" "time" ) @@ -53,10 +54,13 @@ type Response struct { // PendingClarification represents a clarification request waiting for a response. type PendingClarification struct { - Request *Request - ResponseCh chan *Response - CancelCh chan struct{} // Closed when session's turn completes (agent moved on) - CreatedAt time.Time + Request *Request + done chan struct{} // Closed when a response is submitted (broadcast to all waiters) + resp *Response + mu sync.Mutex + resolved bool + CancelCh chan struct{} // Closed when session's turn completes (agent moved on) + CreatedAt time.Time } // Status represents the status of a clarification request. diff --git a/apps/backend/internal/integration/mcp_integration_test.go b/apps/backend/internal/integration/mcp_integration_test.go index fa5578e257..4bc00d682d 100644 --- a/apps/backend/internal/integration/mcp_integration_test.go +++ b/apps/backend/internal/integration/mcp_integration_test.go @@ -18,7 +18,7 @@ func setupMCPTestServer(t *testing.T) (*TestServer, string, string, string, stri ts := NewTestServer(t) // Register MCP handlers on the gateway dispatcher - mcpH := mcphandlers.NewHandlers(ts.TaskSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, ts.Logger) + mcpH := mcphandlers.NewHandlers(ts.TaskSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, ts.Logger) mcpH.RegisterHandlers(ts.Gateway.Dispatcher) client := NewWSClient(t, ts.Server.URL) diff --git a/apps/backend/internal/mcp/handlers/handlers.go b/apps/backend/internal/mcp/handlers/handlers.go index 5f8d9144e8..2084c6c033 100644 --- a/apps/backend/internal/mcp/handlers/handlers.go +++ b/apps/backend/internal/mcp/handlers/handlers.go @@ -31,10 +31,15 @@ import ( // ClarificationService defines the interface for clarification operations. type ClarificationService interface { - CreateRequest(req *clarification.Request) string + CreateRequest(req *clarification.Request) (string, bool) WaitForResponse(ctx context.Context, pendingID string) (*clarification.Response, error) CancelRequest(pendingID string) bool - CancelSession(sessionID string) []string +} + +// SessionCanceller cancels all pending clarifications for a session and marks +// the related chat messages as expired. Used by the MCP-timeout handler. +type SessionCanceller interface { + CancelSessionAndNotify(ctx context.Context, sessionID string) int } // MessageCreator creates messages for clarification requests. @@ -83,6 +88,7 @@ type Handlers struct { taskSvc *service.Service workflowCtrl *workflowctrl.Controller clarificationSvc ClarificationService + sessionCanceller SessionCanceller messageCreator MessageCreator sessionRepo SessionRepository taskRepo TaskRepository @@ -108,6 +114,7 @@ func NewHandlers( taskSvc *service.Service, workflowCtrl *workflowctrl.Controller, clarificationSvc ClarificationService, + sessionCanceller SessionCanceller, messageCreator MessageCreator, sessionRepo SessionRepository, taskRepo TaskRepository, @@ -121,6 +128,7 @@ func NewHandlers( taskSvc: taskSvc, workflowCtrl: workflowCtrl, clarificationSvc: clarificationSvc, + sessionCanceller: sessionCanceller, messageCreator: messageCreator, sessionRepo: sessionRepo, taskRepo: taskRepo, @@ -1114,13 +1122,14 @@ func (h *Handlers) handleAskUserQuestion(ctx context.Context, msg *ws.Message) ( Questions: req.Questions, Context: req.Context, } - pendingID := h.clarificationSvc.CreateRequest(clarificationReq) + pendingID, isNew := h.clarificationSvc.CreateRequest(clarificationReq) // Create one chat message per question (triggers WS events to frontend). // If the create fails, the in-store pending entry must be cancelled too — // otherwise the agent's WaitForResponse would block for the full 2-hour // timeout while the user never sees clarification cards. - if h.messageCreator != nil { + // When dedup fires (isNew=false) the messages already exist, so skip creation. + if isNew && h.messageCreator != nil { if _, err := h.messageCreator.CreateClarificationRequestMessages( ctx, taskID, req.SessionID, pendingID, req.Questions, req.Context, ); err != nil { @@ -1357,7 +1366,7 @@ func (h *Handlers) handleDeleteTaskPlan(ctx context.Context, msg *ws.Message) (* // disconnects while waiting for a clarification response. It cancels the pending // clarification so the user's eventual answer goes through the event fallback path // (new turn) instead of the primary path (which would be dropped). -func (h *Handlers) handleClarificationTimeout(_ context.Context, msg *ws.Message) (*ws.Message, error) { +func (h *Handlers) handleClarificationTimeout(ctx context.Context, msg *ws.Message) (*ws.Message, error) { var req struct { SessionID string `json:"session_id"` } @@ -1368,10 +1377,13 @@ func (h *Handlers) handleClarificationTimeout(_ context.Context, msg *ws.Message return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeValidation, "session_id is required", nil) } - cancelled := h.clarificationSvc.CancelSession(req.SessionID) + cancelled := 0 + if h.sessionCanceller != nil { + cancelled = h.sessionCanceller.CancelSessionAndNotify(ctx, req.SessionID) + } h.logger.Info("cancelled pending clarifications on agent MCP timeout", zap.String("session_id", req.SessionID), - zap.Int("count", len(cancelled))) + zap.Int("count", cancelled)) - return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{"ok": true, "cancelled": len(cancelled)}) + return ws.NewResponse(msg.ID, msg.Action, map[string]interface{}{"ok": true, "cancelled": cancelled}) } diff --git a/apps/backend/internal/mcp/handlers/handlers_test.go b/apps/backend/internal/mcp/handlers/handlers_test.go index ae06f79e2c..a3f505b17b 100644 --- a/apps/backend/internal/mcp/handlers/handlers_test.go +++ b/apps/backend/internal/mcp/handlers/handlers_test.go @@ -9,6 +9,7 @@ import ( "time" "github.com/jmoiron/sqlx" + "github.com/kandev/kandev/internal/clarification" "github.com/kandev/kandev/internal/common/logger" "github.com/kandev/kandev/internal/db" "github.com/kandev/kandev/internal/events/bus" @@ -706,3 +707,135 @@ func TestHandleCreateTask_BlockedBy_Accepted(t *testing.T) { // Fails on workspace_id, not on blocked_by parsing assertWSError(t, resp, ws.ErrorCodeValidation) } + +func TestHandleClarificationTimeout_MarksMessagesExpired(t *testing.T) { + svc, repo := newTestTaskService(t) + ctx := context.Background() + + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ID: "ws-1", Name: "Test"})) + require.NoError(t, repo.CreateWorkflow(ctx, &models.Workflow{ID: "wf-1", WorkspaceID: "ws-1", Name: "Board"})) + task, err := svc.CreateTask(ctx, &service.CreateTaskRequest{ + WorkspaceID: "ws-1", + WorkflowID: "wf-1", + Title: "Task", + }) + require.NoError(t, err) + + sess := &models.TaskSession{ + ID: "sess-1", + TaskID: task.ID, + IsPrimary: true, + State: models.TaskSessionStateRunning, + } + require.NoError(t, repo.CreateTaskSession(ctx, sess)) + turn := &models.Turn{ID: "turn-1", TaskSessionID: sess.ID, TaskID: task.ID} + require.NoError(t, repo.CreateTurn(ctx, turn)) + + store := clarification.NewStore(time.Minute) + eventBus := bus.NewMemoryEventBus(testLogger(t)) + t.Cleanup(func() { eventBus.Close() }) + canceller := clarification.NewCanceller(store, repo, eventBus, testLogger(t)) + + pendingID := "test-pending-id" + require.NoError(t, repo.CreateMessage(ctx, &models.Message{ + TaskSessionID: sess.ID, + TaskID: task.ID, + TurnID: turn.ID, + AuthorType: "agent", + Type: "clarification_request", + Content: "Q?", + Metadata: map[string]interface{}{ + "pending_id": pendingID, + "question_id": "q1", + "status": "pending", + }, + })) + + // Drain the in-memory store so the handler must fall back to DB-driven cleanup + store.CancelSession(sess.ID) + + h := NewHandlers(svc, nil, store, canceller, nil, repo, repo, eventBus, nil, nil, nil, testLogger(t)) + msg := makeWSMessage(t, ws.ActionMCPClarificationTimeout, map[string]interface{}{"session_id": sess.ID}) + resp, err := h.handleClarificationTimeout(ctx, msg) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, resp.Type) + + msgs, err := repo.FindPendingClarificationMessagesBySessionID(ctx, sess.ID) + require.NoError(t, err) + require.Empty(t, msgs, "all clarification messages should be expired") +} + +func TestHandleClarificationTimeout_WithoutCanceller_ReturnsZero(t *testing.T) { + h := &Handlers{sessionCanceller: nil, logger: testLogger(t).WithFields()} + msg := makeWSMessage(t, ws.ActionMCPClarificationTimeout, map[string]interface{}{"session_id": "s1"}) + resp, err := h.handleClarificationTimeout(context.Background(), msg) + require.NoError(t, err) + require.Equal(t, ws.MessageTypeResponse, resp.Type) +} + +func TestHandleAskUserQuestion_Dedup_CreatesOnePendingBundle(t *testing.T) { + svc, repo := newTestTaskService(t) + ctx := context.Background() + + require.NoError(t, repo.CreateWorkspace(ctx, &models.Workspace{ID: "ws-1", Name: "Test"})) + require.NoError(t, repo.CreateWorkflow(ctx, &models.Workflow{ID: "wf-1", WorkspaceID: "ws-1", Name: "Board"})) + task, err := svc.CreateTask(ctx, &service.CreateTaskRequest{ + WorkspaceID: "ws-1", + WorkflowID: "wf-1", + Title: "Task", + }) + require.NoError(t, err) + + sess := &models.TaskSession{ + ID: "sess-dedup", + TaskID: task.ID, + IsPrimary: true, + State: models.TaskSessionStateRunning, + } + require.NoError(t, repo.CreateTaskSession(ctx, sess)) + + store := clarification.NewStore(time.Minute) + log := testLogger(t) + h := NewHandlers(svc, nil, store, nil, nil, repo, repo, nil, nil, nil, nil, log) + + payload := map[string]interface{}{ + "session_id": sess.ID, + "task_id": task.ID, + "questions": []map[string]interface{}{ + {"prompt": "What colour?", "options": []map[string]interface{}{ + {"label": "Red", "description": "R"}, + {"label": "Blue", "description": "B"}, + }}, + }, + } + + // Fire two identical requests concurrently. Because handleAskUserQuestion + // blocks, we run them in goroutines and cancel the session after a short + // delay so both return. + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + msg := makeWSMessage(t, ws.ActionMCPAskUserQuestion, payload) + if _, err := h.handleAskUserQuestion(ctx, msg); err != nil { + t.Errorf("handleAskUserQuestion returned unexpected error: %v", err) + } + }() + } + + // Wait until the single deduped pending bundle is visible in the store + // (confirming both goroutines have passed the CreateRequest gate). + require.Eventually(t, func() bool { + return len(store.ListPending()) == 1 + }, time.Second, 5*time.Millisecond) + store.CancelSession(sess.ID) + wg.Wait() + + // After dedup there should be exactly one pending entry even though two + // handlers were invoked. + pending := store.ListPending() + require.Len(t, pending, 0) // cancelled + // The actual dedup invariant is verified at the Store level; this test + // ensures the handler path does not panic or leak with concurrent calls. +} diff --git a/apps/backend/internal/task/handlers/process_handlers_test.go b/apps/backend/internal/task/handlers/process_handlers_test.go index 5f4a884ab0..4d8822cd3e 100644 --- a/apps/backend/internal/task/handlers/process_handlers_test.go +++ b/apps/backend/internal/task/handlers/process_handlers_test.go @@ -181,6 +181,9 @@ func (m *mockRepository) FindMessagesByPendingID(ctx context.Context, pendingID func (m *mockRepository) FindMessageByPendingIDAndQuestion(ctx context.Context, sessionID, pendingID, questionID string) (*models.Message, error) { return nil, nil } +func (m *mockRepository) FindPendingClarificationMessagesBySessionID(ctx context.Context, sessionID string) ([]*models.Message, error) { + return nil, nil +} func (m *mockRepository) UpdateMessage(ctx context.Context, message *models.Message) error { return nil } diff --git a/apps/backend/internal/task/repository/interface.go b/apps/backend/internal/task/repository/interface.go index 003932993a..9306a65fab 100644 --- a/apps/backend/internal/task/repository/interface.go +++ b/apps/backend/internal/task/repository/interface.go @@ -102,6 +102,7 @@ type MessageRepository interface { FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error) FindMessagesByPendingID(ctx context.Context, pendingID string) ([]*models.Message, error) FindMessageByPendingIDAndQuestion(ctx context.Context, sessionID, pendingID, questionID string) (*models.Message, error) + FindPendingClarificationMessagesBySessionID(ctx context.Context, sessionID string) ([]*models.Message, error) UpdateMessage(ctx context.Context, message *models.Message) error ListMessages(ctx context.Context, sessionID string) ([]*models.Message, error) ListMessagesByTurnID(ctx context.Context, turnID string) ([]*models.Message, error) diff --git a/apps/backend/internal/task/repository/sqlite/message.go b/apps/backend/internal/task/repository/sqlite/message.go index 2c8c7972a0..70fffcfde6 100644 --- a/apps/backend/internal/task/repository/sqlite/message.go +++ b/apps/backend/internal/task/repository/sqlite/message.go @@ -353,6 +353,27 @@ func (r *Repository) FindMessagesByPendingID(ctx context.Context, pendingID stri return result, err } +// FindPendingClarificationMessagesBySessionID returns every clarification_request +// message for the session whose metadata.status is still "pending". Used by the +// canceller as a fallback when the in-memory store entry has already been drained +// by a racing timeout path. +func (r *Repository) FindPendingClarificationMessagesBySessionID(ctx context.Context, sessionID string) ([]*models.Message, error) { + drv := r.ro.DriverName() + query := fmt.Sprintf(` + SELECT id, task_session_id, task_id, turn_id, author_type, author_id, content, requests_input, type, metadata, created_at + FROM task_session_messages + WHERE task_session_id = ? AND type = 'clarification_request' AND %s = 'pending' + ORDER BY created_at ASC + `, dialect.JSONExtract(drv, "metadata", "status")) + rows, err := r.ro.QueryContext(ctx, r.ro.Rebind(query), sessionID) + if err != nil { + return nil, err + } + defer func() { _ = rows.Close() }() + result, _, err := scanMessageRows(rows, 0) + return result, err +} + // FindMessageByPendingIDAndQuestion finds the message for a specific (pending_id, // question_id) pair within a session. Used to flip per-question status (answered / // rejected) on multi-question clarification bundles. The persistence layer stores