feat(spec): resolve agent instructions from MLflow Prompt Registry

omnigent/inner/loader.py

@@ -76,6 +76,7 @@ def load_agent_def(
     path_or_dict: str | Path | YamlData,
     *,
     enforce_handler_allowlist: bool = False,
+    expand_env: bool = True,
 ) -> AgentDef:
     """Load an AgentDef from a YAML file path or a raw dict.
 
@@ -102,6 +103,12 @@ def load_agent_def(
         ``omnigent run``, operator specs, the CLI) keep working with
         custom handlers — the operator already has code execution, so
         the restriction would add no security there.
+    :param expand_env: Whether this is a runtime/deploy load (``True``,
+        the default) versus a scaffolding/validation load (``False``).
+        Forwarded to :func:`_resolve_instructions` to gate the MLflow
+        Prompt Registry fetch: untrusted / HTTP-uploaded bundles are
+        validated with ``expand_env=False`` and must not trigger a
+        registry network fetch during validation.
     """
     if isinstance(path_or_dict, (str, Path)):
         path = Path(path_or_dict)
@@ -113,7 +120,7 @@ def load_agent_def(
         instructions_root = None
     if enforce_handler_allowlist:
         _reject_unregistered_policy_handlers(data)
-    return _parse_agent_def(data, instructions_root=instructions_root)
+    return _parse_agent_def(data, instructions_root=instructions_root, expand_env=expand_env)
 
 
 def _reject_unregistered_policy_handlers(data: YamlData) -> None:
@@ -186,6 +193,8 @@ def _read_contained_file(root: Path, value: str) -> str | None:
 def _resolve_instructions(
     raw_value: object,
     instructions_root: Path | None,
+    *,
+    expand_env: bool = True,
 ) -> str | None:
     """
     Resolve the ``instructions:`` field to a system-prompt string.
@@ -194,24 +203,56 @@ def _resolve_instructions(
     so omnigent-flavored YAMLs and native Omnigent YAMLs treat the
     field identically:
 
+    - MLflow Prompt Registry reference (``source: mlflow`` mapping or
+      ``mlflow+prompts:/...`` string): the prompt text is loaded from
+      the registry, but only when *expand_env* is ``True`` (see below).
     - Single-line value that names an existing file relative to
       *instructions_root*: the file's contents are read.
     - Multi-line value (contains ``\\n``): treated as inline text.
     - File-path-shaped value that doesn't resolve: silently
       treated as inline text. Matches native Omnigent behavior so users
       who type ``instructions: AGENTS.md`` and forget the file
       get the literal string back rather than a misleading error.
-    - ``None`` / non-string raw values return ``None``.
+    - ``None`` / non-string raw values (other than an MLflow mapping)
+      return ``None``.
 
     :param raw_value: The raw YAML value at the ``instructions:``
         key, or ``None`` if the key was absent.
     :param instructions_root: Directory to anchor relative path
         lookups. ``None`` skips the file-read attempt and treats
         every value as inline text — used when the loader has no
         on-disk anchor (raw-dict input path).
+    :param expand_env: Whether this is a runtime/deploy parse
+        (``True``) versus a scaffolding/validation parse (``False``).
+        It gates the MLflow registry fetch the same way
+        :func:`omnigent.spec.parser._resolve_instructions` does: an
+        untrusted / HTTP-uploaded bundle is validated with
+        ``expand_env=False`` and must not trigger a network fetch on
+        the author's behalf, so an MLflow reference is left as its
+        literal string.
     :returns: The resolved instruction text, or ``None`` if no
         ``instructions:`` was supplied.
     """
+    # Reuse the single canonical MLflow helper from the spec layer
+    # (imported lazily to avoid an import-time cycle: the spec package
+    # imports this loader module) so the two parser copies can never
+    # drift on registry-resolution rules.
+    from omnigent.spec.mlflow_prompts import (
+        parse_mlflow_instructions,
+        resolve_mlflow_prompt,
+    )
+
+    mlflow_ref = parse_mlflow_instructions(raw_value)
+    if mlflow_ref is not None:
+        if not expand_env:
+            return mlflow_ref.reference
+        return resolve_mlflow_prompt(
+            mlflow_ref.reference,
+            tracking_uri=mlflow_ref.tracking_uri,
+            registry_uri=mlflow_ref.registry_uri,
+            vars=mlflow_ref.vars,
+            cache_ttl=mlflow_ref.cache_ttl,
+        )
     if raw_value is None:
         return None
     if not isinstance(raw_value, str):
@@ -229,18 +270,22 @@ def _parse_agent_def(
     data: YamlData,
     *,
     instructions_root: Path | None = None,
+    expand_env: bool = True,
 ) -> AgentDef:
     agent = AgentDef()
     # ``AgentDef.name`` and ``AgentDef.prompt`` are both ``str | None``;
     # missing YAML keys flow through as ``None``.
     agent.name = data.get("name")
     agent.prompt = data.get("prompt")
     # ``instructions:`` is resolved into the agent's system-prompt
-    # text right here (file path → contents, or inline string).
-    # Translator code downstream only sees the resolved string
-    # in ``agent.instructions`` and the raw user-supplied
-    # ``agent.prompt`` — it doesn't have to re-walk the path.
-    agent.instructions = _resolve_instructions(data.get("instructions"), instructions_root)
+    # text right here (file path → contents, inline string, or MLflow
+    # Prompt Registry text). Translator code downstream only sees the
+    # resolved string in ``agent.instructions`` and the raw
+    # user-supplied ``agent.prompt`` — it doesn't have to re-walk the
+    # path.
+    agent.instructions = _resolve_instructions(
+        data.get("instructions"), instructions_root, expand_env=expand_env
+    )
     agent.input_type = data.get("input_type")
     agent.output_type = data.get("output_type")
     if "async_enabled" in data:

omnigent/spec/__init__.py

@@ -247,6 +247,7 @@ def load(
                 source,
                 enforce_handler_allowlist=enforce_handler_allowlist,
                 prune_invalid_sub_agents=prune_invalid_sub_agents,
+                expand_env=expand_env,
             )
         if source.suffix.lower() in {".yaml", ".yml"}:
             # The path is a YAML file but failed the omnigent check
@@ -281,13 +282,16 @@ def load(
     # 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.
+    # *expand_env* for env expansion. The flag is still forwarded so
+    # ``expand_env=False`` (untrusted upload validation) also gates the
+    # MLflow Prompt Registry fetch for an ``instructions:`` reference.
     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,
+            expand_env=expand_env,
         )
 
     spec = parse(root, expand_env=expand_env)

omnigent/spec/_omnigent_compat.py

@@ -310,6 +310,7 @@ def load_omnigent_yaml(
     *,
     enforce_handler_allowlist: bool = False,
     prune_invalid_sub_agents: bool = False,
+    expand_env: bool = True,
 ) -> AgentSpec:
     """
     Load an omnigent YAML and translate it to an
@@ -334,6 +335,13 @@ def load_omnigent_yaml(
         WARNING logged per drop. The root agent must still validate.
         See :func:`omnigent.spec.load` for the full rationale — this
         is the execution-path backwards-compatibility guard.
+    :param expand_env: Forwarded to
+        :func:`omnigent.inner.loader.load_agent_def` to gate the MLflow
+        Prompt Registry fetch. ``False`` (untrusted / HTTP-uploaded
+        bundle validation) leaves an MLflow ``instructions:`` reference
+        as its literal string instead of contacting the registry. Does
+        not affect ``${VAR}`` handling on this path — the omnigent
+        loader never expands env vars.
     :returns: A validated :class:`AgentSpec` with
         ``executor.type == OMNIGENT_EXECUTOR_TYPE``.
     :raises OmnigentError: If the synthesized spec fails
@@ -365,7 +373,11 @@ def load_omnigent_yaml(
     from omnigent.spec.omnigent import agent_def_to_agent_spec
     from omnigent.spec.validator import validate
 
-    agent_def = load_agent_def(path, enforce_handler_allowlist=enforce_handler_allowlist)
+    agent_def = load_agent_def(
+        path,
+        enforce_handler_allowlist=enforce_handler_allowlist,
+        expand_env=expand_env,
+    )
     # Read the raw YAML alongside so the translator can preserve
     # policy-level YAML fields that the omnigent loader drops
     # (label policies in particular compile to synthetic