iCode-TUI 0.29.1__py3-none-any.whl

chrys/orchestration/workflows/hooks.py

@@ -0,0 +1,176 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Workflow hook lifetimes: session attachments outlive individual run resources.
+
+Browsing history is passive. First admitted execution attaches a session (start
+for a new session, restored for an existing one); later runs may reload hook
+configuration without synthesizing another session lifetime. Host shutdown or
+explicit session release ends the attachment. Run hooks are observers, once per
+admitted run, including cancellation and failure.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from chrys.foundation.trajectory.context import trajectory_scope
+from chrys.foundation.trajectory.event_types import RuntimeFinishReason
+from chrys.foundation.util.once_close import finish_close
+from chrys.service.hooks.events import HookEvent
+from chrys.service.trajectory.session import SessionStartInfo, SessionTrajectory
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from chrys.orchestration.workflows.session import WorkflowSessionOwner
+    from chrys.service.hooks.manager import HookManager
+    from chrys.service.state.locks import ActiveSessionGuard
+
+
+@dataclass
+class _Attachment:
+    manager: HookManager | None
+    payload: dict[str, Any]
+    recorder: SessionTrajectory
+    guard: ActiveSessionGuard
+    write_lock_path: Path | None
+    recovery: asyncio.Task[int] | None = None
+
+    def new_recorder(self) -> SessionTrajectory:
+        return SessionTrajectory(
+            session_id=self.recorder.session_id,
+            session_dir=self.recorder.session_dir,
+            write_lock_path=self.write_lock_path,
+            session_start_info=lambda: SessionStartInfo(self.payload["cwd"], "", ""),
+        )
+
+    async def settle_recovery(self, *, cancel: bool = False) -> None:
+        task = self.recovery
+        if task is None:
+            return
+        if cancel and not task.done():
+            task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            if not cancel:
+                raise
+        except Exception:
+            logger.exception("Workflow hook outbox recovery failed")
+        finally:
+            self.recovery = None
+
+
+class WorkflowSessionHooks:
+    def __init__(self) -> None:
+        self._sessions: dict[str, _Attachment] = {}
+
+    async def attach(self, resources: WorkflowSessionOwner) -> None:
+        session = resources.session
+        if not session.session_id or not session.workspace or (not resources.identity):
+            raise RuntimeError("Attaching workflow hooks requires a session id, workspace and workflow identity.")
+        manager = session.hook_manager
+        recorder = resources.trajectory
+        if recorder is None:
+            raise RuntimeError("Attaching workflow hooks requires a trajectory recorder.")
+        context = recorder.context()
+        if manager is not None:
+            manager.trajectory_context_provider = lambda: context
+        payload = {
+            "session_id": session.session_id,
+            "cwd": session.workspace.primary_cwd,
+            "profile": "",
+            "session_kind": "workflow",
+            "workflow_id": resources.identity.workflow_id,
+        }
+        previous = self._sessions.get(session.session_id)
+        # Transfer ownership before the first suspension, so cancellation during
+        # attachment still leaves a valid owner for run-end and session-end.
+        attachment = _Attachment(
+            manager, payload, recorder, session.guard, session.session_write_lock_path(session.session_id)
+        )
+        self._sessions[session.session_id] = attachment
+        resources.hooks_transferred = True
+        if previous is not None:
+            # Recovery belongs to the attachment, not the per-run manager.
+            # Finish dispatch before retiring the manager that owns it.
+            try:
+                await previous.settle_recovery()
+            finally:
+                if previous.manager is not None:
+                    await finish_close(asyncio.create_task(previous.manager.drain_session()))
+        if manager is not None and (previous is None or previous.manager is None):
+            with trajectory_scope(context):
+                attachment.recovery = asyncio.create_task(manager.recover_outbox())
+        if previous is not None:
+            return
+        if manager is not None:
+            event = HookEvent.SESSION_RESTORED if resources.restoring else HookEvent.SESSION_START
+            if resources.restoring:
+                payload = {**payload, "restored_session_id": session.session_id}
+            with trajectory_scope(context):
+                await manager.fire(event, payload, scope="session")
+
+    async def run_event(self, session_id: str, event: HookEvent, *, run_id: str, **fields: Any) -> None:
+        attachment = self._sessions[session_id]
+        if attachment.manager is not None:
+            with trajectory_scope(attachment.recorder.context(run_operation_id=run_id)):
+                await attachment.manager.fire(
+                    event, {**attachment.payload, "run_id": run_id, **fields}, scope="session"
+                )
+                await attachment.manager.drain_session(close=False)
+
+    async def release(self, session_id: str) -> None:
+        attachment = self._sessions.pop(session_id, None)
+        if attachment is None or attachment.manager is None:
+            return
+        await finish_close(asyncio.create_task(self._end_session(attachment)))
+
+    async def _end_session(self, attachment: _Attachment) -> None:
+        """Idle end hooks borrow the same session guard before opening a fresh writer."""
+        manager = attachment.manager
+        if manager is None:
+            raise RuntimeError("Ending an attached hook session requires its hook manager.")
+        session_id = attachment.recorder.session_id
+        recorder = None
+        acquired = False
+        manager.trajectory_context_provider = None
+        try:
+            await attachment.settle_recovery(cancel=True)
+            if attachment.recorder.is_closed:
+                try:
+                    lock = await asyncio.to_thread(attachment.guard.acquire_for_restore, session_id, timeout=0)
+                except TimeoutError:
+                    logger.warning("Session-end hook trajectory unavailable: workflow session %s is in use", session_id)
+                else:
+                    attachment.guard.install(session_id, lock)
+                    acquired = True
+                    directory = attachment.recorder.session_dir
+                    # SessionDeleted also releases attachments. Never recreate
+                    # the deleted directory just to record its end hooks.
+                    if directory is not None and directory.is_dir():
+                        recorder = attachment.new_recorder()
+            else:
+                recorder = attachment.recorder
+            context = recorder.context() if recorder is not None else None
+            with trajectory_scope(context):
+                await manager.fire(HookEvent.SESSION_END, attachment.payload, scope="session")
+        finally:
+            try:
+                await manager.drain_session()
+            finally:
+                try:
+                    if acquired and recorder is not None:
+                        await recorder.close(reason=RuntimeFinishReason.GRACEFUL_SHUTDOWN)
+                finally:
+                    if acquired:
+                        attachment.guard.release()
+
+    async def close(self) -> None:
+        for session_id in tuple(self._sessions):
+            await self.release(session_id)

chrys/orchestration/workflows/preview.py

@@ -0,0 +1,322 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""Loading a workflow file in a worker before a run: environment, manifest and digests.
+
+A preview is the same two steps the run itself starts with, environment
+preparation and a worker load, done on a throwaway worker. What comes back
+is the data-only manifest plus the digests the confirmation ledger pins, so
+the CLI can show and confirm exactly what a run would execute.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Final
+
+from chrys.foundation.events.types import WorkflowRunRequest
+from chrys.foundation.models.workflow_session import WorkflowPins, WorkflowTarget
+from chrys.foundation.platform.files import surrogate_safe_text
+from chrys.foundation.util.once_close import finish_close
+from chrys.orchestration.workflows.worker_client import (
+    LOAD_TIMED_OUT,
+    AskHandler,
+    CapturedOutput,
+    EmitHandler,
+    LoadDiagnostic,
+    LoadResult,
+    WorkerLostError,
+    WorkerRpcError,
+    WorkerStartError,
+    WorkflowWorkerClient,
+    load_diagnostics,
+)
+from chrys.service.workflows.admission import spec_digest
+from chrys.service.workflows.discovery import WORKFLOWS_DIR_NAME, WorkflowSource
+from chrys.service.workflows.environment import (
+    EnvironmentPlan,
+    PreparedEnvironment,
+    WorkflowEnvironmentError,
+    WorkflowEnvironmentManager,
+    parse_environment_request,
+    plan_environment,
+)
+from chrys.service.workflows.graph import ManifestWarning, manifest_warnings
+from chrys.service.workflows.ledger import LedgerEntry
+from chrys.service.workflows.sdk_artifact import SdkArtifact, materialize_sdk_artifact
+
+PREVIEW_ENVIRONMENT_ERROR: Final = "environment_error"
+PREVIEW_WORKER_START_FAILED: Final = "worker_start_failed"
+PREVIEW_LOAD_FAILED: Final = "load_failed"
+PREVIEW_WORKER_LOST: Final = "worker_lost"
+PREVIEW_WORKING_DIR_MISSING: Final = "working_dir_missing"
+REJECT_SPEC_CHANGED: Final = "spec_changed"
+REJECT_NOT_CONFIRMED: Final = "not_confirmed"
+
+SDK_ARTIFACT_DIR_NAME: Final = "sdk"
+BYTECODE_CACHE_DIR_NAME: Final = ".pycache"
+
+
+class WorkflowPreviewError(Exception):
+    """The workflow cannot be loaded; ``code`` is the deterministic rejection reason.
+
+    A ``diagnose`` load also says where it failed (``diagnostics``, whether
+    any were left out) and whether it ran out of time (``timed_out``).
+    """
+
+    def __init__(
+        self,
+        code: str,
+        message: str,
+        *,
+        traceback: str = "",
+        stdout: CapturedOutput | None = None,
+        diagnostics: tuple[LoadDiagnostic, ...] = (),
+        diagnostics_truncated: bool = False,
+        timed_out: bool = False,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.traceback = traceback
+        self.stdout = stdout or CapturedOutput("", False)
+        self.diagnostics = diagnostics
+        self.diagnostics_truncated = diagnostics_truncated
+        self.timed_out = timed_out
+
+
+@dataclass(frozen=True, slots=True)
+class WorkflowInspection:
+    """Source bytes and a lexical environment plan; creating this never runs user code."""
+
+    source: WorkflowSource
+    environment: EnvironmentPlan
+    prepared_environment: PreparedEnvironment | None = None
+
+    @classmethod
+    def read(cls, source: WorkflowSource) -> WorkflowInspection:
+        try:
+            request = parse_environment_request(source.source)
+            return cls(source, plan_environment(request, entry_path=Path(source.canonical_path)))
+        except WorkflowEnvironmentError as exc:
+            raise WorkflowPreviewError(PREVIEW_ENVIRONMENT_ERROR, _located(exc)) from exc
+
+
+class WorkflowTrustDeclined(Exception):
+    """The user declined execution of the inspected source; no worker was launched."""
+
+
+@dataclass(frozen=True, slots=True)
+class LoadedWorkflow:
+    """A worker that has executed the workflow file; the caller owns ``client``."""
+
+    client: WorkflowWorkerClient
+    load: LoadResult
+    manifest: dict[str, Any]
+    spec_digest: str
+
+
+@dataclass(frozen=True, slots=True)
+class WorkflowPreview:
+    """What a run of ``source`` would execute, as reported by a throwaway worker."""
+
+    source: WorkflowSource
+    environment: PreparedEnvironment
+    load: LoadResult
+    manifest: dict[str, Any]
+    spec_digest: str
+
+    @property
+    def title(self) -> str:
+        title = self.manifest.get("title")
+        return title if isinstance(title, str) else ""
+
+    @property
+    def warnings(self) -> tuple[ManifestWarning, ...]:
+        return manifest_warnings(self.manifest)
+
+    def ledger_entry(self) -> LedgerEntry:
+        return ledger_entry_for(self.source, self.load, self.spec_digest, self.environment, title=self.title)
+
+
+@dataclass(frozen=True, slots=True)
+class PreparedWorkflow:
+    """A transient execution target and the preview whose pins it submits."""
+
+    target: WorkflowTarget
+    preview: WorkflowPreview
+
+    def run_request(self, *, input_text: str, request_id: str, timeout: float = 0.0) -> WorkflowRunRequest:
+        preview = self.preview
+        return WorkflowRunRequest(
+            target=self.target,
+            pins=WorkflowPins(
+                preview.source.identity, preview.spec_digest, preview.environment.environment_fingerprint
+            ),
+            input_text=input_text,
+            request_id=request_id,
+            timeout=timeout,
+        )
+
+
+def ledger_entry_for(
+    source: WorkflowSource, load: LoadResult, digest: str, environment: PreparedEnvironment, *, title: str
+) -> LedgerEntry:
+    return LedgerEntry(
+        title=title,
+        canonical_path=source.canonical_path,
+        source_kind=source.source_kind,
+        workflow_id=source.workflow_id,
+        entry_digest=source.source_digest,
+        manifest_digest=load.manifest_digest,
+        schema_version=load.manifest["schema_version"],
+        spec_digest=digest,
+        environment_fingerprint=environment.environment_fingerprint,
+    )
+
+
+def sdk_artifact_dir(config_dir: Path) -> Path:
+    """Where the injected SDK lives; discovery reserves the name ``sdk`` in the global directory, so it never collides."""
+    return config_dir / WORKFLOWS_DIR_NAME / SDK_ARTIFACT_DIR_NAME
+
+
+def worker_bytecode_cache_dir(config_dir: Path) -> Path:
+    """The private bytecode cache of every workflow worker; discovery ignores hidden names, so it never collides."""
+    return config_dir / WORKFLOWS_DIR_NAME / BYTECODE_CACHE_DIR_NAME
+
+
+async def materialize_runtime_sdk(config_dir: Path) -> SdkArtifact:
+    return await asyncio.to_thread(materialize_sdk_artifact, sdk_artifact_dir(config_dir))
+
+
+async def prepare_workflow_environment(source: WorkflowSource, *, sdk: SdkArtifact) -> PreparedEnvironment:
+    """Parse the file's environment request, choose the interpreter, and probe it."""
+    try:
+        request = parse_environment_request(source.source)
+        plan = plan_environment(request, entry_path=Path(source.canonical_path))
+        return await WorkflowEnvironmentManager(sdk_digest=sdk.digest).prepare(plan)
+    except WorkflowEnvironmentError as exc:
+        raise WorkflowPreviewError(PREVIEW_ENVIRONMENT_ERROR, _located(exc)) from exc
+
+
+async def load_workflow(
+    source: WorkflowSource,
+    *,
+    environment: PreparedEnvironment,
+    sdk: SdkArtifact,
+    workspace: Path,
+    bytecode_cache: Path,
+    ask_handler: AskHandler | None = None,
+    emit_handler: EmitHandler | None = None,
+    diagnose: bool = False,
+) -> LoadedWorkflow:
+    """Start a fresh worker and execute the file in it; on any failure the worker is closed before raising.
+
+    *diagnose* (validation) also compiles a folder's other Python files first
+    and places a failure in the error's ``diagnostics``.
+    """
+    _require_workspace(workspace)
+    try:
+        client = await WorkflowWorkerClient.launch(
+            environment=environment,
+            sdk=sdk,
+            workspace=workspace,
+            bytecode_cache=bytecode_cache,
+            ask_handler=ask_handler,
+            emit_handler=emit_handler,
+        )
+    except WorkerStartError as exc:
+        raise WorkflowPreviewError(PREVIEW_WORKER_START_FAILED, str(exc)) from exc
+    package_dir = source.package.directory if source.package is not None else None
+    precompile = source.package.python_files(Path(source.canonical_path).name) if source.package is not None else ()
+    try:
+        load = await client.load(
+            source.source,
+            filename=source.canonical_path,
+            workspace=workspace,
+            package_dir=package_dir,
+            diagnose=diagnose,
+            precompile=precompile if diagnose else (),
+        )
+        manifest = load.manifest
+        digest = spec_digest(source.source_digest, load.manifest_digest, manifest["schema_version"])
+    except WorkerRpcError as exc:
+        await _close_worker(client)
+        diagnostics, truncated = load_diagnostics(exc.data)
+        raise WorkflowPreviewError(
+            PREVIEW_LOAD_FAILED,
+            f"{exc.code}: {exc.message}",
+            traceback=exc.traceback,
+            stdout=exc.stdout,
+            diagnostics=diagnostics,
+            diagnostics_truncated=truncated,
+            timed_out=exc.data.get("reason") == LOAD_TIMED_OUT,
+        ) from exc
+    except WorkerLostError as exc:
+        await _close_worker(client)
+        raise WorkflowPreviewError(PREVIEW_WORKER_LOST, str(exc)) from exc
+    except BaseException:
+        await _close_worker(client)
+        raise
+    return LoadedWorkflow(client=client, load=load, manifest=manifest, spec_digest=digest)
+
+
+async def preview_workflow(
+    source: WorkflowSource,
+    *,
+    sdk: SdkArtifact,
+    workspace: Path,
+    bytecode_cache: Path,
+    on_environment_ready: Callable[[PreparedEnvironment], Awaitable[None]] | None = None,
+    diagnose: bool = False,
+) -> WorkflowPreview:
+    """Prepare the environment and load the file on a worker that is closed before returning."""
+    _require_workspace(workspace)
+    environment = await prepare_workflow_environment(source, sdk=sdk)
+    if on_environment_ready is not None:
+        await on_environment_ready(environment)
+    loaded = await load_workflow(
+        source,
+        environment=environment,
+        sdk=sdk,
+        workspace=workspace,
+        bytecode_cache=bytecode_cache,
+        diagnose=diagnose,
+    )
+    await _close_worker(loaded.client)
+    return WorkflowPreview(
+        source=source,
+        environment=environment,
+        load=loaded.load,
+        manifest=loaded.manifest,
+        spec_digest=loaded.spec_digest,
+    )
+
+
+def _require_workspace(workspace: Path) -> None:
+    """Refuse before probing or spawning anything when the worker's cwd is gone."""
+    if not workspace.is_dir():
+        raise WorkflowPreviewError(
+            PREVIEW_WORKING_DIR_MISSING,
+            f"The working directory no longer exists: {surrogate_safe_text(os.fspath(workspace))}",
+        )
+
+
+def _located(exc: WorkflowEnvironmentError) -> str:
+    """The message with its place in the file, for surfaces that show only the message."""
+    if exc.line is None:
+        return str(exc)
+    place = f"line {exc.line}" if exc.column is None else f"line {exc.line}, column {exc.column}"
+    return f"{exc} ({place})"
+
+
+async def _close_worker(client: WorkflowWorkerClient) -> None:
+    try:
+        await client.close()
+    except asyncio.CancelledError:
+        # close owns a shielded task; drain it even if the caller is cancelled repeatedly.
+        await finish_close(asyncio.create_task(client.close()))
+        raise

chrys/orchestration/workflows/rollback.py

@@ -0,0 +1,151 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
+
+"""File-only Workflow rollback under the caller's execution and session fences."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+from dataclasses import asdict
+from pathlib import Path
+from time import time
+from typing import TYPE_CHECKING, Any
+
+from chrys.foundation.events.types import WorkflowRollbackRequest, WorkflowRollbackResult
+from chrys.foundation.models.mutation_scope import WorkflowRunScope
+from chrys.foundation.util.once_close import finish_close
+from chrys.service.mutations.coordination import ATTRIBUTION_DIR_NAME, MutationCoordinator
+from chrys.service.mutations.store import SnapshotPolicy, SnapshotStore
+from chrys.service.mutations.tracker import MutationTracker
+from chrys.service.mutations.types import FileMutation, MutationOp, MutationProvenance, MutationSource
+
+if TYPE_CHECKING:
+    from chrys.foundation.config.settings import Settings
+    from chrys.orchestration.workflows.session import WorkflowSessionOwner
+
+
+def _file_revision(path: str) -> tuple[int, ...] | None:
+    """Fence editor writes and endpoint replacements without reading large files."""
+    try:
+        entry = Path(path).lstat()
+    except FileNotFoundError, NotADirectoryError:
+        return None
+    return (entry.st_dev, entry.st_ino, entry.st_mode, entry.st_size, entry.st_mtime_ns, entry.st_ctime_ns)
+
+
+async def rollback_files(
+    owner: WorkflowSessionOwner, request: WorkflowRollbackRequest, *, settings: Settings
+) -> WorkflowRollbackResult:
+    """Reload and recheck the plan on both preview and commit; keep Run history intact."""
+    state = owner.require_state()
+    directory, workspace = owner.require_session_dir(), owner.require_workspace()
+    if state.mutations is None:
+        raise ValueError("This session has no recorded file changes.")
+    mutations = state.mutations
+    coordinator = MutationCoordinator(
+        registry_root=directory.parent / ATTRIBUTION_DIR_NAME, session_id=request.session_id
+    )
+
+    def apply() -> tuple[WorkflowRollbackResult, dict[str, Any] | None]:
+        snapshots = SnapshotStore(directory, policy=SnapshotPolicy.from_settings(settings))
+        tracker = MutationTracker.deserialize(mutations, snapshots)
+        cwd = workspace.primary_cwd
+        scope_paths = [cwd, *(directory.path for directory in workspace.working_dirs)]
+        reclassified = coordinator.reclassify(tracker, force=True, fallback_root=cwd)
+        periods = tracker.get_all_periods()
+        target = tracker.get_period(WorkflowRunScope(request.run_id))
+        scopes = {period.scope for period in periods if period.period_index >= target.period_index}
+        plan = coordinator.augment_rollback_plan(
+            tracker, tracker.get_rollback_plan_for_periods(scopes), scope_paths=scope_paths, fallback_root=cwd
+        )
+        revisions = {path: _file_revision(path) for path, _ in plan.entries}
+        token = hashlib.sha256(
+            json.dumps(
+                {
+                    "ledger": tracker.serialize(),
+                    "plan": asdict(plan),
+                    "run": request.run_id,
+                    "files": revisions,
+                },
+                sort_keys=True,
+                default=str,
+            ).encode()
+        ).hexdigest()
+        result = WorkflowRollbackResult(
+            session_id=request.session_id,
+            request_id=request.request_id,
+            run_id=request.run_id,
+            token=token,
+            paths=tuple(path for path, _ in plan.entries),
+            exclusions=tuple((path, reason.value) for path, reason in plan.exclusions),
+            warnings=tuple(warning.message for warning in plan.warnings),
+        )
+        if request.token:
+            if request.token != token:
+                raise ValueError("The rollback plan changed. Review the updated plan before confirming.")
+            # An open window makes concurrent peers aware of the explicit user operation.
+            window = coordinator.open_window(scope_paths, fallback_root=cwd)
+            try:
+                started = time()
+                before = {path: snapshots.probe(path) for path, _ in plan.entries}
+                if any(_file_revision(path) != revision for path, revision in revisions.items()):
+                    raise ValueError(
+                        "Files changed while preparing rollback. Review the updated plan before confirming."
+                    )
+                restored = tracker.restore_files(plan)
+                finished = time()
+                restored_paths = {item.path for item in restored if item.changed}
+                coordinator.publish_claims(
+                    [
+                        FileMutation(
+                            path=path,
+                            operation=(
+                                MutationOp.DELETE
+                                if not snapshot.existed
+                                else MutationOp.MODIFY
+                                if before[path].existed
+                                else MutationOp.CREATE
+                            ),
+                            source=MutationSource.IMPLICIT,
+                            provenance=MutationProvenance.PROVEN,
+                            tool_call_id=request.request_id,
+                            timestamp=finished,
+                            t_start=started,
+                            t_end=finished,
+                            before_hash=before[path].content_hash,
+                            before_skip=before[path].skip_reason,
+                            after_hash=snapshot.content_hash,
+                            after_skip=snapshot.skip_reason,
+                        )
+                        for path, snapshot in plan.entries
+                        if path in restored_paths
+                    ],
+                    fallback_root=cwd,
+                )
+            finally:
+                coordinator.close_window(window)
+            result.applied = True
+            result.changed = sum(item.changed for item in restored)
+            result.error = "\n".join(f"{item.path}: {item.reason}" for item in restored if not item.ok)
+        return result, tracker.serialize() if reclassified or result.applied else None
+
+    result: WorkflowRollbackResult | None = None
+
+    async def complete() -> None:
+        nonlocal result
+        result, updated = await asyncio.to_thread(apply)
+        if updated is not None:
+            try:
+                await owner.save_mutations(updated)
+            except OSError as exc:
+                result.error = "\n".join(part for part in (result.error, f"Session checkpoint failed: {exc}") if part)
+
+    try:
+        # Drain both the file writes and checkpoint before releasing the lock.
+        await finish_close(asyncio.create_task(complete()))
+        if result is None:
+            raise RuntimeError("Workflow rollback completed without a result.")
+        return result
+    finally:
+        await asyncio.to_thread(coordinator.close)