okstra 0.167.0 → 0.169.0

package/runtime/python/okstra_ctl/dispatch_state.py

@@ -17,14 +17,22 @@ Enforced by `tests/contract/test_dispatcher_shared_helpers.py`.
 """
 from __future__ import annotations
 
+import contextlib
+import fcntl
 import json
 import os
+import tempfile
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Callable, Mapping, Sequence
 
 from . import cmux
+from .agent_invocation import (
+    AgentInvocationError,
+    agent_model_assignment_from_payload,
+    verify_agent_invocation,
+)
 from .worker_prompt_contract import (
     PromptRecord,
     validate_initial_prompt_records,
@@ -95,6 +103,17 @@ class WorkerJob:
     role: str
     idle_timeout_seconds: int
     dispatch_kind: str
+    invocation_id: str = ""
+    audience: str = ""
+    assignment_ref: str = ""
+    prompt_metadata_path: Path = Path()
+    catalog_digest: str = ""
+    assignment_digest: str = ""
+    duty_digest: str = ""
+    instruction_digest: str = ""
+    prompt_digest: str = ""
+    host_model_value: str | None = None
+    enforcement_mode: str = ""
 
     @property
     def command(self) -> list[str]:
@@ -118,7 +137,7 @@ class WorkerJob:
         return QUIET if self.backend == BACKEND_CLI_WRAPPER else LIVE
 
     def to_payload(self) -> dict[str, Any]:
-        return {
+        payload = {
             "workerId": self.worker_id,
             "provider": self.provider,
             "backend": self.backend,
@@ -133,6 +152,27 @@ class WorkerJob:
             "dispatchKind": self.dispatch_kind,
             "command": self.command,
         }
+        if self.invocation_id:
+            payload.update({
+                "invocationId": self.invocation_id,
+                "audience": self.audience,
+                "assignmentRef": self.assignment_ref,
+                "promptMetadataPath": str(self.prompt_metadata_path),
+                "digests": self.digests,
+                "hostModelValue": self.host_model_value,
+                "enforcementMode": self.enforcement_mode,
+            })
+        return payload
+
+    @property
+    def digests(self) -> dict[str, str]:
+        return {
+            "catalogDigest": self.catalog_digest,
+            "assignmentDigest": self.assignment_digest,
+            "dutyDigest": self.duty_digest,
+            "instructionDigest": self.instruction_digest,
+            "promptDigest": self.prompt_digest,
+        }
 
 
 # --- run-manifest reads -------------------------------------------------------
@@ -201,11 +241,293 @@ def write_json(path: Path, payload: Mapping[str, Any]) -> None:
     # Write via temp file + os.replace so a crash mid-write cannot leave a
     # truncated team-state.json that every later load_json_object rejects;
     # os.replace is atomic on POSIX, matching run_context._atomic_write_json.
-    tmp = path.with_suffix(path.suffix + ".tmp")
-    tmp.write_text(
-        json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
+    path.parent.mkdir(parents=True, exist_ok=True)
+    temp_path: Path | None = None
+    try:
+        with tempfile.NamedTemporaryFile(
+            mode="w",
+            encoding="utf-8",
+            dir=path.parent,
+            prefix=f".{path.name}.",
+            suffix=".tmp",
+            delete=False,
+        ) as handle:
+            temp_path = Path(handle.name)
+            handle.write(json.dumps(payload, ensure_ascii=False, indent=2) + "\n")
+            handle.flush()
+            os.fsync(handle.fileno())
+        os.replace(temp_path, path)
+        temp_path = None
+    finally:
+        if temp_path is not None:
+            temp_path.unlink(missing_ok=True)
+
+
+@contextlib.contextmanager
+def _team_state_lock(path: Path):
+    lock_path = path.with_name(path.name + ".lock")
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    handle = lock_path.open("a+")
+    try:
+        fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+        yield
+    finally:
+        fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+        handle.close()
+
+
+def mutate_team_state(
+    team_state_path: Path,
+    mutation: Callable[[dict[str, Any]], bool],
+) -> bool:
+    """Apply one arbitrary team-state mutation under the shared file lock."""
+    with _team_state_lock(team_state_path):
+        payload = load_json_object(team_state_path, "team-state")
+        changed = mutation(payload)
+        if changed:
+            write_json(team_state_path, payload)
+        return changed
+
+
+def append_worker_dispatch(
+    team_state_path: Path,
+    record: Mapping[str, Any],
+) -> None:
+    """Append one worker dispatch without losing another state writer's update."""
+    with _team_state_lock(team_state_path):
+        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(dict(record))
+        write_json(team_state_path, payload)
+
+
+def update_worker_dispatch_status(
+    team_state_path: Path,
+    *,
+    prompt_path: Path,
+    attempt: int,
+    status: str,
+    reason: str,
+) -> bool:
+    """Settle the newest matching dispatch row under the shared state lock."""
+    with _team_state_lock(team_state_path):
+        payload = load_json_object(team_state_path, "team-state")
+        dispatches = payload.get("workerDispatches", [])
+        if not isinstance(dispatches, list):
+            raise DispatchError(
+                f"team-state workerDispatches must be an array: {team_state_path}"
+            )
+        matches = [
+            record
+            for record in dispatches
+            if dispatch_record_matches(record, prompt_path, attempt)
+        ]
+        if not matches:
+            return False
+        # A re-send can reuse prompt + attempt. The newest row is the live one.
+        matches[-1]["status"] = status
+        matches[-1]["reason"] = reason
+        write_json(team_state_path, payload)
+        return True
+
+
+_AGENT_DIGEST_KEYS = (
+    "catalogDigest",
+    "assignmentDigest",
+    "dutyDigest",
+    "instructionDigest",
+    "promptDigest",
+)
+_AGENT_ENFORCEMENT_MODES = frozenset({
+    "core-pre-dispatch",
+    "host-native-spec-link-gate",
+})
+
+
+def record_verified_agent_dispatch(
+    *,
+    project_root: Path,
+    run_manifest_path: Path,
+    metadata_path: Path,
+    enforcement_mode: str,
+    allow_native_core: bool = False,
+) -> dict[str, Any]:
+    """Append one immutable dispatch-to-invocation association.
+
+    Host-native calls use this record as a specification link; it deliberately
+    does not claim that Okstra observed the bytes delivered by the host. A
+    code-owned lead process may set ``allow_native_core`` because lead model
+    assignments use the host's native model vocabulary even when Okstra owns
+    the outer provider CLI process.
+    """
+    project_root = project_root.resolve()
+    run_manifest_path = resolve_project_path(
+        project_root, str(run_manifest_path)
+    ).resolve(strict=True)
+    metadata_path = resolve_project_path(
+        project_root, str(metadata_path)
+    ).resolve(strict=True)
+    if enforcement_mode not in _AGENT_ENFORCEMENT_MODES:
+        raise DispatchError(f"invalid agent enforcement mode: {enforcement_mode}")
+    manifest = load_json_object(run_manifest_path, "run manifest")
+    metadata = load_json_object(metadata_path, "agent invocation metadata")
+    assignment_ref = require_string(metadata, "assignmentRef")
+    assignments = manifest.get("invocationAssignments")
+    if not isinstance(assignments, Mapping):
+        raise DispatchError("run manifest has no invocationAssignments object")
+    try:
+        assignment = agent_model_assignment_from_payload(assignments.get(assignment_ref))
+    except AgentInvocationError as exc:
+        raise DispatchError(str(exc)) from exc
+    worker_id = require_string(metadata, "workerId")
+    audience = require_string(metadata, "audience")
+    invocation_id = require_string(metadata, "invocationId")
+    errors = verify_agent_invocation(
+        metadata_path,
+        project_root=project_root,
+        expected_run_manifest_path=run_manifest_path,
+        expected_assignment=assignment,
+        expected_invocation_id=invocation_id,
+        expected_worker_id=worker_id,
+        expected_assignment_ref=assignment_ref,
+        expected_audience=audience,
     )
-    os.replace(tmp, path)
+    if errors:
+        raise DispatchError("agent invocation verification failed: " + "; ".join(errors))
+    if (
+        enforcement_mode == "host-native-spec-link-gate"
+        and assignment.runner != "native-session"
+    ):
+        raise DispatchError(
+            "host-native enforcement requires a native-session assignment"
+        )
+    if (
+        enforcement_mode == "core-pre-dispatch"
+        and assignment.runner == "native-session"
+        and not allow_native_core
+    ):
+        raise DispatchError(
+            "native-session assignment must use host-native-spec-link-gate"
+        )
+    digests = metadata.get("digests")
+    if not isinstance(digests, Mapping):
+        raise DispatchError("agent invocation digests are missing")
+    digest_values = {
+        key: require_string(digests, key) for key in _AGENT_DIGEST_KEYS
+    }
+    prompt = metadata.get("prompt")
+    if not isinstance(prompt, Mapping):
+        raise DispatchError("agent invocation prompt is missing")
+    prompt_path = require_string(prompt, "path")
+    dispatch_id = f"{invocation_id}:attempt-1"
+    record = {
+        "dispatchId": dispatch_id,
+        "workerId": worker_id,
+        "audience": audience,
+        "invocationId": invocation_id,
+        "assignmentRef": assignment_ref,
+        "promptPath": prompt_path,
+        "promptMetadataPath": _relative_project_path(project_root, metadata_path),
+        **digest_values,
+        "modelExecutionValue": assignment.model_execution_value,
+        "hostModelValue": assignment.host_model_value,
+        "enforcementMode": enforcement_mode,
+        "promptDeliveryVerified": False,
+        "status": "dispatched",
+    }
+    team_state_path = resolve_required_path(project_root, manifest, "teamStatePath")
+    with _team_state_lock(team_state_path):
+        team_state = load_json_object(team_state_path, "team-state")
+        dispatches = team_state.setdefault("agentDispatches", [])
+        if not isinstance(dispatches, list):
+            raise DispatchError("team-state agentDispatches must be an array")
+        existing = [
+            item for item in dispatches
+            if isinstance(item, Mapping) and item.get("dispatchId") == dispatch_id
+        ]
+        if existing:
+            if len(existing) != 1 or dict(existing[0]) != record:
+                raise DispatchError(f"agent dispatch ID conflicts: {dispatch_id}")
+            return record
+        dispatches.append(record)
+        write_json(team_state_path, team_state)
+    return record
+
+
+def link_agent_dispatch_result(
+    *,
+    project_root: Path,
+    run_manifest_path: Path,
+    dispatch_id: str,
+    result_path: Path,
+) -> dict[str, str]:
+    """Link an existing result artifact to exactly one verified dispatch."""
+    project_root = project_root.resolve()
+    manifest_path = resolve_project_path(
+        project_root, str(run_manifest_path)
+    ).resolve(strict=True)
+    manifest = load_json_object(manifest_path, "run manifest")
+    run_root = resolve_required_path(
+        project_root, manifest, "runDirectoryPath"
+    ).resolve(strict=True)
+    resolved_result = resolve_project_path(
+        project_root, str(result_path)
+    ).resolve(strict=True)
+    if not resolved_result.is_relative_to(run_root):
+        raise DispatchError("agent result link must stay inside the current run")
+    result_relative = _relative_project_path(project_root, resolved_result)
+    team_state_path = resolve_required_path(project_root, manifest, "teamStatePath")
+    link = {"dispatchId": dispatch_id, "resultPath": result_relative}
+    with _team_state_lock(team_state_path):
+        team_state = load_json_object(team_state_path, "team-state")
+        dispatches = [
+            item
+            for collection in (
+                team_state.get("agentDispatches") or [],
+                team_state.get("workerDispatches") or [],
+            )
+            for item in collection
+            if isinstance(item, Mapping) and item.get("dispatchId") == dispatch_id
+        ]
+        if len(dispatches) != 1:
+            raise DispatchError(
+                f"agent result link requires exactly one dispatch record: {dispatch_id}"
+            )
+        links = team_state.setdefault("agentResultLinks", [])
+        if not isinstance(links, list):
+            raise DispatchError("team-state agentResultLinks must be an array")
+        same_result = [
+            item for item in links
+            if isinstance(item, Mapping) and item.get("resultPath") == result_relative
+        ]
+        if same_result and any(dict(item) != link for item in same_result):
+            raise DispatchError(
+                f"agent result is already linked to another dispatch: {result_relative}"
+            )
+        if any(
+            isinstance(item, Mapping)
+            and item.get("dispatchId") == dispatch_id
+            and item.get("resultPath") != result_relative
+            for item in links
+        ):
+            raise DispatchError(
+                f"agent dispatch is already linked to another result: {dispatch_id}"
+            )
+        if link not in links:
+            links.append(link)
+            write_json(team_state_path, team_state)
+    return link
+
+
+def _relative_project_path(project_root: Path, path: Path) -> str:
+    try:
+        return path.resolve(strict=False).relative_to(project_root).as_posix()
+    except ValueError as exc:
+        raise DispatchError(f"path is outside project root: {path}") from exc
 
 
 def worker_state(team_state: Mapping[str, Any], worker_id: str) -> Mapping[str, Any]:
@@ -236,31 +558,32 @@ def transition_worker_status(
     if status in REASON_REQUIRED_STATUSES and not reason:
         raise DispatchError(f"worker status `{status}` requires a non-empty reason")
     timestamp = _utc_timestamp(at)
-    payload = load_json_object(team_state_path, "team-state")
-    workers = payload.get("workers")
-    if not isinstance(workers, list):
-        raise DispatchError(f"team-state workers must be an array: {team_state_path}")
-    for worker in workers:
-        if isinstance(worker, dict) and worker.get("workerId") == worker_id:
-            worker["status"] = status
-            worker["reason"] = reason
-            if status == "in-progress":
-                worker["startedAt"] = timestamp
-                worker.pop("endedAt", None)
-            elif status == "not-run":
-                worker.pop("startedAt", None)
-                worker.pop("endedAt", None)
-            else:
-                worker["endedAt"] = timestamp
-            if model_execution_value:
-                # `model` is the catalog display name the task-manifest
-                # declares; only the execution identifier belongs here. Writing
-                # both from one value destroys the display name on every worker
-                # whose two differ, and `validate-run` compares team-state's
-                # `model` against the manifest.
-                worker["modelExecutionValue"] = model_execution_value
-            write_json(team_state_path, payload)
-            return
+    with _team_state_lock(team_state_path):
+        payload = load_json_object(team_state_path, "team-state")
+        workers = payload.get("workers")
+        if not isinstance(workers, list):
+            raise DispatchError(f"team-state workers must be an array: {team_state_path}")
+        for worker in workers:
+            if isinstance(worker, dict) and worker.get("workerId") == worker_id:
+                worker["status"] = status
+                worker["reason"] = reason
+                if status == "in-progress":
+                    worker["startedAt"] = timestamp
+                    worker.pop("endedAt", None)
+                elif status == "not-run":
+                    worker.pop("startedAt", None)
+                    worker.pop("endedAt", None)
+                else:
+                    worker["endedAt"] = timestamp
+                if model_execution_value:
+                    # `model` is the catalog display name the task-manifest
+                    # declares; only the execution identifier belongs here. Writing
+                    # both from one value destroys the display name on every worker
+                    # whose two differ, and `validate-run` compares team-state's
+                    # `model` against the manifest.
+                    worker["modelExecutionValue"] = model_execution_value
+                write_json(team_state_path, payload)
+                return
     raise DispatchError(f"team-state has no workerId={worker_id}: {team_state_path}")
 
 
@@ -294,11 +617,12 @@ def record_dispatch_facts(team_state_path: Path, dispatch_mode: str) -> None:
     rule out. A marker the launch prompt pre-recorded (the concurrent-run
     `skipped` decision) belongs to that run and is never overwritten.
     """
-    payload = load_json_object(team_state_path, "team-state")
-    payload["dispatchMode"] = dispatch_mode
-    if not _has_recorded_team_create(payload):
-        payload["teamCreate"] = {"attempted": False, "status": "implicit"}
-    write_json(team_state_path, payload)
+    with _team_state_lock(team_state_path):
+        payload = load_json_object(team_state_path, "team-state")
+        payload["dispatchMode"] = dispatch_mode
+        if not _has_recorded_team_create(payload):
+            payload["teamCreate"] = {"attempted": False, "status": "implicit"}
+        write_json(team_state_path, payload)
 
 
 def _has_recorded_team_create(team_state: Mapping[str, Any]) -> bool:
@@ -324,7 +648,13 @@ def validate_initial_prompts(
     manifest: Mapping[str, Any], jobs: Sequence[WorkerJob]
 ) -> None:
     records = [
-        PromptRecord(job.worker_id, job.dispatch_kind, job.prompt_path)
+        PromptRecord(
+            job.worker_id,
+            job.dispatch_kind,
+            job.prompt_path,
+            metadata_path=(job.prompt_metadata_path if job.invocation_id else None),
+            expected_duty_audience=(job.audience if job.invocation_id else None),
+        )
         for job in jobs
     ]
     errors = validate_initial_prompt_records(
@@ -342,6 +672,8 @@ def validate_dispatch_prompts(
     active_context: Mapping[str, Any],
     jobs: Sequence[WorkerJob],
 ) -> None:
+    if isinstance(manifest.get("agentContract"), Mapping):
+        _validate_agent_invocations(manifest, jobs)
     initial_jobs = [
         job for job in jobs if not job.dispatch_kind.startswith("reverify-r")
     ]
@@ -388,6 +720,80 @@ def validate_dispatch_prompts(
         raise DispatchError("reverify prompt contract: " + "; ".join(errors))
 
 
+def _validate_agent_invocations(
+    manifest: Mapping[str, Any],
+    jobs: Sequence[WorkerJob],
+) -> None:
+    manifest_path_value = require_string(manifest, "runManifestPath")
+    assignments = manifest.get("invocationAssignments")
+    if not isinstance(assignments, Mapping):
+        raise DispatchError("run manifest has no invocationAssignments object")
+    errors: list[str] = []
+    for job in jobs:
+        missing = [
+            name
+            for name, value in (
+                ("invocationId", job.invocation_id),
+                ("audience", job.audience),
+                ("assignmentRef", job.assignment_ref),
+                ("promptMetadataPath", str(job.prompt_metadata_path)),
+                ("catalogDigest", job.catalog_digest),
+                ("assignmentDigest", job.assignment_digest),
+                ("dutyDigest", job.duty_digest),
+                ("instructionDigest", job.instruction_digest),
+                ("promptDigest", job.prompt_digest),
+                ("enforcementMode", job.enforcement_mode),
+            )
+            if not value or value == "."
+        ]
+        if missing:
+            errors.append(
+                f"{job.worker_id}: agent invocation metadata fields are missing: "
+                + ", ".join(missing)
+            )
+            continue
+        if job.enforcement_mode != "core-pre-dispatch":
+            errors.append(
+                f"{job.worker_id}: code-owned dispatch requires core-pre-dispatch"
+            )
+        assignment_payload = assignments.get(job.assignment_ref)
+        try:
+            assignment = agent_model_assignment_from_payload(assignment_payload)
+        except AgentInvocationError as exc:
+            errors.append(f"{job.worker_id}: {exc}")
+            continue
+        if assignment.model_execution_value != job.model_execution_value:
+            errors.append(
+                f"{job.worker_id}: model execution value does not match invocation assignment"
+            )
+        if assignment.host_model_value != job.host_model_value:
+            errors.append(
+                f"{job.worker_id}: host model value does not match invocation assignment"
+            )
+        manifest_path = resolve_project_path(job.project_root, manifest_path_value)
+        verification = verify_agent_invocation(
+            job.prompt_metadata_path,
+            project_root=job.project_root,
+            expected_run_manifest_path=manifest_path,
+            expected_assignment=assignment,
+            expected_invocation_id=job.invocation_id,
+            expected_worker_id=job.worker_id,
+            expected_assignment_ref=job.assignment_ref,
+            expected_audience=job.audience,
+        )
+        errors.extend(f"{job.worker_id}: {error}" for error in verification)
+        try:
+            metadata = load_json_object(job.prompt_metadata_path, "agent metadata")
+        except DispatchError as exc:
+            errors.append(f"{job.worker_id}: {exc}")
+            continue
+        digests = metadata.get("digests")
+        if not isinstance(digests, Mapping) or dict(digests) != job.digests:
+            errors.append(f"{job.worker_id}: job digests do not match agent metadata")
+    if errors:
+        raise DispatchError("agent invocation contract: " + "; ".join(errors))
+
+
 def worker_jobs_from_file(
     project_root: Path,
     jobs_file: Path,
@@ -446,6 +852,11 @@ def _worker_job_from_file(
         resolve_project_path(project_root, path)
         for path in string_list(item.get("completionPaths"))
     ) or (result_path,)
+    digests = item.get("digests")
+    digest_values = digests if isinstance(digests, Mapping) else {}
+    host_model_value = item.get("hostModelValue")
+    if host_model_value is not None and not isinstance(host_model_value, str):
+        raise DispatchError("jobs file hostModelValue must be a string or null")
     return WorkerJob(
         worker_id=worker_id,
         provider=provider,
@@ -461,6 +872,20 @@ def _worker_job_from_file(
         role=require_string(item, "role"),
         idle_timeout_seconds=idle_timeout_seconds,
         dispatch_kind=dispatch_kind,
+        invocation_id=string_value(item.get("invocationId")),
+        audience=string_value(item.get("audience")),
+        assignment_ref=string_value(item.get("assignmentRef")),
+        prompt_metadata_path=resolve_project_path(
+            project_root,
+            string_value(item.get("promptMetadataPath")),
+        ) if string_value(item.get("promptMetadataPath")) else Path(),
+        catalog_digest=string_value(digest_values.get("catalogDigest")),
+        assignment_digest=string_value(digest_values.get("assignmentDigest")),
+        duty_digest=string_value(digest_values.get("dutyDigest")),
+        instruction_digest=string_value(digest_values.get("instructionDigest")),
+        prompt_digest=string_value(digest_values.get("promptDigest")),
+        host_model_value=host_model_value,
+        enforcement_mode=string_value(item.get("enforcementMode")),
     )