diff --git a/CHANGELOG.md b/CHANGELOG.md index b753fa8..fa1dbe7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -52,6 +52,13 @@ Every call the SDK makes now matches a route the CAS egress proxy serves (plan engine's `PreToolUse` hook passes it, `PolicyToolAuthorizer` puts it in the policy input as `call_id`, and a call without one is denied. - A rewritten message carries the gateway's masked text (`[EMAIL]`). +- **The Lyzr key is read at every turn.** The engine's `ANTHROPIC_API_KEY`, + `OPENAI_API_KEY` and `LYZR_API_KEY` and the bound MCP servers' `x-api-key` + are set from `Settings.api_key()` when each turn starts (and at boot), so a + key replaced in `LYZR_API_KEY` between turns, or a reissued token, is used + by the next turn. A missing key fails the turn with `SettingsError` before + the engine starts. `build_gateway_envs` no longer takes `api_key`; + `key_envs(api_key)` returns the three variables. - **BREAKING: the agent document is read from its gateway sections.** The model is `llm.config.model` (the gateway agent's kong model; `null` → the image's own model) and the MCP servers are `mcp.servers`, every one attached diff --git a/docs/USING_THE_SDK.md b/docs/USING_THE_SDK.md index dcb1d18..42e8ff1 100644 --- a/docs/USING_THE_SDK.md +++ b/docs/USING_THE_SDK.md @@ -105,7 +105,7 @@ request is made. A URL setting must be an `http(s)` URL. |---|---|---|---| | `agent_id` | `COMPUTERAGENT_AGENT_ID` | **yes** | The agent's **24-hex Mongo `_id`** in AgentOS (the id in the agent's URL). Not the stable `agentId` field. | | `engine` | `COMPUTERAGENT_ENGINE` | **yes** | `"claude-agent-sdk"` or `"gitagent"`; an `EngineDriver` instance only as an argument. | -| `api_key` | `LYZR_API_KEY` | **yes** | The key the SDK presents on every call as `x-api-key`. In a sandbox it is injected at creation and the egress proxy swaps it for the real one. Read from the environment **at each use**, so it can be replaced while the agent runs. | +| `api_key` | `LYZR_API_KEY` | **yes** | The key the SDK presents on every call as `x-api-key`. In a sandbox it is injected at creation and the egress proxy swaps it for the real one. Read **at every turn** (and at each other use): the engine's key variables and the bound MCP servers' `x-api-key` follow it, so the app server can set it from each message's `credential`. An explicit `api_key=` is kept and replaced only by its own reissued successors (§8a). | | `base_url` | `COMPUTERAGENT_BASE_URL` | **yes** | Route to the AgentOS agent config: the model and its MCP bindings. | | `guardrails_url` | `COMPUTERAGENT_GUARDRAILS_URL` | **yes** | Route to the guardrails check, for the user's message and every tool call. | | `llm_url` | `COMPUTERAGENT_LLM_URL` | **yes** | Route for every model call the engine makes. | @@ -318,7 +318,7 @@ Every engine gets the same environment: ANTHROPIC_BASE_URL = llm_url GITCLAW_MODEL_BASE_URL = llm_url GITAGENT_MODEL_BASE_URL = llm_url -ANTHROPIC_API_KEY = OPENAI_API_KEY = LYZR_API_KEY = api_key +ANTHROPIC_API_KEY = OPENAI_API_KEY = LYZR_API_KEY = api_key # read at every turn CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS = 1 CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC = 1 DISABLE_AUTOUPDATER = 1 diff --git a/src/computeragent/harness/__init__.py b/src/computeragent/harness/__init__.py index 504dabe..055b48e 100644 --- a/src/computeragent/harness/__init__.py +++ b/src/computeragent/harness/__init__.py @@ -342,9 +342,7 @@ def __init__( ) # 3. Where the model calls go: always llm_url. - gateway_envs = build_gateway_envs( - api_key=self._settings.api_key(), llm_url=self._settings.llm_url - ) + gateway_envs = build_gateway_envs(llm_url=self._settings.llm_url) merged_envs = merge_envs(envs, gateway_envs) merged_options = build_options(self.agent, engine_name=engine_name, caller_options=options) # `budget_max_usd` used to be carried to the engine context and read by @@ -459,6 +457,9 @@ def __init__( mcp_allowed_tools=mcp_allowed_tools, turn_guard=turn_guard, forward_cli=engine_name == "claude-agent-sdk", + # The key is read at every turn: the engine's key variables and + # the bound MCP servers' x-api-key follow it. + api_key=self._settings.api_key, ) async def __aenter__(self) -> ComputerAgent: diff --git a/src/computeragent/harness/coordinator.py b/src/computeragent/harness/coordinator.py index 0b213b8..216f23f 100644 --- a/src/computeragent/harness/coordinator.py +++ b/src/computeragent/harness/coordinator.py @@ -134,6 +134,9 @@ def __init__( # through a loopback forwarder (see `egress_forwarder`), so they carry # the newest sandbox token and their answers are observed. forward_cli: bool = False, + # Read at every turn (and at boot): the engine's ANTHROPIC_API_KEY / + # OPENAI_API_KEY / LYZR_API_KEY and the bound MCP servers' x-api-key. + api_key: Callable[[], str] | None = None, ) -> None: self._engine = engine self._substrate = substrate @@ -162,6 +165,7 @@ def __init__( self._mcp_allowed_tools: list[str] = list(mcp_allowed_tools or []) self._turn_guard = turn_guard self._forward_cli = forward_cli + self._api_key = api_key self._forwarder: Any = None # Substrate state — booted lazily on first chat. @@ -211,7 +215,9 @@ async def _ensure_booted(self) -> None: ) if self._substrate_handle is not None: return - self._substrate_handle = await self._substrate.boot(source=self._source, envs=self._envs) + self._substrate_handle = await self._substrate.boot( + source=self._source, envs=self._turn_envs() + ) self._workspace_path = self._substrate_handle.workspace_path loaded = await self._identity_loader.load( source=self._source, @@ -325,23 +331,55 @@ async def _emit_lifecycle_session_ended(self) -> None: # ── one chat turn ────────────────────────────────────────────────────── + def _turn_envs(self) -> dict[str, str]: + """The static envs with this turn's key variables over them.""" + if self._api_key is None: + return self._envs + from .gateway_envs import key_envs + + return {**self._envs, **key_envs(self._api_key())} + + def _turn_options(self, key: str | None) -> Any: + """The manifest options with this turn's key on the bound MCP servers. + A server the loader or the caller supplied is left as it is.""" + options = self._manifest.options + if key is None or not isinstance(options, dict) or not self._mcp_servers: + return options + servers = options.get("mcp_servers") + if not isinstance(servers, dict): + return options + updated = dict(servers) + for name, bound in self._mcp_servers.items(): + server = updated.get(name) + if isinstance(server, dict) and server.get("url") == bound.get("url"): + updated[name] = { + **server, + "headers": {**dict(server.get("headers") or {}), "x-api-key": key}, + } + return {**options, "mcp_servers": updated} + async def _engine_envs_and_manifest(self) -> tuple[dict[str, str], Any]: - """The env and manifest this turn's engine starts with: for the Claude - CLI, egress URLs point at the loopback forwarder.""" - if not self._forward_cli: - return self._envs, self._manifest + """The env and manifest this turn's engine starts with: the current key + in the env and on the bound MCP servers, and for the Claude CLI, egress + URLs pointed at the loopback forwarder.""" from dataclasses import replace - from .egress_forwarder import EgressForwarder + envs = self._turn_envs() + options = self._turn_options(envs.get("LYZR_API_KEY") if self._api_key else None) + if self._forward_cli: + from .egress_forwarder import EgressForwarder - if self._forwarder is None: - self._forwarder = EgressForwarder() - options = self._manifest.options - envs, rewritten = await self._forwarder.rewrite( - self._envs, options if isinstance(options, dict) else None - ) + if self._forwarder is None: + self._forwarder = EgressForwarder() + envs, rewritten = await self._forwarder.rewrite( + envs, options if isinstance(options, dict) else None + ) + if rewritten is not None: + options = rewritten manifest = ( - self._manifest if rewritten is None else replace(self._manifest, options=rewritten) + self._manifest + if options is self._manifest.options + else replace(self._manifest, options=options) ) return envs, manifest @@ -366,6 +404,9 @@ async def run_chat( await self._ensure_booted() assert self._workspace_path is not None assert self._manifest is not None + # This turn's key, read now: a missing key fails the turn before the + # guardrails check or the engine runs. + envs, manifest = await self._engine_envs_and_manifest() # Reuse the same id across every chat on this coordinator so a # multi-turn conversation lands as ONE session doc in Mongo (with @@ -394,7 +435,6 @@ async def run_chat( if blocked_reply is None: await _enqueue_user_message(state.user_queue, message) - envs, manifest = await self._engine_envs_and_manifest() ctx = EngineContext( session_id=session_id, manifest=manifest, diff --git a/src/computeragent/harness/gateway_envs.py b/src/computeragent/harness/gateway_envs.py index f6012e1..d778b83 100644 --- a/src/computeragent/harness/gateway_envs.py +++ b/src/computeragent/harness/gateway_envs.py @@ -60,8 +60,20 @@ def dispatch_model(agent: ResolvedAgent, *, engine_name: str) -> str | None: return strip_provider_prefix(model) -def build_gateway_envs(*, api_key: str, llm_url: str) -> dict[str, str]: - """Environment variables that send the engine's model calls to ``llm_url``.""" +def key_envs(api_key: str) -> dict[str, str]: + """The key variables for one turn: one credential under three names. The + engine presents it on every model call and the egress proxy swaps it for + the real key.""" + return { + "ANTHROPIC_API_KEY": api_key, + "OPENAI_API_KEY": api_key, + "LYZR_API_KEY": api_key, + } + + +def build_gateway_envs(*, llm_url: str) -> dict[str, str]: + """Environment variables that send the engine's model calls to ``llm_url``. + Holds no key: :func:`key_envs` is added at every turn.""" if not isinstance(llm_url, str) or not llm_url.startswith(("http://", "https://")): raise GatewayConfigError( f"llm_url must be an http(s) URL, got {llm_url!r}. Every model call is sent to it, " @@ -72,11 +84,6 @@ def build_gateway_envs(*, api_key: str, llm_url: str) -> dict[str, str]: "ANTHROPIC_BASE_URL": base, "GITAGENT_MODEL_BASE_URL": base, "GITCLAW_MODEL_BASE_URL": base, - # One credential under three names: the engine presents it on every - # model call and the egress proxy swaps it for the real key. - "ANTHROPIC_API_KEY": api_key, - "OPENAI_API_KEY": api_key, - "LYZR_API_KEY": api_key, # The Claude CLI otherwise injects Claude-Code-only experimental betas # (a `context_management` body field plus beta headers) that a # spec-compliant gateway rejects outright with diff --git a/tests/test_computeragent_init.py b/tests/test_computeragent_init.py index beb72d8..77bc5e8 100644 --- a/tests/test_computeragent_init.py +++ b/tests/test_computeragent_init.py @@ -246,10 +246,12 @@ def test_gap_dir_as_a_string_is_accepted(agent_doc, tmp_path): @respx.mock def test_gateway_envs_reach_the_coordinator(agent_doc): agent = build(agent_doc) - envs = agent._coordinator._envs # type: ignore[attr-defined] + coordinator = agent._coordinator # type: ignore[attr-defined] + envs = coordinator._envs - assert envs["ANTHROPIC_API_KEY"] == KEY - assert envs["LYZR_API_KEY"] == KEY + for name in ("ANTHROPIC_API_KEY", "OPENAI_API_KEY", "LYZR_API_KEY"): + assert name not in envs # added at every turn + assert coordinator._api_key() == KEY assert envs["CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS"] == "1" assert envs["ANTHROPIC_BASE_URL"] == LLM_URL diff --git a/tests/test_gateway_envs.py b/tests/test_gateway_envs.py index aec099a..fd3dfef 100644 --- a/tests/test_gateway_envs.py +++ b/tests/test_gateway_envs.py @@ -15,6 +15,7 @@ build_gateway_envs, build_options, dispatch_model, + key_envs, merge_envs, strip_provider_prefix, ) @@ -44,9 +45,6 @@ def expected_envs(base: str = LLM_URL) -> dict[str, str]: "ANTHROPIC_BASE_URL": base, "GITAGENT_MODEL_BASE_URL": base, "GITCLAW_MODEL_BASE_URL": base, - "ANTHROPIC_API_KEY": KEY, - "OPENAI_API_KEY": KEY, - "LYZR_API_KEY": KEY, "CLAUDE_CODE_DISABLE_EXPERIMENTAL_BETAS": "1", "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", "DISABLE_AUTOUPDATER": "1", @@ -54,12 +52,12 @@ def expected_envs(base: str = LLM_URL) -> dict[str, str]: def test_llm_url_envs_exactly(): - assert build_gateway_envs(api_key=KEY, llm_url=LLM_URL) == expected_envs() + assert build_gateway_envs(llm_url=LLM_URL) == expected_envs() def test_llm_url_is_passed_verbatim_with_no_proxy_or_version_segments(): """The egress proxy owns its routes: no /up/, no /v4 appended.""" - envs = build_gateway_envs(api_key=KEY, llm_url=LLM_URL) + envs = build_gateway_envs(llm_url=LLM_URL) for name in ("ANTHROPIC_BASE_URL", "GITCLAW_MODEL_BASE_URL", "GITAGENT_MODEL_BASE_URL"): assert envs[name] == LLM_URL @@ -68,13 +66,13 @@ def test_llm_url_is_passed_verbatim_with_no_proxy_or_version_segments(): def test_a_trailing_slash_on_llm_url_is_dropped(): - assert build_gateway_envs(api_key=KEY, llm_url=LLM_URL + "/") == expected_envs() + assert build_gateway_envs(llm_url=LLM_URL + "/") == expected_envs() def test_the_cli_background_traffic_is_switched_off(): """The Claude CLI otherwise talks to Anthropic and Datadog itself, around the egress proxy.""" - envs = build_gateway_envs(api_key=KEY, llm_url=LLM_URL) + envs = build_gateway_envs(llm_url=LLM_URL) assert envs["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] == "1" assert envs["DISABLE_AUTOUPDATER"] == "1" @@ -83,7 +81,7 @@ def test_the_cli_background_traffic_is_switched_off(): @pytest.mark.parametrize("bad", ["", "egress.internal/llm", "ftp://egress/llm", None]) def test_a_malformed_llm_url_is_refused(bad): with pytest.raises(GatewayConfigError, match="llm_url"): - build_gateway_envs(api_key=KEY, llm_url=bad) # type: ignore[arg-type] + build_gateway_envs(llm_url=bad) # type: ignore[arg-type] # ── model ──────────────────────────────────────────────────────────────────── @@ -126,14 +124,21 @@ def test_no_model_stays_none(): # ── env merge ──────────────────────────────────────────────────────────────── -def test_gateway_envs_beat_caller_envs(): +def test_key_envs(): + assert key_envs("k") == {"ANTHROPIC_API_KEY": "k", "OPENAI_API_KEY": "k", "LYZR_API_KEY": "k"} + assert not set(key_envs("k")) & set(build_gateway_envs(llm_url=LLM_URL)) + + +def test_gateway_and_turn_key_envs_beat_caller_envs(): merged = merge_envs( - {"ANTHROPIC_API_KEY": "caller-key", "MY_VAR": "kept"}, - build_gateway_envs(api_key=KEY, llm_url=LLM_URL), + {"ANTHROPIC_API_KEY": "caller-key", "ANTHROPIC_BASE_URL": "http://x", "MY_VAR": "kept"}, + build_gateway_envs(llm_url=LLM_URL), ) + turn = {**merged, **key_envs(KEY)} # what the coordinator hands the engine - assert merged["ANTHROPIC_API_KEY"] == KEY - assert merged["MY_VAR"] == "kept" + assert turn["ANTHROPIC_API_KEY"] == KEY + assert turn["ANTHROPIC_BASE_URL"] == LLM_URL + assert turn["MY_VAR"] == "kept" def test_bedrock_flag_is_dropped(): @@ -142,7 +147,7 @@ def test_bedrock_flag_is_dropped(): egress proxy.""" merged = merge_envs( {"CLAUDE_CODE_USE_BEDROCK": "1"}, - build_gateway_envs(api_key=KEY, llm_url=LLM_URL), + build_gateway_envs(llm_url=LLM_URL), ) assert "CLAUDE_CODE_USE_BEDROCK" not in merged diff --git a/tests/test_turn_key.py b/tests/test_turn_key.py new file mode 100644 index 0000000..f50f768 --- /dev/null +++ b/tests/test_turn_key.py @@ -0,0 +1,139 @@ +"""The Lyzr key is read at every turn: the engine's key variables and the +bound MCP servers' ``x-api-key`` follow it.""" + +from __future__ import annotations + +from collections.abc import AsyncIterator +from datetime import datetime, timedelta, timezone + +import httpx +import pytest +import respx + +from computeragent.harness import ComputerAgent +from computeragent.harness.egress_token import CREDENTIAL +from computeragent.harness.settings import SettingsError +from computeragent.protocol_types import EngineCapabilities, EngineContext, EngineEvent + +AGENT_ID = "6aa93f1b3817a13d3103d2c1" +BASE = "http://egress.internal/computer-agent-api/agentos/api" +URL = f"{BASE}/v1/agents/{AGENT_ID}" +LLM_URL = "http://egress.internal/llm" +MCP_URL = "http://egress.internal/mcp/v1" +GUARDRAILS = "http://egress.internal/guardrails" +KEY_NAMES = ("ANTHROPIC_API_KEY", "OPENAI_API_KEY", "LYZR_API_KEY") + + +class RecordingEngine: + # Not the Claude CLI's name, so the loopback forwarder stays out of it. + name = "recording" + capabilities = EngineCapabilities( + streaming_input=True, + partial_messages=False, + permission_callback=True, + sessions=True, + budget=False, + ) + + def __init__(self) -> None: + self.contexts: list[EngineContext] = [] + + async def start_session(self, ctx: EngineContext) -> AsyncIterator[EngineEvent]: + self.contexts.append(ctx) + async for _ in ctx.user_message_queue: + break + yield EngineEvent(kind="sdk_message", payload={"type": "assistant"}) + + +def build(agent_doc, engine, **kwargs) -> ComputerAgent: + agent_doc["mcp"] = {"servers": [{"server_id": "srv_1", "tool_name": "jira"}]} + respx.get(URL).mock(return_value=httpx.Response(200, json=agent_doc)) + respx.post(f"{GUARDRAILS}/v1/check").mock( + return_value=httpx.Response(200, json={"data": {"decision": "allow"}}) + ) + params: dict = { + "agent_id": AGENT_ID, + "base_url": BASE, + "engine": engine, + "llm_url": LLM_URL, + "mcp_url": MCP_URL, + "guardrails_url": GUARDRAILS, + } + params.update(kwargs) + return ComputerAgent(**params) + + +def keys_of(ctx: EngineContext) -> set[str]: + return {ctx.envs[name] for name in KEY_NAMES} + + +@respx.mock +async def test_engine_env_follows_the_key(agent_doc, monkeypatch): + monkeypatch.setenv("LYZR_API_KEY", "sbt_1") + engine = RecordingEngine() + + async with build(agent_doc, engine) as agent: + await agent.chat("one") + monkeypatch.setenv("LYZR_API_KEY", "sbt_2") + await agent.chat("two") + + first, second = engine.contexts + assert keys_of(first) == {"sbt_1"} + assert keys_of(second) == {"sbt_2"} + for name in ("ANTHROPIC_BASE_URL", "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"): + assert first.envs[name] == second.envs[name] + + +@respx.mock +async def test_mcp_headers_follow_the_key(agent_doc, monkeypatch): + monkeypatch.setenv("LYZR_API_KEY", "sbt_1") + engine = RecordingEngine() + mine = {"type": "http", "url": "http://elsewhere/mcp", "headers": {"x-api-key": "own"}} + + async with build(agent_doc, engine) as agent: + agent._coordinator._options_overlay["mcp_servers"] = { # type: ignore[attr-defined] + "jira": agent._coordinator._mcp_servers["jira"], # type: ignore[attr-defined] + "mine": mine, + } + await agent.chat("one") + monkeypatch.setenv("LYZR_API_KEY", "sbt_2") + await agent.chat("two") + + servers = engine.contexts[1].manifest.options["mcp_servers"] + assert servers["jira"]["headers"]["x-api-key"] == "sbt_2" + assert servers["mine"]["headers"]["x-api-key"] == "own" + assert ( + engine.contexts[0].manifest.options["mcp_servers"]["jira"]["headers"]["x-api-key"] + == "sbt_1" + ) + + +@respx.mock +async def test_explicit_api_key_follows_its_reissue(agent_doc, monkeypatch): + monkeypatch.setenv("LYZR_API_KEY", "sbt_env") + engine = RecordingEngine() + + async with build(agent_doc, engine, api_key="sbt_k") as agent: + await agent.chat("one") + monkeypatch.setenv("LYZR_API_KEY", "sbt_env_2") + await agent.chat("two") + expiry = (datetime.now(timezone.utc) + timedelta(hours=1)).isoformat() + assert CREDENTIAL.adopt("sbt_k", "sbt_k2", expiry) + await agent.chat("three") + + assert [keys_of(c) for c in engine.contexts] == [{"sbt_k"}, {"sbt_k"}, {"sbt_k2"}] + + +@respx.mock +async def test_missing_key_fails_the_turn(agent_doc, monkeypatch): + monkeypatch.setenv("LYZR_API_KEY", "sbt_1") + engine = RecordingEngine() + + async with build(agent_doc, engine) as agent: + await agent.chat("one") + monkeypatch.delenv("LYZR_API_KEY") + with pytest.raises(SettingsError, match="LYZR_API_KEY") as exc: + await agent.chat("two") + + assert len(engine.contexts) == 1 + assert "sbt_1" not in str(exc.value)