ags-cli 0.1.0__py3-none-any.whl

harness/orchestrator/engine.py

@@ -0,0 +1,821 @@
+"""Orchestration engine.
+
+Responsibilities:
+- Materialize a session's runs from a compiled policy (one Run row per graph node).
+- Execute the run graph wave-by-wave: independent nodes run concurrently; a node
+  receives the outputs of its dependencies as additional context.
+- For each node: assemble context (objective + retrieved memory + upstream
+  outputs), drive the adapter, persist every NormalizedEvent to the append-only
+  ``run_events`` ledger, publish to the live bus, and project terminal state onto
+  the Run row.
+- Tolerate partial failure: one node failing does not abort independent nodes.
+
+The same ``execute_session`` coroutine is used by the in-process asyncio dispatch
+worker (``harness/workers/inprocess.py``) and directly by tests/CLI (inline), keeping
+behaviour identical across paths.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import uuid
+from dataclasses import dataclass
+from datetime import UTC, datetime
+
+from harness_sdk.base import BaseAdapter, SessionContext
+from harness_sdk.events import NormalizedEvent
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from harness.adapters.registry import build_adapter
+from harness.eventbus.inprocess_bus import InProcessEventBus
+from harness.memory.context_budget import (
+    build_memory_context,
+    context_window,
+    estimate_tokens,
+)
+from harness.memory.service import MemoryService
+from harness.models.artifact import Artifact
+from harness.models.enums import (
+    AdapterKind,
+    ArtifactKind,
+    RunEventType,
+    RunStatus,
+    SessionStatus,
+)
+from harness.models.provider_config import ProviderConfig
+from harness.models.run import Run
+from harness.models.run_event import RunEvent
+from harness.models.session import Session
+from harness.models.workspace import Workspace
+from harness.observability.logging import get_logger
+from harness.observability.metrics import (
+    ACTIVE_RUNS,
+    ADAPTER_LATENCY,
+    RUN_FAILURES,
+    RUN_TOTAL,
+    TOKENS_TOTAL,
+)
+from harness.orchestrator.lifecycle import InMemoryCancellationRegistry
+from harness.orchestrator.policies import CompiledPolicy, RunNodeSpec
+from harness.orchestrator.run_graph import RunGraph
+from harness.services.workspaces import (
+    effective_host,
+    effective_project,
+    effective_write_mode,
+)
+from harness.settings import get_settings
+
+log = get_logger("orchestrator")
+
+
+@dataclass
+class NodeResult:
+    node_id: str
+    run_id: uuid.UUID
+    provider: str
+    status: RunStatus
+    output: str
+    error: str | None = None
+
+
+async def create_runs(
+    db: AsyncSession,
+    *,
+    session_id: uuid.UUID,
+    policy: CompiledPolicy,
+    provider_configs: dict[str, ProviderConfig],
+) -> dict[str, Run]:
+    """Create one queued Run per graph node for THIS turn.
+
+    A session is a conversation: the same provider may be invoked across multiple
+    turns (e.g. claude → agy → claude). Node ids (``claude-0``) are stable per
+    policy, so the idempotency key carries a per-turn token to keep each turn's runs
+    distinct while preserving the human-readable ``graph_node_id`` used for
+    within-turn dependency wiring."""
+    turn = uuid.uuid4().hex[:12]
+    runs: dict[str, Run] = {}
+    for node in policy.nodes:
+        pc = provider_configs.get(node.provider)
+        kind = pc.kind if pc else AdapterKind.mock
+        run = Run(
+            session_id=session_id,
+            provider_config_id=pc.id if pc else None,
+            adapter_kind=kind,
+            role=node.role,
+            status=RunStatus.queued,
+            graph_node_id=node.node_id,
+            idempotency_key=f"{session_id}:{node.node_id}:{turn}",
+        )
+        db.add(run)
+        runs[node.node_id] = run
+    await db.flush()
+    return runs
+
+
+@dataclass
+class _Collected:
+    """In-memory result of driving an adapter — produced WITHOUT touching the DB so
+    multiple nodes in a wave can run concurrently. Live events still stream to the
+    in-process event bus (shared within the process, safe under concurrency)."""
+
+    events: list[NormalizedEvent]
+    final_output: str
+    status: RunStatus
+    error: str | None
+    cost: dict
+    started_at: datetime
+    finished_at: datetime
+
+
+async def _collect_adapter(
+    *,
+    adapter: BaseAdapter,
+    run_id: str,
+    ctx: SessionContext,
+    message: str,
+    bus: InProcessEventBus | None,
+    canceller: InMemoryCancellationRegistry | None,
+    publish_user: bool = True,
+    publish_done: bool = True,
+    attachments_meta: list[dict] | None = None,
+) -> _Collected:
+    """Collect events from *adapter* for one attempt.
+
+    Args:
+        publish_user: When False, skip publishing the initial user-message event
+            to the bus.  Pass False on retry attempts where subscribers have
+            already received the user message from the first attempt.
+        publish_done: When False, skip calling ``bus.publish_done`` at the end.
+            Pass False when the caller will conditionally publish done itself
+            (e.g. only after a successful retry completes).
+    """
+    output_parts: list[str] = []
+    final_message = ""
+    status = RunStatus.succeeded
+    error: str | None = None
+    cost: dict = {}
+    events: list[NormalizedEvent] = []
+    started = datetime.now(UTC)
+    ACTIVE_RUNS.inc()
+    try:
+        from harness_sdk.events import MessageEvent
+        user_msg = MessageEvent(seq=0, role="user", text=message,
+                                attachments=attachments_meta or [])
+        events.append(user_msg)
+        if bus is not None and publish_user:
+            await bus.publish(
+                run_id,
+                {"type": user_msg.type, "seq": user_msg.seq,
+                 **user_msg.model_dump(exclude={"type", "seq"})},
+            )
+
+        provider_session = await adapter.start_session(ctx)
+        seq_counter = 1
+        # Per-run deadline: a non-responding adapter (stuck remote API, wedged CLI) is
+        # cancelled and the run fails cleanly instead of hanging the whole stack.
+        deadline = get_settings().run_timeout_s
+        stall = get_settings().stall_timeout_s
+        async with asyncio.timeout(deadline):
+            it = adapter.stream_events(provider_session, message).__aiter__()
+            _it_exhausted = False  # set True only on the normal StopAsyncIteration path
+            try:
+                while True:
+                    try:
+                        if stall > 0:
+                            event = await asyncio.wait_for(it.__anext__(), timeout=stall)
+                        else:
+                            event = await it.__anext__()
+                    except StopAsyncIteration:
+                        _it_exhausted = True
+                        break
+                    except TimeoutError:
+                        # Per-event stall detected: the adapter stopped yielding events.
+                        # This is caught here (inside asyncio.wait_for) so it cannot be
+                        # confused with the outer asyncio.timeout(deadline) TimeoutError.
+                        status = RunStatus.failed
+                        error = f"no output for {stall}s (stalled)"
+                        with contextlib.suppress(Exception):
+                            await adapter.cancel_run(run_id)
+                        break
+                    event.seq = seq_counter
+                    seq_counter += 1
+                    events.append(event)
+                    if bus is not None:
+                        await bus.publish(
+                            run_id,
+                            {"type": getattr(event.type, "value", event.type), "seq": event.seq,
+                             **event.model_dump(exclude={"type", "seq"}, mode="json")},
+                        )
+                    if event.type == "token":
+                        output_parts.append(event.text)
+                    elif event.type == "message":
+                        final_message = event.text
+                    elif event.type == "cost":
+                        k = adapter.kind.value
+                        TOKENS_TOTAL.labels(k, "prompt").inc(event.prompt_tokens)
+                        TOKENS_TOTAL.labels(k, "completion").inc(event.completion_tokens)
+                        cost = event.model_dump(exclude={"type", "seq"})
+                    elif event.type == "error":
+                        status = RunStatus.failed
+                        error = event.message
+                    elif event.type == "status" and event.status == "needs_approval":
+                        # A tool call hit the approval gate (see fs_tools.execute /
+                        # ApprovalRequired). The adapter already stopped emitting
+                        # events; the accompanying adapter_meta event carries the
+                        # pending request, materialized into an Approval row by
+                        # _create_pending_approval once this node is persisted.
+                        status = RunStatus.needs_approval
+                    if canceller is not None and await canceller.is_cancelled(run_id):
+                        status = RunStatus.cancelled
+                        await adapter.cancel_run(run_id)
+                        break
+            finally:
+                # Close the iterator whenever the loop exits abnormally (stall break,
+                # cancel break).  aclose() on an already-exhausted generator is a no-op,
+                # so this is safe on the happy path too.  Closing deterministically
+                # triggers the inner generator's finally/_reap instead of leaving
+                # it to GC — which in turn lets adapter backstops run synchronously.
+                if not _it_exhausted:
+                    with contextlib.suppress(Exception):
+                        await it.aclose()
+    except TimeoutError:  # adapter didn't respond within the per-run deadline
+        from harness_sdk.events import ErrorEvent
+        status = RunStatus.failed
+        error = f"provider '{adapter.config.name}' did not respond within {deadline}s"
+        RUN_FAILURES.labels(adapter.kind.value, "timeout").inc()
+        log.error("node_timeout", run_id=run_id, error=error)
+        with contextlib.suppress(Exception):  # best-effort cleanup of the stuck adapter
+            await adapter.cancel_run(run_id)
+        err = ErrorEvent(seq=len(events), message=error, retriable=True)
+        events.append(err)
+        if bus is not None:
+            await bus.publish(
+                run_id, {"type": err.type, "seq": err.seq,
+                         **err.model_dump(exclude={"type", "seq"})},
+            )
+    except Exception as exc:  # adapter blew up — contain it to this node
+        status = RunStatus.failed
+        error = f"{type(exc).__name__}: {exc}"
+        RUN_FAILURES.labels(adapter.kind.value, type(exc).__name__).inc()
+        log.error("node_failed", run_id=run_id, error=error)
+    finally:
+        ACTIVE_RUNS.dec()
+    if bus is not None and publish_done:
+        await bus.publish_done(run_id)
+    return _Collected(
+        events=events,
+        final_output=final_message or "".join(output_parts).strip(),
+        status=status,
+        error=error,
+        cost=cost,
+        started_at=started,
+        finished_at=datetime.now(UTC),
+    )
+
+
+def _persist_artifacts(
+    db: AsyncSession,
+    *,
+    run: Run,
+    workspace_id: uuid.UUID,
+    session_id: uuid.UUID,
+    c: _Collected,
+) -> None:
+    """Materialize Artifact rows. Provides shared transcript persistence: every run
+    gets a canonical transcript artifact (provider-independent), plus a row for any
+    ArtifactEvent the adapter emitted (files/patches/plans/docs)."""
+    import hashlib
+
+    for ev in c.events:
+        if ev.type == "artifact":
+            db.add(
+                Artifact(
+                    workspace_id=workspace_id, session_id=session_id, run_id=run.id,
+                    kind=ArtifactKind(ev.kind), uri=ev.uri,
+                    sha256=ev.sha256 or hashlib.sha256(ev.uri.encode()).hexdigest(),
+                    meta=ev.meta or {},
+                )
+            )
+    transcript = c.final_output or ""
+    if transcript:
+        db.add(
+            Artifact(
+                workspace_id=workspace_id, session_id=session_id, run_id=run.id,
+                kind=ArtifactKind.doc,
+                uri=f"harness://transcript/{run.id}",
+                sha256=hashlib.sha256(transcript.encode()).hexdigest(),
+                size=len(transcript.encode()),
+                meta={"kind": "transcript", "provider": run.adapter_kind.value,
+                      "content": transcript[:20_000]},
+            )
+        )
+
+
+def _persist_collected(
+    db: AsyncSession, *, run: Run, adapter_kind: AdapterKind, provider: str, c: _Collected
+) -> NodeResult:
+    """Write a collected node's events + terminal state onto the shared session.
+    Called sequentially (one node at a time) so the single AsyncSession is never
+    flushed concurrently."""
+    run.started_at = c.started_at
+    for ev in c.events:
+        db.add(
+            RunEvent(
+                run_id=run.id, seq=ev.seq, type=RunEventType(ev.type),
+                payload=ev.model_dump(exclude={"type", "seq"}),
+            )
+        )
+    run.final_output = c.final_output
+    run.status = c.status
+    run.error = c.error
+    run.cost = c.cost
+    run.finished_at = c.finished_at
+    ADAPTER_LATENCY.labels(adapter_kind.value, provider).observe(
+        (c.finished_at - c.started_at).total_seconds()
+    )
+    RUN_TOTAL.labels(adapter_kind.value, c.status.value).inc()
+    return NodeResult(
+        node_id=run.graph_node_id, run_id=run.id, provider=provider,
+        status=c.status, output=c.final_output, error=c.error,
+    )
+
+
+async def _create_pending_approval(
+    db: AsyncSession, *, run: Run, session_id: uuid.UUID, collected: _Collected
+) -> None:
+    """When a node paused for a human decision, materialize the Approval row
+    the webgui/CLI poll for, from the adapter_meta payload the adapter emitted
+    alongside its ``needs_approval`` status event."""
+    from harness.models.enums import ApprovalActionClass
+    from harness.services.approvals import create_approval
+
+    meta = next(
+        (
+            ev.data.get("approval_request") for ev in collected.events
+            if ev.type == "adapter_meta" and "approval_request" in ev.data
+        ),
+        None,
+    )
+    if meta is None:
+        return
+    action_class_raw = meta.get("action_class")
+    action_class = (
+        ApprovalActionClass(action_class_raw) if action_class_raw else ApprovalActionClass.command
+    )
+    await create_approval(
+        db, session_id=session_id, run_id=run.id, action_class=action_class, payload=meta,
+    )
+
+
+async def execute_node(
+    db: AsyncSession,
+    *,
+    adapter: BaseAdapter,
+    run: Run,
+    ctx: SessionContext,
+    message: str,
+    bus: InProcessEventBus | None = None,
+    canceller: InMemoryCancellationRegistry | None = None,
+) -> NodeResult:
+    """Drive a single adapter run to completion and persist it. Convenience wrapper
+    over collect+persist for the single-node path."""
+    run.status = RunStatus.running
+    await db.flush()
+    collected = await _collect_adapter(
+        adapter=adapter, run_id=str(run.id), ctx=ctx, message=message,
+        bus=bus, canceller=canceller,
+    )
+    result = _persist_collected(
+        db, run=run, adapter_kind=adapter.kind, provider=adapter.config.name, c=collected
+    )
+    if collected.status == RunStatus.needs_approval:
+        await _create_pending_approval(db, run=run, session_id=run.session_id, collected=collected)
+    await db.flush()
+    return result
+
+
+async def execute_session(
+    db: AsyncSession,
+    *,
+    session_id: uuid.UUID,
+    policy: CompiledPolicy,
+    bus: InProcessEventBus | None = None,
+    canceller: InMemoryCancellationRegistry | None = None,
+    message_override: str | None = None,
+    model_override: str | None = None,
+    model_overrides: dict[str, str] | None = None,
+    attachments: list[dict] | None = None,
+) -> list[NodeResult]:
+    """Execute the full run graph for a session, wave by wave.
+
+    Within a wave the slow adapter/LLM calls run concurrently (true parallel
+    provider execution), but their DB writes are applied sequentially on the
+    shared session — SQLAlchemy's AsyncSession is not safe for concurrent flush.
+
+    ``message_override`` supplies a per-turn prompt (e.g. a follow-up on resume);
+    when set it is the message sent to each node instead of the session's original
+    objective, while shared memory is still retrieved and injected as usual.
+
+    ``model_override`` applies one model to every node; ``model_overrides`` maps a
+    provider name to its own model and wins for the providers it names. Together
+    they let a fanout wave run each agent on a different model."""
+    import asyncio
+
+    sess = (await db.execute(select(Session).where(Session.id == session_id))).scalar_one()
+    workspace = (
+        await db.execute(select(Workspace).where(Workspace.id == sess.workspace_id))
+    ).scalar_one()
+
+    # Per-workspace runtime settings: which project folder the agents operate in and
+    # whether they may edit files. Set via `ags project set/mode`; fall back to the
+    # process cwd / HARNESS_WRITE default. Injected into every node's context below.
+    # A project_host means the folder lives on a remote ssh host — the adapter wraps
+    # its CLI invocation in ssh, and local-filesystem features are skipped.
+    project_dir, _ = effective_project(workspace.settings or {})
+    # Per-session override: a session may pin its own plan/write mode (set via
+    # `/mode` in chat or the GUI toggle); absent/None inherits the workspace value.
+    write_mode = effective_write_mode(sess.settings or {}, workspace.settings or {})
+    project_host = effective_host(workspace.settings or {})
+
+    # Private workspace: an isolated session runs in its own git worktree, not
+    # the shared project dir. The service owns all downgrade cases (remote host,
+    # non-git dir, git failure) and records the reason on the session.
+    if (sess.settings or {}).get("workspace_mode") == "worktree":
+        from harness.services.worktrees import ensure_worktree
+
+        provisioned = await ensure_worktree(
+            db, session=sess, workspace=workspace, project_dir=project_dir,
+            project_host=project_host,
+        )
+        if provisioned is not None:
+            project_dir = provisioned[0]
+
+    provider_configs = {
+        pc.name: pc
+        for pc in (
+            await db.execute(
+                select(ProviderConfig).where(ProviderConfig.enabled.is_(True))
+            )
+        ).scalars()
+    }
+    # Execute only THIS turn's runs: the ones still queued. Earlier turns are
+    # terminal and must not be re-run (the same provider can recur across turns).
+    all_runs = list((
+        await db.execute(
+            select(Run).where(Run.session_id == session_id).order_by(Run.created_at.asc())
+        )
+    ).scalars())
+    queued = [r for r in all_runs if r.status == RunStatus.queued]
+
+    # Provider of the most recent terminal run — drives native-resume vs handoff.
+    terminal = [r for r in all_runs if r.status in (RunStatus.succeeded, RunStatus.failed)]
+    last_provider: str | None = None
+    if terminal:
+        last_run = max(terminal, key=lambda r: r.created_at)
+        pc_by_id = {pc.id: pc.name for pc in provider_configs.values()}
+        last_provider = pc_by_id.get(last_run.provider_config_id) or last_run.adapter_kind.value
+
+    if queued:
+        runs = {r.graph_node_id: r for r in queued}
+    elif not all_runs:
+        runs = await create_runs(
+            db, session_id=session_id, policy=policy, provider_configs=provider_configs
+        )
+    else:
+        # No queued work for this turn — nothing to execute.
+        return []
+
+    sess.status = SessionStatus.running
+    await db.flush()
+
+    # The message for this turn: a per-turn override (resume follow-up) or the
+    # session's original objective.
+    turn_message = message_override or sess.canonical_objective
+
+    # Node-independent attachment metadata for the user MessageEvent (chips /
+    # thumbnails in the transcript) — computed once, reused for every node.
+    attachments_meta = [
+        {k: a[k] for k in ("id", "name", "media_type", "size", "image")}
+        for a in (attachments or [])
+    ] or None
+
+    # Retrieve shared memory once per session execution; injected into every node.
+    semantic_block = ""
+    try:
+        result = await MemoryService(db).retrieve(
+            workspace_slug=workspace.slug,
+            query=turn_message,
+            session_id=str(session_id),
+        )
+        semantic_block = result.as_context_block()
+    except Exception as exc:  # memory must never sink a run
+        log.warning("memory_retrieval_failed", error=str(exc))
+
+    # Build the conversational transcript blocks (the variable part of the context).
+    # They're assembled here once; each node then budgets them to ITS model's window
+    # (see build_memory_context in _collect_one) — so a small-window model trims while
+    # a 1M-window model keeps the full history.
+    from harness.models.run_event import RunEvent
+    session_summary = sess.summary or ""
+    transcript_blocks: list[str] = []
+    completed_runs = [r for r in all_runs if r.status in (RunStatus.succeeded, RunStatus.failed)]
+    if completed_runs:
+        run_ids = [r.id for r in completed_runs]
+        stmt = (
+            select(RunEvent)
+            .where(RunEvent.run_id.in_(run_ids))
+            .order_by(RunEvent.run_id, RunEvent.seq)
+        )
+        past_events = (await db.execute(stmt)).scalars().all()
+        provider_of = {r.id: r.adapter_kind.value for r in completed_runs}
+        for ev in past_events:
+            provider = provider_of.get(ev.run_id, "Assistant")
+            if ev.type.value == "message":
+                role = ev.payload.get("role", "assistant")
+                text = ev.payload.get("text", "")
+                if text:
+                    if role == "user":
+                        transcript_blocks.append(f"User: {text}")
+                    else:
+                        transcript_blocks.append(f"[{provider}]: {text}")
+            elif ev.type.value == "tool_call":
+                name = ev.payload.get("name", "")
+                args = ev.payload.get("arguments", {})
+                transcript_blocks.append(f"[{provider}] ran tool '{name}' with {args}")
+            elif ev.type.value == "tool_result":
+                out = ev.payload.get("output", "")
+                if len(out) > 500:
+                    out = out[:500] + "... (truncated)"
+                transcript_blocks.append(f"[Tool Result]: {out}")
+
+    # Load enabled MCP servers once; each node filters them to its workspace/provider
+    # and mode in-memory (no DB access on the concurrent fan-out path). Servers fail
+    # soft — a load error simply means no MCP tools this run.
+    from harness.mcp.registry import filter_servers, load_enabled_servers
+    mcp_all_servers: list[dict] = []
+    try:
+        mcp_all_servers = await load_enabled_servers(db)
+    except Exception as exc:  # MCP must never sink a run
+        log.warning("mcp_load_failed", error=str(exc))
+
+    graph = RunGraph.from_specs(policy.nodes)
+    node_by_id: dict[str, RunNodeSpec] = {n.node_id: n for n in policy.nodes}
+    results: dict[str, NodeResult] = {}
+
+    for wave in graph.waves():
+        # Phase 1: drive every ready node's adapter concurrently (no DB writes).
+        async def _collect_one(
+            node: RunNodeSpec,
+        ) -> tuple[RunNodeSpec, BaseAdapter | None, _Collected | None, dict, str | None]:
+            pc = provider_configs.get(node.provider)
+            if pc is None:
+                return node, None, None, {}, None
+            # Per-provider pick wins over the wave-wide one; without either the
+            # provider keeps its configured default.
+            node_model = (model_overrides or {}).get(node.provider) or model_override
+            if node_model:
+                from harness_sdk.base import AdapterConfig
+
+                from harness.adapters.registry import build_adapter_from_config
+                from harness.security.secrets import resolve_secret_ref
+                cfg = AdapterConfig(
+                    name=pc.name,
+                    kind=pc.kind,
+                    base_url=pc.base_url,
+                    model=node_model,
+                    api_key=resolve_secret_ref(pc.secret_ref),
+                    params=pc.params or {},
+                    capabilities=pc.capabilities or {},
+                )
+                adapter = build_adapter_from_config(cfg)
+            else:
+                adapter = build_adapter(pc)
+            upstream = "\n\n".join(
+                f"[{node_by_id[d].role.value}:{results[d].provider}]\n{results[d].output}"
+                for d in node.depends_on if d in results
+            )
+            run = runs[node.node_id]
+            message = turn_message
+            if upstream:
+                message += f"\n\n## Upstream provider outputs\n{upstream}"
+            # Attachments: images ride as native content blocks only where the
+            # adapter supports them (claude or a vision-capable provider, local
+            # project); everything else — documents, and images demoted for
+            # other providers — is referenced by absolute path for the agent to
+            # read with its own tools.
+            def _native_image(a: dict) -> bool:
+                vision = pc.kind == AdapterKind.claude_code or bool(
+                    (pc.capabilities or {}).get("vision")
+                )
+                return bool(a.get("image")) and vision and not project_host
+
+            image_atts = [a for a in (attachments or []) if _native_image(a)]
+            file_atts = [a for a in (attachments or []) if not _native_image(a)]
+            if file_atts:
+                listing = "\n".join(
+                    f"- {a['path']} ({a['name']}, {a['size']} bytes)" for a in file_atts
+                )
+                message += f"\n\n## Attached files (read with your own tools)\n{listing}"
+            # Budget the injected context to THIS node's model window: recent turns
+            # verbatim, older ones folded into the session summary. Big-window models
+            # (1M) keep the full history untouched.
+            params = pc.params or {}
+            from harness.orchestrator.native_resume import decide_resume
+            # For antigravity: probe the CLI binary once and thread the supported
+            # flags into decide_resume so native-resume auto-enables when --conversation
+            # is advertised. We use getattr so non-CLI adapters (BaseAdapter subclasses
+            # without ensure_probe) are skipped without touching the SDK.
+            # Probe errors must never sink the run — failures produce probed=None.
+            probed: set[str] | None = None
+            if pc.kind == AdapterKind.antigravity:
+                ensure = getattr(adapter, "ensure_probe", None)
+                if ensure is not None:
+                    try:
+                        probe_result = await ensure()
+                        probed = probe_result.get("flags")
+                    except Exception:
+                        probed = None
+            resume = decide_resume(
+                provider_name=pc.name, adapter_kind=pc.kind,
+                native_sessions=sess.native_sessions or {},
+                last_provider=last_provider, params=params,
+                probed_flags=probed,
+            )
+            window = context_window(pc.kind, node_model or pc.model, params)
+            node_transcript = [] if resume.native_id else transcript_blocks
+            # On a switch the handoff block carries the summary — don't inject it twice.
+            node_summary = "" if (resume.native_id or resume.is_switch) else session_summary
+            node_memory = build_memory_context(
+                semantic_block, node_transcript, node_summary,
+                window=window, message_tokens=estimate_tokens(message),
+                reserve_tokens=params.get("max_tokens"),
+            )
+            if resume.is_switch and last_provider:
+                from harness.orchestrator.handoff import build_handoff_block, files_touched_from_blocks
+                node_memory = build_handoff_block(
+                    from_provider=last_provider, summary=session_summary,
+                    recent_blocks=transcript_blocks,
+                    files_touched=files_touched_from_blocks(transcript_blocks),
+                ) + "\n\n" + node_memory
+            # MCP servers offered to THIS node (workspace+provider bind, mode-gated).
+            node_mcp_servers = filter_servers(
+                mcp_all_servers,
+                workspace_slug=workspace.slug,
+                provider=pc.name,
+                read_only=not write_mode,
+            )
+            # Managed-skills awareness for the CLI agents (claude/agy), injected into the
+            # system prompt. The CLIs natively load the skills; this tells them what's
+            # available. Skipped for non-CLI providers (no native Skill tool).
+            system_prompt = None
+            if pc.kind in (AdapterKind.claude_code, AdapterKind.antigravity):
+                from harness.skills import index_for
+                system_prompt = index_for(provider=pc.name, workspace=workspace.slug) or None
+            extra: dict = {}
+            if project_host:
+                extra["ssh_host"] = project_host
+            if node_mcp_servers:
+                extra["mcp_servers"] = node_mcp_servers
+            if resume.native_id:
+                if pc.kind == AdapterKind.antigravity:
+                    extra["conversation_id"] = resume.native_id
+                else:
+                    extra["native_session_id"] = resume.native_id
+            if image_atts:
+                extra["attachments_images"] = [
+                    {"path": a["path"], "media_type": a["media_type"], "name": a["name"]}
+                    for a in image_atts
+                ]
+            ctx = SessionContext(
+                run_id=str(run.id),
+                session_id=str(session_id),
+                workspace_slug=workspace.slug,
+                objective=turn_message,
+                constraints=sess.constraints or {},
+                system_prompt=system_prompt,
+                memory_context=node_memory,
+                workspace_dir=project_dir,
+                read_only=not write_mode,
+                trust_policy=workspace.trust_policy.value,
+                extra=extra,
+            )
+            # Build context metadata (safe: pure function, no DB write).
+            # Returned in the tuple and assigned to run.context_meta in the
+            # sequential Phase-2 loop below to avoid concurrent ORM writes.
+            from harness.orchestrator.native_resume import context_meta_for
+            ctx_meta = context_meta_for(
+                resume=resume,
+                node_memory=node_memory,
+                transcript_blocks=node_transcript,
+                window=window,
+            )
+            # Record the model this node actually ran on — with per-provider
+            # overrides the answer differs per run, and the UI/API have no other
+            # way to tell which one was used.
+            ctx_meta = {**ctx_meta, "model": node_model or pc.model}
+            # Caller is the sole publisher of done — always suppress inside _collect_adapter
+            # so that the else-branch below is the single publish_done site for every path.
+            snapshot_sha: str | None = None
+            # Remote projects are never snapshotted here — project_dir names a dir on
+            # the remote host; a same-named local dir would be the wrong repo.
+            if not ctx.read_only and project_dir and not project_host:
+                from harness.orchestrator.snapshot import git_snapshot
+                snapshot_sha = await git_snapshot(project_dir)
+            collected = await _collect_adapter(
+                adapter=adapter, run_id=str(run.id), ctx=ctx, message=message,
+                bus=bus, canceller=canceller,
+                publish_done=False,
+                attachments_meta=attachments_meta,
+            )
+            if resume.native_id and collected.status == RunStatus.failed:
+                log.warning("native_resume_failed_retrying", provider=pc.name,
+                            native_id=resume.native_id, error=collected.error)
+                extra.pop("native_session_id", None)
+                extra.pop("conversation_id", None)
+                retry_memory = build_memory_context(
+                    semantic_block, transcript_blocks, session_summary,
+                    window=window, message_tokens=estimate_tokens(message),
+                    reserve_tokens=params.get("max_tokens"),
+                )
+                ctx = ctx.model_copy(update={
+                    "memory_context": retry_memory,
+                    "extra": extra,
+                })
+                # Recompute ctx_meta for the retry: native-resume is abandoned so mode
+                # must be "replay" with the full transcript and correct memory_chars.
+                from harness.orchestrator.native_resume import ResumeDecision
+                ctx_meta = context_meta_for(
+                    resume=ResumeDecision(native_id=None, is_switch=False),
+                    node_memory=retry_memory,
+                    transcript_blocks=transcript_blocks,
+                    window=window,
+                )
+                ctx_meta = {**ctx_meta, "model": node_model or pc.model}
+                # Retry: subscriber already saw the user message; publish_done=True (default)
+                # so the stream terminates after this final attempt.
+                collected = await _collect_adapter(
+                    adapter=adapter, run_id=str(run.id), ctx=ctx, message=message,
+                    bus=bus, canceller=canceller,
+                    publish_user=False,
+                    attachments_meta=attachments_meta,
+                )
+            else:
+                # First attempt succeeded (or no retry possible): emit done now.
+                if bus is not None:
+                    await bus.publish_done(str(run.id))
+            return node, adapter, collected, ctx_meta, snapshot_sha
+
+        collected_wave = await asyncio.gather(*[_collect_one(n) for n in wave])
+
+        # Phase 2: persist sequentially on the shared session, then flush once.
+        for node, adapter, collected, ctx_meta, snapshot_sha in collected_wave:
+            run = runs[node.node_id]
+            if adapter is None or collected is None:
+                run.status = RunStatus.failed
+                run.error = f"provider '{node.provider}' not enabled/configured"
+                run.finished_at = datetime.now(UTC)
+                results[node.node_id] = NodeResult(
+                    node.node_id, run.id, node.provider, RunStatus.failed, "", run.error
+                )
+                continue
+            # Assign context_meta here (sequential) — safe for the shared session.
+            run.context_meta = ctx_meta
+            results[node.node_id] = _persist_collected(
+                db, run=run, adapter_kind=adapter.kind, provider=adapter.config.name, c=collected
+            )
+            if collected.status == RunStatus.needs_approval:
+                await _create_pending_approval(
+                    db, run=run, session_id=session_id, collected=collected
+                )
+            _persist_artifacts(
+                db, run=run, workspace_id=workspace.id, session_id=session_id, c=collected
+            )
+            if snapshot_sha:
+                db.add(
+                    Artifact(
+                        workspace_id=workspace.id, session_id=session_id, run_id=run.id,
+                        kind=ArtifactKind.patch,
+                        uri=f"git-snapshot://{snapshot_sha}",
+                        sha256=snapshot_sha,
+                        size=0,
+                        meta={"git_snapshot": snapshot_sha, "pre_run": True},
+                    )
+                )
+            from harness.orchestrator.native_resume import extract_native_id
+            nsid = extract_native_id(collected.events or [], adapter.kind)
+            if nsid:
+                sess.native_sessions = {**(sess.native_sessions or {}), adapter.config.name: nsid}
+        await db.flush()
+
+    any_ok = any(r.status == RunStatus.succeeded for r in results.values())
+    any_parked = any(r.status == RunStatus.needs_approval for r in results.values())
+    if any_parked:
+        # Waiting on a human decision — not failed, not done. Stays "running" so
+        # reconcile.py doesn't sweep it up as orphaned/interrupted on restart.
+        sess.status = SessionStatus.running
+    else:
+        sess.status = SessionStatus.completed if any_ok else SessionStatus.failed
+    await db.flush()
+    return list(results.values())