diff --git a/orchestrator/agent_launcher.py b/orchestrator/agent_launcher.py index 0b8bb84..78616bd 100644 --- a/orchestrator/agent_launcher.py +++ b/orchestrator/agent_launcher.py @@ -123,14 +123,82 @@ def _extract_model_key(secret_value: str) -> Optional[str]: def _resolve_secret_ref(secret_ref: str) -> Optional[str]: """Resolve an azkv:// ref to the model key. - Production: an Azure Key Vault adapter (out of this repo; wire via deployment) reads the secret - named by the ref's trailing segment and returns its value. Dev/CI: read env - `HEICODE_SECRET_`. Either way the value is the HM #60 JSON ``{"openai_api_key": "sk-..."}`` - (bare string also accepted). Returns None when unavailable — we never fabricate a key. + Order: (1) dev/CI env map ``HEICODE_SECRET_`` — keeps tests/dev hermetic and offline; + (2) production Azure Key Vault read via the pod's **workload identity** (``_resolve_from_keyvault``, + agent_swarm#56). Either source yields the HM #60 JSON ``{"openai_api_key": "sk-..."}`` (bare string + also accepted). Returns None when unavailable — we never fabricate a key. """ name = secret_ref.rstrip("/").rsplit("/", 1)[-1] raw = os.getenv(f"HEICODE_SECRET_{name}") - return _extract_model_key(raw) if raw is not None else None + if raw is not None: + return _extract_model_key(raw) + return _resolve_from_keyvault(secret_ref) + + +def _azkv_enabled() -> bool: + """Whether to attempt a real Key Vault read. + + Only when the orchestrator pod has **workload identity** injected (the AKS webhook sets + ``AZURE_FEDERATED_TOKEN_FILE`` when the SA is annotated + the pod is labelled + ``azure.workload.identity/use: "true"``), or a deployment explicitly opts in with + ``SECRET_RESOLVER=azkv``. Keeps dev/CI/tests hermetic: without these the resolver never imports + the azure SDK and never touches the network. + """ + return bool( + os.getenv("AZURE_FEDERATED_TOKEN_FILE") + or os.getenv("SECRET_RESOLVER", "").strip().lower() == "azkv" + ) + + +def _parse_azkv_ref(secret_ref: str): + """``azkv:///secrets/[/]`` → ``(vault_url, secret_name, version|None)``. + + ```` may be a bare name (→ ``https://.vault.azure.net``) or a full host. Also + tolerates the short ``azkv:///`` form. Returns None if it can't parse. + """ + rest = secret_ref[len("azkv://"):].strip("/") if secret_ref.startswith("azkv://") else "" + parts = [p for p in rest.split("/") if p] + if len(parts) >= 3 and parts[1] == "secrets": + host, name, version = parts[0], parts[2], (parts[3] if len(parts) > 3 else None) + elif len(parts) == 2: + host, name, version = parts[0], parts[1], None + else: + return None + if not host or not name: + return None + vault_url = host if host.startswith("http") else ( + f"https://{host}" if "." in host else f"https://{host}.vault.azure.net") + return vault_url, name, version + + +def _resolve_from_keyvault(secret_ref: str) -> Optional[str]: + """Read the model key from Azure Key Vault using the pod's workload identity. + + `DefaultAzureCredential` picks up the federated token the AKS workload-identity webhook injects + (see `_azkv_enabled`). Lazy-imports the azure SDK so dev/CI without it are unaffected. Returns + None on ANY failure (not enabled / unparseable ref / SDK missing / no credential / network / + secret absent) — never fabricates, never raises. + """ + if not _azkv_enabled(): + return None + parsed = _parse_azkv_ref(secret_ref) + if not parsed: + logger.warning("azkv resolver: unparseable secret_ref") + return None + vault_url, name, version = parsed + try: + from azure.identity import DefaultAzureCredential + from azure.keyvault.secrets import SecretClient + except Exception as exc: # SDK not installed + logger.warning("azkv resolver: azure SDK unavailable (%s); add azure-identity + azure-keyvault-secrets", exc) + return None + try: + client = SecretClient(vault_url=vault_url, credential=DefaultAzureCredential()) + secret = client.get_secret(name, version) if version else client.get_secret(name) + return _extract_model_key(secret.value) + except Exception as exc: # no credential / RBAC / network / missing secret + logger.warning("azkv resolver: failed to read secret '%s' from %s: %s", name, vault_url, exc) + return None def model_api_base() -> str: diff --git a/orchestrator/requirements.txt b/orchestrator/requirements.txt index 09f6622..01d6f8d 100644 --- a/orchestrator/requirements.txt +++ b/orchestrator/requirements.txt @@ -19,3 +19,6 @@ opentelemetry-instrumentation-logging==0.45b0 azure-cosmos==4.7.0 azure-storage-blob==12.23.1 azure-identity==1.19.0 +# model-key resolution from Key Vault via Pod workload identity (lazy-imported in +# agent_launcher._resolve_from_keyvault; only when AZURE_FEDERATED_TOKEN_FILE / SECRET_RESOLVER=azkv) +azure-keyvault-secrets==4.9.0 diff --git a/scripts/test-agent-launcher.py b/scripts/test-agent-launcher.py index 6eb1257..14bfa8d 100644 --- a/scripts/test-agent-launcher.py +++ b/scripts/test-agent-launcher.py @@ -80,6 +80,36 @@ def test_resolve_model_key(): check("unresolved -> None (never fabricated)", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/missing"}}) is None) +def test_azkv_resolver(): + # ── parse azkv:// refs (pure) ── + check("azkv parse: bare vault name -> https URL", + al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1") + == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None)) + check("azkv parse: full host preserved", + al._parse_azkv_ref("azkv://heicode-vault.vault.azure.net/secrets/res_git_1") + == ("https://heicode-vault.vault.azure.net", "res_git_1", None)) + check("azkv parse: version captured", + al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1/abc") + == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", "abc")) + check("azkv parse: short form azkv:///", + al._parse_azkv_ref("azkv://heicode-vault/swarm-model-key-u1") + == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None)) + check("azkv parse: bad ref -> None", al._parse_azkv_ref("azkv://heicode-vault") is None) + + # ── gating: hermetic unless workload identity injected / opt-in ── + for k in ("AZURE_FEDERATED_TOKEN_FILE", "SECRET_RESOLVER"): + os.environ.pop(k, None) + check("azkv disabled without workload identity / opt-in", al._azkv_enabled() is False) + check("disabled -> _resolve_from_keyvault returns None (no SDK/network touched)", + al._resolve_from_keyvault("azkv://heicode-vault/secrets/swarm-model-key-u1") is None) + os.environ["SECRET_RESOLVER"] = "azkv" + check("azkv enabled via SECRET_RESOLVER=azkv", al._azkv_enabled() is True) + os.environ.pop("SECRET_RESOLVER") + os.environ["AZURE_FEDERATED_TOKEN_FILE"] = "/var/run/secrets/azure/tokens/azure-identity-token" + check("azkv enabled via AZURE_FEDERATED_TOKEN_FILE", al._azkv_enabled() is True) + os.environ.pop("AZURE_FEDERATED_TOKEN_FILE") + + def test_command_backend_build(): os.environ["AGENT_LAUNCH_CMD"] = "launch-agent.sh --id {agent_id} --caps {capabilities}" spec = al.AgentLaunchSpec(agent_id="swarm-abc-agent-1", capabilities="python,general", @@ -143,6 +173,7 @@ def main(): test_launch_count() test_plan_specs() test_resolve_model_key() + test_azkv_resolver() test_command_backend_build() test_k8s_manifests() asyncio.run(test_backend_none_noop())