"""Agent image spec: parsing, validation, and safe extraction.""" from __future__ import annotations import logging import shutil from pathlib import Path from omnigent.errors import ErrorCode, OmnigentError # Omnigent compat: imported surgically from a dedicated module so # the integration's tech debt is removable in one shot. See # omnigent/spec/_omnigent_compat.py. from omnigent.spec._omnigent_compat import ( diagnose_yaml_rejection, is_omnigent_yaml, load_omnigent_yaml, ) from omnigent.spec.parser import expand_env_vars, parse, parse_default_policies, parse_server_llm from omnigent.spec.tar_utils import ExtractionError, extract_safe from omnigent.spec.types import ( DEFAULT_ASK_TIMEOUT, DEFAULT_POLICY_CLASSIFIER_TIMEOUT, AgentSpec, ApiKeyAuth, BuiltinToolConfig, DatabricksAuth, ExecutorAuth, ExecutorSpec, FunctionPolicySpec, FunctionRef, GuardrailsSpec, InteractionConfig, LabelDef, LLMConfig, LocalToolInfo, MCPServerConfig, ModalityConfig, Phase, PhaseSelector, PolicyAction, PolicySpec, ProviderAuth, RetryPolicy, SkillSpec, ToolRuntime, ToolsConfig, ) from omnigent.spec.validator import ValidationResult, validate _logger = logging.getLogger(__name__) __all__ = [ "DEFAULT_ASK_TIMEOUT", "DEFAULT_POLICY_CLASSIFIER_TIMEOUT", "AgentSpec", "ApiKeyAuth", "BuiltinToolConfig", "DatabricksAuth", "ExecutorAuth", "ExecutorSpec", "ExtractionError", "FunctionPolicySpec", "FunctionRef", "GuardrailsSpec", "InteractionConfig", "LLMConfig", "LabelDef", "LocalToolInfo", "MCPServerConfig", "ModalityConfig", "Phase", "PhaseSelector", "PolicyAction", "PolicySpec", "ProviderAuth", "RetryPolicy", "SkillSpec", "ToolRuntime", "ToolsConfig", "ValidationResult", "expand_env_vars", "extract_safe", "load", "materialize_bundle", "parse", "parse_default_policies", "parse_server_llm", "validate", ] def materialize_bundle(source: Path, dest: Path) -> Path: """ Copy a spec source into *dest* as a uniform bundle directory. Agent-plane accepts two source shapes: an agent-image directory (``config.yaml`` + bundled assets) or a standalone omnigent YAML file. Downstream code (``_preregister_agent``, ``_prepare_omnigent_yaml_bundle``) always wants to operate on a directory it can tar, mutate, or hand to the in-process omnigent server. Taking the file-vs-directory branch once — here — means every caller downstream is uniform: "materialize, then operate on the returned path as a directory." No caller has to reinspect the input shape. :param source: The spec source. Either a directory containing ``config.yaml`` (standard omnigent shape) or a standalone omnigent YAML file (e.g. ``examples/coding_supervisor.yaml``). Must exist. :param dest: Destination directory to populate. Created if it does not exist; may be empty or already contain the copied contents from a prior call (``shutil.copytree`` is invoked with ``dirs_exist_ok=True``). :returns: *dest*, always as a populated directory. For the directory case the contents are a recursive copy of *source*. For the file case the YAML is placed at the root of *dest* so :func:`omnigent.spec._find_omnigent_yaml_in_dir` picks it up on a subsequent :func:`load` call. :raises FileNotFoundError: If *source* does not exist. """ if source.is_dir(): shutil.copytree(source, dest, dirs_exist_ok=True) return dest if source.is_file(): dest.mkdir(parents=True, exist_ok=True) shutil.copy2(source, dest / source.name) return dest raise FileNotFoundError(f"source not found: {source}") def _find_omnigent_yaml_in_dir(root: Path) -> Path | None: """ Return the omnigent YAML inside *root* when the directory is a single-file omnigent bundle, or ``None`` otherwise. A directory qualifies as a single-file omnigent bundle when there is no ``config.yaml`` at the root and exactly one file at the root is recognised by :func:`is_omnigent_yaml`. Any other shape (``config.yaml`` present, zero or multiple omnigent YAMLs, YAMLs in subdirectories) returns ``None`` and the caller falls through to the standard omnigent directory-parse path. :param root: Path to an extracted bundle directory. :returns: The matched omnigent YAML path, or ``None`` if the directory is not a single-file omnigent bundle. """ if (root / "config.yaml").exists(): return None matches = [p for p in root.iterdir() if p.is_file() and is_omnigent_yaml(p)] if len(matches) == 1: return matches[0] return None def load( source: Path | bytes, *, dest: Path | None = None, expand_env: bool = True, enforce_handler_allowlist: bool = False, prune_invalid_sub_agents: bool = False, ) -> AgentSpec: """ Load an agent spec from a directory, tarball path, or raw bytes. If *source* is a directory, parse and validate it directly. If *source* is a file path (tarball or omnigent YAML) or raw bytes, dispatch accordingly. :param source: Path to an agent image directory, ``.tar.gz`` bundle, omnigent single-file YAML, or raw tarball bytes (e.g. from an HTTP upload). :param dest: Extraction destination -- required when *source* is a tarball or bytes, ignored when *source* is a directory or omnigent YAML. :param expand_env: Whether to expand ``${VAR}`` references in connection blocks, MCP headers, and MCP ``env`` against the current process environment. ``True`` (the default) is for operator-authored specs whose author is the process owner (local ``omnigent run``, ``--agent`` preregistration). It MUST be ``False`` for tenant-supplied / HTTP-uploaded bundles: expanding their ``${VAR}`` against the server or runner process environment leaks server-side secrets into a spec-controlled MCP/LLM connection. See ``.claude/skills/code-review/security-guidelines.md``. :param enforce_handler_allowlist: When ``True``, reject any ``type: function`` policy whose handler dotted path is not a registered policy handler. This is the guard for the untrusted agent-bundle upload path — set by :func:`omnigent.server.bundles.validate_agent_bundle`. It is applied for both bundle shapes: the omnigent single-file YAML path (before the inner loader resolves/calls the handler at parse time) and the ``config.yaml`` path (post-parse, since that parser does not resolve handlers). Defaults to ``False`` so trusted spec loading (local ``omnigent run``, operator configs) keeps supporting custom handlers. :param prune_invalid_sub_agents: When ``True``, a sub-agent that fails validation is **dropped** from the spec (removed from ``sub_agents`` and from any parent's ``tools.agents`` reference) and loading continues, instead of failing the whole load. A WARNING is logged for each dropped sub-agent. The root agent must still validate — a genuine root-level error always raises. This is the backwards-compatibility guard for the **execution** paths (runner spec resolution, server :class:`~omnigent.runtime.agent_cache.AgentCache`): a bundle reaching those paths was already validated by the server that produced it, so a sub-agent that fails *here* means this client is older than that server and can't run that sub-agent (e.g. it names a harness this version doesn't know). Dropping the sub-agent lets the parent agent launch with the capabilities this client *does* support, rather than the whole agent failing to start. Defaults to ``False`` so authoring/upload paths (``omnigent run``, :func:`omnigent.server.bundles.validate_agent_bundle`) stay strict and surface real authoring mistakes to the author. :returns: A validated :class:`AgentSpec`. :raises OmnigentError: If the spec fails validation, if a policy names an unregistered handler under *enforce_handler_allowlist*, or if *source* is a tarball/bytes and *dest* is not provided. :raises FileNotFoundError: If *source* is a :class:`Path` that does not exist, or if the extracted directory is missing ``config.yaml``. :raises ExtractionError: If the tarball fails safety checks. """ if isinstance(source, bytes): if dest is None: raise OmnigentError( "dest is required when loading from bytes", code=ErrorCode.INVALID_INPUT, ) extract_safe(source, dest) root = dest elif source.is_dir(): root = source elif source.is_file(): # Omnigent single-file YAML dispatch — see # omnigent.spec._omnigent_compat. Tech-debt aside; # remove this branch when omnigent compat ends. if is_omnigent_yaml(source): return load_omnigent_yaml( source, enforce_handler_allowlist=enforce_handler_allowlist, prune_invalid_sub_agents=prune_invalid_sub_agents, ) if source.suffix.lower() in {".yaml", ".yml"}: # The path is a YAML file but failed the omnigent check # (missing required key, ``spec_version`` set, malformed, # etc.). Falling through to the tarball-extraction branch # would surface the misleading "dest is required when # loading from a tarball" — the file's a YAML, not a # tarball. Diagnose the actual reason and surface it. reason = diagnose_yaml_rejection(source) raise OmnigentError( f"{source}: not a valid omnigent YAML — {reason}", code=ErrorCode.INVALID_INPUT, ) if dest is None: raise OmnigentError( "dest is required when loading from a tarball", code=ErrorCode.INVALID_INPUT, ) extract_safe(source, dest) root = dest else: raise FileNotFoundError(f"source not found: {source}") # Omnigent single-file YAML dispatch (extracted-bundle variant) # — when a bundle (directory, tarball, or raw bytes) resolves to # a root that contains exactly one omnigent YAML and no # ``config.yaml``, route to the omnigent adapter. This is the # shape produced by ``omnigent server --agent`` when it wraps a # YAML into a single-file tarball for the agent store. See # omnigent/spec/_omnigent_compat.py for the tech-debt # ownership; remove when omnigent compat ends. # The omnigent single-file path (load_omnigent_yaml) copies # headers / connection verbatim and never expands ``${VAR}`` # against the process env, so it is safe regardless of # *expand_env* and the flag does not apply. candidate = _find_omnigent_yaml_in_dir(root) if candidate is not None: return load_omnigent_yaml( candidate, enforce_handler_allowlist=enforce_handler_allowlist, prune_invalid_sub_agents=prune_invalid_sub_agents, ) spec = parse(root, expand_env=expand_env) if prune_invalid_sub_agents: _prune_invalid_sub_agents(spec) result = validate(spec) if not result.valid: errors = "; ".join(f"{e.path}: {e.message}" for e in result.errors) raise OmnigentError( f"invalid agent spec: {errors}", code=ErrorCode.INVALID_INPUT, ) if enforce_handler_allowlist: # config.yaml path: the parser stores handler dotted paths as # strings without resolving them (resolution happens later at # engine build), so a post-parse scan is safe and sufficient — # it stops an unregistered handler from ever being stored/run. # The single-file omnigent YAML path is guarded earlier, inside # the loader, because that loader executes factories at parse. _reject_unregistered_spec_policy_handlers(spec) return spec def _prune_invalid_sub_agents(spec: AgentSpec) -> list[str]: """Drop sub-agents that fail validation so the parent can still load. Walks *spec*'s sub-agent tree depth-first and removes any sub-agent whose own subtree fails :func:`validate`, mutating *spec* in place: the failing sub-agent is removed from ``sub_agents`` and its name (if any) is removed from the parent's ``tools.agents`` reference list so the dangling reference does not itself fail validation. A WARNING is logged for each drop. Depth-first ordering matters: a child's invalid *grandchild* is pruned before the child is judged, so a single bad leaf does not take out an otherwise-valid sub-tree above it. This is the mechanism behind ``load(..., prune_invalid_sub_agents= True)`` — see :func:`load` for when and why it is enabled (the execution paths only). It deliberately does **not** touch the root spec: if the root itself is invalid, the caller's subsequent :func:`validate` still fails loud. :param spec: The spec to prune in place. :returns: The names (or ``""``) of every sub-agent dropped anywhere in the tree, parent-most first within each level. """ dropped: list[str] = [] surviving: list[AgentSpec] = [] for sa in spec.sub_agents: # Prune this child's own invalid descendants before deciding # whether the child itself survives. dropped.extend(_prune_invalid_sub_agents(sa)) sa_result = validate(sa) if sa_result.valid: surviving.append(sa) continue name = sa.name or "" errors = "; ".join(f"{e.path}: {e.message}" for e in sa_result.errors) _logger.warning( "Dropping sub-agent %r from agent %r: it failed validation on this " "client and will be unavailable. This usually means the spec was " "produced by a newer Omnigent server with a feature this client does " "not support (e.g. a harness it does not recognize) — upgrade this " "client to use it. Validation errors: %s", name, spec.name or "", errors, ) dropped.append(name) # Remove the now-dangling reference so the parent still validates. if sa.name is not None and sa.name in spec.tools.agents: spec.tools.agents.remove(sa.name) spec.sub_agents = surviving return dropped def _reject_unregistered_spec_policy_handlers(spec: AgentSpec) -> None: """Reject function policies whose handler is not registered. Scans a parsed :class:`AgentSpec`'s guardrail policies for :class:`~omnigent.spec.types.FunctionPolicySpec` entries whose ``function.path`` is not in the policy registry. Recurses into ``sub_agents`` — the ``config.yaml`` parser discovers child agents from ``agents/`` subdirectories, each with its own ``guardrails``, and those handlers are resolved + called at engine build just like the root's, so a clean root with a malicious sub-agent would otherwise bypass the upload allowlist. Used on the untrusted agent-bundle upload path only (see :func:`load`). :param spec: The parsed agent spec (or sub-agent) to scan. :raises OmnigentError: If a function policy names an unregistered handler, e.g. ``"subprocess.Popen"``. """ from omnigent.policies.registry import is_registered_handler guardrails = spec.guardrails if guardrails is not None: for policy in guardrails.policies or []: if ( isinstance(policy, FunctionPolicySpec) and policy.function is not None and not is_registered_handler(policy.function.path) ): raise OmnigentError( f"Policy {policy.name!r}: handler {policy.function.path!r} is not a " f"registered policy handler. Uploaded agent bundles may only use " f"handlers from the policy registry; a server admin must add custom " f"handlers via the 'policy_modules' config.", code=ErrorCode.INVALID_INPUT, ) for sub_agent in spec.sub_agents: _reject_unregistered_spec_policy_handlers(sub_agent)