fix(backend): close clarification overlay after agent MCP timeout and dedup retried questions (#1185)

This commit is contained in:
Carlos Florêncio
2026-05-31 12:36:13 +01:00
committed by GitHub
parent a10e81c0d7
commit dea04b2bed
14 changed files with 490 additions and 83 deletions
+5 -3
View File
@@ -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)
@@ -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.
@@ -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)
}
}
@@ -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,
@@ -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{
+68 -21
View File
@@ -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
}
@@ -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)
}
}
+8 -4
View File
@@ -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.
@@ -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)
+20 -8
View File
@@ -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})
}
@@ -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.
}
@@ -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
}
@@ -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)
@@ -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