CHANGELOG.md

@@ -9,15 +9,17 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm
 
 ### Added
 
-- **LLM Agent Token-Budget Awareness and Tool Output Pruning (issue #61)** (`llm.py`): Two complementary mechanisms to prevent context-window exhaustion in the multi-turn tool-call loop.
+- **LLM Agent Token-Budget Awareness and Tool Output Pruning (issue #61, #65)** (`llm.py`): Two complementary mechanisms to prevent context-window exhaustion in the multi-turn tool-call loop, with accurate token tracking for billing.
 
-  **Token-budget awareness** (`token_budget_percent` config, default off): tracks estimated token consumption before each LLM invocation. Emits a warning at 75% budget consumption. When the budget ceiling is exceeded, injects a synthesis prompt, permits one final tool-call round, then forces a text-only response. Token estimation uses `chat_model.get_num_tokens()` when available, falling back to a `len(content)//4` heuristic.
+  **Token-budget awareness** (`token_budget_percent` config, default off): tracks actual token consumption from LLM response metadata before each invocation. Emits a warning at 75% budget consumption. When the budget ceiling is exceeded, injects a synthesis prompt, permits one final tool-call round, then forces a text-only response. Token counts from all rounds (including budget-exhaustion synthesis, stuck-model synthesis, and pruning passes) are accumulated for accurate billing.
 
-  **Tool output pruning** (`allow_tool_output_pruning` config, default `false`): an implicit LLM extraction pass that runs between tool execution and ToolMessage construction, scoped initially to `file_read` calls. The main model can opt-out per call via the `output_prune` meta-argument, or provide a focused search directive via `output_prune_context`. The pruning model responds with `[PRUNE_INFO_START/END]` and `[PRUNE_OUTPUT_START/END]` markers (§4.4.9) to separate pruning metadata from extracted content. An optional `pruning_model` config allows using a lighter model for extraction. On failure, the raw tool output is used as-is.
+  **Tool output pruning** (`allow_tool_output_pruning` config, default `false`): an implicit LLM extraction pass that runs between tool execution and ToolMessage construction, scoped initially to `file_read` calls. The main model can opt-out per call via the `output_prune` meta-argument, or provide a focused search directive via `output_prune_context`. The pruning model responds with `[PRUNE_INFO_START/END]` and `[PRUNE_OUTPUT_START/END]` markers (§4.4.9) to separate pruning metadata from extracted content. On failure, the raw tool output is used as-is.
+
+  **Token tracking:** `_extract_token_counts()` helper extracts token metadata via a three-tier fallback (`usage_metadata` → `response_metadata["token_usage"]` → `(0, 0, False)`), returning a `(prompt, completion, metadata_present)` tuple so callers can distinguish absent metadata from legitimate zero-count responses. Missing token-usage metadata on the pruning model raises `MissingUsageMetadataError` instead of silently zeroing billing data. Tokens from prior loop rounds are preserved in `ActorResult` even when a pruning-pass error occurs mid-loop.
 
   **ADR:** `docs/adr/ADR-2031-tool-loop-token-budget-and-pruning.md` documents the specification extensions (three new §4.4 config fields, two runtime behaviors, one normative marker format).
 
-  **Module:** `src/cleveractors/agents/llm.py`. 1 source file changed, BDD: 21 new scenarios in `features/llm_agent_tool_calling.feature`, Robot: 5 new integration tests in `robot/llm_token_budget.robot`, ASV: 16 benchmarks in `benchmarks/token_budget_benchmark.py`.
+  **Module:** `src/cleveractors/agents/llm.py`, `src/cleveractors/core/exceptions.py`. BDD: scenarios in `features/llm_agent_tool_calling.feature`, Robot: integration tests in `robot/llm_token_budget.robot`, ASV: benchmarks in `benchmarks/token_budget_benchmark.py`.
 
 - **Tool Schema Descriptions (`llm_tools.py`) (issue #59)** (`cleveractors.agents.llm_tools`): New module providing `normalize_tool_entry()` and `_BUILTIN_TOOL_SCHEMAS` — a registry of OpenAI function-calling parameter schemas for all built-in tools (`echo`, `math`, `json_parse`, `http_request`, `file_read`, `file_write`, `progress_bar`, `shell`, `python_exec`). Each schema includes targeted LLM guidance: `file_read` requires `max_chars` for files >10KB to prevent context overflow; `file_write` clarifies that response text is NOT persisted to disk; `shell` directs the LLM toward `file_read` for reading files; `python_exec` lists available sandbox builtins and instructs the LLM to use `file_read`/`file_write` instead of `open()`. String tool names in agent config (e.g. `"file_read"`) are normalized to full parameter schemas so the LLM knows about parameters like `max_chars`. Config dicts with a `"name"` key are also normalized; already-formatted OpenAI tool dicts pass through unchanged.
 
@@ -29,6 +31,7 @@ and this project adheres to [Clever Semantic Versioning](https://www.w3.org/subm
 
 ### Fixed
 
+
 - **LLM Agent Synthesis-Round Tool Name Validation (issue #63)** (`llm.py`): When the model hallucinates a tool name not present in the declared tools list during the synthesis round, the invalid call is now rejected with a `ToolMessage` error stating `"Tool '<name>' is not available."` instead of crashing or silently discarding the response. A `_declared_names` set is built once per synthesis round from `self._lc_tools` for O(1) lookup.
 
 - **ToolAgent Tool Name Dispatch Bug (issue #63)** (`tool.py`): `_execute_tool()` was checking `tool_name not in self.tools` where `self.tools` is a heterogeneous list of strings and dicts — a string tool name like `"file_read"` would never match the list because the list contains strings, not dicts with that name. Fixed by extracting `string_tool_names` and checking against that list instead.

benchmarks/token_budget_benchmark.py

@@ -36,46 +36,45 @@ def _make_agent(**extra_config: Any) -> LLMAgent:
     return agent
 
 
-def _messages_with_size(chars: int) -> list[Any]:
-    """Build a message list with approximately *chars* characters of content."""
-    content = "x" * max(1, chars - 70)
-    return [
-        SystemMessage(content="You are a helpful assistant."),
-        HumanMessage(content=content),
-    ]
+# ── Token Extraction Benchmarks ───────────────────────────────────────
 
 
-# ── Token Budget Benchmarks ────────────────────────────────────────────
-
-
-class TokenBudgetBenchmark:
-    """Benchmarks for token-budget awareness (§4.4.7)."""
+class TokenExtractionBenchmark:
+    """Benchmarks for _extract_token_counts (§4.4.7)."""
 
     def setup(self) -> None:
         self.agent = _make_agent()
 
+        self.resp_with_usage = MagicMock(spec=AIMessage)
+        self.resp_with_usage.usage_metadata = {
+            "input_tokens": 150,
+            "output_tokens": 50,
+        }
+
+        self.resp_via_metadata = MagicMock(spec=AIMessage)
+        self.resp_via_metadata.usage_metadata = None
+        self.resp_via_metadata.response_metadata = {
+            "token_usage": {"prompt_tokens": 200, "completion_tokens": 75},
+        }
+
+        self.resp_no_metadata = MagicMock(spec=AIMessage)
+        self.resp_no_metadata.usage_metadata = None
+        self.resp_no_metadata.response_metadata = None
+
     def teardown(self) -> None:
         self.agent = None
 
-    def time_estimate_small_messages(self) -> None:
-        """Estimate tokens for a small (100 char) conversation."""
-        msgs = _messages_with_size(100)
-        self.agent._estimate_token_count(msgs)
+    def time_extract_via_usage_metadata(self) -> None:
+        """Extract token counts via usage_metadata path."""
+        self.agent._extract_token_counts(self.resp_with_usage)
 
-    def time_estimate_medium_messages(self) -> None:
-        """Estimate tokens for a 10K char conversation."""
-        msgs = _messages_with_size(10_000)
-        self.agent._estimate_token_count(msgs)
+    def time_extract_via_response_metadata(self) -> None:
+        """Extract token counts via response_metadata path."""
+        self.agent._extract_token_counts(self.resp_via_metadata)
 
-    def time_estimate_large_messages(self) -> None:
-        """Estimate tokens for a 50K char conversation."""
-        msgs = _messages_with_size(50_000)
-        self.agent._estimate_token_count(msgs)
-
-    def time_estimate_very_large_messages(self) -> None:
-        """Estimate tokens for a 200K char conversation near context limit."""
-        msgs = _messages_with_size(200_000)
-        self.agent._estimate_token_count(msgs)
+    def time_extract_no_metadata(self) -> None:
+        """Extract token counts when no metadata is available."""
+        self.agent._extract_token_counts(self.resp_no_metadata)
 
     def time_context_window_lookup(self) -> None:
         """Look up the context window for the configured model."""
@@ -207,6 +206,7 @@ class PruningPassBenchmark:
                 f"[PRUNE_OUTPUT_START]\n{self.extracted}\n"
                 f"[PRUNE_OUTPUT_END]"
             )
+            resp.usage_metadata = {"input_tokens": 50, "output_tokens": 20}
             return resp
 
         mock_prune_model.ainvoke = _mock_prune_ainvoke

features/llm_agent_tool_calling.feature

@@ -120,6 +120,19 @@ Feature: LLMAgent Tool Calling
     When I create the LLM agent
     Then the agent _pruning_model should be None
 
+  # ── pruning_model matching enforcement (issue #65) ──────────────────
+
+  Scenario: pruning_model mismatch raises ConfigurationError
+    Given I have an LLM agent config with model "gpt-4" and pruning_model "gpt-3.5-turbo"
+    When I attempt to create the LLM agent
+    Then ConfigurationError should be raised containing "must match"
+
+  Scenario: pruning_model same as model succeeds
+    Given I have an LLM agent config with model "gpt-4" and pruning_model "gpt-4"
+    When I create the LLM agent
+    Then the agent _pruning_model should be "gpt-4"
+    And the agent model should be "gpt-4"
+
   # ── pruning_threshold (§4.4.8 D-2) ────────────────────────────────────
 
   Scenario: LLMAgent stores pruning_threshold when configured
@@ -243,18 +256,11 @@ Feature: LLMAgent Tool Calling
     When I create the LLM agent
     Then _get_model_context_window should return 128000
 
-  # ── Token Estimation (§4.4.7) ───────────────────────────────────────
-
-  Scenario: _estimate_token_count uses heuristic fallback
-    Given I have a basic LLM agent configuration (no tools)
-    When I create the LLM agent
-    And I set a mock chat model that returns a plain text response
-    And I estimate token count for messages with content "hello"
-    Then the estimated token count should be greater than 0
+  # ── _run_pruning_pass with mocked pruning model (§4.4.8) ────────────
 
   # ── _run_pruning_pass with mocked pruning model (§4.4.8) ────────────
 
-  Scenario: _run_pruning_pass returns marked-up response
+  Scenario: _run_pruning_pass returns marked-up response with token counts
     Given I have an LLM agent config with allow_tool_output_pruning true
     When I create the LLM agent
     And I set a mock pruning model that returns a marked-up response
@@ -262,6 +268,7 @@ Feature: LLMAgent Tool Calling
     Then the pruned result should contain PRUNE_INFO_START
     And the pruned result should contain PRUNE_OUTPUT_START
     And the pruned result should contain "extracted content"
+    And the pruning token counts should be greater than zero
 
   Scenario: _run_pruning_pass falls back to raw output on error
     Given I have an LLM agent config with allow_tool_output_pruning true
@@ -276,6 +283,28 @@ Feature: LLMAgent Tool Calling
     And I run the pruning pass on tool "echo" with empty output
     Then the pruned result should be empty
 
+  # ── Token usage metadata enforcement (issue #65) ────────────────────
+
+  Scenario: _run_pruning_pass fails when response lacks usage_metadata
+    Given I have an LLM agent config with allow_tool_output_pruning true
+    When I create the LLM agent
+    And I set a mock pruning model that returns a response without usage_metadata
+    And I attempt to run the pruning pass on tool "file_read" with output "50KB of logs"
+    Then a MissingUsageMetadataError should be raised containing "token usage metadata"
+
+  Scenario: MissingUsageMetadataError in tool loop preserves accumulated tokens
+    Given I have an LLM agent configuration with tools [{"name": "file_read"}]
+    And I set allow_tool_output_pruning to true
+    And I set pruning_threshold to 1
+    And I set tool_max_rounds to 3
+    And I set unsafe_mode to true
+    When I create the LLM agent
+    And I set a mock chat model that returns one file_read tool call then text
+    And I set a mock pruning model that returns a response without usage_metadata
+    And I attempt to tool_calling process a message "Test billing preservation"
+    Then an ExecutionError should be raised containing "token usage metadata"
+    And the last token usage should have accumulated counts from prior rounds
+
   # ── Token budget exhaustion in tool loop (§4.4.7) ──────────────────
 
   Scenario: Token budget exhaustion triggers synthesis flow
@@ -355,3 +384,115 @@ Feature: LLMAgent Tool Calling
     And I set a mock chat model that gets stuck then returns hallucinated tool name in synthesis
     And I tool_calling process a message "Get stuck"
     Then the synthesis messages should contain a ToolMessage with content "Tool 'nonexistent_tool' is not available."
+
+  # ── Additional coverage: token_budget_percent type validation ──────
+
+  Scenario: LLMAgent validates token_budget_percent must be numeric
+    Given I have a basic LLM agent configuration (no tools)
+    And I set token_budget_percent to string "not_a_number"
+    When I attempt to create the LLM agent
+    Then ConfigurationError should be raised containing "token_budget_percent must be a number"
+
+  # ── Additional coverage: pruning response missing end tag ──────────
+
+  Scenario: Parse pruning response with missing end tag falls back to full text
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I parse a pruning response that has start tag but no end tag
+    Then the parsed result should contain "Partial info without end tag"
+
+  # ── Additional coverage: tool_max_rounds parse error ──────────────
+
+  Scenario: LLMAgent rejects invalid tool_max_rounds on process_message
+    Given I have an LLM agent configuration with tools [{"name": "echo"}]
+    And I set tool_max_rounds to string "abc"
+    When I create the LLM agent
+    And I set a mock chat model that returns a plain text response
+    And I attempt to tool_calling process a message "Hello"
+    Then ConfigurationError should be raised containing "tool_max_rounds must be an integer"
+
+  # ── Additional coverage: direct-format tool calls ──────────────────
+
+  Scenario: LLMAgent handles tool_calls in direct format without function key
+    Given I have an LLM agent configuration with tools [{"name": "echo"}]
+    And I set tool_max_rounds to 2
+    When I create the LLM agent
+    And I set a mock chat model that returns a direct-format tool call then text
+    And I tool_calling process a message "Test direct format"
+    Then the tool_calling result should contain "Direct format answer"
+
+  # ── _extract_token_counts coverage: response_metadata fallback ─────
+
+  Scenario: _extract_token_counts extracts tokens via response_metadata fallback
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I extract token counts from a response using response_metadata fallback
+    Then the extracted token counts should be (30, 15, True)
+
+  Scenario: _extract_token_counts returns zero for empty usage_metadata
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I extract token counts from a response with empty usage_metadata
+    Then the extracted token counts should be (0, 0, False)
+
+  Scenario: _extract_token_counts returns zero when response_metadata has no token_usage key
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I extract token counts from a response with response_metadata without token_usage key
+    Then the extracted token counts should be (0, 0, False)
+
+  Scenario: _extract_token_counts returns zero for non-dict response_metadata
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I extract token counts from a response with non-dict response_metadata
+    Then the extracted token counts should be (0, 0, False)
+
+  Scenario: _extract_token_counts returns zero for missing response_metadata
+    Given I have a basic LLM agent configuration (no tools)
+    When I create the LLM agent
+    And I extract token counts from a response with missing response_metadata
+    Then the extracted token counts should be (0, 0, False)
+
+  # ── Budget synthesis with tool calls and pruning ──────────────────
+
+  Scenario: Budget synthesis tool calls trigger pruning pass
+    Given I have an LLM agent configuration with tools [{"name": "file_read"}]
+    And I set token_budget_percent to 0.000001
+    And I set allow_tool_output_pruning to true
+    And I set pruning_threshold to 1
+    And I set tool_max_rounds to 5
+    And I set unsafe_mode to true
+    When I create the LLM agent
+    And I set a mock chat model that triggers budget exhaustion then returns tool calls in synthesis
+    And I set a mock pruning model that returns a marked-up response
+    And I tool_calling process a message "Test budget synth pruning"
+    Then the tool_calling result should contain "Synthesis answer with pruning"
+    And the pruning pass should have been called
+
+  # ── MissingUsageMetadataError in budget synthesis ─────────────────
+
+  Scenario: MissingUsageMetadataError in budget synthesis preserves accumulated tokens
+    Given I have an LLM agent configuration with tools [{"name": "file_read"}]
+    And I set token_budget_percent to 0.000001
+    And I set allow_tool_output_pruning to true
+    And I set pruning_threshold to 1
+    And I set tool_max_rounds to 5
+    And I set unsafe_mode to true
+    When I create the LLM agent
+    And I set a mock chat model that triggers budget exhaustion then returns tool calls in synthesis
+    And I set a mock pruning model that returns a response without usage_metadata
+    And I attempt to tool_calling process a message "Test budget synth error"
+    Then an ExecutionError should be raised containing "token usage metadata"
+    And the last token usage should have accumulated counts from prior rounds
+
+  # ── Stuck-model synthesis tool execution error ─────────────────────
+
+  Scenario: Stuck-model synthesis tool error preserves accumulated tokens
+    Given I have an LLM agent configuration with tools [{"name": "echo"}]
+    And I set tool_max_rounds to 2
+    When I create the LLM agent
+    And I set a mock chat model that gets stuck then returns a tool call that raises ExecutionError
+    And I tool_calling process a message "Get stuck with error"
+    Then the tool_calling result should contain "Final answer after error"
+    And the accumulated prompt tokens should be greater than 0
+    And the accumulated completion tokens should be greater than 0

features/steps/llm_agent_tool_calling_steps.py

@@ -17,7 +17,11 @@ else:
     from langchain_core.language_models.chat_models import BaseChatModel
 
 from cleveractors.agents.llm import LLMAgent
-from cleveractors.core.exceptions import ConfigurationError
+from cleveractors.core.exceptions import (
+    ConfigurationError,
+    ExecutionError,
+    MissingUsageMetadataError,
+)
 from cleveractors.templates.renderer import TemplateRenderer
 
 
@@ -253,6 +257,26 @@ def step_set_tool_max_rounds(context, n):
     context.llm_config["tool_max_rounds"] = n
 
 
+@given("I set pruning_threshold to {n:d}")
+def step_set_pruning_threshold(context, n):
+    context.llm_config["pruning_threshold"] = n
+
+
+@given("I set unsafe_mode to true")
+def step_set_unsafe_mode_true(context):
+    context.llm_config["unsafe_mode"] = True
+
+
+@given('I set tool_max_rounds to string "{bad_value}"')
+def step_set_tool_max_rounds_bad(context, bad_value):
+    context.llm_config["tool_max_rounds"] = bad_value
+
+
+@given('I set token_budget_percent to string "{bad_value}"')
+def step_set_budget_percent_string(context, bad_value):
+    context.llm_config["token_budget_percent"] = bad_value
+
+
 @when(
     "I set a mock chat model that returns tool_calls with empty content "
     "for {n:d} rounds then text"
@@ -422,6 +446,20 @@ def step_config_with_pruning_model(context, model):
     }
 
 
+@given(
+    'I have an LLM agent config with model "{agent_model}" '
+    'and pruning_model "{prune_model}"'
+)
+def step_config_with_model_and_pruning_model(context, agent_model, prune_model):
+    context.llm_config = {
+        "name": "pm_model_agent",
+        "provider": "openai",
+        "api_key": "test_key",
+        "model": agent_model,
+        "pruning_model": prune_model,
+    }
+
+
 @given("I have an LLM agent config with pruning_threshold {n:d}")
 def step_config_with_pruning_threshold(context, n):
     context.llm_config = {
@@ -546,12 +584,11 @@ def step_parse_no_markers(context, text):
     context._parsed_pruning = context.llm_agent._parse_pruning_response(text)
 
 
-@when('I estimate token count for messages with content "{text}"')
-def step_estimate_tokens(context, text):
-    from langchain_core.messages import HumanMessage as HM
-
-    msgs = [HM(content=text)]
-    context._estimated_tokens = context.llm_agent._estimate_token_count(msgs)
+@when("I parse a pruning response that has start tag but no end tag")
+def step_parse_missing_end_tag(context):
+    agent = context.llm_agent
+    response = agent._PRUNE_INFO_START + "\nPartial info without end tag"
+    context._parsed_pruning = agent._parse_pruning_response(response)
 
 
 def _cleanup_prune_patch(context):
@@ -572,6 +609,7 @@ def step_mock_pruning_model(context):
             "[PRUNE_INFO_START]\nRemoved 48K chars of logs\n[PRUNE_INFO_END]\n"
             "[PRUNE_OUTPUT_START]\nextracted content\n[PRUNE_OUTPUT_END]"
         )
+        response.usage_metadata = {"input_tokens": 50, "output_tokens": 20}
         return response
 
     mock_prune_model.ainvoke = _mock_prune_ainvoke
@@ -597,11 +635,39 @@ def step_mock_pruning_model_error(context):
     context._prune_patch.start()
 
 
+@when("I set a mock pruning model that returns a response without usage_metadata")
+def step_mock_pruning_model_no_usage(context):
+    mock_prune_model = Mock(spec=BaseChatModel)
+
+    async def _mock_prune_ainvoke(messages, **kw):
+        context._pruning_was_called = True
+        response = Mock(spec=AIMessage)
+        response.content = (
+            "[PRUNE_INFO_START]\nRemoved 48K chars of logs\n[PRUNE_INFO_END]\n"
+            "[PRUNE_OUTPUT_START]\nextracted content\n[PRUNE_OUTPUT_END]"
+        )
+        # Deliberately omit usage_metadata to trigger the enforcement
+        # (issue #65 Gap 2 enforcement).
+        if hasattr(response, "usage_metadata"):
+            del response.usage_metadata
+        return response
+
+    mock_prune_model.ainvoke = _mock_prune_ainvoke
+    context._prune_patch = patch(
+        "cleveractors.agents.llm.build_chat_model",
+        return_value=mock_prune_model,
+    )
+    context._prune_patch.start()
+
+
 @when('I run the pruning pass on tool "{tool_name}" with output "{output}"')
 def step_run_pruning_pass(context, tool_name, output):
-    context._prune_result = asyncio.run(
+    _prune_str, _prune_pt, _prune_ct = asyncio.run(
         context.llm_agent._run_pruning_pass(tool_name, output, [])
     )
+    context._prune_result = _prune_str
+    context._prune_prompt_tokens = _prune_pt
+    context._prune_completion_tokens = _prune_ct
 
 
 @when(
@@ -609,7 +675,7 @@ def step_run_pruning_pass(context, tool_name, output):
     '"{output}" and prune_context "{prune_context}"'
 )
 def step_run_pruning_pass_with_context(context, tool_name, output, prune_context):
-    context._prune_result = asyncio.run(
+    _prune_str, _prune_pt, _prune_ct = asyncio.run(
         context.llm_agent._run_pruning_pass(
             tool_name,
             output,
@@ -617,13 +683,28 @@ def step_run_pruning_pass_with_context(context, tool_name, output, prune_context
             prune_context=prune_context,
         )
     )
+    context._prune_result = _prune_str
+    context._prune_prompt_tokens = _prune_pt
+    context._prune_completion_tokens = _prune_ct
 
 
 @when('I run the pruning pass on tool "{tool_name}" with empty output')
 def step_run_pruning_pass_empty(context, tool_name):
-    context._prune_result = asyncio.run(
+    _prune_str, _prune_pt, _prune_ct = asyncio.run(
         context.llm_agent._run_pruning_pass(tool_name, "", [])
     )
+    context._prune_result = _prune_str
+    context._prune_prompt_tokens = _prune_pt
+    context._prune_completion_tokens = _prune_ct
+
+
+@when('I attempt to run the pruning pass on tool "{tool_name}" with output "{output}"')
+def step_attempt_run_pruning_pass(context, tool_name, output):
+    try:
+        asyncio.run(context.llm_agent._run_pruning_pass(tool_name, output, []))
+        context._prune_error = None
+    except MissingUsageMetadataError as e:
+        context._prune_error = e
 
 
 @when("I set a mock chat model that returns tool_calls then text for budget exhaustion")
@@ -681,6 +762,45 @@ def step_mock_one_tool_call(context):
     context.llm_agent.chat_model = mock_model
 
 
+@when("I set a mock chat model that returns one file_read tool call then text")
+def step_mock_one_file_read_tool_call(context):
+    import os
+
+    cwd = os.getcwd()
+    file_path = os.path.join(cwd, "test_prune_file.txt")
+    with open(file_path, "w") as f:
+        f.write("A" * 600)
+    context._prune_file_path = file_path
+    if not hasattr(context, "_cleanup_files"):
+        context._cleanup_files = []
+    context._cleanup_files.append(file_path)
+
+    call_counter = [0]
+    context._pruning_was_called = False
+
+    async def _mock_ainvoke(messages, **invoke_kwargs):
+        call_counter[0] += 1
+        if call_counter[0] == 1:
+            return _create_mock_response(
+                "Thinking...",
+                tool_calls=[
+                    {
+                        "id": "call_prune_1",
+                        "type": "function",
+                        "function": {
+                            "name": "file_read",
+                            "arguments": _json.dumps({"file": file_path}),
+                        },
+                    }
+                ],
+            )
+        return _create_mock_response("Final answer")
+
+    mock_model = Mock(spec=BaseChatModel)
+    mock_model.ainvoke = _mock_ainvoke
+    context.llm_agent.chat_model = mock_model
+
+
 @when(
     "I set a mock chat model that returns tool call with output_prune false then text"
 )
@@ -727,18 +847,100 @@ def step_mock_output_prune_false(context):
     context.llm_agent.chat_model = mock_model
 
 
+@when("I set a mock chat model that returns a direct-format tool call then text")
+def step_mock_direct_format_tool_call(context):
+    call_counter = [0]
+
+    async def _mock_ainvoke(messages, **invoke_kwargs):
+        call_counter[0] += 1
+        if call_counter[0] == 1:
+            return _create_mock_response(
+                "Thinking...",
+                tool_calls=[
+                    {
+                        "id": "call_direct_1",
+                        "name": "echo",
+                        "args": {"text": "hello"},
+                    }
+                ],
+            )
+        return _create_mock_response("Direct format answer")
+
+    mock_model = Mock(spec=BaseChatModel)
+    mock_model.ainvoke = _mock_ainvoke
+    context.llm_agent.chat_model = mock_model
+
+
+@when('I attempt to tool_calling process a message "{message_txt}"')
+def step_attempt_process_message(context, message_txt):
+    try:
+        context.result = asyncio.run(context.llm_agent.process_message(message_txt))
+        context._process_error = None
+    except (ConfigurationError, ExecutionError) as e:
+        context._process_error = e
+
+
 # ─── Then steps (token-budget / pruning) ─────────────────────────────
 
 
 @then('ConfigurationError should be raised containing "{text}"')
 def step_config_error_raised(context, text):
-    err = context._create_error
+    err = getattr(context, "_process_error", None)
+    if err is None:
+        err = getattr(context, "_create_error", None)
     assert err is not None, "Expected ConfigurationError but none was raised"
     assert text in str(err), (
         f"ConfigurationError message did not contain {text!r}: {err}"
     )
 
 
+@then('a MissingUsageMetadataError should be raised containing "{text}"')
+def step_missing_metadata_error_raised(context, text):
+    err = context._prune_error
+    assert err is not None, "Expected MissingUsageMetadataError but none was raised"
+    assert isinstance(err, MissingUsageMetadataError), (
+        f"Expected MissingUsageMetadataError, got {type(err).__name__}: {err}"
+    )
+    assert text in str(err), (
+        f"MissingUsageMetadataError message did not contain {text!r}: {err}"
+    )
+
+
+@then('an ExecutionError should be raised containing "{text}"')
+def step_execution_error_raised(context, text):
+    err = context._process_error
+    assert err is not None, "Expected ExecutionError but none was raised"
+    assert isinstance(err, ExecutionError), (
+        f"Expected ExecutionError, got {type(err).__name__}: {err}"
+    )
+    assert text in str(err), f"ExecutionError message did not contain {text!r}: {err}"
+
+
+@then("the last token usage should have accumulated counts from prior rounds")
+def step_last_token_usage_accumulated(context):
+    pt, ct = context.llm_agent._last_token_usage
+    assert pt > 0, f"Expected prompt_tokens > 0, got {pt}"
+    assert ct > 0, f"Expected completion_tokens > 0, got {ct}"
+    _cleanup_prune_patch(context)
+    context._process_error = None
+
+
+@then("the pruning token counts should be greater than zero")
+def step_prune_token_counts_gt_zero(context):
+    pt = context._prune_prompt_tokens
+    ct = context._prune_completion_tokens
+    assert pt is not None, "prune_prompt_tokens was not set"
+    assert ct is not None, "prune_completion_tokens was not set"
+    assert pt > 0, f"Expected prompt_tokens > 0, got {pt}"
+    assert ct > 0, f"Expected completion_tokens > 0, got {ct}"
+
+
+@then('the agent model should be "{expected}"')
+def step_agent_model(context, expected):
+    actual = context.llm_agent.model
+    assert actual == expected, f"Agent model = {actual!r} != {expected!r}"
+
+
 @then("the agent _token_budget_percent should be {value:g}")
 def step_budget_stored(context, value):
     actual = context.llm_agent._token_budget_percent
@@ -850,19 +1052,19 @@ def step_parsed_equals(context, text):
     )
 
 
+@then('the parsed result should contain "{text}"')
+def step_parsed_contains(context, text):
+    assert text in context._parsed_pruning, (
+        f"Expected {text!r} in parsed result: {context._parsed_pruning!r}"
+    )
+
+
 @then("_get_model_context_window should return {n:d}")
 def step_context_window(context, n):
     actual = context.llm_agent._get_model_context_window()
     assert actual == n, f"_get_model_context_window = {actual} != {n}"
 
 
-@then("the estimated token count should be greater than 0")
-def step_estimate_positive(context):
-    assert context._estimated_tokens > 0, (
-        f"Expected >0 tokens, got {context._estimated_tokens}"
-    )
-
-
 @then("the pruned result should contain PRUNE_INFO_START")
 def step_pruned_has_info(context):
     assert context.llm_agent._PRUNE_INFO_START in context._prune_result
@@ -1063,3 +1265,193 @@ def step_synthesis_messages_have_toolmessage(context, text):
         f"No ToolMessage found with exact content {text!r}. "
         f"ToolMessages: {[str(m.content) for m in tool_msgs]}"
     )
+
+
+# ─── _extract_token_counts direct testing steps ──────────────────────
+
+
+@when("I extract token counts from a response with usage_metadata only")
+def step_extract_via_usage_metadata(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = {"input_tokens": 50, "output_tokens": 25}
+    resp.response_metadata = None
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@when("I extract token counts from a response using response_metadata fallback")
+def step_extract_via_response_metadata(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = None
+    resp.response_metadata = {
+        "token_usage": {"prompt_tokens": 30, "completion_tokens": 15}
+    }
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@when("I extract token counts from a response with empty usage_metadata")
+def step_extract_empty_usage(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = {}
+    resp.response_metadata = None
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@when(
+    "I extract token counts from a response with response_metadata "
+    "without token_usage key"
+)
+def step_extract_no_token_usage(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = None
+    resp.response_metadata = {"other_key": "value"}
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@when("I extract token counts from a response with non-dict response_metadata")
+def step_extract_non_dict_response_metadata(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = None
+    resp.response_metadata = [1, 2, 3]
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@when("I extract token counts from a response with missing response_metadata")
+def step_extract_missing_response_metadata(context):
+    resp = Mock(spec=AIMessage)
+    resp.usage_metadata = None
+    resp.response_metadata = None
+    context._extracted_tokens = context.llm_agent._extract_token_counts(resp)
+
+
+@then("the extracted token counts should be ({pt:d}, {ct:d}, {mp})")
+def step_assert_extracted_tokens(context, pt, ct, mp):
+    expected_mp = mp == "True"
+    actual_pt, actual_ct, actual_mp = context._extracted_tokens
+    assert actual_pt == pt, f"Expected prompt_tokens={pt}, got {actual_pt}"
+    assert actual_ct == ct, f"Expected completion_tokens={ct}, got {actual_ct}"
+    assert actual_mp is expected_mp, (
+        f"Expected metadata_present={expected_mp}, got {actual_mp}"
+    )
+
+
+# ─── Budget synthesis with tool calls and pruning ────────────────────
+
+
+@when(
+    "I set a mock chat model that triggers budget exhaustion "
+    "then returns tool calls in synthesis"
+)
+def step_mock_budget_then_synth_tools(context):
+    import os
+
+    cwd = os.getcwd()
+    file_path = os.path.join(cwd, "test_budget_synth_prune.txt")
+    with open(file_path, "w") as f:
+        f.write("A" * 700)
+    if not hasattr(context, "_cleanup_files"):
+        context._cleanup_files = []
+    context._cleanup_files.append(file_path)
+
+    call_counter = [0]
+    context._pruning_was_called = False
+
+    async def _mock_ainvoke(messages, **invoke_kwargs):
+        call_counter[0] += 1
+        if call_counter[0] == 1:
+            # Use echo (not in pruning filter) to avoid pruning in main loop
+            return _create_mock_response(
+                "",
+                tool_calls=[
+                    {
+                        "id": "call_budget_synth_1",
+                        "type": "function",
+                        "function": {
+                            "name": "echo",
+                            "arguments": '{"text": "hello"}',
+                        },
+                    }
+                ],
+            )
+        if call_counter[0] == 2:
+            return _create_mock_response(
+                "Let me read the file",
+                tool_calls=[
+                    {
+                        "id": "call_budget_synth_2",
+                        "type": "function",
+                        "function": {
+                            "name": "file_read",
+                            "arguments": _json.dumps({"file": file_path}),
+                        },
+                    }
+                ],
+            )
+        return _create_mock_response("Synthesis answer with pruning")
+
+    mock_model = Mock(spec=BaseChatModel)
+    mock_model.ainvoke = _mock_ainvoke
+    context.llm_agent.chat_model = mock_model
+
+
+# ─── Stuck-model synthesis tool error ────────────────────────────────
+
+
+@when(
+    "I set a mock chat model that gets stuck then returns a tool call "
+    "that raises ExecutionError"
+)
+def step_mock_stuck_then_tool_error(context):
+    call_counter = [0]
+
+    async def _mock_ainvoke(messages, **invoke_kwargs):
+        call_counter[0] += 1
+        if call_counter[0] <= 2:
+            return _create_mock_response(
+                "",
+                tool_calls=[
+                    {
+                        "id": f"call_stuck_err_{call_counter[0]}",
+                        "type": "function",
+                        "function": {
+                            "name": "echo",
+                            "arguments": '{"text": "stuck"}',
+                        },
+                    }
+                ],
+            )
+        if call_counter[0] == 3:
+            return _create_mock_response(
+                "Let me try a different tool",
+                tool_calls=[
+                    {
+                        "id": "call_synth_err",
+                        "type": "function",
+                        "function": {
+                            "name": "nonexistent_tool",
+                            "arguments": '{"text": "boom"}',
+                        },
+                    }
+                ],
+            )
+        return _create_mock_response("Final answer after error")
+
+    mock_model = Mock(spec=BaseChatModel)
+    mock_model.ainvoke = _mock_ainvoke
+    context.llm_agent.chat_model = mock_model
+
+
+@then("the accumulated prompt tokens should be greater than {n:d}")
+def step_accumulated_prompt_gt(context, n):
+    pt, ct = context.llm_agent._last_token_usage
+    assert pt > n, f"Expected accumulated prompt_tokens > {n}, got {pt}"
+
+
+@then("the accumulated completion tokens should be greater than {n:d}")
+def step_accumulated_completion_gt(context, n):
+    pt, ct = context.llm_agent._last_token_usage
+    assert ct > n, f"Expected accumulated completion_tokens > {n}, got {ct}"
+
+
+@then("the tool_calling result should be empty")
+def step_result_empty(context):
+    assert not context.result, f"Expected empty result, got {context.result!r}"

src/cleveractors/agents/llm.py

@@ -42,6 +42,7 @@ from cleveractors.core.exceptions import (
     AgentCreationError,
     ConfigurationError,
     ExecutionError,
+    MissingUsageMetadataError,
 )
 from cleveractors.result import MAX_REASONABLE_TOKENS
 from cleveractors.templates.renderer import TemplateRenderer
@@ -221,6 +222,17 @@ class LLMAgent(AgentWithMemory):
         )
         self._pruning_model: str | None = config.get("pruning_model") or None
 
+        # Temporary limitation: pruning_model must match the agent's model
+        # (issue #65).  Support for a separate pruning model is disabled
+        # because multi-model token tracking across separate LangChain
+        # instances is not yet plumbed through NodeUsage.
+        if self._pruning_model is not None and self._pruning_model != self.model:
+            raise ConfigurationError(
+                f"pruning_model ({self._pruning_model}) must match the "
+                f"agent model ({self.model}). Using a different model "
+                f"for pruning is temporarily disabled (issue #65)."
+            )
+
         # pruning_threshold: character length threshold for pruning (§4.4.8 D-2).
         # When raw tool output length is at or below this value, no pruning
         # occurs even when allow_tool_output_pruning is true and tool matches.
@@ -368,22 +380,54 @@ class LLMAgent(AgentWithMemory):
             return self.max_tokens
         return 128_000
 
-    def _estimate_token_count(self, messages: list[Any]) -> int:
-        """Estimate total token count using get_num_tokens or heuristic fallback.
+    def _extract_token_counts(self, response: object) -> tuple[int, int, bool]:
+        """Extract (prompt_tokens, completion_tokens, metadata_present) from an LLM response.
 
-        Note: the heuristic fallback (chars//4) does not count AIMessage.tool_calls
-        JSON. Per ADR D-6 this is a known limitation — tool-call payloads can be
-        hundreds of characters and are excluded from the estimate."""
-        try:
-            _estimator = getattr(self.chat_model, "get_num_tokens", None)
-            if _estimator is not None:
-                return int(_estimator(messages))
-        except Exception:  # nosec B110 — intentional fallback to heuristic
-            pass
-        total_chars = 0
-        for msg in messages:
-            total_chars += len(str(getattr(msg, "content", "")))
-        return max(1, total_chars // 4)
+        Uses the same three-tier fallback chain as the main token extraction
+        in process_message():
+          1. usage_metadata (LangChain standard field)
+          2. response_metadata["token_usage"] (provider-specific)
+          3. (0, 0, False) if neither is available
+
+        The third element (metadata_present) is True when token usage metadata
+        was found and parsed, allowing callers to distinguish "no metadata"
+        from "metadata that resolved to zero/zero".
+        """
+        _prompt: int = 0
+        _completion: int = 0
+        _metadata_present: bool = False
+        _usage_raw: object = getattr(response, "usage_metadata", None)
+        _usage: dict[str, Any] | None = (
+            _usage_raw if isinstance(_usage_raw, dict) else None
+        )
+        if _usage is not None and _usage:
+            _prompt = self._safe_int(_usage.get("input_tokens"), "input_tokens")
+            _completion = self._safe_int(_usage.get("output_tokens"), "output_tokens")
+            _metadata_present = True
+        elif _usage is not None:
+            self._log_no_usage_metadata(_CAUSE_USAGE_METADATA_EMPTY)
+        elif (
+            hasattr(response, "response_metadata")
+            and response.response_metadata is not None
+        ):
+            _rm: object = response.response_metadata
+            if isinstance(_rm, dict):
+                _token_usage: dict[str, Any] = _rm.get("token_usage", {})
+                if _token_usage:
+                    _prompt = self._safe_int(
+                        _token_usage.get("prompt_tokens"), "prompt_tokens"
+                    )
+                    _completion = self._safe_int(
+                        _token_usage.get("completion_tokens"), "completion_tokens"
+                    )
+                    _metadata_present = True
+                else:
+                    self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_NO_TOKEN_USAGE)
+            else:
+                self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_NOT_DICT)
+        else:
+            self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_MISSING)
+        return (_prompt, _completion, _metadata_present)
 
     def _augment_tool_schemas_for_pruning(
         self, tools: list[dict[str, Any]]
@@ -471,14 +515,21 @@ class LLMAgent(AgentWithMemory):
         raw_output: str,
         task_context_msgs: list[Any],
         prune_context: str | None = None,
-    ) -> str:
+    ) -> tuple[str, int, int]:
         """Execute the tool-output pruning extraction pass (§4.4.8).
 
         When *prune_context* is a non-empty string, it replaces the
         broader task context in the extraction prompt.
+
+        Returns:
+            Tuple of (parsed_content, prompt_tokens, completion_tokens)
+            from the pruning LLM call.  On failure returns (raw_output, 0, 0).
         """
         if not raw_output:
-            return raw_output
+            return (raw_output, 0, 0)
+        # Phase 1: Make the pruning LLM call. Failures here fall back to
+        # raw output (non-functional model, network error, missing API key,
+        # invalid model name, template-rendering exception, etc.).
         try:
             prune_model_name = self._pruning_model or self.model
             prune_model = build_chat_model(
@@ -525,15 +576,34 @@ class LLMAgent(AgentWithMemory):
                 )
                 prune_messages = [_SM(content=system_content)] + user_msgs
             prune_response = await prune_model.ainvoke(prune_messages)
-            return self._parse_pruning_response(str(prune_response.content))
         except Exception:
             logger.warning(
-                "Agent %s: pruning pass failed for tool %r; using raw output",
+                "Agent %s: pruning pass LLM call failed for tool %r; using raw output",
                 self.name,
                 tool_name,
                 exc_info=True,
             )
-            return raw_output
+            return (raw_output, 0, 0)
+        # Phase 2: Parse and validate token counts. Missing token usage
+        # metadata is a hard failure so the caller is alerted that billing
+        # data would be incomplete (issue #65).
+        parsed = self._parse_pruning_response(str(prune_response.content))
+        _pp, _pc, _metadata_present = self._extract_token_counts(prune_response)
+        if not _metadata_present:
+            raise MissingUsageMetadataError(
+                f"Pruning model response for tool {tool_name!r} did not "
+                f"include token usage metadata. Token tracking is required "
+                f"when pruning is enabled. Use a model that reports "
+                f"usage_metadata or disable allow_tool_output_pruning."
+            )
+        logger.info(
+            "Agent %s: pruning pass for %r consumed %d prompt + %d completion tokens",
+            self.name,
+            tool_name,
+            _pp,
+            _pc,
+        )
+        return (parsed, _pp, _pc)
 
     # ------------------------------------------------------------------
     # Public credentials property
@@ -751,6 +821,8 @@ class LLMAgent(AgentWithMemory):
                     invoke_kwargs["tools"] = self._lc_tools
 
             _budget_exhausted: bool = False
+            _accumulated_prompt: int = 0
+            _accumulated_completion: int = 0
 
             _declared_names = {
                 t.get("function", {}).get("name", "") for t in (self._lc_tools or [])
@@ -761,6 +833,8 @@ class LLMAgent(AgentWithMemory):
                     break
 
                 # ── Token-budget check (§4.4.7) ─────────────────────────
+                # Uses accumulated actual token counts from prior rounds
+                # instead of heuristic estimation (issue #65).
                 if self._token_budget_percent is not None and _has_tools:
                     _budget_ceiling = max(
                         1,
@@ -769,7 +843,7 @@ class LLMAgent(AgentWithMemory):
                             * self._get_model_context_window()
                         ),
                     )
-                    _est = self._estimate_token_count(messages)
+                    _est = _accumulated_prompt + _accumulated_completion
                     _remaining = max(0, _budget_ceiling - _est)
                     if _est > int(0.75 * _budget_ceiling):
                         logger.warning(
@@ -800,6 +874,9 @@ class LLMAgent(AgentWithMemory):
                         response = await self.chat_model.ainvoke(
                             messages, **invoke_kwargs
                         )
+                        _bp, _bc, _ = self._extract_token_counts(response)
+                        _accumulated_prompt += _bp
+                        _accumulated_completion += _bc
                         _synth_tool_calls: list[dict[str, Any]] = (
                             getattr(response, "tool_calls", None) or []
                         )
@@ -885,12 +962,36 @@ class LLMAgent(AgentWithMemory):
                                         and tool_name in self._pruning_tool_filter
                                         and len(raw_out) > self._pruning_threshold
                                     ):
-                                        tool_output = await self._run_pruning_pass(
-                                            tool_name,
-                                            raw_out,
-                                            list(messages),
-                                            prune_context=_bq_prune_ctx,
-                                        )
+                                        try:
+                                            (
+                                                tool_output,
+                                                _bp_tok,
+                                                _bc_tok,
+                                            ) = await self._run_pruning_pass(
+                                                tool_name,
+                                                raw_out,
+                                                list(messages),
+                                                prune_context=_bq_prune_ctx,
+                                            )
+                                        except MissingUsageMetadataError:
+                                            logger.warning(
+                                                "Agent %s: pruning pass for %r "
+                                                "raised MissingUsageMetadataError; "
+                                                "%d prompt + %d completion tokens "
+                                                "accumulated before failure will be "
+                                                "preserved for billing",
+                                                self.name,
+                                                tool_name,
+                                                _accumulated_prompt,
+                                                _accumulated_completion,
+                                            )
+                                            _captured_prompt = _accumulated_prompt
+                                            _captured_completion = (
+                                                _accumulated_completion
+                                            )
+                                            raise
+                                        _accumulated_prompt += _bp_tok
+                                        _accumulated_completion += _bc_tok
                                         logger.info(
                                             "Agent %s: pruning pass for %r: "
                                             "%d -> %d chars",
@@ -924,11 +1025,17 @@ class LLMAgent(AgentWithMemory):
                                         )
                                     )
                             response = await self.chat_model.ainvoke(messages)
+                            _bp3, _bc3, _ = self._extract_token_counts(response)
+                            _accumulated_prompt += _bp3
+                            _accumulated_completion += _bc3
                         response_text = str(response.content)
                         _budget_exhausted = True
                         break
 
                 response = await self.chat_model.ainvoke(messages, **invoke_kwargs)
+                _mp, _mc, _ = self._extract_token_counts(response)
+                _accumulated_prompt += _mp
+                _accumulated_completion += _mc
 
                 response_tool_calls: list[dict[str, Any]] = (
                     getattr(response, "tool_calls", None) or []
@@ -1016,12 +1123,34 @@ class LLMAgent(AgentWithMemory):
                             and len(raw_tool_output) > self._pruning_threshold
                         ):
                             _task_context = list(messages)
-                            tool_output = await self._run_pruning_pass(
-                                tool_name,
-                                raw_tool_output,
-                                _task_context,
-                                prune_context=_prune_context,
-                            )
+                            try:
+                                (
+                                    tool_output,
+                                    _mlp_tok,
+                                    _mlc_tok,
+                                ) = await self._run_pruning_pass(
+                                    tool_name,
+                                    raw_tool_output,
+                                    _task_context,
+                                    prune_context=_prune_context,
+                                )
+                            except MissingUsageMetadataError:
+                                logger.warning(
+                                    "Agent %s: pruning pass for %r "
+                                    "raised MissingUsageMetadataError; "
+                                    "%d prompt + %d completion tokens "
+                                    "accumulated before failure will be "
+                                    "preserved for billing",
+                                    self.name,
+                                    tool_name,
+                                    _accumulated_prompt,
+                                    _accumulated_completion,
+                                )
+                                _captured_prompt = _accumulated_prompt
+                                _captured_completion = _accumulated_completion
+                                raise
+                            _accumulated_prompt += _mlp_tok
+                            _accumulated_completion += _mlc_tok
                             logger.info(
                                 "Agent %s: pruning pass for %r: %d -> %d chars",
                                 self.name,
@@ -1072,6 +1201,9 @@ class LLMAgent(AgentWithMemory):
                     )
                 )
                 response = await self.chat_model.ainvoke(messages, tools=self._lc_tools)
+                _sp, _sc, _ = self._extract_token_counts(response)
+                _accumulated_prompt += _sp
+                _accumulated_completion += _sc
 
                 # Allow *one* final tool-call round so the model can
                 # write files or perform other last-minute operations
@@ -1165,77 +1297,31 @@ class LLMAgent(AgentWithMemory):
                                 )
                             )
                     response = await self.chat_model.ainvoke(messages)
+                    _sfp, _sfc, _ = self._extract_token_counts(response)
+                    _accumulated_prompt += _sfp
+                    _accumulated_completion += _sfc
 
                 response_text = str(response.content)
 
-            # Extract real token usage from LangChain response metadata (AC2).
-            # Primary source: usage_metadata (LangChain standard field).
-            # Fallback: response_metadata["token_usage"] (provider-specific).
-            # If neither is available: log a warning and use 0.
-            #
-            # Token variables are hoisted before the if/elif/else block so
-            # that the type checker sees a single consistent type regardless
-            # of which branch is taken (N1).
-            _prompt_tokens: int = 0
-            _completion_tokens: int = 0
+            # Use accumulated token counts from all ainvoke() calls
+            # (main model rounds + pruning passes) instead of only the
+            # final response (fix for issue #65 Gap 1 & Gap 2).
+            # If no ainvoke() calls were made (pre-ainvoke failure),
+            # _accumulated_prompt / _accumulated_completion remain 0
+            # which is correct (no tokens consumed).
+            _prompt_tokens = _accumulated_prompt
+            _completion_tokens = _accumulated_completion
 
-            _usage_raw: object = getattr(response, "usage_metadata", None)
-            # Guard: usage_metadata must be a dict (not just any truthy value)
-            # before calling .get() on it.  A provider or test mock returning a
-            # non-dict truthy value (e.g. "1" or [42]) would otherwise raise
-            # AttributeError and discard the LLM response (issue #2).
-            _usage: dict[str, Any] | None = (
-                _usage_raw if isinstance(_usage_raw, dict) else None
-            )
-            if _usage is not None and _usage:
-                # usage_metadata is present and non-empty: extract token counts.
-                # Wrap int() in try/except so a non-numeric provider value
-                # (e.g. a string "abc" or a list) falls back to 0 and logs a
-                # warning instead of propagating an exception that would discard
-                # the LLM response (m1).
-                _prompt_tokens = self._safe_int(
-                    _usage.get("input_tokens"), "input_tokens"
+            # Log a summary of accumulated token usage for observability.
+            if _accumulated_prompt > 0 or _accumulated_completion > 0:
+                logger.info(
+                    "Agent %s: accumulated token usage: "
+                    "%d prompt + %d completion tokens across all rounds",
+                    self.name,
+                    _accumulated_prompt,
+                    _accumulated_completion,
                 )
-                _completion_tokens = self._safe_int(
-                    _usage.get("output_tokens"), "output_tokens"
-                )
-            elif _usage is not None:
-                # usage_metadata is present but empty ({}) — treat as no data
-                # and log a warning (issue #7: empty dict is falsy but not None,
-                # so we must distinguish it from None to avoid silently falling
-                # through to the response_metadata branch).
-                self._log_no_usage_metadata(_CAUSE_USAGE_METADATA_EMPTY)
-            elif (
-                hasattr(response, "response_metadata")
-                and response.response_metadata is not None
-            ):
-                # Guard against response_metadata being a non-dict truthy value
-                # (e.g. a list) — only call .get() when it is actually a dict
-                # (issue #2).
-                _rm: object = response.response_metadata
-                if not isinstance(_rm, dict):
-                    # response_metadata is present but not a dict (e.g. a list).
-                    # Log a distinct cause so operators can distinguish this from
-                    # the "no token_usage key" case (review issue #6).
-                    self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_NOT_DICT)
-                else:
-                    _token_usage: dict[str, Any] = _rm.get("token_usage", {})
-                    if _token_usage:
-                        # Wrap int() in try/except for non-numeric values (m1).
-                        _prompt_tokens = self._safe_int(
-                            _token_usage.get("prompt_tokens"), "prompt_tokens"
-                        )
-                        _completion_tokens = self._safe_int(
-                            _token_usage.get("completion_tokens"), "completion_tokens"
-                        )
-                    else:
-                        # response_metadata is a dict but has no token_usage key.
-                        self._log_no_usage_metadata(
-                            _CAUSE_RESPONSE_METADATA_NO_TOKEN_USAGE
-                        )
-            else:
-                # response_metadata is either absent or explicitly None.
-                self._log_no_usage_metadata(_CAUSE_RESPONSE_METADATA_MISSING)
+
             # Capture token counts immediately after ainvoke() succeeds.
             # Setting the sentinel variables here marks the "ainvoke succeeded"
             # boundary: any exception raised after this point is a post-ainvoke

src/cleveractors/core/exceptions.py

@@ -65,5 +65,16 @@ class ApplicationError(CleverAgentsException):
     """Exception raised for application-level errors."""
 
 
+class MissingUsageMetadataError(CleverAgentsException):
+    """
+    Raised when an LLM response lacks token usage metadata and the
+    caller requires it for accurate billing/token tracking.
+
+    This is raised instead of a generic RuntimeError so that callers
+    can distinguish "no metadata was returned by the provider" from
+    other runtime failures.
+    """
+
+
 class StreamRoutingError(CleverAgentsException):
     """Exception raised for errors in stream routing."""