From fafa177d2ef30c09707638311252f852f74a52b3 Mon Sep 17 00:00:00 2001 From: Trecek Date: Thu, 11 Jun 2026 16:47:40 -0700 Subject: [PATCH] fix: eager session-id persistence via on_session_id_resolved callback (#4072) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The dispatch system silently dropped dispatched_session_id (and related identity fields) on CancelledError: only status/reason/ended_at were written via mark_dispatch_interrupted(). The session ID was known in acc.stdout_session_id/acc.channel_b_session_id inside the task group but never persisted to state. Add an on_session_id_resolved callback fired from _extract_stdout_session_id inside the live-execution task group, mirroring the existing on_pid_resolved pattern. The closure calls mark_dispatch_session_identity to eagerly persist the session ID while the dispatch is still RUNNING. Three layers of defense: 1. Eager callback writes to state during execution 2. CancelledError handler passes captured_session_id + marker_dir to mark_dispatch_interrupted (additive — only writes empty fields) 3. Normal completion path prefers eagerly-captured value over skill_result mark_dispatch_interrupted now accepts optional dispatched_session_id, session_chain, and dispatched_session_log_dir. New mark_dispatch_session_identity function is idempotent and status-guarded (RUNNING only). Co-Authored-By: Claude Opus 4.6 (1M context) --- .../core/types/_type_protocols_execution.py | 1 + .../core/types/_type_subprocess.py | 1 + .../execution/headless/__init__.py | 2 + .../execution/headless/_headless_execute.py | 2 + src/autoskillit/execution/process/__init__.py | 4 + .../execution/process/_process_race.py | 5 + src/autoskillit/execution/recording.py | 2 + src/autoskillit/fleet/__init__.py | 2 + src/autoskillit/fleet/_api.py | 24 ++++- src/autoskillit/fleet/state.py | 43 ++++++++- tests/contracts/test_protocol_satisfaction.py | 2 + tests/execution/AGENTS.md | 1 + .../test_process_session_id_callback.py | 82 ++++++++++++++++ tests/fakes.py | 3 + tests/fleet/test_api.py | 94 +++++++++++++++++++ tests/fleet/test_state.py | 42 +++++++++ tests/fleet/test_state_lock_contract.py | 6 ++ 17 files changed, 314 insertions(+), 2 deletions(-) create mode 100644 tests/execution/test_process_session_id_callback.py diff --git a/src/autoskillit/core/types/_type_protocols_execution.py b/src/autoskillit/core/types/_type_protocols_execution.py index 2a8daaec4e..189e73c7db 100644 --- a/src/autoskillit/core/types/_type_protocols_execution.py +++ b/src/autoskillit/core/types/_type_protocols_execution.py @@ -109,6 +109,7 @@ async def dispatch_food_truck( session_id: str | None = None, resume_message: str | None = None, backend_override: str | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SkillResult: ... diff --git a/src/autoskillit/core/types/_type_subprocess.py b/src/autoskillit/core/types/_type_subprocess.py index 538d18a15d..e4bc353fd7 100644 --- a/src/autoskillit/core/types/_type_subprocess.py +++ b/src/autoskillit/core/types/_type_subprocess.py @@ -175,4 +175,5 @@ def __call__( session_record_types: frozenset[str] = frozenset({"assistant"}), inspector_callback: InspectorCallback | None = None, workload_basenames: frozenset[str] | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> Awaitable[SubprocessResult]: ... diff --git a/src/autoskillit/execution/headless/__init__.py b/src/autoskillit/execution/headless/__init__.py index afc0b42717..66693fae58 100644 --- a/src/autoskillit/execution/headless/__init__.py +++ b/src/autoskillit/execution/headless/__init__.py @@ -380,6 +380,7 @@ async def dispatch_food_truck( session_id: str | None = None, resume_message: str | None = None, backend_override: str | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SkillResult: if self._ctx.backend is not None and not self._ctx.backend.capabilities.food_truck_capable: raise RuntimeError( @@ -483,4 +484,5 @@ async def dispatch_food_truck( marker_dir=effective_marker_dir, session_id=session_id, model_identity=model_identity, + on_session_id_resolved=on_session_id_resolved, ) diff --git a/src/autoskillit/execution/headless/_headless_execute.py b/src/autoskillit/execution/headless/_headless_execute.py index 5425504c4e..26c8d6dcb1 100644 --- a/src/autoskillit/execution/headless/_headless_execute.py +++ b/src/autoskillit/execution/headless/_headless_execute.py @@ -111,6 +111,7 @@ async def _execute_claude_headless( model_identity: ModelIdentity = ModelIdentity.unknown(), inspector_eligible: bool = False, inspector_model: str = "", + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SkillResult: """Shared subprocess execution for headless Claude sessions. @@ -262,6 +263,7 @@ async def _execute_claude_headless( max_extension_seconds=max_extension_seconds, marker_dir=marker_dir, session_id=session_id, + on_session_id_resolved=on_session_id_resolved, stream_parser=_stream_parser, completion_record_types=_step_backend.capabilities.completion_record_types, session_record_types=_step_backend.capabilities.session_record_types, diff --git a/src/autoskillit/execution/process/__init__.py b/src/autoskillit/execution/process/__init__.py index 1d9ba037f9..c9768c1ac9 100644 --- a/src/autoskillit/execution/process/__init__.py +++ b/src/autoskillit/execution/process/__init__.py @@ -220,6 +220,7 @@ async def run_managed_async( stream_parser: StreamParser | None = None, inspector_callback: InspectorCallback | None = None, workload_basenames: frozenset[str] | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SubprocessResult: """Async subprocess execution with temp file I/O and process tree cleanup. @@ -352,6 +353,7 @@ async def run_managed_async( functools.partial( _extract_stdout_session_id, stream_parser=stream_parser, + on_session_id_resolved=on_session_id_resolved, ), stdout_path, acc, @@ -632,6 +634,7 @@ async def __call__( session_record_types: frozenset[str] = frozenset({"assistant"}), inspector_callback: InspectorCallback | None = None, workload_basenames: frozenset[str] | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SubprocessResult: return await run_managed_async( cmd, @@ -657,4 +660,5 @@ async def __call__( session_record_types=session_record_types, inspector_callback=inspector_callback, workload_basenames=workload_basenames, + on_session_id_resolved=on_session_id_resolved, ) diff --git a/src/autoskillit/execution/process/_process_race.py b/src/autoskillit/execution/process/_process_race.py index e1256bbb85..2192292788 100644 --- a/src/autoskillit/execution/process/_process_race.py +++ b/src/autoskillit/execution/process/_process_race.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import asdict, dataclass, field from pathlib import Path from typing import TYPE_CHECKING, assert_never @@ -366,6 +367,8 @@ async def _extract_stdout_session_id( stream_parser: StreamParser | None = None, _poll_interval: float = 0.3, _timeout: float = 10.0, + *, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> None: """Extract session ID from stdout type=system record and deposit on accumulator. @@ -413,6 +416,8 @@ async def _extract_stdout_session_id( if sid: acc.stdout_session_id = sid logger.debug("stdout_session_id_extracted", session_id=sid) + if on_session_id_resolved is not None: + on_session_id_resolved(sid) ready.set() return logger.debug("stdout_session_id_extraction_timeout", timeout=_timeout) diff --git a/src/autoskillit/execution/recording.py b/src/autoskillit/execution/recording.py index bd3e190201..f34549c746 100644 --- a/src/autoskillit/execution/recording.py +++ b/src/autoskillit/execution/recording.py @@ -140,6 +140,7 @@ async def __call__( session_record_types: frozenset[str] = frozenset({"assistant"}), inspector_callback: InspectorCallback | None = None, workload_basenames: frozenset[str] | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SubprocessResult: step_name = (env or {}).get(SCENARIO_STEP_NAME_ENV, "") @@ -430,6 +431,7 @@ async def __call__( session_record_types: frozenset[str] = frozenset({"assistant"}), inspector_callback: InspectorCallback | None = None, workload_basenames: frozenset[str] | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SubprocessResult: step_name = (env or {}).get(SCENARIO_STEP_NAME_ENV, "") diff --git a/src/autoskillit/fleet/__init__.py b/src/autoskillit/fleet/__init__.py index 467e88480c..b883fc115f 100644 --- a/src/autoskillit/fleet/__init__.py +++ b/src/autoskillit/fleet/__init__.py @@ -61,6 +61,7 @@ mark_dispatch_interrupted, mark_dispatch_resumable, mark_dispatch_running, + mark_dispatch_session_identity, normalize_dispatch_token_usage, read_all_campaign_captures, read_state, @@ -146,6 +147,7 @@ "record_gate_outcome", "mark_dispatch_interrupted", "mark_dispatch_resumable", + "mark_dispatch_session_identity", "mark_dispatch_running", "read_all_campaign_captures", "read_state", diff --git a/src/autoskillit/fleet/_api.py b/src/autoskillit/fleet/_api.py index 1b29478f60..73a2fcc1f5 100644 --- a/src/autoskillit/fleet/_api.py +++ b/src/autoskillit/fleet/_api.py @@ -629,6 +629,7 @@ def _reject_with_state(error_code: FleetErrorCode, message: str) -> DispatchResu _dispatched_ticks: list[int] = [] _dispatched_create_time: list[float] = [] _dispatched_boot_id: list[str] = [] + _dispatched_session_id: list[str] = [] # Collect prior dispatch_ids from attempt_history for defense-in-depth parsing prior_ids: list[str] = [] @@ -678,6 +679,14 @@ def _on_spawn(pid: int, ticks: int) -> None: issue_url=_issue_urls_raw, ) + def _on_session_id(session_id: str) -> None: + from autoskillit.fleet.state import mark_dispatch_session_identity + + _dispatched_session_id.append(session_id) + mark_dispatch_session_identity( + state_path, effective_name, dispatched_session_id=session_id + ) + marker_dir: Path | None = None if _locator is not None: try: @@ -710,6 +719,7 @@ def _on_spawn(pid: int, ticks: int) -> None: project_dir=str(tool_ctx.project_dir), marker_dir=marker_dir, session_id=caller_session_id, + on_session_id_resolved=_on_session_id, timeout=resolved_timeout, idle_output_timeout=float(idle_output_timeout) if idle_output_timeout is not None @@ -733,11 +743,21 @@ def _on_spawn(pid: int, ticks: int) -> None: try: from autoskillit.fleet.state import mark_dispatch_interrupted # noqa: PLC0415 + captured_session_id = _dispatched_session_id[0] if _dispatched_session_id else "" + if not captured_session_id: + sr = locals().get("skill_result") + if sr is not None: + captured_session_id = getattr(sr, "session_id", "") or "" + with anyio.CancelScope(shield=True): mark_dispatch_interrupted( state_path, effective_name, reason="signal_induced_cancellation", + dispatched_session_id=captured_session_id, + dispatched_session_log_dir=str(marker_dir) + if marker_dir is not None + else "", ) except Exception: logger.warning( @@ -894,7 +914,9 @@ def _on_spawn(pid: int, ticks: int) -> None: dispatch_id=dispatch_id, campaign_id=campaign_id, caller_session_id=caller_session_id, - dispatched_session_id=skill_result.session_id, + dispatched_session_id=_dispatched_session_id[0] + if _dispatched_session_id + else skill_result.session_id, session_chain=extended_chain, dispatched_session_log_dir=project_log_dir, dispatched_pid=_dispatched_pid[0] if _dispatched_pid else 0, diff --git a/src/autoskillit/fleet/state.py b/src/autoskillit/fleet/state.py index 0f0d50a8d8..8ab473a048 100644 --- a/src/autoskillit/fleet/state.py +++ b/src/autoskillit/fleet/state.py @@ -82,6 +82,7 @@ "read_state", "mark_dispatch_running", "mark_dispatch_interrupted", + "mark_dispatch_session_identity", "mark_dispatch_resumable", "reset_blocking_dispatch", "append_dispatch_record", @@ -395,8 +396,17 @@ def mark_dispatch_interrupted( dispatch_name: str, *, reason: str, + dispatched_session_id: str = "", + session_chain: list[str] | None = None, + dispatched_session_log_dir: str = "", ) -> None: - """Atomically mark a dispatch as interrupted with a reason.""" + """Atomically mark a dispatch as interrupted with a reason. + + Identity fields (dispatched_session_id, session_chain, dispatched_session_log_dir) + are additive — they are only written if the corresponding field on the record is + still empty, preventing the cancellation handler from overwriting values already + eagerly persisted by the on_session_id_resolved callback. + """ with CampaignStateMutator(state_path) as m: if m.state is None: raise FileNotFoundError(f"State file not found or corrupted: {state_path}") @@ -406,12 +416,43 @@ def mark_dispatch_interrupted( d.status = DispatchStatus.INTERRUPTED d.reason = reason d.ended_at = time.time() + if dispatched_session_id and not d.dispatched_session_id: + d.dispatched_session_id = dispatched_session_id + if session_chain is not None and not d.session_chain: + d.session_chain = session_chain + if dispatched_session_log_dir and not d.dispatched_session_log_dir: + d.dispatched_session_log_dir = dispatched_session_log_dir m.mark_dirty() return else: raise ValueError(f"Dispatch '{dispatch_name}' not found in state") +def mark_dispatch_session_identity( + state_path: Path, + dispatch_name: str, + *, + dispatched_session_id: str, +) -> None: + """Eagerly persist the dispatched_session_id once it is discovered during execution. + + Idempotent and status-guarded: only writes when the dispatch is RUNNING and + dispatched_session_id is still empty. If the dispatch has already transitioned + to INTERRUPTED (race condition), this is a no-op. + """ + with CampaignStateMutator(state_path) as m: + if m.state is None: + raise FileNotFoundError(f"State file not found or corrupted: {state_path}") + for d in m.state.dispatches: + if d.name == dispatch_name: + if d.status == DispatchStatus.RUNNING and not d.dispatched_session_id: + d.dispatched_session_id = dispatched_session_id + m.mark_dirty() + return + else: + raise ValueError(f"Dispatch '{dispatch_name}' not found in state") + + def mark_dispatch_resumable( state_path: Path, dispatch_name: str, diff --git a/tests/contracts/test_protocol_satisfaction.py b/tests/contracts/test_protocol_satisfaction.py index 0de78f51e1..9c9ff942ea 100644 --- a/tests/contracts/test_protocol_satisfaction.py +++ b/tests/contracts/test_protocol_satisfaction.py @@ -410,6 +410,7 @@ def test_req_api_001_public_params_present(self): "stream_parser", "inspector_callback", "workload_basenames", + "on_session_id_resolved", } assert expected == public_params, ( f"run_managed_async public params changed.\n" @@ -481,6 +482,7 @@ def test_req_api_002_call_params_match_protocol(self): "session_record_types", "inspector_callback", "workload_basenames", + "on_session_id_resolved", } assert expected == actual, ( f"DefaultSubprocessRunner.__call__ params changed.\n" diff --git a/tests/execution/AGENTS.md b/tests/execution/AGENTS.md index e554c8d3db..0cb0187f4e 100644 --- a/tests/execution/AGENTS.md +++ b/tests/execution/AGENTS.md @@ -87,6 +87,7 @@ Subprocess integration, headless session, process lifecycle, and session result | `test_process_session_log_monitor.py` | Unit tests for _session_log_monitor — core detection + session ID + wiring | | `test_process_session_log_monitor_dispatch_marker.py` | Dispatch marker suppression gate tests for _session_log_monitor | | `test_process_session_log_monitor_stale_suppression.py` | TCP/CPU stale suppression gate and bounded suppression tests | +| `test_process_session_id_callback.py` | Tests for the on_session_id_resolved callback fired from inside the task group | | `test_process_submodules.py` | Tests verifying process.py decomposition into focused sub-modules (P8-2) | | `test_provider_outcome_container.py` | Tests for ProviderOutcome typed container construction — required fields and TypeError on omission | | `test_push_trigger_applies.py` | Unit tests for _push_trigger_applies_to_branch and _has_merge_group_trigger | diff --git a/tests/execution/test_process_session_id_callback.py b/tests/execution/test_process_session_id_callback.py new file mode 100644 index 0000000000..b7215cfc31 --- /dev/null +++ b/tests/execution/test_process_session_id_callback.py @@ -0,0 +1,82 @@ +"""Tests for the on_session_id_resolved callback fired from inside the task group.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest + +pytestmark = [pytest.mark.layer("execution"), pytest.mark.medium, pytest.mark.feature("execution")] + + +class TestOnSessionIdResolvedCallback: + """on_session_id_resolved callback fires during live subprocess execution.""" + + @pytest.mark.anyio + async def test_callback_fires_with_resolved_session_id( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """When stdout emits a system init record, on_session_id_resolved fires.""" + import anyio + + from autoskillit.execution.process._process_race import ( + RaceAccumulator, + _extract_stdout_session_id, + ) + + stdout_path = tmp_path / "stdout.jsonl" + session_id = "sess-callback-123" + init_record: dict[str, Any] = { + "type": "system", + "subtype": "init", + "session_id": session_id, + } + stdout_path.write_text(json.dumps(init_record) + "\n") + + acc = RaceAccumulator() + ready = anyio.Event() + captured: list[str] = [] + + def on_sid(sid: str) -> None: + captured.append(sid) + + await _extract_stdout_session_id( + stdout_path, + acc, + ready, + on_session_id_resolved=on_sid, + ) + + assert ready.is_set() + assert acc.stdout_session_id == session_id + assert captured == [session_id] + + @pytest.mark.anyio + async def test_callback_not_called_when_no_session_id_found(self, tmp_path: Path) -> None: + """When no init record is present, callback is never invoked.""" + import anyio + + from autoskillit.execution.process._process_race import ( + RaceAccumulator, + _extract_stdout_session_id, + ) + + stdout_path = tmp_path / "stdout.jsonl" + stdout_path.write_text("") + + acc = RaceAccumulator() + ready = anyio.Event() + captured: list[str] = [] + + await _extract_stdout_session_id( + stdout_path, + acc, + ready, + on_session_id_resolved=lambda sid: captured.append(sid), + _timeout=0.5, + ) + + assert not acc.stdout_session_id + assert captured == [] diff --git a/tests/fakes.py b/tests/fakes.py index fb73b449a6..4d1df0b8cf 100644 --- a/tests/fakes.py +++ b/tests/fakes.py @@ -132,6 +132,7 @@ class DispatchFoodTruckCall: marker_dir: Path | None = None session_id: str | None = None resume_message: str | None = None + on_session_id_resolved: Callable[[str], None] | None = None _DEFAULT_SKILL_RESULT = SkillResult( @@ -282,6 +283,7 @@ async def dispatch_food_truck( marker_dir: Path | None = None, session_id: str | None = None, resume_message: str | None = None, + on_session_id_resolved: Callable[[str], None] | None = None, ) -> SkillResult: self.dispatch_calls.append( DispatchFoodTruckCall( @@ -312,6 +314,7 @@ async def dispatch_food_truck( marker_dir=marker_dir, session_id=session_id, resume_message=resume_message, + on_session_id_resolved=on_session_id_resolved, ) ) if self._queue: diff --git a/tests/fleet/test_api.py b/tests/fleet/test_api.py index 003a7cf658..ba2fce2ffd 100644 --- a/tests/fleet/test_api.py +++ b/tests/fleet/test_api.py @@ -645,4 +645,98 @@ async def _spawn_then_cancel(**kwargs): d = state["dispatches"][0] assert d["status"] == "interrupted" assert d["reason"] == "signal_induced_cancellation" + assert ( + d["dispatched_session_id"] == "" + ) # No session_id available when CancelledError fires before dispatch completes assert tool_ctx.fleet_lock.active_count == 0 + + +class TestSessionIdEagerPersistence: + """Tests for the on_session_id_resolved callback fired during live execution. + + These tests verify that the session identity is eagerly persisted to state + when discovered during the subprocess execution task group, so that it + survives a CancelledError that arrives before normal completion. + """ + + @pytest.mark.anyio + async def test_post_completion_cancel_persists_session_id(self, tool_ctx, monkeypatch) -> None: + """CancelledError after session completes must persist dispatched_session_id.""" + from tests.fleet._helpers import _setup_dispatch + + _setup_dispatch(tool_ctx, monkeypatch) + + async def _resolve_session_then_cancel(**kwargs): + on_spawn = kwargs.get("on_spawn") + if on_spawn: + on_spawn(99999, 0) + on_session_id_resolved = kwargs.get("on_session_id_resolved") + if on_session_id_resolved: + on_session_id_resolved("test-session-abc") + raise asyncio.CancelledError + + tool_ctx.executor.dispatch_food_truck = _resolve_session_then_cancel + + with pytest.raises(asyncio.CancelledError): + from autoskillit.fleet import execute_dispatch + + await execute_dispatch( + tool_ctx=tool_ctx, + recipe="test-recipe", + task="do something", + ingredients=None, + dispatch_name="test-dispatch", + timeout_sec=None, + prompt_builder=lambda **kw: "prompt", + quota_checker=_no_sleep_quota_checker, + quota_refresher=_noop_quota_refresher, + ) + + state_files = list((tool_ctx.temp_dir / "dispatches").glob("*.json")) + assert len(state_files) == 1 + state = json.loads(state_files[0].read_text()) + d = state["dispatches"][0] + assert d["dispatched_session_id"] == "test-session-abc" + assert d["status"] == "interrupted" + + @pytest.mark.anyio + async def test_mid_session_cancel_persists_session_id_via_callback( + self, tool_ctx, monkeypatch + ) -> None: + """CancelledError during execution must still persist session_id via eager callback.""" + from tests.fleet._helpers import _setup_dispatch + + _setup_dispatch(tool_ctx, monkeypatch) + + async def _spawn_resolve_session_then_cancel(**kwargs): + on_spawn = kwargs.get("on_spawn") + if on_spawn: + on_spawn(99999, 0) + on_session_id_resolved = kwargs.get("on_session_id_resolved") + if on_session_id_resolved: + on_session_id_resolved("early-session-xyz") + raise asyncio.CancelledError + + tool_ctx.executor.dispatch_food_truck = _spawn_resolve_session_then_cancel + + with pytest.raises(asyncio.CancelledError): + from autoskillit.fleet import execute_dispatch + + await execute_dispatch( + tool_ctx=tool_ctx, + recipe="test-recipe", + task="do something", + ingredients=None, + dispatch_name="test-dispatch", + timeout_sec=None, + prompt_builder=lambda **kw: "prompt", + quota_checker=_no_sleep_quota_checker, + quota_refresher=_noop_quota_refresher, + ) + + state_files = list((tool_ctx.temp_dir / "dispatches").glob("*.json")) + assert len(state_files) == 1 + state = json.loads(state_files[0].read_text()) + d = state["dispatches"][0] + assert d["dispatched_session_id"] == "early-session-xyz" + assert d["status"] == "interrupted" diff --git a/tests/fleet/test_state.py b/tests/fleet/test_state.py index 3ae991a9d6..eb1cf51d51 100644 --- a/tests/fleet/test_state.py +++ b/tests/fleet/test_state.py @@ -19,8 +19,10 @@ build_protected_campaign_ids, classify_stale_dispatch, has_failed_dispatch, + mark_dispatch_interrupted, mark_dispatch_resumable, mark_dispatch_running, + mark_dispatch_session_identity, read_all_campaign_captures, read_state, resume_campaign_from_state, @@ -1761,3 +1763,43 @@ def test_upsert_failure_to_failure_preserves_existing_attempt_history( assert len(disp.attempt_history) == 2 assert disp.attempt_history[0]["reason"] == "first_failure" assert disp.attempt_history[1]["dispatch_id"] == "prior-attempt" + + +class TestInterruptedPayloadCompleteness: + """Regression guard: mark_dispatch_interrupted must preserve identity fields.""" + + def test_interrupted_does_not_clear_already_persisted_session_id(self, tmp_path: Path) -> None: + sp = _state_path(tmp_path) + write_initial_state(sp, "cid", "camp", "/m.yaml", _make_dispatches("d1")) + mark_dispatch_running(sp, "d1", dispatch_id="d-1", dispatched_pid=42) + mark_dispatch_session_identity(sp, "d1", dispatched_session_id="sess-123") + + mark_dispatch_interrupted(sp, "d1", reason="signal_induced_cancellation") + + state = read_state(sp) + assert state is not None + disp = next(d for d in state.dispatches if d.name == "d1") + assert disp.dispatched_session_id == "sess-123" + assert disp.status == DispatchStatus.INTERRUPTED + + +class TestMarkDispatchInterruptedAcceptsIdentity: + """mark_dispatch_interrupted should accept identity fields for late capture.""" + + def test_interrupted_accepts_dispatched_session_id(self, tmp_path: Path) -> None: + sp = _state_path(tmp_path) + write_initial_state(sp, "cid", "camp", "/m.yaml", _make_dispatches("d1")) + mark_dispatch_running(sp, "d1", dispatch_id="d-1", dispatched_pid=42) + + mark_dispatch_interrupted( + sp, + "d1", + reason="signal_induced_cancellation", + dispatched_session_id="late-capture-id", + ) + + state = read_state(sp) + assert state is not None + disp = next(d for d in state.dispatches if d.name == "d1") + assert disp.dispatched_session_id == "late-capture-id" + assert disp.status == DispatchStatus.INTERRUPTED diff --git a/tests/fleet/test_state_lock_contract.py b/tests/fleet/test_state_lock_contract.py index a6600b2f7e..7d0a1d97ad 100644 --- a/tests/fleet/test_state_lock_contract.py +++ b/tests/fleet/test_state_lock_contract.py @@ -24,6 +24,7 @@ mark_dispatch_interrupted, mark_dispatch_resumable, mark_dispatch_running, + mark_dispatch_session_identity, read_state, reset_blocking_dispatch, update_orchestrator_session_id, @@ -50,6 +51,7 @@ _MUTATION_FUNCTIONS: dict[str, object] = { "mark_dispatch_running": mark_dispatch_running, "mark_dispatch_interrupted": mark_dispatch_interrupted, + "mark_dispatch_session_identity": mark_dispatch_session_identity, "mark_dispatch_resumable": mark_dispatch_resumable, "append_dispatch_record": append_dispatch_record, "write_captured_values": write_captured_values, @@ -176,6 +178,8 @@ def tracking_flock(fd: int, op: int) -> None: fn(sp, "d1", dispatch_id="x", dispatched_pid=42) # type: ignore[operator] elif fn_name == "mark_dispatch_interrupted": fn(sp, "d1", reason="test") # type: ignore[operator] + elif fn_name == "mark_dispatch_session_identity": + fn(sp, "d1", dispatched_session_id="sess-test") # type: ignore[operator] elif fn_name == "mark_dispatch_resumable": fn(sp, "d1", sidecar_path="/tmp/sidecar") # type: ignore[operator] elif fn_name == "append_dispatch_record": @@ -322,6 +326,8 @@ def tracking_flock(fd: int | IO[Any], op: int) -> None: fn(sp, "d1", dispatch_id="x", dispatched_pid=42) # type: ignore[operator] elif fn_name == "mark_dispatch_interrupted": fn(sp, "d1", reason="test") # type: ignore[operator] + elif fn_name == "mark_dispatch_session_identity": + fn(sp, "d1", dispatched_session_id="sess-test") # type: ignore[operator] elif fn_name == "mark_dispatch_resumable": fn(sp, "d1", sidecar_path="/tmp/sidecar") # type: ignore[operator] elif fn_name == "append_dispatch_record":