From ce5b7ffea10a40a495e05dd0b0e4a4568d9f0a35 Mon Sep 17 00:00:00 2001 From: BiFangKNT <1320414964@qq.com> Date: Fri, 23 Oct 2026 17:23:44 +0800 Subject: [PATCH] fix(agent): preserve terminal events when cancelling runner execution Emit and persist a Host cancellation event before closing the runner stream. Preserve existing terminal events and record cancelled runs consistently. Serialize Host event sequence allocation to prevent conflicts with admin cancellation audits, and notify observers with the persisted sequence. Add deterministic cancellation and concurrent event allocation regressions. --- src/langbot/pkg/agent/runner/orchestrator.py | 47 ++++++++-- src/langbot/pkg/agent/runner/run_journal.py | 2 +- .../pkg/agent/runner/run_ledger_store.py | 48 +++++------ .../agent/test_orchestrator_integration.py | 86 ++++++++++++++++++- .../unit_tests/agent/test_run_ledger_store.py | 20 +++++ 5 files changed, 169 insertions(+), 34 deletions(-) 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(