Report agent token usage from sandbox jobs

sandbox/models.py

@@ -8,12 +8,18 @@ from typing import Any
 
 @dataclass(frozen=True, slots=True)
 class JobResult:
-    """Structured result returned by a direct ReMe job invocation."""
+    """Structured result returned by a direct ReMe job invocation.
+
+    ``token_usage`` contains the input/output/total token delta for each
+    configured agent wrapper during this invocation. Missing provider usage is
+    represented by ``None`` at the metric level.
+    """
 
     success: bool
     answer: Any = ""
     metadata: dict[str, Any] = field(default_factory=dict)
     error: str | None = None
+    token_usage: dict[str, dict[str, int | None]] = field(default_factory=dict)
 
     @classmethod
     def from_dict(cls, value: dict[str, Any]) -> "JobResult":
@@ -22,6 +28,7 @@ class JobResult:
             success=bool(value.get("success")),
             answer=value.get("answer", ""),
             metadata=value.get("metadata") if isinstance(value.get("metadata"), dict) else {},
+            token_usage=value.get("token_usage") if isinstance(value.get("token_usage"), dict) else {},
             error=value.get("error") if isinstance(value.get("error"), str) else None,
         )
 

sandbox/worker.py

@@ -80,8 +80,10 @@ async def _run(request: dict[str, Any]) -> dict[str, Any]:
     # Import after setting TMPDIR so ReMe and its subprocesses cannot leave
     # case-specific temporary files outside the disposable case root.
     from reme.config import resolve_app_config
+    from reme.enumeration import ComponentEnum
     from reme.reme import ReMe
     from reme.schema import ApplicationConfig
+    from reme.utils.evaluation_interface import track_agent_token_usage
 
     app_config = resolve_app_config(
         config=config,
@@ -95,11 +97,24 @@ async def _run(request: dict[str, Any]) -> dict[str, Any]:
     app = ReMe(**app_config)
     try:
         await app.start()
-        response = await app.run_job(job, **job_args)
+        agent_names = list(app.context.components.get(ComponentEnum.AGENT_WRAPPER, {}))
+        try:
+            with track_agent_token_usage(agent_names, app.context) as token_usage:
+                response = await app.run_job(job, **job_args)
+        except Exception as exc:  # Preserve usage accumulated before a job failure.
+            return {
+                "success": False,
+                "answer": "",
+                "metadata": {},
+                "token_usage": token_usage,
+                "error": f"{type(exc).__name__}: {exc}",
+                "traceback": traceback.format_exc(),
+            }
         return {
             "success": bool(response.success),
             "answer": response.answer,
             "metadata": response.metadata,
+            "token_usage": token_usage,
             "error": None if response.success else str(response.answer),
         }
     finally:
@@ -125,6 +140,7 @@ def main() -> int:
             "success": False,
             "answer": "",
             "metadata": {},
+            "token_usage": {},
             "error": f"{type(exc).__name__}: {exc}",
             "traceback": traceback.format_exc(),
         }

tests/unit/test_sandbox.py

@@ -10,7 +10,9 @@ from types import SimpleNamespace
 
 import pytest
 
+from reme.enumeration import ComponentEnum
 from reme.schema import ApplicationConfig
+from reme.utils import global_counter_add
 from sandbox import DockerReMeSandboxFactory, ImageCandidate, SourceCandidate, SourceSnapshot
 from sandbox.direct_docker_workspace import DirectDockerWorkspace
 from sandbox import worker
@@ -79,7 +81,15 @@ class FakeBackend:
                 },
             ).encode()
             self.files[response_path] = json.dumps(
-                {"success": True, "answer": "yes", "metadata": {"job": "ok"}, "error": None},
+                {
+                    "success": True,
+                    "answer": "yes",
+                    "metadata": {"job": "ok"},
+                    "token_usage": {
+                        "bench": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
+                    },
+                    "error": None,
+                },
             ).encode()
         if command[:2] == ["tar", "-czf"]:
             self.files[command[2]] = b"case-archive"
@@ -204,6 +214,9 @@ async def test_image_candidate_skips_source_upload_and_runs_jobs(tmp_path):
 
     assert result.success is True
     assert result.answer == "yes"
+    assert result.token_usage == {
+        "bench": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
+    }
     assert "/workspace/candidate.tar.gz" not in workspaces[0].backend.files
     assert exported.read_bytes() == b"case-archive"
     assert json.loads(workspaces[0].backend.files["/workspace/case/manifest.json"])["candidate_mode"] == "image"
@@ -477,6 +490,10 @@ async def test_worker_calls_application_directly_and_always_closes(monkeypatch,
         def __init__(self, **config):
             """Record construction."""
             events.append(("construct", config["workspace_dir"]))
+            self.context = SimpleNamespace(
+                metadata={},
+                components={ComponentEnum.AGENT_WRAPPER: {"bench": SimpleNamespace()}},
+            )
 
         async def start(self):
             """Record startup."""
@@ -485,6 +502,9 @@ async def test_worker_calls_application_directly_and_always_closes(monkeypatch,
         async def run_job(self, name, **arguments):
             """Return a successful direct job response."""
             events.append(("run_job", name, arguments))
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "input_tokens"], 12)
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "output_tokens"], 3)
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "total_tokens"], 15)
             return SimpleNamespace(success=True, answer="answer", metadata={"direct": True})
 
         async def close(self):
@@ -503,7 +523,15 @@ async def test_worker_calls_application_directly_and_always_closes(monkeypatch,
         },
     )
 
-    assert result == {"success": True, "answer": "answer", "metadata": {"direct": True}, "error": None}
+    assert result == {
+        "success": True,
+        "answer": "answer",
+        "metadata": {"direct": True},
+        "token_usage": {
+            "bench": {"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
+        },
+        "error": None,
+    }
     assert worker.os.environ["TMPDIR"] == str(tmp_path / "tmp")
     layout = json.loads((tmp_path / "runtime-layout.json").read_text())
     assert layout["workspace_root"] == "workspace"
@@ -517,3 +545,48 @@ async def test_worker_calls_application_directly_and_always_closes(monkeypatch,
         ("run_job", "agentic_answer", {"query": "hello"}),
         "close",
     ]
+
+
+@pytest.mark.asyncio
+async def test_worker_returns_usage_accumulated_before_job_failure(monkeypatch, tmp_path):
+    """A failed step still reports tokens spent before the exception."""
+
+    class FailingReMe:
+        """Record token usage and then fail the requested job."""
+
+        def __init__(self, **_config):
+            self.context = SimpleNamespace(
+                metadata={},
+                components={ComponentEnum.AGENT_WRAPPER: {"bench": SimpleNamespace()}},
+            )
+
+        async def start(self):
+            """Start without side effects."""
+
+        async def run_job(self, _name, **_arguments):
+            """Consume tokens before surfacing a step error."""
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "input_tokens"], 8)
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "output_tokens"], 2)
+            global_counter_add(self.context.metadata, ["__token_counter", "bench", "total_tokens"], 10)
+            raise RuntimeError("step failed")
+
+        async def close(self):
+            """Close without side effects."""
+
+    monkeypatch.setattr("reme.config.resolve_app_config", lambda **kwargs: kwargs)
+    monkeypatch.setattr("reme.reme.ReMe", FailingReMe)
+
+    result = await worker._run(  # pylint: disable=protected-access
+        {
+            "job": "agentic_answer",
+            "config": "lme.yaml",
+            "case_root": str(tmp_path),
+            "workspace_dir": str(tmp_path / "workspace"),
+        },
+    )
+
+    assert result["success"] is False
+    assert result["error"] == "RuntimeError: step failed"
+    assert result["token_usage"] == {
+        "bench": {"input_tokens": 8, "output_tokens": 2, "total_tokens": 10},
+    }