From 32b2c26734c3ee5cfa5a9d136f067b68cbefc469 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=85=B1=E7=89=9B=E8=82=89?= Date: Tue, 4 Aug 2026 17:25:02 +0800 Subject: [PATCH] Report agent token usage from sandbox jobs --- sandbox/models.py | 9 ++++- sandbox/worker.py | 18 ++++++++- tests/unit/test_sandbox.py | 77 +++++++++++++++++++++++++++++++++++++- 3 files changed, 100 insertions(+), 4 deletions(-) diff --git a/sandbox/models.py b/sandbox/models.py index 310e12e1..35204b0b 100644 --- a/sandbox/models.py +++ b/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, ) diff --git a/sandbox/worker.py b/sandbox/worker.py index 23d33654..57301584 100644 --- a/sandbox/worker.py +++ b/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(), } diff --git a/tests/unit/test_sandbox.py b/tests/unit/test_sandbox.py index ac196f5e..1ee77869 100644 --- a/tests/unit/test_sandbox.py +++ b/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}, + }