okstra 0.167.0 → 0.169.0

package/runtime/python/okstra_ctl/dispatch_core.py

@@ -16,19 +16,22 @@ from .dispatch_state import (
     BACKEND_CMUX_PANE,
     BACKEND_MIXED,
     BACKEND_TMUX_PANE,
+    append_worker_dispatch as _append_worker_dispatch,
     DispatchError,
     WorkerJob,
     dispatch_mode as _dispatch_mode,
-    dispatch_record_matches,
     LIVENESS_AUDIT_HEARTBEAT,
     LIVENESS_WRAPPER_STATUS,
     load_json_object as _load_json_object,
+    link_agent_dispatch_result as _link_agent_dispatch_result,
     missing_completion_paths as _missing_completion_paths,
+    mutate_team_state as _mutate_team_state,
     require_string as _require_string,
     resolve_project_path as _resolve_project_path,
     resolve_required_path as _resolve_required_path,
     record_dispatch_facts as _record_dispatch_facts,
     transition_worker_status as _transition_worker_status,
+    update_worker_dispatch_status as _update_worker_dispatch_status,
     string_list as _string_list,
     string_value as _string_value,
     TEARDOWN_BEFORE_TERMINAL_REASON,
@@ -37,7 +40,6 @@ from .dispatch_state import (
     worker_jobs_from_file as _worker_jobs_from_file,
     worker_state as _worker_state,
     worktree_path as _worktree_path,
-    write_json as _write_json,
 )
 from .final_report_paths import (
     final_report_data_path as _final_report_data_path,
@@ -74,6 +76,7 @@ class WorkerHandle:
     completed_process: subprocess.CompletedProcess[str] | None
     status_sidecar_path: Path | None
     degraded_from: str
+    running_process: subprocess.Popen[str] | None = None
 
 
 @dataclass(frozen=True)
@@ -174,7 +177,6 @@ def build_dispatch_plan(
     )
     if jobs_file:
         jobs = _jobs_from_file(project_root, workspace_root, jobs_file, options)
-        _validate_dispatch_prompts(manifest, active_context, jobs)
     else:
         jobs = _jobs_from_roster(
             project_root,
@@ -186,6 +188,7 @@ def build_dispatch_plan(
             requested_workers,
             options,
         )
+    _validate_dispatch_prompts(manifest, active_context, jobs)
     return DispatchPlan(
         project_root=project_root,
         workspace_root=workspace_root.resolve(),
@@ -199,6 +202,7 @@ def build_dispatch_plan(
 
 
 def dispatch_plan(plan: DispatchPlan, *, wait: bool = True) -> int:
+    _validate_report_writer_isolation(plan.jobs)
     if wait:
         pane_backends = sorted(
             {
@@ -225,6 +229,161 @@ def dispatch_plan(plan: DispatchPlan, *, wait: bool = True) -> int:
     return 0
 
 
+def dispatch_cli_wrapper_plan(plan: DispatchPlan) -> int:
+    """Start one dependency-free CLI batch before collecting any worker."""
+    if any(job.backend != BACKEND_CLI_WRAPPER for job in plan.jobs):
+        raise DispatchError("concurrent CLI dispatch requires cli-wrapper jobs only")
+    _validate_report_writer_isolation(plan.jobs)
+    _record_dispatch_facts(plan.team_state_path, _dispatch_mode(plan.jobs))
+    final_codes = _dispatch_cli_wrapper_batch(plan, plan.jobs)
+    return next((final_codes[job.worker_id] for job in plan.jobs
+                 if final_codes.get(job.worker_id, 0) != 0), 0)
+
+
+def _validate_report_writer_isolation(jobs: Sequence[WorkerJob]) -> None:
+    if (
+        any(job.worker_id == REPORT_WRITER_WORKER_ID for job in jobs)
+        and len(jobs) > 1
+    ):
+        raise DispatchError(
+            "report-writer must run in a separate Phase 6 dispatch after "
+            "analysis and convergence"
+        )
+
+
+def _dispatch_cli_wrapper_batch(
+    plan: DispatchPlan,
+    jobs: Sequence[WorkerJob],
+) -> dict[str, int]:
+    """Run one concurrent batch through all retries before the next dependency."""
+    pending = [(job, 1) for job in jobs]
+    final_codes: dict[str, int] = {}
+    while pending:
+        active: list[tuple[WorkerHandle, int]] = []
+        starting_job: WorkerJob | None = None
+        starting_attempt = 1
+        try:
+            for job, attempt in pending:
+                starting_job = job
+                starting_attempt = attempt
+                active.append((_spawn_cli_job_nonblocking(plan, job, attempt), attempt))
+        except (OSError, DispatchError) as exc:
+            if starting_job is not None:
+                _abort_cli_batch_start(
+                    plan,
+                    active,
+                    starting_job,
+                    starting_attempt,
+                    pending,
+                    exc,
+                )
+            raise
+        outcomes: list[tuple[WorkerHandle, int, WorkerOutcome]] = []
+        for handle, attempt in active:
+            process = handle.running_process
+            if process is None:
+                raise DispatchError("concurrent CLI worker process is missing")
+            returncode = process.wait()
+            completed = subprocess.CompletedProcess(handle.job.command, returncode)
+            settled = replace(handle, completed_process=completed, running_process=None)
+            outcomes.append((settled, attempt, _outcome_from_completed(settled)))
+
+        next_round: list[tuple[WorkerJob, int]] = []
+        for handle, attempt, outcome in outcomes:
+            job = handle.job
+            _finish_attempt(plan, job, attempt, outcome)
+            if _worker_terminal_status(plan.team_state_path, job.worker_id) == "completed":
+                final_codes[job.worker_id] = 0
+                continue
+            if _should_retry(outcome, attempt):
+                _append_event(
+                    plan,
+                    "worker-retry-scheduled",
+                    _retry_details(job, attempt, outcome),
+                )
+                next_round.append((job, attempt + 1))
+                continue
+            final_codes[job.worker_id] = outcome.returncode or 1
+        pending = next_round
+    return final_codes
+
+
+def _abort_cli_batch_start(
+    plan: DispatchPlan,
+    active: Sequence[tuple[WorkerHandle, int]],
+    failed_job: WorkerJob,
+    failed_attempt: int,
+    pending: Sequence[tuple[WorkerJob, int]],
+    error: Exception,
+) -> None:
+    reason = f"concurrent CLI batch launch failed: {error}"
+    for handle, _attempt in active:
+        process = handle.running_process
+        if process is not None and process.poll() is None:
+            process.terminate()
+    for handle, attempt in active:
+        process = handle.running_process
+        if process is not None:
+            process.wait()
+        _transition_worker_status(
+            plan.team_state_path, handle.job.worker_id, "error", reason
+        )
+        _update_dispatch_status(
+            plan.team_state_path, handle.job, attempt, "error", reason
+        )
+    _transition_worker_status(
+        plan.team_state_path, failed_job.worker_id, "error", reason
+    )
+    if not _update_dispatch_status(
+        plan.team_state_path,
+        failed_job,
+        failed_attempt,
+        "error",
+        reason,
+    ):
+        failed_handle = WorkerHandle(
+            failed_job,
+            "",
+            None,
+            status_path_for_prompt(failed_job.prompt_path),
+            "",
+        )
+        _record_dispatch(
+            plan.team_state_path,
+            failed_handle,
+            failed_attempt,
+            "error",
+            reason,
+        )
+    closed = {
+        (handle.job.worker_id, attempt) for handle, attempt in active
+    }
+    closed.add((failed_job.worker_id, failed_attempt))
+    for job, attempt in pending:
+        if (job.worker_id, attempt) in closed:
+            continue
+        _transition_worker_status(
+            plan.team_state_path, job.worker_id, "error", reason
+        )
+        if not _update_dispatch_status(
+            plan.team_state_path, job, attempt, "error", reason
+        ):
+            handle = WorkerHandle(
+                job,
+                "",
+                None,
+                status_path_for_prompt(job.prompt_path),
+                "",
+            )
+            _record_dispatch(
+                plan.team_state_path,
+                handle,
+                attempt,
+                "error",
+                reason,
+            )
+
+
 def await_dispatches(
     plan: DispatchPlan,
     *,
@@ -312,7 +471,15 @@ def _select_workers(
             options,
         )
     supported = set(options.supported_worker_wrappers)
-    selected = list(requested_workers) if requested_workers else [w for w in recommended if w in supported]
+    selected = (
+        list(requested_workers)
+        if requested_workers
+        else [
+            worker
+            for worker in recommended
+            if worker in supported and worker != REPORT_WRITER_WORKER_ID
+        ]
+    )
     unknown = [worker for worker in selected if worker not in recommended]
     if unknown:
         raise DispatchError("requested worker(s) are not in this run roster: " + ", ".join(unknown))
@@ -339,7 +506,11 @@ def _select_cli_wrapper_assignments(
         if _is_cli_wrapper_assignment(team_state, worker_id, options)
     }
     if not requested_workers:
-        selected = [worker for worker in recommended if worker in dispatchable]
+        selected = [
+            worker
+            for worker in recommended
+            if worker in dispatchable and worker != REPORT_WRITER_WORKER_ID
+        ]
         if selected:
             return selected
         raise DispatchError("run roster has no cli-wrapper assignments for codex-dispatch")
@@ -400,6 +571,8 @@ def _job_from_roster_worker(
     options: _BuildOptions,
 ) -> WorkerJob:
     result_path = _resolve_project_path(project_root, fact.result_path)
+    metadata = _agent_metadata(prompt_path)
+    invocation = _invocation_job_fields(metadata, prompt_path)
     return WorkerJob(
         worker_id=fact.worker_id,
         provider=fact.provider,
@@ -419,9 +592,47 @@ def _job_from_roster_worker(
         role=fact.role,
         idle_timeout_seconds=options.idle_timeout_seconds,
         dispatch_kind=options.dispatch_kind,
+        **invocation,
     )
 
 
+def _agent_metadata(prompt_path: Path) -> Mapping[str, Any] | None:
+    metadata_path = prompt_path.with_name(prompt_path.name + ".meta.json")
+    if not metadata_path.is_file():
+        return None
+    return _load_json_object(metadata_path, "agent invocation metadata")
+
+
+def _invocation_job_fields(
+    metadata: Mapping[str, Any] | None,
+    prompt_path: Path,
+) -> dict[str, Any]:
+    if metadata is None:
+        return {}
+    digests = metadata.get("digests")
+    assignment = metadata.get("modelAssignment")
+    if not isinstance(digests, Mapping) or not isinstance(assignment, Mapping):
+        raise DispatchError("agent invocation metadata is incomplete")
+    host_model_value = assignment.get("hostModelValue")
+    if host_model_value is not None and not isinstance(host_model_value, str):
+        raise DispatchError("agent invocation hostModelValue is invalid")
+    return {
+        "invocation_id": _require_string(metadata, "invocationId"),
+        "audience": _require_string(metadata, "audience"),
+        "assignment_ref": _require_string(metadata, "assignmentRef"),
+        "prompt_metadata_path": prompt_path.with_name(
+            prompt_path.name + ".meta.json"
+        ),
+        "catalog_digest": _require_string(digests, "catalogDigest"),
+        "assignment_digest": _require_string(digests, "assignmentDigest"),
+        "duty_digest": _require_string(digests, "dutyDigest"),
+        "instruction_digest": _require_string(digests, "instructionDigest"),
+        "prompt_digest": _require_string(digests, "promptDigest"),
+        "host_model_value": host_model_value,
+        "enforcement_mode": "core-pre-dispatch",
+    }
+
+
 def _materialize_roster_prompts(
     project_root: Path,
     workspace_root: Path,
@@ -510,6 +721,29 @@ def _spawn_job(plan: DispatchPlan, job: WorkerJob, attempt: int) -> WorkerHandle
     return handle
 
 
+def _spawn_cli_job_nonblocking(
+    plan: DispatchPlan, job: WorkerJob, attempt: int,
+) -> WorkerHandle:
+    _transition_worker_status(
+        plan.team_state_path, job.worker_id, "in-progress", "",
+        model_execution_value=job.model_execution_value,
+    )
+    process: subprocess.Popen[str] | None = None
+    try:
+        process = subprocess.Popen(job.command, cwd=plan.project_root, text=True)
+        handle = WorkerHandle(
+            job, "", None, status_path_for_prompt(job.prompt_path), "", process
+        )
+        _record_dispatch(plan.team_state_path, handle, attempt, "running", "")
+        _append_event(plan, "worker-dispatched", _attempt_details(job, attempt, handle))
+        return handle
+    except (OSError, DispatchError):
+        if process is not None and process.poll() is None:
+            process.terminate()
+            process.wait()
+        raise
+
+
 def _dispatch_job_with_retry(plan: DispatchPlan, job: WorkerJob) -> int:
     for attempt in range(1, MAX_WORKER_ATTEMPTS + 1):
         handle = _spawn_job(plan, job, attempt)
@@ -700,6 +934,13 @@ def _retry_from_record(
 
 def _finish_attempt(plan: DispatchPlan, job: WorkerJob, attempt: int, outcome: WorkerOutcome) -> None:
     if outcome.returncode == 0 and not outcome.missing_completion_paths and not outcome.timeout:
+        if job.invocation_id:
+            _link_agent_dispatch_result(
+                project_root=plan.project_root,
+                run_manifest_path=plan.manifest_path,
+                dispatch_id=f"{job.invocation_id}:attempt-{attempt}",
+                result_path=job.worker_result_path,
+            )
         post_process = _post_process_report_writer_result(plan, job)
         if not post_process["ok"]:
             reason = _require_string(post_process, "reason")
@@ -779,18 +1020,23 @@ def _worker_terminal_status(team_state_path: Path, worker_id: str) -> str:
 def _record_dispatch(
     team_state_path: Path, handle: WorkerHandle, attempt: int, status: str, reason: str
 ) -> None:
-    payload = _load_json_object(team_state_path, "team-state")
-    dispatches = payload.setdefault("workerDispatches", [])
-    if not isinstance(dispatches, list):
-        raise DispatchError(f"team-state workerDispatches must be an array: {team_state_path}")
-    dispatches.append(_dispatch_record(handle.job, attempt, status, handle.pane_id, handle.degraded_from, reason))
-    _write_json(team_state_path, payload)
+    _append_worker_dispatch(
+        team_state_path,
+        _dispatch_record(
+            handle.job,
+            attempt,
+            status,
+            handle.pane_id,
+            handle.degraded_from,
+            reason,
+        ),
+    )
 
 
 def _dispatch_record(
     job: WorkerJob, attempt: int, status: str, pane_id: str, degraded_from: str, reason: str = ""
 ) -> dict[str, Any]:
-    return {
+    record = {
         "workerId": job.worker_id,
         "role": job.role,
         "kind": job.dispatch_kind,
@@ -814,6 +1060,18 @@ def _dispatch_record(
         "completionPaths": [str(path) for path in job.completion_paths],
         "worktreePath": job.worktree_path,
     }
+    if job.invocation_id:
+        record.update({
+            "dispatchId": f"{job.invocation_id}:attempt-{attempt}",
+            "invocationId": job.invocation_id,
+            "audience": job.audience,
+            "assignmentRef": job.assignment_ref,
+            "promptMetadataPath": str(job.prompt_metadata_path),
+            **job.digests,
+            "hostModelValue": job.host_model_value,
+            "enforcementMode": job.enforcement_mode,
+        })
+    return record
 
 
 def _liveness_mode(backend: str) -> str:
@@ -824,22 +1082,14 @@ def _liveness_mode(backend: str) -> str:
 
 def _update_dispatch_status(
     team_state_path: Path, job: WorkerJob, attempt: int, status: str, reason: str
-) -> None:
-    payload = _load_json_object(team_state_path, "team-state")
-    matches = [
-        record
-        for record in payload.get("workerDispatches", [])
-        if dispatch_record_matches(record, job.prompt_path, attempt)
-    ]
-    if not matches:
-        return
-    # A lead may re-send a round's prompt without advancing `attempt`, which
-    # appends a second row the prompt cannot tell from the first. The newest is
-    # the one still running, so an outcome written to the older row would leave
-    # the live one running for good — the shape that wedged an await loop.
-    matches[-1]["status"] = status
-    matches[-1]["reason"] = reason
-    _write_json(team_state_path, payload)
+) -> bool:
+    return _update_worker_dispatch_status(
+        team_state_path,
+        prompt_path=job.prompt_path,
+        attempt=attempt,
+        status=status,
+        reason=reason,
+    )
 
 
 def _teardown_marked_dispatches(team_state_path: Path) -> list[Mapping[str, Any]]:
@@ -884,15 +1134,23 @@ def _mark_roster_skips(
     if not reasons:
         return
     team_state_path = _resolve_required_path(project_root, manifest, "teamStatePath")
-    current_team_state = _load_json_object(team_state_path, "team-state")
-    for worker_id, reason in reasons.items():
-        worker = _worker_state(current_team_state, worker_id)
-        if (
-            _string_value(worker.get("status")) != "not-run"
-            or _string_value(worker.get("reason"))
-        ):
-            continue
-        _transition_worker_status(team_state_path, worker_id, "not-run", reason)
+
+    def mark_pristine_skips(current_team_state: dict[str, Any]) -> bool:
+        changed = False
+        for worker_id, reason in reasons.items():
+            worker = _worker_state(current_team_state, worker_id)
+            if (
+                _string_value(worker.get("status")) != "not-run"
+                or _string_value(worker.get("reason"))
+            ):
+                continue
+            worker["reason"] = reason
+            worker.pop("startedAt", None)
+            worker.pop("endedAt", None)
+            changed = True
+        return changed
+
+    _mutate_team_state(team_state_path, mark_pristine_skips)
 
 
 def _skip_reasons(
@@ -1068,6 +1326,23 @@ def _job_from_record(project_root: Path, record: Mapping[str, Any]) -> WorkerJob
         role=_require_string(record, "role"),
         idle_timeout_seconds=600,
         dispatch_kind=_require_string(record, "kind"),
+        invocation_id=_string_value(record.get("invocationId")),
+        audience=_string_value(record.get("audience")),
+        assignment_ref=_string_value(record.get("assignmentRef")),
+        prompt_metadata_path=Path(
+            _string_value(record.get("promptMetadataPath"))
+        ),
+        catalog_digest=_string_value(record.get("catalogDigest")),
+        assignment_digest=_string_value(record.get("assignmentDigest")),
+        duty_digest=_string_value(record.get("dutyDigest")),
+        instruction_digest=_string_value(record.get("instructionDigest")),
+        prompt_digest=_string_value(record.get("promptDigest")),
+        host_model_value=(
+            record.get("hostModelValue")
+            if isinstance(record.get("hostModelValue"), str)
+            else None
+        ),
+        enforcement_mode=_string_value(record.get("enforcementMode")),
     )