diff --git a/src/langbot/pkg/agent/runner/orchestrator.py b/src/langbot/pkg/agent/runner/orchestrator.py index 48ca75cf9..0e04cc67a 100644 --- a/src/langbot/pkg/agent/runner/orchestrator.py +++ b/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 diff --git a/src/langbot/pkg/agent/runner/run_journal.py b/src/langbot/pkg/agent/runner/run_journal.py index c2ea15df8..ed288ff55 100644 --- a/src/langbot/pkg/agent/runner/run_journal.py +++ b/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]: diff --git a/src/langbot/pkg/agent/runner/run_ledger_store.py b/src/langbot/pkg/agent/runner/run_ledger_store.py index 06ec1d969..6be6aaab2 100644 --- a/src/langbot/pkg/agent/runner/run_ledger_store.py +++ b/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.""" diff --git a/tests/unit_tests/agent/test_orchestrator_integration.py b/tests/unit_tests/agent/test_orchestrator_integration.py index fa5fb18d1..602b7a93c 100644 --- a/tests/unit_tests/agent/test_orchestrator_integration.py +++ b/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): diff --git a/tests/unit_tests/agent/test_run_ledger_store.py b/tests/unit_tests/agent/test_run_ledger_store.py index 2195f86f7..04d41a26a 100644 --- a/tests/unit_tests/agent/test_run_ledger_store.py +++ b/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(