diff --git a/.gitignore b/.gitignore index b239c5b0..8529c8f5 100644 --- a/.gitignore +++ b/.gitignore @@ -72,3 +72,4 @@ Thumbs.db *.tmp *.log *.bak +.codewiki/ diff --git "a/CodeWiki\344\273\213\347\273\215.md" "b/CodeWiki\344\273\213\347\273\215.md" new file mode 100644 index 00000000..2f11eccd --- /dev/null +++ "b/CodeWiki\344\273\213\347\273\215.md" @@ -0,0 +1,95 @@ +## CodeWiki：用 AI 为代码仓库自动生成结构化文档的开源框架 + +对于任何一个开发者来说，理解一个陌生代码仓库都是一项挑战。无论是刚加入团队的新人、做代码审查的同事，还是想要复用开源项目的开发者，都需要面对"读代码"这道坎。而代码文档往往是稀缺资源——要么缺失，要么早已过时，与实际代码脱节。 + +CodeWiki 正是为了解决这个问题而生的开源项目。它由 FPT Software AI Center（FSoft）团队开发，已被 ACL 2026 收录，目标是自动为大规模代码仓库生成全面、结构化、架构感知的文档。项目开源地址：[https://github.com/FSoft-AI4Code/CodeWiki](https://github.com/FSoft-AI4Code/CodeWiki) + +### 它解决了什么问题？ + +传统的代码文档工具大多只关注函数和类级别的 docstring 提取，例如 Sphinx、Javadoc 这类工具，它们能生成 API 参考文档，但无法回答更高层次的问题：这个项目的整体架构是什么？模块之间是如何协作的？数据流是怎样的？ + +与此同时，市面上也出现了一些 AI 驱动的代码文档工具，如 DeepWiki 等商业方案。但这些方案大多是闭源的，且在处理大型仓库时存在文档质量参差不齐、对多语言支持不足等问题。 + +CodeWiki 的定位很明确：做一个开源的、支持多语言的、能生成仓库级全景文档的框架，并且在文档质量上对标甚至超越商业方案。 + +### 核心设计理念 + +CodeWiki 的设计围绕三个核心思想展开。 + +第一个是**层次化分解（Hierarchical Decomposition）**。面对一个动辄数十万行代码的大型仓库，CodeWiki 不会一股脑地把所有代码喂给大模型，而是先通过 AST（抽象语法树）解析和依赖分析，将仓库拆解为模块、子模块、文件、类、函数等多层结构。这种分而治之的策略既降低了大模型的上下文压力，也让最终生成的文档具有清晰的层次感。 + +第二个是**多智能体协作（Multi-Agent Processing）**。CodeWiki 内部设计了多个专职 Agent，每个 Agent 负责不同层级或不同类型的文档生成任务。例如有的 Agent 负责分析模块间的依赖关系，有的负责总结某个类的职责，有的负责生成架构图。这些 Agent 各司其职、协同工作，最终将各自的输出汇总成一份完整的文档。 + +第三个是**多模态合成（Multi-Modal Synthesis）**。CodeWiki 生成的文档不仅仅是文字描述，还包括架构图、依赖关系图、模块树等可视化内容。文字与图表相结合，让读者既能从宏观上把握项目全貌，也能深入理解具体模块的实现细节。 + +### 支持的语言和模型 + +CodeWiki 目前支持 8 种主流编程语言：Python、Java、JavaScript、TypeScript、C、C++、C# 和 Kotlin。这个覆盖面已经涵盖了绝大多数企业和开源项目的技术栈。 + +在 LLM 后端方面，CodeWiki 提供了灵活的适配层，支持两种接入模式：API 模式支持 OpenAI 兼容接口（可对接 OpenAI、Azure OpenAI 及任意兼容服务）、Anthropic（Claude）、AWS Bedrock 等主流大模型服务；订阅模式则支持通过 Claude Code 和 Codex 的本地 CLI 运行，无需单独申请 API Key。开发者可以根据自己的偏好和预算选择合适的模型后端。 + +### 技术架构一览 + +从代码结构来看，CodeWiki 采用了前后端分离的架构： + +- **前端（fe/）**：提供 Web 界面，用户可以输入仓库地址、配置参数、查看生成的文档。 +- **后端（be/）**：核心处理引擎，包含多个关键模块： + - `dependency_analyzer/`：负责 AST 解析和模块依赖分析，支持多种语言的语法树解析。 + - `cluster_modules.py`：将分析出的模块进行聚类分组，形成有逻辑层次的模块结构。 + - `agent_tools/`：多 Agent 系统的工具集，定义了各个 Agent 可以调用的能力。 + - `llm_services.py`：LLM 服务的统一抽象层，屏蔽不同模型提供商的差异。 + - `documentation_generator.py`：文档生成器，负责将各 Agent 的输出整合为最终文档。 + - `prompt_template.py`：Prompt 模板管理，确保文档生成的一致性和质量。 + +### 效果如何？与 DeepWiki 的对比 + +为了评估代码文档的质量，CodeWiki 团队还推出了配套的评测基准 **[CodeWikiBench](https://github.com/FSoft-AI4Code/CodeWikiBench)**，涵盖多种编程语言的仓库级文档质量评估。 + +根据论文中的实验数据，CodeWiki 在 CodeWikiBench 上的整体平均得分比 DeepWiki 高出约 **4.73%**，尤其在 Python、JavaScript 等脚本语言上表现突出。在 C/C++ 等系统级语言上，由于代码结构更加复杂（宏定义、指针、内存管理等），DeepWiki 在该类别上略占优势（高出约 3.15%），但 CodeWiki 在整体表现上仍然领先。 + +这个成绩说明 CodeWiki 作为一个开源方案，在文档生成质量上已经可以与商业产品一较高下。 + +### 快速上手 + +安装 CodeWiki 非常简单，前置依赖为 Python 3.12+ 和 Node.js（用于 Mermaid 图表验证），然后运行： + +```bash +pip install git+https://github.com/FSoft-AI4Code/CodeWiki.git +``` + +安装完成后，进入目标项目目录，通过 CLI 命令行直接生成文档： + +```bash +cd /path/to/your/project +codewiki generate --output docs +``` + +也可以通过 Web 界面进行交互式操作，支持配置模型后端、选择目标语言、调整文档粒度等参数。此外，CodeWiki 还支持一些实用的高级功能：使用 `--update` 参数可以只重新生成发生变更的模块，大幅提升大型项目的更新效率；通过 `--github-pages --create-branch` 可以一键生成可部署到 GitHub Pages 的 HTML 文档页面；项目还支持 Docker 容器化部署。 + +### 适用场景 + +CodeWiki 特别适合以下几种场景： + +**新人 Onboarding**。当一个新成员加入团队时，面对一个有几十万行代码的仓库，CodeWiki 可以帮他快速建立对项目的整体认知，了解模块划分和核心逻辑，大幅缩短上手时间。 + +**开源项目维护**。开源项目的文档质量直接影响社区的参与度。用 CodeWiki 自动生成并更新文档，可以降低维护者的文档负担，让更多人愿意参与贡献。 + +**代码审查与重构**。在进行大规模重构或代码审查时，CodeWiki 生成的架构文档可以帮助团队更好地理解模块间的依赖关系，避免"牵一发而动全身"的风险。 + +**技术选型与调研**。当你需要评估一个陌生的开源项目是否适合你的需求时，CodeWiki 可以帮你快速生成一份项目概览，省去大量阅读源码的时间。 + +### 总结 + +值得一提的是，CodeWiki 还支持作为 MCP（Model Context Protocol）服务器运行，可以集成到 Claude Desktop、Cursor 等 AI 编程工具中，让 AI 在理解代码时能够直接参考生成的文档，实现更精准的代码问答和辅助开发。 + +CodeWiki 代表了代码文档自动化领域的一个重要进展。它不是简单的 docstring 提取工具，而是一个能够理解代码架构、生成全景文档的智能系统。作为 ACL 2026 收录的学术工作，它既有扎实的理论基础，又提供了开箱即用的工程实现。对于任何需要理解和维护大型代码仓库的团队来说，CodeWiki 都是一个值得关注和尝试的工具。 + +--- + +**项目信息** + +- GitHub：[https://github.com/FSoft-AI4Code/CodeWiki](https://github.com/FSoft-AI4Code/CodeWiki) +- 官网：[https://fsoft-ai4code.github.io/CodeWiki/](https://fsoft-ai4code.github.io/CodeWiki/) +- 论文：ACL 2026 收录 +- 开发语言：Python（需要 3.12+） +- 支持语言：Python、Java、JavaScript、TypeScript、C、C++、C#、Kotlin diff --git a/IDE_DRIVEN_GUIDE.md b/IDE_DRIVEN_GUIDE.md new file mode 100644 index 00000000..ec673d30 --- /dev/null +++ b/IDE_DRIVEN_GUIDE.md @@ -0,0 +1,377 @@ +# CodeWiki IDE-Driven Mode: Refactoring Process & Usage Guide + +## Background & Motivation + +The original CodeWiki design required users to configure their own LLM API (API Key + base_url), then generate documentation via a one-shot CLI command. This introduced two problems: + +1. **Configuration barrier**: Users need to obtain API Keys, understand provider differences, and handle model compatibility issues +2. **Inflexibility**: The generation process is a black box — users cannot intervene in clustering strategies or documentation style during generation + +**Refactoring goal**: Reduce CodeWiki to a **pure toolchain MCP Server**, fully driven by AI IDE agents (CodeBuddy, Cursor, etc.) to execute the Wiki generation pipeline with **zero LLM configuration**. + +--- + +## Refactoring Process + +### Architecture Analysis + +Through source code analysis, CodeWiki's Wiki generation pipeline depends on LLM in 4 stages: + +| Stage | Code Location | Invocation | LLM Role | +|-------|---------------|------------|----------| +| Module clustering | `cluster_modules.py` | `backend.complete()` | Group components into logical modules | +| Per-module documentation | `pydantic_ai_backend.py` | `agent.run()` multi-turn | Read code, write docs, draw Mermaid diagrams | +| Sub-module recursion | `generate_sub_module_documentations.py` | Sub-agent loop | Recursively handle nested modules | +| Parent module overview | `documentation_generator.py` | `backend.complete()` | Synthesize overviews from child documents | + +Key finding: **dependency analysis (Tree-sitter AST parsing), dependency graph construction, topological sorting, and Mermaid validation** — the core toolchain — requires no LLM at all. + +### Refactoring Strategy + +Transform the MCP Server from "black-box one-shot generation" into a "fine-grained toolset": + +``` +Before refactoring: + IDE → generate_docs(repo) → [CodeWiki internally calls LLM] → result + +After refactoring: + IDE Agent → analyze_repo → read_code → (Agent reasons clustering) → write_doc → overview + ↑ Pure tool call ↑ Pure tool call ↑ IDE's own LLM ↑ Pure tool call +``` + +### File Side-Channel Architecture + +A key design decision in the refactoring: instead of transmitting large payloads (component indexes, source code, processing order) through the MCP stdio channel — which required aggressive truncation and caused overflow errors — the server writes all bulky data to **per-session workspace files** on disk. The MCP response returns only file paths and a compact summary. The IDE agent then reads those files directly using its own file-access capabilities. + +This approach eliminates truncation limits entirely: component indexes, source code files, and processing orders are written in full, no matter how large the repository. + +### New File Inventory + +``` +codewiki/mcp/ +├── server.py # Refactored: 10 tool registrations (8 fine-grained + 2 legacy) +├── session.py # Session state management (SessionStore, thread-safe) +├── workspace.py # Per-session file workspace (write/read/cleanup) +└── tools/ + ├── __init__.py # Tool package entry point + ├── analysis.py # analyze_repo with incremental change detection + ├── code_reader.py # read_code_components (writes .src files to workspace) + ├── doc_writer.py # write_doc_file + edit_doc_file (with path traversal guards) + ├── module_tree.py # save_module_tree + get_processing_order + └── prompt_server.py # get_prompt template service +``` + +### MCP Toolset + +The server exposes **8 fine-grained tools** (zero LLM config) plus **2 legacy tools**: + +| Tool | Purpose | Data Flow | Requires LLM | +|------|---------|-----------|:---:| +| `analyze_repo` | Parse repo, build dependency graph, detect incremental changes | Writes workspace files (component index, leaf nodes, languages, changes), returns paths + stats | No | +| `read_code_components` | Write component source code to workspace `.src` files | Each component → `sources/{sanitized_id}.src`, returns file paths | No | +| `write_doc_file` | Create .md documents with auto Mermaid validation | Writes file directly to output dir | No | +| `edit_doc_file` | Edit documents: `str_replace` / `insert` / `undo` | Modifies file in place, keeps edit history (capped at 20/file) | No | +| `save_module_tree` | Persist IDE agent's module clustering | Writes `module_tree.json` + `first_module_tree.json` + `processing_order.json` | No | +| `get_processing_order` | Compute leaf-first processing order | Writes `processing_order.json` to workspace, returns path | No | +| `get_prompt` | Retrieve prompt templates for each pipeline stage | Returns inline (small payload) | No | +| `close_session` | Write `metadata.json`, clean up workspace files, free memory | Cleans workspace dir + prunes empty parent dirs | No | +| `generate_docs` | [Legacy] One-shot generation (requires `codewiki config set`) | Full pipeline | **Yes** | +| `get_module_tree` | [Legacy] Get existing module clustering tree | Reads from disk | No | + +### Thread Safety & Concurrency + +Synchronous tool handlers (file I/O, Tree-sitter parsing) run via `asyncio.to_thread()` to prevent blocking the MCP stdio event loop. The exception is `analyze_repo` — Tree-sitter C extensions are not thread-safe, so it runs on the main thread (acceptable for a one-time heavy operation). + +Session management is fully thread-safe: `SessionStore` uses a mutex lock for all read/write operations, supports up to **10 concurrent sessions** (oldest evicted at capacity), and sessions auto-expire after **2 hours** of inactivity. + +### Security Hardening + +The `doc_writer` module enforces path traversal guards: all file paths are resolved and verified to stay within the session's `output_dir`. Filenames that attempt directory escape are rejected. Edit operations are tracked in session-scoped history (capped at 20 entries per file to prevent unbounded memory growth). + +### Backward Compatibility + +- Existing CLI (`codewiki generate`, `codewiki config`) remains completely unchanged +- Existing Web App remains completely unchanged +- Legacy MCP tools (`generate_docs`, `get_module_tree`) are preserved — users with configured LLMs can still use them +- The `codewiki/__init__.py` unconditional CLI import was removed, so MCP Server can now start without installing CLI-specific dependencies (`keyring`, `click`, etc.) + +--- + +## Usage + +### Prerequisites + +```bash +# 1. Clone the project +git clone https://github.com/mambo-wang/CodeWiki-CN.git +cd CodeWiki-CN + +# 2. Install dependencies +pip install -e . + +# 3. Verify +python -c "from codewiki.mcp.server import server; print('MCP Server OK')" +``` + +### CodeBuddy Configuration + +**Step 1**: Configure the MCP Server in CodeBuddy. + +Add to CodeBuddy's MCP configuration: + +```json +{ + "mcpServers": { + "codewiki": { + "command": "python", + "args": ["-m", "codewiki.mcp.server"], + "cwd": "/path/to/CodeWiki-CN" + } + } +} +``` + +**Step 2**: Project rules are automatically configured in `.codebuddy/rules/codewiki-wiki-generator/RULE.mdc`. When you mention "generate documentation" or "Wiki" in Agent mode, CodeBuddy automatically loads this rule. + +**Step 3**: Open CodeBuddy Agent mode and enter: + +``` +Analyze this repository and generate Wiki documentation for me +``` + +### Cursor Configuration + +**Step 1**: Add the Server in Cursor Settings → MCP: + +```json +{ + "mcpServers": { + "codewiki": { + "command": "python", + "args": ["-m", "codewiki.mcp.server"], + "cwd": "/path/to/CodeWiki-CN" + } + } +} +``` + +**Step 2**: Project rules are configured in `.cursorrules` and automatically loaded when Cursor opens the project. + +**Step 3**: In Cursor Agent mode, enter: + +``` +Please generate Wiki documentation for the current repository, output to the docs directory. +``` + +### Claude Desktop Configuration + +Add to `~/Library/Application Support/Claude/claude_desktop_config.json` (macOS): + +```json +{ + "mcpServers": { + "codewiki": { + "command": "python", + "args": ["-m", "codewiki.mcp.server"], + "cwd": "/path/to/CodeWiki-CN" + } + } +} +``` + +### Other MCP-Capable IDEs + +Any AI IDE supporting the MCP stdio protocol can be used with similar configuration — specify `command: python`, `args: ["-m", "codewiki.mcp.server"]`. + +--- + +## IDE Agent Workflow + +When you trigger Wiki generation in an AI IDE, the Agent works through the following 5 phases: + +``` +Phase 1: analyze_repo + │ → Get session_id, workspace_dir, stats, file paths + │ → Read workspace files: component_index.json, leaf_nodes.json, languages.json + │ +Phase 2: get_prompt("cluster") + read_code_components + save_module_tree + │ → Agent reasons independently, groups components into 3-8 logical modules + │ → Source code written to workspace sources/ dir, agent reads .src files directly + │ → Get leaf-first processing order from processing_order.json + │ +Phase 3: Per-module generation + │ For each leaf module: + │ ├── get_prompt("system_leaf") → Get documentation writing instructions + │ ├── read_code_components → Source written to sources/*.src, read directly + │ └── write_doc_file → Write .md (auto Mermaid validation) + │ + │ For each parent module: + │ ├── Read child module .md files + │ ├── get_prompt("overview_module") → Get overview instructions + │ └── write_doc_file → Write overview + │ +Phase 4: get_prompt("overview_repo") → Generate repository overview overview.md + │ +Phase 5: close_session → Write metadata.json, clean up workspace, release resources +``` + +--- + +## Incremental Updates + +### Problems with the Original `--update` + +The original CodeWiki CLI provided a `codewiki generate --update` incremental update command, but had a bug: the CLI adapter did not pass `commit_id` when creating `DocumentationGenerator`, causing `commit_id` in `metadata.json` to always be `null`. When `_detect_changed_files()` reads `null`, it falls back to full generation. Only Web mode (`background_worker.py`) correctly writes `commit_id`, so under CLI, `--update` is effectively **always equivalent to full generation**. + +This was fixed: the CLI adapter now passes `commit_id` correctly, and the MCP `close_session` tool writes `metadata.json` (with current git commit + timestamp) before cleaning up the workspace, establishing the baseline for future incremental detection. + +### MCP Incremental Update Solution + +Incremental detection is built into the `analyze_repo` tool with a dual-strategy approach: + +``` +First call to analyze_repo: + → Generate full documentation (changes field is null) + → close_session writes metadata.json with commit_id + timestamp + +Subsequent call to analyze_repo after code changes: + → Automatically detect changes via git diff or mtime comparison + → Return changes field with affected_modules + cascade_modules + → AI Agent only updates affected module documentation +``` + +**Change detection strategies** (by priority): + +1. **Git strategy**: Read `commit_id` from `metadata.json`, run `git diff` against current HEAD, also check `git status` to capture uncommitted changes (modified + untracked files) +2. **Mtime strategy** (fallback for non-git repos): Walk source files and compare modification times against `timestamp` in `metadata.json` + +**Return structure**: + +```json +{ + "changes": { + "has_previous": true, + "no_changes": false, + "method": "git", + "changed_files": ["auth.py"], + "affected_modules": ["Authentication Module"], + "cascade_modules": ["Core System", "overview"], + "hint": "Only 1 module(s) need updating: ..." + } +} +``` + +- `affected_modules`: Directly affected modules that need documentation updates +- `cascade_modules`: Indirectly affected parent modules (child docs changed, so overviews must refresh) and `overview` + +### Agent Incremental Update Workflow + +When `analyze_repo` returns `changes` with `no_changes: false`, the Agent executes: + +``` +1. Only process modules in affected_modules: + ├── read_code_components → Read changed component source code from workspace + └── edit_doc_file(str_replace) → Partially modify documentation (instead of full rewrite) + +2. Process parent modules in cascade_modules: + ├── Read updated child documents + └── edit_doc_file → Refresh overview sections + +3. Finally update overview.md +``` + +Compared to the 5-phase full generation workflow, incremental updates typically only need to process 1-3 modules, significantly reducing time. + +### Implementation Details + +Core code is in `codewiki/mcp/tools/analysis.py`, with 4 dedicated functions (~170 lines): + +| Function | Responsibility | +|----------|----------------| +| `_detect_changes()` | Main entry point, coordinates git/mtime strategies, calls module mapping | +| `_detect_via_git()` | Git detection: commit diff + uncommitted changes (modified + untracked) | +| `_detect_via_mtime()` | Mtime fallback: walk source files, compare mtime against generation timestamp | +| `_find_affected_modules()` | Substring matching changed files → module mapping (reuses original CLI logic) | + +`handle_analyze_repo()` calls `_detect_changes()` after building the component index, appending results to the `changes` field in the return JSON. On first run (no existing docs), `changes` is `null`, behavior is identical to before. + +--- + +## Workspace Lifecycle + +Each `analyze_repo` call creates a session workspace at `{repo_path}/.codewiki/sessions/{session_id}/`: + +``` +.codewiki/sessions/{session_id}/ +├── component_index.json # Full component index (id, type, file for each) +├── leaf_nodes.json # Complete leaf node ID list +├── languages.json # Language statistics +├── changes.json # Incremental change info (optional) +├── summary.json # Compact analysis summary +├── processing_order.json # Leaf-first generation order (after save_module_tree) +└── sources/ + └── {sanitized_id}.src # Individual component source files +``` + +The workspace is automatically cleaned up when `close_session` is called. Empty parent directories are pruned as well. Sessions that expire (2-hour TTL) or are evicted (max 10 concurrent) also trigger workspace cleanup. + +--- + +## Output Structure + +The generated documentation structure is consistent with the original CodeWiki: + +``` +docs/ +├── overview.md # Repository overview (start reading here) +├── module1.md # Individual module documentation +├── module2.md # ... +├── module_tree.json # Module hierarchy structure +├── first_module_tree.json # Initial clustering result (immutable snapshot) +└── metadata.json # Generation metadata (commit_id + timestamp) +``` + +--- + +## Original CLI Mode (Still Available) + +If you prefer one-shot command-line generation, the original method is completely unaffected: + +```bash +# Configure LLM +codewiki config set \ + --provider openai-compatible \ + --api-key YOUR_KEY \ + --base-url https://api.example.com \ + --main-model claude-sonnet-4 + +# One-shot generation +codewiki generate +``` + +See the Quick Start section in [README.md](README.md) for details. + +--- + +## FAQ + +**Q: MCP Server fails to start with missing dependencies?** +A: Make sure you have run `pip install -e .` to install CodeWiki and its dependencies. The MCP Server no longer requires CLI-specific packages like `keyring` or `click`. + +**Q: analyze_repo is slow?** +A: Tree-sitter parsing for large repositories (>100K lines) takes some time, usually completing within 30 seconds. Use `include_patterns` / `exclude_patterns` to narrow the analysis scope. There are no component count or source code length truncation limits. + +**Q: Mermaid validation errors?** +A: The Agent will automatically correct syntax based on validation results. If failures persist, check that `mermaid-py` is properly installed. + +**Q: How to have the Agent write documentation in a specific language?** +A: Specify explicitly in the conversation: "Please generate the Wiki documentation in English." or "Please use Chinese for the documentation." + +**Q: What to do when a session times out?** +A: Sessions default to a 2-hour TTL with a maximum of 10 concurrent sessions. After timeout or eviction, simply re-call `analyze_repo` to create a new session. + +**Q: How to incrementally update documentation after code changes?** +A: Simply tell the AI Agent "update the Wiki documentation". When the Agent calls `analyze_repo`, it automatically detects changes and the returned `changes` field indicates which modules are affected. The Agent only updates affected module documentation instead of regenerating everything. Supports both git and non-git repository detection. + +**Q: What is the granularity of incremental updates?** +A: Module-level. If any component's source file in a module changes, that module's entire documentation is marked for update. Its parent module's overview is also marked (cascading update). `overview.md` is refreshed whenever any change occurs. diff --git a/codewiki/__init__.py b/codewiki/__init__.py index 77f63b9a..e078a354 100644 --- a/codewiki/__init__.py +++ b/codewiki/__init__.py @@ -1,14 +1,13 @@ """ CodeWiki: Transform codebases into comprehensive documentation using AI-powered analysis. -This package provides a CLI tool for generating documentation from code repositories. +This package provides a CLI tool for generating documentation from code repositories, +and an MCP server for IDE-driven documentation generation. """ __version__ = "1.0.1" __author__ = "CodeWiki Contributors" __license__ = "MIT" -from codewiki.cli.main import cli - -__all__ = ["cli", "__version__"] +__all__ = ["__version__"] diff --git a/codewiki/cli/adapters/doc_generator.py b/codewiki/cli/adapters/doc_generator.py index 61b1c1b7..cbda31ca 100644 --- a/codewiki/cli/adapters/doc_generator.py +++ b/codewiki/cli/adapters/doc_generator.py @@ -37,7 +37,8 @@ def __init__( output_dir: Path, config: Dict[str, Any], verbose: bool = False, - generate_html: bool = False + generate_html: bool = False, + commit_id: str = None, ): """ Initialize the CLI documentation generator. @@ -48,12 +49,14 @@ def __init__( config: LLM configuration verbose: Enable verbose output generate_html: Whether to generate HTML viewer + commit_id: Git commit SHA for incremental update tracking """ self.repo_path = repo_path self.output_dir = output_dir self.config = config self.verbose = verbose self.generate_html = generate_html + self.commit_id = commit_id self.progress_tracker = ProgressTracker(total_stages=5, verbose=verbose) self.job = DocumentationJob() @@ -178,7 +181,7 @@ async def _run_backend_generation(self, backend_config: BackendConfig): self.progress_tracker.update_stage(0.2, "Initializing dependency analyzer...") # Create documentation generator - doc_generator = DocumentationGenerator(backend_config) + doc_generator = DocumentationGenerator(backend_config, commit_id=self.commit_id) if self.verbose: self.progress_tracker.update_stage(0.5, "Parsing source files...") diff --git a/codewiki/cli/commands/generate.py b/codewiki/cli/commands/generate.py index 1c370cb8..d8c9afe8 100644 --- a/codewiki/cli/commands/generate.py +++ b/codewiki/cli/commands/generate.py @@ -525,6 +525,8 @@ def generate_command( agent_instructions_dict = config.agent_instructions.to_dict() # Create generator + # Get commit_id early so it can be stored in metadata.json for --update support + commit_id = get_git_commit_hash(repo_path) generator = CLIDocumentationGenerator( repo_path=repo_path, output_dir=output_dir, @@ -545,7 +547,8 @@ def generate_command( 'max_depth': max_depth if max_depth is not None else config.max_depth, }, verbose=verbose, - generate_html=github_pages + generate_html=github_pages, + commit_id=commit_id, ) # Run generation @@ -556,7 +559,6 @@ def generate_command( # Get repository info repo_url = None - commit_hash = get_git_commit_hash(repo_path) current_branch = get_git_branch(repo_path) if is_git_repository(repo_path): diff --git a/codewiki/mcp/server.py b/codewiki/mcp/server.py index d54539ea..e0b708ca 100644 --- a/codewiki/mcp/server.py +++ b/codewiki/mcp/server.py @@ -1,16 +1,30 @@ """ CodeWiki MCP Server. -Exposes documentation generation as MCP tools: - - generate_docs: Generate full documentation for a repository - - analyze_repo: Analyze repository structure and dependencies - - get_module_tree: Get the module clustering for a repository +Provides two sets of tools: + +**Fine-grained tools (IDE-driven, zero LLM config):** + - ``analyze_repo`` — Parse a repo and build a dependency graph (session-based) + - ``read_code_components`` — Write component source code to workspace files + - ``write_doc_file`` — Create a documentation .md file with Mermaid validation + - ``edit_doc_file`` — Edit a documentation file (str_replace / insert / undo) + - ``save_module_tree`` — Persist IDE agent's module clustering + - ``get_processing_order`` — Get leaf-first documentation order + - ``get_prompt`` — Retrieve CodeWiki's prompt templates + - ``close_session`` — Clean up a session and workspace files + +Large analysis results (component index, source code, processing order) are +written to workspace files on disk. The IDE agent reads these files directly +instead of receiving large payloads through the MCP stdio channel. + +**Legacy tools (require CodeWiki LLM config):** + - ``generate_docs`` — Full documentation generation (black-box) + - ``get_module_tree`` — Retrieve existing module clustering Usage: - # Run as standalone MCP server (stdio transport) python -m codewiki.mcp.server - # Or register in your MCP client config: + # Cursor / Claude Desktop config: { "mcpServers": { "codewiki": { @@ -30,54 +44,55 @@ from mcp.server import Server from mcp.server.stdio import stdio_server -from mcp.types import ( - TextContent, - Tool, -) +from mcp.types import TextContent, Tool + +from codewiki.mcp.session import SessionState, SessionStore logger = logging.getLogger(__name__) -# Create the MCP server +# --------------------------------------------------------------------------- +# Global session store (lives for the lifetime of the MCP server process) +# --------------------------------------------------------------------------- +_store = SessionStore() + +# --------------------------------------------------------------------------- +# MCP Server instance +# --------------------------------------------------------------------------- server = Server("codewiki") -def _load_config(): - """Load CodeWiki configuration from ~/.codewiki/config.json + keyring.""" - from codewiki.cli.config_manager import ConfigManager - manager = ConfigManager() - if not manager.load(): - raise RuntimeError( - "CodeWiki not configured. Run 'codewiki config set' first." - ) - return manager - +# =================================================================== +# Tool definitions +# =================================================================== -@server.list_tools() -async def list_tools() -> list[Tool]: - """List available CodeWiki MCP tools.""" +def _fine_grained_tools() -> list[Tool]: + """Return the zero-config, IDE-driven tool set.""" return [ Tool( - name="generate_docs", + name="analyze_repo", description=( - "Generate comprehensive AI-powered documentation for a code repository. " - "Analyzes dependencies, clusters modules, and generates markdown documentation." + "Analyze a code repository's structure, dependencies, and components " + "using Tree-sitter AST parsing. No LLM required. " + "Writes the full component index, leaf nodes, and language stats to " + "workspace files on disk, and returns file paths plus a compact summary. " + "Read the workspace files for complete data. " + "This is the entry point for the wiki generation pipeline. " + "After calling this, use get_prompt('cluster') to learn clustering rules, " + "then save_module_tree to persist your grouping. " + "INCREMENTAL UPDATE: If docs already exist in output_dir (metadata.json + " + "module_tree.json), the response includes a 'changes' field showing which " + "files changed and which modules need updating." ), inputSchema={ "type": "object", "properties": { "repo_path": { "type": "string", - "description": "Absolute path to the repository to document", + "description": "Absolute path to the repository to analyze", }, "output_dir": { "type": "string", - "description": "Output directory for generated docs (default: ./docs)", - "default": "docs", - }, - "doc_type": { - "type": "string", - "enum": ["api", "architecture", "user-guide", "developer"], - "description": "Type of documentation to generate", + "description": "Output directory for generated docs (default: /docs)", }, "include_patterns": { "type": "string", @@ -92,18 +107,218 @@ async def list_tools() -> list[Tool]: }, ), Tool( - name="analyze_repo", + name="read_code_components", + description=( + "Write the source code for a list of component IDs to workspace files. " + "Component IDs have the form 'file_path::ComponentName'. " + "Each component's full source is written to an individual .src file " + "in the session's sources/ directory. Returns file paths — no truncation." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID from analyze_repo", + }, + "component_ids": { + "type": "array", + "items": {"type": "string"}, + "description": "List of component IDs to read", + }, + }, + "required": ["session_id", "component_ids"], + }, + ), + Tool( + name="write_doc_file", + description=( + "Create a new markdown documentation file in the output directory. " + "Automatically validates Mermaid diagrams after writing." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID from analyze_repo", + }, + "filename": { + "type": "string", + "description": "Filename for the doc (e.g., 'auth_module.md')", + }, + "content": { + "type": "string", + "description": "Markdown content to write", + }, + }, + "required": ["session_id", "filename", "content"], + }, + ), + Tool( + name="edit_doc_file", + description=( + "Edit an existing documentation file. Supports str_replace (find-and-replace), " + "insert (add text at a line), and undo (revert last edit). " + "Automatically validates Mermaid diagrams after editing." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID from analyze_repo", + }, + "filename": { + "type": "string", + "description": "Filename of the doc to edit", + }, + "command": { + "type": "string", + "enum": ["str_replace", "insert", "undo"], + "description": "Edit command to run", + }, + "old_str": { + "type": "string", + "description": "String to find (required for str_replace)", + }, + "new_str": { + "type": "string", + "description": "Replacement string (for str_replace/insert)", + }, + "insert_line": { + "type": "integer", + "description": "Line number for insert (0-indexed)", + }, + }, + "required": ["session_id", "filename", "command"], + }, + ), + Tool( + name="save_module_tree", + description=( + "Save the IDE agent's module clustering result. " + "Accepts a JSON module tree and persists it to disk. " + "Computes the leaf-first processing order and writes it to a workspace file. " + "Returns the file path for the processing order." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID from analyze_repo", + }, + "module_tree": { + "type": "object", + "description": ( + "Module tree dict. Each key is a module name with value " + "{'components': [component_ids], 'children': {nested modules}}" + ), + }, + }, + "required": ["session_id", "module_tree"], + }, + ), + Tool( + name="get_processing_order", + description=( + "Compute and write the leaf-first processing order to a workspace file. " + "Returns the file path. Process leaf modules (is_leaf=true) before parent modules." + ), + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID from analyze_repo", + }, + }, + "required": ["session_id"], + }, + ), + Tool( + name="get_prompt", + description=( + "Retrieve CodeWiki's prompt templates for each pipeline stage. " + "Available types: cluster, system_complex, system_leaf, user, " + "overview_module, overview_repo. Optionally pass variables to " + "fill in template placeholders." + ), + inputSchema={ + "type": "object", + "properties": { + "prompt_type": { + "type": "string", + "enum": [ + "cluster", + "system_complex", + "system_leaf", + "user", + "overview_module", + "overview_repo", + ], + "description": "Which prompt template to retrieve", + }, + "variables": { + "type": "object", + "description": "Optional template variables to fill in", + }, + }, + "required": ["prompt_type"], + }, + ), + Tool( + name="close_session", + description="Close and clean up an analysis session to free memory.", + inputSchema={ + "type": "object", + "properties": { + "session_id": { + "type": "string", + "description": "Session ID to close", + }, + }, + "required": ["session_id"], + }, + ), + ] + + +def _legacy_tools() -> list[Tool]: + """Return the legacy tools that require CodeWiki LLM configuration.""" + return [ + Tool( + name="generate_docs", description=( - "Analyze a repository's structure, dependencies, and component hierarchy " - "without generating full documentation. Returns file counts, languages, " - "and dependency information." + "[LEGACY — requires 'codewiki config set' first] " + "Generate full documentation for a repository in one shot. " + "For IDE-driven generation, use the fine-grained tools instead." ), inputSchema={ "type": "object", "properties": { "repo_path": { "type": "string", - "description": "Absolute path to the repository to analyze", + "description": "Absolute path to the repository to document", + }, + "output_dir": { + "type": "string", + "description": "Output directory for generated docs (default: ./docs)", + "default": "docs", + }, + "doc_type": { + "type": "string", + "enum": ["api", "architecture", "user-guide", "developer"], + "description": "Type of documentation to generate", + }, + "include_patterns": { + "type": "string", + "description": "Comma-separated file patterns to include", + }, + "exclude_patterns": { + "type": "string", + "description": "Comma-separated patterns to exclude", }, }, "required": ["repo_path"], @@ -111,10 +326,7 @@ async def list_tools() -> list[Tool]: ), Tool( name="get_module_tree", - description=( - "Get the module clustering tree for a repository. " - "Shows how source files are grouped into logical modules." - ), + description="Get the existing module clustering tree for a repository.", inputSchema={ "type": "object", "properties": { @@ -134,32 +346,108 @@ async def list_tools() -> list[Tool]: ] +# =================================================================== +# Tool dispatch +# =================================================================== + +@server.list_tools() +async def list_tools() -> list[Tool]: + """List all available CodeWiki MCP tools.""" + return _fine_grained_tools() + _legacy_tools() + + @server.call_tool() async def call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]: - """Handle MCP tool calls.""" + """Route tool calls to the appropriate handler.""" try: - if name == "generate_docs": - return await _handle_generate_docs(arguments) - elif name == "analyze_repo": - return await _handle_analyze_repo(arguments) + # --- Fine-grained tools (no LLM config needed) --- + # Synchronous handlers run via asyncio.to_thread() so they never + # block the event loop (which would hang the MCP stdio server). + if name == "analyze_repo": + from codewiki.mcp.tools.analysis import handle_analyze_repo + # NOTE: Tree-sitter C extensions are not thread-safe, so this + # must run on the main thread (blocking the event loop is + # acceptable for this one-time heavy operation). + return [_text(handle_analyze_repo(arguments, _store))] + + elif name == "read_code_components": + from codewiki.mcp.tools.code_reader import handle_read_code_components + return [_text(await asyncio.to_thread(handle_read_code_components, arguments, _store))] + + elif name == "write_doc_file": + from codewiki.mcp.tools.doc_writer import handle_write_doc_file + result = await handle_write_doc_file(arguments, _store) + return [_text(result)] + + elif name == "edit_doc_file": + from codewiki.mcp.tools.doc_writer import handle_edit_doc_file + result = await handle_edit_doc_file(arguments, _store) + return [_text(result)] + + elif name == "save_module_tree": + from codewiki.mcp.tools.module_tree import handle_save_module_tree + return [_text(await asyncio.to_thread(handle_save_module_tree, arguments, _store))] + + elif name == "get_processing_order": + from codewiki.mcp.tools.module_tree import handle_get_processing_order + return [_text(await asyncio.to_thread(handle_get_processing_order, arguments, _store))] + + elif name == "get_prompt": + from codewiki.mcp.tools.prompt_server import handle_get_prompt + return [_text(await asyncio.to_thread(handle_get_prompt, arguments, _store))] + + elif name == "close_session": + sid = arguments["session_id"] + session = _store.get(sid) + if session: + _write_generation_metadata(session) + # Clean up workspace files on disk + if session.workspace is not None: + session.workspace.cleanup() + removed = _store.remove(sid) + return [_text(json.dumps({ + "status": "closed" if removed else "not_found", + "session_id": sid, + }))] + + # --- Legacy tools (require CodeWiki LLM config) --- + elif name == "generate_docs": + return await _legacy_generate_docs(arguments) + elif name == "get_module_tree": - return await _handle_get_module_tree(arguments) + return await _legacy_get_module_tree(arguments) + else: - return [TextContent(type="text", text=f"Unknown tool: {name}")] + return [_text(json.dumps({"error": f"Unknown tool: {name}"}))] + except Exception as e: logger.error("Tool %s failed: %s", name, e, exc_info=True) - return [TextContent(type="text", text=f"Error: {e}")] + return [_text(json.dumps({"error": str(e)}))] + + +# =================================================================== +# Legacy tool handlers (require _load_config) +# =================================================================== + +def _load_config(): + """Load CodeWiki configuration from ~/.codewiki/config.json + keyring.""" + from codewiki.cli.config_manager import ConfigManager + manager = ConfigManager() + if not manager.load(): + raise RuntimeError( + "CodeWiki not configured. Run 'codewiki config set' first." + ) + return manager -async def _handle_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: - """Handle generate_docs tool call.""" +async def _legacy_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: + """Legacy generate_docs — requires CodeWiki LLM configuration.""" repo_path = Path(arguments["repo_path"]).expanduser().resolve() output_dir = Path(arguments.get("output_dir", "docs")).expanduser().resolve() if not repo_path.exists(): - return [TextContent(type="text", text=f"Repository not found: {repo_path}")] + return [_text(json.dumps({"error": f"Repository not found: {repo_path}"}))] - # Load config manager = _load_config() config = manager.get_config() api_key = manager.get_api_key() @@ -167,9 +455,8 @@ async def _handle_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: from codewiki.src.be.backend import is_caw_provider caw_mode = bool(config) and is_caw_provider(getattr(config, "provider", "")) if not api_key and not caw_mode: - return [TextContent(type="text", text="API key not configured. Run 'codewiki config set --api-key '")] + return [_text(json.dumps({"error": "API key not configured. Run 'codewiki config set --api-key '"}))] - # Build agent instructions from arguments agent_instructions = {} if arguments.get("doc_type"): agent_instructions["doc_type"] = arguments["doc_type"] @@ -197,11 +484,8 @@ async def _handle_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: from codewiki.src.be.documentation_generator import DocumentationGenerator doc_gen = DocumentationGenerator(backend_config) - - # Run generation await doc_gen.run() - # Collect results generated_files = [] for f in output_dir.iterdir(): if f.suffix in (".md", ".json", ".html"): @@ -213,74 +497,23 @@ async def _handle_generate_docs(arguments: dict[str, Any]) -> list[TextContent]: "files_generated": sorted(generated_files), "file_count": len(generated_files), } - return [TextContent(type="text", text=json.dumps(result, indent=2))] + return [_text(json.dumps(result, indent=2))] -async def _handle_analyze_repo(arguments: dict[str, Any]) -> list[TextContent]: - """Handle analyze_repo tool call — lightweight dependency analysis only.""" - repo_path = Path(arguments["repo_path"]).expanduser().resolve() - - if not repo_path.exists(): - return [TextContent(type="text", text=f"Repository not found: {repo_path}")] - - manager = _load_config() - config = manager.get_config() - api_key = manager.get_api_key() - - from codewiki.src.config import Config as BackendConfig, set_cli_context - set_cli_context(True) - - # Create a minimal backend config (no LLM calls needed for analysis) - backend_config = BackendConfig.from_cli( - repo_path=str(repo_path), - output_dir=str(repo_path / ".codewiki_temp"), - llm_base_url=config.base_url or "http://localhost", - llm_api_key=api_key or "not-needed", - main_model=config.main_model or "unused", - cluster_model=config.cluster_model or "unused", - fallback_model=config.fallback_model or "unused", - ) - - from codewiki.src.be.dependency_analyzer import DependencyGraphBuilder - graph_builder = DependencyGraphBuilder(backend_config) - components, leaf_nodes = graph_builder.build_dependency_graph() - - # Aggregate statistics - languages = {} - files = set() - for comp in components.values(): - lang = getattr(comp, "language", "unknown") - languages[lang] = languages.get(lang, 0) + 1 - files.add(getattr(comp, "relative_path", "")) - - result = { - "status": "success", - "repo_path": str(repo_path), - "total_components": len(components), - "total_files": len(files), - "leaf_nodes": len(leaf_nodes), - "languages": languages, - "sample_components": sorted(list(components.keys()))[:20], - } - return [TextContent(type="text", text=json.dumps(result, indent=2))] - - -async def _handle_get_module_tree(arguments: dict[str, Any]) -> list[TextContent]: - """Handle get_module_tree tool call — returns existing module tree.""" +async def _legacy_get_module_tree(arguments: dict[str, Any]) -> list[TextContent]: + """Legacy get_module_tree.""" repo_path = Path(arguments["repo_path"]).expanduser().resolve() output_dir = Path(arguments.get("output_dir", "docs")).expanduser().resolve() module_tree_path = output_dir / "module_tree.json" if not module_tree_path.exists(): - return [TextContent( - type="text", - text=f"Module tree not found at {module_tree_path}. Run 'codewiki generate' first." - )] + return [_text(json.dumps({ + "error": f"Module tree not found at {module_tree_path}. Run 'codewiki generate' first." + }))] - module_tree = json.loads(module_tree_path.read_text()) + module_tree = json.loads(module_tree_path.read_text(encoding="utf-8")) def _summarize_tree(tree, depth=0): - """Create a readable summary of the module tree.""" lines = [] for name, info in tree.items(): indent = " " * depth @@ -299,13 +532,62 @@ def _summarize_tree(tree, depth=0): "total_modules": len(module_tree), "tree_summary": summary, } - return [TextContent(type="text", text=json.dumps(result, indent=2))] + return [_text(json.dumps(result, indent=2))] + +# =================================================================== +# Helpers +# =================================================================== + +def _text(content: str) -> TextContent: + return TextContent(type="text", text=content) + + +def _write_generation_metadata(session: SessionState) -> None: + """Write ``metadata.json`` to the session's output directory. + + Records the current git commit and timestamp so that + :func:`_detect_changes` can diff against this baseline on the next + ``analyze_repo`` call, enabling incremental updates. + """ + try: + output_dir = Path(session.output_dir) + repo_path = Path(session.repo_path) + + commit_id: str | None = None + try: + import git + repo = git.Repo(repo_path, search_parent_directories=True) + commit_id = repo.head.commit.hexsha + except Exception: + pass + + from datetime import datetime + metadata = { + "generation_info": { + "commit_id": commit_id, + "timestamp": datetime.now().isoformat(), + }, + } + (output_dir / "metadata.json").write_text( + json.dumps(metadata, indent=2, ensure_ascii=False) + ) + except Exception as e: + logger.warning("Failed to write metadata.json: %s", e) + + +# =================================================================== +# Entry point +# =================================================================== async def main(): """Run the MCP server with stdio transport.""" async with stdio_server() as (read_stream, write_stream): - await server.run(read_stream, write_stream, server.create_initialization_options()) + await server.run( + read_stream, + write_stream, + server.create_initialization_options(), + ) if __name__ == "__main__": diff --git a/codewiki/mcp/session.py b/codewiki/mcp/session.py new file mode 100644 index 00000000..4bac0c0d --- /dev/null +++ b/codewiki/mcp/session.py @@ -0,0 +1,123 @@ +"""Session state management for the CodeWiki MCP Server. + +Each ``analyze_repo`` call creates a new session that caches the analysis +results (components, leaf nodes, etc.) in memory. Subsequent tool calls +reference the session by ``session_id`` to read code, write docs, and +manage the module tree without re-parsing the repository. +""" + +from __future__ import annotations + +import threading +import time +import uuid +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from codewiki.src.be.dependency_analyzer.models.core import Node + +if TYPE_CHECKING: + from codewiki.mcp.workspace import SessionWorkspace + + +# Sessions auto-expire after this many seconds of inactivity. +_SESSION_TTL_SECONDS = 2 * 60 * 60 # 2 hours + +# Maximum concurrent sessions to prevent unbounded memory growth. +_MAX_SESSIONS = 10 + + +@dataclass +class SessionState: + """Mutable state shared across all MCP tool calls within a session.""" + + session_id: str + repo_path: str + output_dir: str + components: Dict[str, Node] + leaf_nodes: List[str] + module_tree: Dict[str, Any] = field(default_factory=dict) + registry: Dict[str, Any] = field(default_factory=dict) + workspace: Optional[SessionWorkspace] = field(default=None) + created_at: float = field(default_factory=time.time) + last_accessed: float = field(default_factory=time.time) + + def touch(self) -> None: + """Update the last-accessed timestamp.""" + self.last_accessed = time.time() + + @property + def is_expired(self) -> bool: + return (time.time() - self.last_accessed) > _SESSION_TTL_SECONDS + + +class SessionStore: + """In-memory store for all active MCP sessions (thread-safe).""" + + def __init__(self) -> None: + self._sessions: Dict[str, SessionState] = {} + self._lock = threading.Lock() + + def create( + self, + repo_path: str, + output_dir: str, + components: Dict[str, Node], + leaf_nodes: List[str], + workspace: Optional[SessionWorkspace] = None, + ) -> SessionState: + """Create a new session and return it.""" + with self._lock: + self._purge_expired_locked() + # Evict oldest if at capacity + if len(self._sessions) >= _MAX_SESSIONS: + oldest_id = min( + self._sessions, + key=lambda sid: self._sessions[sid].last_accessed, + ) + evicted = self._sessions[oldest_id] + if evicted.workspace is not None: + evicted.workspace.cleanup() + del self._sessions[oldest_id] + session_id = uuid.uuid4().hex[:12] + # Ensure no collision + while session_id in self._sessions: + session_id = uuid.uuid4().hex[:12] + state = SessionState( + session_id=session_id, + repo_path=repo_path, + output_dir=output_dir, + components=components, + leaf_nodes=leaf_nodes, + workspace=workspace, + ) + self._sessions[session_id] = state + return state + + def get(self, session_id: str) -> Optional[SessionState]: + """Return the session or ``None`` if not found / expired.""" + with self._lock: + state = self._sessions.get(session_id) + if state is None: + return None + if state.is_expired: + if state.workspace is not None: + state.workspace.cleanup() + del self._sessions[session_id] + return None + state.touch() + return state + + def remove(self, session_id: str) -> bool: + """Remove a session. Returns True if it existed.""" + with self._lock: + return self._sessions.pop(session_id, None) is not None + + def _purge_expired_locked(self) -> None: + """Remove all expired sessions. Caller must hold _lock.""" + expired = [sid for sid, s in self._sessions.items() if s.is_expired] + for sid in expired: + state = self._sessions[sid] + if state.workspace is not None: + state.workspace.cleanup() + del self._sessions[sid] diff --git a/codewiki/mcp/tools/__init__.py b/codewiki/mcp/tools/__init__.py new file mode 100644 index 00000000..62ea687d --- /dev/null +++ b/codewiki/mcp/tools/__init__.py @@ -0,0 +1,6 @@ +"""CodeWiki MCP Tools package. + +Each module in this package implements one or more MCP tools that operate +on a :class:`~codewiki.mcp.session.SessionState`. The tools are registered +by the MCP server in ``codewiki/mcp/server.py``. +""" diff --git a/codewiki/mcp/tools/analysis.py b/codewiki/mcp/tools/analysis.py new file mode 100644 index 00000000..a8e292ca --- /dev/null +++ b/codewiki/mcp/tools/analysis.py @@ -0,0 +1,351 @@ +"""MCP tool: analyze_repo — parse a repository and build the dependency graph. + +This is the entry-point tool for the IDE-driven wiki generation pipeline. +It runs CodeWiki's Tree-sitter-based dependency analyzer (no LLM needed), +caches the results in a new session, and writes the full component index, +leaf nodes, and other analysis data to files on disk. The IDE agent reads +those files directly instead of receiving large payloads over stdio. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +from codewiki.mcp.session import SessionState, SessionStore +from codewiki.mcp.workspace import SessionWorkspace + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Incremental update: detect changes since last generation +# --------------------------------------------------------------------------- + +def _detect_changes( + repo_path: Path, + output_dir: Path, +) -> Optional[Dict[str, Any]]: + """Detect changes since last documentation generation. + + Returns a changes dict with affected modules, or None if no previous + generation exists (first run). + + Detection strategy: + 1. Git-based: compare stored commit_id with current HEAD, plus check + uncommitted changes via ``git status``. + 2. Fallback: compare file mtime with stored ``timestamp`` in metadata. + """ + metadata_path = output_dir / "metadata.json" + module_tree_path = output_dir / "module_tree.json" + + if not metadata_path.exists() or not module_tree_path.exists(): + return None + + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + module_tree = json.loads(module_tree_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError, UnicodeDecodeError): + return None + + # Try git-based detection first + changes = _detect_via_git(repo_path, metadata) + + # Fallback to mtime-based detection + if changes is None: + changes = _detect_via_mtime(repo_path, metadata) + + if changes is None: + return None + + changed_files = changes["changed_files"] + if not changed_files: + return { + "has_previous": True, + "no_changes": True, + "method": changes.get("method", "unknown"), + "message": "No changes detected since last generation. Documentation is up to date.", + } + + affected, cascade = _find_affected_modules(module_tree, changed_files) + + return { + "has_previous": True, + "no_changes": False, + "method": changes.get("method", "unknown"), + "changed_files": changed_files, + "affected_modules": sorted(affected), + "cascade_modules": sorted(cascade), + "hint": ( + f"Only {len(affected)} module(s) need updating: {sorted(affected)}. " + f"Parent modules to refresh: {sorted(cascade)}. " + "Use edit_doc_file for targeted updates, write_doc_file for new modules." + ), + } + + +def _detect_via_git( + repo_path: Path, + metadata: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + """Detect changes via git. Returns None if not in a git repo. + + Checks both committed changes (diff against stored commit_id) and + uncommitted changes (``git status``). + """ + try: + import git + repo = git.Repo(repo_path, search_parent_directories=True) + except Exception: + return None + + prev_commit = metadata.get("generation_info", {}).get("commit_id") + try: + current_commit = repo.head.commit.hexsha + except Exception: + return None + + changed: list[str] = [] + method = "git" + + # 1) Committed changes since last generation + if prev_commit and prev_commit != current_commit: + try: + diff_index = repo.commit(prev_commit).diff(current_commit) + seen: set[str] = set() + for diff in diff_index: + if diff.a_path and diff.a_path not in seen: + changed.append(diff.a_path) + seen.add(diff.a_path) + if diff.b_path and diff.b_path not in seen: + changed.append(diff.b_path) + seen.add(diff.b_path) + except Exception: + pass + + # 2) Uncommitted changes (user may have edited but not committed) + try: + for item in repo.untracked_files: + if item not in changed: + changed.append(item) + for file_path in [d.a_path for d in repo.index.diff(None)]: + if file_path and file_path not in changed: + changed.append(file_path) + except Exception: + pass + + return {"changed_files": changed, "method": method} + + +def _detect_via_mtime( + repo_path: Path, + metadata: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + """Fallback: detect changed files by comparing mtime with generation timestamp.""" + timestamp_str = metadata.get("generation_info", {}).get("timestamp") + if not timestamp_str: + return None + + try: + from datetime import datetime + prev_time = datetime.fromisoformat(timestamp_str).timestamp() + except (ValueError, TypeError): + return None + + # Language extensions recognized by CodeWiki + source_extensions = { + ".py", ".java", ".js", ".jsx", ".ts", ".tsx", + ".c", ".h", ".cpp", ".hpp", ".cc", ".hh", + ".cs", ".kt", ".kts", + } + + changed: list[str] = [] + for dirpath, dirnames, filenames in os.walk(repo_path): + # Skip hidden dirs and common non-source dirs + dirnames[:] = [ + d for d in dirnames + if not d.startswith(".") and d not in ("node_modules", "__pycache__", "venv", ".venv") + ] + for filename in filenames: + filepath = Path(dirpath) / filename + if filepath.suffix.lower() not in source_extensions: + continue + try: + if filepath.stat().st_mtime > prev_time: + rel_path = str(filepath.relative_to(repo_path)) + changed.append(rel_path) + except OSError: + continue + + return {"changed_files": changed, "method": "mtime"} + + +def _find_affected_modules( + module_tree: Dict[str, Any], + changed_files: List[str], +) -> Tuple[set, set]: + """Map changed files to affected modules using module_tree.json. + + Uses substring matching (same as the CLI ``_invalidate_affected_modules``). + Returns (affected_modules, cascade_parent_modules). + """ + affected: set[str] = set() + cascade: set[str] = set() + + def _walk(tree: Dict, parents: list[str] | None = None): + if parents is None: + parents = [] + for mod_name, mod_info in tree.items(): + components = mod_info.get("components", []) + hit = False + for comp in components: + if any(cf in comp or comp in cf for cf in changed_files): + hit = True + break + if hit: + affected.add(mod_name) + cascade.update(parents) + + children = mod_info.get("children", {}) + if isinstance(children, dict) and children: + _walk(children, parents + [mod_name]) + + _walk(module_tree) + + # overview.md depends on all child docs, always refresh if anything changed + if affected: + cascade.add("overview") + + return affected, cascade + + +def handle_analyze_repo( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Run the dependency analysis, write results to workspace files, + and return a compact summary with file paths.""" + repo_path = Path(arguments["repo_path"]).expanduser().resolve() + if not repo_path.exists(): + return json.dumps({"error": f"Repository not found: {repo_path}"}) + + output_dir = Path(arguments.get("output_dir", str(repo_path / "docs"))).expanduser().resolve() + output_dir.mkdir(parents=True, exist_ok=True) + + # Build a minimal Config for the dependency analyzer (no LLM fields used) + from codewiki.src.config import Config + config = Config( + repo_path=str(repo_path), + output_dir=str(output_dir / "temp"), + dependency_graph_dir=str(output_dir / "temp" / "dependency_graphs"), + docs_dir=str(output_dir), + max_depth=2, + llm_base_url="not-needed", + llm_api_key="not-needed", + main_model="unused", + cluster_model="unused", + ) + + # Apply optional include/exclude patterns + include = arguments.get("include_patterns") + exclude = arguments.get("exclude_patterns") + if include or exclude: + agent_instructions: Dict[str, Any] = {} + if include: + agent_instructions["include_patterns"] = [p.strip() for p in include.split(",")] + if exclude: + agent_instructions["exclude_patterns"] = [p.strip() for p in exclude.split(",")] + config.agent_instructions = agent_instructions + + from codewiki.src.be.dependency_analyzer import DependencyGraphBuilder + builder = DependencyGraphBuilder(config) + components, leaf_nodes = builder.build_dependency_graph() + + # Create the session (generates session_id) + session = store.create( + repo_path=str(repo_path), + output_dir=str(output_dir), + components=components, + leaf_nodes=leaf_nodes, + ) + + # Create the workspace with the real session_id + workspace = SessionWorkspace(repo_path, session.session_id) + session.workspace = workspace + + # -- Write full data to workspace files -- + + # 1. Full component index (no pagination) + component_index: list[dict] = [] + for comp_id, node in components.items(): + component_index.append({ + "id": comp_id, + "type": getattr(node, "component_type", "unknown"), + "file": getattr(node, "relative_path", ""), + }) + workspace.write_json("component_index.json", component_index) + + # 2. Full leaf nodes list + workspace.write_json("leaf_nodes.json", leaf_nodes) + + # 3. Language stats + languages: Dict[str, int] = {} + for node in components.values(): + lang = getattr(node, "language", "unknown") + languages[lang] = languages.get(lang, 0) + 1 + workspace.write_json("languages.json", languages) + + # 4. Incremental update: detect changes since last generation + changes = _detect_changes(repo_path, output_dir) + if changes is not None: + workspace.write_json("changes.json", changes) + + # 5. Summary with preview for quick reference + summary = { + "session_id": session.session_id, + "repo_name": repo_path.name, + "repo_path": str(repo_path), + "output_dir": str(output_dir), + "total_components": len(components), + "total_leaf_nodes": len(leaf_nodes), + "languages": languages, + "leaf_nodes_preview": leaf_nodes[:20], + } + workspace.write_json("summary.json", summary) + + # -- Return compact MCP response -- + result = { + "session_id": session.session_id, + "workspace_dir": str(workspace.root), + "repo_name": repo_path.name, + "output_dir": str(output_dir), + "stats": { + "total_components": len(components), + "total_leaf_nodes": len(leaf_nodes), + "languages": languages, + }, + "files": { + "component_index": str(workspace.root / "component_index.json"), + "leaf_nodes": str(workspace.root / "leaf_nodes.json"), + "languages": str(workspace.root / "languages.json"), + "summary": str(workspace.root / "summary.json"), + }, + "changes": changes, + "hint": ( + "Read the files above for full data. " + "Use read_code_components(session_id, component_ids) to read source code. " + "Use save_module_tree(session_id, module_tree) after clustering. " + "Call get_prompt('cluster') for clustering rules." + ), + } + if changes and not changes.get("no_changes"): + result["hint"] = ( + "Incremental update detected. Only update affected modules listed in " + "'changes.affected_modules'. Use edit_doc_file for targeted updates. " + "Refresh cascade parent modules in 'changes.cascade_modules'." + ) + return json.dumps(result, indent=2, ensure_ascii=False) diff --git a/codewiki/mcp/tools/code_reader.py b/codewiki/mcp/tools/code_reader.py new file mode 100644 index 00000000..ee6a0392 --- /dev/null +++ b/codewiki/mcp/tools/code_reader.py @@ -0,0 +1,61 @@ +"""MCP tool: read_code_components — write component source code to disk. + +Instead of transmitting source code through the MCP stdio channel (which +required aggressive truncation), this tool writes complete, untruncated +source files to the session workspace. The IDE agent then reads them +directly from disk. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Dict, List + +from codewiki.mcp.session import SessionStore + +logger = logging.getLogger(__name__) + + +def handle_read_code_components( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Write the source code for given component IDs to workspace files. + + Returns a compact JSON with file paths — no source code inline. + """ + session_id = arguments["session_id"] + session = store.get(session_id) + if session is None: + return json.dumps({"error": f"Session {session_id} not found or expired."}) + + if session.workspace is None: + return json.dumps({"error": "Session workspace not initialized."}) + + component_ids: List[str] = arguments["component_ids"] + components = session.components + workspace = session.workspace + + written_files: Dict[str, str] = {} # filename -> component_id + not_found: List[str] = [] + + for cid in component_ids: + node = components.get(cid) + if node is None: + not_found.append(cid) + continue + lang = getattr(node, "language", "") + source = getattr(node, "source_code", "").strip() + file_path = workspace.write_component_source(cid, source, lang) + written_files[file_path.name] = cid + + result = { + "written": len(written_files), + "not_found_count": len(not_found), + "not_found": not_found, + "source_dir": str(workspace.root / "sources"), + "files": written_files, + } + return json.dumps(result, indent=2, ensure_ascii=False) diff --git a/codewiki/mcp/tools/doc_writer.py b/codewiki/mcp/tools/doc_writer.py new file mode 100644 index 00000000..14bad977 --- /dev/null +++ b/codewiki/mcp/tools/doc_writer.py @@ -0,0 +1,208 @@ +"""MCP tools: write_doc_file + edit_doc_file. + +These tools create and edit markdown documentation files in the output +directory, with automatic Mermaid diagram validation after every write. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path +from typing import Any, Dict + +from codewiki.mcp.session import SessionState, SessionStore + +logger = logging.getLogger(__name__) + +# Max edit history entries per file (prevent unbounded memory growth) +_MAX_HISTORY_PER_FILE = 20 + + +def _is_within(path: Path, base: Path) -> bool: + """Return True if *path* resolves to somewhere inside *base*.""" + try: + path.resolve().relative_to(base.resolve()) + return True + except ValueError: + return False + + +def _safe_doc_path(session: SessionState, filename: str) -> Path | None: + """Resolve *filename* within session.output_dir, guarding against traversal.""" + if not filename.endswith(".md"): + filename += ".md" + output_base = Path(session.output_dir).resolve() + doc_path = (output_base / filename).resolve() + if not _is_within(doc_path, output_base): + return None + return doc_path + + +def _ensure_parent_dirs(path: Path) -> None: + """Create parent directories if they don't exist.""" + path.parent.mkdir(parents=True, exist_ok=True) + + +async def _validate_mermaid(file_path: str, relative_path: str) -> str: + """Run Mermaid validation and return the result string.""" + try: + from codewiki.src.be.utils import validate_mermaid_diagrams + return await validate_mermaid_diagrams(file_path, relative_path) + except Exception as e: + return f"Mermaid validation skipped: {e}" + + +def _save_history(session: SessionState, doc_path: Path, content: str) -> None: + """Append *content* to edit history for *doc_path*, capped at _MAX_HISTORY_PER_FILE.""" + history = session.registry.get("file_history") + if history is None: + history = {} + elif isinstance(history, str): + history = json.loads(history) + key = str(doc_path) + entry = history.setdefault(key, []) + entry.append(content) + # Trim to last N entries + if len(entry) > _MAX_HISTORY_PER_FILE: + del entry[: len(entry) - _MAX_HISTORY_PER_FILE] + session.registry["file_history"] = history # keep as native dict + + +async def handle_write_doc_file( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Create a new documentation file in the output directory.""" + session_id = arguments["session_id"] + session = store.get(session_id) + if session is None: + return json.dumps({"error": f"Session {session_id} not found or expired."}) + + filename = arguments["filename"] + doc_path = _safe_doc_path(session, filename) + if doc_path is None: + return json.dumps({"error": "Filename escapes output directory."}) + + content = arguments["content"] + + _ensure_parent_dirs(doc_path) + + if doc_path.exists(): + return json.dumps({ + "error": f"File already exists: {filename}. Use edit_doc_file to modify it." + }) + + doc_path.write_text(content, encoding="utf-8") + + # Mermaid validation + mermaid_result = await _validate_mermaid(str(doc_path), filename) + + result = { + "status": "created", + "path": str(doc_path), + "filename": filename, + "lines": content.count("\n") + 1, + "mermaid_validation": mermaid_result, + } + return json.dumps(result, indent=2, ensure_ascii=False) + + +async def handle_edit_doc_file( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Edit an existing documentation file (str_replace, insert, or undo).""" + session_id = arguments["session_id"] + session = store.get(session_id) + if session is None: + return json.dumps({"error": f"Session {session_id} not found or expired."}) + + filename = arguments["filename"] + doc_path = _safe_doc_path(session, filename) + if doc_path is None: + return json.dumps({"error": "Filename escapes output directory."}) + + command = arguments["command"] + + if command == "undo": + # Undo via registry history + history = session.registry.get("file_history", {}) + if isinstance(history, str): + history = json.loads(history) + path_history = history.get(str(doc_path), []) + if not path_history: + return json.dumps({"error": f"No edit history found for {filename}."}) + old_content = path_history.pop() + history[str(doc_path)] = path_history + session.registry["file_history"] = history + doc_path.write_text(old_content, encoding="utf-8") + + # Validate Mermaid after undo + mermaid_result = await _validate_mermaid(str(doc_path), filename) + return json.dumps({ + "status": "undone", + "filename": filename, + "mermaid_validation": mermaid_result, + }, ensure_ascii=False) + + if not doc_path.exists(): + return json.dumps({"error": f"File not found: {filename}. Use write_doc_file to create it."}) + + # Save current content to history before editing + current_content = doc_path.read_text(encoding="utf-8") + _save_history(session, doc_path, current_content) + + if command == "str_replace": + old_str = arguments.get("old_str") + new_str = arguments.get("new_str", "") + if old_str is None: + return json.dumps({"error": "old_str is required for str_replace."}) + + occurrences = current_content.count(old_str) + if occurrences == 0: + return json.dumps({"error": f"old_str not found in {filename}."}) + if occurrences > 1: + return json.dumps({"error": f"old_str appears {occurrences} times in {filename}. Make it unique."}) + + new_content = current_content.replace(old_str, new_str, 1) + doc_path.write_text(new_content, encoding="utf-8") + + # Snippet around the edit + replacement_line = current_content.split(old_str)[0].count("\n") + lines = new_content.split("\n") + start = max(0, replacement_line - 4) + end = min(len(lines), replacement_line + new_str.count("\n") + 5) + snippet = "\n".join(f"{i + start + 1:6}\t{lines[i]}" for i in range(start, end)) + + elif command == "insert": + insert_line = arguments.get("insert_line", 0) + new_str = arguments.get("new_str", "") + if not new_str: + return json.dumps({"error": "new_str is required for insert."}) + + lines = current_content.split("\n") + insert_line = max(0, min(insert_line, len(lines))) + new_str_lines = new_str.split("\n") + lines = lines[:insert_line] + new_str_lines + lines[insert_line:] + new_content = "\n".join(lines) + doc_path.write_text(new_content, encoding="utf-8") + + start = max(0, insert_line - 4) + end = min(len(lines), insert_line + len(new_str_lines) + 4) + snippet = "\n".join(f"{i + start + 1:6}\t{lines[i]}" for i in range(start, end)) + + else: + return json.dumps({"error": f"Unknown command: {command}. Use str_replace, insert, or undo."}) + + # Mermaid validation + mermaid_result = await _validate_mermaid(str(doc_path), filename) + + result = { + "status": "edited", + "command": command, + "filename": filename, + "snippet": snippet, + "mermaid_validation": mermaid_result, + } + return json.dumps(result, indent=2, ensure_ascii=False) diff --git a/codewiki/mcp/tools/module_tree.py b/codewiki/mcp/tools/module_tree.py new file mode 100644 index 00000000..f0d6ca3e --- /dev/null +++ b/codewiki/mcp/tools/module_tree.py @@ -0,0 +1,144 @@ +"""MCP tools: save_module_tree + get_processing_order. + +The IDE agent decides how to group components into modules (clustering) +using its own LLM. These tools persist that decision and compute the +leaf-first processing order for documentation generation. +""" + +from __future__ import annotations + +import json +import logging +import os +from pathlib import Path +from typing import Any, Dict, List, Tuple + +from codewiki.mcp.session import SessionState, SessionStore +from codewiki.src.config import FIRST_MODULE_TREE_FILENAME, MODULE_TREE_FILENAME + +logger = logging.getLogger(__name__) + + +def _get_processing_order(module_tree: Dict[str, Any], parent_path: List[str] = []) -> List[Dict[str, Any]]: + """Compute leaf-first processing order from a module tree. + + Returns a list of dicts with module path, name, leaf status, and + component/children info. + """ + order: List[Dict[str, Any]] = [] + + def _collect(tree: Dict[str, Any], path: List[str]) -> None: + for module_name, module_info in tree.items(): + current_path = path + [module_name] + children = module_info.get("children", {}) + has_children = isinstance(children, dict) and len(children) > 0 + + if has_children: + _collect(children, current_path) + order.append({ + "module": module_name, + "path": current_path, + "is_leaf": False, + "children": list(children.keys()), + "components": module_info.get("components", []), + }) + else: + order.append({ + "module": module_name, + "path": current_path, + "is_leaf": True, + "components": module_info.get("components", []), + }) + + _collect(module_tree, parent_path) + return order + + +def handle_save_module_tree( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Persist the IDE agent's clustering result as the module tree.""" + session_id = arguments["session_id"] + session = store.get(session_id) + if session is None: + return json.dumps({"error": f"Session {session_id} not found or expired."}) + + module_tree = arguments["module_tree"] + output_dir = session.output_dir + + # Save both immutable snapshot and mutable working copy + first_path = os.path.join(output_dir, FIRST_MODULE_TREE_FILENAME) + working_path = os.path.join(output_dir, MODULE_TREE_FILENAME) + + os.makedirs(output_dir, exist_ok=True) + + with open(first_path, "w", encoding="utf-8") as f: + json.dump(module_tree, f, indent=2, ensure_ascii=False) + with open(working_path, "w", encoding="utf-8") as f: + json.dump(module_tree, f, indent=2, ensure_ascii=False) + + # Cache in session + session.module_tree = module_tree + + # Compute processing order and write to workspace file + order = _get_processing_order(module_tree) + order_file = None + if session.workspace is not None: + order_path = session.workspace.write_json("processing_order.json", order) + order_file = str(order_path) + + result = { + "status": "saved", + "module_count": len(module_tree), + "tree_path": working_path, + "first_tree_path": first_path, + "processing_order_file": order_file, + "hint": ( + "Read the processing_order.json file for the leaf-first generation order. " + "Process leaf modules first (is_leaf=true), then parent modules. " + "For each leaf module: get_prompt('system_leaf') + read_code_components + write_doc_file. " + "For each parent module: get_prompt('overview_module') + write_doc_file." + ), + } + return json.dumps(result, indent=2, ensure_ascii=False) + + +def handle_get_processing_order( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Write the leaf-first processing order to a workspace file and return its path.""" + session_id = arguments["session_id"] + session = store.get(session_id) + if session is None: + return json.dumps({"error": f"Session {session_id} not found or expired."}) + + # Try session cache first, then disk + module_tree = session.module_tree + if not module_tree: + tree_path = os.path.join(session.output_dir, MODULE_TREE_FILENAME) + if os.path.exists(tree_path): + with open(tree_path, encoding="utf-8") as f: + module_tree = json.load(f) + session.module_tree = module_tree + else: + return json.dumps({ + "error": "Module tree not found. Call save_module_tree first." + }) + + order = _get_processing_order(module_tree) + + # Write to workspace file + order_file = None + if session.workspace is not None: + order_path = session.workspace.write_json("processing_order.json", order) + order_file = str(order_path) + + result = { + "session_id": session_id, + "module_count": len(module_tree), + "processing_order_file": order_file, + "hint": "Read the processing_order.json file for the full leaf-first order.", + } + return json.dumps(result, indent=2, ensure_ascii=False) diff --git a/codewiki/mcp/tools/prompt_server.py b/codewiki/mcp/tools/prompt_server.py new file mode 100644 index 00000000..04fe2347 --- /dev/null +++ b/codewiki/mcp/tools/prompt_server.py @@ -0,0 +1,176 @@ +"""MCP tool: get_prompt — serve CodeWiki's prompt templates to the IDE agent. + +CodeWiki ships with carefully designed prompt templates for each stage of +the wiki generation pipeline. This tool lets the IDE agent retrieve them +(with optional variable substitution) so it can follow the same proven +methodology without needing its own copy of the prompts. +""" + +from __future__ import annotations + +import json +import logging +from typing import Any, Dict, Optional + +from codewiki.mcp.session import SessionStore +from codewiki.src.be.prompt_template import ( + CLUSTER_REPO_PROMPT, + CLUSTER_MODULE_PROMPT, + SYSTEM_PROMPT, + LEAF_SYSTEM_PROMPT, + USER_PROMPT, + REPO_OVERVIEW_PROMPT, + MODULE_OVERVIEW_PROMPT, + format_system_prompt, + format_leaf_system_prompt, + format_cluster_prompt, + format_user_prompt, +) + +logger = logging.getLogger(__name__) + + +# Prompt catalog: maps prompt_type to (raw_template, usage_hint, variables_doc) +_PROMPT_CATALOG: Dict[str, Dict[str, str]] = { + "cluster": { + "description": "Prompt for grouping components into modules. The LLM receives a component list and returns a JSON module tree.", + "usage_hint": ( + "Use this prompt to cluster components into logical modules. " + "The response should contain JSON. " + "Pass the component list from analyze_repo's component_index." + ), + }, + "system_complex": { + "description": "System prompt for documenting a complex (multi-file, parent) module. Includes sub-module delegation instructions.", + "usage_hint": ( + "Use as the system prompt when generating docs for a parent module. " + "The agent should create {module_name}.md with architecture overview " + "and cross-references to sub-module docs." + ), + }, + "system_leaf": { + "description": "System prompt for documenting a leaf (single-file or simple) module.", + "usage_hint": ( + "Use as the system prompt when generating docs for a leaf module. " + "The agent should create {module_name}.md with detailed documentation " + "including Mermaid diagrams." + ), + }, + "user": { + "description": "User prompt template that provides the module tree and core component source code.", + "usage_hint": ( + "Use as the user/assistant prompt alongside system_leaf or system_complex. " + "It provides the module tree context and the actual source code of core components." + ), + }, + "overview_module": { + "description": "Prompt for generating a parent module overview from its children's documentation.", + "usage_hint": ( + "Use this after all child modules are documented. " + "Provide the module tree with children's docs embedded. " + "The response should be wrapped in tags." + ), + }, + "overview_repo": { + "description": "Prompt for generating the final repository overview.", + "usage_hint": ( + "Use this as the LAST step after all modules are documented. " + "Provide the full module tree with child docs. " + "Save the result as overview.md." + ), + }, +} + + +def handle_get_prompt( + arguments: Dict[str, Any], + store: SessionStore, +) -> str: + """Return a prompt template, optionally with variables filled in.""" + prompt_type = arguments["prompt_type"] + variables = arguments.get("variables", {}) + + if prompt_type not in _PROMPT_CATALOG: + available = list(_PROMPT_CATALOG.keys()) + return json.dumps({ + "error": f"Unknown prompt_type: {prompt_type}", + "available_types": available, + }) + + catalog_entry = _PROMPT_CATALOG[prompt_type] + + # Resolve the prompt content + content = _resolve_prompt(prompt_type, variables) + + result = { + "prompt_type": prompt_type, + "description": catalog_entry["description"], + "usage_hint": catalog_entry["usage_hint"], + "content": content, + } + return json.dumps(result, indent=2, ensure_ascii=False) + + +def _resolve_prompt(prompt_type: str, variables: Dict[str, Any]) -> str: + """Resolve a prompt template with optional variable substitution.""" + + if prompt_type == "cluster": + potential_core_components = variables.get("potential_core_components", "") + module_tree = variables.get("module_tree", {}) + module_name = variables.get("module_name", None) + return format_cluster_prompt( + potential_core_components=potential_core_components, + module_tree=module_tree, + module_name=module_name, + ) + + elif prompt_type == "system_complex": + module_name = variables.get("module_name", "MODULE_NAME") + custom_instructions = variables.get("custom_instructions", None) + return format_system_prompt(module_name, custom_instructions) + + elif prompt_type == "system_leaf": + module_name = variables.get("module_name", "MODULE_NAME") + custom_instructions = variables.get("custom_instructions", None) + return format_leaf_system_prompt(module_name, custom_instructions) + + elif prompt_type == "user": + # If full variables are provided, use the full formatter + session_id = variables.get("session_id") + module_name = variables.get("module_name", "MODULE_NAME") + core_component_ids = variables.get("core_component_ids", []) + module_tree = variables.get("module_tree", {}) + + if session_id and core_component_ids: + # Try to resolve from session + from codewiki.mcp.session import SessionStore + # We can't easily access the store here, so fall back to template + pass + + # Return the template with placeholders filled as possible + return USER_PROMPT.format( + module_name=module_name, + module_tree=json.dumps(module_tree, indent=2) if module_tree else "", + formatted_core_component_codes=variables.get( + "formatted_core_component_codes", + "" + ), + ) + + elif prompt_type == "overview_module": + module_name = variables.get("module_name", "MODULE_NAME") + repo_structure = variables.get("repo_structure", "") + return MODULE_OVERVIEW_PROMPT.format( + module_name=module_name, + repo_structure=repo_structure if isinstance(repo_structure, str) else json.dumps(repo_structure, indent=4), + ) + + elif prompt_type == "overview_repo": + repo_name = variables.get("repo_name", "REPO_NAME") + repo_structure = variables.get("repo_structure", "") + return REPO_OVERVIEW_PROMPT.format( + repo_name=repo_name, + repo_structure=repo_structure if isinstance(repo_structure, str) else json.dumps(repo_structure, indent=4), + ) + + return f"Unknown prompt type: {prompt_type}" diff --git a/codewiki/mcp/workspace.py b/codewiki/mcp/workspace.py new file mode 100644 index 00000000..ef0e20da --- /dev/null +++ b/codewiki/mcp/workspace.py @@ -0,0 +1,99 @@ +"""Session file workspace -- write large analysis artifacts to disk. + +Instead of transmitting bulky data through the MCP stdio channel, the +server writes analysis results (component index, leaf nodes, source code, +etc.) to a per-session directory on disk. The IDE agent then reads these +files directly using its own file-access capabilities. + +Directory layout (relative to ``repo_path``):: + + .codewiki/sessions/{session_id}/ + component_index.json + leaf_nodes.json + languages.json + changes.json + summary.json + processing_order.json + sources/ + {sanitized_component_id}.src +""" + +from __future__ import annotations + +import json +import logging +import re +import shutil +from pathlib import Path +from typing import Any + +logger = logging.getLogger(__name__) + +# Base directory under repo_path +_WORKSPACE_REL = Path(".codewiki") / "sessions" + + +def _safe_filename(component_id: str) -> str: + """Sanitize a component ID for use as a filename. + + Component IDs look like ``src/main.py::MyClass``. We replace any + character that is not a word char, hyphen, or dot with ``__``. + """ + return re.sub(r"[^\w\-.]", "__", component_id) + ".src" + + +class SessionWorkspace: + """Manages the on-disk workspace for a single MCP session.""" + + def __init__(self, repo_path: Path, session_id: str) -> None: + self.root = repo_path / _WORKSPACE_REL / session_id + self.root.mkdir(parents=True, exist_ok=True) + (self.root / "sources").mkdir(exist_ok=True) + logger.debug("Workspace created at %s", self.root) + + # -- writers ---------------------------------------------------------- + + def write_json(self, name: str, data: Any) -> Path: + """Write *data* as pretty-printed JSON and return the file path.""" + p = self.root / name + p.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8") + return p + + def write_component_source( + self, + component_id: str, + source: str, + language: str = "", + ) -> Path: + """Write a single component's source code to the ``sources/`` dir.""" + p = self.root / "sources" / _safe_filename(component_id) + header = f"// Component: {component_id}\n// Language: {language}\n" + p.write_text(header + source, encoding="utf-8") + return p + + # -- readers ---------------------------------------------------------- + + def read_json(self, name: str) -> Any: + """Read a JSON file from the workspace. Returns ``None`` if missing.""" + p = self.root / name + if not p.exists(): + return None + return json.loads(p.read_text(encoding="utf-8")) + + # -- cleanup ---------------------------------------------------------- + + def cleanup(self) -> None: + """Remove the session directory and try to prune empty parents.""" + if self.root.exists(): + shutil.rmtree(self.root, ignore_errors=True) + # Walk up and remove empty parent directories + try: + sessions_dir = self.root.parent # .codewiki/sessions + if sessions_dir.exists() and not any(sessions_dir.iterdir()): + sessions_dir.rmdir() + base_dir = sessions_dir.parent # .codewiki + if base_dir.exists() and not any(base_dir.iterdir()): + base_dir.rmdir() + except OSError: + pass + logger.debug("Workspace cleaned up: %s", self.root) diff --git a/codewiki/src/be/utils.py b/codewiki/src/be/utils.py index 02fac183..f3e6e038 100644 --- a/codewiki/src/be/utils.py +++ b/codewiki/src/be/utils.py @@ -1,4 +1,5 @@ import asyncio +import os import re import sys import threading @@ -150,6 +151,12 @@ def extract_mermaid_blocks(content: str) -> List[Tuple[int, str]]: # Skip it proactively so SpiderMonkey is never loaded into the process. _PYTHONMONKEY_BROKEN = sys.version_info >= (3, 12) +# mermaid-py spawns a Node.js subprocess that can hang indefinitely (e.g. when +# Node.js is missing or the mermaid CLI is misconfigured). Default to +# disabled; set MERMAID_VALIDATE=1 to enable. +_MERMAID_PY_BROKEN = os.environ.get("MERMAID_VALIDATE", "0") != "1" +_MERMAID_PY_PROBED = True # Skip probing — rely on env var + async def _try_pythonmonkey_parse(diagram_content: str) -> str | None: """Attempt to parse via PythonMonkey/mermaid-parser-py. @@ -232,8 +239,15 @@ async def validate_single_diagram(diagram_content: str, diagram_num: int, line_s """ core_error = await _try_pythonmonkey_parse(diagram_content) if core_error is None: + if _MERMAID_PY_BROKEN: + return f" Diagram {diagram_num}: validation skipped (set MERMAID_VALIDATE=1 to enable)" try: - core_error = _parse_via_mermaid_py(diagram_content) + core_error = await asyncio.wait_for( + asyncio.to_thread(_parse_via_mermaid_py, diagram_content), + timeout=15.0, + ) + except asyncio.TimeoutError: + return f" Diagram {diagram_num}: validation timed out (15s) — diagram may be invalid" except Exception as e: return f" Diagram {diagram_num}: Exception during validation - {str(e)}" diff --git a/codewiki/src/utils.py b/codewiki/src/utils.py index a1801726..7f03b7a6 100644 --- a/codewiki/src/utils.py +++ b/codewiki/src/utils.py @@ -18,28 +18,28 @@ def ensure_directory(path: str) -> None: @staticmethod def save_json(data: Any, filepath: str) -> None: """Save data as JSON to file.""" - with open(filepath, 'w') as f: - json.dump(data, f, indent=4) - + with open(filepath, 'w', encoding='utf-8') as f: + json.dump(data, f, indent=4, ensure_ascii=False) + @staticmethod def load_json(filepath: str) -> Optional[Dict[str, Any]]: """Load JSON from file, return None if file doesn't exist.""" if not os.path.exists(filepath): return None - - with open(filepath, 'r') as f: + + with open(filepath, 'r', encoding='utf-8') as f: return json.load(f) - + @staticmethod def save_text(content: str, filepath: str) -> None: """Save text content to file.""" - with open(filepath, 'w') as f: + with open(filepath, 'w', encoding='utf-8') as f: f.write(content) - + @staticmethod def load_text(filepath: str) -> str: """Load text content from file.""" - with open(filepath, 'r') as f: + with open(filepath, 'r', encoding='utf-8') as f: return f.read() file_manager = FileManager() diff --git a/img/logo-banner.png b/img/logo-banner.png new file mode 100644 index 00000000..700f3b8b Binary files /dev/null and b/img/logo-banner.png differ diff --git a/pyproject.toml b/pyproject.toml index b618d572..74f26a09 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,7 +57,8 @@ dependencies = [ "python-multipart>=0.0.20", "colorama>=0.4.6", "logfire>=4.1.0", - "coding-agent-wrapper>=0.1.2" + "coding-agent-wrapper>=0.1.2", + "mcp>=1.0.0" ] [external] @@ -101,7 +102,9 @@ packages = [ "codewiki.src.be.dependency_analyzer.analyzers", "codewiki.src.be.dependency_analyzer.models", "codewiki.src.be.dependency_analyzer.utils", - "codewiki.src.fe" + "codewiki.src.fe", + "codewiki.mcp", + "codewiki.mcp.tools" ] [tool.setuptools.package-data] diff --git a/skills/codewiki-wiki-generator/SKILL.md b/skills/codewiki-wiki-generator/SKILL.md new file mode 100644 index 00000000..425bab11 --- /dev/null +++ b/skills/codewiki-wiki-generator/SKILL.md @@ -0,0 +1,181 @@ +--- +name: codewiki-wiki-generator +description: "Generate Wiki documentation for code repositories using CodeWiki-CN MCP tools. Use this skill when the user asks to generate a Wiki, code documentation, repository documentation, or analyze codebase structure. Requires CodeWiki-CN MCP server to be configured." +version: 2.0.0 +--- + +# CodeWiki Documentation Generator + +You are a code documentation generation expert. Use CodeWiki-CN's MCP tools to generate comprehensive Wiki documentation for code repositories. All 8 tools require **no LLM configuration** — you provide all the intelligence and reasoning, CodeWiki provides the toolchain. + +## Core Mechanism: File Side-Channel + +CodeWiki MCP uses a **file side-channel** architecture: large data payloads (component index, source code, processing order) are written to disk files, and MCP returns only file paths and concise summaries. You need to **use your own file reading capabilities** to read workspace files for complete data. + +The workspace directory is located at `{repo_path}/.codewiki/sessions/{session_id}/`, containing: + +- `component_index.json` — Complete component index (each entry includes id/type/file) +- `leaf_nodes.json` — Complete list of leaf node IDs +- `languages.json` — Language statistics +- `summary.json` — Analysis summary +- `changes.json` — Incremental change information (optional) +- `processing_order.json` — Documentation generation order +- `sources/` — Component source files (one `.src` file per component) + +## Prerequisites + +Before starting, confirm the CodeWiki MCP server is available. The MCP tool list should include the following 8 tools: `analyze_repo`, `read_code_components`, `write_doc_file`, `edit_doc_file`, `save_module_tree`, `get_processing_order`, `get_prompt`, `close_session`. + +If tools are unavailable, prompt the user to install and configure CodeWiki-CN: + +```bash +git clone https://github.com/FSoft-AI4Code/CodeWiki.git +cd CodeWiki && pip install -e . +``` + +Then add to MCP configuration: + +```json +{"mcpServers":{"codewiki":{"command":"python","args":["-m","codewiki.mcp.server"],"cwd":"/path/to/CodeWiki"}}} +``` + +## Five-Phase Workflow + +Execute strictly in the following order. All tool calls after Phase 1 require the `session_id` returned by `analyze_repo`. + +### Phase 1: Analyze Repository + +Call `analyze_repo`: + +```json +{ "repo_path": "", "output_dir": "/repowiki" } +``` + +Returns: `session_id`, `workspace_dir` (workspace root path), `stats` (component/leaf counts, language statistics), `files` (paths to each data file), `changes` (incremental change information). + +**You must read workspace files next for complete data:** + +1. Read `{workspace_dir}/component_index.json` — Complete component list +2. Read `{workspace_dir}/leaf_nodes.json` — Leaf node ID list +3. Review `stats` to understand repository scale and plan clustering strategy + +**Remember the `session_id`** — every subsequent step requires it. + +### Phase 2: Module Clustering + +This is the phase that requires the most comprehension. You need to group components into logical modules. + +1. **Get clustering rules**: Call `get_prompt` with `{"prompt_type": "cluster"}` +2. **Read source code**: Call `read_code_components` with component ID lists; source code is written to the workspace `sources/` directory, then read these `.src` files directly to understand each component's functionality and relationships. You can pass any number of component IDs per batch (no limit, no truncation) +3. **Read additional repository files if needed**: Use your file reading tools directly to read source files within the repository +4. **Group by the following principles**: + - Functional cohesion: closely related components go into the same module + - File proximity: components in the same file/directory tend to belong to the same module + - Scale control: typically 3-8 top-level modules, each with 5-30 components + - Component IDs must be preserved exactly (including `::` separators) +5. **Save module tree**: Call `save_module_tree`: + +```json +{ + "session_id": "", + "module_tree": { + "ModuleName": { + "components": ["file.py::ClassA", "file.py::func_b"], + "children": {} + } + } +} +``` + +The return result includes the `processing_order_file` path — read this file to get the leaf-first documentation generation order. + +### Phase 3: Per-Module Documentation Generation + +Read `processing_order.json` to get the processing order. **Process leaf modules first**, then parent modules. + +**For each leaf module** (is_leaf=true): + +1. Get system prompt: `get_prompt` → `{"prompt_type": "system_leaf", "variables": {"module_name": ""}}` +2. Read source code: `read_code_components` → all component IDs in this module, then read files under `sources/` +3. For additional context, use your file reading tools directly to read relevant source files in the repository +4. Write documentation including: module introduction and core functionality, architecture diagram (at least 1 Mermaid diagram), component responsibility descriptions, cross-references `[Module Name](module_name.md)` +5. Save: `write_doc_file` → `{"session_id": "...", "filename": ".md", "content": "..."}` + +If Mermaid validation fails, fix the syntax and retry with `edit_doc_file` (`command: "str_replace"`). + +**For each parent module** (is_leaf=false): + +1. Read all child modules' generated `.md` files directly using your file reading tools +2. Get overview prompt: `get_prompt` → `{"prompt_type": "overview_module", "variables": {"module_name": ""}}` +3. Synthesize child module documentation into a parent module overview +4. Save with `write_doc_file` + +### Phase 4: Generate Repository Overview + +1. Get prompt: `get_prompt` → `{"prompt_type": "overview_repo", "variables": {"repo_name": ""}}` +2. Read all generated module documentation using your file reading tools +3. Write a repository-level overview including: project introduction, end-to-end architecture diagram (Mermaid), reference links to each module's documentation +4. Save: `write_doc_file` → `filename: "overview.md"` + +### Phase 5: Cleanup + +Call `close_session` → `{"session_id": ""}` to release memory and clean up workspace files. + +## Incremental Update Mode + +When documentation has already been generated for a repository (`metadata.json` and `module_tree.json` exist under `output_dir`), the `analyze_repo` return result includes a `changes` field, with complete data written to the `changes.json` file (the changed_files list is no longer truncated). + +**Change detection strategy**: Prefers `git diff` (comparing commit SHA + checking for uncommitted workspace changes); falls back to comparing file modification times for non-git repositories. + +**Incremental update workflow**: + +1. Call `analyze_repo` and check the returned `changes` field or read the `changes.json` file +2. If `no_changes: true`, inform the user that documentation is up-to-date, no action needed +3. If `no_changes: false`, **only update modules listed in `affected_modules`**: + - Use `read_code_components` to read new source code for changed components (written to workspace files, then read) + - Use `edit_doc_file` (`str_replace`) to partially modify the corresponding documentation instead of rewriting the entire document +4. For parent modules in `cascade_modules`, read updated child documents and refresh the overview accordingly +5. Finally update `overview.md` + +The granularity of incremental updates is **module-level** — if any component in a module changes, that module's documentation needs updating. + +## Tool Quick Reference + +| Tool | Purpose | Data Flow | +|------|---------|-----------| +| `analyze_repo` | Analyze repository, build dependency graph | Writes files to workspace, returns paths + stats | +| `read_code_components` | Get component source code | Each component written to `sources/*.src`, returns paths | +| `write_doc_file` | Create .md documents (auto Mermaid validation) | Writes file directly | +| `edit_doc_file` | Edit documents: `str_replace` / `insert` / `undo` | Modifies file directly | +| `save_module_tree` | Save module clustering results | Writes module_tree.json + processing_order.json | +| `get_processing_order` | Get leaf-first processing order | Writes processing_order.json, returns path | +| `get_prompt` | Get prompt templates | Returns inline (small data) | +| `close_session` | Close session and release resources | Cleans up workspace files | + +## Documentation Quality Standards + +- **Language**: Write in English by default (unless the user specifies another language) +- **Mermaid diagrams**: At least 1 architecture diagram per module, prefer `graph TD` or `graph LR` +- **Cross-references**: Use `[Module Name](module_name.md)` format when referencing other modules +- **Code examples**: Show signatures and brief usage for key functions/classes +- **Length**: Leaf module docs 200-500 lines, parent module overviews 100-300 lines, repository overview 80-200 lines + +## Mermaid Syntax Guidelines + +```mermaid +graph TD + A[ComponentA] --> B[ComponentB] + A --> C[ComponentC] +``` + +- Node IDs use only letters and digits (avoid Chinese characters, spaces, colons) +- Node labels wrapped in square brackets: `A[display text]` +- Subgraph syntax: `subgraph title ... end` +- Do not use interactive syntax like `click`, `linkStyle`, etc. + +## Error Handling + +- **Mermaid validation failure**: The tool returns validation error details; fix the syntax and retry with `edit_doc_file` + `str_replace` +- **Session expired** (2-hour timeout): Re-call `analyze_repo` to create a new session +- **Large repositories**: `analyze_repo` may take ~30 seconds; use `include_patterns`/`exclude_patterns` to narrow the analysis scope. There are no longer any component count or source code length truncation limits +- **Component ID format**: Always use the original IDs from `component_index.json` (e.g., `src/main.py::MyClass`), preserving the `::` separator diff --git a/tests/smoke_test_mcp.py b/tests/smoke_test_mcp.py new file mode 100644 index 00000000..7832ca84 --- /dev/null +++ b/tests/smoke_test_mcp.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +"""Smoke test for CodeWiki MCP tools — verifies core functionality after +the file-side-channel optimization. + +Run: python3 tests/smoke_test_mcp.py +""" + +import asyncio +import json +import os +import sys +import tempfile +from pathlib import Path + +# Ensure codewiki is importable +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from codewiki.mcp.session import SessionStore, SessionState +from codewiki.mcp.tools.analysis import handle_analyze_repo +from codewiki.mcp.tools.code_reader import handle_read_code_components +from codewiki.mcp.tools.doc_writer import handle_write_doc_file, handle_edit_doc_file + +# Use the repo itself as a test target +REPO_PATH = str(Path(__file__).resolve().parent.parent) + +_passed = 0 +_failed = 0 + + +def check(name: str, condition: bool, detail: str = ""): + global _passed, _failed + if condition: + _passed += 1 + print(f" PASS: {name}") + else: + _failed += 1 + print(f" FAIL: {name} — {detail}") + + +def main(): + print("=== CodeWiki MCP Smoke Test (File-Side-Channel) ===\n") + + store = SessionStore() + output_dir = tempfile.mkdtemp(prefix="codewiki_smoke_") + + # -- 1. analyze_repo -- + print("[1] analyze_repo") + result = json.loads(handle_analyze_repo({ + "repo_path": REPO_PATH, + "output_dir": output_dir, + }, store)) + check("returns session_id", "session_id" in result, str(result)[:200]) + check("returns workspace_dir", "workspace_dir" in result, str(result.keys())) + check("returns stats", "stats" in result, str(result.keys())) + check("returns files", "files" in result, str(result.keys())) + check("stats has total_components", + "total_components" in result.get("stats", {}), + str(result.get("stats"))) + check("stats has total_leaf_nodes", + "total_leaf_nodes" in result.get("stats", {}), + str(result.get("stats"))) + + session_id = result.get("session_id") + workspace_dir = result.get("workspace_dir") + check("session_id is non-empty", session_id and len(session_id) == 12, str(session_id)) + check("workspace_dir exists on disk", + workspace_dir and Path(workspace_dir).is_dir(), + str(workspace_dir)) + + # -- 2. Workspace files verification -- + print("\n[2] Workspace files verification") + ws = Path(workspace_dir) + check("component_index.json exists", (ws / "component_index.json").exists(), "") + check("leaf_nodes.json exists", (ws / "leaf_nodes.json").exists(), "") + check("languages.json exists", (ws / "languages.json").exists(), "") + check("summary.json exists", (ws / "summary.json").exists(), "") + check("sources/ directory exists", (ws / "sources").is_dir(), "") + + # Read component_index.json and verify structure + comp_index = json.loads((ws / "component_index.json").read_text(encoding="utf-8")) + check("component_index is a list", isinstance(comp_index, list), type(comp_index).__name__) + check("component_index non-empty", len(comp_index) > 0, f"len={len(comp_index)}") + if comp_index: + first = comp_index[0] + check("component has id/type/file", + all(k in first for k in ("id", "type", "file")), + str(first.keys())) + + # Read leaf_nodes.json + leaf_nodes = json.loads((ws / "leaf_nodes.json").read_text(encoding="utf-8")) + check("leaf_nodes is a list", isinstance(leaf_nodes, list), type(leaf_nodes).__name__) + total_leaf = result["stats"]["total_leaf_nodes"] + check("leaf_nodes matches stats count", + len(leaf_nodes) == total_leaf, + f"file={len(leaf_nodes)} vs stats={total_leaf}") + + # -- 3. read_code_components (writes to workspace files) -- + print("\n[3] read_code_components") + if comp_index: + ids = [c["id"] for c in comp_index[:5]] + read_result = json.loads(handle_read_code_components({ + "session_id": session_id, + "component_ids": ids, + }, store)) + check("returns written count", "written" in read_result, str(read_result.keys())) + check("returns source_dir", "source_dir" in read_result, str(read_result.keys())) + check("returns files mapping", "files" in read_result, str(read_result.keys())) + check("written == requested", + read_result.get("written") == len(ids), + f"written={read_result.get('written')}, requested={len(ids)}") + + # Verify source files exist on disk + source_dir = Path(read_result["source_dir"]) + check("source_dir exists", source_dir.is_dir(), str(source_dir)) + for fname, cid in read_result.get("files", {}).items(): + src_file = source_dir / fname + if src_file.exists(): + content = src_file.read_text(encoding="utf-8") + check(f"source file has content ({fname})", + len(content) > 0, f"empty: {fname}") + check(f"source file has header ({fname})", + "Component:" in content, f"no header: {fname[:50]}") + break # just check first one + + # -- 4. read_code_components no cap (removed 20-component limit) -- + print("\n[4] read_code_components no cap") + if len(comp_index) > 20: + many_ids = [c["id"] for c in comp_index[:30]] + many_result = json.loads(handle_read_code_components({ + "session_id": session_id, + "component_ids": many_ids, + }, store)) + check("no 20-component cap", + many_result.get("written") == 30, + f"written={many_result.get('written')}") + + # -- 5. write_doc_file path traversal guard -- + print("\n[5] write_doc_file path traversal guard") + traversal_write = json.loads(asyncio.run(handle_write_doc_file_wrapper({ + "session_id": session_id, + "filename": "../../evil.md", + "content": "pwned", + }, store))) + check("rejects ../../evil.md", "error" in traversal_write, str(traversal_write)) + + # -- 6. write_doc_file normal write -- + print("\n[6] write_doc_file normal write") + normal_write = json.loads(asyncio.run(handle_write_doc_file_wrapper({ + "session_id": session_id, + "filename": "test_doc.md", + "content": "# Test\n\n```mermaid\ngraph TD\n A[Hello] --> B[World]\n```\n", + }, store))) + check("creates test_doc.md", normal_write.get("status") == "created", str(normal_write)) + check("file exists on disk", (Path(output_dir) / "test_doc.md").exists(), "") + + # -- 7. edit_doc_file str_replace -- + print("\n[7] edit_doc_file str_replace") + edit_result = json.loads(asyncio.run(handle_edit_doc_file_wrapper({ + "session_id": session_id, + "filename": "test_doc.md", + "command": "str_replace", + "old_str": "# Test", + "new_str": "# Test Edited", + }, store))) + check("edits file", edit_result.get("status") == "edited", str(edit_result)) + edited_content = (Path(output_dir) / "test_doc.md").read_text() + check("content updated", "# Test Edited" in edited_content, edited_content[:100]) + + # -- 8. edit_doc_file undo -- + print("\n[8] edit_doc_file undo") + undo_result = json.loads(asyncio.run(handle_edit_doc_file_wrapper({ + "session_id": session_id, + "filename": "test_doc.md", + "command": "undo", + }, store))) + check("undone", undo_result.get("status") == "undone", str(undo_result)) + check("mermaid_validation in undo", "mermaid_validation" in undo_result, str(undo_result.keys())) + undone_content = (Path(output_dir) / "test_doc.md").read_text() + check("content reverted", "# Test\n" in undone_content, undone_content[:100]) + + # -- 9. close_session with workspace cleanup -- + print("\n[9] close_session with workspace cleanup") + check("workspace exists before close", ws.exists(), "") + # Simulate close_session cleanup + session = store.get(session_id) + if session and session.workspace: + session.workspace.cleanup() + removed = store.remove(session_id) + check("session removed", removed, "") + check("workspace dir cleaned up", not ws.exists(), f"still exists: {ws}") + + # -- 10. SessionStore thread safety -- + print("\n[10] SessionStore thread safety") + import threading + errors = [] + def worker(): + try: + for _ in range(20): + s = store.create("a", "b", {}, []) + store.get(s.session_id) + store.remove(s.session_id) + except Exception as e: + errors.append(str(e)) + threads = [threading.Thread(target=worker) for _ in range(5)] + for t in threads: + t.start() + for t in threads: + t.join() + check("no concurrent access errors", len(errors) == 0, str(errors[:3])) + + # -- 11. SessionStore max sessions -- + print("\n[11] SessionStore max sessions") + store2 = SessionStore() + created = [] + for i in range(15): + s = store2.create(f"repo{i}", f"out{i}", {}, []) + created.append(s.session_id) + check("max 10 sessions enforced", len(store2._sessions) <= 10, f"got {len(store2._sessions)}") + + # -- Summary -- + print(f"\n=== Results: {_passed} passed, {_failed} failed ===") + return 1 if _failed else 0 + + +async def handle_write_doc_file_wrapper(args, store): + return await handle_write_doc_file(args, store) + + +async def handle_edit_doc_file_wrapper(args, store): + return await handle_edit_doc_file(args, store) + + +if __name__ == "__main__": + sys.exit(main())