fix: resolve session stuck issues and workflow transition bugs (#206)
This commit is contained in:
@@ -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))
|
||||
|
||||
@@ -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>
|
||||
|
||||
Reference in New Issue
Block a user