diff --git a/aggregate.py b/aggregate.py index 8157880..4da1133 100644 --- a/aggregate.py +++ b/aggregate.py @@ -17,6 +17,7 @@ import glob from pathlib import Path, PureWindowsPath import summarize import project_manager +from file_cache import ASTParser def find_next_increment(output_dir: Path, namespace: str) -> int: pattern = re.compile(rf"^{re.escape(namespace)}_(\d+)\.md$") @@ -255,7 +256,17 @@ def build_tier3_context(file_items: list[dict], screenshot_base_dir: Path, scree sections.append("### `" + (entry or path_str) + "`\n\n" + f"```{path.suffix.lstrip('.') if path and path.suffix else 'text'}\n{item.get('content', '')}\n```") else: - sections.append(summarize.summarise_file(path, item.get("content", ""))) + content = item.get("content", "") + if path and path.suffix == ".py" and not item.get("error"): + try: + parser = ASTParser("python") + skeleton = parser.get_skeleton(content) + sections.append(f"### `{entry or path_str}` (AST Skeleton)\n\n```python\n{skeleton}\n```") + except Exception as e: + # Fallback to summary if AST parsing fails + sections.append(summarize.summarise_file(path, content)) + else: + sections.append(summarize.summarise_file(path, content)) parts.append("## Files (Tier 3 - Focused)\n\n" + "\n\n---\n\n".join(sections)) diff --git a/mcp_client.py b/mcp_client.py index 7d31019..2ca791f 100644 --- a/mcp_client.py +++ b/mcp_client.py @@ -300,7 +300,7 @@ def get_definition(path: str, name: str) -> str: try: import ast - code = p.read_text(encoding="utf-8") + code = p.read_text(encoding="utf-8").lstrip(chr(0xFEFF)) lines = code.splitlines() tree = ast.parse(code) diff --git a/tests/test_tiered_context.py b/tests/test_tiered_context.py index 6ad8712..c673095 100644 --- a/tests/test_tiered_context.py +++ b/tests/test_tiered_context.py @@ -24,13 +24,45 @@ def test_build_tier2_context_exists(): result = build_tier2_context(file_items, Path("."), [], history) assert "Other content" in result +def test_build_tier3_context_ast_skeleton(monkeypatch): + from unittest.mock import MagicMock + import aggregate + import file_cache + + # Mock ASTParser + mock_parser_instance = MagicMock() + mock_parser_instance.get_skeleton.return_value = "def other():\n ..." + mock_parser_class = MagicMock(return_value=mock_parser_instance) + + # Mock file_cache.ASTParser in aggregate module + monkeypatch.setattr("aggregate.ASTParser", mock_parser_class) + + file_items = [ + {"path": Path("other.py"), "entry": "other.py", "content": "def other():\n pass", "error": False} + ] + history = [] + + # New behavior check: it should use ASTParser for .py files not in focus + result = build_tier3_context(file_items, Path("."), [], history, focus_files=[]) + + assert "def other():" in result + assert "..." in result + assert "Python" not in result # summarize.py output should not be there if AST skeleton is used + mock_parser_class.assert_called_once_with("python") + mock_parser_instance.get_skeleton.assert_called_once_with("def other():\n pass") + def test_build_tier3_context_exists(): file_items = [ - {"path": Path("focus.py"), "entry": "focus.py", "content": "Focus content", "error": False}, - {"path": Path("other.py"), "entry": "other.py", "content": "Other content", "error": False} + {"path": Path("focus.py"), "entry": "focus.py", "content": "def focus():\n pass", "error": False}, + {"path": Path("other.py"), "entry": "other.py", "content": "def other():\n pass", "error": False} ] history = ["User: hello"] result = build_tier3_context(file_items, Path("."), [], history, focus_files=["focus.py"]) - assert "Focus content" in result - assert "Other content" not in result + assert "def focus():" in result + assert "pass" in result + # other.py should have skeletonized content, not full "pass" (if get_skeleton works) + # However, for a simple "pass", the skeleton might be the same or similar. + # Let's check for the header + assert "other.py" in result + assert "AST Skeleton" in result