diff --git a/hindsight-integrations/claude-code/scripts/lib/config.py b/hindsight-integrations/claude-code/scripts/lib/config.py index 1978771d..f1dba531 100644 --- a/hindsight-integrations/claude-code/scripts/lib/config.py +++ b/hindsight-integrations/claude-code/scripts/lib/config.py @@ -29,6 +29,7 @@ DEFAULTS = { "retainRoles": ["user", "assistant"], "retainEveryNTurns": 10, "retainOverlapTurns": 2, + "retainToolCalls": True, "retainContext": "claude-code", "retainTags": [], "retainMetadata": {}, diff --git a/hindsight-integrations/claude-code/scripts/lib/content.py b/hindsight-integrations/claude-code/scripts/lib/content.py index 9a723ab6..fa4ac886 100644 --- a/hindsight-integrations/claude-code/scripts/lib/content.py +++ b/hindsight-integrations/claude-code/scripts/lib/content.py @@ -232,20 +232,80 @@ def format_current_time() -> str: # --------------------------------------------------------------------------- +def _extract_message_blocks(content, role: str = "") -> list: + """Extract structured content blocks from a message for JSON retention. + + Returns a list of dicts, each representing a content block: + - {"type": "text", "text": "..."} for text blocks + - {"type": "tool_use", "name": "...", "input": {...}} for tool calls + - Channel message tool_use blocks get their text extracted inline. + """ + if isinstance(content, str): + cleaned = strip_channel_envelope(strip_memory_tags(content)).strip() + return [{"type": "text", "text": cleaned}] if cleaned else [] + + if not isinstance(content, list): + return [] + + blocks = [] + for block in content: + if not isinstance(block, dict): + continue + block_type = block.get("type", "") + + if block_type == "text": + text = strip_channel_envelope(strip_memory_tags(block.get("text", ""))).strip() + if text: + blocks.append({"type": "text", "text": text}) + + elif block_type == "tool_use" and role == "assistant": + if _is_channel_message_tool(block): + # Channel messages: extract the outgoing text + tool_input = block.get("input", {}) + for field in _MESSAGE_TEXT_FIELDS: + val = tool_input.get(field) + if isinstance(val, str) and val.strip(): + blocks.append({"type": "text", "text": val.strip()}) + break + else: + name = block.get("name", "unknown") + inp = block.get("input", {}) + # Skip Hindsight MCP tools to avoid feedback loops + if name.startswith("mcp__") and _OPERATIONAL_TOOL_PATTERN.search(name.split("__")[-1]): + continue + blocks.append({"type": "tool_use", "name": name, "input": inp}) + + elif block_type == "tool_result": + # Include tool results for context + result_content = block.get("content", "") + if isinstance(result_content, str) and result_content.strip(): + text = result_content.strip() + # Truncate very long results + if len(text) > 2000: + text = text[:2000] + "... (truncated)" + blocks.append({"type": "tool_result", "tool_use_id": block.get("tool_use_id", ""), "content": text}) + + return blocks + + def prepare_retention_transcript( messages: list, retain_roles: list = None, retain_full_window: bool = False, + include_tool_calls: bool = False, ) -> tuple: """Format messages into a retention transcript. - Port of: prepareRetentionTranscript() in index.js + When include_tool_calls is True, outputs JSON with full message structure + including tool calls and their inputs. Otherwise outputs the legacy + text format with [role: ...]...[role:end] markers. Args: messages: List of message dicts with 'role' and 'content'. retain_roles: Roles to include (default: ['user', 'assistant']). retain_full_window: If True, retain all messages (chunked mode). If False, retain only the last turn (last user msg + responses). + include_tool_calls: If True, output JSON format with full tool call data. Returns: (transcript_text, message_count) or (None, 0) if nothing to retain. @@ -267,9 +327,43 @@ def prepare_retention_transcript( target_messages = messages[last_user_idx:] allowed_roles = set(retain_roles or ["user", "assistant"]) + + if include_tool_calls: + return _prepare_json_transcript(target_messages, allowed_roles) + return _prepare_text_transcript(target_messages, allowed_roles) + + +def _prepare_json_transcript(messages: list, allowed_roles: set) -> tuple: + """Format messages as JSON with full tool call data.""" + import json + + structured_messages = [] + for msg in messages: + role = msg.get("role", "unknown") + if role not in allowed_roles: + continue + + blocks = _extract_message_blocks(msg.get("content", ""), role=role) + if not blocks: + continue + + structured_messages.append({"role": role, "content": blocks}) + + if not structured_messages: + return None, 0 + + transcript = json.dumps(structured_messages, indent=None, ensure_ascii=False) + if len(transcript.strip()) < 10: + return None, 0 + + return transcript, len(structured_messages) + + +def _prepare_text_transcript(messages: list, allowed_roles: set) -> tuple: + """Format messages as legacy text with [role:]...[role:end] markers.""" parts = [] - for msg in target_messages: + for msg in messages: role = msg.get("role", "unknown") if role not in allowed_roles: continue diff --git a/hindsight-integrations/claude-code/scripts/retain.py b/hindsight-integrations/claude-code/scripts/retain.py index 77819663..c748a029 100755 --- a/hindsight-integrations/claude-code/scripts/retain.py +++ b/hindsight-integrations/claude-code/scripts/retain.py @@ -127,7 +127,10 @@ def main(): # Format transcript retain_roles = config.get("retainRoles", ["user", "assistant"]) - transcript, message_count = prepare_retention_transcript(messages_to_retain, retain_roles, retain_full_window) + include_tool_calls = config.get("retainToolCalls", True) + transcript, message_count = prepare_retention_transcript( + messages_to_retain, retain_roles, retain_full_window, include_tool_calls=include_tool_calls + ) if not transcript: debug_log(config, "Empty transcript after formatting, skipping retain") diff --git a/hindsight-integrations/claude-code/settings.json b/hindsight-integrations/claude-code/settings.json index 6d3d288a..8cf6210d 100644 --- a/hindsight-integrations/claude-code/settings.json +++ b/hindsight-integrations/claude-code/settings.json @@ -17,6 +17,7 @@ "retainRoles": ["user", "assistant"], "retainEveryNTurns": 10, "retainOverlapTurns": 2, + "retainToolCalls": true, "retainTags": ["{session_id}"], "retainMetadata": {}, "retainContext": "claude-code", diff --git a/hindsight-integrations/claude-code/tests/test_content.py b/hindsight-integrations/claude-code/tests/test_content.py index d9620bfe..807dfa8b 100644 --- a/hindsight-integrations/claude-code/tests/test_content.py +++ b/hindsight-integrations/claude-code/tests/test_content.py @@ -341,3 +341,96 @@ class TestPrepareRetentionTranscript: msgs = [{"role": "assistant", "content": "only assistant"}] result, _ = prepare_retention_transcript(msgs, retain_full_window=False) assert result is None + + def test_json_format_with_tool_calls(self): + """When include_tool_calls=True, output should be JSON with tool_use blocks.""" + import json + + msgs = [ + {"role": "user", "content": "edit the file"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "I'll edit that file."}, + { + "type": "tool_use", + "name": "Edit", + "input": {"file_path": "/tmp/foo.py", "old_string": "old", "new_string": "new"}, + }, + ], + }, + ] + transcript, count = prepare_retention_transcript( + msgs, retain_full_window=True, include_tool_calls=True + ) + assert transcript is not None + data = json.loads(transcript) + assert len(data) == 2 + assert data[0]["role"] == "user" + assert data[1]["role"] == "assistant" + # Should have both text and tool_use blocks + block_types = [b["type"] for b in data[1]["content"]] + assert "text" in block_types + assert "tool_use" in block_types + # Tool input should be preserved + tool_block = next(b for b in data[1]["content"] if b["type"] == "tool_use") + assert tool_block["name"] == "Edit" + assert tool_block["input"]["file_path"] == "/tmp/foo.py" + + def test_json_format_excludes_hindsight_mcp_tools(self): + """Hindsight MCP tools should be excluded even in JSON mode.""" + import json + + msgs = [ + {"role": "user", "content": "recall something"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Let me check."}, + {"type": "tool_use", "name": "mcp__hindsight__recall", "input": {"query": "test"}}, + ], + }, + ] + transcript, _ = prepare_retention_transcript( + msgs, retain_full_window=True, include_tool_calls=True + ) + data = json.loads(transcript) + assistant_blocks = data[1]["content"] + assert len(assistant_blocks) == 1 + assert assistant_blocks[0]["type"] == "text" + + def test_json_format_includes_tool_results(self): + """Tool results should be included in JSON mode.""" + import json + + msgs = [ + {"role": "user", "content": "run ls"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Running ls."}, + {"type": "tool_use", "name": "Bash", "input": {"command": "ls"}}, + ], + }, + { + "role": "assistant", + "content": [ + {"type": "tool_result", "tool_use_id": "123", "content": "file1.py\nfile2.py"}, + {"type": "text", "text": "Here are the files."}, + ], + }, + ] + transcript, _ = prepare_retention_transcript( + msgs, retain_full_window=True, include_tool_calls=True + ) + data = json.loads(transcript) + result_msg = next(m for m in data if any(b.get("type") == "tool_result" for b in m["content"])) + result_block = next(b for b in result_msg["content"] if b["type"] == "tool_result") + assert "file1.py" in result_block["content"] + + def test_without_tool_calls_uses_text_format(self): + """Default (include_tool_calls=False) should use legacy text format.""" + msgs = _msgs(("user", "hello"), ("assistant", "world")) + transcript, _ = prepare_retention_transcript(msgs, retain_full_window=True, include_tool_calls=False) + assert "[role: user]" in transcript + assert "[user:end]" in transcript