fix: address review findings (iteration 1) - F003, F007, F008, F009, F016

- F009: fix I001 import order in playbook_extractor
- F003: drop always-None session_id kwarg from learning billing events (no session_id source on generation path)
- F016: rename _last_resumable_trace -> _last_resumable_token_totals
- F007: add unit tests for token_accounting + count_input_tokens
- F008: add per-endpoint applied-learnings metering tests (4 endpoints x 3 scenarios)

Author: yilu331

reflexio/server/services/base_generation_service.py

@@ -344,6 +344,13 @@ def _record_billing_learning_events(
             platform_llm = platform_llm_from_config(config)
             ctx = self._usage_context()
 
+            # session_id is intentionally not passed: the generation path has no
+            # session_id source. _usage_context() never includes it, and neither
+            # the Profile/Playbook service configs nor their requests carry a
+            # session_id (unlike the Application-line path in server/api.py, which
+            # reads it from the publish payload). Learning events therefore meter
+            # without session attribution by design.
+
             # ② Learning — value: learnings generated (helper no-ops on count <= 0).
             record_learnings_generated(
                 org_id=ctx["org_id"],
@@ -352,12 +359,13 @@ def _record_billing_learning_events(
                 platform_storage=None,
                 pipeline=ctx.get("pipeline"),
                 request_id=ctx.get("request_id"),
-                session_id=ctx.get("session_id"),
             )
 
             # ② Learning — cost: input-anchored extraction tokens + real provider tokens.
             totals = self._last_token_totals or RunTokenTotals()
-            billing_input_tokens = count_input_tokens(self._extraction_input_text(prepared))
+            billing_input_tokens = count_input_tokens(
+                self._extraction_input_text(prepared)
+            )
             record_extraction_tokens(
                 org_id=ctx["org_id"],
                 billing_input_tokens=billing_input_tokens,
@@ -367,7 +375,6 @@ def _record_billing_learning_events(
                 platform_storage=None,
                 pipeline=ctx.get("pipeline"),
                 request_id=ctx.get("request_id"),
-                session_id=ctx.get("session_id"),
             )
         except Exception:
             logger.warning(

reflexio/server/services/playbook/playbook_extractor.py

@@ -10,8 +10,8 @@
 from reflexio.server.api_endpoints.request_context import RequestContext
 from reflexio.server.llm.litellm_client import LiteLLMClient
 from reflexio.server.llm.model_defaults import ModelRole, resolve_model_name
-from reflexio.server.services.extraction.outcome import ExtractionOutcome
 from reflexio.server.llm.token_accounting import RunTokenTotals, sum_trace_tokens
+from reflexio.server.services.extraction.outcome import ExtractionOutcome
 from reflexio.server.services.extraction.resumable_agent import (
     run_resumable_extraction_agent,
 )
@@ -80,7 +80,7 @@ def __init__(
         self.service_config: PlaybookGenerationServiceConfig = service_config
         self.agent_context: str = agent_context
         self._last_resumable_run_id: str | None = None
-        self._last_resumable_trace: RunTokenTotals | None = None
+        self._last_resumable_token_totals: RunTokenTotals | None = None
 
         # Get LLM config overrides from configuration
         config = self.request_context.configurator.get_config()
@@ -225,7 +225,7 @@ def run(self) -> list[UserPlaybook] | ExtractionOutcome[UserPlaybook]:
             return ExtractionOutcome.completed(
                 user_playbooks,
                 run_id=self._last_resumable_run_id,
-                token_totals=self._last_resumable_trace,
+                token_totals=self._last_resumable_token_totals,
             )
         return user_playbooks
 
@@ -321,7 +321,7 @@ def extract_playbook_entries(
             log_label="Playbook extraction",
         )
         self._last_resumable_run_id = result.run_id
-        self._last_resumable_trace = sum_trace_tokens(result.trace)
+        self._last_resumable_token_totals = sum_trace_tokens(result.trace)
         if not isinstance(result.output, StructuredPlaybookList):
             logger.warning(
                 "Playbook extraction did not finish: %s",

reflexio/server/services/profile/profile_extractor.py

@@ -84,7 +84,7 @@ def __init__(
         self.service_config: ProfileGenerationServiceConfig = service_config
         self.agent_context = agent_context
         self._last_resumable_run_id: str | None = None
-        self._last_resumable_trace: RunTokenTotals | None = None
+        self._last_resumable_token_totals: RunTokenTotals | None = None
 
         # Get LLM config overrides from configuration
         config = self.request_context.configurator.get_config()
@@ -245,7 +245,8 @@ def run(self) -> list[UserProfile] | ExtractionOutcome[UserProfile] | None:
                 # `or` is safe: a dataclass instance (even RunTokenTotals(0,0)) is
                 # always truthy, so a real-but-zero total is never overwritten by
                 # the fallback trace.
-                token_totals=raw_profiles.token_totals or self._last_resumable_trace,
+                token_totals=raw_profiles.token_totals
+                or self._last_resumable_token_totals,
             )
         user_profiles = self._convert_raw_to_user_profiles(
             raw_profiles=raw_profiles or [],
@@ -265,7 +266,7 @@ def run(self) -> list[UserProfile] | ExtractionOutcome[UserProfile] | None:
             return ExtractionOutcome.completed(
                 user_profiles,
                 run_id=self._last_resumable_run_id,
-                token_totals=self._last_resumable_trace,
+                token_totals=self._last_resumable_token_totals,
             )
         return user_profiles or None
 
@@ -400,7 +401,7 @@ def _generate_raw_updates_from_sessions(
             log_label="Profile extraction",
         )
         self._last_resumable_run_id = result.run_id
-        self._last_resumable_trace = sum_trace_tokens(result.trace)
+        self._last_resumable_token_totals = sum_trace_tokens(result.trace)
         if not isinstance(result.output, StructuredProfilesOutput):
             logger.warning(
                 "Profile extraction did not finish: %s", result.finished_reason

tests/server/api_endpoints/test_applied_learnings_metering.py

@@ -8,9 +8,14 @@
 from contextlib import contextmanager
 from unittest.mock import MagicMock, patch
 
+import pytest
 from fastapi.testclient import TestClient
 
-from reflexio.models.api_schema.ui.entities import AgentPlaybookView, ProfileView
+from reflexio.models.api_schema.ui.entities import (
+    AgentPlaybookView,
+    ProfileView,
+    UserPlaybookView,
+)
 from reflexio.server.api import create_app
 from reflexio.server.usage_metrics import UsageEvent, configure_usage_event_recorder
 
@@ -29,6 +34,10 @@ def _make_agent_playbook_view() -> AgentPlaybookView:
     return AgentPlaybookView(agent_version="v1", content="content")
 
 
+def _make_user_playbook_view() -> UserPlaybookView:
+    return UserPlaybookView(agent_version="v1", request_id="r1", content="content")
+
+
 def _client(caller_type: str) -> TestClient:
     app = create_app(get_org_id=lambda: "test-org", get_caller_type=lambda: caller_type)
     return TestClient(app, raise_server_exceptions=False)
@@ -90,7 +99,9 @@ def test_production_agent_search_meters_surfaced_count() -> None:
 
     applied = [e for e in events if e.event_name == "learning_applied"]
     assert len(applied) == 1
-    assert applied[0].count_value == 3  # 2 profiles + 1 agent_playbook + 0 user_playbooks
+    assert (
+        applied[0].count_value == 3
+    )  # 2 profiles + 1 agent_playbook + 0 user_playbooks
     assert applied[0].caller_type == "production_agent"
 
 
@@ -99,7 +110,9 @@ def test_dashboard_search_meters_nothing() -> None:
     events = _capture()
     try:
         with _patch_unified_search([_make_profile_view()], [], []):
-            _client("dashboard").post("/api/search", json={"query": "x", "user_id": "u1"})
+            _client("dashboard").post(
+                "/api/search", json={"query": "x", "user_id": "u1"}
+            )
     finally:
         configure_usage_event_recorder(None)
 
@@ -140,11 +153,13 @@ def test_metering_failure_does_not_break_search_response() -> None:
         mock_response.user_playbooks = []
         mock_reflexio_search.unified_search.return_value = mock_response
         # Make get_config raise so metering blows up after the search completes.
-        mock_reflexio_search.request_context.configurator.get_config.side_effect = RuntimeError(
-            "boom"
+        mock_reflexio_search.request_context.configurator.get_config.side_effect = (
+            RuntimeError("boom")
         )
 
-        with patch("reflexio.server.api.get_reflexio", return_value=mock_reflexio_search):
+        with patch(
+            "reflexio.server.api.get_reflexio", return_value=mock_reflexio_search
+        ):
             resp = _client("production_agent").post(
                 "/api/search", json={"query": "x", "user_id": "u1"}
             )
@@ -154,3 +169,142 @@ def test_metering_failure_does_not_break_search_response() -> None:
 
     # Metering failed silently — no learning_applied event should have been emitted.
     assert [e for e in events if e.event_name == "learning_applied"] == []
+
+
+# --- Per-endpoint metering (the four non-unified routes) -----------------------
+#
+# Each of these endpoints calls a distinct service method on get_reflexio and
+# derives surfaced_count from a distinct response list attribute. The cases below
+# exercise the real endpoint handler + view conversion + _meter_applied_learnings
+# wiring for each, asserting the emitted surfaced_count matches that route's shape.
+
+
+@contextmanager
+def _patch_service_method(method_name: str, response_attr: str, items: list):
+    """Patch get_reflexio so ``method_name`` returns a canned service response.
+
+    The response carries ``items`` on ``response_attr`` (e.g. ``user_profiles``)
+    so the endpoint's view conversion and surfaced_count computation run for real.
+    get_config() returns None so platform_llm_from_config(None) is True without
+    iterating a MagicMock.
+    """
+    mock_reflexio = MagicMock()
+    mock_response = MagicMock()
+    mock_response.success = True
+    mock_response.msg = "OK"
+    setattr(mock_response, response_attr, items)
+    getattr(mock_reflexio, method_name).return_value = mock_response
+    mock_reflexio.request_context.configurator.get_config.return_value = None
+
+    with patch("reflexio.server.api.get_reflexio", return_value=mock_reflexio):
+        yield
+
+
+# (path, payload, service method, response attribute, surfaced item factory)
+_ENDPOINT_CASES = [
+    pytest.param(
+        "/api/search_profiles",
+        {"user_id": "u1", "query": "x"},
+        "search_user_profiles",
+        "user_profiles",
+        _make_profile_view,
+        id="search_profiles",
+    ),
+    pytest.param(
+        "/api/search_user_playbooks",
+        {"query": "x"},
+        "search_user_playbooks",
+        "user_playbooks",
+        _make_user_playbook_view,
+        id="search_user_playbooks",
+    ),
+    pytest.param(
+        "/api/search_agent_playbooks",
+        {"query": "x"},
+        "search_agent_playbooks",
+        "agent_playbooks",
+        _make_agent_playbook_view,
+        id="search_agent_playbooks",
+    ),
+    pytest.param(
+        "/api/get_agent_playbooks",
+        {},
+        "get_agent_playbooks",
+        "agent_playbooks",
+        _make_agent_playbook_view,
+        id="get_agent_playbooks",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    ("path", "payload", "method_name", "response_attr", "make_item"),
+    _ENDPOINT_CASES,
+)
+def test_production_agent_per_endpoint_meters_surfaced_count(
+    path: str,
+    payload: dict,
+    method_name: str,
+    response_attr: str,
+    make_item,
+) -> None:
+    """Each non-unified route emits one learning_applied event with its own count."""
+    events = _capture()
+    items = [make_item(), make_item()]
+    try:
+        with _patch_service_method(method_name, response_attr, items):
+            resp = _client("production_agent").post(path, json=payload)
+        assert resp.status_code == 200
+    finally:
+        configure_usage_event_recorder(None)
+
+    applied = [e for e in events if e.event_name == "learning_applied"]
+    assert len(applied) == 1
+    assert applied[0].count_value == 2  # len(items) for this endpoint's response shape
+    assert applied[0].caller_type == "production_agent"
+
+
+@pytest.mark.parametrize(
+    ("path", "payload", "method_name", "response_attr", "make_item"),
+    _ENDPOINT_CASES,
+)
+def test_dashboard_per_endpoint_meters_nothing(
+    path: str,
+    payload: dict,
+    method_name: str,
+    response_attr: str,
+    make_item,
+) -> None:
+    """A dashboard caller never meters, regardless of the route or result size."""
+    events = _capture()
+    try:
+        with _patch_service_method(method_name, response_attr, [make_item()]):
+            resp = _client("dashboard").post(path, json=payload)
+        assert resp.status_code == 200
+    finally:
+        configure_usage_event_recorder(None)
+
+    assert [e for e in events if e.event_name == "learning_applied"] == []
+
+
+@pytest.mark.parametrize(
+    ("path", "payload", "method_name", "response_attr", "make_item"),
+    _ENDPOINT_CASES,
+)
+def test_empty_result_per_endpoint_meters_nothing(
+    path: str,
+    payload: dict,
+    method_name: str,
+    response_attr: str,
+    make_item,
+) -> None:
+    """A production-agent call surfacing zero results meters nothing on any route."""
+    events = _capture()
+    try:
+        with _patch_service_method(method_name, response_attr, []):
+            resp = _client("production_agent").post(path, json=payload)
+        assert resp.status_code == 200
+    finally:
+        configure_usage_event_recorder(None)
+
+    assert [e for e in events if e.event_name == "learning_applied"] == []