fix: resolve session stuck issues and workflow transition bugs (#206)

This commit is contained in:
João Salavisa
2026-02-12 23:34:16 +00:00
committed by GitHub
parent ff96b4ed2d
commit 11ac58d3e0
11 changed files with 172 additions and 46 deletions
@@ -4,17 +4,11 @@ import (
"context"
"encoding/json"
"fmt"
"time"
"github.com/mark3labs/mcp-go/mcp"
"github.com/mark3labs/mcp-go/server"
)
// askUserQuestionTimeout is the maximum time to wait for a user response.
// This needs to be long enough for users to read and respond to questions.
// Note: This must match or exceed the clarification store timeout (10 minutes).
const askUserQuestionTimeout = 12 * time.Minute
// MCP action constants for backend WS requests
const (
ActionMCPListWorkspaces = "mcp.list_workspaces"
@@ -151,7 +145,7 @@ func (s *Server) updateTaskHandler() server.ToolHandlerFunc {
}
func (s *Server) askUserQuestionHandler() server.ToolHandlerFunc {
return func(_ context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
return func(ctx context.Context, req mcp.CallToolRequest) (*mcp.CallToolResult, error) {
prompt, err := req.RequireString("prompt")
if err != nil {
return mcp.NewToolResultError("prompt is required"), nil
@@ -212,15 +206,12 @@ func (s *Server) askUserQuestionHandler() server.ToolHandlerFunc {
"context": questionCtx,
}
// Use a detached context with our own timeout for waiting on user input.
// The MCP request context may have a short timeout from the agent's MCP client,
// but we need to wait longer for users to read and respond to questions.
// This is similar to how script_message_handler.go handles long-running scripts.
askCtx, cancel := context.WithTimeout(context.Background(), askUserQuestionTimeout)
defer cancel()
// Use the MCP request context from the agent. This ensures that if the agent's
// MCP client times out, we'll detect it and not update the session state.
// Previous behavior used a detached context which meant we couldn't tell if the
// agent had given up, leading to sessions stuck in RUNNING state.
var result map[string]interface{}
if err := s.backend.RequestPayload(askCtx, ActionMCPAskUserQuestion, payload, &result); err != nil {
if err := s.backend.RequestPayload(ctx, ActionMCPAskUserQuestion, payload, &result); err != nil {
return mcp.NewToolResultError(err.Error()), nil
}
@@ -175,8 +175,45 @@ func (h *Handlers) httpRespond(c *gin.Context) {
// Get the pending request to find the session ID
pending, ok := h.store.GetRequest(pendingID)
// If request not found in store (likely timed out), mark the message as expired
// so the UI stops showing it as pending
var sessionID string
if !ok {
c.JSON(http.StatusNotFound, gin.H{"error": "clarification request not found"})
h.logger.Warn("clarification request not found in store (likely timed out), marking message as expired",
zap.String("pending_id", pendingID))
// Find the message by pending_id and mark it as expired
message, err := h.repo.FindMessageByPendingID(c.Request.Context(), pendingID)
if err != nil {
h.logger.Warn("failed to find clarification message for expired request",
zap.String("pending_id", pendingID),
zap.Error(err))
c.JSON(http.StatusGone, gin.H{
"error": "clarification request expired",
"code": "CLARIFICATION_EXPIRED",
"message": "This question has timed out. The agent may have already moved on. Please refresh the page.",
})
return
}
// Update the message status to expired
if message.Metadata == nil {
message.Metadata = make(map[string]interface{})
}
message.Metadata["status"] = "expired"
if err := h.repo.UpdateMessage(c.Request.Context(), message); err != nil {
h.logger.Warn("failed to update expired clarification message",
zap.String("pending_id", pendingID),
zap.String("message_id", message.ID),
zap.Error(err))
}
c.JSON(http.StatusGone, gin.H{
"error": "clarification request expired",
"code": "CLARIFICATION_EXPIRED",
"message": "This question has timed out. The agent may have already moved on. The page will refresh automatically.",
})
return
}
@@ -198,16 +235,21 @@ func (h *Handlers) httpRespond(c *gin.Context) {
h.logger.Warn("failed to respond to clarification",
zap.String("pending_id", pendingID),
zap.Error(err))
c.JSON(http.StatusNotFound, gin.H{"error": err.Error()})
c.JSON(http.StatusGone, gin.H{
"error": "clarification request expired",
"code": "CLARIFICATION_EXPIRED",
"message": "This question has timed out. The agent may have already moved on.",
})
return
}
sessionID = pending.SessionID
// Update the message in the database with status and answer
status := "answered"
if body.Rejected {
status = "rejected"
}
if err := h.messageCreator.UpdateClarificationMessage(c.Request.Context(), pending.SessionID, pendingID, status, answer); err != nil {
if err := h.messageCreator.UpdateClarificationMessage(c.Request.Context(), sessionID, pendingID, status, answer); err != nil {
h.logger.Warn("failed to update clarification message",
zap.String("pending_id", pendingID),
zap.Error(err))
+24 -26
View File
@@ -298,36 +298,34 @@ func (h *Handlers) handleAskUserQuestion(ctx context.Context, msg *ws.Message) (
h.logger.Error("failed waiting for clarification response",
zap.String("pending_id", pendingID),
zap.Error(err))
// When clarification times out, keep session in WAITING_FOR_INPUT
// The expired clarification will be marked as expired when user tries to respond
// This allows the UI to stay responsive while clarification is pending
return ws.NewError(msg.ID, msg.Action, ws.ErrorCodeInternalError, "Failed to get user response: "+err.Error(), nil)
}
// Restore session to RUNNING state after user responds
if err := h.sessionRepo.UpdateTaskSessionState(ctx, req.SessionID, models.TaskSessionStateRunning, ""); err != nil {
h.logger.Warn("failed to restore session state to RUNNING",
zap.String("session_id", req.SessionID),
zap.Error(err))
}
// Don't update session state here. Let the agent's normal message flow handle it.
//
// Previously, we would set the session to RUNNING immediately after the user responded.
// However, this causes issues when:
// 1. The agent's MCP client has already timed out (agent stopped waiting)
// 2. The user responds anyway (UI still showing the clarification)
// 3. We set session to RUNNING, but agent won't process the response
// 4. Session gets stuck in RUNNING state with no agent activity
//
// By keeping the session in WAITING_FOR_INPUT, the user can:
// - Send new messages to continue the conversation
// - The agent will process those messages normally
// - Session state will be updated by the agent's actual activity
//
// If the agent successfully receives this response and continues, the agent's
// tool calls or messages will naturally update the session to RUNNING.
// Update task state to IN_PROGRESS
if err := h.taskRepo.UpdateTaskState(ctx, taskID, v1.TaskStateInProgress); err != nil {
h.logger.Warn("failed to restore task state to IN_PROGRESS",
zap.String("task_id", taskID),
zap.Error(err))
}
// Publish session state changed event
if h.eventBus != nil {
eventData := map[string]interface{}{
"task_id": taskID,
"session_id": req.SessionID,
"new_state": string(models.TaskSessionStateRunning),
}
_ = h.eventBus.Publish(ctx, events.TaskSessionStateChanged, bus.NewEvent(
events.TaskSessionStateChanged,
"mcp-handlers",
eventData,
))
}
h.logger.Debug("user responded to clarification, returning response to agent without updating session state",
zap.String("pending_id", pendingID),
zap.String("session_id", req.SessionID),
zap.String("task_id", taskID))
return ws.NewResponse(msg.ID, msg.Action, resp)
}
@@ -972,7 +972,10 @@ func (s *Service) handleToolCallEvent(ctx context.Context, payload *lifecycle.Ag
zap.String("tool_call_id", payload.Data.ToolCallID))
}
s.updateTaskSessionState(ctx, payload.TaskID, payload.SessionID, models.TaskSessionStateRunning, "", false)
// Allow tool calls to wake session from WAITING_FOR_INPUT
// This ensures that when user responds to a clarification and the agent continues,
// the session properly transitions to RUNNING
s.updateTaskSessionState(ctx, payload.TaskID, payload.SessionID, models.TaskSessionStateRunning, "", true)
}
}
@@ -1134,9 +1137,10 @@ func (s *Service) handleToolUpdateEvent(ctx context.Context, payload *lifecycle.
}
// Update session state for completion events
// Allow tool completions to wake session from WAITING_FOR_INPUT
if payload.Data.ToolStatus == "complete" || payload.Data.ToolStatus == "completed" ||
payload.Data.ToolStatus == "success" || payload.Data.ToolStatus == "error" || payload.Data.ToolStatus == "failed" {
s.updateTaskSessionState(ctx, payload.TaskID, payload.SessionID, models.TaskSessionStateRunning, "", false)
s.updateTaskSessionState(ctx, payload.TaskID, payload.SessionID, models.TaskSessionStateRunning, "", true)
}
}
}
@@ -243,6 +243,9 @@ func (m *mockRepository) GetMessageByToolCallID(ctx context.Context, sessionID,
func (m *mockRepository) GetMessageByPendingID(ctx context.Context, sessionID, pendingID string) (*models.Message, error) {
return nil, nil
}
func (m *mockRepository) FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error) {
return nil, nil
}
func (m *mockRepository) UpdateMessage(ctx context.Context, message *models.Message) error {
return nil
}
@@ -123,6 +123,9 @@ func (m *mockRepository) GetMessageByToolCallID(ctx context.Context, sessionID,
func (m *mockRepository) GetMessageByPendingID(ctx context.Context, sessionID, pendingID string) (*models.Message, error) {
return nil, nil
}
func (m *mockRepository) FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error) {
return nil, nil
}
func (m *mockRepository) UpdateMessage(ctx context.Context, message *models.Message) error {
return nil
}
@@ -49,6 +49,7 @@ type Repository interface {
GetMessage(ctx context.Context, id string) (*models.Message, error)
GetMessageByToolCallID(ctx context.Context, sessionID, toolCallID string) (*models.Message, error)
GetMessageByPendingID(ctx context.Context, sessionID, pendingID string) (*models.Message, error)
FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error)
UpdateMessage(ctx context.Context, message *models.Message) error
ListMessages(ctx context.Context, sessionID string) ([]*models.Message, error)
ListMessagesPaginated(ctx context.Context, sessionID string, opts models.ListMessagesOptions) ([]*models.Message, bool, error)
@@ -267,6 +267,30 @@ func (r *Repository) GetMessageByPendingID(ctx context.Context, sessionID, pendi
return message, nil
}
// FindMessageByPendingID finds a message by pending_id alone (without requiring session ID).
// This is useful when we only have the pending ID (e.g., from expired clarification responses).
func (r *Repository) FindMessageByPendingID(ctx context.Context, pendingID string) (*models.Message, error) {
message := &models.Message{}
var requestsInput int
var messageType string
var metadataJSON string
err := r.db.QueryRowContext(ctx, `
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 json_extract(metadata, '$.pending_id') = ?
ORDER BY created_at DESC LIMIT 1
`, pendingID).Scan(&message.ID, &message.TaskSessionID, &message.TaskID, &message.TurnID, &message.AuthorType, &message.AuthorID,
&message.Content, &requestsInput, &messageType, &metadataJSON, &message.CreatedAt)
if err != nil {
return nil, err
}
message.RequestsInput = requestsInput == 1
message.Type = models.MessageType(messageType)
if metadataJSON != "" {
_ = json.Unmarshal([]byte(metadataJSON), &message.Metadata)
}
return message, nil
}
// CompleteRunningToolCallsForTurn marks all "running" tool call messages for a turn as "complete".
// This is a safety net to ensure no tool calls remain stuck in "running" state after a turn completes.
func (r *Repository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) {
@@ -991,6 +991,30 @@ func (s *Service) MoveTask(ctx context.Context, id string, boardID string, workf
return nil, err
}
// Update active session's workflow_step_id to match the new task workflow step
// This ensures workflow transitions work correctly when tasks are manually moved
if task.WorkflowStepID != "" {
activeSession, err := s.repo.GetActiveTaskSessionByTaskID(ctx, id)
if err == nil && activeSession != nil {
// Only update if the session's workflow_step_id doesn't match
if activeSession.WorkflowStepID == nil || *activeSession.WorkflowStepID != task.WorkflowStepID {
if err := s.repo.UpdateSessionWorkflowStep(ctx, activeSession.ID, task.WorkflowStepID); err != nil {
s.logger.Warn("failed to update session workflow step after task move",
zap.String("task_id", id),
zap.String("session_id", activeSession.ID),
zap.String("workflow_step_id", task.WorkflowStepID),
zap.Error(err))
// Don't fail the operation, just log the warning
} else {
s.logger.Info("updated session workflow step to match moved task",
zap.String("task_id", id),
zap.String("session_id", activeSession.ID),
zap.String("workflow_step_id", task.WorkflowStepID))
}
}
}
}
// Publish state_changed event if state changed, otherwise just updated
if oldState != task.State {
s.publishTaskEvent(ctx, events.TaskStateChanged, task, &oldState)
@@ -36,6 +36,15 @@ export function ClarificationInputOverlay({ message, onResolved }: Clarification
});
if (response.ok) {
onResolved();
} else if (response.status === 410) {
// Clarification expired (agent timed out)
// Backend has marked the message as expired, so reload to clear the UI
const data = await response.json().catch(() => ({}));
console.warn('Clarification expired:', data);
// Force reload to fetch updated messages (clarification no longer pending)
window.location.reload();
} else {
console.error('Failed to submit clarification response:', response.status, response.statusText);
}
} catch (error) {
console.error('Failed to submit clarification response:', error);
@@ -62,6 +71,15 @@ export function ClarificationInputOverlay({ message, onResolved }: Clarification
});
if (response.ok) {
onResolved();
} else if (response.status === 410) {
// Clarification expired (agent timed out)
// Backend has marked the message as expired, so reload to clear the UI
const data = await response.json().catch(() => ({}));
console.warn('Clarification expired:', data);
// Force reload to fetch updated messages (clarification no longer pending)
window.location.reload();
} else {
console.error('Failed to submit clarification response:', response.status, response.statusText);
}
} catch (error) {
console.error('Failed to submit clarification response:', error);
@@ -82,6 +100,14 @@ export function ClarificationInputOverlay({ message, onResolved }: Clarification
});
if (response.ok) {
onResolved();
} else if (response.status === 410) {
// Clarification expired (agent timed out)
const data = await response.json().catch(() => ({}));
console.warn('Clarification expired:', data);
alert('This question has timed out. The agent has already moved on. Please refresh the page.');
onResolved(); // Clear the UI even though it failed
} else {
console.error('Failed to skip clarification:', response.status, response.statusText);
}
} catch (error) {
console.error('Failed to skip clarification:', error);
@@ -22,6 +22,7 @@ export function ClarificationRequestMessage({ comment }: ClarificationRequestMes
const status = metadata.status;
const isAnswered = status === 'answered';
const isSkipped = status === 'rejected';
const isExpired = status === 'expired';
const getStatusIndicator = () => {
if (isAnswered) {
@@ -30,6 +31,9 @@ export function ClarificationRequestMessage({ comment }: ClarificationRequestMes
if (isSkipped) {
return <IconX className="h-3.5 w-3.5 text-muted-foreground" />;
}
if (isExpired) {
return <IconX className="h-3.5 w-3.5 text-orange-500" />;
}
return null;
};
@@ -86,6 +90,12 @@ export function ClarificationRequestMessage({ comment }: ClarificationRequestMes
Skipped
</div>
)}
{isExpired && (
<div className="mt-1 ml-3 flex items-center gap-1.5 text-xs text-orange-500/80">
{getStatusIndicator()}
Timed out (agent moved on)
</div>
)}
</div>
</div>
</div>