From 3a71ae2fe683da88897e55e6c8e1d71f8d97b597 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Carlos=20Flor=C3=AAncio?= Date: Sun, 29 Mar 2026 17:36:50 +0100 Subject: [PATCH] fix: complete all non-terminal tool calls when turn ends (#488) --- .../internal/orchestrator/event_handlers.go | 7 ++- .../orchestrator/event_handlers_streaming.go | 8 +-- .../orchestrator/executor/executor_test.go | 2 +- .../task/handlers/process_handlers_test.go | 2 +- .../internal/task/repository/interface.go | 2 +- .../repository/session_repository_test.go | 57 ++++++++++++++++--- .../task/repository/sqlite/message.go | 11 ++-- .../internal/task/service/service_turns.go | 8 +-- apps/web/e2e/tests/tool-completion.spec.ts | 39 +++++++++++++ apps/web/lib/ws/handlers/turns.ts | 4 +- 10 files changed, 110 insertions(+), 30 deletions(-) create mode 100644 apps/web/e2e/tests/tool-completion.spec.ts diff --git a/apps/backend/internal/orchestrator/event_handlers.go b/apps/backend/internal/orchestrator/event_handlers.go index 3d4ddc6934..838ad76082 100644 --- a/apps/backend/internal/orchestrator/event_handlers.go +++ b/apps/backend/internal/orchestrator/event_handlers.go @@ -14,9 +14,10 @@ import ( // Agent event type string constants. const ( - agentEventComplete = "complete" - agentEventError = "error" - agentEventToolCall = "tool_call" + agentEventComplete = "complete" + agentEventCompleted = "completed" + agentEventError = "error" + agentEventToolCall = "tool_call" ) // buildTaskEventPayload builds the standard map payload for TaskUpdated events. diff --git a/apps/backend/internal/orchestrator/event_handlers_streaming.go b/apps/backend/internal/orchestrator/event_handlers_streaming.go index cd5e46f409..664b8a0482 100644 --- a/apps/backend/internal/orchestrator/event_handlers_streaming.go +++ b/apps/backend/internal/orchestrator/event_handlers_streaming.go @@ -309,9 +309,9 @@ func (s *Service) handleToolUpdateEvent(ctx context.Context, payload *lifecycle. // Determine message type from normalized payload for fallback creation msgType := toolKindToMessageType(payload.Data.Normalized) - // Handle all status updates (running, complete, error) + // Handle all status updates (running, complete, error, cancelled, pending, in_progress) switch payload.Data.ToolStatus { - case "running", agentEventComplete, "completed", "success", agentEventError, "failed": + case "running", agentEventComplete, agentEventCompleted, "success", agentEventError, "failed", "cancelled", "pending", "in_progress": if err := s.messageCreator.UpdateToolCallMessage( ctx, payload.TaskID, @@ -333,7 +333,7 @@ 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 == agentEventComplete || payload.Data.ToolStatus == "completed" || + if payload.Data.ToolStatus == agentEventComplete || payload.Data.ToolStatus == agentEventCompleted || payload.Data.ToolStatus == "success" || payload.Data.ToolStatus == agentEventError || payload.Data.ToolStatus == "failed" { s.updateTaskSessionState(ctx, payload.TaskID, payload.SessionID, models.TaskSessionStateRunning, "", true) } @@ -759,7 +759,7 @@ func (s *Service) persistTodoMessage(ctx context.Context, taskID, sessionID stri todos[i] = map[string]interface{}{ "text": e.Description, "status": e.Status, - "done": e.Status == "completed", + "done": e.Status == agentEventCompleted, } } metadata := map[string]interface{}{"todos": todos} diff --git a/apps/backend/internal/orchestrator/executor/executor_test.go b/apps/backend/internal/orchestrator/executor/executor_test.go index 605d105347..34e33681c7 100644 --- a/apps/backend/internal/orchestrator/executor/executor_test.go +++ b/apps/backend/internal/orchestrator/executor/executor_test.go @@ -390,7 +390,7 @@ func (m *mockRepository) GetActiveTurnBySessionID(ctx context.Context, sessionID } func (m *mockRepository) UpdateTurn(ctx context.Context, turn *models.Turn) error { return nil } func (m *mockRepository) CompleteTurn(ctx context.Context, id string) error { return nil } -func (m *mockRepository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) { +func (m *mockRepository) CompletePendingToolCallsForTurn(ctx context.Context, turnID string) (int64, error) { return 0, nil } func (m *mockRepository) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) { diff --git a/apps/backend/internal/task/handlers/process_handlers_test.go b/apps/backend/internal/task/handlers/process_handlers_test.go index 1440a6457e..29ad1792a3 100644 --- a/apps/backend/internal/task/handlers/process_handlers_test.go +++ b/apps/backend/internal/task/handlers/process_handlers_test.go @@ -163,7 +163,7 @@ func (m *mockRepository) UpdateTurn(ctx context.Context, turn *models.Turn) erro func (m *mockRepository) CompleteTurn(ctx context.Context, id string) error { return nil } -func (m *mockRepository) CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) { +func (m *mockRepository) CompletePendingToolCallsForTurn(ctx context.Context, turnID string) (int64, error) { return 0, nil } func (m *mockRepository) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) { diff --git a/apps/backend/internal/task/repository/interface.go b/apps/backend/internal/task/repository/interface.go index 0f80b9f63b..df3ab72844 100644 --- a/apps/backend/internal/task/repository/interface.go +++ b/apps/backend/internal/task/repository/interface.go @@ -78,7 +78,7 @@ type TurnRepository interface { GetActiveTurnBySessionID(ctx context.Context, sessionID string) (*models.Turn, error) UpdateTurn(ctx context.Context, turn *models.Turn) error CompleteTurn(ctx context.Context, id string) error - CompleteRunningToolCallsForTurn(ctx context.Context, turnID string) (int64, error) + CompletePendingToolCallsForTurn(ctx context.Context, turnID string) (int64, error) ListTurnsBySession(ctx context.Context, sessionID string) ([]*models.Turn, error) } diff --git a/apps/backend/internal/task/repository/session_repository_test.go b/apps/backend/internal/task/repository/session_repository_test.go index eac491aa3f..b386675fe6 100644 --- a/apps/backend/internal/task/repository/session_repository_test.go +++ b/apps/backend/internal/task/repository/session_repository_test.go @@ -310,7 +310,7 @@ func TestSQLiteRepository_UpdateTaskSessionState(t *testing.T) { } } -func TestSQLiteRepository_CompleteRunningToolCallsForTurn(t *testing.T) { +func TestSQLiteRepository_CompletePendingToolCallsForTurn(t *testing.T) { repo, cleanup := createTestSQLiteRepo(t) defer cleanup() ctx := context.Background() @@ -351,22 +351,43 @@ func TestSQLiteRepository_CompleteRunningToolCallsForTurn(t *testing.T) { Type: models.MessageTypeToolCall, Metadata: map[string]interface{}{"tool_call_id": "tc-3", "status": "running"}, } + // Create a tool call with status "pending" — should also be completed + pendingTool := &models.Message{ + ID: "msg-pending-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID, + AuthorType: models.MessageAuthorAgent, Content: "Pending tool", + Type: models.MessageTypeToolCall, + Metadata: map[string]interface{}{"tool_call_id": "tc-4", "status": "pending"}, + } + // Create a tool call with status "in_progress" — should also be completed + inProgressTool := &models.Message{ + ID: "msg-inprogress-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID, + AuthorType: models.MessageAuthorAgent, Content: "In-progress tool", + Type: models.MessageTypeToolCall, + Metadata: map[string]interface{}{"tool_call_id": "tc-5", "status": "in_progress"}, + } + // Create a tool call with status "error" — should NOT be affected + errorTool := &models.Message{ + ID: "msg-error-1", TaskSessionID: sessionID, TaskID: task.ID, TurnID: turnID, + AuthorType: models.MessageAuthorAgent, Content: "Error tool", + Type: models.MessageTypeToolCall, + Metadata: map[string]interface{}{"tool_call_id": "tc-6", "status": "error"}, + } - for _, msg := range []*models.Message{runningTool, completeTool, regularMsg, runningTool2} { + for _, msg := range []*models.Message{runningTool, completeTool, regularMsg, runningTool2, pendingTool, inProgressTool, errorTool} { if err := repo.CreateMessage(ctx, msg); err != nil { t.Fatalf("failed to create message %s: %v", msg.ID, err) } } // Execute - affected, err := repo.CompleteRunningToolCallsForTurn(ctx, turnID) + affected, err := repo.CompletePendingToolCallsForTurn(ctx, turnID) if err != nil { - t.Fatalf("CompleteRunningToolCallsForTurn failed: %v", err) + t.Fatalf("CompletePendingToolCallsForTurn failed: %v", err) } - // Should have updated exactly 2 running tool call messages - if affected != 2 { - t.Errorf("expected 2 affected rows, got %d", affected) + // Should have updated 4 non-terminal tool call messages (running x2, pending, in_progress) + if affected != 4 { + t.Errorf("expected 4 affected rows, got %d", affected) } // Verify running tool calls are now "complete" @@ -391,10 +412,28 @@ func TestSQLiteRepository_CompleteRunningToolCallsForTurn(t *testing.T) { t.Errorf("expected msg-regular-1 status 'running' (unchanged), got %v", msg4.Metadata["status"]) } + // Verify pending tool call is now "complete" + msg5, _ := repo.GetMessage(ctx, "msg-pending-1") + if msg5.Metadata["status"] != "complete" { + t.Errorf("expected msg-pending-1 status 'complete', got %v", msg5.Metadata["status"]) + } + + // Verify in_progress tool call is now "complete" + msg6, _ := repo.GetMessage(ctx, "msg-inprogress-1") + if msg6.Metadata["status"] != "complete" { + t.Errorf("expected msg-inprogress-1 status 'complete', got %v", msg6.Metadata["status"]) + } + + // Verify error tool call was NOT affected + msg7, _ := repo.GetMessage(ctx, "msg-error-1") + if msg7.Metadata["status"] != "error" { + t.Errorf("expected msg-error-1 status 'error' (unchanged), got %v", msg7.Metadata["status"]) + } + // Running again should affect 0 rows - affected2, err := repo.CompleteRunningToolCallsForTurn(ctx, turnID) + affected2, err := repo.CompletePendingToolCallsForTurn(ctx, turnID) if err != nil { - t.Fatalf("second CompleteRunningToolCallsForTurn failed: %v", err) + t.Fatalf("second CompletePendingToolCallsForTurn failed: %v", err) } if affected2 != 0 { t.Errorf("expected 0 affected rows on second call, got %d", affected2) diff --git a/apps/backend/internal/task/repository/sqlite/message.go b/apps/backend/internal/task/repository/sqlite/message.go index fdd97b7fd5..39336e67f2 100644 --- a/apps/backend/internal/task/repository/sqlite/message.go +++ b/apps/backend/internal/task/repository/sqlite/message.go @@ -297,22 +297,23 @@ func (r *Repository) FindMessageByPendingID(ctx context.Context, pendingID strin 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) { +// CompletePendingToolCallsForTurn marks all non-terminal tool call messages for a turn as "complete". +// This is a safety net to ensure no tool calls remain stuck in a non-terminal state (pending, +// running, in_progress, etc.) after a turn completes. +func (r *Repository) CompletePendingToolCallsForTurn(ctx context.Context, turnID string) (int64, error) { drv := r.db.DriverName() query := fmt.Sprintf(` UPDATE task_session_messages SET metadata = %s WHERE turn_id = ? - AND %s = 'running' + AND %s NOT IN ('complete', 'error') AND %s `, dialect.JSONSet(drv, "metadata", "status", "complete"), dialect.JSONExtract(drv, "metadata", "status"), dialect.JSONExtractIsNotNull(drv, "metadata", "tool_call_id")) result, err := r.db.ExecContext(ctx, r.db.Rebind(query), turnID) if err != nil { - return 0, fmt.Errorf("failed to complete running tool calls for turn %s: %w", turnID, err) + return 0, fmt.Errorf("failed to complete pending tool calls for turn %s: %w", turnID, err) } rows, _ := result.RowsAffected() diff --git a/apps/backend/internal/task/service/service_turns.go b/apps/backend/internal/task/service/service_turns.go index 4804addb2d..d502abd621 100644 --- a/apps/backend/internal/task/service/service_turns.go +++ b/apps/backend/internal/task/service/service_turns.go @@ -61,11 +61,11 @@ func (s *Service) CompleteTurn(ctx context.Context, turnID string) error { return err } - // Safety net: mark any tool calls still in "running" state as "complete" - if affected, err := s.turns.CompleteRunningToolCallsForTurn(ctx, turnID); err != nil { - s.logger.Warn("failed to complete running tool calls for turn", zap.String("turn_id", turnID), zap.Error(err)) + // Safety net: mark any tool calls still in a non-terminal state as "complete" + if affected, err := s.turns.CompletePendingToolCallsForTurn(ctx, turnID); err != nil { + s.logger.Warn("failed to complete pending tool calls for turn", zap.String("turn_id", turnID), zap.Error(err)) } else if affected > 0 { - s.logger.Info("completed stale running tool calls on turn end", + s.logger.Info("completed stale pending tool calls on turn end", zap.String("turn_id", turnID), zap.Int64("affected", affected)) } diff --git a/apps/web/e2e/tests/tool-completion.spec.ts b/apps/web/e2e/tests/tool-completion.spec.ts new file mode 100644 index 0000000000..9d35e5770b --- /dev/null +++ b/apps/web/e2e/tests/tool-completion.spec.ts @@ -0,0 +1,39 @@ +import { test, expect } from "../fixtures/test-base"; +import { SessionPage } from "../pages/session-page"; + +test.describe("Tool completion on turn end", () => { + test("tool groups have no spinner after turn completes with uncompleted tool calls", async ({ + testPage, + apiClient, + seedData, + }) => { + // Use tool_use WITHOUT tool_result — leaves tools in "pending" status. + // The safety net should mark them as complete when the turn ends. + const script = [ + 'e2e:tool_use("Terminal", {"command": "echo test1"})', + 'e2e:tool_use("Terminal", {"command": "echo test2"})', + 'e2e:message("done")', + ].join("\n"); + + const task = await apiClient.createTaskWithAgent( + seedData.workspaceId, + "Tool Completion Test", + seedData.agentProfileId, + { + description: script, + workflow_id: seedData.workflowId, + workflow_step_id: seedData.startStepId, + repository_ids: [seedData.repositoryId], + }, + ); + + await testPage.goto(`/t/${task.id}`); + const session = new SessionPage(testPage); + await session.waitForLoad(); + await expect(session.idleInput()).toBeVisible({ timeout: 30_000 }); + + // After turn completes, no grid spinners should remain in the chat + const spinners = session.chat.locator('[role="status"][aria-label="Loading"]'); + await expect(spinners).toHaveCount(0); + }); +}); diff --git a/apps/web/lib/ws/handlers/turns.ts b/apps/web/lib/ws/handlers/turns.ts index 9484b46383..5d8d385614 100644 --- a/apps/web/lib/ws/handlers/turns.ts +++ b/apps/web/lib/ws/handlers/turns.ts @@ -37,13 +37,13 @@ export function registerTurnsHandlers(store: StoreApi): WsHandlers { // Clear the active turn when it completes store.getState().setActiveTurn(payload.session_id, null); - // Safety net: mark any tool calls still in "running" state as "complete". + // Safety net: mark any tool calls still in a non-terminal state as "complete". // This handles edge cases where tool_update events were dropped or not processed. const messages = store.getState().messages.bySession[payload.session_id]; if (messages) { for (const msg of messages) { const meta = msg.metadata as Record | undefined; - if (meta?.status === "running" && meta?.tool_call_id) { + if (meta?.tool_call_id && meta?.status !== "complete" && meta?.status !== "error") { store.getState().updateMessage({ ...msg, metadata: { ...meta, status: "complete" },