fix(agent): clear running state on workflow failures

Author: jordan-umusu

packages/tracecat-ee/tracecat_ee/agent/workflows/durable.py

@@ -337,6 +337,7 @@ def _resolve_agent_output(
 LOAD_TERMINAL_MESSAGE_HISTORY_PATCH = "durable-agent-load-terminal-message-history-v1"
 AGENT_ACTIVITY_GRACEFUL_CANCEL_PATCH = "durable-agent-activity-graceful-cancel-v1"
 AGENT_SESSION_STATUS_PATCH = "durable-agent-session-status-v1"
+AGENT_SESSION_FAILURE_STATUS_PATCH = "durable-agent-session-failure-status-v1"
 AGENT_ACTIVITY_GRACEFUL_CANCEL_HEARTBEAT_TIMEOUT_SECONDS = 10
 
 
@@ -718,6 +719,11 @@ async def run(self, args: AgentWorkflowArgs) -> AgentOutput:
                         error=str(emit_error),
                     )
             raise
+        except TemporalCancelledError:
+            raise
+        except Exception:
+            await self._set_agent_session_failed_for_unhandled_failure()
+            raise
 
     @workflow.update
     def request_cancel(self, request: WorkflowCancelRequest) -> None:
@@ -799,6 +805,16 @@ async def _set_agent_session_status(
                 error=str(e),
             )
 
+    async def _set_agent_session_failed_for_unhandled_failure(self) -> None:
+        """Mark failed for post-deploy unhandled workflow failures."""
+        if not workflow.patched(AGENT_SESSION_FAILURE_STATUS_PATCH):
+            return
+
+        await self._set_agent_session_status(
+            AgentSessionStatus.FAILED,
+            clear_curr_run_id=True,
+        )
+
     async def _run_agent_activity_turn(
         self,
         executor_input: AgentExecutorInput,

tests/temporal/test_durable_agent_workflow.py

@@ -667,6 +667,49 @@ async def mock_run_agent_activity(
     assert status_update_inputs[-1].clear_curr_run_id is True
 
 
+@pytest.mark.anyio
+@pytest.mark.integration
+async def test_agent_workflow_marks_failed_when_executor_activity_errors(
+    temporal_client: Client,
+    agent_worker_factory,
+    agent_workflow_args: AgentWorkflowArgs,
+    mock_session_id: uuid.UUID,
+) -> None:
+    """Executor ActivityError failures should release the running session gate."""
+    queue = f"test-agent-queue-{mock_session_id}"
+    status_update_inputs: list[UpdateSessionStatusInput] = []
+
+    def mock_executor(
+        _call_count: int,
+        _input: AgentExecutorInput,
+    ) -> AgentExecutorResult:
+        raise ApplicationError("sandbox worker timed out")
+
+    activities = create_activities_with_mock_executor(
+        mock_executor,
+        status_update_inputs=status_update_inputs,
+    )
+
+    async with agent_worker_factory(
+        temporal_client, task_queue=queue, custom_activities=activities
+    ):
+        with pytest.raises(WorkflowFailureError):
+            await temporal_client.execute_workflow(
+                DurableAgentWorkflow.run,
+                agent_workflow_args,
+                id=AgentWorkflowID(mock_session_id),
+                task_queue=queue,
+                retry_policy=RETRY_POLICIES["workflow:fail_fast"],
+                execution_timeout=timedelta(seconds=30),
+            )
+
+    assert [input.status for input in status_update_inputs] == [
+        AgentSessionStatus.RUNNING,
+        AgentSessionStatus.FAILED,
+    ]
+    assert status_update_inputs[-1].clear_curr_run_id is True
+
+
 @pytest.mark.anyio
 @pytest.mark.integration
 async def test_agent_workflow_preserves_legacy_activity_message_history(

tests/unit/test_agent_activities.py

@@ -47,7 +47,7 @@
     load_session_activity,
     load_session_messages_activity,
 )
-from tracecat.agent.session.types import AgentSessionEntity
+from tracecat.agent.session.types import AgentSessionEntity, AgentSessionStatus
 from tracecat.agent.skill.types import ResolvedSkillRef
 from tracecat.agent.subagents import ResolvedAgentsConfig
 from tracecat.agent.tools import BuildToolsResult
@@ -394,6 +394,41 @@ async def test_idempotent_for_existing_session(
         assert result.success is True
         assert result.session_id == mock_session_id
 
+    @pytest.mark.anyio
+    @patch("tracecat.agent.session.activities.AgentSessionService.with_session")
+    async def test_curr_run_id_does_not_mark_session_running(
+        self, mock_with_session, mock_role: Role, mock_session_id: uuid.UUID
+    ):
+        """Create setup records the run token without taking the running lock."""
+        curr_run_id = uuid.uuid4()
+        input = CreateSessionInput(
+            role=mock_role,
+            session_id=mock_session_id,
+            entity_type=AgentSessionEntity.WORKFLOW,
+            entity_id=uuid.uuid4(),
+            curr_run_id=curr_run_id,
+        )
+
+        mock_agent_session = MagicMock()
+        mock_agent_session.agents_binding = None
+        mock_agent_session.status = AgentSessionStatus.IDLE.value
+        mock_service = AsyncMock()
+        mock_service.get_or_create_session.return_value = (mock_agent_session, False)
+        mock_service.session = MagicMock()
+        mock_service.session.commit = AsyncMock()
+
+        mock_ctx = AsyncMock()
+        mock_ctx.__aenter__.return_value = mock_service
+        mock_with_session.return_value = mock_ctx
+
+        result = await create_session_activity(input)
+
+        assert result.success is True
+        assert mock_agent_session.curr_run_id == curr_run_id
+        assert mock_agent_session.status == AgentSessionStatus.IDLE.value
+        mock_service.session.add.assert_called_once_with(mock_agent_session)
+        mock_service.session.commit.assert_awaited_once()
+
     @pytest.mark.anyio
     @patch("tracecat.agent.session.activities.AgentSessionService.with_session")
     async def test_backfills_disabled_agents_binding_for_legacy_existing_session(

tests/unit/test_agent_session_router.py

@@ -20,7 +20,7 @@
     send_message,
     stream_session_events,
 )
-from tracecat.agent.session.types import AgentSessionEntity
+from tracecat.agent.session.types import AgentSessionEntity, AgentSessionStatus
 from tracecat.auth.types import Role
 from tracecat.chat.schemas import (
     ApprovalDecision,
@@ -87,7 +87,7 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
 
 @pytest.mark.anyio
 async def test_get_session_includes_agents_binding() -> None:
-    session_stub = _agent_session_stub()
+    session_stub = _agent_session_stub(status=AgentSessionStatus.RUNNING.value)
     fake_svc = SimpleNamespace(
         get_session=AsyncMock(return_value=session_stub),
         list_messages=AsyncMock(return_value=[]),
@@ -107,11 +107,14 @@ async def test_get_session_includes_agents_binding() -> None:
         "enabled": False,
         "subagents": [],
     }
+    assert response.turn_status is AgentSessionStatus.RUNNING
 
 
 @pytest.mark.anyio
 async def test_get_session_vercel_includes_agents_binding() -> None:
-    session_stub = _agent_session_stub()
+    session_stub = _agent_session_stub(
+        status=AgentSessionStatus.WAITING_FOR_APPROVAL.value
+    )
     fake_svc = SimpleNamespace(
         get_session=AsyncMock(return_value=session_stub),
         list_messages=AsyncMock(return_value=[]),
@@ -131,6 +134,7 @@ async def test_get_session_vercel_includes_agents_binding() -> None:
         "enabled": False,
         "subagents": [],
     }
+    assert response.turn_status is AgentSessionStatus.WAITING_FOR_APPROVAL
 
 
 @pytest.mark.anyio

tracecat/agent/session/activities.py

@@ -204,8 +204,9 @@ async def create_session_activity(input: CreateSessionInput) -> CreateSessionRes
 
             # Set curr_run_id if provided (for workflow-initiated sessions)
             if input.curr_run_id is not None:
+                # The workflow marks RUNNING only after setup completes, just
+                # before handing control to the executor activity.
                 agent_session.curr_run_id = input.curr_run_id
-                agent_session.status = AgentSessionStatus.RUNNING.value
                 service.session.add(agent_session)
                 await service.session.commit()
 

tracecat/agent/session/schemas.py

@@ -6,7 +6,7 @@
 from datetime import datetime
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 from tracecat.agent.adapter.vercel import UIMessage
 from tracecat.agent.common.stream_types import HarnessType
@@ -103,7 +103,7 @@ class AgentSessionHistoryRead(BaseModel):
     created_at: datetime
     updated_at: datetime
 
-    model_config = {"from_attributes": True}
+    model_config = ConfigDict(from_attributes=True)
 
 
 class AgentSessionRead(BaseModel):
@@ -136,7 +136,7 @@ class AgentSessionRead(BaseModel):
     created_at: datetime
     updated_at: datetime
 
-    model_config = {"from_attributes": True}
+    model_config = ConfigDict(from_attributes=True, validate_by_name=True)
 
 
 class AgentSessionReadWithMessages(AgentSessionRead):