agentenv-framework 0.9.1254__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- agent_env/__init__.py +0 -0
- agent_env/a2a_agent/__init__.py +23 -0
- agent_env/a2a_agent/a2a_agent.py +600 -0
- agent_env/a2a_agent/conversation_store.py +356 -0
- agent_env/a2a_agent/object_transfer.py +609 -0
- agent_env/a2a_agent/protocol.py +238 -0
- agent_env/a2a_agent/store.py +201 -0
- agent_env/a2a_agent/validator.py +901 -0
- agent_env/artifact/__init__.py +48 -0
- agent_env/artifact/artifact.py +74 -0
- agent_env/artifact/artifacts/__init__.py +0 -0
- agent_env/artifact/artifacts/cli.py +83 -0
- agent_env/artifact/artifacts/docker_image.py +457 -0
- agent_env/artifact/artifacts/environment.py +89 -0
- agent_env/artifact/artifacts/environment_universe.py +156 -0
- agent_env/artifact/artifacts/file.py +214 -0
- agent_env/artifact/artifacts/file_artifact_universe.py +206 -0
- agent_env/artifact/artifacts/skill.py +206 -0
- agent_env/artifact/artifacts/vm_image.py +95 -0
- agent_env/artifact/ref.py +15 -0
- agent_env/artifact/registry.py +251 -0
- agent_env/artifact/store.py +295 -0
- agent_env/artifact/universe.py +38 -0
- agent_env/attribution.py +41 -0
- agent_env/bundle/__init__.py +4 -0
- agent_env/bundle/_fs.py +128 -0
- agent_env/bundle/authoring.py +239 -0
- agent_env/bundle/installed.py +175 -0
- agent_env/bundle/ledger.py +269 -0
- agent_env/bundle/materialize.py +204 -0
- agent_env/bundle/parse.py +594 -0
- agent_env/bundle/plan.py +414 -0
- agent_env/bundle/resolve.py +404 -0
- agent_env/bundle/run.py +361 -0
- agent_env/cli/__init__.py +69 -0
- agent_env/cli/__main__.py +10 -0
- agent_env/cli/_installers.py +576 -0
- agent_env/cli/_plugin_changes.py +947 -0
- agent_env/cli/a2a_agent/__init__.py +22 -0
- agent_env/cli/a2a_agent/add_skill.py +97 -0
- agent_env/cli/a2a_agent/deploy.py +60 -0
- agent_env/cli/a2a_agent/get.py +15 -0
- agent_env/cli/a2a_agent/get_instance.py +26 -0
- agent_env/cli/a2a_agent/put.py +78 -0
- agent_env/cli/a2a_agent/validate.py +56 -0
- agent_env/cli/artifact/__init__.py +20 -0
- agent_env/cli/artifact/cli.py +83 -0
- agent_env/cli/artifact/environment.py +44 -0
- agent_env/cli/artifact/environment_universe.py +163 -0
- agent_env/cli/artifact/file_artifact_universe.py +281 -0
- agent_env/cli/artifact/skill.py +128 -0
- agent_env/cli/banner.py +25 -0
- agent_env/cli/config.py +263 -0
- agent_env/cli/env/__init__.py +31 -0
- agent_env/cli/env/deploy.py +112 -0
- agent_env/cli/env/gateway.py +57 -0
- agent_env/cli/env/get_instance.py +40 -0
- agent_env/cli/env/mcp_server.py +410 -0
- agent_env/cli/env/multi.py +282 -0
- agent_env/cli/env/service_db.py +99 -0
- agent_env/cli/env/snapshot.py +35 -0
- agent_env/cli/env/state/__init__.py +6 -0
- agent_env/cli/env/state/init.py +108 -0
- agent_env/cli/env/state/teardown.py +44 -0
- agent_env/cli/env/website.py +216 -0
- agent_env/cli/env/website_browser.py +67 -0
- agent_env/cli/eval/__init__.py +16 -0
- agent_env/cli/eval/add_tasks.py +63 -0
- agent_env/cli/eval/create.py +49 -0
- agent_env/cli/eval/run.py +194 -0
- agent_env/cli/identity/__init__.py +5 -0
- agent_env/cli/identity/client_identity.py +37 -0
- agent_env/cli/plugin.py +438 -0
- agent_env/cli/run.py +285 -0
- agent_env/cli/task/__init__.py +21 -0
- agent_env/cli/task/create.py +109 -0
- agent_env/cli/task/get.py +17 -0
- agent_env/cli/task/get_instance.py +34 -0
- agent_env/cli/task/run.py +615 -0
- agent_env/cli/task/validate.py +31 -0
- agent_env/cli/up.py +123 -0
- agent_env/cli/utils.py +183 -0
- agent_env/config/__init__.py +50 -0
- agent_env/config/describe.py +1067 -0
- agent_env/config/errors.py +13 -0
- agent_env/config/loader.py +245 -0
- agent_env/config/model.py +154 -0
- agent_env/config/paths.py +25 -0
- agent_env/config/plugin_tables.py +101 -0
- agent_env/config/provenance.py +134 -0
- agent_env/config/runtime.py +1064 -0
- agent_env/config/snapshot.py +81 -0
- agent_env/entity_refs.py +186 -0
- agent_env/env/__init__.py +49 -0
- agent_env/env/env.py +433 -0
- agent_env/env/env_artifact_store.py +87 -0
- agent_env/env/envs/__init__.py +7 -0
- agent_env/env/envs/_deployment.py +281 -0
- agent_env/env/envs/gateway_server.py +32 -0
- agent_env/env/envs/mcp_server.py +387 -0
- agent_env/env/envs/multi_env.py +780 -0
- agent_env/env/envs/service_db/Dockerfile +20 -0
- agent_env/env/envs/service_db/Dockerfile.db-mcp +42 -0
- agent_env/env/envs/service_db/Dockerfile.db-web +8 -0
- agent_env/env/envs/service_db.py +92 -0
- agent_env/env/envs/website.py +294 -0
- agent_env/env/envs/website_browser/Dockerfile +40 -0
- agent_env/env/envs/website_browser/__init__.py +17 -0
- agent_env/env/envs/website_browser/entrypoint.sh +13 -0
- agent_env/env/gateway/Dockerfile +14 -0
- agent_env/env/gateway/__init__.py +37 -0
- agent_env/env/gateway/clock.py +184 -0
- agent_env/env/gateway/constants.py +190 -0
- agent_env/env/gateway/entrypoint.py +156 -0
- agent_env/env/gateway/gateway.py +1095 -0
- agent_env/env/gateway/get_time.py +60 -0
- agent_env/env/gateway/requirements.txt +4 -0
- agent_env/env/gateway/triggers.py +1193 -0
- agent_env/env/legacy_protocol.py +146 -0
- agent_env/env/registry.py +93 -0
- agent_env/env/snapshot_store.py +372 -0
- agent_env/env/store.py +411 -0
- agent_env/eval/__init__.py +14 -0
- agent_env/eval/eval.py +68 -0
- agent_env/eval/store.py +116 -0
- agent_env/examples/hello/README.md +9 -0
- agent_env/examples/hello/artifacts/greeting/check.sh +1 -0
- agent_env/examples/hello/artifacts/greeting/hello.txt +1 -0
- agent_env/examples/hello/tasks/hello.json +10 -0
- agent_env/explorer/__init__.py +12 -0
- agent_env/explorer/app.py +277 -0
- agent_env/explorer/openapi_docs.py +242 -0
- agent_env/explorer/plugin.py +95 -0
- agent_env/explorer/routers/__init__.py +1 -0
- agent_env/explorer/routers/common.py +174 -0
- agent_env/explorer/routers/conversations.py +25 -0
- agent_env/explorer/routers/objects.py +112 -0
- agent_env/explorer/routers/runs.py +475 -0
- agent_env/plugins/__init__.py +59 -0
- agent_env/plugins/_cli.py +244 -0
- agent_env/plugins/_discovery.py +97 -0
- agent_env/plugins/_inventory.py +262 -0
- agent_env/plugins/_registration.py +268 -0
- agent_env/plugins/_report.py +47 -0
- agent_env/plugins/_requirements.py +55 -0
- agent_env/providers/__init__.py +37 -0
- agent_env/providers/env_providers/__init__.py +12 -0
- agent_env/providers/env_providers/constants.py +13 -0
- agent_env/providers/env_providers/env_gateway_provider.py +1354 -0
- agent_env/providers/env_providers/env_provider.py +237 -0
- agent_env/providers/env_providers/env_server_provider.py +133 -0
- agent_env/providers/env_state/__init__.py +49 -0
- agent_env/providers/env_state/env_state_provider.py +554 -0
- agent_env/providers/env_state/local_postgres.py +390 -0
- agent_env/providers/env_state/store.py +112 -0
- agent_env/providers/sandbox_providers/__init__.py +29 -0
- agent_env/providers/sandbox_providers/chained_sandbox_provider.py +93 -0
- agent_env/providers/sandbox_providers/e2b/__init__.py +6 -0
- agent_env/providers/sandbox_providers/e2b/provider.py +376 -0
- agent_env/providers/sandbox_providers/e2b/sandbox.py +342 -0
- agent_env/providers/sandbox_providers/e2b/template.py +158 -0
- agent_env/providers/sandbox_providers/local_sandbox.py +408 -0
- agent_env/providers/sandbox_providers/modal_sandbox.py +500 -0
- agent_env/providers/sandbox_providers/modal_vm_sandbox.py +466 -0
- agent_env/providers/sandbox_providers/sandbox.py +387 -0
- agent_env/providers/sandbox_providers/sandbox_provider.py +614 -0
- agent_env/py.typed +0 -0
- agent_env/runner/__init__.py +16 -0
- agent_env/runner/local_runner.py +190 -0
- agent_env/runner/runner.py +127 -0
- agent_env/runner/store.py +94 -0
- agent_env/store/__init__.py +126 -0
- agent_env/store/_google.py +68 -0
- agent_env/store/base.py +33 -0
- agent_env/store/document_store/__init__.py +60 -0
- agent_env/store/document_store/document_store.py +519 -0
- agent_env/store/document_store/dynamodb_document_store.py +306 -0
- agent_env/store/document_store/evaluation.py +199 -0
- agent_env/store/document_store/firestore_mongo_document_store.py +242 -0
- agent_env/store/document_store/mongo_document_store.py +259 -0
- agent_env/store/document_store/sqlite_document_store.py +314 -0
- agent_env/store/ids.py +123 -0
- agent_env/store/image_store/__init__.py +34 -0
- agent_env/store/image_store/ecr_image_store.py +131 -0
- agent_env/store/image_store/google_credentials.py +113 -0
- agent_env/store/image_store/image_store.py +113 -0
- agent_env/store/image_store/local_registry_image_store.py +106 -0
- agent_env/store/image_store/oci_registry_credentials.py +196 -0
- agent_env/store/local_state.py +27 -0
- agent_env/store/object_store/__init__.py +24 -0
- agent_env/store/object_store/gcs_object_store.py +415 -0
- agent_env/store/object_store/local_object_store.py +143 -0
- agent_env/store/object_store/object_store.py +198 -0
- agent_env/store/object_store/s3_object_store.py +354 -0
- agent_env/store/query.py +112 -0
- agent_env/store/routing.py +609 -0
- agent_env/store/secret_store/__init__.py +19 -0
- agent_env/store/secret_store/aws_secrets_manager_secret_store.py +293 -0
- agent_env/store/secret_store/gcp_secret_manager_secret_store.py +319 -0
- agent_env/store/secret_store/local_secret_store.py +47 -0
- agent_env/store/secret_store/secret_store.py +19 -0
- agent_env/task/__init__.py +27 -0
- agent_env/task/interrupts.py +109 -0
- agent_env/task/registry.py +27 -0
- agent_env/task/step_journal.py +99 -0
- agent_env/task/store.py +940 -0
- agent_env/task/task.py +815 -0
- agent_env/task/teardown.py +142 -0
- agent_env/task_step/__init__.py +47 -0
- agent_env/task_step/context.py +195 -0
- agent_env/task_step/context_ops.py +224 -0
- agent_env/task_step/registry.py +182 -0
- agent_env/task_step/review_store.py +118 -0
- agent_env/task_step/snapshot_utils/__init__.py +2 -0
- agent_env/task_step/snapshot_utils/agent_state_capture.py +261 -0
- agent_env/task_step/snapshot_utils/snapshot_series.py +631 -0
- agent_env/task_step/store.py +129 -0
- agent_env/task_step/task_step.py +173 -0
- agent_env/task_step/task_steps/__init__.py +27 -0
- agent_env/task_step/task_steps/a2a_agent_validator/__init__.py +31 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/README.md +158 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/clip.m4a +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/clip.mp3 +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/clip.mp4 +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/clip.ogg +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/clip.wav +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/document.pdf +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/red.gif +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/red.jpg +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/fixtures/red.png +0 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_agent_card.py +103 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_agent_config_identity.py +69 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_agent_mcp.py +178 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_core_protocol.py +67 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_install.py +115 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution.py +124 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_litellm_attribution_runtime.py +173 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_modalities.py +270 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_peer_agents.py +102 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_role.py +151 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_skill_config.py +194 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_snapshot.py +111 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_system_prompt.py +215 -0
- agent_env/task_step/task_steps/a2a_agent_validator/verify_a2a_trajectory.py +164 -0
- agent_env/task_step/task_steps/add_skills.py +240 -0
- agent_env/task_step/task_steps/apply_server_config.py +288 -0
- agent_env/task_step/task_steps/collect_artifacts.py +940 -0
- agent_env/task_step/task_steps/deploy_agent.py +626 -0
- agent_env/task_step/task_steps/deploy_env.py +197 -0
- agent_env/task_step/task_steps/deploy_human_agent.py +100 -0
- agent_env/task_step/task_steps/deploy_sandbox.py +170 -0
- agent_env/task_step/task_steps/env_card_validator/__init__.py +4 -0
- agent_env/task_step/task_steps/env_card_validator/verify_env_card.py +109 -0
- agent_env/task_step/task_steps/env_card_validator/verify_env_core_protocol.py +108 -0
- agent_env/task_step/task_steps/install_agent.py +376 -0
- agent_env/task_step/task_steps/load_artifact.py +775 -0
- agent_env/task_step/task_steps/mcp_cli_builder/__init__.py +4 -0
- agent_env/task_step/task_steps/mcp_cli_builder/build_mcp_cli.py +201 -0
- agent_env/task_step/task_steps/mcp_cli_builder/codegen.py +627 -0
- agent_env/task_step/task_steps/mcp_env_validator/__init__.py +44 -0
- agent_env/task_step/task_steps/mcp_env_validator/validation_gate_aggregator.py +102 -0
- agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_env_assessment.py +78 -0
- agent_env/task_step/task_steps/mcp_env_validator/verify_mcp_tool_schema.py +152 -0
- agent_env/task_step/task_steps/mcp_env_validator/verify_spec_conformance.py +486 -0
- agent_env/task_step/task_steps/modify_env_tool_access.py +80 -0
- agent_env/task_step/task_steps/multienv_validator/__init__.py +3 -0
- agent_env/task_step/task_steps/multienv_validator/combine_universe_verdicts.py +107 -0
- agent_env/task_step/task_steps/multienv_validator/universe_comparison.py +253 -0
- agent_env/task_step/task_steps/multienv_validator/verify_universe_agent_judge.py +180 -0
- agent_env/task_step/task_steps/multienv_validator/verify_universe_roundtrip.py +241 -0
- agent_env/task_step/task_steps/peer_agents.py +92 -0
- agent_env/task_step/task_steps/prompt_agent.py +900 -0
- agent_env/task_step/task_steps/register_agent_triggers.py +113 -0
- agent_env/task_step/task_steps/register_env_triggers.py +97 -0
- agent_env/task_step/task_steps/reset_env.py +71 -0
- agent_env/task_step/task_steps/review.py +152 -0
- agent_env/task_step/task_steps/run_code.py +461 -0
- agent_env/task_step/task_steps/run_code_runner.py +52 -0
- agent_env/task_step/task_steps/run_docker_container.py +380 -0
- agent_env/task_step/task_steps/sandbox_utils/__init__.py +0 -0
- agent_env/task_step/task_steps/sandbox_utils/sandbox_utils.py +79 -0
- agent_env/task_step/task_steps/snapshot_agent_state.py +205 -0
- agent_env/task_step/task_steps/snapshot_env.py +637 -0
- agent_env/task_step/task_steps/sync_env_clock.py +132 -0
- agent_env/task_step/task_steps/teardown_sandboxes.py +170 -0
- agent_env/task_step/task_steps/verifiers/__init__.py +23 -0
- agent_env/task_step/task_steps/verifiers/agent_prompt_response_verifier.py +166 -0
- agent_env/task_step/task_steps/verifiers/aggregate_verifiers.py +90 -0
- agent_env/task_step/task_steps/verifiers/env_outcome_verifier.py +127 -0
- agent_env/task_step/task_steps/verifiers/judge_utils/__init__.py +2 -0
- agent_env/task_step/task_steps/verifiers/judge_utils/frame_selection.py +377 -0
- agent_env/task_step/task_steps/verifiers/judge_utils/judge_output_format.py +1293 -0
- agent_env/task_step/task_steps/verifiers/judge_utils/trajectory_filter.py +514 -0
- agent_env/task_step/task_steps/verifiers/rubrics_verifier.py +1318 -0
- agent_env/task_step/task_steps/verifiers/run_container_unit_tests_verifier.py +434 -0
- agent_env/task_step/task_steps/verifiers/scoring.py +40 -0
- agent_env/task_step/task_steps/verifiers/verify_sandbox.py +308 -0
- agent_env/task_step/thread_work.py +61 -0
- agent_env/utils/__init__.py +0 -0
- agent_env/utils/card_naming.py +108 -0
- agent_env/utils/deprecation.py +109 -0
- agent_env/utils/docker_build.py +29 -0
- agent_env/utils/exec_retry.py +77 -0
- agent_env/utils/litellm_attribution.py +59 -0
- agent_env/utils/paths.py +20 -0
- agentenv_framework-0.9.1254.dist-info/METADATA +394 -0
- agentenv_framework-0.9.1254.dist-info/RECORD +312 -0
- agentenv_framework-0.9.1254.dist-info/WHEEL +4 -0
- agentenv_framework-0.9.1254.dist-info/entry_points.txt +5 -0
- agentenv_framework-0.9.1254.dist-info/licenses/LICENSE +202 -0
- agentenv_framework-0.9.1254.dist-info/licenses/NOTICE +4 -0
- agentenv_framework-0.9.1254.dist-info/licenses/THIRD_PARTY_NOTICES.md +7665 -0
|
@@ -0,0 +1,631 @@
|
|
|
1
|
+
"""The capture series: freezes agent state mid-run, on a cadence or per turn.
|
|
2
|
+
|
|
3
|
+
``prompt_agent`` optionally drives ``SnapshotSeries`` *while the prompt is still in
|
|
4
|
+
flight* — on a wall-clock cadence (``interval_seconds``), at every conversation-turn
|
|
5
|
+
boundary (``per_turn``), or both — so a run yields a curve of gradable points instead
|
|
6
|
+
of a single endpoint. Nothing here grades: each capture appends one row to
|
|
7
|
+
``context.metadata['agent_snapshots']``, and a row is gradable when it has a
|
|
8
|
+
``bundle_object_url`` (``capture_status`` is provenance, not a gate).
|
|
9
|
+
"""
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import asyncio
|
|
13
|
+
import logging
|
|
14
|
+
import time
|
|
15
|
+
import uuid
|
|
16
|
+
from dataclasses import asdict, dataclass
|
|
17
|
+
from datetime import datetime, timezone
|
|
18
|
+
from typing import Optional
|
|
19
|
+
|
|
20
|
+
from agent_env.task_step.context import TaskStepContext
|
|
21
|
+
from agent_env.task_step.snapshot_utils import agent_state_capture as capture
|
|
22
|
+
|
|
23
|
+
logger = logging.getLogger(__name__)
|
|
24
|
+
|
|
25
|
+
# How much longer than a capture's own budget teardown waits for the lock, so the
|
|
26
|
+
# two deadlines cannot land in the same loop iteration. See `_acquire_for_final`.
|
|
27
|
+
_FINAL_LOCK_GRACE_SECONDS = 15
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _is_cancelling() -> bool:
|
|
31
|
+
task = asyncio.current_task()
|
|
32
|
+
return bool(task.cancelling()) if task is not None else False
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@dataclass
|
|
36
|
+
class SnapshotConfig:
|
|
37
|
+
"""Periodic capture settings for a ``prompt_agent`` step.
|
|
38
|
+
|
|
39
|
+
Its presence on the step is the on/off switch, so there is no half-configured
|
|
40
|
+
state: every field defaults usefully or is rejected in ``__post_init__``.
|
|
41
|
+
"""
|
|
42
|
+
|
|
43
|
+
interval_seconds: Optional[int] = None
|
|
44
|
+
at_end: bool = True
|
|
45
|
+
# One capture per AGENT turn, taken after the agent replies and before the user
|
|
46
|
+
# sim's next message is sent. Independent of the ticker; both may be on.
|
|
47
|
+
per_turn: bool = False
|
|
48
|
+
env_id: Optional[str] = None
|
|
49
|
+
timeout_seconds: int = 300
|
|
50
|
+
# Interior captures only — the final one is exempt, so a series can yield
|
|
51
|
+
# `max_snapshots + 1` rows.
|
|
52
|
+
max_snapshots: int = 24
|
|
53
|
+
|
|
54
|
+
def __post_init__(self) -> None:
|
|
55
|
+
if self.timeout_seconds <= 0:
|
|
56
|
+
raise ValueError("snapshot timeout_seconds must be > 0")
|
|
57
|
+
if self.max_snapshots < 1:
|
|
58
|
+
raise ValueError("max_snapshots must be >= 1")
|
|
59
|
+
if not self.at_end and not self.ticks and not self.per_turn:
|
|
60
|
+
raise ValueError(
|
|
61
|
+
"snapshotting with at_end=False needs a positive interval_seconds "
|
|
62
|
+
f"(got {self.interval_seconds!r}) or per_turn=True; otherwise no "
|
|
63
|
+
"capture ever runs"
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
@property
|
|
67
|
+
def ticks(self) -> bool:
|
|
68
|
+
"""Non-positive starts no ticker: ``wait_for`` would fire at once."""
|
|
69
|
+
return self.interval_seconds is not None and self.interval_seconds > 0
|
|
70
|
+
|
|
71
|
+
def to_dict(self) -> dict:
|
|
72
|
+
return asdict(self)
|
|
73
|
+
|
|
74
|
+
@classmethod
|
|
75
|
+
def from_dict(cls, data: dict) -> "SnapshotConfig":
|
|
76
|
+
defaults = cls()
|
|
77
|
+
return cls(
|
|
78
|
+
interval_seconds=data.get("interval_seconds"),
|
|
79
|
+
at_end=data.get("at_end", defaults.at_end),
|
|
80
|
+
per_turn=data.get("per_turn", defaults.per_turn),
|
|
81
|
+
env_id=data.get("env_id"),
|
|
82
|
+
timeout_seconds=data.get("timeout_seconds", defaults.timeout_seconds),
|
|
83
|
+
max_snapshots=data.get("max_snapshots", defaults.max_snapshots),
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def _rollout_discriminator(step_id: str, instance_id: Optional[str]) -> str:
|
|
88
|
+
"""The per-rollout token every artifact id in one series is scoped by.
|
|
89
|
+
|
|
90
|
+
Instance-scoped, so concurrent rollouts never contend on one id and retries
|
|
91
|
+
reuse it. Deliberately not derived from the context id: a caller-supplied
|
|
92
|
+
``context_id`` is shared by every rollout, so ids built from it would collide.
|
|
93
|
+
"""
|
|
94
|
+
if instance_id:
|
|
95
|
+
return instance_id.rsplit("-", 1)[-1][:16]
|
|
96
|
+
generated = uuid.uuid4().hex[:8]
|
|
97
|
+
logger.warning(
|
|
98
|
+
"%s: no instance_id in context; using random artifact discriminator %s "
|
|
99
|
+
"(retries will not be idempotent)", step_id, generated,
|
|
100
|
+
)
|
|
101
|
+
return generated
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
class SnapshotSeries:
|
|
105
|
+
"""One rollout's capture series: a ticker, the in-flight lock, and teardown.
|
|
106
|
+
|
|
107
|
+
One instance per rollout, held as a local: k concurrent rollouts share one
|
|
108
|
+
``TaskStep``, so none of this state may live on the step.
|
|
109
|
+
|
|
110
|
+
The cadence is the gap BETWEEN captures — the ticker awaits its own, so captures
|
|
111
|
+
never overlap and the lock's only other contender is teardown.
|
|
112
|
+
"""
|
|
113
|
+
|
|
114
|
+
def __init__(
|
|
115
|
+
self,
|
|
116
|
+
*,
|
|
117
|
+
step_id: str,
|
|
118
|
+
agent_name: str,
|
|
119
|
+
prompt_id: str,
|
|
120
|
+
a2a_context_id: str,
|
|
121
|
+
config: SnapshotConfig,
|
|
122
|
+
trajectory_output_prefix: str,
|
|
123
|
+
instance_id: Optional[str] = None,
|
|
124
|
+
):
|
|
125
|
+
self.step_id = step_id
|
|
126
|
+
self.agent_name = agent_name
|
|
127
|
+
self.prompt_id = prompt_id
|
|
128
|
+
# The id that went on the wire (`prompt_agent`'s `solver_context_id`), never
|
|
129
|
+
# its conversation-doc id: every question a capture asks is agent-side, and
|
|
130
|
+
# that is the only name the sidecar knows the session by.
|
|
131
|
+
self.a2a_context_id = a2a_context_id
|
|
132
|
+
self.config = config
|
|
133
|
+
self.trajectory_output_prefix = trajectory_output_prefix
|
|
134
|
+
self._discriminator = _rollout_discriminator(step_id, instance_id)
|
|
135
|
+
self.workspace_artifact_id = f"{step_id}-workspace-{self._discriminator}"
|
|
136
|
+
self._lock = asyncio.Lock()
|
|
137
|
+
self._stop = asyncio.Event()
|
|
138
|
+
self._ticker: Optional[asyncio.Task] = None
|
|
139
|
+
self._started_monotonic = time.monotonic()
|
|
140
|
+
# Interior captures attempted — what `max_snapshots` bounds. Not a
|
|
141
|
+
# row ordinal: failed rows consume attempts too, and the final is exempt.
|
|
142
|
+
self._attempts = 0
|
|
143
|
+
# See `_refuse_turn`.
|
|
144
|
+
self._turn_limit_recorded = False
|
|
145
|
+
self._rows: list[dict] = []
|
|
146
|
+
|
|
147
|
+
# ------------------------------------------------------------- lifecycle
|
|
148
|
+
|
|
149
|
+
def start(self, context: TaskStepContext) -> None:
|
|
150
|
+
"""Start ticking, if a cadence was configured."""
|
|
151
|
+
if self.config.ticks:
|
|
152
|
+
self._ticker = asyncio.create_task(self._tick_loop(context))
|
|
153
|
+
|
|
154
|
+
async def finish(self, context: TaskStepContext) -> None:
|
|
155
|
+
"""Take the final capture, then publish the series.
|
|
156
|
+
|
|
157
|
+
Raises only ``CancelledError``, which must propagate — a caller that
|
|
158
|
+
swallowed it would complete a cancelled step silently.
|
|
159
|
+
"""
|
|
160
|
+
cancelling = _is_cancelling()
|
|
161
|
+
try:
|
|
162
|
+
self._stop.set()
|
|
163
|
+
if cancelling:
|
|
164
|
+
# Waiting for the lock and capturing are each bounded by
|
|
165
|
+
# `timeout_seconds`, so honouring them here would delay the
|
|
166
|
+
# cancellation by up to twice that (10 min at defaults).
|
|
167
|
+
self._record(status="skipped", reason="cancelled", is_final=True)
|
|
168
|
+
return
|
|
169
|
+
# Waited out even when capturing nothing here, so an in-flight interior
|
|
170
|
+
# capture's S3 object and burnt version end up referenced by a row.
|
|
171
|
+
acquired = await self._acquire_for_final()
|
|
172
|
+
try:
|
|
173
|
+
if not self.config.at_end:
|
|
174
|
+
return
|
|
175
|
+
if not acquired:
|
|
176
|
+
self._record(status="skipped", is_final=True,
|
|
177
|
+
reason="final_capture_lock_timeout")
|
|
178
|
+
return
|
|
179
|
+
self._append(await self._capture_never_raising(context, is_final=True))
|
|
180
|
+
finally:
|
|
181
|
+
if acquired:
|
|
182
|
+
self._lock.release()
|
|
183
|
+
except asyncio.CancelledError:
|
|
184
|
+
# Cancellation that arrived DURING the lock wait or the capture, so the
|
|
185
|
+
# flag sampled at entry is stale — the `finally` must still force the
|
|
186
|
+
# ticker down or a mid-capture one survives teardown.
|
|
187
|
+
cancelling = True
|
|
188
|
+
self._record(status="skipped", reason="cancelled", is_final=True)
|
|
189
|
+
raise
|
|
190
|
+
except Exception:
|
|
191
|
+
logger.exception("%s: final capture raised; see its manifest row", self.step_id)
|
|
192
|
+
finally:
|
|
193
|
+
try:
|
|
194
|
+
# Re-sampled as well: a cancellation can be pending without having
|
|
195
|
+
# been delivered at an await yet.
|
|
196
|
+
await self._retire_ticker(force=cancelling or _is_cancelling())
|
|
197
|
+
finally:
|
|
198
|
+
# Nested: retirement re-raises a cancellation aimed at us, and the
|
|
199
|
+
# rows live nowhere else. Safe here because `_publish` never awaits.
|
|
200
|
+
self._publish(context)
|
|
201
|
+
|
|
202
|
+
async def _acquire_for_final(self) -> bool:
|
|
203
|
+
"""Take the lock for the final capture, outlasting the capture it waits on.
|
|
204
|
+
|
|
205
|
+
Strictly LONGER than ``timeout_seconds``, never equal: an interior capture is
|
|
206
|
+
bounded by that same value, so equal bounds put both deadlines in one loop
|
|
207
|
+
iteration and this waiter times out before the interior capture can release —
|
|
208
|
+
recording a fatal ``skipped`` final row for a rollout that in fact completed.
|
|
209
|
+
"""
|
|
210
|
+
timeout = self.config.timeout_seconds + _FINAL_LOCK_GRACE_SECONDS
|
|
211
|
+
try:
|
|
212
|
+
await asyncio.wait_for(self._lock.acquire(), timeout=timeout)
|
|
213
|
+
return True
|
|
214
|
+
except asyncio.TimeoutError:
|
|
215
|
+
logger.warning(
|
|
216
|
+
"%s: capture for %s still in flight after %ss; closing the series",
|
|
217
|
+
self.step_id, self.a2a_context_id, timeout,
|
|
218
|
+
)
|
|
219
|
+
return False
|
|
220
|
+
|
|
221
|
+
async def _retire_ticker(self, *, force: bool = False) -> None:
|
|
222
|
+
"""Stop the ticker before ``execute`` returns.
|
|
223
|
+
|
|
224
|
+
Idle: cancel and await — ``cancel()`` only schedules, and an unawaited task
|
|
225
|
+
outlives the step holding this rollout's context.
|
|
226
|
+
|
|
227
|
+
Mid-capture: leave it unless ``force``. Cancelling stops our client but not
|
|
228
|
+
the sidecar, which keeps tarring and orphans its object; ``stop`` is set and
|
|
229
|
+
the capture bounded, so it ends on its own. One inside ``asyncio.to_thread``
|
|
230
|
+
cannot be stopped regardless.
|
|
231
|
+
"""
|
|
232
|
+
if self._ticker is None or self._ticker.done():
|
|
233
|
+
return
|
|
234
|
+
if self._lock.locked() and not force:
|
|
235
|
+
logger.warning(
|
|
236
|
+
"%s: ticker for %s is mid-capture; letting it finish rather than "
|
|
237
|
+
"orphaning a sidecar upload", self.step_id, self.a2a_context_id,
|
|
238
|
+
)
|
|
239
|
+
return
|
|
240
|
+
if self._lock.locked():
|
|
241
|
+
logger.warning(
|
|
242
|
+
"%s: cancelling a mid-capture ticker for %s under cancellation; its "
|
|
243
|
+
"sidecar upload may be orphaned", self.step_id, self.a2a_context_id,
|
|
244
|
+
)
|
|
245
|
+
# Swallow the ticker's own cancellation, propagate one aimed at US — else a
|
|
246
|
+
# cancelled step completes silently. `.cancelled()` reads True either way.
|
|
247
|
+
parent = asyncio.current_task()
|
|
248
|
+
cancels_before = parent.cancelling() if parent is not None else 0
|
|
249
|
+
self._ticker.cancel()
|
|
250
|
+
try:
|
|
251
|
+
await self._ticker
|
|
252
|
+
except asyncio.CancelledError:
|
|
253
|
+
cancels_after = parent.cancelling() if parent is not None else 0
|
|
254
|
+
if cancels_after > cancels_before:
|
|
255
|
+
raise
|
|
256
|
+
|
|
257
|
+
@property
|
|
258
|
+
def _at_limit(self) -> bool:
|
|
259
|
+
"""Whether the interior budget is spent. Rechecked under the lock by both
|
|
260
|
+
contenders: the increment is only safe unlocked because no await separates it
|
|
261
|
+
from `acquire()`, and one added there would let both spend the last slot."""
|
|
262
|
+
return self._attempts >= self.config.max_snapshots
|
|
263
|
+
|
|
264
|
+
def _refuse_turn(self, turn: int) -> None:
|
|
265
|
+
"""Note that this turn boundary was skipped because the budget is spent.
|
|
266
|
+
Recorded once, not once per remaining turn, so the row count stays bounded
|
|
267
|
+
no matter how long the conversation runs."""
|
|
268
|
+
if not self._turn_limit_recorded:
|
|
269
|
+
self._turn_limit_recorded = True
|
|
270
|
+
self._record(status="skipped", reason="limit_reached", turn=turn)
|
|
271
|
+
|
|
272
|
+
async def _tick_loop(self, context: TaskStepContext) -> None:
|
|
273
|
+
# Nothing awaits this task, so an escaping exception would vanish into
|
|
274
|
+
# asyncio's "never retrieved" path and snapshotting would just stop.
|
|
275
|
+
try:
|
|
276
|
+
await self._ticks(context)
|
|
277
|
+
except asyncio.CancelledError:
|
|
278
|
+
raise
|
|
279
|
+
except Exception as exc:
|
|
280
|
+
logger.exception(
|
|
281
|
+
"%s: snapshot ticker for %s stopped unexpectedly",
|
|
282
|
+
self.step_id, self.a2a_context_id,
|
|
283
|
+
)
|
|
284
|
+
self._record(
|
|
285
|
+
status="failed",
|
|
286
|
+
reason=f"ticker_error: {type(exc).__name__}: {exc}"[:300],
|
|
287
|
+
)
|
|
288
|
+
|
|
289
|
+
async def _ticks(self, context: TaskStepContext) -> None:
|
|
290
|
+
"""Capture every ``interval_seconds`` until teardown or the attempt limit."""
|
|
291
|
+
while True:
|
|
292
|
+
try:
|
|
293
|
+
# Not sleep(): teardown wakes the ticker immediately rather than
|
|
294
|
+
# sleeping out the remaining interval.
|
|
295
|
+
await asyncio.wait_for(
|
|
296
|
+
self._stop.wait(), timeout=self.config.interval_seconds
|
|
297
|
+
)
|
|
298
|
+
return
|
|
299
|
+
except asyncio.TimeoutError:
|
|
300
|
+
pass # interval elapsed → a capture is due
|
|
301
|
+
|
|
302
|
+
if self._at_limit:
|
|
303
|
+
# Stop ticking, but leave the final capture to teardown: counting
|
|
304
|
+
# it here would let ticks suppress the grade of record.
|
|
305
|
+
self._record(status="skipped", reason="limit_reached")
|
|
306
|
+
return
|
|
307
|
+
|
|
308
|
+
async with self._lock:
|
|
309
|
+
# Rechecked: teardown can set `stop` and take the lock while this
|
|
310
|
+
# tick waits on it, and a capture started afterwards would burn a
|
|
311
|
+
# universe version teardown has already stopped waiting for.
|
|
312
|
+
if self._stop.is_set():
|
|
313
|
+
return
|
|
314
|
+
# The budget is rechecked too — see `_at_limit`.
|
|
315
|
+
if self._at_limit:
|
|
316
|
+
self._record(status="skipped", reason="limit_reached")
|
|
317
|
+
return
|
|
318
|
+
self._attempts += 1
|
|
319
|
+
self._append(await self._capture_never_raising(context, is_final=False))
|
|
320
|
+
|
|
321
|
+
async def capture_turn(self, context: TaskStepContext, turn: int) -> None:
|
|
322
|
+
"""Capture at a conversation-turn boundary. No-op unless ``per_turn``.
|
|
323
|
+
|
|
324
|
+
Awaited by the turn loop, not run on the ticker's task: the agent has stopped
|
|
325
|
+
writing only until the user sim's reply starts it again, so the capture has to
|
|
326
|
+
finish before that reply goes out. Never raises except ``CancelledError``.
|
|
327
|
+
"""
|
|
328
|
+
if not self.config.per_turn or self._stop.is_set():
|
|
329
|
+
return
|
|
330
|
+
if self._at_limit:
|
|
331
|
+
self._refuse_turn(turn)
|
|
332
|
+
return
|
|
333
|
+
async with self._lock:
|
|
334
|
+
# Rechecked under the lock for the same reason `_ticks` does: teardown
|
|
335
|
+
# may have closed the series while this waited on an in-flight tick.
|
|
336
|
+
if self._stop.is_set():
|
|
337
|
+
return
|
|
338
|
+
if self._at_limit:
|
|
339
|
+
self._refuse_turn(turn)
|
|
340
|
+
return
|
|
341
|
+
self._attempts += 1
|
|
342
|
+
self._append(
|
|
343
|
+
await self._capture_never_raising(context, is_final=False, turn=turn)
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
# --------------------------------------------------------------- capture
|
|
347
|
+
|
|
348
|
+
async def _capture_never_raising(
|
|
349
|
+
self, context: TaskStepContext, *, is_final: bool, turn: Optional[int] = None
|
|
350
|
+
) -> dict:
|
|
351
|
+
"""Run one capture under a hard timeout. Returns a row; never raises
|
|
352
|
+
except ``CancelledError``.
|
|
353
|
+
|
|
354
|
+
A timeout abandons the capture, possibly orphaning an S3 object — safe,
|
|
355
|
+
since each capture writes to a unique prefix.
|
|
356
|
+
"""
|
|
357
|
+
# Owned here, not built inside `_capture`, so a failure part-way through
|
|
358
|
+
# still returns whatever landed before it.
|
|
359
|
+
row = self._row(is_final=is_final, status="ok", turn=turn)
|
|
360
|
+
try:
|
|
361
|
+
return await asyncio.wait_for(
|
|
362
|
+
self._capture(context, row, is_final=is_final),
|
|
363
|
+
timeout=self.config.timeout_seconds,
|
|
364
|
+
)
|
|
365
|
+
except asyncio.CancelledError:
|
|
366
|
+
raise
|
|
367
|
+
except asyncio.TimeoutError:
|
|
368
|
+
logger.warning(
|
|
369
|
+
"%s: capture timed out after %ss (is_final=%s turn=%s)",
|
|
370
|
+
self.step_id, self.config.timeout_seconds, is_final, turn,
|
|
371
|
+
)
|
|
372
|
+
return self._degrade(row, f"capture_timeout_{self.config.timeout_seconds}s")
|
|
373
|
+
except Exception as exc:
|
|
374
|
+
logger.warning(
|
|
375
|
+
"%s: capture failed (is_final=%s turn=%s, run continues): %s",
|
|
376
|
+
self.step_id, is_final, turn, exc, exc_info=True,
|
|
377
|
+
)
|
|
378
|
+
return self._degrade(row, f"{type(exc).__name__}: {exc}"[:300])
|
|
379
|
+
|
|
380
|
+
@staticmethod
|
|
381
|
+
def _degrade(row: dict, reason: str) -> dict:
|
|
382
|
+
"""``partial`` when the bundle landed (still gradable), ``failed`` when it
|
|
383
|
+
did not — which is fatal for the final capture."""
|
|
384
|
+
if row.get("bundle_object_url"):
|
|
385
|
+
row["capture_status"] = "partial"
|
|
386
|
+
row["capture_reason"] = "; ".join(
|
|
387
|
+
p for p in (row.get("capture_reason"), reason) if p
|
|
388
|
+
)
|
|
389
|
+
else:
|
|
390
|
+
row["capture_status"] = "failed"
|
|
391
|
+
row["capture_reason"] = reason
|
|
392
|
+
return row
|
|
393
|
+
|
|
394
|
+
async def _capture(
|
|
395
|
+
self, context: TaskStepContext, row: dict, *, is_final: bool
|
|
396
|
+
) -> dict:
|
|
397
|
+
"""The three reads: workspace tar, trajectory-so-far, env service state.
|
|
398
|
+
|
|
399
|
+
Not atomic and the agent writes throughout, so a degraded trajectory or env
|
|
400
|
+
read yields ``partial``; only a missing tar is ``failed``. Mutates ``row``
|
|
401
|
+
in place.
|
|
402
|
+
"""
|
|
403
|
+
# One budget across all three reads, so the outer `wait_for` can bound them.
|
|
404
|
+
deadline = time.monotonic() + self.config.timeout_seconds
|
|
405
|
+
|
|
406
|
+
def remaining() -> float:
|
|
407
|
+
return max(1.0, deadline - time.monotonic()) # never a zero-timeout read
|
|
408
|
+
|
|
409
|
+
reasons: list[str] = []
|
|
410
|
+
agent = next(
|
|
411
|
+
(a for a in context.deployed_agents if a.agent_name == self.agent_name), None
|
|
412
|
+
)
|
|
413
|
+
if agent is None:
|
|
414
|
+
raise RuntimeError(f"No deployed agent named '{self.agent_name}' found in context")
|
|
415
|
+
a2a_url = agent.a2a_url or agent.api_url
|
|
416
|
+
if not a2a_url:
|
|
417
|
+
raise RuntimeError(f"Deployed agent '{self.agent_name}' has no a2a_url")
|
|
418
|
+
a2a_card = agent.a2a_card or {}
|
|
419
|
+
|
|
420
|
+
workspace = await capture.capture_workspace(
|
|
421
|
+
a2a_url=a2a_url,
|
|
422
|
+
a2a_card=a2a_card,
|
|
423
|
+
agent_name=self.agent_name,
|
|
424
|
+
a2a_context_id=self.a2a_context_id,
|
|
425
|
+
artifact_id=self.workspace_artifact_id,
|
|
426
|
+
timeout_seconds=remaining(),
|
|
427
|
+
)
|
|
428
|
+
row["id"] = workspace.universe_id
|
|
429
|
+
row["version"] = workspace.universe_version
|
|
430
|
+
row["bundle_object_url"] = workspace.bundle_object_url
|
|
431
|
+
|
|
432
|
+
# Every point reads its trajectory the same way, cumulatively.
|
|
433
|
+
# `PromptResponse.agent_trajectory_s3_uri` is only the LAST turn's, so using
|
|
434
|
+
# it for the final row would make that point narrower than its predecessors.
|
|
435
|
+
# It stays the fallback: on a failed run the live read may be gone while the
|
|
436
|
+
# per-turn upload already landed.
|
|
437
|
+
traj = await capture.read_partial_trajectory(
|
|
438
|
+
a2a_url=a2a_url,
|
|
439
|
+
a2a_card=a2a_card,
|
|
440
|
+
context_id=self.a2a_context_id,
|
|
441
|
+
timeout_seconds=remaining(),
|
|
442
|
+
trajectory_output_prefix=self.trajectory_output_prefix,
|
|
443
|
+
)
|
|
444
|
+
if traj.reason and is_final:
|
|
445
|
+
recorded = self._recorded_trajectory_uri(context)
|
|
446
|
+
row["trajectory_s3_uri"] = recorded
|
|
447
|
+
reasons.append(traj.reason if recorded else f"{traj.reason}; trajectory_missing")
|
|
448
|
+
elif traj.reason:
|
|
449
|
+
reasons.append(traj.reason)
|
|
450
|
+
elif traj.object_url:
|
|
451
|
+
row["trajectory_s3_uri"] = traj.object_url
|
|
452
|
+
else:
|
|
453
|
+
try:
|
|
454
|
+
row["trajectory_s3_uri"] = await asyncio.to_thread(
|
|
455
|
+
capture.upload_trajectory, traj.trajectory,
|
|
456
|
+
self.trajectory_output_prefix,
|
|
457
|
+
)
|
|
458
|
+
except Exception as exc:
|
|
459
|
+
logger.warning("%s: trajectory upload failed: %s", self.step_id, exc)
|
|
460
|
+
reasons.append("trajectory_upload_failed")
|
|
461
|
+
|
|
462
|
+
if self.config.env_id:
|
|
463
|
+
reason = await self._capture_env_state(context, row, remaining())
|
|
464
|
+
if reason:
|
|
465
|
+
reasons.append(reason)
|
|
466
|
+
|
|
467
|
+
if reasons:
|
|
468
|
+
# Never `failed`: disqualifying on an auxiliary read would discard the
|
|
469
|
+
# whole curve the moment one degrades for the run's lifetime.
|
|
470
|
+
row["capture_status"] = "partial"
|
|
471
|
+
row["capture_reason"] = "; ".join(reasons)
|
|
472
|
+
return row
|
|
473
|
+
|
|
474
|
+
async def _capture_env_state(
|
|
475
|
+
self, context: TaskStepContext, row: dict, timeout_seconds: float
|
|
476
|
+
) -> Optional[str]:
|
|
477
|
+
"""Freeze the env's service state as a ``EnvironmentUniverseArtifact`` version,
|
|
478
|
+
recording its ``{id, version}`` on the row. Returns a reason on failure —
|
|
479
|
+
the export is all-or-nothing, so one unreachable service costs this leg of
|
|
480
|
+
the tick, not the workspace bundle."""
|
|
481
|
+
from agent_env.env.env import Env, gateway_url_of
|
|
482
|
+
from agent_env.task_step.task_steps.snapshot_env import SnapshotEnvTaskStep
|
|
483
|
+
|
|
484
|
+
deployed = next(
|
|
485
|
+
(d for d in context.deployed_envs if d.env_id == self.config.env_id), None
|
|
486
|
+
)
|
|
487
|
+
if deployed is None:
|
|
488
|
+
logger.warning(
|
|
489
|
+
"%s: env_id=%s not in deployed_envs (%s); skipping env capture",
|
|
490
|
+
self.step_id, self.config.env_id, [d.env_id for d in context.deployed_envs],
|
|
491
|
+
)
|
|
492
|
+
return "env_not_deployed"
|
|
493
|
+
if not gateway_url_of(deployed):
|
|
494
|
+
return "env_has_no_gateway_url"
|
|
495
|
+
|
|
496
|
+
try:
|
|
497
|
+
# Pin the env version the sandbox actually runs, so service
|
|
498
|
+
# enumeration can't drift if the env was re-registered since deploy.
|
|
499
|
+
env = await asyncio.to_thread(Env.get, deployed.env_id, deployed.env_version)
|
|
500
|
+
result = await SnapshotEnvTaskStep.snapshot_env_state(
|
|
501
|
+
env=env,
|
|
502
|
+
gateway_url=deployed.gateway_url,
|
|
503
|
+
# Instance-scoped like the workspace id, never context-scoped: a
|
|
504
|
+
# pinned `context_id` is shared by every concurrent rollout.
|
|
505
|
+
snapshot_id=f"snapshot-{self.config.env_id}-{self._discriminator}",
|
|
506
|
+
deployed=deployed,
|
|
507
|
+
export_timeout_seconds=int(timeout_seconds),
|
|
508
|
+
)
|
|
509
|
+
except Exception as exc:
|
|
510
|
+
logger.warning(
|
|
511
|
+
"%s: env state capture failed: %s", self.step_id, exc, exc_info=True,
|
|
512
|
+
)
|
|
513
|
+
return f"env_capture_failed: {type(exc).__name__}"
|
|
514
|
+
|
|
515
|
+
row["env_universe_id"] = result.environment_universe_artifact_id
|
|
516
|
+
row["env_universe_version"] = result.environment_universe_artifact_version
|
|
517
|
+
return None
|
|
518
|
+
|
|
519
|
+
def _recorded_trajectory_uri(self, context: TaskStepContext) -> Optional[str]:
|
|
520
|
+
"""This rollout's per-turn trajectory uri from ``prompt_responses``.
|
|
521
|
+
|
|
522
|
+
Matched on ``a2a_context_id``, newest-first: a retry can restore a context
|
|
523
|
+
snapshot taken during an earlier attempt, so that attempt's response for the
|
|
524
|
+
same ``prompt_id`` can still be present and would otherwise attach a
|
|
525
|
+
previous conversation's trajectory to this capture.
|
|
526
|
+
"""
|
|
527
|
+
for prompt in reversed(context.prompt_responses):
|
|
528
|
+
if (
|
|
529
|
+
prompt.prompt_id == self.prompt_id
|
|
530
|
+
and prompt.a2a_context_id == self.a2a_context_id
|
|
531
|
+
):
|
|
532
|
+
return prompt.agent_trajectory_s3_uri
|
|
533
|
+
# None is legitimate: the response is recorded only after the turn loop, so
|
|
534
|
+
# a capture during an in-flight or failed run has no entry yet.
|
|
535
|
+
return None
|
|
536
|
+
|
|
537
|
+
# ------------------------------------------------------------------ rows
|
|
538
|
+
|
|
539
|
+
def _row(
|
|
540
|
+
self,
|
|
541
|
+
*,
|
|
542
|
+
is_final: bool,
|
|
543
|
+
status: str,
|
|
544
|
+
reason: Optional[str] = None,
|
|
545
|
+
turn: Optional[int] = None,
|
|
546
|
+
) -> dict:
|
|
547
|
+
"""Build a manifest row. No stored ordinal: rank by ``version`` within a
|
|
548
|
+
``source_context_id``, which survives a filtered subset where array
|
|
549
|
+
position would not. ``turn`` is 1-based and only set on turn-boundary rows."""
|
|
550
|
+
return {
|
|
551
|
+
"turn": turn,
|
|
552
|
+
"id": None,
|
|
553
|
+
"version": None,
|
|
554
|
+
"bundle_object_url": None,
|
|
555
|
+
"source_agent_name": self.agent_name,
|
|
556
|
+
"source_context_id": self.a2a_context_id,
|
|
557
|
+
"elapsed_s": int(time.monotonic() - self._started_monotonic),
|
|
558
|
+
"captured_at_utc": datetime.now(timezone.utc).isoformat(),
|
|
559
|
+
"is_final": is_final,
|
|
560
|
+
"capture_status": status,
|
|
561
|
+
"capture_reason": reason,
|
|
562
|
+
"trajectory_s3_uri": None,
|
|
563
|
+
"env_universe_id": None,
|
|
564
|
+
"env_universe_version": None,
|
|
565
|
+
}
|
|
566
|
+
|
|
567
|
+
def _record(
|
|
568
|
+
self, *, status: str, reason: str, is_final: bool = False,
|
|
569
|
+
turn: Optional[int] = None,
|
|
570
|
+
) -> None:
|
|
571
|
+
self._append(
|
|
572
|
+
self._row(is_final=is_final, status=status, reason=reason, turn=turn)
|
|
573
|
+
)
|
|
574
|
+
|
|
575
|
+
def _append(self, row: dict) -> None:
|
|
576
|
+
for key in [k for k, v in row.items() if v is None]:
|
|
577
|
+
del row[key]
|
|
578
|
+
self._rows.append(row)
|
|
579
|
+
logger.info(
|
|
580
|
+
"%s: recorded snapshot (v=%s is_final=%s turn=%s status=%s elapsed_s=%s "
|
|
581
|
+
"env_universe_version=%s reason=%s)",
|
|
582
|
+
self.step_id, row.get("version"), row["is_final"], row.get("turn"),
|
|
583
|
+
row["capture_status"], row["elapsed_s"], row.get("env_universe_version"),
|
|
584
|
+
row.get("capture_reason"),
|
|
585
|
+
)
|
|
586
|
+
|
|
587
|
+
def _publish(self, context: TaskStepContext) -> None:
|
|
588
|
+
"""Publish the collected rows to the context, once, at teardown: collecting on
|
|
589
|
+
the series means a capture racing teardown cannot land after the final row."""
|
|
590
|
+
context.metadata.setdefault("agent_snapshots", []).extend(self._rows)
|
|
591
|
+
final = next(
|
|
592
|
+
(r for r in reversed(self._rows)
|
|
593
|
+
if r.get("is_final") and r.get("env_universe_id")),
|
|
594
|
+
None,
|
|
595
|
+
)
|
|
596
|
+
if final:
|
|
597
|
+
# Under the step id, so a judge step can be wired to this one via
|
|
598
|
+
# `load_artifact`'s `artifact_from_step_id` without knowing it snapshots.
|
|
599
|
+
context.metadata.setdefault("env_snapshotted_universes", {})[self.step_id] = {
|
|
600
|
+
"id": final["env_universe_id"],
|
|
601
|
+
"version": final["env_universe_version"],
|
|
602
|
+
}
|
|
603
|
+
|
|
604
|
+
# ------------------------------------------------------------- fatality
|
|
605
|
+
|
|
606
|
+
def raise_if_final_capture_missing(self) -> None:
|
|
607
|
+
"""Fail the step when the agent succeeded but its final capture did not —
|
|
608
|
+
otherwise the instance reports `completed` with no gradable point.
|
|
609
|
+
|
|
610
|
+
Keyed on a landed bundle, not `capture_status`: a `partial` row with a bundle
|
|
611
|
+
is gradable, while `final_capture_lock_timeout` is `skipped` yet fatal.
|
|
612
|
+
"""
|
|
613
|
+
if not self.config.at_end:
|
|
614
|
+
return
|
|
615
|
+
final = next((r for r in reversed(self._rows) if r.get("is_final")), None)
|
|
616
|
+
if final is None:
|
|
617
|
+
# Unreachable after a successful run, and no recorded failure to
|
|
618
|
+
# attribute it to — so loud, but not fatal.
|
|
619
|
+
logger.warning(
|
|
620
|
+
"%s: agent run succeeded but no is_final snapshot row was recorded",
|
|
621
|
+
self.step_id,
|
|
622
|
+
)
|
|
623
|
+
return
|
|
624
|
+
if final.get("bundle_object_url"):
|
|
625
|
+
return
|
|
626
|
+
raise RuntimeError(
|
|
627
|
+
f"{self.step_id}: final capture did not land "
|
|
628
|
+
f"(status={final.get('capture_status')} reason={final.get('capture_reason')}) — "
|
|
629
|
+
"the rollout completed but has no gradable end state, so failing the step "
|
|
630
|
+
"rather than reporting a complete instance with nothing to grade"
|
|
631
|
+
)
|