Merge pull request #2764 from langbot-app/fix/runner-cancel-terminal-event

fix(agent): preserve terminal events when cancelling runner execution

src/langbot/pkg/agent/runner/orchestrator.py

@@ -15,6 +15,8 @@ from langbot_plugin.api.entities.builtin.pipeline import query as pipeline_query
 from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction
 from .reply_stream import ReplyStreamSession
 from ...core import app
+from ...i18n import t
+from ...i18n.workspace import get_workspace_language
 from ...api.http.context import ExecutionContext
 from ...pipeline.pool import get_query_execution_context
 from .binding_resolver import AgentBindingResolver
@@ -408,8 +410,8 @@ class AgentRunOrchestrator:
                             if isinstance(usage, dict):
                                 terminal_usage = usage
                         elif result_type == 'run.failed':
-                            terminal_status = 'failed'
                             data = result_dict.get('data') if isinstance(result_dict.get('data'), dict) else {}
+                            terminal_status = 'cancelled' if data.get('code') == 'cancelled' else 'failed'
                             terminal_reason = data.get('error') or data.get('code')
                             usage = result_dict.get('usage')
                             if isinstance(usage, dict):
@@ -445,10 +447,37 @@ class AgentRunOrchestrator:
                             yield result
 
                         run_snapshot = await self.journal.get_run(run_id)
-                        if run_snapshot and run_snapshot.get('cancel_requested_at') is not None:
+                        if (
+                            terminal_status is None
+                            and run_snapshot
+                            and run_snapshot.get('cancel_requested_at') is not None
+                        ):
                             terminal_status = 'cancelled'
                             terminal_reason = run_snapshot.get('status_reason') or 'cancel_requested'
-                            break
+                            # The Host closes the stream here, so it owns the missing terminal event.
+                            cancelled_result = {
+                                'type': 'run.failed',
+                                'run_id': run_id,
+                                'sequence': last_sequence + 1,
+                                'data': {
+                                    'code': 'cancelled',
+                                    'error': t(
+                                        'Task was cancelled',
+                                        locale=await get_workspace_language(self.ap, execution_context.workspace_uuid),
+                                    ),
+                                    'retryable': False,
+                                },
+                            }
+                            cancelled_event = await self.journal.append_run_result(
+                                result_dict=cancelled_result,
+                                run_id=run_id,
+                                sequence=None,
+                                source='host',
+                            )
+                            cancelled_result['sequence'] = cancelled_event['sequence']
+                            if result_observer is not None:
+                                await result_observer(cancelled_result)
+                            await self.result_normalizer.normalize(cancelled_result, descriptor)
             await self.journal.finalize_run(
                 run_id=run_id,
                 status=terminal_status or 'completed',
@@ -470,14 +499,20 @@ class AgentRunOrchestrator:
             )
             raise
         except Exception as exc:
-            execution_outcome = 'timeout' if self._is_deadline_exhausted(context) else 'failed'
+            execution_outcome = (
+                'cancelled'
+                if terminal_status == 'cancelled'
+                else 'timeout'
+                if self._is_deadline_exhausted(context)
+                else 'failed'
+            )
             # A chain that broke must always carry a reason, even for a bare exception.
             execution_error = str(exc) or type(exc).__name__
             failed_usage = terminal_usage
             await self.journal.finalize_run(
                 run_id=run_id,
-                status='timeout' if self._is_deadline_exhausted(context) else 'failed',
-                status_reason=str(exc),
+                status=execution_outcome,
+                status_reason=terminal_reason if execution_outcome == 'cancelled' else str(exc),
                 usage=failed_usage,
             )
             raise

src/langbot/pkg/agent/runner/run_journal.py

@@ -152,7 +152,7 @@ class AgentRunJournal:
         *,
         result_dict: dict[str, typing.Any],
         run_id: str,
-        sequence: int,
+        sequence: int | None,
         source: str = 'runner',
         metadata: dict[str, typing.Any] | None = None,
     ) -> dict[str, typing.Any]:

src/langbot/pkg/agent/runner/run_ledger_store.py

@@ -305,7 +305,7 @@ class RunLedgerStore:
         self,
         *,
         run_id: str,
-        sequence: int,
+        sequence: int | None,
         event_type: str,
         data: dict[str, typing.Any] | None = None,
         usage: dict[str, typing.Any] | None = None,
@@ -318,6 +318,17 @@ class RunLedgerStore:
         returned. This supports retrying append calls idempotently.
         """
         async with self._session_factory() as session:
+            if sequence is None:
+                # Serialize Host allocation across store instances and database connections.
+                locked = await session.execute(
+                    sqlalchemy.update(AgentRun).where(AgentRun.run_id == run_id).values(updated_at=AgentRun.updated_at)
+                )
+                if not locked.rowcount:
+                    raise ValueError(f'Run {run_id} not found')
+                maximum = await session.execute(
+                    sqlalchemy.select(sqlalchemy.func.max(AgentRunEvent.sequence)).where(AgentRunEvent.run_id == run_id)
+                )
+                sequence = int(maximum.scalar_one_or_none() or 0) + 1
             result = await session.execute(
                 sqlalchemy.select(AgentRunEvent).where(
                     AgentRunEvent.run_id == run_id,
@@ -352,31 +363,16 @@ class RunLedgerStore:
         metadata: dict[str, typing.Any] | None = None,
     ) -> dict[str, typing.Any] | None:
         """Append a Host-authored audit event after the current max sequence."""
-        async with self._session_factory() as session:
-            run = await self._get_run_row(session, run_id)
-            if run is None:
-                return None
-
-            result = await session.execute(
-                sqlalchemy.select(sqlalchemy.func.max(AgentRunEvent.sequence)).where(
-                    AgentRunEvent.run_id == run_id,
-                )
-            )
-            next_sequence = int(result.scalar_one_or_none() or 0) + 1
-            row = AgentRunEvent(
-                run_id=run_id,
-                sequence=next_sequence,
-                type=event_type,
-                data_json=_json_dumps(data or {}),
-                usage_json=None,
-                created_at=as_naive_utc(_utc_now()),
-                source='host',
-                metadata_json=_json_dumps(metadata or {}),
-            )
-            session.add(row)
-            await session.commit()
-            inflight_hub.notify_run(run_id, event_type)
-            return self._event_to_dict(row)
+        if await self.get_run(run_id) is None:
+            return None
+        return await self.append_event(
+            run_id=run_id,
+            sequence=None,
+            event_type=event_type,
+            data=data,
+            source='host',
+            metadata=metadata or {},
+        )
 
     async def pin_skill_revision(self, run_id: str, workspace_id: str, name: str, revision: str) -> str:
         """Pin a Host-owned revision without overwriting concurrent run metadata."""

tests/unit_tests/agent/test_orchestrator_integration.py

@@ -593,8 +593,12 @@ async def test_orchestrator_stops_after_cancel_request(clean_agent_state):
 
     orchestrator.journal.append_run_result = append_and_cancel_once
 
-    messages = [message async for message in orchestrator.run_from_query(make_query())]
+    messages = []
+    with pytest.raises(RunnerExecutionError) as caught:
+        async for message in orchestrator.run_from_query(make_query()):
+            messages.append(message)
 
+    assert caught.value.error_code == 'cancelled'
     assert [message.content for message in messages] == ['first']
     run_id = plugin_connector.contexts[0]['run_id']
     run = await RunLedgerStore(db_engine).get_run(run_id)
@@ -603,6 +607,86 @@ async def test_orchestrator_stops_after_cancel_request(clean_agent_state):
     assert run['status_reason'] == 'user stopped'
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize('result_type', ['tool.call.started', 'run.failed', 'run.completed'])
+@pytest.mark.parametrize('locale', ['en_US', 'zh_Hans'])
+@pytest.mark.parametrize('admin_audit', [False, True])
+async def test_cancel_request_preserves_one_terminal_event(clean_agent_state, result_type, locale, admin_audit):
+    """Cancellation must reach both observers and the journal regardless of event timing."""
+    from collections import OrderedDict
+
+    from langbot.pkg.i18n import t, language_scope
+
+    store = RunLedgerStore(clean_agent_state)
+    resumed = []
+
+    class CancellingConnector(FakePluginConnector):
+        async def run_runner(self, **kwargs):
+            context = kwargs['context']
+            self.contexts.append(context)
+            await store.request_cancel(run_id=context['run_id'], status_reason='cancel during export recovery')
+            data = {
+                'tool.call.started': {'tool_name': 'export_box_files', 'tool_call_id': 'export-2', 'parameters': {}},
+                'run.failed': {'code': 'cancelled', 'error': 'Runner cancellation diagnostic', 'retryable': False},
+                'run.completed': {'finish_reason': 'stop'},
+            }[result_type]
+            yield {'type': result_type, 'sequence': 7, 'data': data}
+            resumed.append(True)
+
+    connector = CancellingConnector()
+    ap = FakeApplication(connector, clean_agent_state)
+    ap._message_language_cache = OrderedDict({TEST_CONTEXT.workspace_uuid: (locale, float('inf'))})
+    orchestrator = AgentRunOrchestrator(ap, FakeRegistry(make_descriptor()))
+    plan = orchestrator.query_bridge.build_plan(make_query())
+    observed = []
+
+    async def observe(result):
+        observed.append(result)
+        if admin_audit and result['type'] == 'tool.call.started':
+            await store.append_audit_event(run_id=connector.contexts[0]['run_id'], event_type='admin.run_cancel')
+
+    failure = None
+    # The trusted Workspace language must take precedence over the caller's language.
+    with language_scope('ja_JP'):
+        try:
+            async for _ in orchestrator.run(
+                plan.event,
+                plan.binding,
+                adapter_context={'_execution_context': TEST_CONTEXT, '_result_observer': observe},
+            ):
+                pass
+        except RunnerExecutionError as exc:
+            failure = exc
+
+    run_id = connector.contexts[0]['run_id']
+    run = await store.get_run(run_id)
+    events, *_ = await store.page_run_events(run_id=run_id, limit=10)
+    terminals = [event for event in observed if event['type'] in {'run.completed', 'run.failed'}]
+    assert len(terminals) == 1
+    assert [event['type'] for event in events if event['type'] != 'admin.run_cancel'] == [
+        event['type'] for event in observed
+    ]
+    if result_type == 'run.completed':
+        assert failure is None
+        assert run['status'] == 'completed'
+    else:
+        assert failure is not None and failure.error_code == 'cancelled'
+        assert run['status'] == 'cancelled'
+        assert terminals[0]['data']['code'] == 'cancelled'
+    if result_type == 'tool.call.started':
+        assert not resumed
+        assert run['status_reason'] == 'cancel during export recovery'
+        assert [event['sequence'] for event in events] == ([7, 8, 9] if admin_audit else [7, 8])
+        assert terminals[0]['sequence'] == events[-1]['sequence']
+        assert events[-1]['source'] == 'host'
+        assert terminals[0]['data']['error'] == t('Task was cancelled', locale=locale)
+    else:
+        assert len(events) == 1
+        assert events[0]['source'] == 'runner'
+        if result_type == 'run.failed':
+            assert terminals[0]['data']['error'] == 'Runner cancellation diagnostic'
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize('cancel_while_waiting', [False, True])
 async def test_consumer_exit_finalizes_cancelled_ledger(clean_agent_state, cancel_while_waiting):

tests/unit_tests/agent/test_run_ledger_store.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import datetime
+import asyncio
 import json
 
 import pytest
@@ -418,6 +419,25 @@ async def test_append_audit_event_uses_next_sequence(store):
     assert await store.append_audit_event(run_id='missing', event_type='admin.missing') is None
 
 
+@pytest.mark.asyncio
+async def test_host_events_allocate_sequences_across_stores(store, db_engine):
+    await store.create_run(run_id='concurrent-host', event_id=None, binding_id=None, runner_id='runner')
+    other = RunLedgerStore(db_engine)
+    events = await asyncio.gather(
+        store.append_audit_event(run_id='concurrent-host', event_type='admin.run_cancel'),
+        other.append_event(
+            run_id='concurrent-host',
+            sequence=None,
+            event_type='run.failed',
+            data={'code': 'cancelled'},
+            source='host',
+        ),
+    )
+    assert sorted(event['sequence'] for event in events) == [1, 2]
+    persisted, *_ = await store.page_run_events(run_id='concurrent-host', limit=10)
+    assert {event['type'] for event in persisted} == {'admin.run_cancel', 'run.failed'}
+
+
 @pytest.mark.asyncio
 async def test_runtime_register_heartbeat_list_and_mark_stale(store):
     registered = await store.register_runtime(