feat(studio): fabric agent will publish tokens

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_eval/runtimes/fabric/runtime.py

@@ -55,13 +55,14 @@
 from nemo_evaluator_sdk.agent_eval.tasks import AgentEvalRunConfig, AgentEvalTask
 from nemo_evaluator_sdk.agent_eval.trials import AgentEvalTrial, AgentEvalTrialStatus, AgentOutput, RunnerInfo
 from nemo_evaluator_sdk.agent_eval.workspace_seeds import SEED_FILES_INPUT_KEY, seed_workspace
+from nemo_evaluator_sdk.values.atif import FinalMetrics
 from nemo_evaluator_sdk.values.evidence import (
     EVIDENCE_FORMAT_ATIF,
     EVIDENCE_TRACE,
     CandidateEvidence,
     EvidenceDescriptor,
 )
-from pydantic import JsonValue
+from pydantic import JsonValue, ValidationError
 
 if TYPE_CHECKING:
     # Annotations use nemo_fabric's real types (single source of truth). nemo_fabric is an optional
@@ -358,9 +359,13 @@ async def _run_task(
                         raise
                     hook_extras = None
         except TimeoutError as exc:
-            return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances))
+            return self._failed_trial(
+                task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir)
+            )
         except Exception as exc:  # noqa: BLE001 - a task failure must not abort the whole run
-            return self._failed_trial(task, evidence_dir, exc, extra_metadata=self._skill_metadata(skill_provenances))
+            return self._failed_trial(
+                task, evidence_dir, exc, extra_metadata=self._failed_metadata(skill_provenances, evidence_dir)
+            )
         finally:
             if self._task_hook is not None:
                 try:
@@ -396,6 +401,18 @@ def _skill_metadata(provenances: list[SkillProvenance]) -> dict[str, Any]:
         """
         return {"skill": provenances[0] if len(provenances) == 1 else None, "skills": provenances}
 
+    @staticmethod
+    def _failed_metadata(provenances: list[SkillProvenance], evidence_dir: Path) -> dict[str, Any]:
+        """Trial metadata for a timed-out/errored task: skill provenance plus whatever tokens Relay flushed.
+
+        Timeouts never reach ``_to_trial``, and there is no ``RunResult`` here, so the trajectory is read
+        straight from the relay dir — these are the long, expensive rows the token count matters most for.
+        """
+        return {
+            **FabricAgentRuntime._skill_metadata(provenances),
+            **_atif_token_metadata(_relay_atif_path(evidence_dir)),
+        }
+
     def _to_trial(
         self,
         task: AgentEvalTask,
@@ -421,6 +438,9 @@ def _to_trial(
             # Skill provenance (name + content hash + injection mode) for the A/B diff.
             **self._skill_metadata(skill_provenances or []),
             **extras,
+            # Token usage from the Relay ATIF trajectory; Fabric's RunResult carries no usage of its
+            # own. Merged last so a hook extra can't shadow it.
+            **_atif_token_metadata(_atif_artifact_path(result)),
         }
 
         if result.status != "succeeded":
@@ -715,6 +735,83 @@ def _result_error(result: RunResult) -> Mapping[str, Any]:
     return {"stage": error.stage, "code": error.code, "message": error.message}
 
 
+def _atif_artifact_path(result: RunResult) -> Path | None:
+    """Path of the ATIF trajectory Fabric promoted as an artifact, if any."""
+    for artifact in result.artifacts.artifacts:
+        if artifact.kind == _ATIF_ARTIFACT_KIND:
+            return Path(artifact.path)
+    return None
+
+
+def _relay_atif_path(evidence_dir: Path) -> Path | None:
+    """Path of the Relay-written ATIF trajectory, used when no ``RunResult`` exists (timeout/error).
+
+    Relay's filename template is per-session, so more than one file can land when subagents emit
+    their own sessions. Picking one under-reports and summing double-counts a root that already
+    aggregates, so anything other than a single match reports nothing rather than a wrong number.
+    """
+    matches = sorted((evidence_dir / _RELAY_SUBDIR).glob(_ATIF_FILENAME_TEMPLATE.format(session_id="*")))
+    if len(matches) == 1:
+        return matches[0]
+    if matches:
+        logger.warning("Fabric token capture: %d ATIF trajectories under %s; skipping", len(matches), evidence_dir)
+    return None
+
+
+def _atif_token_metadata(path: Path | None) -> dict[str, int]:
+    """Project an ATIF trajectory's token totals onto the trial-metadata ``TOKEN_KEYS``.
+
+    Each token field is resolved on its own: the trajectory-level ``final_metrics`` aggregate when it
+    reports that field, else the sum of the matching per-step ``metrics``. Every ``final_metrics``
+    field is optional, so a block carrying only ``total_steps`` or a cost — or one that fails to
+    validate — must not suppress counts the steps do carry. That partial shape is likeliest on the
+    timeout path, where the trajectory was flushed mid-run and the counts matter most.
+
+    ``total_tokens`` and ``cache_creation_tokens`` have no ATIF source and stay unset — Intake
+    recomputes the total. A missing or unreadable trajectory yields ``{}``: an absent token count
+    must not fail the trial.
+    """
+    if path is None:
+        return {}
+    try:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError) as exc:
+        logger.warning("Fabric token capture: unreadable ATIF trajectory %s (%s)", path, exc)
+        return {}
+    if not isinstance(payload, Mapping):
+        return {}
+
+    totals = FinalMetrics()
+    final_metrics = payload.get("final_metrics")
+    if isinstance(final_metrics, Mapping):
+        try:
+            totals = FinalMetrics.model_validate(final_metrics)
+        except ValidationError as exc:
+            logger.warning("Fabric token capture: invalid final_metrics in %s (%s)", path, exc)
+
+    captured = {
+        "prompt_tokens": (totals.total_prompt_tokens, "prompt_tokens"),
+        "completion_tokens": (totals.total_completion_tokens, "completion_tokens"),
+        "cache_read_tokens": (totals.total_cached_tokens, "cached_tokens"),
+    }
+    resolved = {
+        key: total if total is not None else _sum_step_metric(payload, step_key)
+        for key, (total, step_key) in captured.items()
+    }
+    return {key: value for key, value in resolved.items() if value is not None}
+
+
+def _sum_step_metric(payload: Mapping[str, Any], key: str) -> int | None:
+    """Sum one per-step ATIF metric across the trajectory, or ``None`` when no step reported it."""
+    total: int | None = None
+    for step in payload.get("steps") or []:
+        metrics = step.get("metrics") if isinstance(step, Mapping) else None
+        value = metrics.get(key) if isinstance(metrics, Mapping) else None
+        if isinstance(value, int) and not isinstance(value, bool):
+            total = value if total is None else total + value
+    return total
+
+
 def _safe_path_name(value: str) -> str:
     return "".join(char if char.isalnum() or char in "._-" else "-" for char in value).strip(".-")[:120]
 

packages/nemo_evaluator_sdk/src/nemo_evaluator_sdk/agent_inference.py

@@ -369,6 +369,12 @@ async def _invoke_post() -> dict[str, Any]:
             field_name=invocation.response_path_field,
         )
         response = _openai_response(str(response_value))
+        # The synthesized response keeps only the extracted text, so carry the agent's own token
+        # usage across: it is the sole token source for a row evaluation, whose trials are otherwise
+        # built with no measurements at all. Model targets already return their full completion.
+        usage = result_data.get("usage") if isinstance(result_data, Mapping) else None
+        if isinstance(usage, Mapping):
+            response["usage"] = dict(usage)
         if invocation.trajectory_path:
             trajectory = _extract_jsonpath(
                 result_data,

packages/nemo_evaluator_sdk/tests/agent_eval/test_fabric_runtime.py

@@ -1100,3 +1100,107 @@ def handler(agent: Any, kwargs: dict[str, Any]) -> _FakeResult:
     assert not (workspace / ".agents").exists()
     # Provenance is still stamped on the failed trial for the A/B diff.
     assert [prov["name"] for prov in trials[0].metadata["skills"]] == list(names)
+
+
+# --- ATIF token capture -----------------------------------------------------
+
+
+def _atif(tmp_path: Path, payload: Mapping[str, Any]) -> Path:
+    path = tmp_path / "trajectory-abc.atif.json"
+    path.write_text(json.dumps(payload), encoding="utf-8")
+    return path
+
+
+def _step(**metrics: int) -> dict[str, Any]:
+    return {"source": "agent", "message": "", "metrics": metrics}
+
+
+def test_final_metrics_totals_are_projected_onto_the_token_keys(tmp_path: Path) -> None:
+    path = _atif(
+        tmp_path,
+        {
+            "schema_version": "ATIF-v1.7",
+            "steps": [_step(prompt_tokens=1)],
+            "final_metrics": {
+                "total_prompt_tokens": 358,
+                "total_completion_tokens": 19324,
+                "total_cached_tokens": 3984621,
+            },
+        },
+    )
+    assert fabric_runtime._atif_token_metadata(path) == {
+        "prompt_tokens": 358,
+        "completion_tokens": 19324,
+        "cache_read_tokens": 3984621,
+    }
+
+
+def test_steps_are_summed_when_the_trajectory_has_no_aggregate_block(tmp_path: Path) -> None:
+    path = _atif(
+        tmp_path,
+        {
+            "schema_version": "ATIF-v1.7",
+            "steps": [_step(prompt_tokens=100, completion_tokens=10), _step(prompt_tokens=58, cached_tokens=7)],
+        },
+    )
+    assert fabric_runtime._atif_token_metadata(path) == {
+        "prompt_tokens": 158,
+        "completion_tokens": 10,
+        "cache_read_tokens": 7,
+    }
+
+
+def test_a_partial_aggregate_does_not_suppress_the_fields_only_the_steps_report(tmp_path: Path) -> None:
+    # Every final_metrics field is optional, so a block reporting only steps/cost validates cleanly.
+    # Resolving per field is what keeps a mid-run flush (the timeout path) from publishing nothing.
+    path = _atif(
+        tmp_path,
+        {
+            "schema_version": "ATIF-v1.7",
+            "steps": [_step(prompt_tokens=100, completion_tokens=10), _step(prompt_tokens=58)],
+            "final_metrics": {"total_steps": 2, "total_cost_usd": 0.12},
+        },
+    )
+    assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 158, "completion_tokens": 10}
+
+
+def test_a_reported_total_wins_over_the_step_sum_for_that_field_alone(tmp_path: Path) -> None:
+    # The aggregate is authoritative where it speaks; the steps fill only the fields it omits.
+    path = _atif(
+        tmp_path,
+        {
+            "schema_version": "ATIF-v1.7",
+            "steps": [_step(prompt_tokens=1, completion_tokens=10)],
+            "final_metrics": {"total_prompt_tokens": 358},
+        },
+    )
+    assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 358, "completion_tokens": 10}
+
+
+def test_an_unvalidatable_aggregate_still_falls_back_to_the_steps(tmp_path: Path) -> None:
+    path = _atif(
+        tmp_path,
+        {
+            "schema_version": "ATIF-v1.7",
+            "steps": [_step(prompt_tokens=158)],
+            "final_metrics": {"total_prompt_tokens": "not-an-int"},
+        },
+    )
+    assert fabric_runtime._atif_token_metadata(path) == {"prompt_tokens": 158}
+
+
+def test_a_trajectory_with_no_counts_anywhere_records_nothing(tmp_path: Path) -> None:
+    path = _atif(tmp_path, {"schema_version": "ATIF-v1.7", "steps": [_step()], "final_metrics": {}})
+    assert fabric_runtime._atif_token_metadata(path) == {}
+
+
+@pytest.mark.parametrize("payload", ["[]", "{ not json", '"a string"'])
+def test_an_unreadable_trajectory_records_nothing_rather_than_failing(tmp_path: Path, payload: str) -> None:
+    path = tmp_path / "trajectory-abc.atif.json"
+    path.write_text(payload, encoding="utf-8")
+    assert fabric_runtime._atif_token_metadata(path) == {}
+
+
+def test_a_missing_trajectory_records_nothing(tmp_path: Path) -> None:
+    assert fabric_runtime._atif_token_metadata(None) == {}
+    assert fabric_runtime._atif_token_metadata(tmp_path / "absent.atif.json") == {}

plugins/nemo-evaluator/src/nemo_evaluator/intake/row_adapter.py

@@ -18,7 +18,10 @@
 
 import hashlib
 import json
+import logging
+from collections.abc import Mapping
 from datetime import datetime
+from typing import Any
 
 from nemo_evaluator_sdk.agent_eval.results import AgentEvalResult, AgentEvalSummary, RunMetadata
 from nemo_evaluator_sdk.agent_eval.scores import (
@@ -32,6 +35,8 @@
 from nemo_evaluator_sdk.values.multi_metric_results import BenchmarkEvaluationResult
 from nemo_evaluator_sdk.values.results import EvaluationResult, RowScore
 
+logger = logging.getLogger(__name__)
+
 #: Key ``sample`` carries when generation itself failed, rather than the metric.
 _INFERENCE_ERROR = "inference_error"
 
@@ -80,6 +85,61 @@ def _output(row: RowScore) -> AgentOutput | None:
     return AgentOutput(output_text=output_text, response=response)
 
 
+def _first_int(usage: Mapping[str, Any], *keys: str) -> int | None:
+    """First key in ``keys`` holding a token count, or ``None``.
+
+    A count is a non-negative int that is not a bool. Negatives are rejected because they are used
+    as an unknown-value sentinel rather than a measurement, and nothing downstream would catch one:
+    Intake's ``total_prompt_tokens`` is an unconstrained ``int | None``, so a negative would be
+    summed into the evaluation rollup and shown as a real total.
+    """
+    for key in keys:
+        value = usage.get(key)
+        if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
+            return value
+    return None
+
+
+def _token_metadata(row: RowScore) -> dict[str, int]:
+    """Project the generation response's ``usage`` block onto the trial-metadata token keys.
+
+    A row's generation response is the only place its token usage survives: ``row.requests`` mixes
+    the generation call with each metric's judge calls, so summing that would credit judge tokens to
+    the agent. ``total_tokens`` is left unset — Intake recomputes it from the parts.
+
+    Both usage schemas a target can return are read, OpenAI's first: a ``GenericAgent`` points at an
+    arbitrary URL, so the response is whatever that endpoint emits, and an Anthropic-shaped block
+    would otherwise be dropped whole. Note the two disagree on whether cache reads are already
+    counted in the prompt total (OpenAI includes them, Anthropic does not); the values are recorded
+    as reported rather than reconciled, since nothing downstream adds them together.
+
+    No key list covers every provider, and a renamed key would fail the same silent way this
+    function exists to fix, so a usage block that yields nothing is logged with the keys it actually
+    carried. An unrecognized schema is then a log line naming what to add, not an unexplained blank.
+    """
+    response = row.sample.get("response")
+    usage = response.get("usage") if isinstance(response, Mapping) else None
+    if not isinstance(usage, Mapping):
+        return {}
+    details = usage.get("prompt_tokens_details")
+    cache_read = _first_int(details, "cached_tokens") if isinstance(details, Mapping) else None
+    if cache_read is None:
+        cache_read = _first_int(usage, "cache_read_input_tokens")
+    captured = {
+        "prompt_tokens": _first_int(usage, "prompt_tokens", "input_tokens"),
+        "completion_tokens": _first_int(usage, "completion_tokens", "output_tokens"),
+        "cache_read_tokens": cache_read,
+        "cache_creation_tokens": _first_int(usage, "cache_creation_input_tokens"),
+    }
+    recorded = {key: value for key, value in captured.items() if value is not None}
+    if not recorded:
+        logger.warning(
+            "No token counts recognized in the generation response's usage block; keys present: %s",
+            sorted(str(key) for key in usage),
+        )
+    return recorded
+
+
 def _scores(row: RowScore, *, run_id: str, task_id: str, trial_id: str) -> list[AgentEvalTaskScore]:
     """One score per metric key on the row; ``metrics`` values are already ``MetricOutput``."""
     errors = row.metric_errors or {}
@@ -152,6 +212,7 @@ def row_result_to_agent_eval_result(
                 task_id=task_id,
                 status=AgentEvalTrialStatus.COMPLETED if output is not None else AgentEvalTrialStatus.FAILED,
                 output=output,
+                metadata=_token_metadata(row),
             )
         )
         scores.extend(_scores(row, run_id=run_id, task_id=task_id, trial_id=trial_id))