agent-hitch 0.2.9 → 0.2.11

package/integrations/harbor/hitch_harbor_agent.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import asyncio
 import hashlib
 import json
+import math
 import os
 import re
 import secrets
@@ -127,6 +128,7 @@ class HitchHarborAgent(BaseAgent):
         verifier_identity: str | None = None,
         logical_attempt: int | None = None,
         model_capture: dict[str, Any] | None = None,
+        managed_local_inference: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(logs_dir=logs_dir, **kwargs)
@@ -165,6 +167,14 @@ class HitchHarborAgent(BaseAgent):
             raise ValueError("logical_attempt must be a positive integer")
         self.logical_attempt = logical_attempt
         self.model_capture = _validate_model_capture(model_capture)
+        if managed_local_inference is not None and not _valid_managed_model_identity(managed_local_inference):
+            raise ValueError("managed_local_inference identity is invalid")
+        if managed_local_inference is not None and self.model_capture is None:
+            raise ValueError("managed_local_inference requires model_capture")
+        self.managed_local_inference = dict(managed_local_inference) if managed_local_inference else None
+        capture_inference = self.model_capture.get("managed_inference") if self.model_capture else None
+        if capture_inference != self.managed_local_inference:
+            raise ValueError("managed_local_inference must match model_capture identity")
         self._hitch_version: str | None = None
         self._entrypoint: str | None = None
         self._artifact_manifest: dict[str, Any] | None = None
@@ -884,17 +894,38 @@ class HitchHarborAgent(BaseAgent):
             if timeout_ms <= 0:
                 raise RuntimeError("candidate whole-task budget expired during phase binding/upload")
         elif session:
-            preparation_ms = (time.monotonic_ns() - invocation_started_ns) // 1_000_000
-            timeout_ms = task_budget_ms - preparation_ms
-            (self.logs_dir / "hitch-agent-budget.json").write_text(json.dumps({
-                "schema_version": "hitch-agent-budget@1", "run_id": run_id,
-                "task_budget_ms": task_budget_ms, "preparation_ms": preparation_ms,
-                "hitch_timeout_ms": max(0, timeout_ms),
-                "collection_timeout_ms": session.config["profile"]["budget"]["collection_timeout_ms"],
-                "scope": "invocation-budget-and-collection-allowance",
-            }))
+            timeout_ms = self._record_session_budget(
+                session, run_id, invocation_started_ns, task_budget_ms
+            )
             if timeout_ms <= 0:
                 raise RuntimeError("candidate budget expired during input preparation; no model was launched")
+        private_environment, native_deadline_ns = await self._host_task_credentials(
+            environment,
+            context,
+            run_id,
+            timeout_ms,
+            invocation_started_ns=(
+                invocation_started_ns
+                if prepared_phase is None and session is None
+                else None
+            ),
+        )
+        if private_environment is not None:
+            if prepared_phase is not None:
+                timeout_ms = (prepared_phase.deadline_ns - time.monotonic_ns()) // 1_000_000
+            if native_deadline_ns is not None:
+                native_remaining_ms = (native_deadline_ns - time.monotonic_ns()) // 1_000_000
+                if native_remaining_ms <= 0:
+                    raise RuntimeError("candidate native Harbor budget expired during host credential preparation; no model was launched")
+            elif timeout_ms <= 0:
+                raise RuntimeError("candidate budget expired during host credential preparation; no model was launched")
+        if prepared_phase is None and session:
+            timeout_ms = self._record_session_budget(
+                session, run_id, invocation_started_ns, task_budget_ms
+            )
+            if timeout_ms <= 0:
+                stage = "host credential preparation" if private_environment is not None else "input preparation"
+                raise RuntimeError(f"candidate budget expired during {stage}; no model was launched")
         arguments = [
             self._node_prefix(),
             "HITCH_ROOT=/tmp/hitch-state",
@@ -934,7 +965,14 @@ class HitchHarborAgent(BaseAgent):
             arguments.extend(["--internal-credential-name", shlex.quote(name)])
         command = self._logged_run_command(" ".join(arguments))
         try:
-            execution = await environment.exec(command, cwd=workdir)
+            if private_environment is None:
+                execution = await environment.exec(command, cwd=workdir)
+            else:
+                execution = await environment.hitch_exec_with_private_env(
+                    command,
+                    cwd=workdir,
+                    env=private_environment,
+                )
         except (Exception, asyncio.CancelledError):
             # Harbor cancels this await on its agent deadline. Collect what is
             # already durable before it removes the container, without turning
@@ -959,6 +997,153 @@ class HitchHarborAgent(BaseAgent):
             (self.logs_dir / "hitch-collection-timeout.json").write_text(json.dumps(receipt))
             raise RuntimeError("hitch_run_collection_timeout: terminal evidence export exceeded its allowance") from error
 
+    async def _host_task_credentials(
+        self,
+        environment: BaseEnvironment,
+        context: AgentContext,
+        run_id: str,
+        remaining_timeout_ms: int,
+        *,
+        invocation_started_ns: int | None,
+    ) -> tuple[dict[str, str] | None, int | None]:
+        # Preserve the original single-file bridge behavior when this optional
+        # host feature is not configured. An empty value remains configured and
+        # reaches the strict parser below.
+        if "HITCH_HOST_CREDENTIAL_HELPER_JSON" not in os.environ:
+            return None, None
+        from hitch_host_credentials import (
+            HOST_CREDENTIAL_VALIDITY_MARGIN_MS,
+            HostCredentialHelperError,
+            load_host_credential_helper,
+            prepare_host_credentials,
+        )
+
+        try:
+            config = load_host_credential_helper(self.credential_names)
+            if config is None:
+                return None, None
+            if not callable(getattr(environment, "hitch_exec_with_private_env", None)):
+                raise HostCredentialHelperError(
+                    "host_credential_transport_unsupported",
+                    "Harbor environment does not support private Target credentials",
+                )
+            native_deadline_ns: int | None = None
+            if invocation_started_ns is not None and remaining_timeout_ms <= 0:
+                native_timeout_ms = self._native_harbor_agent_timeout_ms(environment)
+                if native_timeout_ms is None:
+                    raise HostCredentialHelperError(
+                        "host_credential_helper_request_invalid",
+                        "native Harbor agent timeout is not finite",
+                    )
+                native_deadline_ns = invocation_started_ns + native_timeout_ms * 1_000_000
+                native_remaining_ms = (native_deadline_ns - time.monotonic_ns()) // 1_000_000
+                if native_remaining_ms <= 0:
+                    raise RuntimeError("candidate native Harbor budget expired during input preparation; no model was launched")
+                remaining_timeout_ms = native_remaining_ms
+            credentials = await prepare_host_credentials(
+                config,
+                remaining_timeout_ms + HOST_CREDENTIAL_VALIDITY_MARGIN_MS,
+            )
+            return credentials, native_deadline_ns
+        except HostCredentialHelperError as error:
+            trial_id, task_id, attempt = self._trial_identity()
+            evidence = {
+                "schema_version": "1",
+                "code": error.code,
+                "message": error.message,
+                "recorded_at": datetime.now(timezone.utc).isoformat(),
+                "scope": "host-task-credential-helper",
+                "eval_id": self.eval_id,
+                "trial_id": trial_id,
+                "task_id": task_id,
+                "attempt": attempt,
+                "assigned_run_id": run_id,
+            }
+            if context.metadata is None:
+                context.metadata = {}
+            context.metadata["hitch_bridge_error_code"] = error.code
+            context.metadata["hitch_bridge_error_artifact"] = "hitch-bridge-error.json"
+            try:
+                await self._write_bridge_error(environment, evidence)
+            except Exception:
+                pass
+            raise HitchBridgeError(error.code, error.message, evidence) from None
+
+    def _native_harbor_agent_timeout_ms(self, environment: BaseEnvironment) -> int | None:
+        """Resolve Harbor 0.21's effective native agent timeout from its own inputs."""
+        try:
+            from harbor.models.task.config import TaskConfig
+            from harbor.models.trial.config import TrialConfig
+
+            environment_dir = Path(getattr(environment, "environment_dir"))
+            trial_paths = getattr(environment, "trial_paths")
+            config_path = Path(getattr(trial_paths, "config_path"))
+            task_config = TaskConfig.model_validate_toml(
+                (environment_dir.parent / "task.toml").read_text(encoding="utf-8")
+            )
+            trial_config = TrialConfig.model_validate_json(
+                config_path.read_text(encoding="utf-8")
+            )
+
+            def resolve(default_timeout_sec: float | None) -> float | None:
+                base_timeout_sec = trial_config.agent.override_timeout_sec or default_timeout_sec
+                if base_timeout_sec is None:
+                    return None
+                multiplier = trial_config.agent_timeout_multiplier
+                if multiplier is None:
+                    multiplier = trial_config.timeout_multiplier
+                return min(
+                    base_timeout_sec,
+                    trial_config.agent.max_timeout_sec or float("inf"),
+                ) * multiplier
+
+            if task_config.steps:
+                resolved = [
+                    resolve(
+                        step.agent.timeout_sec
+                        if step.agent.timeout_sec is not None
+                        else task_config.agent.timeout_sec
+                    )
+                    for step in task_config.steps
+                ]
+                # Harbor does not identify the current step in BaseAgent.run().
+                # Requiring the largest finite step budget prevents a credential
+                # from expiring during any native step without guessing order.
+                timeout_sec = None if any(value is None for value in resolved) else max(resolved)
+            else:
+                timeout_sec = resolve(task_config.agent.timeout_sec)
+            if timeout_sec is None:
+                return None
+            timeout_ms = math.ceil(timeout_sec * 1_000)
+            if not math.isfinite(timeout_sec) or not 1 <= timeout_ms <= 9_007_199_254_740_991:
+                raise ValueError("native Harbor timeout is outside the supported range")
+            return timeout_ms
+        except Exception:
+            # Task/config contents and host paths are not safe diagnostic text.
+            from hitch_host_credentials import HostCredentialHelperError
+            raise HostCredentialHelperError(
+                "host_credential_helper_request_invalid",
+                "native Harbor agent timeout could not be resolved",
+            ) from None
+
+    def _record_session_budget(
+        self,
+        session,
+        run_id: str,
+        invocation_started_ns: int,
+        task_budget_ms: int,
+    ) -> int:
+        preparation_ms = (time.monotonic_ns() - invocation_started_ns) // 1_000_000
+        timeout_ms = task_budget_ms - preparation_ms
+        (self.logs_dir / "hitch-agent-budget.json").write_text(json.dumps({
+            "schema_version": "hitch-agent-budget@1", "run_id": run_id,
+            "task_budget_ms": task_budget_ms, "preparation_ms": preparation_ms,
+            "hitch_timeout_ms": max(0, timeout_ms),
+            "collection_timeout_ms": session.config["profile"]["budget"]["collection_timeout_ms"],
+            "scope": "invocation-budget-and-collection-allowance",
+        }))
+        return timeout_ms
+
     @staticmethod
     def _logged_run_command(invocation: str) -> str:
         # Bash process substitution leaves an extra pipe FD in the harness and
@@ -1231,9 +1416,25 @@ mv "$stage_dir" "$target_dir"
                 raise RuntimeError("hitch-model-proxy-health: required model proxy is unreachable")
             return [], "degraded-unreachable"
         base = self.model_capture["base_url_template"].replace("{run_id}", run_id)
+        training = self.model_capture.get("training_external")
         return [
             f"OPENAI_BASE_URL={shlex.quote(base.replace('{provider}', 'openai'))}",
             f"ANTHROPIC_BASE_URL={shlex.quote(base.replace('{provider}', 'anthropic'))}",
+            *([
+                "HITCH_TRAINING_EXTERNAL=1",
+                f"HITCH_TRAINING_RUN_ID={run_id}",
+                f"HITCH_TRAINING_BINDING={shlex.quote(json.dumps(training, separators=(',', ':')))}",
+                "OPENAI_API_KEY=hitch-training-external",
+            ] if training else []),
+            *([
+                "HITCH_MANAGED_LOCAL_INFERENCE=1",
+                f"HITCH_MANAGED_RUN_ID={run_id}",
+                f"HITCH_MANAGED_INFERENCE_ID={self.managed_local_inference['inference_id']}",
+                f"HITCH_MANAGED_MODEL_ID={self.managed_local_inference['model_id']}",
+                "OPENAI_API_KEY=hitch-managed-local",
+                *([f"HITCH_MANAGED_NODE_BINDING={shlex.quote(json.dumps(self.managed_local_inference['model_node'], separators=(',', ':')))}"]
+                  if self.managed_local_inference.get("model_node") else []),
+            ] if self.managed_local_inference else []),
         ], "healthy"
 
     @staticmethod
@@ -1695,15 +1896,16 @@ esac""", "hitch_node_runtime_incompatible")
 def _validate_model_capture(value: dict[str, Any] | None) -> dict[str, Any] | None:
     if value is None:
         return None
-    if not isinstance(value, dict) or set(value) != {
+    required_fields = {
         "schema_version", "mode", "required", "topology", "base_url_template", "health_url_template"
-    }:
+    }
+    if not isinstance(value, dict) or not required_fields.issubset(value) or set(value) - required_fields - {"managed_inference", "training_external"}:
         raise ValueError("model_capture fields are invalid")
     if (
         value.get("schema_version") != "1"
         or value.get("mode") not in {"proxy", "hybrid"}
         or not isinstance(value.get("required"), bool)
-        or value.get("topology") != "host-side"
+        or value.get("topology") not in {"host-side", "in-sandbox"}
     ):
         raise ValueError("model_capture identity is invalid")
     for field, provider_count in (("base_url_template", 1), ("health_url_template", 0)):
@@ -1722,9 +1924,35 @@ def _validate_model_capture(value: dict[str, Any] | None) -> dict[str, Any] | No
         )
         if parsed.scheme not in {"http", "https"} or not parsed.hostname or parsed.username or parsed.password:
             raise ValueError(f"model_capture {field} URL is invalid")
+    managed = value.get("managed_inference")
+    training = value.get("training_external")
+    if training is not None:
+        if (managed is not None or not isinstance(training, dict)
+            or set(training) != {"binding_id", "training_run_id", "policy_version", "generation_contract_digest", "api", "max_output_tokens", "max_episode_steps"}
+            or training.get("api") != "chat-completions" or value.get("required") is not True
+            or any(not isinstance(training.get(k), str) or not training[k] for k in ("binding_id", "training_run_id", "policy_version"))
+            or re.fullmatch(r"sha256:[a-f0-9]{64}", str(training.get("generation_contract_digest"))) is None
+            or any(type(training.get(k)) is not int or training[k] <= 0 for k in ("max_output_tokens", "max_episode_steps"))):
+            raise ValueError("model_capture training identity is invalid")
+    if managed is not None and not _valid_managed_model_identity(managed):
+        raise ValueError("model_capture managed inference identity is invalid")
     return dict(value)
 
 
+def _valid_managed_model_identity(value: Any) -> bool:
+    if (not isinstance(value, dict) or not {"inference_id", "model_id"}.issubset(value)
+        or set(value) - {"inference_id", "model_id", "model_node"}
+        or any(not isinstance(value[k], str) or re.fullmatch(r"sha256:[a-f0-9]{64}", value[k]) is None for k in ("inference_id", "model_id"))):
+        return False
+    if "model_node" not in value:
+        return True
+    node = value["model_node"]
+    return (isinstance(node, dict) and set(node) == {"schema_version", "node_id", "generation", "runtime_digest", "launcher"}
+        and node["schema_version"] == "2" and node["launcher"] == "process"
+        and all(isinstance(node[k], str) and re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.-]{0,127}", node[k]) is not None for k in ("node_id", "generation"))
+        and isinstance(node["runtime_digest"], str) and re.fullmatch(r"sha256:[a-f0-9]{64}", node["runtime_digest"]) is not None)
+
+
 def canonical_manifest_json(manifest: dict[str, Any]) -> str:
     """Canonically encode the runtime identity `{ schema_version, node_range,
     entrypoints, files }` with sorted object keys and no insignificant

package/integrations/harbor/hitch_harbor_environment.py

@@ -31,6 +31,7 @@ _REQUIRED_LABELS = {
 _ALLOWED_LABELS = _REQUIRED_LABELS | {_LABEL_TASK}
 _COMPOSE_RESET_MARKER = "__HITCH_COMPOSE_RESET_NULL__"
 _GPU_OVERRIDE_PREFIX = "__HITCH_GPU_OVERRIDE_"
+_HOST_CREDENTIAL_HELPER_ENV = "HITCH_HOST_CREDENTIAL_HELPER_JSON"
 
 
 class HitchHarborDockerEnvironment(DockerEnvironment):
@@ -128,6 +129,25 @@ class HitchHarborDockerEnvironment(DockerEnvironment):
             return result.model_copy(update={"disable_internet": True, "dynamic_network_policy": False})
         return result
 
+    def _compose_env_vars(self, include_os_env=True):
+        result = super()._compose_env_vars(include_os_env=include_os_env)
+        # The host-only helper command is not a task Compose interpolation
+        # input. The agent reads it directly from its trusted host process.
+        result.pop(_HOST_CREDENTIAL_HELPER_ENV, None)
+        return result
+
+    async def hitch_exec_with_private_env(self, command, *, env, cwd=None, timeout_sec=None, user=None):
+        """Execute one Target without placing per-exec credential values in argv."""
+        from hitch_private_exec import exec_with_private_environment
+        return await exec_with_private_environment(
+            self,
+            command,
+            env=env,
+            cwd=cwd,
+            timeout_sec=timeout_sec,
+            user=user,
+        )
+
     async def _apply_network_policy(self, network_policy):
         if self._hitch_static_no_network:
             if network_policy != self.network_policy:
@@ -375,7 +395,8 @@ def _validate_labels(value: Mapping[str, str] | None) -> dict[str, str]:
     labels = dict(value)
     if (
         not re.fullmatch(r"[a-f0-9]{24}", labels.get(_LABEL_ROOT, ""))
-        or labels.get(_LABEL_PROVIDER) != "local-docker"
+        or not isinstance(labels.get(_LABEL_PROVIDER), str)
+        or not re.fullmatch(r"[a-z0-9][a-z0-9._-]{0,127}", labels[_LABEL_PROVIDER])
         or not re.fullmatch(r"eval_[a-f0-9]{32}", labels.get(_LABEL_EVAL, ""))
         or not re.fullmatch(r"work_[a-f0-9]{32}", labels.get(_LABEL_WORK, ""))
         or not re.fullmatch(r"lease_[a-f0-9]{32}", labels.get(_LABEL_LEASE, ""))