iCode-TUI 0.29.1__py3-none-any.whl

chrys/orchestration/engine/run/runtime_skills.py

@@ -0,0 +1,146 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Runtime skill refresh helpers for main-agent turns."""
+
+from __future__ import annotations
+
+import copy
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from chrys.foundation.events.types import AgentRuntimeUpdated, Warning
+
+if TYPE_CHECKING:
+    from chrys.foundation.events.bus import EventBus
+    from chrys.foundation.events.types import RuntimeSkillDetails
+    from chrys.orchestration.engine.build.loaded import AgentManifest
+    from chrys.orchestration.engine.loader import AgentLoader
+    from chrys.orchestration.engine.state.active_session import ActiveSession
+    from chrys.orchestration.engine.state.current_agent import CurrentAgent
+    from chrys.service.skills.model import SkillProviderWarning
+    from chrys.service.skills.provider import ChrysSkillsProvider, StagedSkillRefresh
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class StagedRuntimeSkillRefresh:
+    """Runtime skill refresh staged before an active-injection owner is revalidated."""
+
+    provider: ChrysSkillsProvider
+    context: StagedSkillRefresh
+    warnings: list[SkillProviderWarning]
+    skill_names: list[str]
+    skill_sources: dict[str, list[str]]
+    skill_details: list[RuntimeSkillDetails]
+    skill_catalog: str | None
+
+
+@dataclass(frozen=True)
+class CommittedSkillRefresh:
+    """Warnings and the exact manifest installed by one synchronous refresh commit."""
+
+    warnings: list[SkillProviderWarning]
+    manifest: AgentManifest
+
+
+class RuntimeSkillRefresher:
+    """Refresh runtime skills through the engine host."""
+
+    def __init__(
+        self,
+        *,
+        current: CurrentAgent,
+        loader: AgentLoader,
+        session: ActiveSession,
+        bus: EventBus,
+    ) -> None:
+        self._current = current
+        self._loader = loader
+        self._session = session
+        self._bus = bus
+
+    async def refresh(self, *, update_active_turn: bool = False) -> None:
+        """Refresh runtime skills and optionally update the active reminder snapshot."""
+        session_id = self._session.session_id
+        staged = await self.stage_refresh()
+        if staged is None:
+            return
+        committed = self.commit_staged_refresh(staged, suppress_provider_errors=True)
+        if committed is None:
+            return
+        await self.publish_committed_refresh(committed, session_id=session_id)
+        if update_active_turn and self._current.loaded is not None:
+            self._current.loaded.reminder_middleware.update_skill_catalog_for_active_turn()
+
+    async def stage_refresh(self) -> StagedRuntimeSkillRefresh | None:
+        """Discover runtime skills without mutating live provider or host runtime state."""
+        provider = self._current.loaded.skills_provider if self._current.loaded is not None else None
+        if provider is None:
+            return None
+        try:
+            staged = await provider.stage_context_refresh()
+        except Exception:
+            logger.exception("Failed to stage runtime skill refresh")
+            return None
+        return StagedRuntimeSkillRefresh(
+            provider=provider,
+            context=staged,
+            warnings=list(staged.warnings),
+            skill_names=staged.skill_names(),
+            skill_sources=staged.skill_sources(),
+            skill_details=staged.skill_details(),
+            skill_catalog=staged.render_catalog_reminder(),
+        )
+
+    def commit_staged_refresh(
+        self,
+        staged: StagedRuntimeSkillRefresh | None,
+        *,
+        suppress_provider_errors: bool = False,
+    ) -> CommittedSkillRefresh | None:
+        """Commit a staged runtime skill refresh synchronously to the captured owner."""
+        if staged is None:
+            return None
+        try:
+            warnings = staged.provider.commit_context_refresh(staged.context)
+        except Exception:
+            if not suppress_provider_errors:
+                raise
+            logger.exception("Failed to refresh runtime skills")
+            return None
+        manifest = self._loader.apply_skill_refresh(
+            skill_names=list(staged.skill_names),
+            skill_sources={source: list(names) for source, names in staged.skill_sources.items()},
+            skill_details=list(staged.skill_details),
+        )
+        return CommittedSkillRefresh(warnings=warnings, manifest=manifest)
+
+    async def publish_committed_refresh(
+        self,
+        committed: CommittedSkillRefresh | None,
+        *,
+        session_id: str | None,
+    ) -> None:
+        """Publish captured warnings and manifest values for one committed refresh."""
+        if committed is None:
+            return
+        for warning in committed.warnings:
+            await self._bus.publish(Warning(code=warning.code, message=warning.message, session_id=session_id))
+        await self._publish_runtime_update(committed.manifest, session_id=session_id)
+
+    async def _publish_runtime_update(self, manifest: AgentManifest, *, session_id: str | None) -> None:
+        """Publish the current runtime catalog."""
+        await self._bus.publish(
+            AgentRuntimeUpdated(
+                session_id=session_id,
+                model_profile_id=manifest.runtime_details.model.profile_id,
+                max_context_tokens=manifest.runtime_details.model.max_context_tokens,
+                tool_names=list(manifest.tool_names),
+                skill_names=list(manifest.skill_names),
+                memory_files=list(manifest.memory_files),
+                runtime_details=copy.deepcopy(manifest.runtime_details),
+            )
+        )

chrys/orchestration/engine/run/sub_agent_coordination.py

@@ -0,0 +1,126 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Engine-internal sub-agent pause coordination."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from chrys.orchestration.engine.state.machine import Trigger
+
+if TYPE_CHECKING:
+    from chrys.foundation.events.types import (
+        InvocationAborted,
+        InvocationAbortRequested,
+        InvocationCascadeAborted,
+        InvocationPaused,
+        InvocationResumed,
+        InvocationRetryRequested,
+    )
+    from chrys.orchestration.engine.run.turn_state import TurnRuntimeState
+    from chrys.orchestration.engine.state.current_agent import CurrentAgent
+    from chrys.orchestration.engine.state.machine import EngineStateMachine
+    from chrys.orchestration.engine.trajectory import TrajectoryRecorder
+    from chrys.service.session.history import SessionHistoryManager
+
+
+logger = logging.getLogger(__name__)
+
+
+async def on_sub_agent_retry(current: CurrentAgent, event: InvocationRetryRequested) -> None:
+    """Route a user's per-card Retry click to the owning controller."""
+    if current.loaded is None or current.loaded.sub_agent_tools is None:
+        return
+    current.loaded.sub_agent_tools.request_retry(event.invocation_id)
+
+
+async def on_sub_agent_abort(current: CurrentAgent, event: InvocationAbortRequested) -> None:
+    """Route a user's per-card Abort click to the owning controller."""
+    if current.loaded is None or current.loaded.sub_agent_tools is None:
+        return
+    current.loaded.sub_agent_tools.request_abort(event.invocation_id)
+
+
+async def on_sub_agent_paused(
+    turn_state: TurnRuntimeState,
+    history: SessionHistoryManager,
+    fsm: EngineStateMachine,
+    trajectory_recorder: TrajectoryRecorder,
+    event: InvocationPaused,
+) -> None:
+    """Track a newly paused sub-agent and drive FSM / marker.
+
+    Idempotent on the paused set — if the same id arrives twice
+    (controller re-publishes after retry exhaustion, for example)
+    the FSM transition only fires on the 0→1 edge.
+
+    Defensive FSM guard: if the parent run has already terminated
+    (e.g. the bus is dispatching a stale pause event that was queued
+    before the parent's task was cancelled and ``_post_run`` ran),
+    drop the event silently rather than re-inserting a marker on top
+    of an already-terminal ``interrupted``/``error`` marker.  The
+    parent invariant ("parent run ends only after all sub-agents
+    resolve") should make this unreachable, but defense in depth
+    keeps history well-formed if that invariant ever breaks.
+    """
+    if event.origin.kind != "sub_agent" or event.origin.root.kind == "workflow_node":
+        return
+    if not fsm.is_running():
+        logger.debug(
+            "Ignoring InvocationPaused for %s: FSM=%s (parent run already terminated)",
+            event.origin.invocation_id,
+            fsm.state.name,
+        )
+        return
+    first = not turn_state.paused_sub_agents
+    turn_state.paused_sub_agents.add(event.origin.invocation_id)
+    if first:
+        fsm.try_transition(Trigger.SUB_AGENT_PAUSED)
+        await trajectory_recorder.turn_suspended()
+    if history.is_bound:
+        history.upsert_awaiting_sub_agents_marker(sorted(turn_state.paused_sub_agents))
+
+
+async def on_sub_agent_unpaused(
+    turn_state: TurnRuntimeState,
+    history: SessionHistoryManager,
+    fsm: EngineStateMachine,
+    trajectory_recorder: TrajectoryRecorder,
+    event: InvocationResumed | InvocationAborted | InvocationCascadeAborted,
+) -> None:
+    """Common handler — retry/abort/cascade all remove the invocation from the paused set.
+
+    FSM transitions only on the N→0 edge (last paused sub-agent
+    resolved).  Marker is updated or stripped accordingly.
+    """
+    if event.origin.kind != "sub_agent" or event.origin.root.kind == "workflow_node":
+        return
+    if event.origin.invocation_id not in turn_state.paused_sub_agents:
+        return
+    turn_state.paused_sub_agents.discard(event.origin.invocation_id)
+    if not turn_state.paused_sub_agents:
+        fsm.try_transition(Trigger.SUB_AGENT_RESOLVED)
+        await trajectory_recorder.turn_resumed()
+        if history.is_bound:
+            history.remove_awaiting_sub_agents_marker()
+    elif history.is_bound:
+        history.upsert_awaiting_sub_agents_marker(sorted(turn_state.paused_sub_agents))
+    # Defensive invariant: FSM ``AWAITING_SUB_AGENTS`` must mirror a
+    # non-empty paused set.  A desync here would mean either a stale
+    # pause event leaked past the ``_fsm.is_running()`` guard in
+    # ``_on_sub_agent_paused`` or a ``try_transition`` silently no-op'd.
+    # Log instead of raising so a surprise in production doesn't crash
+    # the engine mid-run — the log points straight at the bug.
+    if bool(turn_state.paused_sub_agents) != fsm.is_awaiting_sub_agents():
+        logger.error(
+            "FSM/paused-set desync: paused=%s FSM=%s",
+            sorted(turn_state.paused_sub_agents),
+            fsm.state.name,
+        )
+
+
+def clear_parent_paused_state(turn_state: TurnRuntimeState, history: SessionHistoryManager) -> None:
+    """Clear engine-side paused sub-agent state after a parent run ends."""
+    turn_state.paused_sub_agents.clear()
+    history.remove_awaiting_sub_agents_marker()

chrys/orchestration/engine/run/turn_hooks.py

@@ -0,0 +1,258 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Lifecycle hook dispatch helpers for main-agent turns."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from chrys.foundation.events.types import Error
+from chrys.foundation.i18n import msg
+from chrys.foundation.platform import safe_getcwd
+from chrys.service.hooks.events import HookEvent
+from chrys.service.hooks.schema import HookDecision
+
+if TYPE_CHECKING:
+    from chrys.foundation.events.bus import EventBus
+    from chrys.foundation.trajectory.context import TrajectoryContext
+    from chrys.orchestration.engine.state.active_session import ActiveSession
+    from chrys.orchestration.engine.state.current_agent import CurrentAgent
+    from chrys.orchestration.engine.state.machine import EngineStateMachine
+    from chrys.service.context.compaction import PreCompactInfo
+    from chrys.service.hooks.manager import HookManager
+
+
+logger = logging.getLogger(__name__)
+
+_TURN_HOOKS_PROMPT_BLOCKED = msg(
+    "turn_hooks.prompt_blocked",
+    fallback="Prompt blocked by hook.",
+)
+
+
+class TurnHookDispatcher:
+    """Dispatch non-prompt turn lifecycle hooks."""
+
+    def __init__(
+        self,
+        *,
+        session: ActiveSession,
+        current: CurrentAgent,
+    ) -> None:
+        self._session = session
+        self._current = current
+        self._pending: set[asyncio.Task[None]] = set()
+
+    async def fire_before_turn(
+        self,
+        user_text: str,
+        *,
+        is_retry: bool = False,
+        target_operation_id: str | None = None,
+    ) -> None:
+        """Publish a ``before_turn`` hook event."""
+        if self._session.hook_manager is None or not self._session.hook_manager.has_hooks_for(HookEvent.BEFORE_TURN):
+            return
+        profile_name = self._session.agent_profile.name if self._session.agent_profile is not None else ""
+        await self._session.hook_manager.fire(
+            HookEvent.BEFORE_TURN,
+            {
+                "session_id": self._session.session_id,
+                "profile": profile_name,
+                "cwd": _workspace_cwd(self._session),
+                "turn": self._session.turn_number,
+                "user_text": user_text,
+                "is_retry": is_retry,
+            },
+            target_operation_id=target_operation_id,
+        )
+
+    async def fire_after_turn(self, *, failed: bool) -> None:
+        """Publish an ``after_turn`` hook event."""
+        if self._session.hook_manager is None or not self._session.hook_manager.has_hooks_for(HookEvent.AFTER_TURN):
+            return
+        profile_name = self._session.agent_profile.name if self._session.agent_profile is not None else ""
+        status = (
+            "failed"
+            if (self._current.loaded is not None and self._current.loaded.bindings.state.run_failed)
+            else (
+                "interrupted"
+                if (self._current.loaded is not None and self._current.loaded.bindings.state.was_interrupted)
+                else "ok"
+            )
+        )
+        await self._session.hook_manager.fire(
+            HookEvent.AFTER_TURN,
+            {
+                "session_id": self._session.session_id,
+                "profile": profile_name,
+                "cwd": _workspace_cwd(self._session),
+                "turn": self._session.turn_number,
+                "status": status,
+                "failed": failed,
+            },
+        )
+
+    def schedule_user_interrupt(self) -> None:
+        """Schedule observer dispatch with the interrupt's session identity."""
+        manager = self._session.hook_manager
+        if manager is None or not manager.has_hooks_for(HookEvent.USER_INTERRUPT):
+            return
+        session_id = self._session.session_id
+        profile_name = self._session.agent_profile.name if self._session.agent_profile is not None else ""
+        cwd = _workspace_cwd(self._session)
+        task = asyncio.create_task(
+            self.fire_user_interrupt(manager, session_id=session_id, profile_name=profile_name, cwd=cwd)
+        )
+        self._pending.add(task)
+        task.add_done_callback(self._pending.discard)
+        task.add_done_callback(self.log_user_interrupt_hook_error)
+
+    async def fire_user_interrupt(
+        self,
+        manager: HookManager,
+        *,
+        session_id: str | None,
+        profile_name: str,
+        cwd: str,
+    ) -> None:
+        """Publish an interrupt hook using only its captured manager and identity."""
+        await manager.fire(
+            HookEvent.USER_INTERRUPT,
+            {"session_id": session_id, "profile": profile_name, "cwd": cwd},
+            scope="detached",
+        )
+
+    @staticmethod
+    def log_user_interrupt_hook_error(task: asyncio.Task[None]) -> None:
+        """Log a failed background interrupt-hook dispatch."""
+        try:
+            task.result()
+        except asyncio.CancelledError:
+            pass
+        except Exception:
+            logger.exception("user_interrupt hook dispatch failed")
+
+    async def fire_pre_compact(self, info: PreCompactInfo) -> None:
+        """Callback from UnifiedContextStrategy — fire pre_compact hooks."""
+        if self._session.hook_manager is None:
+            return
+        from chrys.service.hooks.events import HookEvent
+
+        if not self._session.hook_manager.has_hooks_for(HookEvent.PRE_COMPACT):
+            return
+        profile_name = self._session.agent_profile.name if self._session.agent_profile is not None else ""
+        await self._session.hook_manager.fire(
+            HookEvent.PRE_COMPACT,
+            {
+                "session_id": self._session.session_id,
+                "profile": profile_name,
+                "cwd": self._session.workspace_cwd(),
+                "trigger": info.trigger,
+                "usage_pct": info.usage_pct,
+                "tokens_before": info.tokens_before,
+            },
+            target_operation_id=info.trajectory_operation_id,
+        )
+
+
+class PromptSubmitGate:
+    """Evaluate ``user_prompt_submit`` hooks and apply accepted reminders."""
+
+    def __init__(
+        self,
+        *,
+        session: ActiveSession,
+        current: CurrentAgent,
+        bus: EventBus,
+        fsm: EngineStateMachine,
+    ) -> None:
+        self._session = session
+        self._current = current
+        self._bus = bus
+        self._fsm = fsm
+
+    async def fire(self, text: str, *, injected: bool | None = None) -> bool:
+        """Return True when a blocking hook denied the prompt."""
+        decision = await self.evaluate(text, injected=injected)
+        if await self.handle_decision(decision, injected=injected):
+            return True
+        self.queue_reminders(decision, injected=injected)
+        return False
+
+    async def evaluate(
+        self,
+        text: str,
+        *,
+        injected: bool | None,
+        target_operation_id: str | None = None,
+        trajectory_context: TrajectoryContext | None = None,
+    ) -> HookDecision | None:
+        """Run ``user_prompt_submit`` hooks without applying reminder side effects."""
+        if self._session.hook_manager is None or not self._session.hook_manager.has_hooks_for(
+            HookEvent.USER_PROMPT_SUBMIT
+        ):
+            return None
+        profile_name = self._session.agent_profile.name if self._session.agent_profile is not None else ""
+        is_injected = self._fsm.is_running() if injected is None else injected
+        return await self._session.hook_manager.fire(
+            HookEvent.USER_PROMPT_SUBMIT,
+            {
+                "session_id": self._session.session_id,
+                "profile": profile_name,
+                "cwd": _workspace_cwd(self._session),
+                "text": text,
+                "injected": is_injected,
+            },
+            target_operation_id=target_operation_id,
+            trajectory_context=trajectory_context,
+        )
+
+    async def handle_decision(
+        self,
+        decision: HookDecision | None,
+        *,
+        injected: bool | None,
+        session_id: str | None = None,
+    ) -> bool:
+        """Return True after publishing when a prompt-submit hook blocked."""
+        _ = injected
+        if decision is None:
+            return False
+        if not decision.blocked:
+            return False
+        await self._bus.publish(
+            Error(
+                code="hook_blocked",
+                message=decision.block_reason or "Prompt blocked by hook.",
+                display_message=None if decision.block_reason else _TURN_HOOKS_PROMPT_BLOCKED.bind(),
+                session_id=self._session.session_id if session_id is None else session_id,
+            )
+        )
+        return True
+
+    def queue_reminders(self, decision: HookDecision | None, *, injected: bool | None) -> None:
+        """Apply non-blocking prompt-submit hook reminders after prompt validation passes."""
+        if decision is None or not decision.system_reminders or self._current.loaded is None:
+            return
+        is_injected = self._fsm.is_running() if injected is None else injected
+        self._current.loaded.reminder_middleware.queue_hook_reminders(
+            decision.system_reminders,
+            for_next_turn=not is_injected,
+        )
+
+    @staticmethod
+    def reminder_texts(decision: HookDecision | None) -> list[str]:
+        """Return non-empty prompt-hook reminders from an accepted decision."""
+        if decision is None:
+            return []
+        return [reminder for reminder in decision.system_reminders if reminder]
+
+
+def _workspace_cwd(session: ActiveSession) -> str:
+    """Return the current workspace cwd, falling back only before workspace initialization."""
+    if session.workspace is not None:
+        return session.workspace.primary_cwd
+    return safe_getcwd()

chrys/orchestration/engine/run/turn_state.py

@@ -0,0 +1,149 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Turn-specific recovery input and injection state, borrowing an ExecutionLease."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, Literal
+
+from chrys.foundation.models.turns import UserMessageKind
+from chrys.orchestration.engine.execution import ExecutionLease, RunTaskDrainOutcome
+
+if TYPE_CHECKING:
+    from chrys.orchestration.engine.execution import CurrentRunInjectionWindow, CurrentRunScope
+    from chrys.orchestration.engine.run.bindings import TurnBindings
+    from chrys.orchestration.invoker.kernel import KernelConversation
+    from chrys.service.agent_middleware.system_reminder import CurrentRunReminderTarget, SystemReminderMiddleware
+
+
+@dataclass
+class CurrentTurnInput:
+    """Current prompt data used by crash-recovery checkpoint writes.
+
+    ``kind`` records how a recovery checkpoint must re-create the message
+    (``"opener"`` unflagged, ``"injected"`` flagged mid-turn input) — see
+    :func:`chrys.foundation.models.turns.user_text_matches`. Empty-input
+    resumes clear this field; opener replay registers the popped anchor.
+    """
+
+    text: str = ""
+    contents: list[Any] | None = None
+    created_at: datetime | str | None = None
+    kind: UserMessageKind = "opener"
+
+
+@dataclass(frozen=True)
+class ActiveInjectionTarget:
+    """Operation-local target captured before awaited active-injection work."""
+
+    route: Literal["fsm_active", "executor_fallback"]
+    session_id: str | None
+    session_generation: int
+    build_generation: int
+    current_run_scope: CurrentRunScope
+    run_task: asyncio.Task[None]
+    bindings: TurnBindings
+    conversation: KernelConversation
+    reminder_middleware: SystemReminderMiddleware
+    reminder_target: CurrentRunReminderTarget
+    injection_window: CurrentRunInjectionWindow
+    trajectory_turn_id: str | None = None
+
+
+@dataclass
+class TurnRuntimeState:
+    """Turn product state; lease owns the unique task and all lifetime fences."""
+
+    lease: ExecutionLease = field(default_factory=ExecutionLease)
+    current_input: CurrentTurnInput = field(default_factory=CurrentTurnInput)
+
+    paused_sub_agents: set[str] = field(default_factory=set)
+    shutdown_used_cancel_fallback: bool = False
+
+    history_start_index: int = 0
+
+    cancelled_injection_ids: set[str] = field(default_factory=set)
+
+    inflight_injection_ids: dict[str, int] = field(default_factory=dict)
+
+    async def drain_for_boundary(self) -> RunTaskDrainOutcome:
+        """Observe the active run-task chain for a rebuild/session boundary."""
+        return await self.lease.observe_run_task_chain(propagate_inner_cancel=False)
+
+    def begin_inflight_injection(self, injection_id: str | None) -> None:
+        """Track one in-flight admission of a frontend-identified injection.
+
+        While an id is tracked, a user cancel for it is recorded as a mark
+        (:meth:`mark_injection_cancelled`) that the admission observes at its
+        cancellation checkpoints. No-op for id-less submits.
+        """
+        if not injection_id:
+            return
+        self.inflight_injection_ids[injection_id] = self.inflight_injection_ids.get(injection_id, 0) + 1
+
+    def finish_inflight_injection(self, injection_id: str | None) -> None:
+        """Stop tracking one in-flight admission for *injection_id*.
+
+        When the last admission for the id exits, any cancel mark it did not
+        observe (e.g. the admission aborted for an unrelated reason first) is
+        dropped with it — nothing can observe the mark afterwards.
+        """
+        if not injection_id:
+            return
+        count = self.inflight_injection_ids.get(injection_id, 0)
+        if count > 1:
+            self.inflight_injection_ids[injection_id] = count - 1
+            return
+        self.inflight_injection_ids.pop(injection_id, None)
+        self.cancelled_injection_ids.discard(injection_id)
+
+    def is_injection_inflight(self, injection_id: str) -> bool:
+        """Return whether an admission for *injection_id* is currently in flight."""
+        return injection_id in self.inflight_injection_ids
+
+    def mark_injection_cancelled(self, injection_id: str) -> None:
+        """Record a user cancel for an injection whose admission is in flight.
+
+        Callers must only mark ids that something can still observe (an
+        in-flight admission — see :meth:`is_injection_inflight`); the mark is
+        consumed at the admission's next cancellation checkpoint and dropped
+        at admission end otherwise, so it can never outlive its injection.
+        """
+        self.cancelled_injection_ids.add(injection_id)
+
+    def is_injection_cancelled(self, injection_id: str | None) -> bool:
+        """Return whether *injection_id* carries an unconsumed cancel mark."""
+        return injection_id is not None and injection_id in self.cancelled_injection_ids
+
+    def discard_injection_cancellation(self, injection_id: str | None) -> bool:
+        """Consume the cancel mark for *injection_id*; returns whether it was set."""
+        if injection_id is None or injection_id not in self.cancelled_injection_ids:
+            return False
+        self.cancelled_injection_ids.discard(injection_id)
+        return True
+
+    def set_current_input(
+        self,
+        text: str,
+        contents: list[Any] | None,
+        created_at: datetime | str | None,
+        kind: UserMessageKind = "opener",
+    ) -> None:
+        """Set current prompt data for recovery checkpointing."""
+        self.current_input = CurrentTurnInput(text=text, contents=contents, created_at=created_at, kind=kind)
+
+    def clear_current_input(self) -> None:
+        """Clear current prompt data after run finalization."""
+        self.current_input = CurrentTurnInput()
+
+    def reset_after_session_shutdown(self, *, prompt_admission_owner: str | None = None) -> CurrentRunScope | None:
+        """Reset Turn-only data after the execution owner has drained and reset."""
+        old_scope = self.lease.reset_after_session_shutdown(prompt_admission_owner=prompt_admission_owner)
+        self.current_input = CurrentTurnInput()
+        self.history_start_index = 0
+        self.cancelled_injection_ids.clear()
+        self.inflight_injection_ids.clear()
+        return old_scope