Merge pull request #1 from jinliyl/fix/statistics-counter-contract

fix(evaluation): simplify job counters and make token accounting atomic

benchmark/beam/run.py

@@ -737,7 +737,7 @@ def main(  # pylint: disable=too-many-statements
         print(f"    {'OVERALL':<40s}: {overall:.3f}  binary={binary_overall:.3f}  ({len(all_scores)} Qs)")
         tool_call_mean, tool_call_std = _mean_and_std(all_tool_call_totals)
         print(f"    Tool calls/query: mean={tool_call_mean:.2f} std={tool_call_std:.2f}")
-        print("    Bench tokens/query:")
+        print("    Bench reported tokens/query:")
         for metric in _TOKEN_USAGE_METRICS:
             values = [usage[metric] for usage in all_token_usages if usage.get(metric) is not None]
             if values:

benchmark/longmemeval/run.py

@@ -735,7 +735,7 @@ def _print_summary(results: list[dict], start_time: float) -> None:
     tool_call_mean, tool_call_std = _mean_and_std(tool_call_totals)
     print(f"  Tool calls/query: mean={tool_call_mean:.2f} std={tool_call_std:.2f}")
     token_usages = [r.get("agentic_token_usage", {}) for r in results]
-    print("  Bench tokens/query:")
+    print("  Bench reported tokens/query:")
     for metric in _TOKEN_USAGE_METRICS:
         values = [usage[metric] for usage in token_usages if usage.get(metric) is not None]
         if values:

reme/components/agent_wrapper/base_agent_wrapper.py

@@ -12,7 +12,7 @@ from ..base_component import BaseComponent
 from ..outbound_proxy import BaseOutboundProxy
 from ...enumeration import ChunkEnum, ComponentEnum
 from ...schema import StreamChunk, TokenUsage
-from ...utils import global_counter_add
+from ...utils import global_counter_add_many
 
 if TYPE_CHECKING:
     from ..job.base_job import BaseJob
@@ -236,12 +236,15 @@ class BaseAgentWrapper(BaseComponent):
         counters = getattr(self.app_context, "metadata", None)
         if not isinstance(counters, dict):
             return
-        for field in ("input_tokens", "output_tokens", "total_tokens"):
-            global_counter_add(
-                counters,
-                [self.TOKEN_COUNTER_PREFIX, self.name, field],
-                getattr(usage, field),
-            )
+        prefix = (self.TOKEN_COUNTER_PREFIX, self.name)
+        global_counter_add_many(
+            counters,
+            {
+                (*prefix, "input_tokens"): usage.input_tokens,
+                (*prefix, "output_tokens"): usage.output_tokens,
+                (*prefix, "total_tokens"): usage.total_tokens,
+            },
+        )
 
     @abstractmethod
     async def reply(self, inputs: Any, **kwargs) -> dict:

reme/components/job/background_job.py

@@ -123,7 +123,7 @@ class BackgroundJob(BaseJob):
 
     async def __call__(self, **kwargs) -> Response:
         """Default body: run steps in order; errors propagate to supervisor."""
-        self._record_call(BackgroundJob)
+        self._record_call()
         merged = {**self.kwargs, **kwargs}
         context = RuntimeContext(stop_event=self._stop_event, **merged)
         for step in self._build_steps():

reme/components/job/base_job.py

@@ -58,27 +58,15 @@ class BaseJob(BaseComponent):
         # dict(params) copies kwargs so steps cannot mutate the shared spec.
         return [step_cls(**dict(params)) for step_cls, params in self.step_specs]
 
-    def _counter_key(self, entry_class: type["BaseJob"]) -> list[str]:
-        """Build this job's application-lifetime call counter key.
-
-        The counter path groups calls by the built-in job implementation that
-        accepted the call, then preserves any project-specific subclasses
-        between that implementation and the concrete job class.
-        """
-        mro = type(self).mro()
-        entry_index = mro.index(entry_class)
-        subclass_path = [cls.__name__ for cls in mro[:entry_index]]
-        return ["__job_counter", entry_class.__name__, *subclass_path, self.name]
-
-    def _record_call(self, entry_class: type["BaseJob"]) -> None:
+    def _record_call(self) -> None:
         """Increment this job's application-lifetime call counter."""
         metadata = getattr(self.app_context, "metadata", None)
         if isinstance(metadata, dict):
-            global_counter_inc(metadata, self._counter_key(entry_class))
+            global_counter_inc(metadata, ["__job_counter", self.name])
 
     async def __call__(self, **kwargs) -> Response:
         """Run all steps in order, capturing any failure into the response."""
-        self._record_call(BaseJob)
+        self._record_call()
         merged = {**self.kwargs, **kwargs}
         context = RuntimeContext(**merged)
         try:

reme/components/job/cron_job.py

@@ -47,7 +47,7 @@ class CronJob(BackgroundJob):
             if self._stop_event.is_set():
                 break
             try:
-                self._record_call(CronJob)
+                self._record_call()
                 await self._execute_steps()
             except Exception as exc:
                 self.logger.exception(f"Cron job '{self.name}' failed: {exc}")

reme/components/job/stream_job.py

@@ -12,7 +12,7 @@ class StreamJob(BaseJob):
 
     async def __call__(self, **kwargs) -> None:
         """Run steps; emit failures as ERROR chunks, then a terminal DONE marker."""
-        self._record_call(StreamJob)
+        self._record_call()
         merged = {**self.kwargs, **kwargs}
         context = RuntimeContext(**merged)
         try:

reme/schema/token_usage.py

@@ -6,10 +6,12 @@ from pydantic import BaseModel, Field, model_validator
 
 
 class TokenUsage(BaseModel):
-    """Token usage for one completed agent invocation.
+    """Portable token usage reported for one completed agent invocation.
 
-    ``input_tokens`` and ``output_tokens`` are the complete usage reported by
-    one backend invocation. ``total_tokens`` is their derived sum.
+    Only the provider's top-level input and output counters are retained.
+    Provider-specific cache and reasoning breakdowns are intentionally
+    excluded, so values may not share identical billing semantics across
+    providers. ``total_tokens`` is always their derived sum.
     """
 
     input_tokens: int = Field(default=0, ge=0)
@@ -26,7 +28,7 @@ class TokenUsage(BaseModel):
         cls,
         usage: Any,
     ) -> "TokenUsage":
-        """Normalize a provider usage object or mapping to input/output totals."""
+        """Keep only a provider's portable top-level input/output counters."""
 
         def get(*names: str) -> int | None:
             for name in names:

reme/utils/__init__.py

@@ -15,7 +15,13 @@ from .service_utils import find_reme, locate_reme, precheck_start, cli_find_reme
 from .similarity_utils import cosine_similarity, batch_cosine_similarity
 from .token_utils import estimate_token_count
 from .agent_state_io import AsStateHandler
-from .counter import global_counter_add, global_counter_get, global_counter_get_all, global_counter_inc
+from .counter import (
+    global_counter_add,
+    global_counter_add_many,
+    global_counter_get,
+    global_counter_get_all,
+    global_counter_inc,
+)
 
 __all__ = [
     "hash_text",
@@ -39,6 +45,7 @@ __all__ = [
     "estimate_token_count",
     "AsStateHandler",
     "global_counter_add",
+    "global_counter_add_many",
     "global_counter_get",
     "global_counter_get_all",
     "global_counter_inc",

reme/utils/counter.py

@@ -2,10 +2,70 @@
 
 import copy
 import threading
+from collections.abc import Mapping
 from typing import Any
 
 COUNTER_TREE_KEY = "_counter_tree"
 COUNTER_LOCK_KEY = "_counter_tree_lock"
+_COUNTER_INIT_LOCK = threading.Lock()
+
+
+def _get_counter_lock(metadata: dict[str, Any]) -> Any:
+    """Return the metadata-scoped lock, creating it once when needed."""
+    lock = metadata.get(COUNTER_LOCK_KEY)
+    if lock is not None:
+        return lock
+
+    # Two threads may reach the first counter operation concurrently. Guard
+    # initialization so they cannot install and then use different locks.
+    with _COUNTER_INIT_LOCK:
+        lock = metadata.get(COUNTER_LOCK_KEY)
+        if lock is None:
+            lock = threading.Lock()
+            metadata[COUNTER_LOCK_KEY] = lock
+        return lock
+
+
+def global_counter_add_many(
+    metadata: dict[str, Any],
+    updates: Mapping[tuple[str, ...], int],
+) -> dict[tuple[str, ...], int]:
+    """Atomically fetch-and-add multiple counter paths.
+
+    All paths are validated before the counter tree is mutated. The returned
+    mapping contains each path's value immediately before its increment.
+    """
+    normalized = dict(updates)
+    for path, value in normalized.items():
+        if not isinstance(path, tuple) or not all(isinstance(part, str) for part in path):
+            raise TypeError("counter paths must be tuples of strings")
+        if not isinstance(value, int):
+            raise TypeError("counter increments must be integers")
+    if not normalized:
+        return {}
+
+    lock = _get_counter_lock(metadata)
+    with lock:
+        tree = metadata.get(COUNTER_TREE_KEY)
+        if tree is None:
+            tree = {"value": 0, "children": {}}
+            metadata[COUNTER_TREE_KEY] = tree
+
+        nodes: dict[tuple[str, ...], dict[str, Any]] = {}
+        for path in normalized:
+            node = tree
+            for part in path:
+                child = node["children"].get(part)
+                if child is None:
+                    child = {"value": 0, "children": {}}
+                    node["children"][part] = child
+                node = child
+            nodes[path] = node
+
+        previous = {path: node["value"] for path, node in nodes.items()}
+        for path, value in normalized.items():
+            nodes[path]["value"] += value
+        return previous
 
 
 def global_counter_add(metadata: dict[str, Any], key: list[str], val: int) -> int:
@@ -23,28 +83,8 @@ def global_counter_add(metadata: dict[str, Any], key: list[str], val: int) -> in
     If they are missing they are created lazily so the function is safe to
     call with a plain ``dict``.
     """
-    lock = metadata.get(COUNTER_LOCK_KEY)
-    if lock is None:
-        lock = threading.Lock()
-        metadata[COUNTER_LOCK_KEY] = lock
-
-    with lock:
-        tree = metadata.get(COUNTER_TREE_KEY)
-        if tree is None:
-            tree = {"value": 0, "children": {}}
-            metadata[COUNTER_TREE_KEY] = tree
-
-        node: dict[str, Any] = tree
-        for part in key:
-            assert isinstance(part, str)
-            tmp = node["children"].get(part, None)
-            if tmp is None:
-                tmp = {"value": 0, "children": {}}
-                node["children"][part] = tmp
-            node = tmp
-        res = node["value"]
-        node["value"] = res + val
-    return res
+    path = tuple(key)
+    return global_counter_add_many(metadata, {path: val})[path]
 
 
 def global_counter_inc(metadata: dict[str, Any], key: list[str]) -> int:
@@ -63,10 +103,7 @@ def global_counter_get(metadata: dict[str, Any], key: list[str]) -> int:
     path that does not exist yet is reported as 0, matching the value the
     node would hold right before its first increment.
     """
-    lock = metadata.get(COUNTER_LOCK_KEY)
-    if lock is None:
-        lock = threading.Lock()
-        metadata[COUNTER_LOCK_KEY] = lock
+    lock = _get_counter_lock(metadata)
 
     with lock:
         tree = metadata.get(COUNTER_TREE_KEY)
@@ -91,10 +128,7 @@ def global_counter_get_all(metadata: dict[str, Any], key: list[str]) -> dict[str
     Returns ``None`` when the tree or any part of ``key`` does not exist.
     An empty ``key`` returns a copy of the whole tree.
     """
-    lock = metadata.get(COUNTER_LOCK_KEY)
-    if lock is None:
-        lock = threading.Lock()
-        metadata[COUNTER_LOCK_KEY] = lock
+    lock = _get_counter_lock(metadata)
 
     with lock:
         tree = metadata.get(COUNTER_TREE_KEY)

reme/utils/evaluation_interface.py

@@ -9,14 +9,12 @@ work sharing its Application instance.
 
 from typing import TYPE_CHECKING
 
-from ..components.job import BackgroundJob, BaseJob, CronJob, StreamJob
 from .counter import global_counter_get, global_counter_get_all
 
 if TYPE_CHECKING:
     from ..components.application_context import ApplicationContext
 
 
-_JOB_ENTRY_CLASSES = {CronJob, StreamJob, BackgroundJob, BaseJob}
 _TOKEN_METRICS = (
     "input_tokens",
     "output_tokens",
@@ -31,15 +29,9 @@ def check_job_count(job_name: str, app_context: "ApplicationContext") -> int:
     current Application instance. Unknown job names use the same ``KeyError``
     contract as :meth:`Application.run_job`.
     """
-    job = app_context.jobs.get(job_name)
-    if job is None:
+    if job_name not in app_context.jobs:
         raise KeyError(f"Job '{job_name}' not found")
-
-    entry_class = next((cls for cls in type(job).mro() if cls in _JOB_ENTRY_CLASSES), None)
-    if entry_class is None:
-        raise TypeError(f"Job '{job_name}' does not inherit from a supported job implementation")
-    # pylint: disable-next=protected-access
-    return global_counter_get(app_context.metadata, job._counter_key(entry_class))
+    return global_counter_get(app_context.metadata, ["__job_counter", job_name])
 
 
 class JobCountTracker:

tests/unit/test_evaluation_interface.py

@@ -32,8 +32,8 @@ def test_check_job_count_reads_registered_base_job_count():
     asyncio.run(run())
 
 
-def test_check_job_count_resolves_custom_job_inheritance_path():
-    """check_job_count resolves counters recorded under subclassed job paths."""
+def test_check_job_count_reads_custom_job_by_name():
+    """check_job_count reads subclassed job counters without depending on inheritance."""
 
     async def run():
         class ProjectStreamJob(StreamJob):

tests/unit/test_job.py

@@ -144,12 +144,12 @@ def test_base_job_records_calls_by_name():
         await job()
         await job()
 
-        assert global_counter_get(app_context.metadata, ["__job_counter", "BaseJob", "search"]) == 2
+        assert global_counter_get(app_context.metadata, ["__job_counter", "search"]) == 2
 
     asyncio.run(run())
 
 
-def test_stream_job_records_concrete_subclass_path():
+def test_stream_job_subclass_records_calls_by_job_name():
     async def run():
         class ProjectStreamJob(StreamJob):
             pass
@@ -159,7 +159,7 @@ def test_stream_job_records_concrete_subclass_path():
 
         await job(stream_queue=asyncio.Queue())
 
-        assert global_counter_get(app_context.metadata, ["__job_counter", "StreamJob", "ProjectStreamJob", "chat"]) == 1
+        assert global_counter_get(app_context.metadata, ["__job_counter", "chat"]) == 1
 
     asyncio.run(run())
 
@@ -171,7 +171,7 @@ def test_background_job_records_calls_by_name():
 
         await job()
 
-        assert global_counter_get(app_context.metadata, ["__job_counter", "BackgroundJob", "watch"]) == 1
+        assert global_counter_get(app_context.metadata, ["__job_counter", "watch"]) == 1
 
     asyncio.run(run())
 
@@ -195,7 +195,7 @@ def test_cron_job_records_each_triggered_execution():
 
         await job()
 
-        assert global_counter_get(app_context.metadata, ["__job_counter", "CronJob", "nightly"]) == 1
+        assert global_counter_get(app_context.metadata, ["__job_counter", "nightly"]) == 1
 
     asyncio.run(run())
 

tests/unit/test_token_usage.py

@@ -50,6 +50,24 @@ def test_provider_usage_uses_reported_input_without_cache_adjustment():
     assert usage.total_tokens == 64
 
 
+def test_provider_usage_accepts_prompt_and_completion_aliases():
+    """Portable OpenAI-style aliases normalize to the common counters."""
+    usage = TokenUsage.from_provider({"prompt_tokens": 3, "completion_tokens": 2})
+
+    assert usage.model_dump() == {
+        "input_tokens": 3,
+        "output_tokens": 2,
+        "total_tokens": 5,
+    }
+
+
+def test_total_tokens_is_always_derived_from_input_and_output():
+    """Caller-supplied totals cannot violate the portable usage invariant."""
+    usage = TokenUsage(input_tokens=3, output_tokens=2, total_tokens=999)
+
+    assert usage.total_tokens == 5
+
+
 def test_agentscope_usage_keeps_portable_input_and_output(tmp_path):
     """AgentScope usage has the same input/output-only contract."""
     context = ApplicationContext(workspace_dir=str(tmp_path))

tests/unit/test_utils.py

@@ -9,8 +9,10 @@ import pytest
 
 from reme.utils import common_utils
 from reme.utils.counter import (
+    COUNTER_LOCK_KEY,
     COUNTER_TREE_KEY,
     global_counter_add,
+    global_counter_add_many,
     global_counter_get,
     global_counter_get_all,
     global_counter_inc,
@@ -94,6 +96,79 @@ def test_add_returns_old_value_and_adds_val():
     assert global_counter_get(metadata, ["a"]) == 16
 
 
+def test_add_many_returns_old_values_and_updates_all_paths():
+    """``add_many`` updates sibling metrics under one counter-tree lock."""
+    metadata: dict = {}
+    global_counter_add(metadata, ["usage", "input"], 2)
+
+    previous = global_counter_add_many(
+        metadata,
+        {
+            ("usage", "input"): 3,
+            ("usage", "output"): 4,
+            ("usage", "total"): 7,
+        },
+    )
+
+    assert previous == {
+        ("usage", "input"): 2,
+        ("usage", "output"): 0,
+        ("usage", "total"): 0,
+    }
+    assert global_counter_get(metadata, ["usage", "input"]) == 5
+    assert global_counter_get(metadata, ["usage", "output"]) == 4
+    assert global_counter_get(metadata, ["usage", "total"]) == 7
+
+
+def test_add_many_holds_the_counter_lock_once_for_the_batch():
+    """A multi-metric update is one critical section, not several writes."""
+
+    class CountingLock:
+        """Lock test double that counts entered critical sections."""
+
+        def __init__(self):
+            self.lock = threading.Lock()
+            self.entries = 0
+
+        def __enter__(self):
+            self.lock.acquire()
+            self.entries += 1
+            return self
+
+        def __exit__(self, *_args):
+            self.lock.release()
+
+    lock = CountingLock()
+    metadata = {COUNTER_LOCK_KEY: lock}
+
+    global_counter_add_many(
+        metadata,
+        {
+            ("usage", "input"): 3,
+            ("usage", "output"): 4,
+            ("usage", "total"): 7,
+        },
+    )
+
+    assert lock.entries == 1
+
+
+def test_add_many_validates_every_update_before_mutating():
+    """One invalid update leaves every valid counter unchanged."""
+    metadata: dict = {}
+
+    with pytest.raises(TypeError, match="increments must be integers"):
+        global_counter_add_many(
+            metadata,
+            {
+                ("usage", "input"): 3,
+                ("usage", "output"): "invalid",
+            },
+        )
+
+    assert COUNTER_TREE_KEY not in metadata
+
+
 def test_counters_are_isolated_by_key_path():
     """Sibling and nested keys, plus the root, hold independent counters."""
     metadata: dict = {}