refactor(egress): use provisioned_env instead of sentinel for Codex token (PRD 0030)
Add `provisioned_env: dict[str, str]` to `AgentProvisionPlan`. When `forward_host_credentials=True`, `agent_provision_plan` reads the host Codex access token at prepare time and stores it under `CODEX_HOST_CREDENTIAL_TOKEN_REF`. Both backends merge `provisioned_env` over `os.environ` before calling `egress_resolve_token_values`, so the token slot resolves like any other manifest-declared token ref. Removes `egress_resolve_token_values_with_provider` and the sentinel `continue` skip from `egress_resolve_token_values`. The function is now fully generic — it neither knows nor cares about provider identity.
This commit is contained in:
@@ -12,7 +12,7 @@ from dataclasses import dataclass, field
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Literal
|
from typing import Literal
|
||||||
|
|
||||||
from .codex_auth import write_codex_dummy_auth_file
|
from .codex_auth import codex_host_access_token, write_codex_dummy_auth_file
|
||||||
from .egress import CODEX_HOST_CREDENTIAL_TOKEN_REF, EgressRoute
|
from .egress import CODEX_HOST_CREDENTIAL_TOKEN_REF, EgressRoute
|
||||||
|
|
||||||
|
|
||||||
@@ -92,6 +92,7 @@ class AgentProvisionPlan:
|
|||||||
verify: tuple[AgentProvisionCommand, ...] = ()
|
verify: tuple[AgentProvisionCommand, ...] = ()
|
||||||
egress_routes: tuple[EgressRoute, ...] = ()
|
egress_routes: tuple[EgressRoute, ...] = ()
|
||||||
hidden_env_names: frozenset[str] = field(default_factory=frozenset)
|
hidden_env_names: frozenset[str] = field(default_factory=frozenset)
|
||||||
|
provisioned_env: dict[str, str] = field(default_factory=dict)
|
||||||
|
|
||||||
|
|
||||||
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
_REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||||
@@ -139,6 +140,7 @@ def agent_provision_plan(
|
|||||||
runtime = runtime_for(template)
|
runtime = runtime_for(template)
|
||||||
resolved_guest_env = dict(guest_env or {})
|
resolved_guest_env = dict(guest_env or {})
|
||||||
env_vars: dict[str, str] = {}
|
env_vars: dict[str, str] = {}
|
||||||
|
provisioned_env: dict[str, str] = {}
|
||||||
dirs: list[AgentProvisionDir] = []
|
dirs: list[AgentProvisionDir] = []
|
||||||
files: list[AgentProvisionFile] = []
|
files: list[AgentProvisionFile] = []
|
||||||
pre_copy: list[AgentProvisionCommand] = []
|
pre_copy: list[AgentProvisionCommand] = []
|
||||||
@@ -169,8 +171,12 @@ def agent_provision_plan(
|
|||||||
tls_passthrough=True,
|
tls_passthrough=True,
|
||||||
))
|
))
|
||||||
if forward_host_credentials:
|
if forward_host_credentials:
|
||||||
|
_host_env = host_env or dict(os.environ)
|
||||||
|
provisioned_env[CODEX_HOST_CREDENTIAL_TOKEN_REF] = codex_host_access_token(
|
||||||
|
_host_env,
|
||||||
|
)
|
||||||
auth_file = state_dir / "codex-auth.json"
|
auth_file = state_dir / "codex-auth.json"
|
||||||
write_codex_dummy_auth_file(auth_file, host_env or dict(os.environ))
|
write_codex_dummy_auth_file(auth_file, _host_env)
|
||||||
files.append(AgentProvisionFile(auth_file, f"{auth_dir}/auth.json"))
|
files.append(AgentProvisionFile(auth_file, f"{auth_dir}/auth.json"))
|
||||||
pre_copy.append(AgentProvisionCommand((
|
pre_copy.append(AgentProvisionCommand((
|
||||||
"find", auth_dir,
|
"find", auth_dir,
|
||||||
@@ -220,6 +226,7 @@ def agent_provision_plan(
|
|||||||
verify=tuple(verify),
|
verify=tuple(verify),
|
||||||
egress_routes=tuple(egress_routes),
|
egress_routes=tuple(egress_routes),
|
||||||
hidden_env_names=hidden_env_names,
|
hidden_env_names=hidden_env_names,
|
||||||
|
provisioned_env=provisioned_env,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -42,7 +42,7 @@ from contextlib import ExitStack, contextmanager
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Callable, Generator
|
from typing import Callable, Generator
|
||||||
|
|
||||||
from ...egress import egress_resolve_token_values_with_provider
|
from ...egress import egress_resolve_token_values
|
||||||
from ...log import info
|
from ...log import info
|
||||||
from . import network as network_mod
|
from . import network as network_mod
|
||||||
from . import util as docker_mod
|
from . import util as docker_mod
|
||||||
@@ -176,11 +176,9 @@ def launch(
|
|||||||
# Step 7: compose up. Token values + the OAuth placeholder
|
# Step 7: compose up. Token values + the OAuth placeholder
|
||||||
# flow through subprocess env; the compose file holds only
|
# flow through subprocess env; the compose file holds only
|
||||||
# bare names for the secret-carrying entries.
|
# bare names for the secret-carrying entries.
|
||||||
bottle = plan.spec.manifest.bottle_for(plan.spec.agent_name)
|
effective_env = {**dict(os.environ), **plan.agent_provision.provisioned_env}
|
||||||
token_values = egress_resolve_token_values_with_provider(
|
token_values = egress_resolve_token_values(
|
||||||
plan.egress_plan.token_env_map,
|
plan.egress_plan.token_env_map, effective_env,
|
||||||
bottle.agent_provider.forward_host_credentials,
|
|
||||||
dict(os.environ),
|
|
||||||
)
|
)
|
||||||
compose_env: dict[str, str] = {
|
compose_env: dict[str, str] = {
|
||||||
**os.environ,
|
**os.environ,
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ from typing import Callable, Generator
|
|||||||
|
|
||||||
from ...egress import (
|
from ...egress import (
|
||||||
EGRESS_ROUTES_IN_CONTAINER,
|
EGRESS_ROUTES_IN_CONTAINER,
|
||||||
egress_resolve_token_values_with_provider,
|
egress_resolve_token_values,
|
||||||
)
|
)
|
||||||
from ...pipelock import (
|
from ...pipelock import (
|
||||||
PIPELOCK_CA_CERT_IN_CONTAINER,
|
PIPELOCK_CA_CERT_IN_CONTAINER,
|
||||||
@@ -423,12 +423,8 @@ def _resolve_token_env(
|
|||||||
"""Resolve the egress token env-var values from the host's
|
"""Resolve the egress token env-var values from the host's
|
||||||
environ so they reach the bundle's process env via docker's
|
environ so they reach the bundle's process env via docker's
|
||||||
`-e NAME` inheritance. Empty when no routes declare auth."""
|
`-e NAME` inheritance. Empty when no routes declare auth."""
|
||||||
bottle = plan.spec.manifest.bottle_for(plan.spec.agent_name)
|
effective_env = {**host_env, **plan.agent_provision.provisioned_env}
|
||||||
return egress_resolve_token_values_with_provider(
|
return egress_resolve_token_values(plan.egress_plan.token_env_map, effective_env)
|
||||||
plan.egress_plan.token_env_map,
|
|
||||||
bottle.agent_provider.forward_host_credentials,
|
|
||||||
host_env,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _ensure_smolmachine(image_ref: str, *, dockerfile: str = "") -> Path:
|
def _ensure_smolmachine(image_ref: str, *, dockerfile: str = "") -> Path:
|
||||||
|
|||||||
@@ -29,7 +29,6 @@ from dataclasses import dataclass
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
from .codex_auth import codex_host_access_token
|
|
||||||
from .log import die
|
from .log import die
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
@@ -342,8 +341,6 @@ def egress_resolve_token_values(
|
|||||||
a sealed mapping without touching `os.environ`."""
|
a sealed mapping without touching `os.environ`."""
|
||||||
out: dict[str, str] = {}
|
out: dict[str, str] = {}
|
||||||
for token_env, token_ref in token_env_map.items():
|
for token_env, token_ref in token_env_map.items():
|
||||||
if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF:
|
|
||||||
continue
|
|
||||||
value = host_env.get(token_ref)
|
value = host_env.get(token_ref)
|
||||||
if value is None:
|
if value is None:
|
||||||
die(
|
die(
|
||||||
@@ -361,31 +358,6 @@ def egress_resolve_token_values(
|
|||||||
return out
|
return out
|
||||||
|
|
||||||
|
|
||||||
def egress_resolve_token_values_with_provider(
|
|
||||||
token_env_map: dict[str, str],
|
|
||||||
forward_host_credentials: bool,
|
|
||||||
host_env: dict[str, str],
|
|
||||||
) -> dict[str, str]:
|
|
||||||
"""Resolve all egress token env-var values, including the optional
|
|
||||||
Codex host credential slot.
|
|
||||||
|
|
||||||
Combines `egress_resolve_token_values` (manifest-declared token refs)
|
|
||||||
with the `forward_host_credentials` path (Codex ChatGPT bearer).
|
|
||||||
Returns an empty dict when `token_env_map` is empty.
|
|
||||||
|
|
||||||
Pure function: `host_env` is passed in so tests can use a sealed
|
|
||||||
mapping without touching `os.environ`."""
|
|
||||||
if not token_env_map:
|
|
||||||
return {}
|
|
||||||
token_values = egress_resolve_token_values(token_env_map, host_env)
|
|
||||||
if forward_host_credentials:
|
|
||||||
access_token = codex_host_access_token(host_env)
|
|
||||||
for token_env, token_ref in token_env_map.items():
|
|
||||||
if token_ref == CODEX_HOST_CREDENTIAL_TOKEN_REF:
|
|
||||||
token_values[token_env] = access_token
|
|
||||||
return token_values
|
|
||||||
|
|
||||||
|
|
||||||
class Egress(ABC):
|
class Egress(ABC):
|
||||||
"""The per-bottle egress proxy. Encapsulates the host-side prepare
|
"""The per-bottle egress proxy. Encapsulates the host-side prepare
|
||||||
(route lift + routes.yaml render + token-env-map derivation); the
|
(route lift + routes.yaml render + token-env-map derivation); the
|
||||||
@@ -429,7 +401,6 @@ __all__ = [
|
|||||||
"egress_manifest_routes",
|
"egress_manifest_routes",
|
||||||
"egress_render_routes",
|
"egress_render_routes",
|
||||||
"egress_resolve_token_values",
|
"egress_resolve_token_values",
|
||||||
"egress_resolve_token_values_with_provider",
|
|
||||||
"egress_routes_for_bottle",
|
"egress_routes_for_bottle",
|
||||||
"egress_token_env_map",
|
"egress_token_env_map",
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -145,6 +145,37 @@ class TestAgentProviderRuntime(unittest.TestCase):
|
|||||||
self.assertNotIn("CLAUDE_CODE_OAUTH_TOKEN", plan.env_vars)
|
self.assertNotIn("CLAUDE_CODE_OAUTH_TOKEN", plan.env_vars)
|
||||||
self.assertEqual(frozenset(), plan.hidden_env_names)
|
self.assertEqual(frozenset(), plan.hidden_env_names)
|
||||||
|
|
||||||
|
def test_codex_forward_host_credentials_populates_provisioned_env(self):
|
||||||
|
access = _jwt(2000000000)
|
||||||
|
with tempfile.TemporaryDirectory(prefix="bb-provider.") as tmp:
|
||||||
|
home = Path(tmp) / "host-codex"
|
||||||
|
home.mkdir()
|
||||||
|
(home / "auth.json").write_text(json.dumps({
|
||||||
|
"auth_mode": "chatgpt",
|
||||||
|
"tokens": {"access_token": access},
|
||||||
|
}))
|
||||||
|
plan = agent_provision_plan(
|
||||||
|
template="codex",
|
||||||
|
dockerfile="",
|
||||||
|
state_dir=Path(tmp),
|
||||||
|
forward_host_credentials=True,
|
||||||
|
host_env={"CODEX_HOME": str(home)},
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
{CODEX_HOST_CREDENTIAL_TOKEN_REF: access},
|
||||||
|
plan.provisioned_env,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_codex_without_forward_host_credentials_has_empty_provisioned_env(self):
|
||||||
|
with tempfile.TemporaryDirectory(prefix="bb-provider.") as tmp:
|
||||||
|
plan = agent_provision_plan(
|
||||||
|
template="codex",
|
||||||
|
dockerfile="",
|
||||||
|
state_dir=Path(tmp),
|
||||||
|
forward_host_credentials=False,
|
||||||
|
)
|
||||||
|
self.assertEqual({}, plan.provisioned_env)
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
@@ -2,7 +2,6 @@
|
|||||||
resolution (PRD 0017)."""
|
resolution (PRD 0017)."""
|
||||||
|
|
||||||
import unittest
|
import unittest
|
||||||
import unittest.mock
|
|
||||||
|
|
||||||
from bot_bottle.egress import (
|
from bot_bottle.egress import (
|
||||||
CODEX_HOST_CREDENTIAL_TOKEN_REF,
|
CODEX_HOST_CREDENTIAL_TOKEN_REF,
|
||||||
@@ -10,7 +9,6 @@ from bot_bottle.egress import (
|
|||||||
egress_manifest_routes,
|
egress_manifest_routes,
|
||||||
egress_render_routes,
|
egress_render_routes,
|
||||||
egress_resolve_token_values,
|
egress_resolve_token_values,
|
||||||
egress_resolve_token_values_with_provider,
|
|
||||||
egress_routes_for_bottle,
|
egress_routes_for_bottle,
|
||||||
egress_token_env_map,
|
egress_token_env_map,
|
||||||
)
|
)
|
||||||
@@ -343,72 +341,13 @@ class TestResolveTokenValues(unittest.TestCase):
|
|||||||
{"GH_PAT": ""},
|
{"GH_PAT": ""},
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_codex_host_credential_ref_is_resolved_by_launch(self):
|
def test_codex_host_credential_ref_resolved_via_provisioned_env(self):
|
||||||
out = egress_resolve_token_values(
|
out = egress_resolve_token_values(
|
||||||
{"EGRESS_TOKEN_0": CODEX_HOST_CREDENTIAL_TOKEN_REF},
|
{"EGRESS_TOKEN_0": CODEX_HOST_CREDENTIAL_TOKEN_REF},
|
||||||
{},
|
{CODEX_HOST_CREDENTIAL_TOKEN_REF: "codex-access-token"},
|
||||||
)
|
|
||||||
self.assertEqual({}, out)
|
|
||||||
|
|
||||||
|
|
||||||
class TestResolveTokenValuesWithProvider(unittest.TestCase):
|
|
||||||
def test_empty_map_returns_empty(self):
|
|
||||||
out = egress_resolve_token_values_with_provider({}, False, {})
|
|
||||||
self.assertEqual({}, out)
|
|
||||||
|
|
||||||
def test_empty_map_with_forward_credentials_returns_empty(self):
|
|
||||||
# forward_host_credentials=True but no slots → no codex call needed.
|
|
||||||
out = egress_resolve_token_values_with_provider({}, True, {})
|
|
||||||
self.assertEqual({}, out)
|
|
||||||
|
|
||||||
def test_manifest_tokens_resolved_without_forward_credentials(self):
|
|
||||||
out = egress_resolve_token_values_with_provider(
|
|
||||||
{"EGRESS_TOKEN_0": "GH_PAT"},
|
|
||||||
False,
|
|
||||||
{"GH_PAT": "ghp_secret"},
|
|
||||||
)
|
|
||||||
self.assertEqual({"EGRESS_TOKEN_0": "ghp_secret"}, out)
|
|
||||||
|
|
||||||
def test_codex_token_slotted_in_when_forward_credentials_and_matching_ref(self):
|
|
||||||
with unittest.mock.patch(
|
|
||||||
"bot_bottle.egress.codex_host_access_token",
|
|
||||||
return_value="codex-access-token",
|
|
||||||
):
|
|
||||||
out = egress_resolve_token_values_with_provider(
|
|
||||||
{"EGRESS_TOKEN_0": CODEX_HOST_CREDENTIAL_TOKEN_REF},
|
|
||||||
True,
|
|
||||||
{},
|
|
||||||
)
|
)
|
||||||
self.assertEqual({"EGRESS_TOKEN_0": "codex-access-token"}, out)
|
self.assertEqual({"EGRESS_TOKEN_0": "codex-access-token"}, out)
|
||||||
|
|
||||||
def test_codex_token_not_slotted_when_no_matching_ref(self):
|
|
||||||
# forward_host_credentials=True but no CODEX_HOST_CREDENTIAL_TOKEN_REF
|
|
||||||
# slot in the map → manifest tokens only; Codex token is fetched but
|
|
||||||
# nothing to slot it into.
|
|
||||||
with unittest.mock.patch(
|
|
||||||
"bot_bottle.egress.codex_host_access_token",
|
|
||||||
return_value="codex-access-token",
|
|
||||||
):
|
|
||||||
out = egress_resolve_token_values_with_provider(
|
|
||||||
{"EGRESS_TOKEN_0": "GH_PAT"},
|
|
||||||
True,
|
|
||||||
{"GH_PAT": "ghp_secret"},
|
|
||||||
)
|
|
||||||
self.assertEqual({"EGRESS_TOKEN_0": "ghp_secret"}, out)
|
|
||||||
|
|
||||||
def test_codex_not_called_when_forward_credentials_false(self):
|
|
||||||
called = []
|
|
||||||
with unittest.mock.patch(
|
|
||||||
"bot_bottle.egress.codex_host_access_token",
|
|
||||||
side_effect=lambda *_: called.append(1) or "tok",
|
|
||||||
):
|
|
||||||
egress_resolve_token_values_with_provider(
|
|
||||||
{"EGRESS_TOKEN_0": "GH_PAT"},
|
|
||||||
False,
|
|
||||||
{"GH_PAT": "ghp_secret"},
|
|
||||||
)
|
|
||||||
self.assertEqual([], called)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
unittest.main()
|
unittest.main()
|
||||||
|
|||||||
Reference in New Issue
Block a user