feat: Add Logfire error tracking to fleet env

[dzorlu]

Summary

• Adds structured error tracking via Logfire across fleet env (init failures, tool call failures, MCP timeouts, verifier errors)
• New telemetry.py wrapper with 4 functions (fleet_info, fleet_warning, fleet_error, fleet_exception)
• 15 instrumentation sites across task_env.py, client.py, mcp_tools.py
• No-op if configure_fleet_telemetry() is never called — logfire silently drops events

Test plan

• All 41 previously-passing tests still pass
• 4 pre-existing failures unaffected
• Smoke test with configure_fleet_telemetry(send_to_logfire=False) to verify structured output

🤖 Generated with Claude Code

[cursor]

Shared ContextVar causes wrong telemetry context across instances

Medium Severity

set_task_context is called only in __init__, writing to a thread-level ContextVar. When multiple FleetTaskEnv instances exist in the same thread (e.g., via from_json_file_all), each constructor overwrites the shared _session_context. All subsequent telemetry events from earlier instances (during reset_async, step_async, _compute_reward) will carry the last-created instance's env_key/task_key/etc., producing silently incorrect monitoring data. The context needs to be re-established before telemetry-emitting methods, not just once at construction.

Additional Locations (1)

• src/envs/fleet_env/telemetry.py#L16-L18

 

[cursor]

Breaking API change leaves existing callers without required args

High Severity

from_fleet changed data_key, data_version, and image_type from Optional[str] = None to required positional arguments, but existing callers in test_fleet_env.py (lines 83, 92, 105) and examples/fleet_env_example.py still call it with only api_key and env_key. These raise TypeError at runtime, breaking previously passing tests and the documented example.

Additional Locations (1)

• tests/envs/test_fleet_env.py#L82-L83

 

[cursor]

Unconditional 8-second wait on every list_tools call

Medium Severity

The initial_wait sleep in list_tools() fires on every invocation, not just the first one after provisioning. The FleetMCPTools instance is created once in _ensure_provisioned() and kept alive across reset_async() calls (without close()). Each subsequent call to reset_async() re-invokes self._tools.list_tools(), which unconditionally sleeps for initial_wait=8.0 seconds. In a standard multi-episode RL training loop (where reset() is called once per episode on the same env instance), every episode after the first incurs an unnecessary 8-second penalty.

Additional Locations (1)

• src/envs/fleet_env/task_env.py#L295-L299

 

[cursor]

Consecutive resets silently lose rollout tracking events

Medium Severity

If reset_async() is called a second time without the first rollout completing (e.g., the caller decides to restart the episode), a new fleet_rollout_started is emitted and _rollout_completed_emitted is reset to False, but no fleet_rollout_completed is ever emitted for the first rollout. This breaks the invariant documented in the README that "every fleet_rollout_started gets a matching fleet_rollout_completed." The close() handler only emits one compensating fleet_rollout_completed, leaving the first rollout permanently unmatched.

Additional Locations (1)

• src/envs/fleet_env/task_env.py#L616-L631

 

[cursor]

Sync from_fleet missing fleet_provisioning_completed telemetry event

Low Severity

The async from_fleet_async() emits a fleet_provisioning_completed telemetry event with provisioning_time_s and instance_id, but the sync from_fleet() does not. This means provisioning latency data is silently missing for any code path using the sync method (including the example at examples/fleet_env_example.py), breaking the Logfire SQL queries documented in the README.

Additional Locations (1)

• src/envs/fleet_env/client.py#L269-L274

 

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Bugbot Autofix prepared a fix for the issue found in the latest run.

• ✅ Fixed: Telemetry context overwritten by concurrent FleetTaskEnv instances
  • FleetTaskEnv now reapplies its own telemetry context at the start of runtime lifecycle methods instead of setting shared context in init, preventing cross-instance context bleed in concurrent usage.

Or push these changes by commenting:

@cursor push 5469588568

Preview (5469588568)

diff --git a/src/envs/fleet_env/task_env.py b/src/envs/fleet_env/task_env.py
--- a/src/envs/fleet_env/task_env.py
+++ b/src/envs/fleet_env/task_env.py
@@ -125,14 +125,6 @@
         self._rollout_started = False
         self._tools_cache: Optional[List[Dict]] = None
 
-        # Set telemetry context so init failures are tracked with full context
-        set_task_context(
-            env_key=self.env_key,
-            env_version=self.env_version,
-            task_key=self.task_key,
-            modality=self.modality,
-        )
-
         # Provisioning is deferred to _ensure_provisioned() (called from reset_async)
         # to avoid blocking the event loop with sync Fleet.make() calls.
         self._orch = None
@@ -191,6 +183,15 @@
         """
         return self.task.get("env_variables")
 
+    def _set_telemetry_context(self):
+        """Set telemetry context for this environment instance."""
+        set_task_context(
+            env_key=self.env_key,
+            env_version=self.env_version,
+            task_key=self.task_key,
+            modality=self.modality,
+        )
+
     def reset(self, seed: Optional[int] = None) -> Dict[str, Any]:
         """Reset the environment and return initial observation (sync wrapper).
 
@@ -253,6 +254,7 @@
         import logging
 
         logger = logging.getLogger(__name__)
+        self._set_telemetry_context()
 
         # Count this rollout attempt immediately — even if provisioning fails,
         # it's still a rollout attempt (e.g., fostgres health check failures).
@@ -437,6 +439,8 @@
         Returns:
             Tuple of (observation, reward, done, info)
         """
+        self._set_telemetry_context()
+
         if self._done:
             raise RuntimeError("Episode is done. Call reset() to start a new episode.")
 
@@ -526,6 +530,8 @@
         Returns:
             1.0 if verifier passes, 0.0 otherwise
         """
+        self._set_telemetry_context()
+
         # Support both field names: verifier_code (OpenEnv) and verifier_func (Fleet SDK)
         verifier_code = self.task.get("verifier_code") or self.task.get("verifier_func")
         score = 0.0
@@ -617,6 +623,8 @@
         context overflow, job cancellation, TTL expiry).
         """
         try:
+            self._set_telemetry_context()
+
             # Emit rollout_completed for orphaned rollouts (started but never completed).
             # This happens when the caller (SkyRL) stops without telling us why:
             # max_turns hit, context overflow, job cancellation, etc.
@@ -649,6 +657,8 @@
     async def close_async(self):
         """Async close — avoids blocking the event loop on Fleet instance termination."""
         try:
+            self._set_telemetry_context()
+
             if self._rollout_started and not self._rollout_completed_emitted:
                 stop_reason = "max_steps" if self._step_count >= self.max_steps else "abandoned"
                 fleet_info(

diff --git a/tests/envs/test_fleet_task_env.py b/tests/envs/test_fleet_task_env.py
--- a/tests/envs/test_fleet_task_env.py
+++ b/tests/envs/test_fleet_task_env.py
@@ -82,7 +82,51 @@
         with pytest.raises(ValueError, match="Fleet API key required"):
             FleetTaskEnv(sample_task_config)
 
+    @pytest.mark.anyio
+    async def test_reset_uses_instance_telemetry_context(self, mock_fleet_env_client):
+        """Telemetry should match the env instance being reset."""
+        from envs.fleet_env.task_env import FleetTaskEnv
 
+        task_a = {
+            "task_key": "task-a",
+            "prompt": "Task A",
+            "env_key": "env-a",
+            "env_version": "v1",
+            "task_modality": "tool_use",
+        }
+        task_b = {
+            "task_key": "task-b",
+            "prompt": "Task B",
+            "env_key": "env-b",
+            "env_version": "v2",
+            "task_modality": "tool_use",
+        }
+
+        env_a = FleetTaskEnv(task_a, api_key="test")
+        env_b = FleetTaskEnv(task_b, api_key="test")
+
+        for env in (env_a, env_b):
+            env._orch = MagicMock()
+            env._tools = MagicMock()
+            env._tools.list_tools = AsyncMock(
+                return_value=MagicMock(tools=[{"name": "search"}])
+            )
+
+        with patch("envs.fleet_env.telemetry.logfire.info") as mock_info:
+            await env_a.reset_async()
+            await env_b.reset_async()
+
+        started_calls = [
+            call
+            for call in mock_info.call_args_list
+            if call.args and call.args[0] == "fleet_rollout_started"
+        ]
+        assert started_calls[0].kwargs["task_key"] == "task-a"
+        assert started_calls[0].kwargs["env_key"] == "env-a"
+        assert started_calls[1].kwargs["task_key"] == "task-b"
+        assert started_calls[1].kwargs["env_key"] == "env-b"
+
+
 class TestFleetTaskEnvSpecs:
     """Tests for env/data spec building."""

_{This Bugbot Autofix run was free. To enable autofix for future PRs, go to the Cursor dashboard.}

[cursor]

Telemetry context overwritten by concurrent FleetTaskEnv instances

Medium Severity

_session_context is a single ContextVar shared across all FleetTaskEnv instances. set_task_context() is called in __init__, so creating multiple instances in the same async task causes later instances to overwrite earlier ones' telemetry context. Subsequent telemetry events (e.g., from reset_async()) for earlier instances will carry the wrong env_key, task_key, etc.

Additional Locations (1)

• src/envs/fleet_env/task_env.py#L127-L134