diff --git a/.agentenv/config.example.toml b/.agentenv/config.example.toml index a820b98c..50fbb3d3 100644 --- a/.agentenv/config.example.toml +++ b/.agentenv/config.example.toml @@ -14,6 +14,12 @@ secret = "local" # process env vars, or an optional flat YAML file default = "local" # envs run as containers on the host Docker daemon agent_default = "local" # agents too +# Deployment-wide attribution defaults, any keys (values take env:NAME?default references). +# Modal puts them on its App tags; a run's own attribution, over these, goes on each sandbox's +# tags (Modal) or metadata (E2B). The local provider ignores attribution. +# [sandbox.attribution] +# group = "my-group" + [agents] default_a2a_agent_id = "a2a-default" # what a deploy_agent step that names no agent deploys diff --git a/AGENTS.md b/AGENTS.md index 16266609..b82cf676 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -82,7 +82,7 @@ and deserialized through a registry. | `task/`, `task_step/` | A `Task` holds its `TaskStep`s inline; `Task.run()` executes them as a DAG. `depends_on` (None means all prior steps) gates a step, independent steps run concurrently, `fail_task_on_error` makes a failure fatal or tolerated, `retry_config` rolls a failed span back through the step journal and re-dispatches it. Built-in steps live in `task_step/task_steps/` (`deploy_env`, `deploy_agent`, `prompt_agent`, the verifiers under `verifiers/`, and more). | | `store/` | Four store ABCs with local and cloud implementations: `DocumentStore` (SQLite, MongoDB), `ObjectStore` (filesystem, S3, Cloud Storage), `ImageStore` (local OCI registry, ECR), `SecretStore` (env vars or file, AWS Secrets Manager, Google Cloud Secret Manager). `VersionedEntityStore` implements the shared versioned get/put logic, `QueryBuilder` is the immutable chained query API, `store/base.py` holds the error types. A new backend must pass the conformance kits in `tst/store/`. | | `config/` | The `Config` singleton (`get_config`, `configure`, `reset_config`) in `config/runtime.py`, file discovery in `config/loader.py`, and `load_impl`, which resolves `module:Class` pointers. `agent_env.store` re-exports the config names for compatibility. | -| `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`; `[sandbox] default` and `agent_default` accept a comma-separated fallback chain. `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store) and `EnvironmentGatewayProvider`, which renders a docker-compose for the gateway and its MCP servers inside the sandbox; `providers/env_state/` holds env-state providers (`local_postgres` built in). | +| `providers/` | `providers/sandbox_providers/` holds the sandbox providers `local`, `modal`, `modal_vm`, `e2b`; `[sandbox] default` and `agent_default` accept a comma-separated fallback chain, and `[sandbox.attribution]` sets attribution defaults under any keys: Modal App tags carry those defaults, and a run's attribution over them goes on its sandbox tags (Modal) or metadata (E2B). `providers/env_providers/` holds the environment providers: `EnvironmentProvider` (an env's containers and state store) and `EnvironmentGatewayProvider`, which renders a docker-compose for the gateway and its MCP servers inside the sandbox; `providers/env_state/` holds env-state providers (`local_postgres` built in). | | `a2a_agent/` | The `A2AAgent` entity (`a2a_agent`), its stores and the validator steps. The protocol package provides the agent-side framework. | | `runner/` | The `[runner]` seam: `Runner.submit()` returns `(run_id, instance_id)`; `LocalRunner` is built in. | | `explorer/` | Optional local web UI: `agent-env up`, needs the `explorer` extra, binds loopback `:8234`. | diff --git a/src/agent_env/providers/sandbox_providers/e2b/provider.py b/src/agent_env/providers/sandbox_providers/e2b/provider.py index c2aa1df9..4c463dee 100644 --- a/src/agent_env/providers/sandbox_providers/e2b/provider.py +++ b/src/agent_env/providers/sandbox_providers/e2b/provider.py @@ -17,11 +17,11 @@ from agent_env.attribution import Attribution from agent_env.config.errors import ConfigError from agent_env.providers.sandbox_providers.e2b.sandbox import E2B_ALL_TRAFFIC, E2BSandbox -from agent_env.providers.sandbox_providers.modal_sandbox import _build_cost_attribution_tags from agent_env.providers.sandbox_providers.sandbox import NetworkMode, NetworkPolicy from agent_env.providers.sandbox_providers.sandbox_provider import ( SANDBOX_MODE_VM, SandboxProvider, + apply_default_attribution, ) logger = logging.getLogger(__name__) @@ -29,6 +29,16 @@ _EXPOSED_PORTS_METADATA_KEY = "agent_env_exposed_ports" +def _attribution_metadata(attribution: Attribution | None) -> dict[str, str]: + """Sandbox metadata from the attribution, unset keys filled from ``[sandbox.attribution]``. + The exposed-ports key is agent-env's own, so an attribution key of that name never lands.""" + return { + key: str(value) + for key, value in apply_default_attribution(dict(attribution or {})).items() + if value is not None and key != _EXPOSED_PORTS_METADATA_KEY + } + + class E2BSandboxProvider(SandboxProvider): """Provision Docker-capable E2B sandboxes from a configured base template. @@ -143,7 +153,7 @@ async def create_vm( ) effective_policy = self.effective_network_policy(network_policy) template = await self._resolve_template(cpu=cpu, memory=memory) - metadata = _build_cost_attribution_tags(dict(attribution or {})) + metadata = _attribution_metadata(attribution) ports = list(exposed_ports or []) if ports: # E2B can derive a URL for any port, so the adapter's cache is the diff --git a/src/agent_env/providers/sandbox_providers/modal_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_sandbox.py index 8e36be9a..5b874781 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -253,11 +253,9 @@ async def create_container( effective = self.effective_network_policy(network_policy) attribution = dict(attribution or {}) - # Modal billing breaks cost down by App tags; sandbox-level tags don't show up in billing. - app_tags = _build_cost_attribution_tags(attribution) + app_tags = _build_app_tags() app_name = self._app_name app = await self._get_app(app_name, app_tags) - # Not in billing (see above); makes sandboxes filterable by the step that deployed them. sandbox_tags = _build_sandbox_tags(attribution) image_store = get_config().get_image_store_at(image_name) @@ -401,33 +399,27 @@ async def get_sandbox(self, sandbox_id: str) -> Sandbox: return self._sandbox_cls(sb, tunnel_urls, i6pn_address=i6pn_address) -def _build_cost_attribution_tags(attribution: Attribution) -> dict[str, str]: - """Build the Modal App tag set used for cost attribution. +def _build_app_tags() -> dict[str, str]: + """Modal App tags: the deployment-wide config.toml ``[sandbox.attribution]`` defaults. - Unset dimensions fall back to config.toml ``[sandbox.attribution]``; a dimension with - no value anywhere is omitted rather than emitted as a null tag. Returns a flat - ``dict[str, str]`` for ``modal.App.set_tags``. (priority is a scheduling concern, - not attribution, so it is not included here.) - """ - resolved = apply_default_attribution(attribution) - return { - name: resolved[name] - for name in ("product", "customer", "team") - if resolved.get(name) is not None - } + Modal's billing report breaks cost down by App tags, and one App serves every run with its + tags set on first lookup, so per-run attribution goes on sandbox tags instead.""" + return _modal_tags(apply_default_attribution({})) def _build_sandbox_tags(attribution: Attribution) -> dict[str, str]: - """Build the Modal Sandbox tag set: ``pipeline_step`` and ``run_id``, each when set. + """Modal Sandbox tags: the run's attribution, unset keys filled from ``[sandbox.attribution]``, + including ``pipeline_step`` and ``run_id`` when the deploy step set them.""" + return _modal_tags(apply_default_attribution(attribution)) - Modal rejects the whole create on an invalid tag, so values are sanitized to - ``[a-zA-Z0-9._-]``; past Modal's 63-char limit a value keeps a prefix plus a hash of the - full value, so steps of a long-named task still get distinct tags. - """ + +def _modal_tags(attribution: Attribution) -> dict[str, str]: + """Modal rejects a whole create on one invalid tag: keys and values are made Modal-valid and a + value that is unset or empty is left out, since Modal has no empty tag.""" return { - key: _modal_tag_value(attribution[key]) - for key in (PIPELINE_STEP_KEY, RUN_ID_KEY) - if attribution.get(key) + _modal_tag_value(key): _modal_tag_value(str(value)) + for key, value in attribution.items() + if value is not None and str(value) != "" } @@ -457,7 +449,8 @@ async def _log_sandbox_started( "modal_sandbox_id": sb.object_id, "modal_container_id": container_id, "modal_app_name": app_name, - **sandbox_tags, + "modal_sandbox_tags": sandbox_tags, + **{key: sandbox_tags[key] for key in (PIPELINE_STEP_KEY, RUN_ID_KEY) if key in sandbox_tags}, "cpu": cpu, "memory_mb": memory, "gpu": gpu, @@ -466,6 +459,8 @@ async def _log_sandbox_started( def _modal_tag_value(value: str) -> str: + """``value`` within Modal's ``[a-zA-Z0-9._-]{1,63}``: other characters become ``-``, and past + 63 characters a prefix plus a hash of the full value keeps long values distinct.""" slug = re.sub(r"[^a-zA-Z0-9._-]", "-", value) if len(slug) > _MODAL_TAG_MAX_LEN: digest = hashlib.sha256(value.encode()).hexdigest()[:8] diff --git a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py index 45e44a32..68a27218 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -29,7 +29,7 @@ from agent_env.providers.sandbox_providers.modal_sandbox import ( _ModalProcessAdapter, - _build_cost_attribution_tags, + _build_app_tags, _build_sandbox_tags, _log_sandbox_started, _resolve_app_base_name, @@ -308,7 +308,7 @@ async def create_vm( # disk_size_gb: Modal has no disk-size knob on Sandbox.create; a VM sandbox's rootfs is # fixed at the 512 GiB max. https://modal.com/docs/guide/vm-sandboxes attribution = dict(attribution or {}) - app_tags = _build_cost_attribution_tags(attribution) + app_tags = _build_app_tags() app_name = self._app_name app = await self._get_app(app_name, app_tags) sandbox_tags = _build_sandbox_tags(attribution) diff --git a/tst/unit/attribution_test.py b/tst/unit/attribution_test.py index d46e0df2..53e837e0 100644 --- a/tst/unit/attribution_test.py +++ b/tst/unit/attribution_test.py @@ -19,11 +19,11 @@ def test_attribution_of_reads_the_reserved_metadata_key_only(): def test_attribution_of_ignores_flat_attributes_on_the_step(): - step = _step(cost_center="research", metadata={"attribution": {"team": "t"}}) - assert attribution_of(step) == {"team": "t"} + step = _step(cost_center="research", metadata={"attribution": {"group": "g"}}) + assert attribution_of(step) == {"group": "g"} def test_attribution_of_returns_a_copy(): - metadata = {"attribution": {"team": "t"}} - attribution_of(_step(metadata=metadata))["team"] = "mutated" - assert metadata["attribution"] == {"team": "t"} + metadata = {"attribution": {"group": "g"}} + attribution_of(_step(metadata=metadata))["group"] = "mutated" + assert metadata["attribution"] == {"group": "g"} diff --git a/tst/unit/env/envs/test_mcp_server_container_mode.py b/tst/unit/env/envs/test_mcp_server_container_mode.py index 34627bd2..29bbb8f5 100644 --- a/tst/unit/env/envs/test_mcp_server_container_mode.py +++ b/tst/unit/env/envs/test_mcp_server_container_mode.py @@ -145,11 +145,11 @@ async def test_a_deploy_hands_the_gateway_its_sizing_state_and_attribution_and_r env, state, seen = _env(), SimpleNamespace(state_type="remote", instance_id="st-ext"), {} record = await _deploy(env, {"gw": _sandbox("gw")}, seen, state, gateway_mode=GatewayMode.CONSISTENT, env_state_type="remote", - ttl_seconds=60, disk_size_gb=20, cpu=2.0, memory_mb=4096, priority=3, attribution={"team": "t1"}) + ttl_seconds=60, disk_size_gb=20, cpu=2.0, memory_mb=4096, priority=3, attribution={"group": "g1"}) assert ([c.environment_name for c in seen["mcp_servers"]], seen["mcp_server_images"]) == (["email"], [env.docker_image_artifact]) assert {k: seen[k] for k in ("ttl_seconds", "disk_size_gb", "cpu", "memory_mb", "priority", "attribution", "gateway_mode")} == { - "ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"team": "t1"}, + "ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"group": "g1"}, "gateway_mode": GatewayMode.CONSISTENT} assert (seen["gateway_port"], seen["website_configs"], seen["website_images"], seen["existing_sandbox"], seen["sidecars"]) == ( 18765, None, None, None, None) @@ -199,9 +199,9 @@ async def test_a_deploy_without_a_gateway_records_the_servers_own_card(): async def test_a_deploy_without_a_gateway_hands_the_server_its_sizing_and_attribution(): seen = {} await _deploy(_env("server"), {"srv": _sandbox("srv")}, seen, gateway_mode=GatewayMode.PERFORMANCE, # the deploy_env step's default - ttl_seconds=60, disk_size_gb=20, cpu=2.0, memory_mb=4096, priority=3, attribution={"team": "t1"}) + ttl_seconds=60, disk_size_gb=20, cpu=2.0, memory_mb=4096, priority=3, attribution={"group": "g1"}) - assert seen == {"ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"team": "t1"}} + assert seen == {"ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"group": "g1"}} @pytest.mark.asyncio diff --git a/tst/unit/env/envs/test_mcp_server_plugin_provider.py b/tst/unit/env/envs/test_mcp_server_plugin_provider.py index 3fd1d356..edc61154 100644 --- a/tst/unit/env/envs/test_mcp_server_plugin_provider.py +++ b/tst/unit/env/envs/test_mcp_server_plugin_provider.py @@ -108,12 +108,12 @@ def test_reading_a_stored_env_looks_up_no_environment_provider(monkeypatch): @pytest.mark.asyncio async def test_a_plugin_provider_gets_every_option_and_its_record_is_registered(): env = _env() - record = await _deploy(env, ttl_seconds=60, cpu=2.0, priority=1, gateway_mode=GatewayMode.CONSISTENT, attribution={"team": "t"}) + record = await _deploy(env, ttl_seconds=60, cpu=2.0, priority=1, gateway_mode=GatewayMode.CONSISTENT, attribution={"group": "g"}) [(sandbox_provider, options)] = SEEN["deploys"] assert sandbox_provider == "SANDBOXES" assert options == {"ttl_seconds": 60, "disk_size_gb": 10, "gateway_mode": GatewayMode.CONSISTENT, "cpu": 2.0, "memory_mb": None, - "priority": 1, "env_state_type": None, "env_state_instance_id": None, "attribution": {"team": "t"}} + "priority": 1, "env_state_type": None, "env_state_instance_id": None, "attribution": {"group": "g"}} assert (record.instance_id, env._instance_id, env._deployed, env._sandbox, env._gateway_url) == ("inst-1", "inst-1", record, None, None) @@ -252,8 +252,8 @@ async def close(self): async def test_a_plugin_that_names_its_options_gets_those_and_the_rest_left_at_their_defaults(): env = _env("plugin_named") env._env_provider = _NamedOptionsProvider() - await _deploy(env, ttl_seconds=60, attribution={"team": "t"}) - assert SEEN["deploys"] == [("SANDBOXES", {"ttl_seconds": 60, "attribution": {"team": "t"}})] + await _deploy(env, ttl_seconds=60, attribution={"group": "g"}) + assert SEEN["deploys"] == [("SANDBOXES", {"ttl_seconds": 60, "attribution": {"group": "g"}})] @pytest.mark.asyncio diff --git a/tst/unit/providers/env_providers/env_gateway_provider_test.py b/tst/unit/providers/env_providers/env_gateway_provider_test.py index 93a584c3..a9610996 100644 --- a/tst/unit/providers/env_providers/env_gateway_provider_test.py +++ b/tst/unit/providers/env_providers/env_gateway_provider_test.py @@ -1308,7 +1308,7 @@ async def test_a_chained_gateway_that_fails_everywhere_names_every_member(): @pytest.mark.asyncio -@pytest.mark.parametrize("attribution, forwarded", [({"team": "t1"}, {"team": "t1"}), (None, {})], ids=["set", "unset"]) +@pytest.mark.parametrize("attribution, forwarded", [({"group": "g1"}, {"group": "g1"}), (None, {})], ids=["set", "unset"]) async def test_create_gateway_hands_the_deploy_path_every_argument(attribution, forwarded): """The topology entry point loses nothing on its way to the VM path; attribution arrives as a fresh dict.""" from agent_env.env.gateway import GatewayMode diff --git a/tst/unit/providers/env_providers/env_server_provider_test.py b/tst/unit/providers/env_providers/env_server_provider_test.py index 2795a646..1d33ad02 100644 --- a/tst/unit/providers/env_providers/env_server_provider_test.py +++ b/tst/unit/providers/env_providers/env_server_provider_test.py @@ -53,10 +53,10 @@ async def test_a_deploy_returns_a_sandbox_record_of_the_servers_own_card_and_its @pytest.mark.asyncio async def test_a_deploy_starts_only_the_server_on_sqlite_with_its_sizing(): - run = await _run_server_deploy(cpu=2.0, memory_mb=4096, priority=3, attribution={"team": "t1"}) + run = await _run_server_deploy(cpu=2.0, memory_mb=4096, priority=3, attribution={"group": "g1"}) [call] = run.calls assert (call["image_name"], call["port"], call["cpu"], call["memory"], call["timeout"], call["priority"]) == ("mcp-slack", 18765, 2.0, 4096, 60, 3) - assert call["attribution"] == {"team": "t1"} and "i6pn" not in call + assert call["attribution"] == {"group": "g1"} and "i6pn" not in call assert call["env"] == {"ENVIRONMENT_NAME": "slack", "MCP_HOST": "0.0.0.0", "DATABASE_URL": "sqlite:////tmp/slack.sqlite"} diff --git a/tst/unit/providers/sandbox_providers/attribution_wire_test.py b/tst/unit/providers/sandbox_providers/attribution_wire_test.py index e075aa10..3729e2c0 100644 --- a/tst/unit/providers/sandbox_providers/attribution_wire_test.py +++ b/tst/unit/providers/sandbox_providers/attribution_wire_test.py @@ -1,20 +1,17 @@ -"""Byte-identity net for the emitted compute-attribution wire. +"""What attribution emits on the Modal wire. -Attribution lands in the Modal App `tags` dict (the `labels` sink of the out-of-tree platform -sandbox provider is pinned by that provider's tests). These pin what comes out, so a refactor of -how attribution is carried announces itself instead of quietly re-bucketing spend. - -Precedence is `explicit arg -> config.toml [sandbox.attribution]`, tested `is not None`: an -empty string is a value. +The deployment-wide config.toml ``[sandbox.attribution]`` defaults are the Modal App tags; a +run's attribution, with unset keys filled from those defaults, is its sandbox tags. Neither names +a key: any dimension a deployment uses passes through. Precedence is +``explicit value -> [sandbox.attribution]``, tested ``is not None``. """ from types import SimpleNamespace import pytest -from agent_env.providers.sandbox_providers.modal_sandbox import _build_cost_attribution_tags - -_DEFAULTS = {"product": "p0", "customer": "c0", "team": "t0"} +from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY +from agent_env.providers.sandbox_providers.modal_sandbox import _build_app_tags, _build_sandbox_tags def _use_config(tmp_path, monkeypatch, attribution: dict[str, str]) -> None: @@ -25,67 +22,71 @@ def _use_config(tmp_path, monkeypatch, attribution: dict[str, str]) -> None: @pytest.fixture(autouse=True) def _configured_defaults(tmp_path, monkeypatch): - # Pin config discovery to a fixture so no ambient .agentenv leaks into the assertions. - _use_config(tmp_path, monkeypatch, _DEFAULTS) + _use_config(tmp_path, monkeypatch, {"group": "g0", "owner": "o0"}) -def _tags(**kwargs): - return _build_cost_attribution_tags(kwargs) +# --- App tags: the configured defaults only ------------------------------------------ -# --- the emitted dicts, exactly ------------------------------------------------- +def test_app_tags_are_the_configured_defaults(): + assert _build_app_tags() == {"group": "g0", "owner": "o0"} -def test_modal_tags_with_nothing_supplied(): - assert _tags() == {"product": "p0", "customer": "c0", "team": "t0"} +def test_app_tags_without_configured_defaults_are_empty(tmp_path, monkeypatch): + _use_config(tmp_path, monkeypatch, {}) + assert _build_app_tags() == {} -def test_modal_tags_with_everything_supplied(): - assert _tags(product="p", customer="c", team="t") == {"product": "p", "customer": "c", "team": "t"} +def test_a_configured_default_may_be_an_env_reference(tmp_path, monkeypatch): + _use_config(tmp_path, monkeypatch, {"group": "env:ATTRIBUTION_TEST_GROUP?g0"}) + monkeypatch.delenv("ATTRIBUTION_TEST_GROUP", raising=False) + assert _build_app_tags() == {"group": "g0"} + monkeypatch.setenv("ATTRIBUTION_TEST_GROUP", "from-env") + assert _build_app_tags() == {"group": "from-env"} -def test_modal_vm_shares_the_modal_sink(): +def test_modal_vm_shares_the_modal_builders(): from agent_env.providers.sandbox_providers import modal_vm_sandbox - assert modal_vm_sandbox._build_cost_attribution_tags is _build_cost_attribution_tags + assert modal_vm_sandbox._build_app_tags is _build_app_tags + assert modal_vm_sandbox._build_sandbox_tags is _build_sandbox_tags +# --- sandbox tags: the run's attribution over the defaults ---------------------------- -# --- precedence ----------------------------------------------------------------- +def test_any_attribution_key_reaches_the_sandbox_tags(): + tags = _build_sandbox_tags({"batch": "b1", PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1"}) + assert tags == {"batch": "b1", "group": "g0", "owner": "o0", PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1"} -@pytest.mark.parametrize("field", ["product", "customer", "team"]) -def test_explicit_beats_the_configured_default(field): - assert _tags(**{field: "explicit"})[field] == "explicit" +def test_an_explicit_value_beats_the_configured_default(): + assert _build_sandbox_tags({"group": "explicit"})["group"] == "explicit" -def test_a_configured_default_may_be_an_env_reference(tmp_path, monkeypatch): - _use_config(tmp_path, monkeypatch, {**_DEFAULTS, "product": "env:ATTRIBUTION_TEST_PRODUCT?p0"}) - monkeypatch.delenv("ATTRIBUTION_TEST_PRODUCT", raising=False) - assert _tags()["product"] == "p0" - monkeypatch.setenv("ATTRIBUTION_TEST_PRODUCT", "from-env") - assert _tags()["product"] == "from-env" + +def test_an_empty_string_keeps_the_default_out_and_emits_no_tag(): + """``is not None``: an explicit empty value is not replaced by the default, and Modal has no + empty tag, so the key is left out.""" + assert "group" not in _build_sandbox_tags({"group": ""}) -def test_without_configured_defaults_modal_omits_unset_dimensions(tmp_path, monkeypatch): +def test_no_attribution_and_no_defaults_means_no_tags(tmp_path, monkeypatch): _use_config(tmp_path, monkeypatch, {}) - assert _tags() == {} - assert _tags(team="t") == {"team": "t"} + assert _build_sandbox_tags({}) == {} -@pytest.mark.parametrize("field", ["product", "customer", "team"]) -def test_an_empty_string_is_a_value_not_an_absence(field): - """`is not None`, not truthiness — pinned because flipping it is silent on Modal.""" - assert _tags(**{field: ""})[field] == "" +def test_keys_and_values_are_made_modal_valid(): + tags = _build_sandbox_tags({"my key/v2": "a value/with spaces", "long": "x" * 80}) + assert tags["my-key-v2"] == "a-value-with-spaces" + assert len(tags["long"]) == 63 -# --- the dict is the only carried form ------------------------------------------- +# --- the dict is the only carried form ------------------------------------------------- @pytest.mark.asyncio async def test_a_chain_forwards_the_dict_as_is(): - """ChainedSandboxProvider takes **kwargs; the dict reaches the backend untouched, open - keys included.""" + """ChainedSandboxProvider takes **kwargs; the dict reaches the backend untouched.""" from agent_env.providers.sandbox_providers.chained_sandbox_provider import ChainedSandboxProvider seen = {} @@ -98,7 +99,7 @@ async def create_sandbox(self, *, attribution=None, **kwargs): seen["kwargs"] = kwargs return SimpleNamespace(type="modern", network_policy=None) - attribution = {"team": "t"} + attribution = {"group": "g"} await ChainedSandboxProvider([Provider()]).create_sandbox(attribution=attribution, cpu=2.0) assert seen == {"attribution": attribution, "kwargs": {"cpu": 2.0}} diff --git a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py index a881fea6..ca47bd5f 100644 --- a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py @@ -29,6 +29,13 @@ class _AsyncSandboxSdk: connect = AsyncMock(side_effect=lambda *_args, **_kwargs: _RawSandbox()) +@pytest.fixture(autouse=True) +def _no_configured_defaults(tmp_path, monkeypatch): + cfg = tmp_path / "config.toml" + cfg.write_text("") + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + + @pytest.fixture(autouse=True) def reset_sdk_mock_calls(): _AsyncSandboxSdk.create.reset_mock() @@ -60,11 +67,7 @@ async def test_create_vm_uses_derived_template_and_preserves_attribution( disk_size_gb=99, timeout=123, exposed_ports=[8080, 9000], - attribution={ - "product": "product-a", - "customer": "customer-b", - "team": "team-c", - }, + attribution={"group": "g1", "batch": "b2"}, ) assert isinstance(sandbox, E2BSandbox) @@ -74,17 +77,40 @@ async def test_create_vm_uses_derived_template_and_preserves_attribution( template="agent-env-v1-2c-4096m", timeout=123, api_key="e2b-secret", - metadata={ - "product": "product-a", - "customer": "customer-b", - "team": "team-c", - "agent_env_exposed_ports": "8080,9000", - }, + metadata={"group": "g1", "batch": "b2", "agent_env_exposed_ports": "8080,9000"}, network={"allow_public_traffic": True}, ) sandbox_setup.assert_awaited_once_with([8080, 9000]) +@pytest.mark.asyncio +@pytest.mark.parametrize("ports, recorded", [([8080], "8080"), ([], None)], ids=["ports", "no-ports"]) +async def test_an_attribution_key_never_stands_in_for_the_exposed_ports(sandbox_setup, ports, recorded): + provider = E2BSandboxProvider( + api_key="e2b-secret", base_template="agent-env-v1", template_resolver=_Resolver(), + sandbox_cls=_AsyncSandboxSdk, + ) + await provider.create_vm( + exposed_ports=ports, attribution={"group": "g1", "agent_env_exposed_ports": "1,2,3"}, + ) + metadata = _AsyncSandboxSdk.create.call_args.kwargs["metadata"] + assert metadata.get("agent_env_exposed_ports") == recorded + assert metadata["group"] == "g1" + + +@pytest.mark.asyncio +async def test_configured_defaults_fill_unset_attribution_keys(sandbox_setup, tmp_path, monkeypatch): + cfg = tmp_path / "defaults.toml" + cfg.write_text('[sandbox.attribution]\ngroup = "g0"\nowner = "o0"\n') + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + provider = E2BSandboxProvider( + api_key="e2b-secret", base_template="agent-env-v1", template_resolver=_Resolver(), + sandbox_cls=_AsyncSandboxSdk, + ) + await provider.create_vm(attribution={"group": "explicit"}) + assert _AsyncSandboxSdk.create.call_args.kwargs["metadata"] == {"group": "explicit", "owner": "o0"} + + @pytest.mark.asyncio async def test_create_sandbox_provisions_a_vm_without_forwarding_image_or_env( sandbox_setup, diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_started_log_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_started_log_test.py index c249d80f..ca8011f3 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_started_log_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_started_log_test.py @@ -49,6 +49,16 @@ async def test_logs_the_container_id_with_the_sandbox_and_its_tags(caplog): assert (record.cpu, record.memory_mb, record.gpu) == (2.0, 4096, "H100") +@pytest.mark.asyncio +async def test_any_tag_key_is_logged_without_clashing_with_the_log_record(caplog): + tags = {**_TAGS, "name": "n", "group": "g"} + with caplog.at_level(logging.INFO, logger="agent_env.providers.sandbox_providers.modal_sandbox"): + await _log_sandbox_started(_fake_sb(), app_name="a", sandbox_tags=tags, cpu=1.0, memory=1024, gpu=None) + record = _started_record(caplog) + assert record.modal_sandbox_tags == tags + assert getattr(record, RUN_ID_KEY) == "inst-1" + + @pytest.mark.asyncio async def test_an_unreadable_container_id_still_logs_and_never_raises(caplog): with caplog.at_level(logging.INFO, logger="agent_env.providers.sandbox_providers.modal_sandbox"): diff --git a/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py b/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py index db431ac5..796bd42d 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py @@ -1,4 +1,4 @@ -"""The ``pipeline_step`` sandbox tag: built by deploy steps, stamped on Modal sandboxes.""" +"""Sandbox tags: a run's attribution, including ``pipeline_step`` and ``run_id``, stamped on Modal sandboxes.""" from __future__ import annotations @@ -10,12 +10,18 @@ from agent_env.attribution import PIPELINE_STEP_KEY, RUN_ID_KEY, deploy_attribution from agent_env.providers.sandbox_providers.modal_sandbox import ( ModalSandboxProvider, - _build_cost_attribution_tags, _build_sandbox_tags, ) from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider +@pytest.fixture(autouse=True) +def _no_configured_defaults(tmp_path, monkeypatch): + cfg = tmp_path / "config.toml" + cfg.write_text("") + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + + def _step(step_id="deploy_env", attribution=None): metadata = {"attribution": attribution} if attribution is not None else {} return SimpleNamespace(id=step_id, metadata=metadata) @@ -30,9 +36,9 @@ def _context(instance_id=None, **metadata): def test_pipeline_step_is_task_id_and_step_id(): attribution = deploy_attribution( - _step("deploy_env", {"team": "t"}), _context(task_id="my-pipeline"), + _step("deploy_env", {"group": "g"}), _context(task_id="my-pipeline"), ) - assert attribution == {"team": "t", PIPELINE_STEP_KEY: "my-pipeline_deploy_env"} + assert attribution == {"group": "g", PIPELINE_STEP_KEY: "my-pipeline_deploy_env"} def test_no_task_id_means_no_pipeline_step(): @@ -59,21 +65,21 @@ def test_a_task_authored_pipeline_step_wins(): def test_deploy_attribution_does_not_mutate_the_step(): - step = _step(attribution={"team": "t"}) + step = _step(attribution={"group": "g"}) deploy_attribution(step, _context(task_id="t")) - assert step.metadata["attribution"] == {"team": "t"} + assert step.metadata["attribution"] == {"group": "g"} # --- _build_sandbox_tags -------------------------------------------------------- -def test_sandbox_tags_empty_without_pipeline_step(): - assert _build_sandbox_tags({"team": "t"}) == {} +def test_sandbox_tags_empty_without_attribution(): + assert _build_sandbox_tags({}) == {} -def test_sandbox_tags_carry_both_keys(): - assert _build_sandbox_tags({PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "team": "t"}) == { - PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", +def test_sandbox_tags_carry_every_attribution_key(): + assert _build_sandbox_tags({PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "group": "g"}) == { + PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "group": "g", } @@ -96,9 +102,6 @@ def test_sandbox_tags_at_the_limit_are_unchanged(): assert _build_sandbox_tags({PIPELINE_STEP_KEY: "x" * 63}) == {PIPELINE_STEP_KEY: "x" * 63} -def test_pipeline_step_and_run_id_stay_off_the_app_tags(): - app_tags = _build_cost_attribution_tags({PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1"}) - assert PIPELINE_STEP_KEY not in app_tags and RUN_ID_KEY not in app_tags # --- providers pass the tag to Modal -------------------------------------------- @@ -164,3 +167,20 @@ async def test_vm_create_passes_sandbox_tags(): with pytest.raises(RuntimeError): await provider.create_vm(attribution={PIPELINE_STEP_KEY: "t_s"}) assert create.aio.call_args.kwargs["tags"] == {PIPELINE_STEP_KEY: "t_s"} + + +@pytest.mark.asyncio +async def test_app_tags_are_the_configured_defaults_and_never_a_runs_attribution(tmp_path, monkeypatch): + cfg = tmp_path / "defaults.toml" + cfg.write_text('[sandbox.attribution]\ngroup = "g0"\n') + monkeypatch.setenv("AGENT_ENV_CONFIG", str(cfg)) + provider = _patch_clients(ModalSandboxProvider(app_name="agent-env-test")) + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox._experimental_create") as create: + create.aio = AsyncMock(side_effect=_StopAfterCreate()) + with pytest.raises(RuntimeError): + await provider.create_container( + image_name="img:latest", port=8000, env={}, + attribution={"batch": "b1", RUN_ID_KEY: "inst-1"}, + ) + assert provider._get_app.call_args.args == ("agent-env-test", {"group": "g0"}) + assert create.aio.call_args.kwargs["tags"] == {"batch": "b1", "group": "g0", RUN_ID_KEY: "inst-1"} diff --git a/tst/unit/task_step/attribution_serialization_test.py b/tst/unit/task_step/attribution_serialization_test.py index a3e9b1a0..8db88a73 100644 --- a/tst/unit/task_step/attribution_serialization_test.py +++ b/tst/unit/task_step/attribution_serialization_test.py @@ -20,7 +20,7 @@ from agent_env.task_step.task_steps.deploy_env import DeployEnvTaskStep from agent_env.task_step.task_steps.deploy_sandbox import DeploySandboxTaskStep -ATTRIBUTION = {"cost_center": "research", "team": "platform"} +ATTRIBUTION = {"cost_center": "research", "group": "platform"} def _deploy_env(**kwargs):