diff --git a/backend/Dockerfile.scheduled-runs b/backend/Dockerfile.scheduled-runs index d8b486ba9..10d99fa17 100644 --- a/backend/Dockerfile.scheduled-runs +++ b/backend/Dockerfile.scheduled-runs @@ -33,6 +33,9 @@ RUN pip install --no-cache-dir -r /tmp/requirements.txt # content-hash tag notices changes. COPY backend/src/apis/shared/__init__.py ${LAMBDA_TASK_ROOT}/apis/shared/__init__.py COPY backend/src/apis/shared/errors.py ${LAMBDA_TASK_ROOT}/apis/shared/errors.py +# Leaf module (imports only os); reached via apis/shared/sessions/metadata.py, +# which reads the COST_DIAGNOSTICS_ENABLED kill switch at aggregate-bump time. +COPY backend/src/apis/shared/feature_flags.py ${LAMBDA_TASK_ROOT}/apis/shared/feature_flags.py COPY backend/src/apis/shared/harness/ ${LAMBDA_TASK_ROOT}/apis/shared/harness/ COPY backend/src/apis/shared/scheduled_prompts/ ${LAMBDA_TASK_ROOT}/apis/shared/scheduled_prompts/ COPY backend/src/apis/shared/sessions_bff/ ${LAMBDA_TASK_ROOT}/apis/shared/sessions_bff/ diff --git a/backend/src/agents/main_agent/base_agent.py b/backend/src/agents/main_agent/base_agent.py index 79043d3fe..281d45c27 100644 --- a/backend/src/agents/main_agent/base_agent.py +++ b/backend/src/agents/main_agent/base_agent.py @@ -14,6 +14,7 @@ from agents.main_agent.session import SessionFactory from agents.main_agent.session.hooks import ( AgentStatusHook, + ToolCensusHook, DisplayTextHook, SteeringHook, StopHook, @@ -340,6 +341,15 @@ def _create_hooks(self) -> List: self.agent_status_hook = AgentStatusHook() hooks.append(self.agent_status_hook) + # Content-free tool census (tool name → calls/errors per model call). + # Held on the wrapper so the stream coordinator can read each call's + # tally at turn end and persist it on that call's cost row for the + # admin session profile. Non-drained, per-turn only. Registered + # unconditionally; the callbacks return immediately when + # COST_DIAGNOSTICS_ENABLED=false. + self.tool_census_hook = ToolCensusHook() + hooks.append(self.tool_census_hook) + # Per-model-call prompt-cache prefix fingerprints (toolConfig / # system prompt / history hashes). Best-effort; the stream # coordinator persists them on each call's metadata row so avoidable diff --git a/backend/src/agents/main_agent/session/hooks/__init__.py b/backend/src/agents/main_agent/session/hooks/__init__.py index 88efadd8e..1080adc59 100644 --- a/backend/src/agents/main_agent/session/hooks/__init__.py +++ b/backend/src/agents/main_agent/session/hooks/__init__.py @@ -8,6 +8,7 @@ from agents.main_agent.session.hooks.steering import SteeringHook from agents.main_agent.session.hooks.stop import StopHook from agents.main_agent.session.hooks.tool_approval import MCPExternalApprovalHook +from agents.main_agent.session.hooks.tool_census import ToolCensusHook __all__ = [ "AgentStatusHook", @@ -18,4 +19,5 @@ "SteeringHook", "StopHook", "MCPExternalApprovalHook", + "ToolCensusHook", ] diff --git a/backend/src/agents/main_agent/session/hooks/tool_census.py b/backend/src/agents/main_agent/session/hooks/tool_census.py new file mode 100644 index 000000000..a60d87211 --- /dev/null +++ b/backend/src/agents/main_agent/session/hooks/tool_census.py @@ -0,0 +1,127 @@ +"""Hook that tallies tool calls per model call, for the content-free cost profile. + +The admin cost console can say what a conversation *cost* down to the model +call, but until now nothing recorded what the user was *doing* in it — which +tools ran, how often, how often they failed. That signal lives only inside +message content, which the console must never read. This hook counts it out +of band: tool name → ``{calls, errors}``, attributed to the model call whose +output requested the tools, and the stream coordinator writes the tally onto +that call's ``C#`` cost row as ``toolCalls`` (the same way ``turnAgentId`` +and ``prefixFingerprints`` ride that row). + +Tool names are catalog ids and MCP tool names — structural, not content. +Nothing here reads a tool's input or result beyond ``status``. + +Attribution uses the same cycle counter ``AgentStatusHook`` uses: ``cycle`` +increments on every ``BeforeModelCallEvent``, and a tool that runs during +cycle *N* was requested by model call *N* — the coordinator's ``call_index`` +``N-1`` (0-based). ``tally_for_call`` does that translation so the caller +does not have to. + +Deliberately **non-drained** (unlike ``AgentStatusHook``, whose lists the +coordinator empties mid-turn): the tally is read once at turn end, per call, +and reset at the next turn's ``BeforeInvocationEvent``. Per-turn state only — +never anything a later turn could inherit (CLAUDE.md: never cache session +state on an agent instance). + +Costs nothing against the model: it runs on hook events the loop already +crosses and writes to an in-process dict. Gated by +``COST_DIAGNOSTICS_ENABLED`` (default on, ``=false`` kill switch); while off +every callback returns immediately and ``tally_for_call`` is always ``None``, +so the row simply has no ``toolCalls`` and the profile reads "not tracked". +""" + +from __future__ import annotations + +import copy +import logging +from typing import Any, Dict, Optional + +from strands.hooks import ( + AfterToolCallEvent, + BeforeInvocationEvent, + BeforeModelCallEvent, + HookProvider, + HookRegistry, +) + +from apis.shared.feature_flags import cost_diagnostics_enabled + +logger = logging.getLogger(__name__) + +#: Upper bound on distinct tool names tallied per model call. A single call +#: that requests more than this many *distinct* tools is not a shape the +#: platform produces; the cap only guards the row size. +_MAX_TOOLS_PER_CALL = 64 + + +def _tool_name(event: Any) -> Optional[str]: + tool_use = getattr(event, "tool_use", None) or {} + if not isinstance(tool_use, dict): + return None + name = tool_use.get("name") + return str(name) if name else None + + +def _call_failed(event: Any) -> bool: + """A raised exception OR a result reporting ``status == "error"`` — + the same definition ``AgentStatusHook`` uses for ``ok=False``.""" + if getattr(event, "exception", None) is not None: + return True + result = getattr(event, "result", None) + return isinstance(result, dict) and result.get("status") == "error" + + +class ToolCensusHook(HookProvider): + """Per-turn, per-model-call tool tally: ``{cycle: {tool: {calls, errors}}}``.""" + + def __init__(self) -> None: + self._cycle = 0 + self._tally: Dict[int, Dict[str, Dict[str, int]]] = {} + + def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None: + registry.add_callback(BeforeInvocationEvent, self._on_turn_start) + registry.add_callback(BeforeModelCallEvent, self._on_before_model_call) + registry.add_callback(AfterToolCallEvent, self._on_after_tool_call) + + # -- reads (called by the stream coordinator at turn end) --------------- + + def tally_for_call(self, call_index: int) -> Optional[Dict[str, Dict[str, int]]]: + """The tools requested by model call ``call_index`` (0-based), or + ``None`` when that call requested none — or the census is off. + + Returns a copy so the caller can hand it to the persistence layer + without aliasing per-turn state. + """ + if not cost_diagnostics_enabled(): + return None + entry = self._tally.get(call_index + 1) + return copy.deepcopy(entry) if entry else None + + # -- callbacks ---------------------------------------------------------- + + def _on_turn_start(self, event: BeforeInvocationEvent) -> None: + self._cycle = 0 + self._tally = {} + + def _on_before_model_call(self, event: BeforeModelCallEvent) -> None: + self._cycle += 1 + + def _on_after_tool_call(self, event: AfterToolCallEvent) -> None: + if not cost_diagnostics_enabled(): + return + try: + name = _tool_name(event) + if not name: + return + per_call = self._tally.setdefault(self._cycle, {}) + slot = per_call.get(name) + if slot is None: + if len(per_call) >= _MAX_TOOLS_PER_CALL: + return + slot = per_call[name] = {"calls": 0, "errors": 0} + slot["calls"] += 1 + if _call_failed(event): + slot["errors"] += 1 + except Exception as e: # noqa: BLE001 - a census must never break a turn + logger.debug("Tool census skipped a call: %s", e) diff --git a/backend/src/agents/main_agent/session/turn_based_session_manager.py b/backend/src/agents/main_agent/session/turn_based_session_manager.py index 025e598fb..b889a0630 100644 --- a/backend/src/agents/main_agent/session/turn_based_session_manager.py +++ b/backend/src/agents/main_agent/session/turn_based_session_manager.py @@ -564,8 +564,17 @@ def _adopt_persisted_compaction_state(self) -> None: ) self.compaction_state = persisted - def _save_compaction_state(self, state: CompactionState) -> None: - """Save compaction state to DynamoDB session metadata.""" + def _save_compaction_state(self, state: CompactionState, record_event: bool = False) -> None: + """Save compaction state to DynamoDB session metadata. + + ``record_event=True`` means this save *is* a compaction (a new + checkpoint was cut), and bumps the session's ``compactionCount`` in + the same update. That counter is what the admin profile reads for + "how many times did compaction fire" — the persisted ``compaction`` + map is last-write-wins and cannot answer it. A top-level ``ADD`` is + monotonic and race-free: two Agents serving one session can each + increment, and neither can move it backwards. + """ if not self.user_id or not self.compaction_config or not self.compaction_config.enabled: return @@ -585,10 +594,18 @@ def _save_compaction_state(self, state: CompactionState) -> None: return state.updated_at = datetime.now(timezone.utc).isoformat() + update_expression = "SET compaction = :state" + values: Dict[str, Any] = {":state": state.to_dict()} + if record_event: + from apis.shared.feature_flags import cost_diagnostics_enabled + + if cost_diagnostics_enabled(): + update_expression += " ADD compactionCount :one" + values[":one"] = 1 table.update_item( Key={"PK": pk, "SK": sk}, - UpdateExpression="SET compaction = :state", - ExpressionAttributeValues={":state": state.to_dict()}, + UpdateExpression=update_expression, + ExpressionAttributeValues=values, ) logger.debug(f"Saved compaction state: checkpoint={state.checkpoint}") except Exception as e: @@ -812,7 +829,8 @@ async def update_after_turn( # Running total persisted alongside the rest of the compaction state # so a refresh can rehydrate the end-of-conversation summary indicator. self.compaction_state.total_summarized_turns += summarized_turns - self._save_compaction_state(self.compaction_state) + # This save is the compaction event itself — count it. + self._save_compaction_state(self.compaction_state, record_event=True) logger.info( f"Compaction checkpoint set: {new_checkpoint}, " diff --git a/backend/src/agents/main_agent/streaming/stream_coordinator.py b/backend/src/agents/main_agent/streaming/stream_coordinator.py index 0ffd89b3b..70af1875e 100644 --- a/backend/src/agents/main_agent/streaming/stream_coordinator.py +++ b/backend/src/agents/main_agent/streaming/stream_coordinator.py @@ -1171,6 +1171,11 @@ async def stream_response( message_ids_to_store = assistant_message_ids if assistant_message_ids else ([message_id] if message_id is not None else []) if message_ids_to_store: + # Content-free tool census, read (not drained) per call so each + # cost row carries the tools that call requested. None when the + # wrapper has no hook (tests, older agents) or the census is off. + tool_census_hook = getattr(main_agent_wrapper, "tool_census_hook", None) + # Build list of metadata storage tasks for parallel execution metadata_tasks = [] for idx, msg_id in enumerate(message_ids_to_store): @@ -1222,6 +1227,10 @@ async def stream_response( citations=citations_for_message, # Pass citations for persistence call_index=idx, # Nth model call of this turn (prefix fingerprint lookup) turn_agent_id=turn_agent_id, # Which Agent ran this turn (#756) + tool_calls=( + tool_census_hook.tally_for_call(idx) + if tool_census_hook is not None else None + ), ) ) @@ -2711,6 +2720,7 @@ async def _store_message_metadata( citations: Optional[List] = None, call_index: Optional[int] = None, turn_agent_id: Optional[str] = None, + tool_calls: Optional[Dict[str, Dict[str, int]]] = None, ) -> None: """ Store message-level metadata (token usage, latency, model info, citations) @@ -2882,6 +2892,14 @@ async def _store_message_metadata( if turn_agent_id: metadata_kwargs["turnAgentId"] = turn_agent_id + # Content-free tool census for this call (tool name → calls / + # errors), another extra field. Read by the admin session + # profile to show what the user was doing; tool names are + # catalog ids, never content. Absent when the call requested + # no tools or COST_DIAGNOSTICS_ENABLED=false. + if tool_calls: + metadata_kwargs["toolCalls"] = tool_calls + message_metadata = MessageMetadata(**metadata_kwargs) # Store metadata diff --git a/backend/src/apis/shared/feature_flags.py b/backend/src/apis/shared/feature_flags.py index 0c2be954b..d55282911 100644 --- a/backend/src/apis/shared/feature_flags.py +++ b/backend/src/apis/shared/feature_flags.py @@ -284,6 +284,30 @@ def tool_summaries_enabled() -> bool: return os.environ.get("TOOL_SUMMARIES_ENABLED", "").strip().lower() != "false" +def cost_diagnostics_enabled() -> bool: + """Whether the content-free behavioral counters are written at turn end. + + Covers the ``ToolCensusHook`` tally (tool name → calls/errors per model + call, persisted as ``toolCalls`` on the call's ``C#`` cost row), the + ``toolCallCount`` / ``toolErrorCount`` session rollups, and the + ``compactionCount`` session counter. These are what the admin session + profile reads to say *what the user was doing* without reading the + conversation. **Default ON with a kill switch** (house style): unset or + empty resolves to enabled; only the literal ``"false"`` disables. + + Read-side surfaces (``GET /admin/costs/.../profile``) are not gated — + they tolerate the attributes' absence and report "not tracked", which is + exactly what an environment with this switched off should see. + + Cost note (CLAUDE.md token-effectiveness tenet): every write here is + additive to rows the turn already writes (one extra attribute on the + ``C#`` put, two ``ADD`` terms on the existing session-aggregate + ``UpdateItem``, one ``ADD`` on the existing compaction-state update). + Nothing reaches the prompt; the cacheable prefix is untouched. + """ + return os.environ.get("COST_DIAGNOSTICS_ENABLED", "").strip().lower() != "false" + + def config_cache_enabled() -> bool: """Whether tenant-global config catalogs are served from the in-process cache. diff --git a/backend/src/apis/shared/sessions/metadata.py b/backend/src/apis/shared/sessions/metadata.py index 93ae58864..c52056656 100644 --- a/backend/src/apis/shared/sessions/metadata.py +++ b/backend/src/apis/shared/sessions/metadata.py @@ -1729,6 +1729,20 @@ async def _bump_session_aggregates( # A split of :wasted, not a deduction from it. values[":partialWasted"] = wasted_decimal if is_partial_miss else Decimal("0") + # Content-free behavioral rollups for the admin session profile: how + # many tool calls this session has made and how many failed, summed + # from the per-call census the coordinator attached as `toolCalls`. + # Only written while the census is on — an absent attribute is what + # lets the profile say "not tracked" instead of an honest-looking 0. + from apis.shared.feature_flags import cost_diagnostics_enabled + + if cost_diagnostics_enabled(): + tool_calls_total, tool_errors_total = _tool_census_totals(message_metadata) + update_parts_add.append("toolCallCount :toolCalls") + update_parts_add.append("toolErrorCount :toolErrors") + values[":toolCalls"] = tool_calls_total + values[":toolErrors"] = tool_errors_total + update_expression = ( "ADD " + ", ".join(update_parts_add) + " SET " + ", ".join(update_parts_set) ) @@ -1750,6 +1764,29 @@ async def _bump_session_aggregates( logger.debug("bump_session_aggregates failed (will be backfilled on read): %s", e) +def _tool_census_totals(message_metadata: Any) -> tuple[int, int]: + """``(calls, errors)`` summed over the call's ``toolCalls`` extra field. + + The field is ``{tool_name: {"calls": n, "errors": e}}`` when the + coordinator attached one, and absent otherwise; malformed entries count as + zero rather than raising — the aggregate bump must never fail on it. + """ + extra = getattr(message_metadata, "model_extra", None) + tool_calls = extra.get("toolCalls") if isinstance(extra, dict) else None + if not isinstance(tool_calls, dict): + return 0, 0 + calls = errors = 0 + for entry in tool_calls.values(): + if not isinstance(entry, dict): + continue + try: + calls += int(entry.get("calls") or 0) + errors += int(entry.get("errors") or 0) + except (TypeError, ValueError): + continue + return calls, errors + + def _emit_session_cache_rollup_metrics( session_id: str, update_response: Optional[Dict[str, Any]], diff --git a/backend/tests/agents/main_agent/session/test_compaction_stability.py b/backend/tests/agents/main_agent/session/test_compaction_stability.py index bef500736..ef386e602 100644 --- a/backend/tests/agents/main_agent/session/test_compaction_stability.py +++ b/backend/tests/agents/main_agent/session/test_compaction_stability.py @@ -99,7 +99,7 @@ def _restore(state_store: dict, stored_messages: list): state_store.get("compaction") ) - def _save(state: CompactionState) -> None: + def _save(state: CompactionState, record_event: bool = False) -> None: state.updated_at = _iso(_now()) state_store["compaction"] = state.to_dict() @@ -141,7 +141,7 @@ def _make(state_store: dict): state_store.get("compaction") ) - def _save(state: CompactionState) -> None: + def _save(state: CompactionState, record_event: bool = False) -> None: state.updated_at = _iso(_now()) state_store["compaction"] = state.to_dict() diff --git a/backend/tests/agents/main_agent/session/test_tool_census_hook.py b/backend/tests/agents/main_agent/session/test_tool_census_hook.py new file mode 100644 index 000000000..5916cee0f --- /dev/null +++ b/backend/tests/agents/main_agent/session/test_tool_census_hook.py @@ -0,0 +1,131 @@ +"""Tests for ToolCensusHook — the content-free per-call tool tally. + +Properties, in order of how expensive they are to get wrong: + +1. **Attribution.** A tool that runs during model-call cycle N belongs to + coordinator call index N-1. Off by one and every tool lands on the wrong + cost row, and the trajectory view lies about which call requested what. +2. **Per-turn reset.** The interrupt path unwinds without a turn end; without + the ``BeforeInvocationEvent`` reset the next turn inherits a stale tally. +3. **Non-drained reads.** The coordinator reads each call's tally once at turn + end; a read must not empty it (the same call could be read twice on a + retry) and must not alias the hook's state. +4. **Failure counting** matches ``AgentStatusHook``: an exception OR a result + with ``status == "error"``. +5. Fail-soft: the kill switch yields ``None`` everywhere, a malformed event + never raises, and the per-call distinct-tool cap holds. +""" + +from unittest.mock import MagicMock + +import pytest + +from agents.main_agent.session.hooks.tool_census import _MAX_TOOLS_PER_CALL, ToolCensusHook + + +@pytest.fixture(autouse=True) +def census_enabled(monkeypatch): + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + + +@pytest.fixture +def hook(): + return ToolCensusHook() + + +def _after(name="list_courses", *, exception=None, status="success"): + event = MagicMock() + event.tool_use = {"toolUseId": "t", "name": name, "input": {}} + event.exception = exception + event.result = {"status": status, "content": []} + return event + + +def _turn(hook: ToolCensusHook, *cycles): + """Drive a turn: each element of ``cycles`` is the list of after-tool + events that run after that model call.""" + hook._on_turn_start(MagicMock()) + for events in cycles: + hook._on_before_model_call(MagicMock()) + for event in events: + hook._on_after_tool_call(event) + + +def test_tools_are_attributed_to_the_model_call_that_requested_them(hook): + _turn( + hook, + [_after("list_courses"), _after("list_courses"), _after("get_syllabus")], # after call 0 + [], # call 1 ran no tools + [_after("calculator", status="error")], # after call 2 + ) + assert hook.tally_for_call(0) == { + "list_courses": {"calls": 2, "errors": 0}, + "get_syllabus": {"calls": 1, "errors": 0}, + } + assert hook.tally_for_call(1) is None + assert hook.tally_for_call(2) == {"calculator": {"calls": 1, "errors": 1}} + assert hook.tally_for_call(3) is None # the final answer call requested nothing + + +def test_a_new_turn_forgets_the_previous_one(hook): + _turn(hook, [_after("a")]) + assert hook.tally_for_call(0) is not None + _turn(hook, []) + assert hook.tally_for_call(0) is None + + +def test_reads_do_not_drain_and_do_not_alias(hook): + _turn(hook, [_after("a")]) + first = hook.tally_for_call(0) + first["a"]["calls"] = 99 # mutate the copy + second = hook.tally_for_call(0) + assert second == {"a": {"calls": 1, "errors": 0}} + + +def test_failures_count_on_exception_or_error_status(hook): + _turn(hook, [ + _after("t", exception=RuntimeError("boom")), + _after("t", status="error"), + _after("t"), + ]) + assert hook.tally_for_call(0) == {"t": {"calls": 3, "errors": 2}} + + +def test_kill_switch_records_nothing_and_reads_none(monkeypatch): + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "false") + hook = ToolCensusHook() + _turn(hook, [_after("a")]) + assert hook.tally_for_call(0) is None + assert hook._tally == {} + + +def test_empty_string_flag_means_on(monkeypatch): + # A GitHub Actions variable that is unset arrives as "", which must not disable. + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "") + hook = ToolCensusHook() + _turn(hook, [_after("a")]) + assert hook.tally_for_call(0) == {"a": {"calls": 1, "errors": 0}} + + +def test_malformed_events_never_raise(hook): + hook._on_turn_start(MagicMock()) + hook._on_before_model_call(MagicMock()) + bad = MagicMock() + bad.tool_use = "not-a-dict" + hook._on_after_tool_call(bad) + nameless = MagicMock() + nameless.tool_use = {"toolUseId": "t"} + hook._on_after_tool_call(nameless) + assert hook.tally_for_call(0) is None + + +def test_distinct_tool_cap_holds_but_known_tools_keep_counting(hook): + hook._on_turn_start(MagicMock()) + hook._on_before_model_call(MagicMock()) + for i in range(_MAX_TOOLS_PER_CALL + 5): + hook._on_after_tool_call(_after(f"tool_{i}")) + hook._on_after_tool_call(_after("tool_0")) + tally = hook.tally_for_call(0) + assert len(tally) == _MAX_TOOLS_PER_CALL + assert tally["tool_0"]["calls"] == 2 + assert f"tool_{_MAX_TOOLS_PER_CALL + 1}" not in tally diff --git a/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py b/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py index d943cbae7..f7f337065 100644 --- a/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py +++ b/backend/tests/agents/main_agent/session/test_turn_based_session_manager.py @@ -388,6 +388,43 @@ def test_save_noop_when_compaction_disabled(self, make_session_manager, compacti # Should not raise mgr._save_compaction_state(CompactionState(checkpoint=5)) + def test_record_event_bumps_a_monotonic_compaction_count( + self, make_session_manager, compaction_config, dynamodb_sessions_table, monkeypatch + ): + # The persisted `compaction` map is last-write-wins and cannot say how + # many times compaction fired; the top-level counter can, and only a + # save that *is* a compaction (record_event=True) moves it. + monkeypatch.setenv("DYNAMODB_SESSIONS_METADATA_TABLE_NAME", TABLE_NAME) + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + mgr = make_session_manager(compaction_config=compaction_config) + from agents.main_agent.session.turn_based_session_manager import TurnBasedSessionManager + TurnBasedSessionManager._dynamodb_table = dynamodb_sessions_table + seed_session_record(dynamodb_sessions_table, TEST_SESSION_ID, TEST_USER_ID) + + mgr._save_compaction_state(CompactionState(checkpoint=3), record_event=True) + mgr._save_compaction_state(CompactionState(checkpoint=3)) # bookkeeping save, no event + mgr._save_compaction_state(CompactionState(checkpoint=7), record_event=True) + + row = mgr._get_session_via_gsi(dynamodb_sessions_table) + assert row["compactionCount"] == 2 + assert row["compaction"]["checkpoint"] == 7 + + def test_record_event_respects_the_kill_switch( + self, make_session_manager, compaction_config, dynamodb_sessions_table, monkeypatch + ): + monkeypatch.setenv("DYNAMODB_SESSIONS_METADATA_TABLE_NAME", TABLE_NAME) + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "false") + mgr = make_session_manager(compaction_config=compaction_config) + from agents.main_agent.session.turn_based_session_manager import TurnBasedSessionManager + TurnBasedSessionManager._dynamodb_table = dynamodb_sessions_table + seed_session_record(dynamodb_sessions_table, TEST_SESSION_ID, TEST_USER_ID) + + mgr._save_compaction_state(CompactionState(checkpoint=3), record_event=True) + + row = mgr._get_session_via_gsi(dynamodb_sessions_table) + assert "compactionCount" not in row + assert row["compaction"]["checkpoint"] == 3 # the state itself still saves + # =========================================================================== # Task 6 — LTM summary retrieval diff --git a/backend/tests/agents/main_agent/streaming/test_tool_census_attach.py b/backend/tests/agents/main_agent/streaming/test_tool_census_attach.py new file mode 100644 index 000000000..5618c4343 --- /dev/null +++ b/backend/tests/agents/main_agent/streaming/test_tool_census_attach.py @@ -0,0 +1,85 @@ +"""The coordinator → cost-row seam for the tool census. + +`_store_message_metadata` gained a `tool_calls` argument; the tally it +receives must land on the persisted `MessageMetadata` as the `toolCalls` +extra field (the same mechanism `turnAgentId` uses), and must be *absent* — +not an empty dict — when the call requested no tools, so the profile's +coverage flag stays honest. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +from agents.main_agent.streaming.stream_coordinator import StreamCoordinator + + +def _coordinator() -> StreamCoordinator: + return object.__new__(StreamCoordinator) + + +def _usage_metadata(): + return {"usage": {"inputTokens": 100, "outputTokens": 20, "totalTokens": 120}} + + +@pytest.mark.asyncio +async def test_tool_calls_are_attached_to_the_stored_metadata(): + store = AsyncMock() + with patch("apis.shared.sessions.metadata.store_message_metadata", store): + await _coordinator()._store_message_metadata( + session_id="s1", + user_id="u1", + message_id=3, + accumulated_metadata=_usage_metadata(), + stream_start_time=0.0, + stream_end_time=1.0, + first_token_time=0.5, + agent=None, + call_index=0, + tool_calls={"list_courses": {"calls": 2, "errors": 0}}, + ) + + store.assert_awaited_once() + stored = store.await_args.kwargs["message_metadata"] + assert stored.model_extra["toolCalls"] == {"list_courses": {"calls": 2, "errors": 0}} + # And it serializes onto the row exactly as the profile reads it. + assert stored.model_dump(by_alias=True)["toolCalls"]["list_courses"]["calls"] == 2 + + +@pytest.mark.asyncio +async def test_no_tools_means_no_field_at_all(): + store = AsyncMock() + with patch("apis.shared.sessions.metadata.store_message_metadata", store): + await _coordinator()._store_message_metadata( + session_id="s1", + user_id="u1", + message_id=3, + accumulated_metadata=_usage_metadata(), + stream_start_time=0.0, + stream_end_time=1.0, + first_token_time=None, + agent=None, + call_index=0, + tool_calls=None, + ) + + stored = store.await_args.kwargs["message_metadata"] + assert "toolCalls" not in (stored.model_extra or {}) + assert "toolCalls" not in stored.model_dump(by_alias=True) + + +@pytest.mark.asyncio +async def test_the_argument_is_optional_for_the_interrupt_path(): + # `_persist_interruption` calls without tool_calls; the default must hold. + store = AsyncMock() + with patch("apis.shared.sessions.metadata.store_message_metadata", store): + await _coordinator()._store_message_metadata( + session_id="s1", + user_id="u1", + message_id=1, + accumulated_metadata=_usage_metadata(), + stream_start_time=0.0, + stream_end_time=1.0, + first_token_time=None, + ) + store.assert_awaited_once() diff --git a/backend/tests/shared/test_tool_census_persistence.py b/backend/tests/shared/test_tool_census_persistence.py new file mode 100644 index 000000000..3137b8efe --- /dev/null +++ b/backend/tests/shared/test_tool_census_persistence.py @@ -0,0 +1,98 @@ +"""The tool census reaches DynamoDB the way the profile reads it. + +`toolCalls` rides the per-call `C#` cost row as an extra field (like +`turnAgentId`), and the session row's `toolCallCount` / `toolErrorCount` are +bumped in the same aggregate update as `totalCost`. With the kill switch off, +neither attribute exists at all — that absence is what lets the admin profile +say "not tracked" rather than render an honest-looking zero. +""" + +from decimal import Decimal + +import pytest + +from apis.shared.sessions.models import MessageMetadata, ModelInfo, TokenUsage + + +def _meta(tool_calls=None): + kwargs = dict( + token_usage=TokenUsage(inputTokens=100, outputTokens=50, totalTokens=150), + model_info=ModelInfo(modelId="claude-haiku-4-5", modelName="Claude Haiku 4.5"), + cost=0.01, + ) + if tool_calls is not None: + kwargs["toolCalls"] = tool_calls + return MessageMetadata(**kwargs) + + +def _seed_session(table, session_id="s1", user_id="u1"): + table.put_item(Item={ + "PK": f"USER#{user_id}", + "SK": f"S#{session_id}", + "GSI_PK": f"SESSION#{session_id}", + "GSI_SK": "META", + "sessionId": session_id, + "userId": user_id, + "status": "active", + "createdAt": "2026-09-01T00:00:00Z", + "lastMessageAt": "2026-09-01T00:00:00Z", + "messageCount": Decimal(0), + }) + + +def _session_row(table, session_id="s1", user_id="u1"): + return table.get_item(Key={"PK": f"USER#{user_id}", "SK": f"S#{session_id}"})["Item"] + + +@pytest.mark.asyncio +async def test_tool_calls_land_on_the_cost_row_and_roll_up_on_the_session(sessions_metadata_table, monkeypatch): + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + from apis.shared.sessions.metadata import store_message_metadata + + _seed_session(sessions_metadata_table) + census = {"list_courses": {"calls": 2, "errors": 0}, "calculator": {"calls": 1, "errors": 1}} + await store_message_metadata(session_id="s1", user_id="u1", message_id=1, message_metadata=_meta(census)) + await store_message_metadata(session_id="s1", user_id="u1", message_id=2, message_metadata=_meta()) + + cost_rows = [i for i in sessions_metadata_table.scan()["Items"] if i["SK"].startswith("C#")] + assert len(cost_rows) == 2 + with_census = [r for r in cost_rows if "toolCalls" in r] + assert len(with_census) == 1 + assert with_census[0]["toolCalls"]["calculator"] == {"calls": Decimal(1), "errors": Decimal(1)} + + row = _session_row(sessions_metadata_table) + assert row["toolCallCount"] == Decimal(3) + assert row["toolErrorCount"] == Decimal(1) + + +@pytest.mark.asyncio +async def test_kill_switch_leaves_no_trace(sessions_metadata_table, monkeypatch): + monkeypatch.setenv("COST_DIAGNOSTICS_ENABLED", "false") + from apis.shared.sessions.metadata import store_message_metadata + + _seed_session(sessions_metadata_table) + # The coordinator would not attach toolCalls with the flag off; even if a + # stale row carries one, the session rollups must not be written. + await store_message_metadata( + session_id="s1", user_id="u1", message_id=1, + message_metadata=_meta({"a": {"calls": 1, "errors": 0}}), + ) + row = _session_row(sessions_metadata_table) + assert "toolCallCount" not in row and "toolErrorCount" not in row + # The unrelated aggregates still bump — the switch is scoped to the census. + assert row["totalCost"] == Decimal("0.01") + + +@pytest.mark.asyncio +async def test_malformed_census_entries_count_as_zero_not_as_a_failure(sessions_metadata_table, monkeypatch): + monkeypatch.delenv("COST_DIAGNOSTICS_ENABLED", raising=False) + from apis.shared.sessions.metadata import store_message_metadata + + _seed_session(sessions_metadata_table) + await store_message_metadata( + session_id="s1", user_id="u1", message_id=1, + message_metadata=_meta({"ok": {"calls": 2, "errors": 0}, "weird": "not-a-dict", "half": {"calls": "x"}}), + ) + row = _session_row(sessions_metadata_table) + assert row["toolCallCount"] == Decimal(2) + assert row["toolErrorCount"] == Decimal(0) diff --git a/scripts/build/build-one.sh b/scripts/build/build-one.sh index bfd4f1114..075bd3361 100755 --- a/scripts/build/build-one.sh +++ b/scripts/build/build-one.sh @@ -200,6 +200,7 @@ case "$SERVICE" in MANIFESTS=( "backend/src/apis/shared/__init__.py" "backend/src/apis/shared/errors.py" + "backend/src/apis/shared/feature_flags.py" ) # Both scheduled-runs Lambdas are arm64 (see the scheduled-runs # CDK construct).