diff --git a/packages/agentenv-protocol/README.md b/packages/agentenv-protocol/README.md index ac57fcbb..d5529b04 100644 --- a/packages/agentenv-protocol/README.md +++ b/packages/agentenv-protocol/README.md @@ -477,9 +477,8 @@ tools with (agent-env relays the env card's name: the MultiEnv's declared name, `ATTRIBUTION_PROBE_V1.probe` lets a caller check that an agent forwards attribution upstream. Attribution is an open map of string dimensions, like AgentEnv's cost attribution: the SDK assigns no keys, and each deployment uses whatever dimensions -its model gateway understands (the AgentEnv control plane configures `project_id` -and `task_id` on agents that declare those config fields). An agent records the -dimensions it last sent with a model request and returns them from a bodyless +its model gateway understands. An agent records the dimensions it last sent with a +model request and returns them from a bodyless `POST /ext/attribution-probe` as an `AttributionProbeResponse`: `{"last_seen_attribution": {: , ...}, "last_seen_at_utc": ...}`. The SDK does not record attribution itself, because only the agent knows what it sent. diff --git a/packages/agentenv-protocol/tests/test_a2a_agent.py b/packages/agentenv-protocol/tests/test_a2a_agent.py index 56833d75..9fc6d815 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent.py @@ -1359,7 +1359,7 @@ def __init__(self) -> None: self.sent: dict[str, str] = {} async def run(self, request: TaskRequest) -> TaskResult: - self.sent = {"project_id": "project-1", "task_id": "task-1", "team": "evals"} + self.sent = {"task_id": "task-1", "team": "evals"} return TaskResult.text("ok") @extension(ATTRIBUTION_PROBE_V1.probe) @@ -1381,16 +1381,12 @@ async def probe(self): response = client.post("/a2a", json=_message_request()) assert response.json()["result"]["status"]["state"] == "completed" assert _operation(client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {}).json() == { - "last_seen_attribution": { - "project_id": "project-1", - "task_id": "task-1", - "team": "evals", - }, + "last_seen_attribution": {"task_id": "task-1", "team": "evals"}, "last_seen_at_utc": "2026-01-01T00:00:00Z", } assert ( _operation( - client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {"project_id": "x"} + client, card, ATTRIBUTION_PROBE_V1.uri, "probe", {"unexpected": "x"} ).status_code == 400 ) @@ -2658,16 +2654,6 @@ async def add_bundle_skill(self, request: BundleSkillRequest): ).status_code == 400 ) - assert ( - _operation( - client, - card, - "urn:agentenv:agent-config/v1", - "set", - {"project_id": "project"}, - ).status_code - == 400 - ) assert ( _operation( client, @@ -2820,7 +2806,6 @@ async def add_bundle_skill(self, request: BundleSkillRequest): "parts": [{"kind": "text", "text": "hello"}], "metadata": { "trace_id": "trace-1", - "project_id": "caller-project", }, }, "configuration": {"blocking": True}, @@ -2842,7 +2827,6 @@ async def add_bundle_skill(self, request: BundleSkillRequest): assert seen[0].config.timeout_seconds == 600 assert seen[0].metadata == { "trace_id": "trace-1", - "project_id": "caller-project", "role": "auditor", } assert seen[0].mcp_servers diff --git a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py index e2636e4b..800bfe89 100644 --- a/packages/agentenv-protocol/tests/test_a2a_agent_examples.py +++ b/packages/agentenv-protocol/tests/test_a2a_agent_examples.py @@ -159,7 +159,6 @@ def test_basic_example_uses_typed_config_and_framework_trajectory() -> None: ) schema = config["params"]["methods"]["set"]["request"]["schema"] assert schema["properties"]["provider_token"]["writeOnly"] is True - assert "project_id" not in supported assert "task_id" not in supported configured = client.post( diff --git a/src/agent_env/a2a_agent/validator.py b/src/agent_env/a2a_agent/validator.py index f98f0070..985800ba 100644 --- a/src/agent_env/a2a_agent/validator.py +++ b/src/agent_env/a2a_agent/validator.py @@ -158,8 +158,6 @@ async def validate( VerifyA2AAgentCardStep, VerifyA2AAgentConfigIdentityStep, VerifyA2AAgentMCPStep, - VerifyA2ALitellmAttributionRuntimeStep, - VerifyA2ALitellmAttributionStep, VerifyA2AModalitiesStep, VerifyA2APeerAgentsStep, VerifyA2ARoleStep, @@ -216,14 +214,6 @@ async def validate( protocol_prompt_id = f"{task_id}-protocol-prompt" skill_prompt_id = f"{task_id}-skill-prompt" skill_verifier_id = f"{task_id}-skill-rubric" - # LiteLLM attribution runtime probe: plant unique IDs on context.metadata - # so prompt_agent forwards them via /ext/agent-config; the agent records - # the attribution it emits on the resulting LLM call; the post-prompt - # verify step queries /ext/attribution-probe and asserts the planted - # values came through. Unique per validation run to avoid cross-run - # bleed when /ext/attribution-probe is stateful. - probe_project_id = f"probe-project-{uuid.uuid4().hex[:12]}" - probe_task_id = f"probe-task-{uuid.uuid4().hex[:12]}" # Snapshot validation: plant a unique token on agent A, snapshot, deploy # agent B with the snapshot loaded into a known target context_id, prompt # agent B for recall, grade with a rubric. @@ -268,19 +258,11 @@ async def validate( DeployEnvTaskStep(id=f"{task_id}-deploy-env", version=None, env_id=agent.VALIDATION_ENV_ID), DeployAgentTaskStep(id=deploy_agent_id, version=None, env_ids=[agent.VALIDATION_ENV_ID], a2a_agent_id=agent.id, a2a_agent_version=agent.version), VerifyA2AAgentCardStep(id=f"{task_id}-card", version=None, a2a_agent_id=agent.id, a2a_agent_version=agent.version, depends_on=deploy_agent_dep, fail_task_on_error=False), - VerifyA2ALitellmAttributionStep(id=f"{task_id}-litellm-attribution", version=None, a2a_agent_id=agent.id, a2a_agent_version=agent.version, depends_on=deploy_agent_dep, fail_task_on_error=False), VerifyA2AAgentConfigIdentityStep(id=f"{task_id}-agent-config-identity", version=None, a2a_agent_id=agent.id, a2a_agent_version=agent.version, depends_on=deploy_agent_dep, fail_task_on_error=False), VerifyA2ARoleStep(id=f"{task_id}-role", version=None, a2a_agent_id=agent.id, a2a_agent_version=agent.version, depends_on=deploy_agent_dep, fail_task_on_error=False), VerifyA2ASystemPromptStep(id=f"{task_id}-system-prompt", version=None, a2a_agent_id=agent.id, a2a_agent_version=agent.version, depends_on=deploy_agent_dep, fail_task_on_error=False), PromptAgentTaskStep(id=protocol_prompt_id, version=None, prompt="Hello, respond with OK.", prompt_id=protocol_prompt_id, timeout_seconds=120, depends_on=deploy_agent_dep, fail_task_on_error=False), VerifyCoreA2AProtocolStep(id=f"{task_id}-protocol", version=None, a2a_agent_id=agent.id, prompt_id=protocol_prompt_id, depends_on=[TaskStepDependency(task_step_id=protocol_prompt_id)], fail_task_on_error=False), - VerifyA2ALitellmAttributionRuntimeStep( - id=f"{task_id}-litellm-attribution-runtime", version=None, - a2a_agent_id=agent.id, a2a_agent_version=agent.version, - expected_project_id=probe_project_id, expected_task_id=probe_task_id, - depends_on=[TaskStepDependency(task_step_id=protocol_prompt_id)], - fail_task_on_error=False, - ), *modality_steps, VerifyA2AModalitiesStep(id=f"{task_id}-modalities", version=None, a2a_agent_id=agent.id, probes=grading_probes, depends_on=[TaskStepDependency(task_step_id=sid) for sid in modality_step_ids], fail_task_on_error=False), @@ -424,12 +406,7 @@ async def validate( initial_context = TaskStepContext() if litellm_api_key: initial_context.metadata["user_overrides"] = {"litellm_api_key": litellm_api_key} - # Plant LiteLLM attribution probe values so prompt_agent forwards them - # via /ext/agent-config — the runtime-attribution verify step reads them - # back from /ext/attribution-probe and asserts the agent's bridge - # actually wired them through to outbound LLM calls. - initial_context.metadata["project_id"] = probe_project_id - initial_context.metadata["task_id"] = probe_task_id + initial_context.metadata["task_id"] = task_id num_deploy_steps = 2 run_kwargs = {"on_step_start": on_start, "on_step_complete": on_complete} diff --git a/src/agent_env/artifact/artifacts/docker_image.py b/src/agent_env/artifact/artifacts/docker_image.py index cdb323c1..fc86794a 100644 --- a/src/agent_env/artifact/artifacts/docker_image.py +++ b/src/agent_env/artifact/artifacts/docker_image.py @@ -248,10 +248,7 @@ async def put_from_github( log("create_vm", "Creating VM...", 10) logger.info(f"put_from_github: cloning {owner}/{repo} ref={ref} dockerfile={dockerfile_repo_path} context={context_repo_path}") - sandbox = await get_sandbox_provider().create_vm( - disk_size_gb=20, timeout=1800, exposed_ports=[], - attribution={"project_id": "agent-env-image-builder"}, - ) + sandbox = await get_sandbox_provider().create_vm(disk_size_gb=20, timeout=1800, exposed_ports=[]) try: log("install_git", "Installing git...", 20) await sandbox.exec_script("apt-get update -qq && apt-get install -y -qq git >/dev/null 2>&1") diff --git a/src/agent_env/attribution.py b/src/agent_env/attribution.py index 453b4ff9..b53be48a 100644 --- a/src/agent_env/attribution.py +++ b/src/agent_env/attribution.py @@ -1,8 +1,8 @@ """Cost-attribution dimensions carried from the caller down to a sandbox provider. -Attribution is an open ``dict[str, str]``. Core threads it without reading it; each provider -picks out the keys its billing backend understands and ignores the rest, so a deployment can -attribute on whatever dimensions and terminology it uses. +Attribution is an open ``dict[str, str]``. Core threads it without reading it, and providers +carry every key (Modal as sandbox tags, E2B as sandbox metadata), so a deployment can attribute +on whatever dimensions and terminology it uses. """ from __future__ import annotations diff --git a/src/agent_env/cli/a2a_agent/validate.py b/src/agent_env/cli/a2a_agent/validate.py index 270bd18c..f56f03dc 100644 --- a/src/agent_env/cli/a2a_agent/validate.py +++ b/src/agent_env/cli/a2a_agent/validate.py @@ -16,7 +16,7 @@ def run_validation(agent: A2AAgent, litellm_api_key: str | None = None) -> dict: # Re-fetch agent to get the persisted metadata agent = A2AAgent.get(agent.id, agent.version) - for key in ("validated_agent_card", "validated_a2a_protocol", "validated_a2a_extensions", "validated_data_extensions", "validated_litellm_attribution", "validated_litellm_attribution_runtime"): + for key in ("validated_agent_card", "validated_a2a_protocol", "validated_a2a_extensions", "validated_data_extensions"): value = agent.metadata.get(key, {}) click.echo(click.style(f"\n{key}:", bold=True)) click.echo(json.dumps(value, indent=2)) diff --git a/src/agent_env/cli/task/create.py b/src/agent_env/cli/task/create.py index 21fc0a9d..8bc96849 100644 --- a/src/agent_env/cli/task/create.py +++ b/src/agent_env/cli/task/create.py @@ -11,22 +11,12 @@ @click.command() @click.argument("filepath", type=click.Path(exists=True, dir_okay=False, path_type=Path)) @click.option("--id", "task_id", required=True, help="Task id") -@click.option( - "--project-id", - "project_id", - required=True, - help=( - "Project id for LiteLLM cost attribution. Required for new " - "tasks; tasks created before this requirement landed run without " - "one and the runner falls back to manual attribution." - ), -) @click.option( "--skip-validation", is_flag=True, help="Save even if preflight finds problems (they are printed either way).", ) -def create(filepath: Path, task_id: str, project_id: str, skip_validation: bool): +def create(filepath: Path, task_id: str, skip_validation: bool): """Create a task from a JSON file containing step definitions. \b @@ -48,7 +38,7 @@ def create(filepath: Path, task_id: str, project_id: str, skip_validation: bool) built = _build_steps(steps) _report_preflight( - Task(id=task_id, version=None, steps=[s for _, s in built], project_id=project_id), + Task(id=task_id, version=None, steps=[s for _, s in built]), skip_validation, ) @@ -58,8 +48,8 @@ def create(filepath: Path, task_id: str, project_id: str, skip_validation: bool) click.echo(f" Step {i + 1}/{len(built)}: {step.type} id={saved.id} version={saved.version}") task_steps.append(saved) - click.echo(f"Creating task '{task_id}' with {len(task_steps)} steps (project_id={project_id})...") - task = Task.put(id=task_id, steps=task_steps, project_id=project_id) + click.echo(f"Creating task '{task_id}' with {len(task_steps)} steps...") + task = Task.put(id=task_id, steps=task_steps) click.echo(f"Created task: id={task.id} version={task.version} steps={len(task.steps)}") diff --git a/src/agent_env/cli/task/run.py b/src/agent_env/cli/task/run.py index f4344e7d..83da292f 100644 --- a/src/agent_env/cli/task/run.py +++ b/src/agent_env/cli/task/run.py @@ -311,53 +311,6 @@ def _stamp_agent_env_client_metadata(context_metadata: dict, client_id: str | No context_metadata.setdefault("agent_env_hub", {}).setdefault("caller", client_id) -# Step types whose execution path makes LiteLLM calls (directly or via a -# deployed agent that consumes LITELLM_* env vars). Used to decide whether -# the cost-attribution banner is relevant for a given task. -_LITELLM_USING_STEP_TYPES: frozenset[str] = frozenset({ - "deploy_agent", - "prompt_agent", - "rubrics_verifier", -}) - - -def _task_uses_litellm(task) -> bool: - """True if any of the task's steps are known to make LiteLLM calls.""" - return any(getattr(step, "type", None) in _LITELLM_USING_STEP_TYPES for step in task.steps) - - -def _warn_missing_cost_attribution( - project_id: str | None, - task, -) -> None: - """Yellow CLI banner when `--project-id` is missing on a LiteLLM task. - - Project ID is the primary attribution dimension (it becomes the - LiteLLM `user` field that spend reporting joins on to derive the - customer). Skipped when the task has no LiteLLM-using - steps. Never blocks the run. - """ - if project_id: - return - if not _task_uses_litellm(task): - return - click.echo(click.style( - "⚠ LiteLLM cost attribution is incomplete: --project-id not set.", - fg="yellow", bold=True, - )) - click.echo(click.style( - " Runs will still start, but spend won't be tagged to a " - "project.", - fg="yellow", - )) - click.echo(click.style( - " Pass --project-id on your next run so the " - "cost shows up in your team's budget.", - fg="yellow", - )) - click.echo() - - async def _run_single(task, run_index, task_id, output_dir, agent_model=None, agent_artifact_id=None, start_step=0, context=None): """Execute a single parallel run.""" tag = f"[run {run_index}] " @@ -383,8 +336,6 @@ async def _run_single(task, run_index, task_id, output_dir, agent_model=None, ag @click.option("--a2a-agent-id", default=None, type=str, help="Override A2A agent id for deploy_agent steps") @click.option("--agent-sandbox", default=None, help="Override agent sandbox backend(s) for deploy_agent steps; comma-separated for fallback chain (e.g. 'modal,local')") -@click.option("--project-id", "project_id", default=None, type=str, - help="Project id for LiteLLM cost attribution. Forwarded as the `user` field on every LLM call and as `projectId:` in `metadata.tags`. Optional, but spend won't be attributed to a project without it.") @click.option("--env-sandbox", default=None, help="Override env sandbox backend(s) for deploy_env steps; comma-separated for fallback chain (e.g. 'modal,local')") @click.option("--gateway-env-id", default=None, type=str, @@ -395,7 +346,7 @@ async def _run_single(task, run_index, task_id, output_dir, agent_model=None, ag help="Override the env state type for deploy_env steps") @click.option("--env-state-instance-id", default=None, type=str, help="Attach env to an existing EnvStateInstance") -def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, project_id: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None): +def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, agent_model: str | None, agent_artifact_id: str | None, start_step: int, context_json: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, a2a_agent_id: str | None, agent_sandbox: str | None, env_sandbox: str | None, gateway_env_id: str | None, service_db_env_id: str | None, env_state_type: str | None, env_state_instance_id: str | None): """Run a task by executing its steps sequentially.""" if k < 1: raise click.BadParameter("must be at least 1", param_hint="'--k'") @@ -438,15 +389,12 @@ def run(task_id: str, task_version: int | None, output_dir: str | None, k: int, if env_state_instance_id: initial_context.metadata.setdefault("user_overrides", {})["env_state_instance_id"] = env_state_instance_id initial_context.metadata["task_id"] = task_id - if project_id: - initial_context.metadata["project_id"] = project_id click.echo(f"Fetching task: id={task_id} version={task_version or 'latest'}...") task = Task.get(task_id, version=task_version) click.echo(f"Found task: id={task.id} version={task.version} steps={len(task.steps)}") print_banner() - _warn_missing_cost_attribution(project_id, task) if k > 1: # Parallel runs: compact output with [run N] prefixes @@ -503,13 +451,11 @@ async def _run_all(): help="Override agent sandbox backend(s) for deploy_agent steps; comma-separated for fallback chain (e.g. 'modal,local')") @click.option("--env-sandbox", default=None, help="Override env sandbox backend(s) for deploy_env steps; comma-separated for fallback chain (e.g. 'modal,local')") -@click.option("--project-id", "project_id", default=None, type=str, - help="Project id for LiteLLM cost attribution. Forwarded as the `user` field on every LLM call and as `projectId:` in `metadata.tags`. Optional, but spend won't be attributed to a project without it.") @click.option("--env-state-type", default=None, type=str, help="Override the env state type for deploy_env steps") @click.option("--env-state-instance-id", default=None, type=str, help="Attach env to an existing EnvStateInstance") -def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: int, output_dir: str | None, agent_model: str | None, agent_artifact_id: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, agent_sandbox: str | None, env_sandbox: str | None, project_id: str | None, env_state_type: str | None, env_state_instance_id: str | None): +def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: int, output_dir: str | None, agent_model: str | None, agent_artifact_id: str | None, litellm_api_key: str | None, judge_litellm_api_key: str | None, apply_trajectory_filter: bool | None, agent_sandbox: str | None, env_sandbox: str | None, env_state_type: str | None, env_state_instance_id: str | None): """Run a task in batch against multiple seeds from a CSV file. Each row in the CSV becomes a seed dict passed to the task via @@ -538,7 +484,6 @@ def run_batch(task_id: str, task_version: int | None, seeds: str, concurrency: i click.echo() print_banner() - _warn_missing_cost_attribution(project_id, task) sem = asyncio.Semaphore(concurrency) batch_run_group_id = uuid4().hex @@ -576,8 +521,6 @@ async def _run_seed(index: int, seed: dict): if env_state_instance_id: ctx.metadata.setdefault("user_overrides", {})["env_state_instance_id"] = env_state_instance_id ctx.metadata["task_id"] = task_id - if project_id: - ctx.metadata["project_id"] = project_id click.echo(click.style(f"{tag} Starting...", fg="yellow")) context = await task.run( diff --git a/src/agent_env/explorer/routers/runs.py b/src/agent_env/explorer/routers/runs.py index 37338706..ee61f230 100644 --- a/src/agent_env/explorer/routers/runs.py +++ b/src/agent_env/explorer/routers/runs.py @@ -51,7 +51,6 @@ class RunRequest(BaseModel): context_json: Optional[dict] = None step_overrides: Optional[dict] = None overrides: Optional[dict] = None - project_id: Optional[str] = None priority: Optional[int] = None # Accepted but ignored: the local explorer gets models from the [model] config, not a # per-request key; declared so a client that always sends them isn't rejected. @@ -113,10 +112,8 @@ def _run_metadata(body: "RunRequest") -> dict: # Seed from any user_overrides already in metadata — a re-run inherits the prior # instance's — so the explicit fields merge into them rather than replacing the lot. user_overrides = {**(metadata.get("user_overrides") or {}), **(body.overrides or {})} - for key in ("project_id", "priority"): - value = getattr(body, key) - if value is not None: - user_overrides[key] = value + if body.priority is not None: + user_overrides["priority"] = body.priority # The steps read `step_params` (TaskStep.step_param_overrides); the request field is # named differently. Writing the request name through dropped every override. if body.step_overrides is not None: diff --git a/src/agent_env/explorer/ui/src/components/start-runs-panel.tsx b/src/agent_env/explorer/ui/src/components/start-runs-panel.tsx index e2c582b6..b394d97a 100644 --- a/src/agent_env/explorer/ui/src/components/start-runs-panel.tsx +++ b/src/agent_env/explorer/ui/src/components/start-runs-panel.tsx @@ -60,7 +60,6 @@ interface StartRunsPanelProps { taskId: string; taskVersion: number; taskSteps?: ReadonlyArray>; - taskProjectId?: string; onStarted?: () => void; onCompleted?: () => void; } @@ -81,15 +80,11 @@ export function StartRunsPanel({ taskId, taskVersion, taskSteps, - taskProjectId, onStarted, onCompleted, }: StartRunsPanelProps) { const [persisted] = useState(loadPersistedConfig); - const [projectId, setProjectId] = useState( - taskProjectId?.trim() ?? '', - ); const [agentModel, setAgentModel] = useState( persisted.agent_model ?? '', ); @@ -398,7 +393,6 @@ export function StartRunsPanel({ const submitSeeds = requiresSeeds ? seeds ?? undefined : undefined; await startRuns({ version: taskVersion, - projectId: projectId.trim() || undefined, overrides: buildOverrides(), count: submitSeeds ? undefined : parsedCount ?? undefined, seeds: submitSeeds, @@ -408,7 +402,6 @@ export function StartRunsPanel({ canStart, startRuns, taskVersion, - projectId, buildOverrides, requiresSeeds, seeds, diff --git a/src/agent_env/explorer/ui/src/components/task-detail-page.tsx b/src/agent_env/explorer/ui/src/components/task-detail-page.tsx index 0576d4d7..ea7bfa9b 100644 --- a/src/agent_env/explorer/ui/src/components/task-detail-page.tsx +++ b/src/agent_env/explorer/ui/src/components/task-detail-page.tsx @@ -1329,11 +1329,6 @@ export function TaskDetailPage({ taskId={String(task?.id ?? '')} taskVersion={Number(task?.version ?? 1)} taskSteps={steps} - taskProjectId={ - typeof task?.project_id === 'string' - ? (task.project_id as string) - : undefined - } onStarted={handleStartedRuns} onCompleted={fetchRunGroups} /> diff --git a/src/agent_env/explorer/ui/src/components/task-runner-page.tsx b/src/agent_env/explorer/ui/src/components/task-runner-page.tsx index d4107405..b11b8457 100644 --- a/src/agent_env/explorer/ui/src/components/task-runner-page.tsx +++ b/src/agent_env/explorer/ui/src/components/task-runner-page.tsx @@ -148,12 +148,7 @@ export function TaskRunnerPage({ /* --- run task --- */ const buildRunBody = useCallback( - (opts?: RunOptions) => - runBodyFor( - opts, - taskVersion, - new URLSearchParams(window.location.search).get('projectId'), - ), + (opts?: RunOptions) => runBodyFor(opts, taskVersion), [taskVersion], ); diff --git a/src/agent_env/explorer/ui/src/hooks/useRunGroupPolling.ts b/src/agent_env/explorer/ui/src/hooks/useRunGroupPolling.ts index c55416c7..f44b86b5 100644 --- a/src/agent_env/explorer/ui/src/hooks/useRunGroupPolling.ts +++ b/src/agent_env/explorer/ui/src/hooks/useRunGroupPolling.ts @@ -16,7 +16,6 @@ export interface RunOverrides { interface StartRunsConfig { version: number; - projectId?: string; overrides?: RunOverrides; count?: number; seeds?: Record[]; @@ -96,7 +95,6 @@ export function useRunGroupPolling(taskId: string) { const body: Record = { version: config.version, }; - if (config.projectId) body.project_id = config.projectId; if (config.overrides && Object.keys(config.overrides).length > 0) { body.overrides = config.overrides; } diff --git a/src/agent_env/explorer/ui/src/lib/task-runner-run.smoke.ts b/src/agent_env/explorer/ui/src/lib/task-runner-run.smoke.ts index 3642c996..08996336 100644 --- a/src/agent_env/explorer/ui/src/lib/task-runner-run.smoke.ts +++ b/src/agent_env/explorer/ui/src/lib/task-runner-run.smoke.ts @@ -58,13 +58,11 @@ async function main(): Promise { const TASK = '@local/agentenv-framework/hello/hello'; // The request a run sends. - const body = buildRunBody({ start_step: 2, context_from_instance_id: 'i-1', context_json: { x: 1 } }, 3, ''); + const body = buildRunBody({ start_step: 2, context_from_instance_id: 'i-1', context_json: { x: 1 } }, 3); assert(body.version === 3 && body.priority === 0, 'run body: task version, interactive priority'); assert(body.start_step === 2, 'run body: start step'); assert(body.context_from_instance_id === 'i-1' && !('context_json' in body), 'run body: instance context wins over inline context'); - assert(!('project_id' in body), 'run body: an empty projectId is not sent'); - assert(buildRunBody({ version: 5 }, 3, 'p-9').version === 5, 'run body: an explicit version wins'); - assert(buildRunBody(undefined, null, 'p-9').project_id === 'p-9', 'run body: projectId is forwarded'); + assert(buildRunBody({ version: 5 }, 3).version === 5, 'run body: an explicit version wins'); // A run starts. const ok = fakeFetch(200, { workflow_id: 'local-abc', instance_id: 'i-2', status: 'QUEUED' }); diff --git a/src/agent_env/explorer/ui/src/lib/task-runner-run.ts b/src/agent_env/explorer/ui/src/lib/task-runner-run.ts index 6fcc5441..19998689 100644 --- a/src/agent_env/explorer/ui/src/lib/task-runner-run.ts +++ b/src/agent_env/explorer/ui/src/lib/task-runner-run.ts @@ -32,7 +32,6 @@ function taskUrl(base: string, taskId: string): string { export function buildRunBody( opts: RunOptions | undefined, taskVersion: number | null, - projectId: string | null, ): Record { const body: Record = {}; body.version = opts?.version ?? taskVersion ?? undefined; @@ -45,8 +44,6 @@ export function buildRunBody( else if (opts?.context_json) body.context_json = opts.context_json; if (opts?.overrides_from_instance_id) body.overrides_from_instance_id = opts.overrides_from_instance_id; - // Only-if-truthy so "" can't defeat the backend fallback. - if (projectId) body.project_id = projectId; return body; } diff --git a/src/agent_env/providers/sandbox_providers/e2b/provider.py b/src/agent_env/providers/sandbox_providers/e2b/provider.py index c2aa1df9..92f3eb8b 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__) @@ -141,9 +141,18 @@ async def create_vm( "is fixed by the selected template and cannot be configured per sandbox", disk_size_gb, ) + metadata = { + key: value + for key, value in apply_default_attribution(dict(attribution or {})).items() + if value is not None + } + if _EXPOSED_PORTS_METADATA_KEY in metadata: + raise ValueError( + f"attribution key {_EXPOSED_PORTS_METADATA_KEY!r} is reserved: " + "the E2B provider stores the sandbox's exposed ports under it" + ) 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 {})) 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 86808e64..78b46ea6 100644 --- a/src/agent_env/providers/sandbox_providers/modal_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_sandbox.py @@ -36,6 +36,7 @@ ECR_READER_SECRET_KEYS = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_REGION"] _MODAL_TAG_MAX_LEN = 63 +_MODAL_TAG_KEY = re.compile(r"[a-zA-Z0-9._-]{1,63}") def _resolve_app_base_name(app_name: Optional[str]) -> str: @@ -203,9 +204,7 @@ def __init__( # cannot attach a GPU. V1 has no ``i6pn``, so a GPU sandbox has no private-IPv6 # east-west networking (fine for a single sandbox; not for an i6pn mesh). self._gpu = gpu - # One Modal App per attribution scope (keyed by project_id), looked up on demand. - # Billing aggregates per App and the dashboard is browsable by App name, so a - # scope-named, scope-tagged App is what makes cost both attributable and findable. + # Looked up on demand and cached by name. self._apps: dict[str, modal.App] = {} async def _get_client(self) -> modal.Client: @@ -255,15 +254,12 @@ async def create_container( effective = self.effective_network_policy(network_policy) attribution = dict(attribution or {}) - # Place the sandbox under a per-project App: Modal billing aggregates cost per App and - # the dashboard is browsable by App name, so a scope-named, scope-tagged App is what - # makes cost both attributable and findable. We attribute at the App level only — - # sandbox-level tags don't show up in billing. - app_tags = _build_cost_attribution_tags(attribution) - app_name = _app_name_for_project(self._app_name, app_tags.get("project_id")) + # Modal bills by App tags, and the one shared App is tagged once, so it carries only the + # deployment's [sandbox.attribution]; the run's own attribution goes on its sandbox. + app_tags = _attribution_tags({}) + sandbox_tags = _attribution_tags(attribution) + 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) from agent_env.store.image_store import ( @@ -406,56 +402,17 @@ async def get_sandbox(self, sandbox_id: str) -> Sandbox: return self._sandbox_cls(sb, tunnel_urls, i6pn_address=i6pn_address) -def _app_name_for_project(base: str, project_id: Optional[str]) -> str: - """Map a project_id to a stable Modal App name for per-project cost attribution. - - The app name is what shows up (and is searchable) in the Modal dashboard, so it encodes - the project. Names are sanitized to ``[a-zA-Z0-9._-]``, kept under Modal's 64-char limit, - and not double-prefixed when ``project_id`` already starts with ``base``. Falls back to - ``base`` when no project_id is available. - """ - if not project_id: - return base - slug = re.sub(r"[^a-zA-Z0-9._-]", "-", project_id).strip("-") or base - name = slug if (slug == base or slug.startswith(f"{base}-")) else f"{base}-{slug}" - if len(name) > 64: - logger.warning( - f"Modal app name for project_id={project_id!r} exceeds 64 chars; truncating to " - f"{name[:64]!r}. Distinct project_ids that share this prefix will be cost-attributed " - f"together." - ) - name = name[:64] - return name - - -def _build_cost_attribution_tags(attribution: Attribution) -> dict[str, str]: - """Build the Modal App tag set used for cost attribution. - - 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", "project_id") - if resolved.get(name) is not None - } - - -def _build_sandbox_tags(attribution: Attribution) -> dict[str, str]: - """Build the Modal Sandbox tag set: ``pipeline_step`` and ``run_id``, each when set. - - 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. - """ - return { - key: _modal_tag_value(attribution[key]) - for key in (PIPELINE_STEP_KEY, RUN_ID_KEY) - if attribution.get(key) - } +def _attribution_tags(attribution: Attribution) -> dict[str, str]: + """``attribution`` as Modal tags, any keys: unset ones filled from config.toml + ``[sandbox.attribution]``, None omitted. A key Modal can't take is refused, not rewritten.""" + tags = {} + for key, value in apply_default_attribution(attribution).items(): + if value is None: + continue + if not _MODAL_TAG_KEY.fullmatch(key): + raise ValueError(f"attribution key {key!r} can't be a Modal tag: use 1-63 of a-z A-Z 0-9 . _ -") + tags[key] = _modal_tag_value(value) + return tags SANDBOX_STARTED_EVENT = "agent_env.modal_sandbox_started" @@ -484,7 +441,9 @@ 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, + PIPELINE_STEP_KEY: sandbox_tags.get(PIPELINE_STEP_KEY), + RUN_ID_KEY: sandbox_tags.get(RUN_ID_KEY), "cpu": cpu, "memory_mb": memory, "gpu": gpu, @@ -493,6 +452,8 @@ async def _log_sandbox_started( def _modal_tag_value(value: str) -> str: + """Modal rejects the whole create on an invalid tag, so ``value`` is cut to ``[a-zA-Z0-9._-]``; + past 63 chars it keeps a prefix plus a hash of the whole, so long values stay 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 91173f4c..7391b786 100644 --- a/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py +++ b/src/agent_env/providers/sandbox_providers/modal_vm_sandbox.py @@ -29,9 +29,7 @@ from agent_env.providers.sandbox_providers.modal_sandbox import ( _ModalProcessAdapter, - _app_name_for_project, - _build_cost_attribution_tags, - _build_sandbox_tags, + _attribution_tags, _log_sandbox_started, _resolve_app_base_name, _tunnel_url, @@ -309,10 +307,10 @@ 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_name = _app_name_for_project(self._app_name, app_tags.get("project_id")) + app_tags = _attribution_tags({}) + sandbox_tags = _attribution_tags(attribution) + app_name = self._app_name app = await self._get_app(app_name, app_tags) - sandbox_tags = _build_sandbox_tags(attribution) client = await self._get_client() effective = self.effective_network_policy(network_policy) ports = list(exposed_ports or []) diff --git a/src/agent_env/task/task.py b/src/agent_env/task/task.py index 1e638440..5626789d 100644 --- a/src/agent_env/task/task.py +++ b/src/agent_env/task/task.py @@ -305,12 +305,10 @@ def __init__( id: str, version: Optional[int], steps: list[TaskStep] | None = None, - project_id: Optional[str] = None, ): self.id = id self.version = version self.steps: list[TaskStep] = steps or [] - self.project_id = project_id self._validate_dag(self.steps) @classmethod @@ -361,15 +359,12 @@ def preflight(self) -> list[str]: return problems def to_dict(self) -> dict: - doc: dict = { + return { "id": self.id, "type": self.type, "version": self.version, "steps": [step.to_dict() for step in self.steps], } - if self.project_id is not None: - doc["project_id"] = self.project_id - return doc @classmethod def from_dict(cls, data: dict) -> "Task": @@ -391,7 +386,6 @@ def from_dict(cls, data: dict) -> "Task": id=data["id"], version=data.get("version"), steps=steps, - project_id=data.get("project_id"), ) async def run( diff --git a/src/agent_env/task_step/task_steps/a2a_agent_validator/__init__.py b/src/agent_env/task_step/task_steps/a2a_agent_validator/__init__.py index 0c3aab46..b879cad9 100644 --- a/src/agent_env/task_step/task_steps/a2a_agent_validator/__init__.py +++ b/src/agent_env/task_step/task_steps/a2a_agent_validator/__init__.py @@ -3,8 +3,6 @@ from .verify_a2a_agent_mcp import VerifyA2AAgentMCPStep from .verify_a2a_core_protocol import VerifyCoreA2AProtocolStep from .verify_a2a_install import VerifyA2AInstallStep -from .verify_a2a_litellm_attribution import VerifyA2ALitellmAttributionStep -from .verify_a2a_litellm_attribution_runtime import VerifyA2ALitellmAttributionRuntimeStep from .verify_a2a_modalities import VerifyA2AModalitiesStep from .verify_a2a_peer_agents import VerifyA2APeerAgentsStep from .verify_a2a_role import VerifyA2ARoleStep @@ -19,8 +17,6 @@ "VerifyA2AAgentMCPStep", "VerifyCoreA2AProtocolStep", "VerifyA2AInstallStep", - "VerifyA2ALitellmAttributionStep", - "VerifyA2ALitellmAttributionRuntimeStep", "VerifyA2AModalitiesStep", "VerifyA2APeerAgentsStep", "VerifyA2ARoleStep", diff --git a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution.py b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution.py deleted file mode 100644 index 2a7745c0..00000000 --- a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution.py +++ /dev/null @@ -1,124 +0,0 @@ -"""Verify A2A agent declares LiteLLM cost-attribution support on its card. - -Static check: reads the agent's `/.well-known/agent.json` (already fetched -at deploy time and stored on `deployed.a2a_card`) and inspects the -`urn:agentenv:agent-config/v1` extension. An agent "supports LiteLLM -attribution" when its agent-config extension advertises **both** -`project_id` and `task_id` in its `params.methods.set.request.supported` -list. - -The check tells consumers (e.g. claude-code's auto-integration flow, -the hub UI) which agents will honor per-prompt attribution delivered -through `/ext/agent-config` from `PromptAgentTaskStep`. Agents that -don't advertise it will have those fields filtered out at the worker -side with a warning, and spend logs from those agents will not carry -the `projectId:` / `taskId:` tags. - -This is a card-level static check only. It does NOT verify the agent -actually threads attribution into its outbound LLM calls — that's an -agent-implementation concern (e.g. an agent bridge's own -`_apply_litellm_attribution` helper). A future runtime probe could be added -if a self-reporting endpoint is exposed; for now, the static check -mirrors what an external integrator can discover via the AgentCard. -""" - -from __future__ import annotations - -import logging -from typing import ClassVar, Optional - -from agent_env.task_step.context import TaskStepContext -from agent_env.task_step.task_step import TaskStep, TaskStepDependency - -logger = logging.getLogger(__name__) - -# Body fields the worker (PromptAgentTaskStep) sends via /ext/agent-config -# for LiteLLM cost attribution. Both must be in the agent-config extension's -# `supported` list for the agent to be considered attribution-capable. -ATTRIBUTION_FIELDS = ("project_id", "task_id") - - -class VerifyA2ALitellmAttributionStep(TaskStep): - type: ClassVar[str] = "verify_a2a_litellm_attribution" - - def __init__( - self, - id: str, - version: Optional[int], - a2a_agent_id: str, - a2a_agent_version: Optional[int] = None, - depends_on: Optional[list[TaskStepDependency]] = None, - fail_task_on_error: bool = True, - ): - super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) - self.a2a_agent_id = a2a_agent_id - self.a2a_agent_version = a2a_agent_version - - def to_dict(self) -> dict: - base = super().to_dict() - base["a2a_agent_id"] = self.a2a_agent_id - base["a2a_agent_version"] = self.a2a_agent_version - return base - - @classmethod - def from_dict(cls, data: dict) -> "VerifyA2ALitellmAttributionStep": - return cls( - **cls._base_from_dict(data), - a2a_agent_id=data["a2a_agent_id"], - a2a_agent_version=data.get("a2a_agent_version"), - ) - - async def execute(self, context: TaskStepContext) -> TaskStepContext: - from agent_env.a2a_agent import A2AAgent - - deployed = next((a for a in context.deployed_agents), None) - if deployed is None: - raise RuntimeError("No deployed agent found in context") - - card = deployed.a2a_card or {} - agent = A2AAgent.get(self.a2a_agent_id, self.a2a_agent_version) - - ext = A2AAgent.find_extension(card, A2AAgent.EXT_AGENT_CONFIG) - supported: list[str] = [] - if ext: - supported = ( - ((ext.get("params") or {}).get("methods") or {}) - .get("set", {}) - .get("request", {}) - .get("supported") - or [] - ) - - advertises = {field: field in supported for field in ATTRIBUTION_FIELDS} - missing = [field for field, present in advertises.items() if not present] - passed = ext is not None and not missing - - if not ext: - logger.warning( - "Agent '%s' does not advertise the agent-config extension — " - "LiteLLM attribution cannot flow.", self.a2a_agent_id, - ) - elif missing: - logger.warning( - "Agent '%s' advertises agent-config but is missing attribution " - "fields in its `supported` set: %s. Spend logs from this agent " - "will not carry projectId/taskId tags. See an A2A agent bridge's " - "`_apply_litellm_attribution` helper for a reference implementation.", - self.a2a_agent_id, missing, - ) - else: - logger.info( - "Agent '%s' supports LiteLLM attribution: advertises %s in " - "agent-config.supported", self.a2a_agent_id, list(ATTRIBUTION_FIELDS), - ) - - validated = { - "passed": passed, - "extension_present": ext is not None, - "supported_fields": supported, - "advertises": advertises, - "missing": missing, - } - agent.update_metadata({**agent.metadata, "validated_litellm_attribution": validated}) - context.metadata.setdefault("verifications", {})["a2a_litellm_attribution"] = validated - return context diff --git a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution_runtime.py b/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution_runtime.py deleted file mode 100644 index a0f416f1..00000000 --- a/src/agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution_runtime.py +++ /dev/null @@ -1,173 +0,0 @@ -"""Verify A2A agent actually emits LiteLLM attribution on outbound LLM calls. - -Optional behavioral counterpart to the static -`VerifyA2ALitellmAttributionStep`. When a runtime advertises its private probe, -this step asks whether it actually forwards project_id/task_id to the LLM call. -The probe is not part of the general AgentEnv SDK contract. - -How it works: - 1. The validator plants unique probe values in `initial_context.metadata` - before any prompt fires. - 2. The standard `prompt_agent` flow POSTs those values via - `/ext/agent-config` and sends a prompt — the agent processes it, - makes an outbound LLM call, and records the attribution it claims - to have emitted in module-level state. - 3. This step queries the agent's `/ext/attribution-probe` extension to - read back that recorded state. - 4. Passes when the recorded values match the planted probe values. - -Catches: - - Bridge ignores project_id/task_id from agent-config - - Bridge sets the wrong env var / header names - - Agent advertises support but never wires it through - -Does NOT catch: - - CLI eats the env var and never emits the actual HTTP header (probe - trusts the bridge's claim about what it sent) - - Proxy strips the header on receive - -For full end-to-end coverage of the latter two, a spend-log probe would -be needed (requires proxy admin access). -""" - -from __future__ import annotations - -import logging -from typing import ClassVar, Optional - -import httpx - -from agent_env.task_step.context import TaskStepContext -from agent_env.task_step.task_step import TaskStep, TaskStepDependency - -logger = logging.getLogger(__name__) - -EXT_ATTRIBUTION_PROBE_URI = "urn:agentenv:attribution-probe/v1" -EXT_ATTRIBUTION_PROBE_PATH = "/ext/attribution-probe" - - -class VerifyA2ALitellmAttributionRuntimeStep(TaskStep): - type: ClassVar[str] = "verify_a2a_litellm_attribution_runtime" - - def __init__( - self, - id: str, - version: Optional[int], - a2a_agent_id: str, - expected_project_id: str, - expected_task_id: str, - a2a_agent_version: Optional[int] = None, - depends_on: Optional[list[TaskStepDependency]] = None, - fail_task_on_error: bool = True, - ): - super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error) - self.a2a_agent_id = a2a_agent_id - self.a2a_agent_version = a2a_agent_version - self.expected_project_id = expected_project_id - self.expected_task_id = expected_task_id - - def to_dict(self) -> dict: - base = super().to_dict() - base["a2a_agent_id"] = self.a2a_agent_id - base["a2a_agent_version"] = self.a2a_agent_version - base["expected_project_id"] = self.expected_project_id - base["expected_task_id"] = self.expected_task_id - return base - - @classmethod - def from_dict(cls, data: dict) -> "VerifyA2ALitellmAttributionRuntimeStep": - return cls( - **cls._base_from_dict(data), - a2a_agent_id=data["a2a_agent_id"], - a2a_agent_version=data.get("a2a_agent_version"), - expected_project_id=data["expected_project_id"], - expected_task_id=data["expected_task_id"], - ) - - async def execute(self, context: TaskStepContext) -> TaskStepContext: - from agent_env.a2a_agent import A2AAgent - - deployed = next((a for a in context.deployed_agents), None) - if deployed is None: - raise RuntimeError("No deployed agent found in context") - - agent = A2AAgent.get(self.a2a_agent_id, self.a2a_agent_version) - - # This is a runtime-specific diagnostic, not a general SDK capability. - # Its absence is therefore not a conformance failure. - probe_ext = A2AAgent.find_extension(deployed.a2a_card or {}, EXT_ATTRIBUTION_PROBE_URI) - if probe_ext is None: - logger.info( - "Agent '%s' does not advertise the attribution-probe extension " - "(%s) — skipping optional runtime attribution verification.", - self.a2a_agent_id, EXT_ATTRIBUTION_PROBE_URI, - ) - validated = { - "passed": True, - "skipped": True, - "extension_present": False, - "reason": "optional runtime attribution probe not advertised", - "expected_project_id": self.expected_project_id, - "expected_task_id": self.expected_task_id, - } - agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated}) - context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated - return context - - endpoint = (probe_ext.get("params") or {}).get("endpoint", EXT_ATTRIBUTION_PROBE_PATH) - url = (deployed.a2a_url or deployed.api_url) + endpoint - logger.info("Querying attribution probe at %s", url) - - try: - async with httpx.AsyncClient(timeout=10.0) as client: - resp = await client.post(url, json={}) - resp.raise_for_status() - probe_data = resp.json() - except Exception as e: - logger.warning("attribution-probe request failed: %s", e) - validated = { - "passed": False, - "extension_present": True, - "reason": f"probe request failed: {e}", - "expected_project_id": self.expected_project_id, - "expected_task_id": self.expected_task_id, - } - agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated}) - context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated - return context - - last_seen = probe_data.get("last_seen_attribution") or {} - observed_project_id = last_seen.get("project_id") - observed_task_id = last_seen.get("task_id") - - project_matches = observed_project_id == self.expected_project_id - task_matches = observed_task_id == self.expected_task_id - passed = project_matches and task_matches - - if passed: - logger.info( - "Runtime attribution verified for '%s': observed projectId=%s taskId=%s", - self.a2a_agent_id, observed_project_id, observed_task_id, - ) - else: - logger.warning( - "Runtime attribution mismatch for '%s': expected projectId=%s taskId=%s, " - "observed projectId=%s taskId=%s", - self.a2a_agent_id, self.expected_project_id, self.expected_task_id, - observed_project_id, observed_task_id, - ) - - validated = { - "passed": passed, - "extension_present": True, - "expected_project_id": self.expected_project_id, - "expected_task_id": self.expected_task_id, - "observed_project_id": observed_project_id, - "observed_task_id": observed_task_id, - "project_matches": project_matches, - "task_matches": task_matches, - "last_seen_at_utc": probe_data.get("last_seen_at_utc"), - } - agent.update_metadata({**agent.metadata, "validated_litellm_attribution_runtime": validated}) - context.metadata.setdefault("verifications", {})["a2a_litellm_attribution_runtime"] = validated - return context diff --git a/src/agent_env/task_step/task_steps/deploy_agent.py b/src/agent_env/task_step/task_steps/deploy_agent.py index cf5d4433..39bbe298 100644 --- a/src/agent_env/task_step/task_steps/deploy_agent.py +++ b/src/agent_env/task_step/task_steps/deploy_agent.py @@ -300,13 +300,6 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: if resolved_litellm_base_url: env_vars["LITELLM_BASE_URL"] = resolved_litellm_base_url - # Cost-attribution (project_id / task_id) is delivered per-prompt - # via the A2A `/ext/agent-config` extension from prompt_agent.py — - # not baked as container env vars at deploy time. See - # the openai_agents_sdk A2A agent in the universe-generation pipeline - # for the reference implementation that reads agent-config and threads - # `user` / `metadata.tags` into each LiteLLM call. - resolved_sandbox_type = user_overrides.get("agent_sandbox") or self.sandbox_type deploy_kwargs = {"env_vars": env_vars if env_vars else None} if agent_disk_size_gb is not None: diff --git a/src/agent_env/task_step/task_steps/prompt_agent.py b/src/agent_env/task_step/task_steps/prompt_agent.py index e5a76c9c..661ba4c5 100644 --- a/src/agent_env/task_step/task_steps/prompt_agent.py +++ b/src/agent_env/task_step/task_steps/prompt_agent.py @@ -402,7 +402,6 @@ async def _execute_conversation( "harness": harness, "max_turns": max_turns, "max_thinking_tokens": max_thinking_tokens, "output_format": self.output_format, "timeout_seconds": self.timeout_seconds, - "project_id": context.metadata.get("project_id"), "task_id": context.metadata.get("task_id"), "agentenv_tools": self.agentenv_tools, } diff --git a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py index 831b40d7..b8a2725a 100644 --- a/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +++ b/src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py @@ -569,22 +569,9 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext: or config.get_litellm_base_url() ) env_vars = {"LITELLM_API_KEY": api_key, "LITELLM_BASE_URL": litellm_base_url} - judge_project_id = context.metadata.get("project_id") - judge_customer = context.metadata.get("customer_id") - if judge_customer: - env_vars["LITELLM_USER"] = judge_customer - if judge_project_id: - env_vars["LITELLM_PROJECT_ID"] = judge_project_id resolved_judge_id = self.judge_a2a_agent_id or config.get_default_a2a_agent_id() a2a_agent = A2AAgent.get(resolved_judge_id) - deploy_kwargs = { - "env_vars": env_vars, - "attribution": { - k: v for k, v in - (("project_id", judge_project_id), ("customer", judge_customer)) - if v - }, - } + deploy_kwargs = {"env_vars": env_vars} resolved_sandbox_type = overrides.get("agent_sandbox") or self.judge_sandbox_type if resolved_sandbox_type: deploy_kwargs["sandbox_type"] = resolved_sandbox_type @@ -931,11 +918,9 @@ async def _prompt_llm_judge( import litellm from agent_env.config import get_config - from agent_env.utils.litellm_attribution import build_litellm_cost_attribution_kwargs overrides = context.metadata.get("user_overrides", {}) config = get_config() - cost_kwargs = build_litellm_cost_attribution_kwargs(context.metadata) call_config = config.resolve_model_call( model, default_api_key=overrides.get("judge_litellm_api_key") or overrides.get("litellm_api_key"), @@ -968,7 +953,6 @@ async def _prompt_llm_judge( messages=[{"role": "user", "content": user_content}], timeout=self.judge_timeout_seconds, **call_config.client_kwargs(), - **cost_kwargs, ) if use_response_format: request_kwargs[ModelParam.RESPONSE_FORMAT] = response_format @@ -1010,7 +994,6 @@ async def _configure_judge_a2a( "max_thinking_tokens": judge_max_thinking_tokens, "output_format": self._spec().output_format, "timeout_seconds": self.judge_timeout_seconds, - "project_id": context.metadata.get("project_id"), "task_id": context.metadata.get("task_id"), } desired_config = {k: v for k, v in desired_config.items() if v is not None} diff --git a/src/agent_env/utils/litellm_attribution.py b/src/agent_env/utils/litellm_attribution.py deleted file mode 100644 index 2ee408b8..00000000 --- a/src/agent_env/utils/litellm_attribution.py +++ /dev/null @@ -1,59 +0,0 @@ -"""LiteLLM cost-attribution helpers. - -Every LiteLLM request agent-env sends carries two attribution fields so the -proxy's spend logs can be rolled up by who ran the work: - -- `user` = the `project_id`. Spend reporting joins on it to attribute a - request to its owning project (and, downstream, to the customer); callers - never pass a customer identifier directly. -- `metadata.tags` = secondary attribution: `projectId:` (duplicates - `user`, so tag-based reports agree with user-based ones) and `taskId:` - (the agent-env Task), so spend rolls up per project and per task. - -Both fields are sourced from `TaskStepContext.metadata` so the run -group's project/task flows into every in-process LLM call (judge LLMs, etc.) -without each call site having to plumb them by hand. - -Use this helper at every `litellm.completion` / `litellm.acompletion` -call site so attribution is consistent across the codebase. -""" - -from __future__ import annotations - -import logging -from typing import Any - -from agent_env.config.model import ModelParam - -logger = logging.getLogger(__name__) - - -def build_litellm_cost_attribution_kwargs(metadata: dict[str, Any] | None) -> dict[str, Any]: - """Build LiteLLM cost-attribution kwargs (`user`, `metadata`). - - Reads `project_id` and `task_id` from `metadata` (typically - `TaskStepContext.metadata`). Missing `project_id` warns but doesn't - block — the request is sent regardless. `task_id` is best-effort. - - Returns a kwargs dict to spread into `litellm.completion(**kwargs)` / - `litellm.acompletion(**kwargs)`. - """ - metadata = metadata or {} - kwargs: dict[str, Any] = {} - project_id = metadata.get("project_id") - task_id = metadata.get("task_id") - if project_id: - kwargs[ModelParam.USER] = project_id - tags = [f"projectId:{project_id}"] - if task_id: - tags.append(f"taskId:{task_id}") - kwargs[ModelParam.METADATA] = {"tags": tags} - else: - logger.warning( - "context.metadata.project_id is not set; LiteLLM cost " - "attribution will not be tagged with a project. Request " - "is still being sent. Fill the Project ID field in the run UI " - "(or pass project_id at the API call site) so spend is " - "attributed correctly." - ) - return kwargs diff --git a/tst/data/a2a_agent/agent.py b/tst/data/a2a_agent/agent.py index be1193dd..6fc085c4 100644 --- a/tst/data/a2a_agent/agent.py +++ b/tst/data/a2a_agent/agent.py @@ -26,7 +26,6 @@ class EchoAgentConfig(AgentConfig): model: str | None = None system_prompt: str | None = None - project_id: str | None = None task_id: str | None = None model_params: WriteOnly[dict[str, Any] | None] = None diff --git a/tst/installer/test_container_journey.py b/tst/installer/test_container_journey.py index 0c226767..6c0c45ab 100644 --- a/tst/installer/test_container_journey.py +++ b/tst/installer/test_container_journey.py @@ -155,7 +155,7 @@ def test_a_new_user_installs_agent_env_adds_plugins_runs_a_task_and_removes_them assert GRADER not in container.ok("agent-env", "plugin", "list").stdout container.ok("agent-env", "plugin", "add", grader, "--yes") - container.ok("agent-env", "task", "create", "task.json", "--id", "journey", "--project-id", "journey") + container.ok("agent-env", "task", "create", "task.json", "--id", "journey") container.ok("agent-env", "task", "run", "--id", "journey", "--output-dir", "/work/out") (context,) = (journey.work / "out").glob("*.json") metadata = json.loads(context.read_text())["metadata"] @@ -166,7 +166,7 @@ def test_a_new_user_installs_agent_env_adds_plugins_runs_a_task_and_removes_them assert blocked.returncode == 1 and "stored task(s) use its step types (toy_navigate)" in blocked.stderr container.ok("agent-env", "plugin", "remove", BROWSER, "--yes", "--force") - gone = container.run("agent-env", "task", "create", "task.json", "--id", "again", "--project-id", "journey") + gone = container.run("agent-env", "task", "create", "task.json", "--id", "again") assert gone.returncode != 0 and "Unknown task step type: toy_navigate" in gone.stdout + gone.stderr finally: container.remove() diff --git a/tst/unit/cli/litellm_step_types_test.py b/tst/unit/cli/litellm_step_types_test.py deleted file mode 100644 index f3438663..00000000 --- a/tst/unit/cli/litellm_step_types_test.py +++ /dev/null @@ -1,6 +0,0 @@ -from agent_env.cli.task.run import _LITELLM_USING_STEP_TYPES -from agent_env.task_step.registry import _builtin_registry - - -def test_litellm_step_types_are_built_in_step_types(): - assert _LITELLM_USING_STEP_TYPES <= set(_builtin_registry()) diff --git a/tst/unit/cli/test_task_create_preflight.py b/tst/unit/cli/test_task_create_preflight.py index 6449ad4c..3aca34f0 100644 --- a/tst/unit/cli/test_task_create_preflight.py +++ b/tst/unit/cli/test_task_create_preflight.py @@ -41,7 +41,7 @@ def store(monkeypatch): def _run(tmp_path, *extra): return CliRunner().invoke(create, [ - str(_task_file(tmp_path)), "--id", "t", "--project-id", "p", *extra, + str(_task_file(tmp_path)), "--id", "t", *extra, ]) 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 46f96697..34627bd2 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={"project_id": "p1"}) + ttl_seconds=60, disk_size_gb=20, cpu=2.0, memory_mb=4096, priority=3, attribution={"team": "t1"}) 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": {"project_id": "p1"}, + "ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"team": "t1"}, "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={"project_id": "p1"}) + 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": {"project_id": "p1"}} + assert seen == {"ttl_seconds": 60, "disk_size_gb": 20, "cpu": 2.0, "memory_mb": 4096, "priority": 3, "attribution": {"team": "t1"}} @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 553c615a..3fd1d356 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={"project_id": "p"}) + record = await _deploy(env, ttl_seconds=60, cpu=2.0, priority=1, gateway_mode=GatewayMode.CONSISTENT, attribution={"team": "t"}) [(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": {"project_id": "p"}} + "priority": 1, "env_state_type": None, "env_state_instance_id": None, "attribution": {"team": "t"}} 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={"project_id": "p"}) - assert SEEN["deploys"] == [("SANDBOXES", {"ttl_seconds": 60, "attribution": {"project_id": "p"}})] + await _deploy(env, ttl_seconds=60, attribution={"team": "t"}) + assert SEEN["deploys"] == [("SANDBOXES", {"ttl_seconds": 60, "attribution": {"team": "t"}})] @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 880abe09..93a584c3 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", [({"project_id": "p1"}, {"project_id": "p1"}), (None, {})], ids=["set", "unset"]) +@pytest.mark.parametrize("attribution, forwarded", [({"team": "t1"}, {"team": "t1"}), (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 de14edf1..2795a646 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={"project_id": "0123456789abcdef01234567"}) + run = await _run_server_deploy(cpu=2.0, memory_mb=4096, priority=3, attribution={"team": "t1"}) [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"] == {"project_id": "0123456789abcdef01234567"} and "i6pn" not in call + assert call["attribution"] == {"team": "t1"} 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 0ac55a2d..0af60203 100644 --- a/tst/unit/providers/sandbox_providers/attribution_wire_test.py +++ b/tst/unit/providers/sandbox_providers/attribution_wire_test.py @@ -1,27 +1,25 @@ """Byte-identity net for the emitted compute-attribution wire. -Attribution lands in the Modal App `tags` dict and the Modal App *name* (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. +Attribution is an open dict: every key lands on the Modal sandbox `tags`, and the shared Modal +App's `tags` carry only the configured defaults (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. Modal accepts an empty project and `_app_name_for_project` -collapses every project into one App, so that case is pinned. +empty string is a value. """ from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch import pytest -from agent_env.providers.sandbox_providers.modal_sandbox import ( - DEFAULT_APP_NAME, - _app_name_for_project, - _build_cost_attribution_tags, -) +from agent_env.attribution import RUN_ID_KEY +from agent_env.config import get_config +from agent_env.providers.sandbox_providers.modal_sandbox import ModalSandboxProvider, _attribution_tags +from tst.unit.store.fakes import FakeImageStore -_OBJECT_ID = "0123456789abcdef01234567" -_DEFAULT_PROJECT_ID = "fedcba9876543210fedcba98" -_DEFAULTS = {"product": "p0", "customer": "c0", "team": "t0", "project_id": _DEFAULT_PROJECT_ID} +_DEFAULTS = {"org": "o0", "team": "t0"} def _use_config(tmp_path, monkeypatch, attribution: dict[str, str]) -> None: @@ -37,66 +35,91 @@ def _configured_defaults(tmp_path, monkeypatch): def _tags(**kwargs): - return _build_cost_attribution_tags(kwargs) + return _attribution_tags(kwargs) # --- the emitted dicts, exactly ------------------------------------------------- def test_modal_tags_with_nothing_supplied(): - assert _tags() == { - "product": "p0", - "customer": "c0", - "team": "t0", - "project_id": _DEFAULT_PROJECT_ID, - } + assert _tags() == {"org": "o0", "team": "t0"} -def test_modal_tags_with_everything_supplied(): - assert _tags(product="p", customer="c", team="t", project_id=_OBJECT_ID) == { - "product": "p", "customer": "c", "team": "t", "project_id": _OBJECT_ID, +def test_modal_tags_carry_any_key(): + assert _tags(team="t", cost_center="research", run_id="inst-1") == { + "org": "o0", "team": "t", "cost_center": "research", "run_id": "inst-1", } def test_modal_vm_shares_the_modal_sink(): 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._attribution_tags is _attribution_tags -def test_the_default_project_names_its_own_modal_app(): - app_name = _app_name_for_project(DEFAULT_APP_NAME, _tags().get("project_id")) - assert app_name == f"{DEFAULT_APP_NAME}-{_DEFAULT_PROJECT_ID}" - assert len(app_name) <= 64 # Modal's limit; longer names are silently truncated # --- precedence ----------------------------------------------------------------- -@pytest.mark.parametrize("field", ["product", "customer", "team", "project_id"]) +@pytest.mark.parametrize("field", list(_DEFAULTS)) def test_explicit_beats_the_configured_default(field): assert _tags(**{field: "explicit"})[field] == "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" + _use_config(tmp_path, monkeypatch, {**_DEFAULTS, "org": "env:ATTRIBUTION_TEST_ORG?o0"}) + monkeypatch.delenv("ATTRIBUTION_TEST_ORG", raising=False) + assert _tags()["org"] == "o0" + monkeypatch.setenv("ATTRIBUTION_TEST_ORG", "from-env") + assert _tags()["org"] == "from-env" def test_without_configured_defaults_modal_omits_unset_dimensions(tmp_path, monkeypatch): _use_config(tmp_path, monkeypatch, {}) assert _tags() == {} - assert _tags(team="t") == {"team": "t"} - assert _app_name_for_project(DEFAULT_APP_NAME, _tags().get("project_id")) == DEFAULT_APP_NAME + assert _tags(team="t", unset=None) == {"team": "t"} -@pytest.mark.parametrize("field", ["product", "customer", "team"]) -def test_an_empty_string_is_a_value_not_an_absence(field): +def test_an_empty_string_is_a_value_not_an_absence(): """`is not None`, not truthiness — pinned because flipping it is silent on Modal.""" - assert _tags(**{field: ""})[field] == "" + assert _tags(team="")["team"] == "" + + +@pytest.mark.parametrize("key", ["", "cost center", "team/sub", "k" * 64]) +def test_a_key_modal_cannot_take_is_refused_not_rewritten(key): + with pytest.raises(ValueError, match="can't be a Modal tag"): + _tags(**{key: "v"}) + + +# --- what reaches Modal ----------------------------------------------------------- + + +def _provider(): + get_config().set_image_store(FakeImageStore()) + provider = ModalSandboxProvider(app_name="agent-env-test") + provider._get_app = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + provider._get_client = AsyncMock(return_value=MagicMock()) # type: ignore[method-assign] + return provider + + +@pytest.mark.asyncio +async def test_the_shared_app_gets_the_defaults_and_the_sandbox_gets_the_run(): + provider = _provider() + attribution = {"team": "t", "cost_center": "research", RUN_ID_KEY: "inst-1"} + with patch("agent_env.providers.sandbox_providers.modal_sandbox.modal.Sandbox._experimental_create") as create: + create.aio = AsyncMock(side_effect=RuntimeError("stop")) + with pytest.raises(RuntimeError): + await provider.create_container(image_name="img:latest", port=8000, env={}, attribution=attribution) + assert provider._get_app.call_args.args == ("agent-env-test", _DEFAULTS) + assert create.aio.call_args.kwargs["tags"] == {**_DEFAULTS, **attribution} + + +@pytest.mark.asyncio +async def test_a_bad_key_fails_before_the_app_is_looked_up(): + provider = _provider() + with pytest.raises(ValueError, match="'cost center'"): + await provider.create_container(image_name="img:latest", port=8000, env={}, attribution={"cost center": "x"}) + provider._get_app.assert_not_awaited() # --- the dict is the only carried form ------------------------------------------- @@ -118,7 +141,7 @@ async def create_sandbox(self, *, attribution=None, **kwargs): seen["kwargs"] = kwargs return SimpleNamespace(type="modern", network_policy=None) - attribution = {"project_id": _OBJECT_ID, "cost_center": "x"} + attribution = {"team": "t"} await ChainedSandboxProvider([Provider()]).create_sandbox(attribution=attribution, cpu=2.0) assert seen == {"attribution": attribution, "kwargs": {"cpu": 2.0}} @@ -128,5 +151,5 @@ async def test_the_quartet_is_no_longer_a_keyword_on_any_backend(): """A caller still on the pre-dict signature fails loud rather than billing unattributed.""" from agent_env.providers.sandbox_providers.local_sandbox import LocalSandboxProvider - with pytest.raises(TypeError, match="project_id"): - await LocalSandboxProvider().create_vm(project_id=_OBJECT_ID) + with pytest.raises(TypeError, match="team"): + await LocalSandboxProvider().create_vm(team="t") 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 d6569eea..7ec2b44d 100644 --- a/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py +++ b/tst/unit/providers/sandbox_providers/e2b_sandbox_provider_test.py @@ -60,12 +60,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", - "project_id": "project-d", - }, + attribution={"team": "team-c", "cost_center": "research", "run_id": "inst-1"}, ) assert isinstance(sandbox, E2BSandbox) @@ -76,10 +71,9 @@ async def test_create_vm_uses_derived_template_and_preserves_attribution( timeout=123, api_key="e2b-secret", metadata={ - "product": "product-a", - "customer": "customer-b", "team": "team-c", - "project_id": "project-d", + "cost_center": "research", + "run_id": "inst-1", "agent_env_exposed_ports": "8080,9000", }, network={"allow_public_traffic": True}, @@ -252,6 +246,24 @@ async def test_rejects_image_override_before_resolving_or_creating(): _AsyncSandboxSdk.create.assert_not_awaited() +@pytest.mark.asyncio +async def test_rejects_the_exposed_ports_key_as_attribution_before_resolving_or_creating(): + """The provider reads the key back as the sandbox's exposed ports on reconnect.""" + resolver = _Resolver() + provider = E2BSandboxProvider( + api_key="e2b-secret", + base_template="agent-env-v1", + template_resolver=resolver, + sandbox_cls=_AsyncSandboxSdk, + ) + + with pytest.raises(ValueError, match="'agent_env_exposed_ports' is reserved"): + await provider.create_vm(attribution={"agent_env_exposed_ports": "22"}) + + resolver.resolve.assert_not_awaited() + _AsyncSandboxSdk.create.assert_not_awaited() + + @pytest.mark.asyncio async def test_forwards_network_allowlist_to_e2b(sandbox_setup): provider = E2BSandboxProvider( 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..41494a23 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 @@ -44,6 +44,7 @@ async def test_logs_the_container_id_with_the_sandbox_and_its_tags(caplog): assert record.modal_container_id == "ta-01ABC" assert record.modal_sandbox_id == "sb-01XYZ" assert record.modal_app_name == "agent-env-p" + assert record.modal_sandbox_tags == _TAGS assert getattr(record, PIPELINE_STEP_KEY) == "my-pipeline_deploy_env" assert getattr(record, RUN_ID_KEY) == "inst-1" assert (record.cpu, record.memory_mb, record.gpu) == (2.0, 4096, "H100") @@ -59,6 +60,17 @@ async def test_an_unreadable_container_id_still_logs_and_never_raises(caplog): assert _started_record(caplog).modal_container_id is None +@pytest.mark.asyncio +async def test_a_tag_named_like_a_log_record_field_does_not_clobber_it(caplog): + tags = {"name": "n", "module": "m", RUN_ID_KEY: "inst-1"} + 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 record.name == "agent_env.providers.sandbox_providers.modal_sandbox" + assert getattr(record, RUN_ID_KEY) == "inst-1" + + @pytest.fixture(autouse=True) def _local_image_store(): from agent_env.config import get_config 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 16a6cb1b..2eca1e74 100644 --- a/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py +++ b/tst/unit/providers/sandbox_providers/modal_sandbox_tags_test.py @@ -8,11 +8,7 @@ import pytest 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_sandbox import ModalSandboxProvider, _attribution_tags from agent_env.providers.sandbox_providers.modal_vm_sandbox import ModalVmSandboxProvider @@ -30,9 +26,9 @@ def _context(instance_id=None, **metadata): def test_pipeline_step_is_task_id_and_step_id(): attribution = deploy_attribution( - _step("deploy_env", {"project_id": "p"}), _context(task_id="my-pipeline"), + _step("deploy_env", {"team": "t"}), _context(task_id="my-pipeline"), ) - assert attribution == {"project_id": "p", PIPELINE_STEP_KEY: "my-pipeline_deploy_env"} + assert attribution == {"team": "t", PIPELINE_STEP_KEY: "my-pipeline_deploy_env"} def test_no_task_id_means_no_pipeline_step(): @@ -59,46 +55,37 @@ def test_a_task_authored_pipeline_step_wins(): def test_deploy_attribution_does_not_mutate_the_step(): - step = _step(attribution={"project_id": "p"}) + step = _step(attribution={"team": "t"}) deploy_attribution(step, _context(task_id="t")) - assert step.metadata["attribution"] == {"project_id": "p"} - - -# --- _build_sandbox_tags -------------------------------------------------------- + assert step.metadata["attribution"] == {"team": "t"} -def test_sandbox_tags_empty_without_pipeline_step(): - assert _build_sandbox_tags({"project_id": "p"}) == {} +# --- _attribution_tags ---------------------------------------------------------- -def test_sandbox_tags_carry_both_keys(): - assert _build_sandbox_tags({PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "project_id": "p"}) == { - PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", +def test_sandbox_tags_carry_every_key(): + assert _attribution_tags({PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "team": "t"}) == { + PIPELINE_STEP_KEY: "t_s", RUN_ID_KEY: "inst-1", "team": "t", } def test_sandbox_tags_sanitize_invalid_characters(): - assert _build_sandbox_tags({PIPELINE_STEP_KEY: "my task/v2_deploy env"}) == { + assert _attribution_tags({PIPELINE_STEP_KEY: "my task/v2_deploy env"}) == { PIPELINE_STEP_KEY: "my-task-v2_deploy-env", } def test_sandbox_tags_fit_modal_limit_and_keep_steps_distinct(): long_task = "t" * 80 - a = _build_sandbox_tags({PIPELINE_STEP_KEY: f"{long_task}_deploy_env"})[PIPELINE_STEP_KEY] - b = _build_sandbox_tags({PIPELINE_STEP_KEY: f"{long_task}_deploy_agent"})[PIPELINE_STEP_KEY] + a = _attribution_tags({PIPELINE_STEP_KEY: f"{long_task}_deploy_env"})[PIPELINE_STEP_KEY] + b = _attribution_tags({PIPELINE_STEP_KEY: f"{long_task}_deploy_agent"})[PIPELINE_STEP_KEY] assert len(a) == len(b) == 63 assert a != b assert a.startswith("t" * 54) 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 + assert _attribution_tags({PIPELINE_STEP_KEY: "x" * 63}) == {PIPELINE_STEP_KEY: "x" * 63} # --- providers pass the tag to Modal -------------------------------------------- diff --git a/tst/unit/task/task_document_test.py b/tst/unit/task/task_document_test.py new file mode 100644 index 00000000..7272dcd4 --- /dev/null +++ b/tst/unit/task/task_document_test.py @@ -0,0 +1,9 @@ +from agent_env.task import Task + + +def test_a_stored_task_document_with_an_unread_project_id_still_loads(): + """Stored task documents may carry keys Task no longer reads, such as project_id.""" + task = Task.from_dict({"id": "t", "type": "task", "version": 3, "steps": [], "project_id": "p"}) + + assert (task.id, task.version, task.steps) == ("t", 3, []) + assert "project_id" not in task.to_dict() diff --git a/tst/unit/task_step/verifiers/test_rubrics_verifier_screenshots.py b/tst/unit/task_step/verifiers/test_rubrics_verifier_screenshots.py index fdc93833..be2734da 100644 --- a/tst/unit/task_step/verifiers/test_rubrics_verifier_screenshots.py +++ b/tst/unit/task_step/verifiers/test_rubrics_verifier_screenshots.py @@ -401,7 +401,6 @@ async def test_prompt_llm_judge_base_url_precedence(monkeypatch): import litellm import agent_env.config as cfg_mod - import agent_env.utils.litellm_attribution as attr_mod from agent_env.store import LocalSecretStore @@ -409,7 +408,6 @@ async def test_prompt_llm_judge_base_url_precedence(monkeypatch): real_cfg.set_secret_store(LocalSecretStore(values={"litellm_api_key": "k"}, use_env=False)) monkeypatch.setattr(real_cfg, "_model_cfg", ModelConfig()) monkeypatch.setattr(cfg_mod, "get_config", lambda: real_cfg) - monkeypatch.setattr(attr_mod, "build_litellm_cost_attribution_kwargs", lambda md: {}) monkeypatch.delenv("LITELLM_BASE_URL", raising=False) monkeypatch.delenv("LITELLM_API_KEY", raising=False)