fix(backend): close clarification overlay after agent MCP timeout and dedup retried questions (#1185)
This commit is contained in:
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user