diff --git a/headroom/parser.py b/headroom/parser.py index a7158414e..28560e220 100644 --- a/headroom/parser.py +++ b/headroom/parser.py @@ -48,6 +48,22 @@ def compute_hash(text: str) -> str: return hashlib.md5(text.encode()).hexdigest()[:16] # nosec B324 +def _canonical_call_key(name: str, arguments: Any) -> str: + """Canonical identity for a tool invocation: name + arguments with JSON + key order normalized, so semantically identical calls hash equal even + when the provider serializes arguments differently.""" + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + pass + if isinstance(arguments, (dict, list)): + canon = json.dumps(arguments, sort_keys=True, separators=(",", ":"), default=str) + else: + canon = str(arguments) + return compute_hash(f"{name}\x00{canon}") + + def _extract_tool_result_text(payload: dict[str, Any]) -> str: """Extract text from a tool result payload. @@ -157,6 +173,7 @@ def parse_message_to_blocks( content = message.get("content") if content: tool_result_parts: list[dict[str, Any]] = [] + tool_use_parts: list[dict[str, Any]] = [] if isinstance(content, str): text = content elif isinstance(content, list): @@ -172,6 +189,13 @@ def parse_message_to_blocks( elif isinstance(part, dict) and "toolResult" in part: # Strands/Bedrock converse format; same treatment. tool_result_parts.append(part) + elif isinstance(part, dict) and part.get("type") == "tool_use": + # Anthropic Messages format: call side of the tool unit; + # collect for dedicated tool_call blocks below. + tool_use_parts.append(part) + elif isinstance(part, dict) and "toolUse" in part: + # Strands/Bedrock converse format; same treatment. + tool_use_parts.append(part) elif isinstance(part, str): text_parts.append(part) text = "\n".join(text_parts) @@ -242,6 +266,33 @@ def parse_message_to_blocks( ) blocks.extend(tr_blocks) + for part in tool_use_parts: + payload = part["toolUse"] if "toolUse" in part else part + if not isinstance(payload, dict): + continue + tu_name = payload.get("name") or "unknown" + tu_args = payload.get("input", {}) + tu_id = payload.get("toolUseId") if "toolUse" in part else payload.get("id") + try: + tu_args_text = json.dumps(tu_args, sort_keys=True, default=str) + except (TypeError, ValueError): + tu_args_text = str(tu_args) + tu_text = f"{tu_name}({tu_args_text})" + blocks.append( + Block( + kind="tool_call", + text=tu_text, + tokens_est=tokenizer.count_text(tu_text) + 10, + content_hash=compute_hash(tu_text), + source_index=index, + flags={ + "tool_call_id": tu_id, + "function_name": tu_name, + "call_key": _canonical_call_key(tu_name, tu_args), + }, + ) + ) + # Handle tool calls (assistant messages with tool_calls) tool_calls = message.get("tool_calls") if tool_calls: @@ -259,6 +310,9 @@ def parse_message_to_blocks( flags={ "tool_call_id": tc.get("id"), "function_name": func.get("name"), + "call_key": _canonical_call_key( + func.get("name") or "unknown", func.get("arguments", "") + ), }, ) ) @@ -315,6 +369,7 @@ def parse_messages( # at more than one position means the agent re-fetched something already # in context — an over-compression signal (#853). The first serve is # free; every repeat is counted as waste. + counted_results: set[int] = set() reread_groups: dict[str, list[Block]] = {} for block in all_blocks: if block.kind == "tool_result" and block.tokens_est >= REREAD_MIN_TOKENS: @@ -337,6 +392,44 @@ def parse_messages( prev_index = block.source_index if not is_polling: total_waste.reread_tokens += block.tokens_est + counted_results.add(id(block)) + + # Re-issued-call detection: the agent invoking the same tool with the + # same arguments again is a re-fetch even when the result bytes differ + # (timestamps, mtimes, ordering defeat the content-hash pass above). + # Same polling guard and size floor as above, applied to the repeat + # invocation's result; results the content-hash pass already counted + # are skipped so identical-content repeats are never counted twice. + results_by_call_id: dict[str, Block] = {} + for block in all_blocks: + if block.kind == "tool_result": + tc_id = block.flags.get("tool_call_id") + if tc_id and tc_id not in results_by_call_id: + results_by_call_id[tc_id] = block + + call_groups: dict[str, list[Block]] = {} + for block in all_blocks: + if block.kind == "tool_call": + call_key = block.flags.get("call_key") + if call_key: + call_groups.setdefault(call_key, []).append(block) + + for group in call_groups.values(): + prev_index = group[0].source_index + for block in group: + if block.source_index == prev_index: + continue + is_polling = block.source_index - prev_index <= REREAD_ADJACENT_GAP + prev_index = block.source_index + if is_polling: + continue + result = results_by_call_id.get(block.flags.get("tool_call_id") or "") + if result is None or result.tokens_est < REREAD_MIN_TOKENS: + continue + if id(result) in counted_results: + continue + total_waste.reread_tokens += result.tokens_est + counted_results.add(id(result)) # Compute block breakdown breakdown: dict[str, int] = {} diff --git a/tests/test_parser.py b/tests/test_parser.py index fa38c5249..22fe0424a 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -1051,3 +1051,229 @@ class TestStrandsToolResultBlocks: assert anthropic_waste.total() > 0 assert anthropic_waste.total() == openai_waste.total() + + +# --- TestCallArgMatchReread --- + + +class TestCallArgMatchReread: + """Tests for re-issued-call (arg-match) reread detection in parse_messages.""" + + LARGE_CONTENT = "def handler(event):\n return process(event)\n" * 10 # > 200 chars + CHANGED_CONTENT = LARGE_CONTENT + "# mtime 1718000000\n" + + def _expected_tokens(self, text): + """Mirror mock_tokenizer + message overhead used for tool_result blocks.""" + return len(text) // 4 + 1 + 4 + + @staticmethod + def _filler(n): + """Interleaved turns that push a repeat beyond the polling gap.""" + return [ + {"role": "assistant" if i % 2 == 0 else "user", "content": f"step {i} of the task"} + for i in range(n) + ] + + @staticmethod + def _openai_call(call_id, name, arguments): + return { + "role": "assistant", + "content": None, + "tool_calls": [{"id": call_id, "function": {"name": name, "arguments": arguments}}], + } + + @staticmethod + def _openai_result(call_id, content): + return {"role": "tool", "tool_call_id": call_id, "content": content} + + def test_canonical_call_key_normalizes_serialization(self): + """Reordered JSON-string args, dict args, and spaced JSON hash equal.""" + from headroom.parser import _canonical_call_key + + k1 = _canonical_call_key("read_file", '{"path": "a.py", "lines": 100}') + k2 = _canonical_call_key("read_file", '{"lines":100,"path":"a.py"}') + k3 = _canonical_call_key("read_file", {"path": "a.py", "lines": 100}) + assert k1 == k2 == k3 + assert _canonical_call_key("read_file", '{"path": "b.py", "lines": 100}') != k1 + assert _canonical_call_key("grep", '{"path": "a.py", "lines": 100}') != k1 + + def test_reissued_call_changed_result_counts(self, mock_tokenizer): + """Identical call re-issued far apart counts even when result bytes differ.""" + messages = ( + [ + self._openai_call("c1", "read_file", '{"path": "a.py", "lines": 100}'), + self._openai_result("c1", self.LARGE_CONTENT), + ] + + self._filler(4) + + [ + self._openai_call("c2", "read_file", '{"lines":100,"path":"a.py"}'), + self._openai_result("c2", self.CHANGED_CONTENT), + ] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT) + + def test_identical_result_not_double_counted(self, mock_tokenizer): + """Byte-identical repeat is counted once (content-hash pass wins).""" + messages = ( + [ + self._openai_call("c1", "read_file", '{"path": "a.py"}'), + self._openai_result("c1", self.LARGE_CONTENT), + ] + + self._filler(4) + + [ + self._openai_call("c2", "read_file", '{"path": "a.py"}'), + self._openai_result("c2", self.LARGE_CONTENT), + ] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == self._expected_tokens(self.LARGE_CONTENT) + + def test_adjacent_reissue_is_polling(self, mock_tokenizer): + """Back-to-back identical calls (poll loop) are not re-reads.""" + messages = [ + self._openai_call("c1", "check_ci", '{"run": 7}'), + self._openai_result("c1", self.LARGE_CONTENT), + self._openai_call("c2", "check_ci", '{"run": 7}'), + self._openai_result("c2", self.CHANGED_CONTENT), + ] + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == 0 + + def test_different_args_not_matched(self, mock_tokenizer): + """Same tool with different arguments is not a re-issued call.""" + messages = ( + [ + self._openai_call("c1", "read_file", '{"path": "a.py"}'), + self._openai_result("c1", self.LARGE_CONTENT), + ] + + self._filler(4) + + [ + self._openai_call("c2", "read_file", '{"path": "b.py"}'), + self._openai_result("c2", self.CHANGED_CONTENT), + ] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == 0 + + def test_small_result_ignored(self, mock_tokenizer): + """Repeat of a call whose result is trivially small is skipped.""" + messages = ( + [ + self._openai_call("c1", "run_tests", "{}"), + self._openai_result("c1", "ok"), + ] + + self._filler(4) + + [ + self._openai_call("c2", "run_tests", "{}"), + self._openai_result("c2", "ok again"), + ] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == 0 + + def test_repeat_call_without_result_not_counted(self, mock_tokenizer): + """A re-issued call with no recorded result contributes nothing.""" + messages = ( + [ + self._openai_call("c1", "read_file", '{"path": "a.py"}'), + self._openai_result("c1", self.LARGE_CONTENT), + ] + + self._filler(4) + + [self._openai_call("c2", "read_file", '{"path": "a.py"}')] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == 0 + + def test_anthropic_tool_use_produces_tool_call_blocks(self, mock_tokenizer): + """Anthropic tool_use parts become tool_call blocks with call metadata.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Reading the file now."}, + { + "type": "tool_use", + "id": "t1", + "name": "read_file", + "input": {"path": "a.py"}, + }, + ], + } + ] + blocks, _, _ = parse_messages(messages, mock_tokenizer) + tool_calls = [b for b in blocks if b.kind == "tool_call"] + assert len(tool_calls) == 1 + assert tool_calls[0].flags["function_name"] == "read_file" + assert tool_calls[0].flags["tool_call_id"] == "t1" + assert tool_calls[0].flags["call_key"] + + def test_anthropic_reissued_call_changed_result_counts(self, mock_tokenizer): + """Full Anthropic-format flow: re-issued tool_use with drifted result.""" + + def call(uid): + return { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": uid, "name": "read_file", "input": {"path": "a.py"}} + ], + } + + def result(uid, content): + return { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": uid, "content": content}], + } + + messages = ( + [call("t1"), result("t1", self.LARGE_CONTENT)] + + self._filler(4) + + [call("t2"), result("t2", self.CHANGED_CONTENT)] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT) + + def test_strands_tooluse_matched(self, mock_tokenizer): + """Strands/Bedrock toolUse/toolResult format is matched the same way.""" + + def call(uid): + return { + "role": "assistant", + "content": [{"toolUse": {"toolUseId": uid, "name": "search", "input": {"q": "x"}}}], + } + + def result(uid, content): + return { + "role": "user", + "content": [{"toolResult": {"toolUseId": uid, "content": [{"text": content}]}}], + } + + messages = ( + [call("s1"), result("s1", self.LARGE_CONTENT)] + + self._filler(4) + + [call("s2"), result("s2", self.CHANGED_CONTENT)] + ) + _, _, waste = parse_messages(messages, mock_tokenizer) + assert waste.reread_tokens == self._expected_tokens(self.CHANGED_CONTENT) + + def test_cross_format_call_key_parity(self, mock_tokenizer): + """OpenAI JSON-string args and Anthropic dict input produce the same call_key.""" + openai_msgs = [self._openai_call("c1", "read_file", '{"lines": 100, "path": "a.py"}')] + anthropic_msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "t1", + "name": "read_file", + "input": {"path": "a.py", "lines": 100}, + } + ], + } + ] + o_blocks, _, _ = parse_messages(openai_msgs, mock_tokenizer) + a_blocks, _, _ = parse_messages(anthropic_msgs, mock_tokenizer) + o_key = [b for b in o_blocks if b.kind == "tool_call"][0].flags["call_key"] + a_key = [b for b in a_blocks if b.kind == "tool_call"][0].flags["call_key"] + assert o_key == a_key