refactor: apply 5-pass audit optimizations across mcp-memory codebase

This commit is contained in:
Riz Ashraf committed 2026-10-06 06:05:38 +01:00
1 parent 924b6d09fa
commit 5bd8b1587a
43 files changed
+1621 -1413

No files matched your search

Generated
-1
View File
@@ -2365,7 +2365,6 @@ dependencies = [
"blake3", "blake3",
"chrono", "chrono",
"clap", "clap",
"clipboard-win",
"dirs 7.0.0", "dirs 7.0.0",
"fastembed", "fastembed",
"futures-util", "futures-util",
+16
View File
@@ -80,6 +80,12 @@ Records structured diagnostic hypotheses, test evidence, and verification status
### 🔀 Context Workspace Diffing (`context_workspaces`) ### 🔀 Context Workspace Diffing (`context_workspaces`)
Computes structured diffs of pinned files and active task IDs between two saved context workspaces. Computes structured diffs of pinned files and active task IDs between two saved context workspaces.
### 🔒 Resilient Storage & Serde Parameter Tolerances
* **Atomic Store Write Locks**: Retains write guards during both in-memory updates and `redb` database serialization to eliminate lock-release TOCTOU race conditions.
* **Database Quarantine Protection**: Automatically flags corrupted database keys during deserialization (`is_corrupted = true`) to prevent corrupt states from being overwritten with empty defaults on subsequent writes.
* **Serde Relation Parameter Aliases**: `create_relations` accepts flexible aliases (`source` / `target` / `relationType` / `type`) so LLMs never encounter parameter validation errors.
* **Pre-cached Embedding Search**: Reuses pre-computed snippet embeddings (`snippet.embedding`), bypassing ONNX inference latency during in-memory semantic searches.
### 📡 Real-time WebSocket Memory Sync (`ws://127.0.0.1:3000/ws`) ### 📡 Real-time WebSocket Memory Sync (`ws://127.0.0.1:3000/ws`)
Broadcasting event pipeline streams real-time graph, task, and activity mutations directly to the Brain Monitor UI. Broadcasting event pipeline streams real-time graph, task, and activity mutations directly to the Brain Monitor UI.
@@ -154,3 +160,13 @@ To view the dashboard, open your browser at:
* **Kanban Board:** Track active Tasks and trigger status transitions directly from the browser. * **Kanban Board:** Track active Tasks and trigger status transitions directly from the browser.
* **Clipboard Inspector:** Review OS-level clipboard image captures via `/api/clipboard/capture`. * **Clipboard Inspector:** Review OS-level clipboard image captures via `/api/clipboard/capture`.
* **Live WebSocket Telemetry:** Real-time UI updates triggered by server state changes. * **Live WebSocket Telemetry:** Real-time UI updates triggered by server state changes.
---
## High-Performance Concurrency & Resilience Guarantees
* **Async Channel Backpressure (`push_async`)**: `Store::modify_async` uses `DbWriteQueue::push_async` with `tx.send(task).await` backpressure to guarantee database write persistence under heavy async write loads without dropping write transactions.
* **Atomic Search Index Swaps**: `MemoryState::rebuild_index` constructs and populates a new `MemoryIndex` instance in isolation before performing an atomic pointer swap (`*self.search_index.write().await = new_idx`), eliminating transient empty search result windows.
* **SIMD-Friendly Single-Pass Cosine Similarity**: `cosine_similarity` calculates dot product and Euclidean norm squares in a single linear pass over float vectors, enabling SIMD compiler auto-vectorization.
* **Safe Stream Decoding on Log Tails**: Log tail operations (`get_recent_logs`) read raw bytes and decode using lossy UTF-8 conversion (`String::from_utf8_lossy`) to ensure resilience when seeking across multi-byte UTF-8 boundaries.
* **Serde Parameter & Enum Tolerance**: All action enums (`StickyNoteAction`, `SnippetSearchMode`, `Relation`) support case-insensitive variants and field aliases (`source`/`from`, `target`/`to`, `relationType`/`relation_type`) to ensure seamless execution when LLMs pass varied string formatting.
+10
View File
@@ -18,3 +18,13 @@ The MCP Memory server is the central brain. You must be PROACTIVE, not reactive,
## 3. Delegation ## 3. Delegation
Continue to use the `MemoryLibrarian` subagent to log routine code changes (`log_code_change`) in the background to prevent cluttering the main conversation context. Continue to use the `MemoryLibrarian` subagent to log routine code changes (`log_code_change`) in the background to prevent cluttering the main conversation context.
## 4. Performance & Batching Rules
- **Batch Mutating Operations**: When creating or updating multiple graph entities, code snippets, or observations, always batch items into a single tool call array (e.g. `create_entities` with multiple items) to leverage the server's single-pass transaction flush.
- **High-Signal Tool Confirmations**: Tool call execution responses return structured, informative summaries (entity names, types, created counts, and edge paths). Agents DO NOT need to invoke follow-up `open_nodes` calls purely to confirm successful creation.
- **Tantivy Search Reader Refresh**: Search queries (`omni_search`, `search_nodes`) automatically reload pending commits prior to executing searches, ensuring immediate visibility of newly created items.
## 5. Pure Native Rust Invariants & Subprocess Prohibitions
- **Zero External Subprocesses**: Native system handlers (`clipboard`, `ast`, `search`, `db`) MUST use pure native Rust crates (`arboard`, `tree-sitter`, `tantivy`, `psycopg`). Subprocess calls to `powershell.exe`, `wl-paste`, `xclip`, or `cmd.exe` are strictly banned in native handlers.
- **Transient Lock Handling**: Transient OS handle collisions (e.g. Win32 OLE `OpenClipboard` locks) must be handled natively with retry loops and backoffs in Rust.
+10
View File
@@ -53,6 +53,8 @@ The daemon has completely eliminated raw JSON file sprawl and fragmented delta-f
### Key Principles: ### Key Principles:
* **Embedded Database Engine:** All structured components (Tasks, Snippets, Tech Debt, Checklists, etc.) are stored as binary-encoded values inside a unified `redb` database file (`store.redb`). * **Embedded Database Engine:** All structured components (Tasks, Snippets, Tech Debt, Checklists, etc.) are stored as binary-encoded values inside a unified `redb` database file (`store.redb`).
* **ACID Compliance & File Locks:** The Windows daemon holds an exclusive read-write lock on the database file, guaranteeing zero data corruption, race conditions, or lock contention during concurrent access. * **ACID Compliance & File Locks:** The Windows daemon holds an exclusive read-write lock on the database file, guaranteeing zero data corruption, race conditions, or lock contention during concurrent access.
* **Atomic Write-Guard Scope:** Store modification methods (`Store::modify` and `Store::modify_async`) retain the write lock through both the in-memory mutation and JSON serialization phases, eliminating lock-release TOCTOU race conditions.
* **Store Quarantine Mode:** If deserialization fails during `Store::load_from_db`, the store flags `is_corrupted = true` and refuses to overwrite database keys with default values on subsequent writes.
* **Asynchronous Checkpointing:** The core Knowledge Graph (Entities, Relations, Observations) still utilizes a Write-Ahead Logging (WAL) pattern (`wal.jsonl`) and a master snapshot (`master.json`) to allow safe, lock-free memory mutations which are reconciled in the background. * **Asynchronous Checkpointing:** The core Knowledge Graph (Entities, Relations, Observations) still utilizes a Write-Ahead Logging (WAL) pattern (`wal.jsonl`) and a master snapshot (`master.json`) to allow safe, lock-free memory mutations which are reconciled in the background.
## 6. Domain Models & Component Stores ## 6. Domain Models & Component Stores
@@ -67,6 +69,7 @@ Currently implemented persistent stores include:
## 7. Full-Text & Semantic Search Engine (Tantivy + FastEmbed) ## 7. Full-Text & Semantic Search Engine (Tantivy + FastEmbed)
To support blazing-fast, intelligent semantic retrieval across the sprawling knowledge graph, the daemon embeds **Tantivy** (a full-text search engine inspired by Apache Lucene) alongside **FastEmbed** (a local ONNX runtime for vector embeddings). To support blazing-fast, intelligent semantic retrieval across the sprawling knowledge graph, the daemon embeds **Tantivy** (a full-text search engine inspired by Apache Lucene) alongside **FastEmbed** (a local ONNX runtime for vector embeddings).
* **The `MemoryIndex`:** Whenever the graph or auxiliary stores mutate, a background thread dynamically rebuilds the Tantivy index (`tantivy_index/` dir) and computes semantic vectors. * **The `MemoryIndex`:** Whenever the graph or auxiliary stores mutate, a background thread dynamically rebuilds the Tantivy index (`tantivy_index/` dir) and computes semantic vectors.
* **Pre-cached Vector Embeddings:** `SearchService::semantic_search` reuses pre-cached snippet embedding vectors (`snippet.embedding`), bypassing redundant ONNX neural network inference calls during query execution.
* **Global Omni-Search:** This architecture powers the `omni_search` tool, allowing subagents to instantly fuzzy-search and semantically rank documents across Entities, Tasks, Snippets, Error Fixes, and ADRs simultaneously in milliseconds, without loading massive JSON arrays into RAM. * **Global Omni-Search:** This architecture powers the `omni_search` tool, allowing subagents to instantly fuzzy-search and semantically rank documents across Entities, Tasks, Snippets, Error Fixes, and ADRs simultaneously in milliseconds, without loading massive JSON arrays into RAM.
## 8. Webhook Telemetry & Passive Ingestion ## 8. Webhook Telemetry & Passive Ingestion
@@ -187,3 +190,10 @@ It supports reading and writing rich formats natively to the Windows Host OS usi
* **File Drops (CF_HDROP):** The server can parse file lists copied from Windows Explorer, and can inversely synthesize file drops into the clipboard from absolute paths. * **File Drops (CF_HDROP):** The server can parse file lists copied from Windows Explorer, and can inversely synthesize file drops into the clipboard from absolute paths.
* **Images (CF_BITMAP):** The server natively rasterizes clipboard bitmaps to JPEG on read, and can write raw RgbaImage buffers back to the clipboard on write. * **Images (CF_BITMAP):** The server natively rasterizes clipboard bitmaps to JPEG on read, and can write raw RgbaImage buffers back to the clipboard on write.
* **Developer Tooling:** read_file_skeleton (AST), get_active_worktree_context (Git), get_recent_logs, toggle_clipboard_watch_mode. * **Developer Tooling:** read_file_skeleton (AST), get_active_worktree_context (Git), get_recent_logs, toggle_clipboard_watch_mode.
## 19. High-Performance Concurrency & Resilience Guarantees
* **Async Channel Backpressure (`push_async`)**: `Store::modify_async` uses `DbWriteQueue::push_async` with `tx.send(task).await` backpressure to guarantee database write persistence under heavy async write loads without dropping write transactions.
* **Atomic Search Index Swaps**: `MemoryState::rebuild_index` constructs and populates a new `MemoryIndex` instance in isolation before performing an atomic pointer swap (`*self.search_index.write().await = new_idx`), eliminating transient empty search result windows.
* **SIMD-Friendly Single-Pass Cosine Similarity**: `cosine_similarity` calculates dot product and Euclidean norm squares in a single linear pass over float vectors, enabling SIMD compiler auto-vectorization.
* **Safe Stream Decoding on Log Tails**: Log tail operations (`get_recent_logs`) read raw bytes and decode using lossy UTF-8 conversion (`String::from_utf8_lossy`) to ensure resilience when seeking across multi-byte UTF-8 boundaries.
* **Serde Parameter & Enum Tolerance**: All action enums (`StickyNoteAction`, `SnippetSearchMode`, `Relation`) support case-insensitive variants and field aliases (`source`/`from`, `target`/`to`, `relationType`/`relation_type`) to ensure seamless execution when LLMs pass varied string formatting.
-10
View File
@@ -1,10 +0,0 @@
$lines = Get-Content justfile
$valid_lines = @()
foreach ($line in $lines) {
if ($line -match "Building Windows Nvim" -or $line -match "cargo build --release -p mcp-memory-win-nvim" -or $line -match "^build-nvim-win:") {
continue
}
$valid_lines += $line
}
$valid_lines | Set-Content justfile
+64 -10
View File
@@ -11,9 +11,11 @@ To prevent graph fragmentation and ensure seamless LLM context retrieval:
* **Entity Types (`entity_type`)**: MUST ALWAYS be **`PascalCase`** (e.g. `DatabaseTable`, `McpTool`, `ArchitectureComponent`, `File`, `DataStructure`). * **Entity Types (`entity_type`)**: MUST ALWAYS be **`PascalCase`** (e.g. `DatabaseTable`, `McpTool`, `ArchitectureComponent`, `File`, `DataStructure`).
* **Relation Types (`relation_type`)**: MUST ALWAYS be **`snake_case`** (e.g. `depends_on`, `calls`, `implements`, `uses`, `contains`). * **Relation Types (`relation_type`)**: MUST ALWAYS be **`snake_case`** (e.g. `depends_on`, `calls`, `implements`, `uses`, `contains`).
* **Field Keys & Properties**: MUST ALWAYS be **`snake_case`** (e.g. `file_path`, `git_commit`, `created_at`). * **Field Keys & Properties**: MUST ALWAYS be **`snake_case`** (e.g. `file_path`, `git_commit`, `created_at`).
* **Relation Tool Parameters**: `create_relations` supports Serde field aliases (`source` -> `from`, `target` -> `to`, `relationType`/`type` -> `relation_type`) so LLM tool calls succeed seamlessly regardless of parameter naming.
> [!NOTE] > [!NOTE]
> The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation. > The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation.
> Store operations enforce atomic lock scope for serialization/deserialization and enter Quarantine Mode upon database corruption to prevent data overwrite hazards.
--- ---
@@ -92,12 +94,30 @@ The server consolidates granular single-purpose tools into domain-named smart to
--- ---
## 3. Subgraph Expansion & Multi-Hop Navigation ## 4. High-Signal Tool Responses & Performance Guidelines
- **Tool:** `get_subgraph`
- **When to use:** When you need to understand the complete architectural neighborhood surrounding a specific component, module, or database table.
- **Behavior:** Performs a multi-hop Breadth-First Search (BFS) around a `root_node` (or `root_entity`) up to a requested `depth` (e.g. 1 to 3 hops) and returns all connected entities and relations. Pass `format: "markdown_tree"` to generate a compact, token-budgeted Markdown topology tree capped within a requested `max_tokens` budget.
--- To optimize context usage, response times, and LLM reasoning efficiency:
* **High-Signal Feedback**:
- MCP tool calls (such as `create_entities`, `create_relations`, `add_observations`, `pin_file`) return explicit, structured summaries containing created names, types, edge counts, and relation paths.
- LLMs do NOT need to execute follow-up `open_nodes` calls merely to confirm successful creation.
* **Batch Operations**:
- When creating or updating multiple entities, snippets, or observations, always batch items into a single tool call array (e.g. `create_entities` with multiple array items) rather than making separate calls.
- The server handles batch store mutations serially in a single transaction pass with single-permit event-driven flushes.
* **Real-time Tantivy Search Indexing**:
- The Tantivy search engine automatically checks pending commits and reloads search readers prior to executing `omni_search` or `search_nodes`. Search queries always return up-to-date document results immediately following mutations.
- Single-item deletions use targeted document removal rather than global index wipes.
* **Real-Time AST & Workspace Source Code Symbol Scanning**:
- `find_symbol_references`, `get_callers`, and `analyze_impact` scan both stored code snippets and physical workspace source code files on disk (`.rs`, `.ts`, `.py`, `.go`, `.java`, `.c`, `.cpp`), providing accurate AST symbol references and call site tracking.
* **Graph Entity Merge & Self-Loop Protection**:
- `merge_entities` re-points all relations from `source_entity` to `target_entity` and automatically prunes cyclic self-loops (`target -> target`).
* **Safe UTF-8 Token Truncation**:
- Large responses (e.g. `get_active_worktree_context`, `read_graph`, `summarize_subgraph`) are safely truncated along UTF-8 character boundaries (`floor_char_boundary`), ensuring response bounds without runtime panics.
## 4. Automated Error Fix Auto-Matcher ## 4. Automated Error Fix Auto-Matcher
- **Tools:** `log_error_fix`, `search_error_fixes` (and alias `suggest_error_fix`) - **Tools:** `log_error_fix`, `search_error_fixes` (and alias `suggest_error_fix`)
@@ -116,7 +136,7 @@ The server consolidates granular single-purpose tools into domain-named smart to
## 6. Self-Healing Graph Health Sweeper ## 6. Self-Healing Graph Health Sweeper
- **Tool:** `sweep_graph_health` - **Tool:** `sweep_graph_health`
- **When to use:** Periodically or before committing major graph changes to audit entity consistency. - **When to use:** Periodically or before committing major graph changes to audit entity consistency.
- **Behavior:** Detects orphaned nodes (0 relations), computes name similarity to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans. - **Behavior:** Detects orphaned nodes (0 relations), computes name similarity using pre-computed lowercase keys to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans.
--- ---
@@ -127,7 +147,41 @@ The server consolidates granular single-purpose tools into domain-named smart to
--- ---
## 8. LLM Pre-Flight Context Bundle ## 9. Native Rust Invariants & Subprocess Prohibition (CRITICAL)
- **Tool:** `get_preflight_context`
- **When to use:** At the start of a turn or subagent task to gain total situational awareness in 1 call. To maintain maximum security, speed, and cross-platform reliability:
- **Behavior:** Aggregates current active branch, in-progress tasks with acceptance criteria, pinned files, top open tech debts, and active unverified hypotheses into a consolidated executive context bundle. * **Zero Subprocess Fallbacks**: System and server tools (`clipboard`, `ast`, `search`, `db`) MUST strictly use pure native Rust crates (`arboard`, `tree-sitter`, `tantivy`, `psycopg`). Invocations of external shell commands (`powershell.exe`, `wl-paste`, `xclip`, `cmd.exe`) are strictly prohibited in native handlers.
* **Transient Lock Recovery**: Transient OS handle collisions (such as Win32 OLE `OpenClipboard` lock contention) must be handled using native retry loops with backoffs directly in Rust.
* **Automated Static Regression Gates**: Automated AST/source audit tests (e.g. `test_no_subprocess_clipboard_regression`) verify at test time that forbidden subprocess patterns are absent from handler implementations.
---
## 10. High-Performance Concurrency & Resilience Guarantees
* **Zero-Allocation NDJSON Reader**: `mcp-stdio` reclaims line buffers via `std::mem::take` and in-place trimming to eliminate heap allocations during high-frequency NDJSON message parsing.
* **Micro-Batched Fastembed Inference**: `generate_embeddings_async` chunks text arrays into 32-item micro-batches inside `spawn_blocking`, eliminating RAM/CPU spikes during batch indexing.
* **Single-Handle OS Clipboard Retries & Image Downsampling**: `WriteClipboardHandler` initializes OS clipboard handles once per operation and downsamples images exceeding $2048 \times 2048$ resolution before writing raw RGBA bytes.
* **Zero Transaction Drop Persistence Guarantee**: `Store::modify` automatically spawns an async task to execute `push_async` with channel backpressure if `push` encounters queue saturation, ensuring zero data loss under spike write loads.
* **Non-Blocking Read Lock Sticky Notes**: `StickyNotesHandler` (`action = "read"`) queries notes using shared read locks, executing write pruning only when expired items exist.
* **Token-Budgeted Query Projections**: Decision queries (`query_decisions`) support `limit` caps and compact `include_body: false` projections for token budget optimization.
* **Serde Parameter & Enum Ergonomics**: Action enums (`StickyNoteAction`, `SnippetSearchMode`, `Relation`) support case-insensitive variants and common synonyms (`create`/`add`, `remove`/`delete`, `list`/`read`, `source`/`from`, `target`/`to`, `relationType`/`relation_type`) ensuring seamless LLM tool execution.
* **Atomic Search Index Swaps**: `MemoryState::rebuild_index` constructs and populates a new `MemoryIndex` instance in isolation before performing an atomic pointer swap (`*self.search_index.write().await = new_idx`), eliminating transient empty search result windows.
* **Non-Blocking Tantivy Search Queries**: `MemoryIndex::search` queries current index searcher snapshots without executing synchronous disk commit locks, preventing query stalls during heavy background indexing.
* **Zero-Allocation HashSet<&str> Snippet Deduplication**: `indexer.rs` utilizes borrowed `HashSet<&str>` name lookups during snippet batch modifications, eliminating heap string re-allocations inside the store write lock.
* **AST Recursion Depth Safeguard & Zero-Copy Borrowing**: Tree-sitter AST traversal caps recursion depth at 100 to prevent thread stack overflows and borrows string slices (`&str`) during AST node walking.
* **Strongly-Typed SearchResult & Pre-Allocated Search Vectors**: `search.rs` uses a strongly-typed `SearchResult` struct with named fields and pre-allocates result vector capacity (`Vec::with_capacity(top_docs.len())`).
* **BFS Graph Traversal Pre-allocation & Visited Node Upper Bound**: `GraphQueryBuilder::find_shortest_path` pre-allocates adjacency map capacity (`HashMap::with_capacity(relations.len() * 2)`) and enforces a visited node upper bound (10,000 max) to guarantee deterministic BFS runtime.
* **LLM Tool Schema Ergonomics & Context Guidance**: `ReadGraphHandler` schema explicitly instructs LLMs on `namespace` filtering and `search_nodes` / `get_subgraph` tools for large graph discovery.
* **Filesystem Event Debouncing & Proactive State Refresh**: `spawn_watcher` implements a sliding 250ms debouncing window per file path, ignores `.git`, `target`, `.gemini`, and `node_modules`, and broadcasts activity events to `MemoryState`.
* **Buffered Line-by-Line AST Workspace Symbol Scanning**: `scan_workspace_for_symbol` reads workspace files via `BufReader` line streams instead of loading entire files into heap strings, preventing memory spikes when traversing source trees.
* **AST Node Type Aliasing & Skeleton Preallocation**: `replace_ast_node` documents friendly node aliases (`function`, `fn`, `method`, `struct`, `class`, `enum`, `trait`, `type`), and `read_file_skeleton` preallocates string buffer capacity (`code.len() / 2`).
* **Batch Vector Indexing & Similarity Score Guidance**: `VectorDB` provides `index_documents_batch` for single-request multi-point vector upserts and explicit score calibration notes ($\ge 0.75$ high confidence match).
* **Compact JSON MCP Resources & UTF-8 Activity Truncation**: MCP resources serialize using compact JSON (`to_string`), `TerminalHistoryResource` / `MilestonesResource` enforce output bounds, and `format_tool_activity_description` uses `floor_char_boundary` for guaranteed UTF-8 safety.
* **SIMD-Friendly Single-Pass Cosine Similarity**: `cosine_similarity` calculates dot product and Euclidean norm squares in a single linear pass over float vectors, enabling SIMD compiler auto-vectorization.
* **Safe Stream Decoding on Log Tails**: Log tail operations (`get_recent_logs`) read raw bytes and decode using lossy UTF-8 conversion (`String::from_utf8_lossy`) to ensure resilience when seeking across multi-byte UTF-8 boundaries.
* **Task Summary UTF-8 Truncation Safety**: `tasks` tool (`action = "list"`) truncates serialized task text strictly along UTF-8 character boundaries using `floor_char_boundary` when enforcing `max_tokens`.
* **Sequential Snapshot Lock Scope Flattening**: `GenerateStandupReportHandler` reads `tasks`, `ledger`, and `session_summaries` sequentially rather than nesting read locks, preventing multi-lock deadlocks during concurrent store modifications.
* **Directory Tree Depth Safeguard**: `ReadDirectoryArchitectureHandler` caps directory recursion at depth 10 to prevent stack overflow on deep or cyclic directory structures.
* **Deterministic Total-Order Score Ranking**: `OmniSearchHandler` uses `f64::total_cmp` for Reciprocal Rank Fusion (RRF) score sorting, guaranteeing deterministic NaN-safe search result ordering.
* **RPC Timeout Memory Hygiene**: `nvim-core` maintains request hygiene by removing pending request entries from static RPC maps upon timeout or channel drop, eliminating orphan memory leaks.
+14 -21
View File
@@ -99,50 +99,43 @@ build-nvim:
# 3. DEPLOY (Move Binaries and Rules to Global Paths) # 3. DEPLOY (Move Binaries and Rules to Global Paths)
# ========================================================= # =========================================================
# Stage updated server binary, perform graceful shutdown, and auto-restart # Stage updated server binary, perform graceful shutdown, and auto-restart (skips if binary is unchanged)
deploy-server: deploy-server:
@Write-Host "Deploying Windows Server (Staged Hot-Swap)..." -ForegroundColor Cyan @Write-Host "Deploying Windows Server (Staged Hot-Swap)..." -ForegroundColor Cyan
-& { $targets = @("$env:USERPROFILE\.local\bin\mcp-memory-server.exe", "$env:USERPROFILE\.cargo\bin\mcp-memory-server.exe"); foreach ($target in $targets) { if (Test-Path (Split-Path $target)) { if (Test-Path $target) { $bak = "$target.old.$([guid]::NewGuid().ToString().Substring(0,8))"; Move-Item -Force $target $bak -ErrorAction SilentlyContinue }; Copy-Item -Force target\release\mcp-memory-server.exe $target; Get-ChildItem "$(Split-Path $target)\mcp-memory-server.exe.old.*" -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue } } } -& { $src = "target\release\mcp-memory-server.exe"; $targets = @("$env:USERPROFILE\.local\bin\mcp-memory-server.exe", "$env:USERPROFILE\.cargo\bin\mcp-memory-server.exe"); $srcHash = if (Test-Path $src) { (Get-FileHash $src -Algorithm SHA256).Hash } else { $null }; $updated = $false; foreach ($target in $targets) { $parent = Split-Path $target; if (Test-Path $parent) { if ((Test-Path $target) -and $srcHash -and ((Get-FileHash $target -Algorithm SHA256).Hash -eq $srcHash)) { Write-Host "Skipping copy to $target (identical checksum)." -ForegroundColor Yellow } else { if (Test-Path $target) { $bak = "$target.old.$([guid]::NewGuid().ToString().Substring(0,8))"; Move-Item -Force $target $bak -ErrorAction SilentlyContinue }; Copy-Item -Force $src $target; Get-ChildItem "$parent\mcp-memory-server.exe.old.*" -ErrorAction SilentlyContinue | Remove-Item -Force -ErrorAction SilentlyContinue; Write-Host "Copied updated binary to $target" -ForegroundColor Green; $updated = $true } } }; if ($updated -or -not (Get-NetTCPConnection -LocalPort 3000 -State Listen -ErrorAction SilentlyContinue)) { just stop; just start } else { Write-Host "Server binary unchanged and server is running. Skipping restart." -ForegroundColor Green } }
-just stop
just start
# Deploy Windows-native stub binary # Deploy Windows-native stub binary (skips if identical checksum)
deploy-stub-win: deploy-stub-win:
@Write-Host "Deploying Windows Stub..." -ForegroundColor Cyan @Write-Host "Deploying Windows Stub..." -ForegroundColor Cyan
-Stop-Process -Name mcp-memory-stub -Force -ErrorAction SilentlyContinue -& { $src = "target\release\mcp-memory-stub.exe"; $dest = "C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe"; if ((Test-Path $src) -and (Test-Path $dest) -and ((Get-FileHash $src -Algorithm SHA256).Hash -eq (Get-FileHash $dest -Algorithm SHA256).Hash)) { Write-Host "Skipping $dest (identical checksum)." -ForegroundColor Yellow } else { Stop-Process -Name mcp-memory-stub -Force -ErrorAction SilentlyContinue; if (Test-Path $dest) { Move-Item -Force $dest "$dest.old" -ErrorAction SilentlyContinue; Remove-Item -Force "$dest.old" -ErrorAction SilentlyContinue }; Copy-Item -Force $src $dest; Write-Host "Copied updated $src to $dest" -ForegroundColor Green } }
-if (Test-Path "C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe") { Move-Item -Force "C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe" "C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe.old" -ErrorAction SilentlyContinue; Remove-Item -Force "C:\Users\reazul.ashraf\.local\bin\mcp-memory-stub.exe.old" -ErrorAction SilentlyContinue }
Copy-Item -Force target\release\mcp-memory-stub.exe "C:\Users\reazul.ashraf\.local\bin\"
# Deploy Windows-native Neovim plugin binary # Deploy Windows-native Neovim plugin binary (skips if identical checksum)
deploy-nvim-win: deploy-nvim-win:
@Write-Host "Deploying Windows Nvim..." -ForegroundColor Cyan @Write-Host "Deploying Windows Nvim..." -ForegroundColor Cyan
-Stop-Process -Name mcp-memory-nvim -Force -ErrorAction SilentlyContinue -& { $src = "target\release\mcp-memory-nvim.exe"; $dest = "C:\Users\reazul.ashraf\.local\bin\mcp-memory-nvim.exe"; if ((Test-Path $src) -and (Test-Path $dest) -and ((Get-FileHash $src -Algorithm SHA256).Hash -eq (Get-FileHash $dest -Algorithm SHA256).Hash)) { Write-Host "Skipping $dest (identical checksum)." -ForegroundColor Yellow } else { Stop-Process -Name mcp-memory-nvim -Force -ErrorAction SilentlyContinue; if (Test-Path $dest) { Move-Item -Force $dest "$dest.old" -ErrorAction SilentlyContinue; Remove-Item -Force "$dest.old" -ErrorAction SilentlyContinue }; Copy-Item -Force $src $dest; Write-Host "Copied updated $src to $dest" -ForegroundColor Green } }
-if (Test-Path "C:\Users\reazul.ashraf\.local\bin\mcp-memory-nvim.exe") { Move-Item -Force "C:\Users\reazul.ashraf\.local\bin\mcp-memory-nvim.exe" "C:\Users\reazul.ashraf\.local\bin\mcp-memory-nvim.exe.old" -ErrorAction SilentlyContinue; Remove-Item -Force "C:\Users\reazul.ashraf\.local\bin\mcp-memory-nvim.exe.old" -ErrorAction SilentlyContinue }
Copy-Item -Force target\release\mcp-memory-nvim.exe "C:\Users\reazul.ashraf\.local\bin\"
# Deploy global Agent rules to Windows config directory # Deploy global Agent rules to Windows config directory (skips identical files)
deploy-rules-win: deploy-rules-win:
@Write-Host "Deploying global Agent rules (Windows)..." -ForegroundColor Cyan @Write-Host "Deploying global Agent rules (Windows)..." -ForegroundColor Cyan
if (!(Test-Path "C:\Users\reazul.ashraf\.gemini\config\rules")) { New-Item -ItemType Directory -Force -Path "C:\Users\reazul.ashraf\.gemini\config\rules" | Out-Null } -& { $destDir = "C:\Users\reazul.ashraf\.gemini\config\rules"; if (!(Test-Path $destDir)) { New-Item -ItemType Directory -Force -Path $destDir | Out-Null }; $srcBase = (Get-Item "agent-rules").FullName; $updatedCount = 0; $skippedCount = 0; Get-ChildItem -Recurse agent-rules -File | ForEach-Object { $rel = $_.FullName.Substring($srcBase.Length + 1); $target = Join-Path $destDir $rel; $targetDir = Split-Path $target; if (!(Test-Path $targetDir)) { New-Item -ItemType Directory -Force -Path $targetDir | Out-Null }; if ((Test-Path $target) -and ((Get-FileHash $_.FullName -Algorithm SHA256).Hash -eq (Get-FileHash $target -Algorithm SHA256).Hash)) { $skippedCount++ } else { Copy-Item -Force $_.FullName $target; Write-Host "Updated rule: $rel" -ForegroundColor Green; $updatedCount++ } }; Write-Host "Agent rules deployment complete ($updatedCount updated, $skippedCount skipped identical)." -ForegroundColor Cyan }
Copy-Item -Force -Recurse agent-rules\* "C:\Users\reazul.ashraf\.gemini\config\rules\"
# Deploy all WSL-native binaries and agent rules # Deploy all WSL-native binaries and agent rules
deploy-wsl: deploy-stub deploy-nvim deploy-rules-wsl deploy-wsl: deploy-stub deploy-nvim deploy-rules-wsl
# Deploy WSL-native stub binary # Deploy WSL-native stub binary (skips if identical checksum)
deploy-stub: deploy-stub:
@Write-Host "Deploying WSL Stub natively..." -ForegroundColor Cyan @Write-Host "Deploying WSL Stub natively..." -ForegroundColor Cyan
wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && cp target/release/mcp-memory-stub /home/riz/.local/bin/' wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && src="target/release/mcp-memory-stub" && dest="/home/riz/.local/bin/mcp-memory-stub" && if [ -f "$dest" ] && cmp -s "$src" "$dest"; then echo -e "\033[33mSkipping $dest (identical checksum).\033[0m"; else cp "$src" "$dest" && echo -e "\033[32mCopied updated $src to $dest\033[0m"; fi'
# Deploy WSL-native Neovim plugin binary # Deploy WSL-native Neovim plugin binary (skips if identical checksum)
deploy-nvim: deploy-nvim:
@Write-Host "Deploying WSL Nvim natively..." -ForegroundColor Cyan @Write-Host "Deploying WSL Nvim natively..." -ForegroundColor Cyan
wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && cp target/release/mcp-memory-nvim /home/riz/.local/bin/' wsl.exe -d Ubuntu -e bash -c 'export PATH="/home/riz/.cargo/bin:$PATH" && cd /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory && src="target/release/mcp-memory-nvim" && dest="/home/riz/.local/bin/mcp-memory-nvim" && if [ -f "$dest" ] && cmp -s "$src" "$dest"; then echo -e "\033[33mSkipping $dest (identical checksum).\033[0m"; else cp "$src" "$dest" && echo -e "\033[32mCopied updated $src to $dest\033[0m"; fi'
# Deploy global Agent rules to WSL config directory # Deploy global Agent rules to WSL config directory (skips identical files)
deploy-rules-wsl: deploy-rules-wsl:
@Write-Host "Deploying global Agent rules (WSL)..." -ForegroundColor Cyan @Write-Host "Deploying global Agent rules (WSL)..." -ForegroundColor Cyan
wsl.exe -d Ubuntu -e bash -c 'mkdir -p /home/riz/.gemini/config/rules && cp -r /mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory/agent-rules/* /home/riz/.gemini/config/rules/' wsl.exe -d Ubuntu -e bash -c 'mkdir -p /home/riz/.gemini/config/rules && src_dir="/mnt/c/Users/reazul.ashraf/workspace/rust/mcp-memory/agent-rules" && dest_dir="/home/riz/.gemini/config/rules" && updated=0 && skipped=0 && for src in $(find "$src_dir" -type f); do rel="${src#$src_dir/}"; dest="$dest_dir/$rel"; mkdir -p "$(dirname "$dest")"; if [ -f "$dest" ] && cmp -s "$src" "$dest"; then ((skipped++)); else cp "$src" "$dest" && echo -e "\033[32mUpdated rule: $rel\033[0m" && ((updated++)); fi; done && echo -e "\033[36mAgent rules WSL deployment complete ($updated updated, $skipped skipped identical).\033[0m"'
# ========================================================= # =========================================================
# 4. TESTING & COVERAGE # 4. TESTING & COVERAGE
+3
View File
@@ -5,3 +5,6 @@ edition = "2024"
[dependencies] [dependencies]
tokio = { version = "1.53.1", features = ["io-util", "macros"] } tokio = { version = "1.53.1", features = ["io-util", "macros"] }
[dev-dependencies]
tokio = { version = "1.53.1", features = ["rt", "macros"] }
+10 -4
View File
@@ -15,8 +15,14 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
let trimmed_start = line.trim_start(); let trimmed_start = line.trim_start();
if trimmed_start.starts_with('{') { if trimmed_start.starts_with('{') {
let trimmed = line.trim(); let mut val = std::mem::take(&mut line);
return Some(trimmed.to_string()); let trimmed_len = val.trim_end().len();
val.truncate(trimmed_len);
let start = val.len() - val.trim_start().len();
if start > 0 {
val.drain(..start);
}
return Some(val);
} }
let line = line.trim_end(); let line = line.trim_end();
@@ -24,7 +30,7 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
break; break;
} }
if line.len() >= 15 && line[..15].eq_ignore_ascii_case("content-length:") { if line.as_bytes().len() >= 15 && line.as_bytes()[..15].eq_ignore_ascii_case(b"content-length:") {
length = line[15..].trim().parse().unwrap_or(0); length = line[15..].trim().parse().unwrap_or(0);
} }
} }
@@ -35,7 +41,7 @@ pub async fn read_mcp_message<R: tokio::io::AsyncRead + Unpin>(
return None; return None;
} }
let mut buffer = Vec::with_capacity(length); let mut buffer = Vec::with_capacity(length.min(64 * 1024));
if stdin.take(length as u64).read_to_end(&mut buffer).await.is_err() || buffer.len() != length { if stdin.take(length as u64).read_to_end(&mut buffer).await.is_err() || buffer.len() != length {
return None; return None;
} }
+62 -16
View File
@@ -38,8 +38,12 @@ pub async fn send_response(response: JsonRpcResponse) {
payload.extend_from_slice(msg.as_bytes()); payload.extend_from_slice(msg.as_bytes());
payload.push(b'\n'); payload.push(b'\n');
let mut stdout = tokio::io::stdout(); let mut stdout = tokio::io::stdout();
let _ = stdout.write_all(&payload).await; if let Err(e) = stdout.write_all(&payload).await {
let _ = stdout.flush().await; tracing::error!("Failed to write response payload to stdout: {}", e);
}
if let Err(e) = stdout.flush().await {
tracing::error!("Failed to flush stdout: {}", e);
}
} }
macro_rules! handle_lua_result { macro_rules! handle_lua_result {
@@ -177,6 +181,8 @@ static RPC_SEMAPHORE: LazyLock<Arc<tokio::sync::Semaphore>> =
LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(100))); LazyLock::new(|| Arc::new(tokio::sync::Semaphore::new(100)));
static NVIM_CONN: LazyLock<Arc<std::sync::Mutex<Option<mpsc::Sender<NvimRequest>>>>> = static NVIM_CONN: LazyLock<Arc<std::sync::Mutex<Option<mpsc::Sender<NvimRequest>>>>> =
LazyLock::new(|| Arc::new(std::sync::Mutex::new(None))); LazyLock::new(|| Arc::new(std::sync::Mutex::new(None)));
static PENDING_REQUESTS: LazyLock<dashmap::DashMap<u64, oneshot::Sender<Result<rmpv::Value, String>>>> =
LazyLock::new(dashmap::DashMap::new);
#[derive(Debug, Default, Clone)] #[derive(Debug, Default, Clone)]
pub struct NvimState { pub struct NvimState {
@@ -293,12 +299,8 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
let (mut read_half, mut write_half) = tokio::io::split(stream); let (mut read_half, mut write_half) = tokio::io::split(stream);
let (tx, mut rx) = mpsc::channel::<NvimRequest>(32); let (tx, mut rx) = mpsc::channel::<NvimRequest>(32);
type PendingRequestsMap =
Arc<dashmap::DashMap<u64, oneshot::Sender<Result<rmpv::Value, String>>>>;
let pending_requests: PendingRequestsMap = Arc::new(dashmap::DashMap::new());
// Write task // Write task
let pending_clone = Arc::clone(&pending_requests);
tokio::spawn(async move { tokio::spawn(async move {
while let Some(req) = rx.recv().await { while let Some(req) = rx.recv().await {
let mut buf = Vec::new(); let mut buf = Vec::new();
@@ -307,11 +309,11 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
continue; continue;
} }
pending_clone.insert(req.msgid, req.reply); PENDING_REQUESTS.insert(req.msgid, req.reply);
if write_half.write_all(&buf).await.is_err() { if write_half.write_all(&buf).await.is_err() {
tracing::error!("Failed to write to Neovim socket"); tracing::error!("Failed to write to Neovim socket");
if let Some((_, sender)) = pending_clone.remove(&req.msgid) { if let Some((_, sender)) = PENDING_REQUESTS.remove(&req.msgid) {
let _ = sender.send(Err("Connection closed during write".to_string())); let _ = sender.send(Err("Connection closed during write".to_string()));
} }
break; break;
@@ -320,7 +322,6 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
}); });
// Read task // Read task
let pending_clone2 = Arc::clone(&pending_requests);
tokio::spawn(async move { tokio::spawn(async move {
use bytes::{Buf, BytesMut}; use bytes::{Buf, BytesMut};
let mut resp_buf = BytesMut::with_capacity(65536); let mut resp_buf = BytesMut::with_capacity(65536);
@@ -338,7 +339,7 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
rmpv::Value::Integer(i) => i.as_u64().unwrap_or(0), rmpv::Value::Integer(i) => i.as_u64().unwrap_or(0),
_ => 0, _ => 0,
}; };
if let Some((_, reply_sender)) = pending_clone2.remove(&msgid) { if let Some((_, reply_sender)) = PENDING_REQUESTS.remove(&msgid) {
let _ = reply_sender.send(Ok(val)); let _ = reply_sender.send(Ok(val));
} }
} else if arr.len() >= 3 } else if arr.len() >= 3
@@ -382,9 +383,9 @@ async fn get_nvim_connection() -> Result<mpsc::Sender<NvimRequest>, String> {
} }
// Cleanup pending requests on disconnect // Cleanup pending requests on disconnect
let keys: Vec<_> = pending_clone2.iter().map(|kv| *kv.key()).collect(); let keys: Vec<_> = PENDING_REQUESTS.iter().map(|kv| *kv.key()).collect();
for k in keys { for k in keys {
if let Some((_, sender)) = pending_clone2.remove(&k) { if let Some((_, sender)) = PENDING_REQUESTS.remove(&k) {
let _ = sender.send(Err("Connection closed".to_string())); let _ = sender.send(Err("Connection closed".to_string()));
} }
} }
@@ -469,8 +470,14 @@ async fn call_nvim(req: rmpv::Value) -> Result<rmpv::Value, String> {
match tokio::time::timeout(tokio::time::Duration::from_secs(30), reply_rx).await { match tokio::time::timeout(tokio::time::Duration::from_secs(30), reply_rx).await {
Ok(Ok(res)) => res, Ok(Ok(res)) => res,
Ok(Err(_)) => Err("Response channel dropped".to_string()), Ok(Err(_)) => {
Err(_) => Err("Timeout waiting for Neovim response".to_string()), PENDING_REQUESTS.remove(&msgid);
Err("Response channel dropped".to_string())
}
Err(_) => {
PENDING_REQUESTS.remove(&msgid);
Err("Timeout waiting for Neovim response".to_string())
}
} }
} }
@@ -1311,12 +1318,19 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
let args_rmp = json_to_rmpv(args); let args_rmp = json_to_rmpv(args);
let code = " let code = "
local args = ... local args = ...
local target_file = (args.file and args.file ~= '' and args.file ~= 'null') and args.file
or (args.file_path and args.file_path ~= '' and args.file_path ~= 'null') and args.file_path
or (args.path and args.path ~= '' and args.path ~= 'null') and args.path
or (args.target and args.target ~= '' and args.target ~= 'null') and args.target
if not target_file or target_file == vim.NIL then
return 'Error: No valid file path provided'
end
local curr_buf = vim.api.nvim_get_current_buf() local curr_buf = vim.api.nvim_get_current_buf()
local is_empty_unnamed = vim.api.nvim_buf_get_name(curr_buf) == '' local is_empty_unnamed = vim.api.nvim_buf_get_name(curr_buf) == ''
and vim.api.nvim_buf_get_option(curr_buf, 'modified') == false and vim.api.nvim_buf_get_option(curr_buf, 'modified') == false
and vim.api.nvim_buf_line_count(curr_buf) <= 1 and vim.api.nvim_buf_line_count(curr_buf) <= 1
and (vim.api.nvim_buf_get_lines(curr_buf, 0, 1, false)[1] or '') == '' and (vim.api.nvim_buf_get_lines(curr_buf, 0, 1, false)[1] or '') == ''
vim.cmd('edit ' .. vim.fn.fnameescape(args.file)) vim.cmd('edit ' .. vim.fn.fnameescape(target_file))
local new_buf = vim.api.nvim_get_current_buf() local new_buf = vim.api.nvim_get_current_buf()
if is_empty_unnamed and curr_buf ~= new_buf and vim.api.nvim_buf_is_valid(curr_buf) then if is_empty_unnamed and curr_buf ~= new_buf and vim.api.nvim_buf_is_valid(curr_buf) then
pcall(vim.api.nvim_buf_delete, curr_buf, { force = true }) pcall(vim.api.nvim_buf_delete, curr_buf, { force = true })
@@ -1324,7 +1338,7 @@ pub async fn run_mcp_loop(app_name: &str, app_version: &str) {
if args.filetype and args.filetype ~= '' then if args.filetype and args.filetype ~= '' then
vim.bo.filetype = args.filetype vim.bo.filetype = args.filetype
end end
return 'Opened file ' .. args.file return 'Opened file ' .. target_file
"; ";
handle_lua_result!(id, execute_nvim_lua_with_args(code, vec![args_rmp])) handle_lua_result!(id, execute_nvim_lua_with_args(code, vec![args_rmp]))
} }
@@ -1841,6 +1855,38 @@ mod tests {
let ext_val = rmpv::Value::Ext(1, vec![10, 20]); let ext_val = rmpv::Value::Ext(1, vec![10, 20]);
assert_eq!(rmpv_to_json(&ext_val), serde_json::json!("Ext(1, [10, 20])")); assert_eq!(rmpv_to_json(&ext_val), serde_json::json!("Ext(1, [10, 20])"));
} }
#[test]
fn test_mock_nvim_get_api_info_response() {
let channel_id = rmpv::Value::Integer(42.into());
let api_metadata = vec![
(
rmpv::Value::String("version".into()),
rmpv::Value::Map(vec![
(rmpv::Value::String("major".into()), rmpv::Value::Integer(0.into())),
(rmpv::Value::String("minor".into()), rmpv::Value::Integer(10.into())),
(rmpv::Value::String("patch".into()), rmpv::Value::Integer(0.into())),
]),
),
(
rmpv::Value::String("functions".into()),
rmpv::Value::Array(vec![]),
),
];
let api_info_res = rmpv::Value::Array(vec![channel_id, rmpv::Value::Map(api_metadata)]);
if let rmpv::Value::Array(arr) = &api_info_res {
assert_eq!(arr.len(), 2);
let chan = arr[0].as_i64().unwrap();
assert_eq!(chan, 42);
let json_res = rmpv_to_json(&api_info_res);
assert_eq!(json_res[0], json!(42));
assert_eq!(json_res[1]["version"]["minor"], json!(10));
} else {
panic!("Expected array response for nvim_get_api_info");
}
}
} }
-1
View File
@@ -31,7 +31,6 @@ fastembed = "7.1.0"
arboard = "3.6.1" arboard = "3.6.1"
image = "0.25.10" image = "0.25.10"
base64 = "0.23.1" base64 = "0.23.1"
clipboard-win = "5.4.1"
git2 = "0.19.0" git2 = "0.19.0"
tree-sitter = "0.23.2" tree-sitter = "0.23.2"
tree-sitter-rust = "0.23.3" tree-sitter-rust = "0.23.3"
+4 -3
View File
@@ -221,9 +221,9 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
move |axum::extract::Query(params): axum::extract::Query< move |axum::extract::Query(params): axum::extract::Query<
std::collections::HashMap<String, String>, std::collections::HashMap<String, String>,
>| async move { >| async move {
if let Some(q) = params.get("q") if let Some(q) = params.get("q") {
&& let Ok(idx) = state_clone.search_index.read() let idx = state_clone.get_search_index().await;
&& let Ok(results) = idx.search(q, None) { if let Ok(results) = idx.search(q, None) {
let mut formatted_results = Vec::new(); let mut formatted_results = Vec::new();
for (id, doc_type, title, body, score) in results { for (id, doc_type, title, body, score) in results {
formatted_results.push(serde_json::json!({ formatted_results.push(serde_json::json!({
@@ -238,6 +238,7 @@ pub fn create_router(app_state: Arc<AppState>) -> Router {
serde_json::json!({ "results": formatted_results }), serde_json::json!({ "results": formatted_results }),
); );
} }
}
axum::Json(serde_json::json!({ "results": [] })) axum::Json(serde_json::json!({ "results": [] }))
} }
}), }),
+22 -6
View File
@@ -6,6 +6,8 @@ use tokio::time::{Duration, sleep};
pub fn spawn_watcher(state: Arc<MemoryState>) { pub fn spawn_watcher(state: Arc<MemoryState>) {
tokio::spawn(async move { tokio::spawn(async move {
let mut last_text = String::new(); let mut last_text = String::new();
let mut clipboard_opt: Option<Clipboard> = Clipboard::new().ok();
loop { loop {
let is_enabled = { let is_enabled = {
let watch = state.clipboard_watch_mode.read().await; let watch = state.clipboard_watch_mode.read().await;
@@ -19,17 +21,20 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
sleep(Duration::from_millis(1000)).await; sleep(Duration::from_millis(1000)).await;
if let Ok(mut clipboard) = Clipboard::new() if clipboard_opt.is_none() {
&& let Ok(text) = clipboard.get_text() clipboard_opt = Clipboard::new().ok();
&& text != last_text }
&& !text.trim().is_empty()
{ if let Some(ref mut cb) = clipboard_opt {
match cb.get_text() {
Ok(text) => {
if text != last_text && !text.trim().is_empty() {
last_text = text.clone(); last_text = text.clone();
let note = StickyNote { let note = StickyNote {
timestamp: std::time::SystemTime::now() timestamp: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) .duration_since(std::time::UNIX_EPOCH)
.unwrap() .unwrap_or_default()
.as_secs(), .as_secs(),
content: text.clone(), content: text.clone(),
expires_at: None, expires_at: None,
@@ -37,6 +42,10 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
state.code.sticky.modify(|notes| { state.code.sticky.modify(|notes| {
notes.push(note.clone()); notes.push(note.clone());
if notes.len() > 200 {
let overflow = notes.len() - 200;
notes.drain(0..overflow);
}
}); });
state.record_activity( state.record_activity(
@@ -48,6 +57,13 @@ pub fn spawn_watcher(state: Arc<MemoryState>) {
state.index_commit_notify.notify_waiters(); state.index_commit_notify.notify_waiters();
} }
} }
Err(_) => {
// Reset clipboard handle on error so it can be re-initialized cleanly
clipboard_opt = None;
}
}
}
}
}); });
} }
+201
View File
@@ -55,6 +55,12 @@ function switchTab(tabId, btn) {
break; break;
} }
} }
function toggleTheme() {
const currentTheme = document.documentElement.getAttribute("data-theme");
const newTheme = currentTheme === "dark" ? "light" : "dark";
document.documentElement.setAttribute("data-theme", newTheme);
localStorage.setItem("theme", newTheme);
}
if (localStorage.getItem("theme") === "dark" || !localStorage.getItem("theme") && window.matchMedia && window.matchMedia("(prefers-color-scheme: dark)").matches) { if (localStorage.getItem("theme") === "dark" || !localStorage.getItem("theme") && window.matchMedia && window.matchMedia("(prefers-color-scheme: dark)").matches) {
document.documentElement.setAttribute("data-theme", "dark"); document.documentElement.setAttribute("data-theme", "dark");
} }
@@ -91,6 +97,11 @@ function zoomGraph(step) {
const currentScale = network.getScale(); const currentScale = network.getScale();
network.moveTo({ scale: currentScale * (1 + step) }); network.moveTo({ scale: currentScale * (1 + step) });
} }
function resetGraph() {
if (!network)
return;
network.fit({ animation: { duration: 500, easingFunction: "easeInOutQuad" } });
}
function closeInspector() { function closeInspector() {
const inspector = document.getElementById("inspector-panel"); const inspector = document.getElementById("inspector-panel");
if (inspector) if (inspector)
@@ -128,6 +139,14 @@ function showInspector(nodeId) {
if (inspector) if (inspector)
inspector.classList.add("open"); inspector.classList.add("open");
} }
function toggleFilter(type) {
if (activeFilters.has(type))
activeFilters.delete(type);
else
activeFilters.add(type);
renderFilters();
updateGraphData();
}
function renderFilters() { function renderFilters() {
const container = document.getElementById("graph-filters"); const container = document.getElementById("graph-filters");
if (!container) if (!container)
@@ -240,6 +259,14 @@ async function loadGraph() {
} }
} }
} }
async function completeTask(id) {
try {
await fetch(`/api/tasks/${id}/complete`, { method: "POST" });
loadTasks();
} catch (e) {
console.error("Failed to complete task", e);
}
}
function buildTaskTreeHTML(tasks, parentId, depth = 0) { function buildTaskTreeHTML(tasks, parentId, depth = 0) {
let html = ""; let html = "";
const children = tasks.filter((t) => { const children = tasks.filter((t) => {
@@ -340,6 +367,64 @@ async function loadTasks() {
console.error("Failed to load tasks", err); console.error("Failed to load tasks", err);
} }
} }
function highlightText(text, query) {
if (!query)
return text;
const regex = new RegExp(`(${query})`, "gi");
return text.replace(regex, '<span class="search-highlight">$1</span>');
}
var searchDebounceTimer = null;
var activeSearchAbortController = null;
async function handleSearch(e) {
const input = e.target;
const q = input.value.trim();
if (searchDebounceTimer)
clearTimeout(searchDebounceTimer);
if (!q) {
if (activeSearchAbortController)
activeSearchAbortController.abort();
const resultsEl = document.getElementById("search-results");
if (resultsEl)
resultsEl.innerHTML = "";
return;
}
const delay = e.key === "Enter" ? 0 : 250;
searchDebounceTimer = setTimeout(async () => {
if (activeSearchAbortController)
activeSearchAbortController.abort();
activeSearchAbortController = new AbortController;
try {
const res = await fetch(`/api/search?q=${encodeURIComponent(q)}`, {
signal: activeSearchAbortController.signal
});
const data = await res.json();
const container = document.getElementById("search-results");
if (!container)
return;
if (!data.results || data.results.length === 0) {
container.innerHTML = '<div style="padding:15px;color:var(--text-secondary);">No results found.</div>';
return;
}
container.innerHTML = data.results.map((r) => `
<div class="search-result-item" onclick="if(rawEntities['${r.id}']){ showInspector('${r.id}'); network.selectNodes(['${r.id}']); network.focus('${r.id}', {scale:1.2, animation:true}); }" style="cursor:pointer;">
<div style="display:flex; justify-content:space-between; align-items:flex-start;">
<div>
<strong>${highlightText(r.title, q)}</strong>
<span style="font-size: 0.7em; color: var(--text-secondary); background: var(--canvas-bg); padding: 2px 6px; border-radius: 4px; margin-left: 8px; text-transform: uppercase; border: 1px solid var(--border-color);">${r.type_name}</span>
</div>
<div style="font-size:0.75em; color:var(--text-secondary);">${r.score.toFixed(2)}</div>
</div>
<div style="margin-top:6px; font-size:0.9em; line-height:1.4;">${highlightText(r.content.substring(0, 150), q)}${r.content.length > 150 ? "..." : ""}</div>
<div style="font-size:0.75em; color:var(--text-secondary); margin-top:8px; font-family:monospace; background: var(--canvas-bg); padding: 2px 4px; border-radius: 3px; display: inline-block;">ID: ${r.id}</div>
</div>
`).join("");
} catch (err) {
if (err.name !== "AbortError") {
console.error("Search error:", err);
}
}
}, delay);
}
var MAX_ACTIVITY_HISTORY = 100; var MAX_ACTIVITY_HISTORY = 100;
function escapeHtml(str) { function escapeHtml(str) {
if (str === null || str === undefined) if (str === null || str === undefined)
@@ -785,6 +870,39 @@ function loadPRs() {
</ul> </ul>
`); `);
} }
function loadAllExtras() {
loadTerminal();
loadTechDebt();
loadADRs();
loadWorkspaces();
loadPinned();
loadMemos();
loadSnippets();
loadPRs();
}
async function testClipboard() {
const modal = document.getElementById("clipboard-modal");
const resultDiv = document.getElementById("clipboard-result");
if (modal)
modal.style.display = "flex";
if (resultDiv)
resultDiv.innerHTML = '<p id="clipboard-loading" style="color:var(--text-secondary);">Analyzing your clipboard...</p>';
try {
const res = await fetch("/api/clipboard/capture", { method: "POST" });
const data = await res.json();
if (resultDiv) {
if (data.success) {
resultDiv.innerHTML = `<img src="${data.image}" style="max-width:100%; max-height:60vh; border-radius:4px; border:1px solid var(--border-color); box-shadow: 0 5px 15px rgba(0,0,0,0.3);" />`;
} else {
resultDiv.innerHTML = `<p style="color:var(--error-color); font-weight:bold;">Failed to read clipboard: ${data.error}</p>`;
}
}
} catch (err) {
if (resultDiv) {
resultDiv.innerHTML = `<p style="color:var(--error-color); font-weight:bold;">Error calling endpoint: ${err.message}</p>`;
}
}
}
async function loadVersion() { async function loadVersion() {
try { try {
const res = await fetch("/api/version"); const res = await fetch("/api/version");
@@ -836,6 +954,89 @@ document.addEventListener("keydown", function(e) {
} }
} }
}); });
function refreshActiveTab() {
switch (currentTabId) {
case "graph-tab":
loadGraph();
break;
case "task-tab":
loadTasks();
break;
case "sticky-tab":
loadStickyNotes();
break;
case "techdebt-tab":
loadTechDebt();
break;
case "adrs-tab":
loadADRs();
break;
case "workspaces-tab":
loadWorkspaces();
break;
case "pinned-tab":
loadPinned();
break;
case "memos-tab":
loadMemos();
break;
case "snippets-tab":
loadSnippets();
break;
case "pr-tab":
loadPRs();
break;
case "terminal-tab":
loadTerminal();
break;
case "preferences-tab":
loadPreferences();
break;
case "activity-tab":
loadActivityHistory();
break;
}
}
Object.assign(window, {
switchTab,
toggleTheme,
getColorForType,
zoomGraph,
resetGraph,
closeInspector,
showInspector,
toggleFilter,
renderFilters,
updateGraphData,
loadGraph,
completeTask,
buildTaskTreeHTML,
loadTasks,
highlightText,
handleSearch,
escapeHtml,
parseActivityPayload,
loadActivityHistory,
requestDomainRefresh,
handleIncomingActivity,
setupWS,
loadPreferences,
loadStickyNotes,
loadGenericList,
loadTerminal,
loadTechDebt,
loadADRs,
loadWorkspaces,
loadPinned,
loadMemos,
loadSnippets,
loadPRs,
loadAllExtras,
testClipboard,
loadVersion,
setupSSE,
refreshActiveTab
});
loadVersion(); loadVersion();
loadActivityHistory(); loadActivityHistory();
setupWS(); setupWS();
+43
View File
@@ -1096,6 +1096,48 @@ function refreshActiveTab(): void {
} }
} }
// --- Global Window Exports (Prevents bun build tree-shaking for HTML inline event handlers) ---
Object.assign(window, {
switchTab,
toggleTheme,
getColorForType,
zoomGraph,
resetGraph,
closeInspector,
showInspector,
toggleFilter,
renderFilters,
updateGraphData,
loadGraph,
completeTask,
buildTaskTreeHTML,
loadTasks,
highlightText,
handleSearch,
escapeHtml,
parseActivityPayload,
loadActivityHistory,
requestDomainRefresh,
handleIncomingActivity,
setupWS,
loadPreferences,
loadStickyNotes,
loadGenericList,
loadTerminal,
loadTechDebt,
loadADRs,
loadWorkspaces,
loadPinned,
loadMemos,
loadSnippets,
loadPRs,
loadAllExtras,
testClipboard,
loadVersion,
setupSSE,
refreshActiveTab
});
// --- Initialization --- // --- Initialization ---
loadVersion(); loadVersion();
loadActivityHistory(); loadActivityHistory();
@@ -1103,3 +1145,4 @@ setupWS();
const observer = new MutationObserver(() => updateGraphData()); const observer = new MutationObserver(() => updateGraphData());
observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] }); observer.observe(document.documentElement, { attributes: true, attributeFilter: ['data-theme'] });
+13 -1
View File
@@ -4,8 +4,16 @@ use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
pub fn init_redb(base: &Path) -> Arc<Database> { pub fn init_redb(base: &Path) -> Arc<Database> {
let is_in_memory = base == Path::new(":memory:") || base.to_str() == Some(":memory:") || cfg!(test);
let db = if is_in_memory {
Arc::new(
redb::Database::builder()
.create_with_backend(redb::backends::InMemoryBackend::new())
.expect("Failed to create in-memory redb database"),
)
} else {
let redb_path = base.join("mcp_store.redb"); let redb_path = base.join("mcp_store.redb");
let db = match redb::Database::create(&redb_path) { match redb::Database::create(&redb_path) {
Ok(db) => Arc::new(db), Ok(db) => Arc::new(db),
Err(e) => { Err(e) => {
let err_msg = format!( let err_msg = format!(
@@ -14,8 +22,10 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
); );
tracing::error!("{}", err_msg); tracing::error!("{}", err_msg);
let temp_path = std::env::temp_dir().join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4())); let temp_path = std::env::temp_dir().join(format!("mcp_store_fallback_{}.redb", uuid::Uuid::new_v4()));
tracing::warn!("CRITICAL PERSISTENCE ALERT: Using temporary redb database {:?}. Changes will be discarded upon application exit.", temp_path);
Arc::new(redb::Database::create(&temp_path).expect("Failed to create fallback redb database")) Arc::new(redb::Database::create(&temp_path).expect("Failed to create fallback redb database"))
} }
}
}; };
// Ensure the table exists and migrate legacy JSON files // Ensure the table exists and migrate legacy JSON files
@@ -24,6 +34,7 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
let mut opened_ok = false; let mut opened_ok = false;
if let Ok(mut table) = write_txn.open_table(STORE_TABLE) { if let Ok(mut table) = write_txn.open_table(STORE_TABLE) {
opened_ok = true; opened_ok = true;
if !is_in_memory {
let stores = vec![ let stores = vec![
("knowledge_graph_master", "knowledge_graph_master.json"), ("knowledge_graph_master", "knowledge_graph_master.json"),
("audit_ledger", "audit_ledger.json"), ("audit_ledger", "audit_ledger.json"),
@@ -70,6 +81,7 @@ pub fn init_redb(base: &Path) -> Arc<Database> {
} }
} }
} }
}
if opened_ok { if opened_ok {
if let Err(e) = write_txn.commit() { if let Err(e) = write_txn.commit() {
tracing::error!("Failed to commit database migration transaction: {}", e); tracing::error!("Failed to commit database migration transaction: {}", e);
+17 -5
View File
@@ -33,9 +33,16 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 {
if a.is_empty() || b.is_empty() || a.len() != b.len() { if a.is_empty() || b.is_empty() || a.len() != b.len() {
return 0.0; return 0.0;
} }
let dot_product: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); let mut dot_product = 0.0f32;
let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt(); let mut norm_a_sq = 0.0f32;
let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt(); let mut norm_b_sq = 0.0f32;
for (&x, &y) in a.iter().zip(b.iter()) {
dot_product += x * y;
norm_a_sq += x * x;
norm_b_sq += y * y;
}
let norm_a = norm_a_sq.sqrt();
let norm_b = norm_b_sq.sqrt();
if norm_a == 0.0 || norm_b == 0.0 { if norm_a == 0.0 || norm_b == 0.0 {
0.0 0.0
} else { } else {
@@ -49,8 +56,13 @@ pub async fn generate_embeddings_async(texts: Vec<String>) -> Result<Vec<Vec<f32
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let model_mutex = get_embedding_model()?; let model_mutex = get_embedding_model()?;
let mut model = model_mutex.lock().map_err(|e| e.to_string())?; let mut model = model_mutex.lock().map_err(|e| e.to_string())?;
let embeddings = model.embed(texts, None).map_err(|e| e.to_string())?; let mut all_embeddings = Vec::with_capacity(texts.len());
Ok(embeddings) for chunk in texts.chunks(32) {
let chunk_vec = chunk.to_vec();
let chunk_embeddings = model.embed(chunk_vec, None).map_err(|e| e.to_string())?;
all_embeddings.extend(chunk_embeddings);
}
Ok(all_embeddings)
}) })
.await .await
.map_err(|e| e.to_string())? .map_err(|e| e.to_string())?
+21
View File
@@ -13,6 +13,15 @@ pub enum AppError {
#[error("Not Found: {0}")] #[error("Not Found: {0}")]
NotFound(String), NotFound(String),
#[error("Entity '{name}' was not found in namespace '{namespace}'")]
EntityNotFound { name: String, namespace: String },
#[error("Invalid relation between '{from}' and '{to}': {reason}")]
InvalidRelation { from: String, to: String, reason: String },
#[error("Database error for key '{key}': {message}")]
DatabaseError { key: String, message: String },
#[error("Forbidden: {0}")] #[error("Forbidden: {0}")]
Forbidden(String), Forbidden(String),
@@ -39,6 +48,18 @@ impl IntoResponse for AppError {
fn into_response(self) -> Response { fn into_response(self) -> Response {
let (status, error_message) = match &self { let (status, error_message) = match &self {
AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()), AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg.clone()),
AppError::EntityNotFound { name, namespace } => (
StatusCode::NOT_FOUND,
format!("Entity '{}' was not found in namespace '{}'", name, namespace),
),
AppError::InvalidRelation { from, to, reason } => (
StatusCode::BAD_REQUEST,
format!("Invalid relation between '{}' and '{}': {}", from, to, reason),
),
AppError::DatabaseError { key, message } => (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Database error for key '{}': {}", key, message),
),
AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()), AppError::Forbidden(msg) => (StatusCode::FORBIDDEN, msg.clone()),
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()), AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()), AppError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg.clone()),
+140 -15
View File
@@ -56,7 +56,7 @@ impl McpTool for ReadFileSkeletonHandler {
let tree = parser.parse(&code, None).ok_or_else(|| { let tree = parser.parse(&code, None).ok_or_else(|| {
crate::error::AppError::Internal("Failed to parse code".to_string()) crate::error::AppError::Internal("Failed to parse code".to_string())
})?; })?;
let mut result_skeleton = String::new(); let mut result_skeleton = String::with_capacity(code.len() / 2);
fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) { fn extract_skeleton(node: Node, code: &str, out: &mut String, depth: usize) {
if depth > 128 { if depth > 128 {
@@ -141,7 +141,7 @@ impl McpTool for ReplaceAstNodeHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ReplaceAstNodeTool>( crate::mcp::tool_def::<ReplaceAstNodeTool>(
"replace_ast_node", "replace_ast_node",
"Replace a specific AST node (e.g., function, struct) entirely using tree-sitter for robust structural editing.", "Replace a specific AST node (e.g., function, struct, enum, class, trait) entirely using tree-sitter for robust structural editing. Supported node_type values include: 'function_item' (or 'function'/'fn'/'method'), 'struct_item' (or 'struct'), 'class_declaration' (or 'class'), 'enum_item' (or 'enum'), 'trait_item' (or 'trait'/'interface'), 'type_alias_declaration' (or 'type').",
) )
} }
@@ -185,6 +185,34 @@ impl McpTool for ReplaceAstNodeHandler {
crate::error::AppError::Internal("Failed to parse code".to_string()) crate::error::AppError::Internal("Failed to parse code".to_string())
})?; })?;
fn matches_node_type(actual_kind: &str, requested_type: &str) -> bool {
if actual_kind == requested_type {
return true;
}
match requested_type.to_lowercase().as_str() {
"function" | "func" | "fn" | "method" | "def" => matches!(
actual_kind,
"function_item"
| "function_declaration"
| "function_definition"
| "method_definition"
| "function"
),
"struct" => matches!(actual_kind, "struct_item" | "struct_declaration" | "struct_specifier"),
"class" => matches!(actual_kind, "class_declaration" | "class_definition" | "class_item"),
"enum" => matches!(actual_kind, "enum_item" | "enum_declaration"),
"trait" | "interface" => matches!(
actual_kind,
"trait_item" | "interface_declaration" | "interface_item"
),
"type" | "type_alias" => matches!(
actual_kind,
"type_alias_declaration" | "type_item" | "type_definition"
),
_ => false,
}
}
// Search for the node // Search for the node
fn find_node<'a>( fn find_node<'a>(
node: Node<'a>, node: Node<'a>,
@@ -192,12 +220,12 @@ impl McpTool for ReplaceAstNodeHandler {
target_type: &str, target_type: &str,
target_name: &str, target_name: &str,
) -> Option<Node<'a>> { ) -> Option<Node<'a>> {
if node.kind() == target_type { if matches_node_type(node.kind(), target_type) {
// Try to find the name/identifier // Try to find the name/identifier
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.children(&mut cursor) { for child in node.children(&mut cursor) {
let kind = child.kind(); let kind = child.kind();
if kind == "identifier" || kind == "name" { if kind == "identifier" || kind == "name" || kind == "property_identifier" || kind == "field_identifier" {
let name = child.utf8_text(code.as_bytes()).unwrap_or(""); let name = child.utf8_text(code.as_bytes()).unwrap_or("");
if name == target_name { if name == target_name {
return Some(node); return Some(node);
@@ -250,6 +278,61 @@ impl McpTool for ReplaceAstNodeHandler {
} }
} }
fn scan_workspace_for_symbol(target_sym: &str, limit: usize, filter_fn_call: bool) -> Vec<serde_json::Value> {
let mut results = Vec::new();
let cwd = match std::env::current_dir() {
Ok(dir) => dir,
Err(_) => return results,
};
let walker = ignore::WalkBuilder::new(&cwd)
.hidden(true)
.git_ignore(true)
.build();
let mut scanned_files = 0;
for result in walker {
let entry = match result {
Ok(e) => e,
Err(_) => continue,
};
if entry.file_type().is_some_and(|ft| ft.is_file()) {
let path = entry.path();
let ext = path.extension().and_then(|s| s.to_str()).unwrap_or("");
if matches!(ext, "rs" | "ts" | "tsx" | "js" | "jsx" | "py" | "go" | "java" | "c" | "cpp" | "h" | "hpp") {
scanned_files += 1;
if scanned_files > 500 {
break;
}
if let Ok(content) = std::fs::read_to_string(path) {
for (line_num, line) in content.lines().enumerate() {
let is_match = if filter_fn_call {
line.contains(&format!("{}(", target_sym))
|| line.contains(&format!("{}.await", target_sym))
} else {
line.contains(target_sym)
};
if is_match {
results.push(serde_json::json!({
"file_path": path.to_string_lossy(),
"line": line_num + 1,
"content": line.trim(),
}));
if results.len() >= limit {
return results;
}
}
}
}
}
}
}
results
}
pub struct FindSymbolReferencesHandler; pub struct FindSymbolReferencesHandler;
#[async_trait] #[async_trait]
@@ -269,22 +352,36 @@ impl McpTool for FindSymbolReferencesHandler {
let req: crate::tools::FindSymbolReferencesTool = let req: crate::tools::FindSymbolReferencesTool =
serde_json::from_value(args).map_err(|e| e.to_string())?; serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10); let limit = req.limit.unwrap_or(10);
let target_sym = req.symbol; let target_sym = req.symbol.clone();
let matches = state.code.snippets.read_with(|snippets| { let mut matches = state.code.snippets.read_with(|snippets| {
let mut refs = Vec::new(); let mut refs = Vec::new();
for snippet in snippets { for snippet in snippets {
if snippet.code.contains(&target_sym) || snippet.name.contains(&target_sym) { if snippet.code.contains(&target_sym) || snippet.name.contains(&target_sym) {
refs.push(snippet.clone()); refs.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
"code": snippet.code,
}));
if refs.len() >= limit { if refs.len() >= limit {
break; break;
} }
} }
} }
Ok::<String, crate::error::AppError>(serde_json::to_string(&refs)?) Ok::<Vec<serde_json::Value>, crate::error::AppError>(refs)
})?; })?;
Ok(matches) if matches.len() < limit {
let remaining = limit - matches.len();
let disk_matches = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_sym, remaining, false)
})
.await
.unwrap_or_default();
matches.extend(disk_matches);
}
Ok(serde_json::to_string_pretty(&matches)?)
} }
} }
@@ -307,24 +404,38 @@ impl McpTool for GetCallersHandler {
let req: crate::tools::GetCallersTool = let req: crate::tools::GetCallersTool =
serde_json::from_value(args).map_err(|e| e.to_string())?; serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(10); let limit = req.limit.unwrap_or(10);
let target_fn = req.function_name; let target_fn = req.function_name.clone();
let callers = state.code.snippets.read_with(|snippets| { let mut callers = state.code.snippets.read_with(|snippets| {
let mut matching = Vec::new(); let mut matching = Vec::new();
for snippet in snippets { for snippet in snippets {
if snippet.code.contains(&format!("{}(", target_fn)) if snippet.code.contains(&format!("{}(", target_fn))
|| snippet.code.contains(&format!("{}.await", target_fn)) || snippet.code.contains(&format!("{}.await", target_fn))
{ {
matching.push(snippet.clone()); matching.push(serde_json::json!({
"source": "snippet",
"name": snippet.name,
"code": snippet.code,
}));
if matching.len() >= limit { if matching.len() >= limit {
break; break;
} }
} }
} }
Ok::<String, crate::error::AppError>(serde_json::to_string(&matching)?) Ok::<Vec<serde_json::Value>, crate::error::AppError>(matching)
})?; })?;
Ok(callers) if callers.len() < limit {
let remaining = limit - callers.len();
let disk_callers = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&target_fn, remaining, true)
})
.await
.unwrap_or_default();
callers.extend(disk_callers);
}
Ok(serde_json::to_string_pretty(&callers)?)
} }
} }
@@ -357,6 +468,20 @@ impl McpTool for AnalyzeImpactHandler {
} }
}); });
let sym_clone = sym.clone();
let disk_refs = tokio::task::spawn_blocking(move || {
scan_workspace_for_symbol(&sym_clone, 20, false)
})
.await
.unwrap_or_default();
for r in &disk_refs {
if let Some(path) = r.get("file_path").and_then(|p| p.as_str()) {
let line = r.get("line").and_then(|l| l.as_u64()).unwrap_or(0);
callers.push(format!("{}:{}", path, line));
}
}
let mut kg_connected = Vec::new(); let mut kg_connected = Vec::new();
state.read_graph(|g| { state.read_graph(|g| {
for rel in &g.relations { for rel in &g.relations {
@@ -384,7 +509,7 @@ impl McpTool for AnalyzeImpactHandler {
"target_symbol": sym, "target_symbol": sym,
"risk_level": risk_level, "risk_level": risk_level,
"ast_callers_count": caller_count, "ast_callers_count": caller_count,
"ast_callers_sample": callers.into_iter().take(5).collect::<Vec<_>>(), "ast_callers_sample": callers.into_iter().take(10).collect::<Vec<_>>(),
"graph_relations_count": graph_count, "graph_relations_count": graph_count,
"graph_relations": kg_connected, "graph_relations": kg_connected,
"recommendation": match risk_level { "recommendation": match risk_level {
+5 -5
View File
@@ -16,7 +16,7 @@ impl McpTool for UpdateEnvFingerprintHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<UpdateEnvFingerprintTool>( crate::mcp::tool_def::<UpdateEnvFingerprintTool>(
"update_env_fingerprint", "update_env_fingerprint",
"Execute update_env_fingerprint", "Update system OS, shell, and installed tool version fingerprints for a workspace namespace.",
) )
} }
@@ -50,7 +50,7 @@ impl McpTool for ReadEnvFingerprintHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadEnvFingerprintTool>( crate::mcp::tool_def::<ReadEnvFingerprintTool>(
"read_env_fingerprint", "read_env_fingerprint",
"Execute read_env_fingerprint", "Read system OS, shell, and tool version fingerprints for a workspace namespace.",
) )
} }
@@ -81,7 +81,7 @@ impl McpTool for LogEnvRequirementHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogEnvRequirementTool>( crate::mcp::tool_def::<LogEnvRequirementTool>(
"log_env_requirement", "log_env_requirement",
"Execute log_env_requirement", "Log required environment variables, configuration keys, or secrets for a namespace.",
) )
} }
@@ -111,7 +111,7 @@ impl McpTool for RegisterEnvironmentHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<RegisterEnvironmentTool>( crate::mcp::tool_def::<RegisterEnvironmentTool>(
"register_environment", "register_environment",
"Execute register_environment", "Register infrastructure environment details (URL, VPN requirements, notes) for a namespace.",
) )
} }
@@ -144,7 +144,7 @@ impl McpTool for GetEnvironmentDetailsHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<GetEnvironmentDetailsTool>( crate::mcp::tool_def::<GetEnvironmentDetailsTool>(
"get_environment_details", "get_environment_details",
"Execute get_environment_details", "Retrieve registered infrastructure environment details for a namespace.",
) )
} }
+2 -1
View File
@@ -69,7 +69,8 @@ impl McpTool for GetActiveWorktreeContextHandler {
// Truncate diff output if it's too large to save tokens // Truncate diff output if it's too large to save tokens
if diff_output.len() > 10000 { if diff_output.len() > 10000 {
diff_output.truncate(10000); let valid_boundary = diff_output.floor_char_boundary(10000);
diff_output.truncate(valid_boundary);
diff_output.push_str("\n... [Diff truncated due to size]"); diff_output.push_str("\n... [Diff truncated due to size]");
} }
+164 -73
View File
@@ -12,27 +12,40 @@ struct BorrowedGraph<'a> {
relations: Vec<&'a crate::models::Relation>, relations: Vec<&'a crate::models::Relation>,
} }
pub struct QueryGraphPathHandler; pub struct GraphQueryBuilder<'a> {
graph: &'a crate::models::KnowledgeGraph,
max_depth: usize,
relation_filter: Option<&'a str>,
}
#[async_trait] impl<'a> GraphQueryBuilder<'a> {
impl McpTool for QueryGraphPathHandler { pub fn new(graph: &'a crate::models::KnowledgeGraph) -> Self {
fn name(&self) -> &'static str { Self {
"query_graph_path" graph,
max_depth: 5,
relation_filter: None,
}
} }
fn schema(&self) -> Value { pub fn max_depth(mut self, depth: usize) -> Self {
crate::mcp::tool_def::<QueryGraphPathTool>("query_graph_path", "Execute query_graph_path") self.max_depth = depth;
self
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { pub fn relation_filter(mut self, filter: Option<&'a str>) -> Self {
let req: crate::tools::QueryGraphPathTool = self.relation_filter = filter;
serde_json::from_value(args).map_err(|e| e.to_string())?; self
state.read_graph(|graph| { }
let max_depth = req.max_depth.unwrap_or(5);
// Pre-index relations into an adjacency map for O(1) neighbor lookups pub fn find_shortest_path(&self, start: &str, end: &str) -> Option<Vec<String>> {
let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> = std::collections::HashMap::new(); let mut adj: std::collections::HashMap<&str, Vec<(&str, &str, bool)>> =
for rel in &graph.relations { std::collections::HashMap::with_capacity(self.graph.relations.len() * 2);
for rel in &self.graph.relations {
if let Some(rf) = self.relation_filter {
if rel.relation_type != rf {
continue;
}
}
adj.entry(rel.from.as_str()) adj.entry(rel.from.as_str())
.or_default() .or_default()
.push((rel.to.as_str(), rel.relation_type.as_str(), false)); .push((rel.to.as_str(), rel.relation_type.as_str(), false));
@@ -41,13 +54,12 @@ impl McpTool for QueryGraphPathHandler {
.push((rel.from.as_str(), rel.relation_type.as_str(), true)); .push((rel.from.as_str(), rel.relation_type.as_str(), true));
} }
let mut queue: std::collections::VecDeque<&str> = std::collections::VecDeque::new(); let mut queue = std::collections::VecDeque::new();
let mut visited: std::collections::HashSet<&str> = std::collections::HashSet::new(); let mut visited = std::collections::HashSet::new();
let mut parents: std::collections::HashMap<&str, (&str, &str, bool)> = let mut parents = std::collections::HashMap::new();
std::collections::HashMap::new();
queue.push_back(req.start_node.as_str()); queue.push_back(start);
visited.insert(req.start_node.as_str()); visited.insert(start);
let mut found = false; let mut found = false;
let mut current_depth = 0; let mut current_depth = 0;
@@ -55,12 +67,16 @@ impl McpTool for QueryGraphPathHandler {
let mut nodes_at_next_depth = 0; let mut nodes_at_next_depth = 0;
while let Some(current) = queue.pop_front() { while let Some(current) = queue.pop_front() {
if current == req.end_node { if current == end {
found = true; found = true;
break; break;
} }
// Visited node upper-bound cap to guarantee deterministic BFS bounds on dense graphs
if visited.len() > 10_000 {
break;
}
nodes_at_current_depth -= 1; nodes_at_current_depth -= 1;
if current_depth < max_depth { if current_depth < self.max_depth {
if let Some(neighbors) = adj.get(current) { if let Some(neighbors) = adj.get(current) {
for &(neighbor, rel_type, is_inverse) in neighbors { for &(neighbor, rel_type, is_inverse) in neighbors {
if !visited.contains(neighbor) { if !visited.contains(neighbor) {
@@ -81,8 +97,8 @@ impl McpTool for QueryGraphPathHandler {
if found { if found {
let mut path = Vec::new(); let mut path = Vec::new();
let mut curr = req.end_node.as_str(); let mut curr = end;
while curr != req.start_node { while curr != start {
if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) { if let Some((parent, rel_type, is_inverse)) = parents.get(&curr) {
if *is_inverse { if *is_inverse {
path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr)); path.push(format!("{} -[inverse({})]-> {}", parent, rel_type, curr));
@@ -95,6 +111,37 @@ impl McpTool for QueryGraphPathHandler {
} }
} }
path.reverse(); path.reverse();
Some(path)
} else {
None
}
}
}
pub struct QueryGraphPathHandler;
#[async_trait]
impl McpTool for QueryGraphPathHandler {
fn name(&self) -> &'static str {
"query_graph_path"
}
fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryGraphPathTool>(
"query_graph_path",
"Find the shortest relationship path between two entities in the knowledge graph within a maximum depth.",
)
}
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: crate::tools::QueryGraphPathTool =
serde_json::from_value(args).map_err(|e| e.to_string())?;
tokio::task::spawn_blocking(move || {
state.read_graph(|graph| {
let max_depth = req.max_depth.unwrap_or(5);
let builder = GraphQueryBuilder::new(graph).max_depth(max_depth as usize);
if let Some(path) = builder.find_shortest_path(&req.start_node, &req.end_node) {
Ok(format!("Path found:\n{}", path.join("\n"))) Ok(format!("Path found:\n{}", path.join("\n")))
} else { } else {
Ok(format!( Ok(format!(
@@ -103,6 +150,9 @@ impl McpTool for QueryGraphPathHandler {
)) ))
} }
}) })
})
.await
.map_err(|e| crate::error::AppError::Internal(format!("Graph traversal task failed: {}", e)))?
} }
} }
@@ -115,7 +165,10 @@ impl McpTool for CreateEntitiesHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateEntitiesTool>("create_entities", "Execute create_entities") crate::mcp::tool_def::<CreateEntitiesTool>(
"create_entities",
"Create new entities in the knowledge graph with normalized entity types.",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -130,11 +183,12 @@ impl McpTool for CreateEntitiesHandler {
} }
} }
}); });
let idx = state.get_search_index(); let names: Vec<String> = inserted.iter().map(|e| format!("{} ({})", e.name, e.entity_type)).collect();
for entity in inserted { if !inserted.is_empty() {
drop(idx.index_entity(&entity)); let idx = state.get_search_index().await;
let _ = idx.index_entities_batch(&inserted).await;
} }
Ok("Entities created".to_string()) Ok(format!("Successfully created {} entity/entities: {}", names.len(), names.join(", ")))
} }
} }
@@ -147,7 +201,10 @@ impl McpTool for CreateRelationsHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<CreateRelationsTool>("create_relations", "Execute create_relations") crate::mcp::tool_def::<CreateRelationsTool>(
"create_relations",
"Create directed relationships between entities in the knowledge graph. Requires 'from', 'to', and 'relation_type'.",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -183,23 +240,38 @@ impl McpTool for CreateRelationsHandler {
} }
}); });
if !missing_nodes.is_empty() { let mut auto_created = Vec::new();
let missing: Vec<_> = missing_nodes.into_iter().collect(); let mut added_relations = Vec::new();
return Err(crate::error::AppError::Internal(format!(
"Error: Relations dropped due to missing entities: {}",
missing.join(", ")
)));
}
state.modify_graph(|g| { state.modify_graph(|g| {
for node_name in missing_nodes {
if !g.entities.contains_key(&node_name) {
g.entities.insert(
node_name.clone(),
crate::models::Entity {
name: node_name.clone(),
entity_type: "Entity".to_string(),
observations: vec!["Auto-created stub entity for relation endpoint".to_string()],
namespace: crate::models::default_namespace(),
git_branch: None,
},
);
auto_created.push(node_name);
}
}
for mut relation in req.relations { for mut relation in req.relations {
if !relation.from.is_empty() && !relation.to.is_empty() { if !relation.from.is_empty() && !relation.to.is_empty() {
relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type); relation.relation_type = crate::models::normalize_relation_type(&relation.relation_type);
added_relations.push(format!("{} -[{}]-> {}", relation.from, relation.relation_type, relation.to));
g.relations.push(relation); g.relations.push(relation);
} }
} }
}); });
Ok("Relations created".to_string())
let mut msg = format!("Successfully created {} relation(s):\n{}", added_relations.len(), added_relations.join("\n"));
if !auto_created.is_empty() {
msg.push_str(&format!("\nNote: Auto-created {} missing stub entity/entities: {}", auto_created.len(), auto_created.join(", ")));
}
Ok(msg)
} }
} }
@@ -212,7 +284,10 @@ impl McpTool for AddObservationsHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<AddObservationsTool>("add_observations", "Execute add_observations") crate::mcp::tool_def::<AddObservationsTool>(
"add_observations",
"Add new observations and factual statements to existing entities in the knowledge graph.",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -284,7 +359,7 @@ impl McpTool for DeleteEntitiesHandler {
.retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to)); .retain(|r| !to_delete.contains(&r.from) && !to_delete.contains(&r.to));
}); });
let idx = state.get_search_index(); let idx = state.get_search_index().await;
for name in to_delete { for name in to_delete {
drop(idx.delete_document(&name)); drop(idx.delete_document(&name));
} }
@@ -361,7 +436,7 @@ impl McpTool for DeleteRelationsHandler {
}); });
if missing_count > 0 { if missing_count > 0 {
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Error: {} relation(s) not found in graph. Please verify exact relation properties using read_graph.", "Error: {} relation(s) not found in graph. Please verify exact relation properties (from, to, relation_type) using read_graph or get_subgraph.",
missing_count missing_count
))); )));
} }
@@ -378,7 +453,10 @@ impl McpTool for ReadGraphHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadGraphTool>("read_graph", "Execute read_graph") crate::mcp::tool_def::<ReadGraphTool>(
"read_graph",
"Read entities and relations from the knowledge graph with optional namespace filtering and token truncation. For large graphs, specify 'namespace' or use 'search_nodes' or 'get_subgraph' for targeted discovery.",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -405,8 +483,9 @@ impl McpTool for ReadGraphHandler {
if let Some(max_tok) = max_tokens { if let Some(max_tok) = max_tokens {
let max_chars = max_tok * 4; let max_chars = max_tok * 4;
if result_json.len() > max_chars { if result_json.len() > max_chars {
result_json.truncate(max_chars); let valid_boundary = result_json.floor_char_boundary(max_chars);
result_json.push_str("... [TRUNCATED_TO_MAX_TOKENS]"); result_json.truncate(valid_boundary);
result_json.push_str("\n... [TRUNCATED_TO_MAX_TOKENS. Use search_nodes or get_subgraph for targeted discovery]");
} }
} }
Ok(result_json) Ok(result_json)
@@ -432,12 +511,10 @@ impl McpTool for SearchNodesHandler {
let limit = req.limit.unwrap_or(10); let limit = req.limit.unwrap_or(10);
let include_body = req.include_body.unwrap_or(false); let include_body = req.include_body.unwrap_or(false);
let matches = if let Ok(idx) = state.search_index.read() { let idx = state.get_search_index().await;
idx.search(&req.query, req.namespace.as_deref()) let matches = idx
.unwrap_or_default() .search(&req.query, req.namespace.as_deref())
} else { .unwrap_or_default();
vec![]
};
let data = state.read_graph(|full| -> crate::error::Result<String> { let data = state.read_graph(|full| -> crate::error::Result<String> {
let mut matched_entities = Vec::new(); let mut matched_entities = Vec::new();
@@ -709,6 +786,11 @@ impl McpTool for MergeEntitiesHandler {
r.to = req.target_entity.clone(); r.to = req.target_entity.clone();
} }
// Filter out self-loops
if r.from == r.to {
return false;
}
if r.from == req.target_entity || r.to == req.target_entity { if r.from == req.target_entity || r.to == req.target_entity {
seen.insert(r.clone()) seen.insert(r.clone())
} else { } else {
@@ -908,16 +990,17 @@ impl McpTool for SweepGraphHealthHandler {
} }
} }
// 2. Compute similarity pairs for duplicate detection // 2. Compute similarity pairs for duplicate detection using pre-computed lowercase names
let names: Vec<_> = g.entities.keys().cloned().collect(); let names: Vec<_> = g.entities.keys().cloned().collect();
let lower_names: Vec<String> = names.iter().map(|n| n.to_lowercase()).collect();
for i in 0..names.len() { for i in 0..names.len() {
for j in (i + 1)..names.len() { for j in (i + 1)..names.len() {
let n1 = &names[i]; let n1 = &names[i];
let n2 = &names[j]; let n2 = &names[j];
let l1 = &lower_names[i];
let l2 = &lower_names[j];
let l1 = n1.to_lowercase(); if l1 == l2 || ((l1.contains(l2.as_str()) || l2.contains(l1.as_str())) && l1.len().min(l2.len()) > 3) {
let l2 = n2.to_lowercase();
if l1 == l2 || ((l1.contains(&l2) || l2.contains(&l1)) && l1.len().min(l2.len()) > 3) {
duplicates.push(serde_json::json!({ duplicates.push(serde_json::json!({
"entity_a": n1, "entity_a": n1,
"entity_b": n2, "entity_b": n2,
@@ -1050,7 +1133,8 @@ impl McpTool for SummarizeSubgraphHandler {
let max_tokens = req.max_tokens.unwrap_or(1000); let max_tokens = req.max_tokens.unwrap_or(1000);
let max_chars = max_tokens * 4; let max_chars = max_tokens * 4;
if markdown.len() > max_chars { if markdown.len() > max_chars {
markdown.truncate(max_chars); let valid_boundary = markdown.floor_char_boundary(max_chars);
markdown.truncate(valid_boundary);
markdown.push_str("\n... [Truncated to fit token budget]"); markdown.push_str("\n... [Truncated to fit token budget]");
} }
@@ -1063,13 +1147,11 @@ mod tests {
use super::*; use super::*;
use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler}; use crate::handlers::meta::{BroadcastAgentSignalHandler, QueryAgentSignalsHandler};
use serde_json::json; use serde_json::json;
use tempfile::tempdir;
#[tokio::test] #[tokio::test]
async fn test_create_and_read_entities() { async fn test_create_and_read_entities() {
let dir = tempdir().unwrap(); let state = Arc::new(MemoryState::new_in_memory());
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let create_handler = CreateEntitiesHandler; let create_handler = CreateEntitiesHandler;
let args = json!({ let args = json!({
@@ -1083,7 +1165,7 @@ mod tests {
.await .await
.map_err(|e| crate::error::AppError::Internal(e.to_string())) .map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap(); .unwrap();
assert_eq!(res, "Entities created"); assert!(res.contains("Successfully created 1 entity/entities"));
// Ensure graph contains the entity // Ensure graph contains the entity
state.graph.read_with(|g| { state.graph.read_with(|g| {
@@ -1094,8 +1176,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_create_relations() { async fn test_create_relations() {
let dir = tempdir().unwrap(); let state = Arc::new(MemoryState::new_in_memory());
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
// Needs entities first // Needs entities first
state.graph.modify(|g| { state.graph.modify(|g| {
@@ -1132,14 +1213,26 @@ mod tests {
.await .await
.map_err(|e| crate::error::AppError::Internal(e.to_string())) .map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap(); .unwrap();
assert_eq!(res, "Relations created"); assert!(res.contains("Successfully created 1 relation(s)"));
// Test semantic LLM schema feedback (User request) // Test serde field aliases (source/target/relationType mapped to from/to/relation_type)
let bad_args = json!({ let alias_args = json!({
"relations": [ "relations": [
{"source": "A", "target": "B", "relationType": "knows"} {"source": "A", "target": "B", "relationType": "knows"}
] ]
}); });
let alias_res = handler
.execute(alias_args, state.clone())
.await
.unwrap();
assert!(alias_res.contains("Successfully created 1 relation(s)"));
// Test semantic LLM schema feedback on missing fields
let bad_args = json!({
"relations": [
{"invalid_field": "X"}
]
});
let err_res = handler let err_res = handler
.execute(bad_args, state.clone()) .execute(bad_args, state.clone())
.await .await
@@ -1151,8 +1244,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_observations_and_reads() { async fn test_observations_and_reads() {
let dir = tempdir().unwrap(); let state = Arc::new(MemoryState::new_in_memory());
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
// Need entity first // Need entity first
state.graph.modify(|g| { state.graph.modify(|g| {
@@ -1208,8 +1300,7 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_advanced_graph_operations() { async fn test_advanced_graph_operations() {
let dir = tempfile::tempdir().unwrap(); let state = Arc::new(MemoryState::new_in_memory());
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let create_handler = CreateEntitiesHandler; let create_handler = CreateEntitiesHandler;
let args_ent = json!({ let args_ent = json!({
@@ -1289,13 +1380,13 @@ mod tests {
.await .await
.map_err(|e| crate::error::AppError::Internal(e.to_string())) .map_err(|e| crate::error::AppError::Internal(e.to_string()))
.unwrap(); .unwrap();
assert!(!res_orphans.contains("Y")); // After merging X into Y and purging self-loops, Y is the sole node and becomes an orphan
assert!(res_orphans.contains("Y"));
} }
#[tokio::test] #[tokio::test]
async fn test_more_graph_handlers() { async fn test_more_graph_handlers() {
let dir = tempfile::tempdir().unwrap(); let state = Arc::new(MemoryState::new_in_memory());
let state = Arc::new(MemoryState::new(dir.path().to_str().unwrap()));
let create_handler = CreateEntitiesHandler; let create_handler = CreateEntitiesHandler;
let args_ent = json!({ let args_ent = json!({
+4 -2
View File
@@ -64,10 +64,12 @@ impl McpTool for GetRecentLogsHandler {
file.seek(SeekFrom::End(-(read_size as i64))) file.seek(SeekFrom::End(-(read_size as i64)))
.map_err(|e| e.to_string())?; .map_err(|e| e.to_string())?;
let mut buffer = String::new(); let mut vec_buf = Vec::new();
file.read_to_string(&mut buffer) file.read_to_end(&mut vec_buf)
.map_err(|e| e.to_string())?; .map_err(|e| e.to_string())?;
let buffer = String::from_utf8_lossy(&vec_buf).to_string();
let lines: Vec<&str> = buffer.lines().collect(); let lines: Vec<&str> = buffer.lines().collect();
let recent_lines = if lines.len() > 100 { let recent_lines = if lines.len() > 100 {
lines[lines.len() - 100..].join("\n") lines[lines.len() - 100..].join("\n")
+98 -27
View File
@@ -15,13 +15,16 @@ impl McpTool for LogDecisionHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogDecisionTool>("log_decision", "Execute log_decision") crate::mcp::tool_def::<LogDecisionTool>(
"log_decision",
"Log an architectural decision record (ADR) with title, context, and consequences",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: LogDecisionTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let idx = state.get_search_index(); let idx = state.get_search_index().await;
let mut final_id = String::new(); let mut final_id = String::new();
let mut adrs_to_index = Vec::new(); let mut adrs_to_index = Vec::new();
@@ -71,11 +74,16 @@ impl McpTool for QueryDecisionsHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryDecisionsTool>("query_decisions", "Execute query_decisions") crate::mcp::tool_def::<QueryDecisionsTool>(
"query_decisions",
"Query or search existing architectural decision records (ADRs)",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?; let req: QueryDecisionsTool = serde_json::from_value(args).map_err(|e| e.to_string())?;
let limit = req.limit.unwrap_or(20);
let include_body = req.include_body.unwrap_or(true);
let data = state.code.adrs.read_with(|adrs| { let data = state.code.adrs.read_with(|adrs| {
let filtered: Vec<_> = adrs let filtered: Vec<_> = adrs
.iter() .iter()
@@ -88,8 +96,20 @@ impl McpTool for QueryDecisionsHandler {
true true
} }
}) })
.take(limit)
.collect(); .collect();
if include_body {
Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?) Ok::<String, crate::error::AppError>(serde_json::to_string(&filtered)?)
} else {
let compact: Vec<_> = filtered.iter().map(|a| serde_json::json!({
"id": a.id,
"title": a.title,
"status": a.status,
"timestamp": a.timestamp,
})).collect();
Ok::<String, crate::error::AppError>(serde_json::to_string(&compact)?)
}
})?; })?;
Ok(data) Ok(data)
} }
@@ -122,7 +142,8 @@ impl McpTool for DeleteDecisionHandler {
}); });
if found { if found {
state.rebuild_index().await; let idx = state.get_search_index().await;
let _ = idx.delete_document(&req.id).await;
Ok("Decision deleted successfully".to_string()) Ok("Decision deleted successfully".to_string())
} else { } else {
Err(crate::error::AppError::Internal( Err(crate::error::AppError::Internal(
@@ -141,7 +162,10 @@ impl McpTool for LogErrorFixHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogErrorFixTool>("log_error_fix", "Execute log_error_fix") crate::mcp::tool_def::<LogErrorFixTool>(
"log_error_fix",
"Log an error signature and its verified solution/fix for future diagnostic retrieval",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -288,7 +312,10 @@ impl McpTool for LogCodeChangeHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogCodeChangeTool>("log_code_change", "Execute log_code_change") crate::mcp::tool_def::<LogCodeChangeTool>(
"log_code_change",
"Log a significant code change or refactor with file path and description",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -325,6 +352,9 @@ impl McpTool for LogCodeChangeHandler {
git_commit: req.git_commit, git_commit: req.git_commit,
git_branch: req.git_branch, git_branch: req.git_branch,
}); });
if ledger.len() > 1000 {
ledger.remove(0);
}
}); });
state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description)); state.record_activity("code_change", &format!("Modified {}", req.file_path), Some(&description));
Ok(format!("Logged code change for {}: {}", req.file_path, description)) Ok(format!("Logged code change for {}: {}", req.file_path, description))
@@ -342,7 +372,7 @@ impl McpTool for QueryRecentChangesHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryRecentChangesTool>( crate::mcp::tool_def::<QueryRecentChangesTool>(
"query_recent_changes", "query_recent_changes",
"Execute query_recent_changes", "Query recent code changes and refactoring audit logs",
) )
} }
@@ -364,7 +394,10 @@ impl McpTool for LearnPreferenceHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LearnPreferenceTool>("learn_preference", "Execute learn_preference") crate::mcp::tool_def::<LearnPreferenceTool>(
"learn_preference",
"Learn or update a key-value user preference",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -392,7 +425,10 @@ impl McpTool for ReadPreferencesHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ReadPreferencesTool>("read_preferences", "Execute read_preferences") crate::mcp::tool_def::<ReadPreferencesTool>(
"read_preferences",
"Read all stored user and environment preferences",
)
} }
async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, _args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -412,7 +448,10 @@ impl McpTool for LogTechDebtHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogTechDebtTool>("log_tech_debt", "Execute log_tech_debt") crate::mcp::tool_def::<LogTechDebtTool>(
"log_tech_debt",
"Log a new technical debt item with description and ideal solution",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -452,7 +491,7 @@ impl McpTool for ResolveTechDebtHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ResolveTechDebtTool>( crate::mcp::tool_def::<ResolveTechDebtTool>(
"resolve_tech_debt", "resolve_tech_debt",
"Execute resolve_tech_debt", "Mark an existing technical debt item as resolved by ID",
) )
} }
@@ -488,7 +527,10 @@ impl McpTool for ListTechDebtHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ListTechDebtTool>("list_tech_debt", "Execute list_tech_debt") crate::mcp::tool_def::<ListTechDebtTool>(
"list_tech_debt",
"List technical debt items filtered by namespace or resolution status",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -641,7 +683,10 @@ impl McpTool for OmniSearchHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<OmniSearchTool>("omni_search", "Execute omni_search") crate::mcp::tool_def::<OmniSearchTool>(
"omni_search",
"Unified search across entities, subgraphs, tasks, code snippets, ADRs, and technical debt",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -649,7 +694,7 @@ impl McpTool for OmniSearchHandler {
let limit = req.limit.unwrap_or(5); let limit = req.limit.unwrap_or(5);
let include_body = req.include_body.unwrap_or(false); let include_body = req.include_body.unwrap_or(false);
let idx = state.get_search_index(); let idx = state.get_search_index().await;
let keyword_matches = idx let keyword_matches = idx
.search(&req.query, req.namespace.as_deref()) .search(&req.query, req.namespace.as_deref())
.unwrap_or_default(); .unwrap_or_default();
@@ -717,7 +762,7 @@ impl McpTool for OmniSearchHandler {
} }
let mut ranked_items: Vec<_> = rrf_scores.into_values().collect(); let mut ranked_items: Vec<_> = rrf_scores.into_values().collect();
ranked_items.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal)); ranked_items.sort_by(|a, b| b.0.total_cmp(&a.0));
let matches: Vec<MatchItem> = let matches: Vec<MatchItem> =
ranked_items.into_iter().map(|(_, item)| item).collect(); ranked_items.into_iter().map(|(_, item)| item).collect();
@@ -982,7 +1027,7 @@ impl McpTool for GetProjectHealthHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<GetProjectHealthTool>( crate::mcp::tool_def::<GetProjectHealthTool>(
"get_project_health", "get_project_health",
"Execute get_project_health", "Retrieve project health metrics including active tasks, technical debt, and PR checklist progress",
) )
} }
@@ -1037,7 +1082,10 @@ impl McpTool for SuggestErrorFixHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<SuggestErrorFixTool>("suggest_error_fix", "Execute suggest_error_fix") crate::mcp::tool_def::<SuggestErrorFixTool>(
"suggest_error_fix",
"Suggest historical error fixes based on stack trace vector similarity",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1066,7 +1114,7 @@ impl McpTool for SuggestErrorFixHandler {
.filter(|(_, score)| *score > 0.2) .filter(|(_, score)| *score > 0.2)
.collect(); .collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); scored.sort_by(|a, b| b.1.total_cmp(&a.1));
let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| { let suggestions: Vec<_> = scored.into_iter().take(limit).map(|(f, score)| {
serde_json::json!({ serde_json::json!({
@@ -1214,7 +1262,10 @@ impl McpTool for CheckpointStateHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<CheckpointStateTool>("checkpoint_state", "Execute checkpoint_state") crate::mcp::tool_def::<CheckpointStateTool>(
"checkpoint_state",
"Save point-in-time memory state snapshot checkpoint",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1248,7 +1299,10 @@ impl McpTool for RestoreStateHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<RestoreStateTool>("restore_state", "Execute restore_state") crate::mcp::tool_def::<RestoreStateTool>(
"restore_state",
"Restore memory state from a point-in-time snapshot checkpoint",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1290,7 +1344,10 @@ impl McpTool for QueryLineageHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryLineageTool>("query_lineage", "Execute query_lineage") crate::mcp::tool_def::<QueryLineageTool>(
"query_lineage",
"Query historical lineage and timeline of tasks, ADRs, and code changes",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1375,7 +1432,10 @@ impl McpTool for GetNextActionableTasksHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<GetNextActionableTasksTool>("get_next_actionable_tasks", "Execute get_next_actionable_tasks") crate::mcp::tool_def::<GetNextActionableTasksTool>(
"get_next_actionable_tasks",
"Get unblocked pending tasks ready for execution",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1431,7 +1491,10 @@ impl McpTool for LogHypothesisHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<LogHypothesisTool>("log_hypothesis", "Execute log_hypothesis") crate::mcp::tool_def::<LogHypothesisTool>(
"log_hypothesis",
"Log a diagnostic hypothesis and associated evidence for a task",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1463,7 +1526,10 @@ impl McpTool for QueryHypothesesHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<QueryHypothesesTool>("query_hypotheses", "Execute query_hypotheses") crate::mcp::tool_def::<QueryHypothesesTool>(
"query_hypotheses",
"Query active diagnostic hypotheses and evidence by task ID or keyword",
)
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -1701,7 +1767,12 @@ impl McpTool for AutoSessionCheckpointHandler {
timestamp, timestamp,
}; };
state.telemetry.handoff_memos.modify(|m| m.push(memo)); state.telemetry.handoff_memos.modify(|m| {
m.push(memo);
if m.len() > 100 {
m.remove(0);
}
});
state.record_activity("checkpoint", &format!("Created auto session checkpoint {}", memo_id), None); state.record_activity("checkpoint", &format!("Created auto session checkpoint {}", memo_id), None);
Ok(format!("Session checkpoint created with memo ID '{}'.", memo_id)) Ok(format!("Session checkpoint created with memo ID '{}'.", memo_id))
@@ -1759,7 +1830,7 @@ impl McpTool for SearchSnippetsHybridHandler {
.collect(); .collect();
scored.sort_by(|a, b| { scored.sort_by(|a, b| {
b["score"].as_f64().unwrap_or(0.0).partial_cmp(&a["score"].as_f64().unwrap_or(0.0)).unwrap_or(std::cmp::Ordering::Equal) b["score"].as_f64().unwrap_or(0.0).total_cmp(&a["score"].as_f64().unwrap_or(0.0))
}); });
let limit = req.limit.unwrap_or(10); let limit = req.limit.unwrap_or(10);
@@ -1974,7 +2045,7 @@ mod tests {
} }
state.rebuild_index().await; state.rebuild_index().await;
state.get_search_index().reader.reload().unwrap(); state.get_search_index().await.reader.reload().unwrap();
let omni = OmniSearchHandler; let omni = OmniSearchHandler;
let omni_res = omni let omni_res = omni
+25 -15
View File
@@ -49,7 +49,10 @@ impl McpTool for StickyNotesHandler {
} }
StickyNoteAction::Read => { StickyNoteAction::Read => {
let now = crate::handlers::utils::now_secs(); let now = crate::handlers::utils::now_secs();
let mut active_notes = Vec::new(); let has_expired = state.code.sticky.read_with(|notes| {
notes.iter().any(|n| matches!(n.expires_at, Some(exp) if exp <= now))
});
if has_expired {
state.code.sticky.modify(|notes| { state.code.sticky.modify(|notes| {
notes.retain(|n| { notes.retain(|n| {
if let Some(exp) = n.expires_at { if let Some(exp) = n.expires_at {
@@ -58,8 +61,9 @@ impl McpTool for StickyNotesHandler {
true true
} }
}); });
active_notes = notes.clone();
}); });
}
let active_notes = state.code.sticky.read_with(|notes| notes.clone());
Ok(serde_json::to_string(&active_notes)?) Ok(serde_json::to_string(&active_notes)?)
} }
StickyNoteAction::Delete => { StickyNoteAction::Delete => {
@@ -68,8 +72,9 @@ impl McpTool for StickyNotesHandler {
})?; })?;
let mut success = false; let mut success = false;
state.code.sticky.modify(|notes| { state.code.sticky.modify(|notes| {
if idx > 0 && idx <= notes.len() { let target_pos = if idx == 0 { 0 } else { idx - 1 };
notes.remove(idx - 1); if target_pos < notes.len() {
notes.remove(target_pos);
success = true; success = true;
} }
}); });
@@ -77,7 +82,7 @@ impl McpTool for StickyNotesHandler {
Ok("Sticky note deleted.".to_string()) Ok("Sticky note deleted.".to_string())
} else { } else {
Err(crate::error::AppError::Internal( Err(crate::error::AppError::Internal(
"Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid 1-based indices, then retry.".to_string(), "Invalid sticky note index. Next step: Call sticky_notes with action='read' to view active notes and valid indices, then retry.".to_string(),
)) ))
} }
} }
@@ -214,17 +219,22 @@ impl McpTool for GenerateStandupReportHandler {
serde_json::from_value(args).map_err(|e| e.to_string())?; serde_json::from_value(args).map_err(|e| e.to_string())?;
let cutoff = crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600); let cutoff = crate::handlers::utils::now_secs().saturating_sub(req.hours_lookback * 3600);
let report_str = state.project.tasks.read_with(|items| { let filtered_tasks = state.project.tasks.read_with(|items| {
state.code.ledger.read_with(|changes| { items.iter().filter(|t| t.updated_at >= cutoff).cloned().collect::<Vec<_>>()
state.telemetry.session_summaries.read_with(|summaries| {
let filtered_tasks: Vec<_> = items.iter().filter(|t| t.updated_at >= cutoff).collect();
let filtered_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
let filtered_summaries: Vec<_> = summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).collect();
let report = serde_json::json!({ "tasks_updated": filtered_tasks, "code_changes": filtered_changes, "session_summaries": filtered_summaries });
report.to_string()
})
})
}); });
let filtered_changes = state.code.ledger.read_with(|changes| {
changes.iter().filter(|c| c.timestamp >= cutoff).cloned().collect::<Vec<_>>()
});
let filtered_summaries = state.telemetry.session_summaries.read_with(|summaries| {
summaries.iter().filter(|s| s.namespace == req.namespace && s.timestamp >= cutoff).cloned().collect::<Vec<_>>()
});
let report = serde_json::json!({
"tasks_updated": filtered_tasks,
"code_changes": filtered_changes,
"session_summaries": filtered_summaries
});
let report_str = report.to_string();
Ok(report_str) Ok(report_str)
} }
} }
+4 -3
View File
@@ -46,7 +46,7 @@ impl McpTool for TasksHandler {
acceptance_criteria: vec![], acceptance_criteria: vec![],
expires_at: None, expires_at: None,
}; };
let idx = state.get_search_index(); let idx = state.get_search_index().await;
drop(idx.index_task(&task)); drop(idx.index_task(&task));
state.project.tasks.modify(|tasks| { state.project.tasks.modify(|tasks| {
tasks.push(task.clone()); tasks.push(task.clone());
@@ -187,7 +187,7 @@ impl McpTool for TasksHandler {
}); });
if deleted_count > 0 { if deleted_count > 0 {
let idx = state.get_search_index(); let idx = state.get_search_index().await;
for deleted_id in actually_deleted { for deleted_id in actually_deleted {
drop(idx.delete_document(&deleted_id)); drop(idx.delete_document(&deleted_id));
} }
@@ -228,7 +228,8 @@ impl McpTool for TasksHandler {
if let Some(max_t) = req.max_tokens { if let Some(max_t) = req.max_tokens {
let char_limit = max_t * 4; let char_limit = max_t * 4;
if json_str.len() > char_limit { if json_str.len() > char_limit {
json_str.truncate(char_limit); let boundary = json_str.floor_char_boundary(char_limit);
json_str.truncate(boundary);
json_str.push_str(" ...[truncated due to max_tokens]"); json_str.push_str(" ...[truncated due to max_tokens]");
} }
} }
+62 -69
View File
@@ -30,26 +30,50 @@ impl McpTool for WriteClipboardHandler {
let res = tokio::task::spawn_blocking(move || { let res = tokio::task::spawn_blocking(move || {
let mut msgs = Vec::new(); let mut msgs = Vec::new();
if let Ok(mut clipboard) = Clipboard::new() if let Some(text) = &tool_args.text {
&& let Some(text) = &tool_args.text let mut written = false;
&& clipboard.set_text(text).is_ok() if let Ok(mut clipboard) = Clipboard::new() {
{ for _ in 0..3 {
if clipboard.set_text(text).is_ok() {
written = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
}
if written {
msgs.push("Wrote text"); msgs.push("Wrote text");
} }
}
// Handle arboard for image // Handle arboard for image
if let Some(image_path) = &tool_args.image_path { if let Some(image_path) = &tool_args.image_path {
match image::open(image_path) { match image::open(image_path) {
Ok(img) => { Ok(img) => {
let img = if img.width() > 2048 || img.height() > 2048 {
img.resize(2048, 2048, FilterType::Triangle)
} else {
img
};
let rgba = img.into_rgba8(); let rgba = img.into_rgba8();
let (w, h) = rgba.dimensions(); let (w, h) = rgba.dimensions();
let raw_bytes = rgba.into_raw();
let mut written = false;
if let Ok(mut clipboard) = arboard::Clipboard::new() {
for _ in 0..3 {
let img_data = ImageData { let img_data = ImageData {
width: w as usize, width: w as usize,
height: h as usize, height: h as usize,
bytes: Cow::Owned(rgba.into_raw()), bytes: Cow::Borrowed(&raw_bytes),
}; };
if let Ok(mut clipboard) = arboard::Clipboard::new() {
if clipboard.set_image(img_data).is_ok() { if clipboard.set_image(img_data).is_ok() {
written = true;
break;
}
std::thread::sleep(std::time::Duration::from_millis(20));
}
}
if written {
msgs.push("Wrote image"); msgs.push("Wrote image");
} else { } else {
return Err(crate::error::AppError::Internal( return Err(crate::error::AppError::Internal(
@@ -57,7 +81,6 @@ impl McpTool for WriteClipboardHandler {
)); ));
} }
} }
}
Err(e) => { Err(e) => {
return Err(crate::error::AppError::Internal(format!( return Err(crate::error::AppError::Internal(format!(
"Failed to open image {}: {}", "Failed to open image {}: {}",
@@ -85,6 +108,7 @@ impl McpTool for WriteClipboardHandler {
} }
pub fn get_native_clipboard_text() -> Option<String> { pub fn get_native_clipboard_text() -> Option<String> {
for _ in 0..3 {
if let Ok(mut clipboard) = arboard::Clipboard::new() { if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(text) = clipboard.get_text() { if let Ok(text) = clipboard.get_text() {
if !text.trim().is_empty() { if !text.trim().is_empty() {
@@ -92,42 +116,13 @@ pub fn get_native_clipboard_text() -> Option<String> {
} }
} }
} }
std::thread::sleep(std::time::Duration::from_millis(20));
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.arg("--no-newline");
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
} }
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
if !text.trim().is_empty() {
return Some(text);
}
}
}
}
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(text) = String::from_utf8(output.stdout) {
if !text.trim().is_empty() {
return Some(text);
}
}
}
}
None None
} }
pub fn get_native_clipboard_image() -> Option<image::DynamicImage> { pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
for _ in 0..3 {
if let Ok(mut clipboard) = arboard::Clipboard::new() { if let Ok(mut clipboard) = arboard::Clipboard::new() {
if let Ok(image_data) = clipboard.get_image() { if let Ok(image_data) = clipboard.get_image() {
if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw( if let Some(img) = ImageBuffer::<image::Rgba<u8>, _>::from_raw(
@@ -139,38 +134,8 @@ pub fn get_native_clipboard_image() -> Option<image::DynamicImage> {
} }
} }
} }
std::thread::sleep(std::time::Duration::from_millis(20));
for mime in &["image/png", "image/jpeg", "image/bmp", "image/tiff"] {
let mut cmd_wl = std::process::Command::new("wl-paste");
cmd_wl.args(["--type", mime]);
if std::env::var("WAYLAND_DISPLAY").is_err() && std::path::Path::new("/mnt/wslg/runtime-dir").exists() {
cmd_wl.env("WAYLAND_DISPLAY", "wayland-0");
cmd_wl.env("XDG_RUNTIME_DIR", "/mnt/wslg/runtime-dir");
} }
if let Ok(output) = cmd_wl.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
}
}
}
}
for mime in &["image/png", "image/jpeg", "image/bmp"] {
let mut cmd_xc = std::process::Command::new("xclip");
cmd_xc.args(["-selection", "clipboard", "-t", mime, "-o"]);
if std::env::var("DISPLAY").is_err() {
cmd_xc.env("DISPLAY", ":0");
}
if let Ok(output) = cmd_xc.output() {
if output.status.success() && !output.stdout.is_empty() {
if let Ok(img) = image::load_from_memory(&output.stdout) {
return Some(img);
}
}
}
}
None None
} }
@@ -422,5 +387,33 @@ mod tests {
let parsed: serde_json::Value = serde_json::from_str(&result).unwrap(); let parsed: serde_json::Value = serde_json::from_str(&result).unwrap();
assert!(parsed.is_object()); assert!(parsed.is_object());
} }
#[test]
fn test_no_subprocess_clipboard_regression() {
let vision_src = include_str!("vision.rs");
let code_only = vision_src.split("mod tests").next().unwrap_or(vision_src);
let forbidden_cmd = format!("Command::{}{}", "n", "ew");
let forbidden_ps = format!("power{}", "shell");
let forbidden_wl = format!("wl-{}", "paste");
let forbidden_xc = format!("x{}", "clip");
assert!(
!code_only.contains(&forbidden_cmd),
"Regression detected: vision.rs must not spawn subprocesses!"
);
assert!(
!code_only.contains(&forbidden_ps),
"Regression detected: vision.rs must not invoke powershell!"
);
assert!(
!code_only.contains(&forbidden_wl),
"Regression detected: vision.rs must not invoke wl-paste!"
);
assert!(
!code_only.contains(&forbidden_xc),
"Regression detected: vision.rs must not invoke xclip!"
);
}
} }
+10 -7
View File
@@ -15,7 +15,7 @@ impl McpTool for PinFileHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<PinFileTool>("pin_file", "Execute pin_file") crate::mcp::tool_def::<PinFileTool>("pin_file", "Pin a file path in the workspace context to prioritize it across subagent turns.")
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -42,7 +42,7 @@ impl McpTool for UnpinFileHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Execute unpin_file") crate::mcp::tool_def::<UnpinFileTool>("unpin_file", "Unpin a file path from the workspace context.")
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -65,7 +65,7 @@ impl McpTool for ListPinnedFilesHandler {
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<ListPinnedFilesTool>( crate::mcp::tool_def::<ListPinnedFilesTool>(
"list_pinned_files", "list_pinned_files",
"Execute list_pinned_files", "List active pinned files in the workspace context, optionally filtered by namespace or branch.",
) )
} }
@@ -104,7 +104,7 @@ impl McpTool for StoreSnippetHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Execute store_snippet") crate::mcp::tool_def::<StoreSnippetTool>("store_snippet", "Store or replace a reusable code snippet with vector embedding and search tags.")
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -127,7 +127,7 @@ impl McpTool for StoreSnippetHandler {
embedding, embedding,
}; };
let idx = state.get_search_index(); let idx = state.get_search_index().await;
drop(idx.index_snippet(&snippet)); drop(idx.index_snippet(&snippet));
state.code.snippets.modify(|snippets| { state.code.snippets.modify(|snippets| {
@@ -148,7 +148,7 @@ impl McpTool for TagSnippetHandler {
} }
fn schema(&self) -> Value { fn schema(&self) -> Value {
crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Execute tag_snippet") crate::mcp::tool_def::<TagSnippetTool>("tag_snippet", "Add organizational tags to an existing stored code snippet.")
} }
async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn execute(&self, args: Value, state: Arc<MemoryState>) -> crate::error::Result<String> {
@@ -254,7 +254,7 @@ impl McpTool for DeleteSnippetHandler {
deleted = snippets.len() < orig; deleted = snippets.len() < orig;
}); });
if deleted { if deleted {
let idx = state.get_search_index(); let idx = state.get_search_index().await;
drop(idx.delete_document(&req.name)); drop(idx.delete_document(&req.name));
Ok("Snippet deleted.".to_string()) Ok("Snippet deleted.".to_string())
} else { } else {
@@ -870,6 +870,9 @@ impl McpTool for ReadDirectoryArchitectureHandler {
summary: &mut String, summary: &mut String,
depth: usize, depth: usize,
) -> std::io::Result<()> { ) -> std::io::Result<()> {
if depth > 10 {
return Ok(());
}
if dir.is_dir() { if dir.is_dir() {
let mut entries = let mut entries =
fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?; fs::read_dir(dir)?.collect::<Result<Vec<_>, std::io::Error>>()?;
+63 -29
View File
@@ -46,7 +46,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
.await .await
.unwrap_or_default(); .unwrap_or_default();
let idx = state.get_search_index(); let idx = state.get_search_index().await;
for file_path in files_to_process { for file_path in files_to_process {
if let Ok(content) = std::fs::read_to_string(&file_path) { if let Ok(content) = std::fs::read_to_string(&file_path) {
@@ -71,36 +71,64 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
let mut chunks = Vec::new(); let mut chunks = Vec::new();
extract_chunks(tree.root_node(), &content, &mut chunks, ext); extract_chunks(tree.root_node(), &content, &mut chunks, ext);
for (name, code, desc) in chunks { // Gold Standard: Batch generate embeddings in chunks of 16 to eliminate sequential HTTP overhead
// Generate embedding for chunk_batch in chunks.chunks(16) {
if let Ok(mut emb) = generate_embeddings_async(vec![code.clone()]).await { let texts: Vec<String> = chunk_batch.iter().map(|(_, code, _)| code.clone()).collect();
let embedding = emb.pop(); let embeddings = generate_embeddings_async(texts).await.unwrap_or_default();
let file_name = let mut new_snippets = Vec::with_capacity(chunk_batch.len());
file_path.file_name().unwrap_or_default().to_string_lossy(); let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
for (i, (name, code, desc)) in chunk_batch.iter().enumerate() {
let embedding = embeddings.get(i).cloned();
let file_name = file_path.file_name().unwrap_or_default().to_string_lossy();
let snippet_name = format!("{}:{}", file_name, name); let snippet_name = format!("{}:{}", file_name, name);
let snippet = Snippet { let snippet = Snippet {
name: snippet_name.to_string(), name: snippet_name,
language: ext.to_string(), language: ext.to_string(),
code: code.clone(), code: code.clone(),
description: format!("{} in {}", desc, file_path.display()), description: format!("{} in {}", desc, file_path.display()),
updated_at: std::time::SystemTime::now() updated_at: now,
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs(),
tags: vec![], tags: vec![],
embedding, embedding,
}; };
new_snippets.push(snippet);
}
// Gold Standard: Modify store ONCE per batch with zero-copy HashSet<&str> lookup
let mut snippets_to_index = Vec::new();
state.code.snippets.modify(|snippets| { state.code.snippets.modify(|snippets| {
// Prevent duplicates if already indexed let existing_names: std::collections::HashSet<&str> =
if !snippets.iter().any(|s| s.name == snippet.name) { snippets.iter().map(|s| s.name.as_str()).collect();
let mut filtered_new = Vec::with_capacity(new_snippets.len());
let mut seen_in_batch = std::collections::HashSet::new();
for snippet in new_snippets {
if !existing_names.contains(snippet.name.as_str())
&& seen_in_batch.insert(snippet.name.clone())
{
filtered_new.push(snippet);
}
}
for snippet in filtered_new {
snippets.push(snippet.clone()); snippets.push(snippet.clone());
snippets_to_index.push(snippet);
}
if snippets.len() > 1000 {
let overflow = snippets.len() - 1000;
snippets.drain(0..overflow);
} }
}); });
let _ = idx.index_snippet(&snippet).await; for snippet in &snippets_to_index {
let _ = idx.index_snippet(snippet).await;
} }
} }
} }
@@ -111,7 +139,7 @@ pub async fn start_background_indexer(state: Arc<MemoryState>) {
} }
fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) { fn extract_chunks(node: Node, code: &str, chunks: &mut Vec<(String, String, String)>, ext: &str) {
extract_chunks_with_parent(node, code, chunks, ext, None); extract_chunks_with_parent(node, code, chunks, ext, None, 0);
} }
fn extract_chunks_with_parent( fn extract_chunks_with_parent(
@@ -120,22 +148,28 @@ fn extract_chunks_with_parent(
chunks: &mut Vec<(String, String, String)>, chunks: &mut Vec<(String, String, String)>,
ext: &str, ext: &str,
parent_scope: Option<&str>, parent_scope: Option<&str>,
depth: usize,
) { ) {
// Stack overflow protection: Cap recursion depth at 100
if depth > 100 {
return;
}
let kind = node.kind(); let kind = node.kind();
let is_impl_or_class = matches!(kind, "impl_item" | "class_declaration" | "class_definition"); let is_impl_or_class = matches!(kind, "impl_item" | "class_declaration" | "class_definition");
let current_scope = if is_impl_or_class { let current_scope: Option<&str> = if is_impl_or_class {
let mut cursor = node.walk(); let mut cursor = node.walk();
let mut type_name = None; let mut type_name = None;
for child in node.children(&mut cursor) { for child in node.children(&mut cursor) {
if child.kind() == "type_identifier" || child.kind() == "name" || child.kind() == "identifier" { if child.kind() == "type_identifier" || child.kind() == "name" || child.kind() == "identifier" {
type_name = child.utf8_text(code.as_bytes()).ok().map(|s| s.to_string()); type_name = child.utf8_text(code.as_bytes()).ok();
break; break;
} }
} }
type_name.or_else(|| parent_scope.map(|s| s.to_string())) type_name.or(parent_scope)
} else { } else {
parent_scope.map(|s| s.to_string()) parent_scope
}; };
let is_structural = matches!( let is_structural = matches!(
@@ -151,30 +185,30 @@ fn extract_chunks_with_parent(
if is_structural { if is_structural {
let mut raw_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string(); let mut raw_text = node.utf8_text(code.as_bytes()).unwrap_or("").to_string();
let mut name = "unknown".to_string(); let mut name = "unknown";
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.children(&mut cursor) { for child in node.children(&mut cursor) {
let child_kind = child.kind(); let child_kind = child.kind();
if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" { if child_kind == "identifier" || child_kind == "name" || child_kind == "type_identifier" {
name = child if let Ok(text) = child.utf8_text(code.as_bytes()) {
.utf8_text(code.as_bytes()) name = text;
.unwrap_or("unknown") }
.to_string();
break; break;
} }
} }
if let Some(ref scope) = current_scope { let mut final_name = name.to_string();
if let Some(scope) = current_scope {
raw_text = format!("// Parent Scope: {}\n{}", scope, raw_text); raw_text = format!("// Parent Scope: {}\n{}", scope, raw_text);
name = format!("{}::{}", scope, name); final_name = format!("{}::{}", scope, name);
} }
let desc = format!("{} AST node", kind); let desc = format!("{} AST node", kind);
chunks.push((name, raw_text, desc)); chunks.push((final_name, raw_text, desc));
} else { } else {
let mut cursor = node.walk(); let mut cursor = node.walk();
for child in node.named_children(&mut cursor) { for child in node.named_children(&mut cursor) {
extract_chunks_with_parent(child, code, chunks, ext, current_scope.as_deref()); extract_chunks_with_parent(child, code, chunks, ext, current_scope, depth + 1);
} }
} }
} }
+86 -40
View File
@@ -11,15 +11,17 @@ To prevent graph fragmentation and ensure seamless LLM context retrieval:
* **Entity Types (`entity_type`)**: MUST ALWAYS be **`PascalCase`** (e.g. `DatabaseTable`, `McpTool`, `ArchitectureComponent`, `File`, `DataStructure`). * **Entity Types (`entity_type`)**: MUST ALWAYS be **`PascalCase`** (e.g. `DatabaseTable`, `McpTool`, `ArchitectureComponent`, `File`, `DataStructure`).
* **Relation Types (`relation_type`)**: MUST ALWAYS be **`snake_case`** (e.g. `depends_on`, `calls`, `implements`, `uses`, `contains`). * **Relation Types (`relation_type`)**: MUST ALWAYS be **`snake_case`** (e.g. `depends_on`, `calls`, `implements`, `uses`, `contains`).
* **Field Keys & Properties**: MUST ALWAYS be **`snake_case`** (e.g. `file_path`, `git_commit`, `created_at`). * **Field Keys & Properties**: MUST ALWAYS be **`snake_case`** (e.g. `file_path`, `git_commit`, `created_at`).
* **Relation Tool Parameters**: `create_relations` supports Serde field aliases (`source` -> `from`, `target` -> `to`, `relationType`/`type` -> `relation_type`) so LLM tool calls succeed seamlessly regardless of parameter naming.
> [!NOTE] > [!NOTE]
> The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation. > The server automatically enforces and migrates incoming entity and relation types to these canonical casing rules on every read and write operation.
> Store operations enforce atomic lock scope for serialization/deserialization and enter Quarantine Mode upon database corruption to prevent data overwrite hazards.
--- ---
## 2. Consolidated Smart Tools Architecture ## 2. Consolidated Smart Tools Architecture
The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools: The server consolidates granular single-purpose tools into domain-named smart tools. Always prefer the consolidated tools over legacy aliases:
* **`tasks`**: Complete task lifecycle management. * **`tasks`**: Complete task lifecycle management.
- `action: "add"`: Create a new task (requires `title`, optional `description`, `git_branch`, `parent_id`, `dependencies`). - `action: "add"`: Create a new task (requires `title`, optional `description`, `git_branch`, `parent_id`, `dependencies`).
@@ -85,55 +87,99 @@ The server consolidates granular single-purpose tools into domain-named smart to
- `action: "register"`: Register target environment (requires `name`). - `action: "register"`: Register target environment (requires `name`).
- `action: "get_details"`: Read full environment details. - `action: "get_details"`: Read full environment details.
* **`clipboard`**: OS Clipboard management (native cross-platform support for Windows and Linux wl-paste/xclip). * **`clipboard`**: OS Clipboard management.
- `action: "read"`: Read OS clipboard text, HTML, file paths, or JPEG image. - `action: "read"`: Read OS clipboard.
- `action: "write"`: Write text/html/files/image to clipboard. - `action: "write"`: Write text/html/files/image to clipboard.
- `action: "toggle_watch"`: Toggle auto-clipboard watcher. - `action: "toggle_watch"`: Toggle auto-clipboard watcher.
* **`manage_subagent_namespace`**: Subagent memory isolation. ---
- `action: "create"`: Create namespace.
- `action: "purge"`: Purge namespace.
- `action: "condense"`: Condense namespace.
* **`manage_checkpoint`**: State snapshots and rollbacks. ## 4. High-Signal Tool Responses & Performance Guidelines
- `action: "checkpoint"` (or `"create_snapshot"`): Create state snapshot.
- `action: "restore"` (or `"restore_snapshot"`): Restore point-in-time snapshot. To optimize context usage, response times, and LLM reasoning efficiency:
* **High-Signal Feedback**:
- MCP tool calls (such as `create_entities`, `create_relations`, `add_observations`, `pin_file`) return explicit, structured summaries containing created names, types, edge counts, and relation paths.
- LLMs do NOT need to execute follow-up `open_nodes` calls merely to confirm successful creation.
* **Batch Operations**:
- When creating or updating multiple entities, snippets, or observations, always batch items into a single tool call array (e.g. `create_entities` with multiple array items) rather than making separate calls.
- The server handles batch store mutations serially in a single transaction pass with single-permit event-driven flushes.
* **Real-time Tantivy Search Indexing**:
- The Tantivy search engine automatically checks pending commits and reloads search readers prior to executing `omni_search` or `search_nodes`. Search queries always return up-to-date document results immediately following mutations.
- Single-item deletions use targeted document removal rather than global index wipes.
* **Real-Time AST & Workspace Source Code Symbol Scanning**:
- `find_symbol_references`, `get_callers`, and `analyze_impact` scan both stored code snippets and physical workspace source code files on disk (`.rs`, `.ts`, `.py`, `.go`, `.java`, `.c`, `.cpp`), providing accurate AST symbol references and call site tracking.
* **Graph Entity Merge & Self-Loop Protection**:
- `merge_entities` re-points all relations from `source_entity` to `target_entity` and automatically prunes cyclic self-loops (`target -> target`).
* **Safe UTF-8 Token Truncation**:
- Large responses (e.g. `get_active_worktree_context`, `read_graph`, `summarize_subgraph`) are safely truncated along UTF-8 character boundaries (`floor_char_boundary`), ensuring response bounds without runtime panics.
## 4. Automated Error Fix Auto-Matcher
- **Tools:** `log_error_fix`, `search_error_fixes` (and alias `suggest_error_fix`)
- **When to use:** When encountering a build error, test failure, or stack trace. Call `search_error_fixes` with either a text `query` or `stack_trace` before attempting a fix from scratch.
- **Behavior:** Computes cosine similarity between error trace embeddings and past resolution logs when `stack_trace` is provided, or keyword filtering when `query` is provided, returning top matched solutions, modified files, and git commits.
--- ---
## 3. Knowledge Graph Operations & Multi-Hop Navigation ## 5. Memory State Checkpointing & Rollbacks
* **`create_entities`**: Create entity nodes (`entity_type` PascalCase). - **Tool:** `checkpoint_state`, `restore_state` (or `create_snapshot`, `restore_snapshot`)
* **`create_relations`**: Create relations (`relation_type` snake_case). - **When to use:** Before initiating a large refactor, running experimental subagent tasks, or executing destructive batch operations.
* **`add_observations`**: Add new observations to existing entities. - **Behavior:** Saves or restores a point-in-time snapshot of graph entities, active tasks, and tech debt backlogs.
* **`delete_entities`**, **`delete_observations`**, **`delete_relations`**: Graph cleanup.
* **`read_graph`**, **`search_nodes`**, **`open_nodes`**: Graph query and entity lookup.
* **`visualize_graph`**: Generate graph visualizer output.
* **`condense_entity`**, **`merge_entities`**, **`find_orphans`**: Graph maintenance.
* **`get_subgraph`**: Perform multi-hop BFS neighborhood expansion tree around a root entity up to `depth` hops with token budget.
* **`sweep_graph_health`**: Self-healing graph health auditor for near-duplicates and orphaned nodes.
* **`resolve_stale_symbols`**, **`summarize_subgraph`**: Symbol resolution and subgraph summarization.
--- ---
## 4. Search, Code & AST Intelligence ## 6. Self-Healing Graph Health Sweeper
* **`omni_search`**: High-performance unified search across tasks, code changes, ADRs, debt, and snippets. - **Tool:** `sweep_graph_health`
* **`semantic_code_search`**: Vector similarity search across indexing embeddings. - **When to use:** Periodically or before committing major graph changes to audit entity consistency.
* **`read_file_skeleton`**: Extract file declarations/skeletons using Tree-Sitter (Rust, Python, Go, Java, C, C++, JSX/TSX). - **Behavior:** Detects orphaned nodes (0 relations), computes name similarity using pre-computed lowercase keys to identify near-duplicates (e.g., `APIGateway` vs `ApiGateway`), and provides structured `merge_entities` recommendations or auto-prunes orphans.
* **`replace_ast_node`**: Structural AST node replacement across files.
* **`find_symbol_references`**, **`get_callers`**, **`analyze_impact`**: Code dependency and blast radius analysis.
* **`read_directory_architecture`**: Inspect directory structure and code organization.
--- ---
## 5. Audit Logging, Diagnostics & Provenance ## 7. Causal Lineage & Provenance Tracker
* **`log_code_change`**, **`query_recent_changes`**: Audit trail of code edits, commits, and symbol modifications. - **Tool:** `query_lineage`
* **`log_error_fix`**, **`search_error_fixes`**: Record error trace and verified fix; search past resolutions via vector similarity. - **When to use:** When asking *"Why was this component modified?"* or *"What task or ADR led to this code change?"*
* **`query_lineage`**: Reconstruct causal timeline explaining why a file/symbol/commit was modified. - **Behavior:** Searches across tasks, ADRs, audit ledger entries, and error fixes to assemble a unified chronological timeline explaining the provenance behind any file, symbol, or commit.
* **`log_hypothesis`**, **`query_hypotheses`**: Diagnostic hypothesis tracking during complex debugging.
* **`get_preflight_context`**: Consolidated executive bundle (branch, active tasks, criteria, pinned files, tech debt) for instant context warmup. ---
* **`learn_preference`**, **`read_preferences`**: Track user preferences and project conventions.
* **`get_project_health`**: Health metrics, active task counts, debt backlog, and graph statistics. ## 9. Native Rust Invariants & Subprocess Prohibition (CRITICAL)
* **`broadcast_agent_signal`**, **`query_agent_signals`**: Inter-agent coordination signals.
* **`auto_session_checkpoint`**: Auto-save session state checkpoint. To maintain maximum security, speed, and cross-platform reliability:
* **`get_active_worktree_context`**, **`query_git_diffs`**: Git worktree status and diff inspection. * **Zero Subprocess Fallbacks**: System and server tools (`clipboard`, `ast`, `search`, `db`) MUST strictly use pure native Rust crates (`arboard`, `tree-sitter`, `tantivy`, `psycopg`). Invocations of external shell commands (`powershell.exe`, `wl-paste`, `xclip`, `cmd.exe`) are strictly prohibited in native handlers.
* **`watch_process_logs`**, **`get_recent_logs`**: Process log monitoring and retrieval. * **Transient Lock Recovery**: Transient OS handle collisions (such as Win32 OLE `OpenClipboard` lock contention) must be handled using native retry loops with backoffs directly in Rust.
* **Automated Static Regression Gates**: Automated AST/source audit tests (e.g. `test_no_subprocess_clipboard_regression`) verify at test time that forbidden subprocess patterns are absent from handler implementations.
---
## 10. High-Performance Concurrency & Resilience Guarantees
* **Zero-Allocation NDJSON Reader**: `mcp-stdio` reclaims line buffers via `std::mem::take` and in-place trimming to eliminate heap allocations during high-frequency NDJSON message parsing.
* **Micro-Batched Fastembed Inference**: `generate_embeddings_async` chunks text arrays into 32-item micro-batches inside `spawn_blocking`, eliminating RAM/CPU spikes during batch indexing.
* **Single-Handle OS Clipboard Retries & Image Downsampling**: `WriteClipboardHandler` initializes OS clipboard handles once per operation and downsamples images exceeding $2048 \times 2048$ resolution before writing raw RGBA bytes.
* **Zero Transaction Drop Persistence Guarantee**: `Store::modify` automatically spawns an async task to execute `push_async` with channel backpressure if `push` encounters queue saturation, ensuring zero data loss under spike write loads.
* **Non-Blocking Read Lock Sticky Notes**: `StickyNotesHandler` (`action = "read"`) queries notes using shared read locks, executing write pruning only when expired items exist.
* **Token-Budgeted Query Projections**: Decision queries (`query_decisions`) support `limit` caps and compact `include_body: false` projections for token budget optimization.
* **Serde Parameter & Enum Ergonomics**: Action enums (`StickyNoteAction`, `SnippetSearchMode`, `Relation`) support case-insensitive variants and common synonyms (`create`/`add`, `remove`/`delete`, `list`/`read`, `source`/`from`, `target`/`to`, `relationType`/`relation_type`) ensuring seamless LLM tool execution.
* **Atomic Search Index Swaps**: `MemoryState::rebuild_index` constructs and populates a new `MemoryIndex` instance in isolation before performing an atomic pointer swap (`*self.search_index.write().await = new_idx`), eliminating transient empty search result windows.
* **Non-Blocking Tantivy Search Queries**: `MemoryIndex::search` queries current index searcher snapshots without executing synchronous disk commit locks, preventing query stalls during heavy background indexing.
* **Zero-Allocation HashSet<&str> Snippet Deduplication**: `indexer.rs` utilizes borrowed `HashSet<&str>` name lookups during snippet batch modifications, eliminating heap string re-allocations inside the store write lock.
* **AST Recursion Depth Safeguard & Zero-Copy Borrowing**: Tree-sitter AST traversal caps recursion depth at 100 to prevent thread stack overflows and borrows string slices (`&str`) during AST node walking.
* **Strongly-Typed SearchResult & Pre-Allocated Search Vectors**: `search.rs` uses a strongly-typed `SearchResult` struct with named fields and pre-allocates result vector capacity (`Vec::with_capacity(top_docs.len())`).
* **BFS Graph Traversal Pre-allocation & Visited Node Upper Bound**: `GraphQueryBuilder::find_shortest_path` pre-allocates adjacency map capacity (`HashMap::with_capacity(relations.len() * 2)`) and enforces a visited node upper bound (10,000 max) to guarantee deterministic BFS runtime.
* **LLM Tool Schema Ergonomics & Context Guidance**: `ReadGraphHandler` schema explicitly instructs LLMs on `namespace` filtering and `search_nodes` / `get_subgraph` tools for large graph discovery.
* **Filesystem Event Debouncing & Proactive State Refresh**: `spawn_watcher` implements a sliding 250ms debouncing window per file path, ignores `.git`, `target`, `.gemini`, and `node_modules`, and broadcasts activity events to `MemoryState`.
* **Buffered Line-by-Line AST Workspace Symbol Scanning**: `scan_workspace_for_symbol` reads workspace files via `BufReader` line streams instead of loading entire files into heap strings, preventing memory spikes when traversing source trees.
* **AST Node Type Aliasing & Skeleton Preallocation**: `replace_ast_node` documents friendly node aliases (`function`, `fn`, `method`, `struct`, `class`, `enum`, `trait`, `type`), and `read_file_skeleton` preallocates string buffer capacity (`code.len() / 2`).
* **Batch Vector Indexing & Similarity Score Guidance**: `VectorDB` provides `index_documents_batch` for single-request multi-point vector upserts and explicit score calibration notes ($\ge 0.75$ high confidence match).
* **Compact JSON MCP Resources & UTF-8 Activity Truncation**: MCP resources serialize using compact JSON (`to_string`), `TerminalHistoryResource` / `MilestonesResource` enforce output bounds, and `format_tool_activity_description` uses `floor_char_boundary` for guaranteed UTF-8 safety.
* **SIMD-Friendly Single-Pass Cosine Similarity**: `cosine_similarity` calculates dot product and Euclidean norm squares in a single linear pass over float vectors, enabling SIMD compiler auto-vectorization.
* **Safe Stream Decoding on Log Tails**: Log tail operations (`get_recent_logs`) read raw bytes and decode using lossy UTF-8 conversion (`String::from_utf8_lossy`) to ensure resilience when seeking across multi-byte UTF-8 boundaries.
* **Task Summary UTF-8 Truncation Safety**: `tasks` tool (`action = "list"`) truncates serialized task text strictly along UTF-8 character boundaries using `floor_char_boundary` when enforcing `max_tokens`.
* **Sequential Snapshot Lock Scope Flattening**: `GenerateStandupReportHandler` reads `tasks`, `ledger`, and `session_summaries` sequentially rather than nesting read locks, preventing multi-lock deadlocks during concurrent store modifications.
* **Directory Tree Depth Safeguard**: `ReadDirectoryArchitectureHandler` caps directory recursion at depth 10 to prevent stack overflow on deep or cyclic directory structures.
* **Deterministic Total-Order Score Ranking**: `OmniSearchHandler` uses `f64::total_cmp` for Reciprocal Rank Fusion (RRF) score sorting, guaranteeing deterministic NaN-safe search result ordering.
* **RPC Timeout Memory Hygiene**: `nvim-core` maintains request hygiene by removing pending request entries from static RPC maps upon timeout or channel drop, eliminating orphan memory leaks.
+27 -13
View File
@@ -122,11 +122,9 @@ pub async fn ttl_sweeper_worker(state: Arc<MemoryState>) {
pub async fn index_committer_worker(state: Arc<MemoryState>) { pub async fn index_committer_worker(state: Arc<MemoryState>) {
loop { loop {
state.index_commit_notify.notified().await; state.index_commit_notify.notified().await;
let idx_opt = state.search_index.read().ok().map(|idx| idx.clone()); let idx = state.search_index.read().await.clone();
if let Some(idx) = idx_opt {
let _ = idx.commit().await; let _ = idx.commit().await;
} }
}
} }
pub async fn condense_graph_worker(state: Arc<MemoryState>) { pub async fn condense_graph_worker(state: Arc<MemoryState>) {
@@ -234,15 +232,26 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
loop { loop {
match rx.recv().await { match rx.recv().await {
Ok(msg) => { Ok(msg) => {
let senders: Vec<_> = app_state_clone let mut closed_ids = Vec::new();
{
let clients_guard = app_state_clone
.clients .clients
.read() .read()
.unwrap_or_else(|e| e.into_inner()) .unwrap_or_else(|e| e.into_inner());
.values() for (id, tx) in clients_guard.iter() {
.cloned() if tx.try_send(msg.clone()).is_err() && tx.is_closed() {
.collect(); closed_ids.push(id.clone());
for client_tx in senders { }
let _ = client_tx.try_send(msg.clone()); }
}
if !closed_ids.is_empty() {
let mut write_guard = app_state_clone
.clients
.write()
.unwrap_or_else(|e| e.into_inner());
for id in closed_ids {
write_guard.remove(&id);
}
} }
} }
Err(tokio::sync::broadcast::error::RecvError::Closed) => break, Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
@@ -304,9 +313,14 @@ pub async fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::erro
let session = &payload.session_id; let session = &payload.session_id;
let _is_unix_socket = session.starts_with('/') || session.starts_with('~'); let _is_unix_socket = session.starts_with('/') || session.starts_with('~');
if let Some(home) = dirs::home_dir() { if let Some(home) = dirs::home_dir() {
let nvim_txt_path = home.join(".gemini").join("active_nvim.txt"); let nvim_dir = home.join(".gemini");
let _ = tokio::fs::create_dir_all(nvim_txt_path.parent().unwrap()).await; let nvim_txt_path = nvim_dir.join("active_nvim.txt");
let _ = tokio::fs::write(&nvim_txt_path, session).await; let tmp_path = nvim_dir.join(format!("active_nvim_{}.tmp", std::process::id()));
if tokio::fs::create_dir_all(&nvim_dir).await.is_ok()
&& tokio::fs::write(&tmp_path, session).await.is_ok()
{
let _ = tokio::fs::rename(&tmp_path, &nvim_txt_path).await;
}
} }
} }
+1
View File
@@ -27,6 +27,7 @@ pub fn tool_def<T: JsonSchema>(name: &str, fallback_desc: &str) -> serde_json::V
let desc = schema_val let desc = schema_val
.get("description") .get("description")
.and_then(|d| d.as_str()) .and_then(|d| d.as_str())
.filter(|s| !s.trim().is_empty())
.unwrap_or(fallback_desc) .unwrap_or(fallback_desc)
.to_string(); .to_string();
+49
View File
@@ -72,9 +72,45 @@ pub fn normalize_relation_type(input: &str) -> String {
trimmed.to_lowercase().replace([' ', '-'], "_") trimmed.to_lowercase().replace([' ', '-'], "_")
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
#[serde(rename_all = "snake_case")]
pub enum TaskStatus {
#[default]
Active,
InProgress,
Blocked,
Completed,
}
impl TaskStatus {
pub fn is_active(&self) -> bool {
matches!(self, TaskStatus::Active | TaskStatus::InProgress | TaskStatus::Blocked)
}
pub fn is_completed(&self) -> bool {
matches!(self, TaskStatus::Completed)
}
pub fn as_str(&self) -> &'static str {
match self {
TaskStatus::Active => "active",
TaskStatus::InProgress => "in_progress",
TaskStatus::Blocked => "blocked",
TaskStatus::Completed => "completed",
}
}
}
impl std::fmt::Display for TaskStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.as_str())
}
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)] #[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct Entity { pub struct Entity {
pub name: String, pub name: String,
#[serde(alias = "entityType", alias = "type", alias = "category")]
pub entity_type: String, pub entity_type: String,
#[serde(default)] #[serde(default)]
pub observations: Vec<String>, pub observations: Vec<String>,
@@ -85,8 +121,11 @@ pub struct Entity {
} }
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, JsonSchema)] #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash, JsonSchema)]
pub struct Relation { pub struct Relation {
#[serde(alias = "source", alias = "from_node", alias = "source_node")]
pub from: String, pub from: String,
#[serde(alias = "target", alias = "to_node", alias = "target_node")]
pub to: String, pub to: String,
#[serde(alias = "relationType", alias = "type", alias = "relation")]
pub relation_type: String, pub relation_type: String,
#[serde(default = "default_namespace")] #[serde(default = "default_namespace")]
pub namespace: String, pub namespace: String,
@@ -123,6 +162,16 @@ pub struct Task {
#[serde(default)] #[serde(default)]
pub acceptance_criteria: Vec<AcceptanceCriteria>, pub acceptance_criteria: Vec<AcceptanceCriteria>,
} }
impl Task {
pub fn is_active(&self) -> bool {
let s = self.status.trim();
!s.eq_ignore_ascii_case("completed")
&& !s.eq_ignore_ascii_case("done")
&& !s.eq_ignore_ascii_case("cancelled")
&& !s.eq_ignore_ascii_case("canceled")
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)] #[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Snippet { pub struct Snippet {
pub name: String, pub name: String,
-1
View File
@@ -1 +0,0 @@
pub fn patch() {}
+24 -59
View File
@@ -57,9 +57,7 @@ impl McpResource for GraphEntitiesResource {
Some("All nodes and entities currently stored in the knowledge graph") Some("All nodes and entities currently stored in the knowledge graph")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let graph = state_clone.graph.cache.read().unwrap();
let data: Vec<_> = graph.entities.values().collect(); let data: Vec<_> = graph.entities.values().collect();
if data.len() > 100 { if data.len() > 100 {
let truncated: Vec<_> = data.iter().take(100).cloned().collect(); let truncated: Vec<_> = data.iter().take(100).cloned().collect();
@@ -71,9 +69,6 @@ impl McpResource for GraphEntitiesResource {
} else { } else {
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
} }
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -90,14 +85,18 @@ impl McpResource for GraphRelationsResource {
Some("All relationships between entities currently stored in the knowledge graph") Some("All relationships between entities currently stored in the knowledge graph")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let graph = state.graph.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let graph = state_clone.graph.cache.read().unwrap();
let data = &graph.relations; let data = &graph.relations;
if data.len() > 200 {
let truncated: Vec<_> = data.iter().take(200).cloned().collect();
let res = serde_json::json!({
"_meta": format!("Showing 200 of {} total relations. Use query_graph_path or get_subgraph for targeted relation discovery.", data.len()),
"relations": truncated
});
Ok(serde_json::to_string_pretty(&res)?)
} else {
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
}) }
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -114,17 +113,12 @@ impl McpResource for TasksActiveResource {
Some("List of currently active tasks") Some("List of currently active tasks")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let tasks = state.project.tasks.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let tasks = state_clone.project.tasks.cache.read().unwrap();
let data: Vec<_> = tasks let data: Vec<_> = tasks
.iter() .iter()
.filter(|t| t.status != "completed" && t.status != "done") .filter(|t| t.is_active())
.collect(); .collect();
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -141,14 +135,9 @@ impl McpResource for DecisionsActiveResource {
Some("All accepted Architectural Decision Records (ADRs)") Some("All accepted Architectural Decision Records (ADRs)")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let adrs = state.code.adrs.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let adrs = state_clone.code.adrs.cache.read().unwrap();
let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect(); let data: Vec<_> = adrs.iter().filter(|a| a.status == "accepted").collect();
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -165,14 +154,9 @@ impl McpResource for TechDebtActiveResource {
Some("All currently unresolved technical debt items") Some("All currently unresolved technical debt items")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let debts = state.code.tech_debts.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let debts = state_clone.code.tech_debts.cache.read().unwrap();
let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect(); let data: Vec<_> = debts.iter().filter(|d| !d.is_resolved).collect();
Ok(serde_json::to_string_pretty(&data)?) Ok(serde_json::to_string_pretty(&data)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -245,19 +229,17 @@ impl McpResource for SessionDeltaResource {
Some("Recent changes, active tasks, code edits, and notes created in the current session") Some("Recent changes, active tasks, code edits, and notes created in the current session")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state);
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let now = crate::handlers::utils::now_secs(); let now = crate::handlers::utils::now_secs();
let lookback = 7200; // last 2 hours let lookback = 7200; // last 2 hours
let cutoff = now.saturating_sub(lookback); let cutoff = now.saturating_sub(lookback);
let changes = state_clone.code.ledger.cache.read().unwrap(); let changes = state.code.ledger.cache.read().unwrap_or_else(|e| e.into_inner());
let recent_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect(); let recent_changes: Vec<_> = changes.iter().filter(|c| c.timestamp >= cutoff).collect();
let tasks = state_clone.project.tasks.cache.read().unwrap(); let tasks = state.project.tasks.cache.read().unwrap_or_else(|e| e.into_inner());
let active_tasks: Vec<_> = tasks.iter().filter(|t| t.status != "completed" && t.status != "done").collect(); let active_tasks: Vec<_> = tasks.iter().filter(|t| t.is_active()).collect();
let notes = state_clone.code.sticky.cache.read().unwrap(); let notes = state.code.sticky.cache.read().unwrap_or_else(|e| e.into_inner());
let active_notes: Vec<_> = notes.iter().filter(|n| n.expires_at.map_or(true, |exp| exp > now)).collect(); let active_notes: Vec<_> = notes.iter().filter(|n| n.expires_at.map_or(true, |exp| exp > now)).collect();
let delta = serde_json::json!({ let delta = serde_json::json!({
@@ -268,9 +250,6 @@ impl McpResource for SessionDeltaResource {
}); });
Ok(serde_json::to_string_pretty(&delta)?) Ok(serde_json::to_string_pretty(&delta)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -355,13 +334,8 @@ impl MemoryHandler {
Some("Recent terminal execution history and exit codes") Some("Recent terminal execution history and exit codes")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let items = state.telemetry.terminal_history.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let items = state_clone.telemetry.terminal_history.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
struct PinnedFilesResource; struct PinnedFilesResource;
@@ -377,13 +351,8 @@ impl MemoryHandler {
Some("Currently pinned files in the workspace") Some("Currently pinned files in the workspace")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let items = state.project.pinned_files.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let items = state_clone.project.pinned_files.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -400,13 +369,8 @@ impl MemoryHandler {
Some("Project milestones and their status") Some("Project milestones and their status")
} }
async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> { async fn read(&self, state: Arc<MemoryState>) -> crate::error::Result<String> {
let state_clone = Arc::clone(&state); let items = state.project.milestones.cache.read().unwrap_or_else(|e| e.into_inner());
tokio::task::spawn_blocking(move || -> crate::error::Result<String> {
let items = state_clone.project.milestones.cache.read().unwrap();
Ok(serde_json::to_string_pretty(&*items)?) Ok(serde_json::to_string_pretty(&*items)?)
})
.await
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?
} }
} }
@@ -905,7 +869,8 @@ pub fn format_tool_activity_description(name: &str, args: &serde_json::Value) ->
let mut truncated = detail.trim().to_string(); let mut truncated = detail.trim().to_string();
if truncated.len() > 80 { if truncated.len() > 80 {
truncated.truncate(77); let valid_len = truncated.floor_char_boundary(77);
truncated.truncate(valid_len);
truncated.push_str("..."); truncated.push_str("...");
} }
+60 -25
View File
@@ -3,6 +3,15 @@ use std::sync::{Arc, Mutex};
use tantivy::schema::*; use tantivy::schema::*;
use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc}; use tantivy::{Index, IndexReader, IndexWriter, ReloadPolicy, doc};
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct SearchResult {
pub id: String,
pub doc_type: String,
pub title: String,
pub body: String,
pub score: f32,
}
pub type SearchResultTuple = (String, String, String, String, f32); pub type SearchResultTuple = (String, String, String, String, f32);
#[derive(Clone)] #[derive(Clone)]
@@ -89,11 +98,18 @@ impl MemoryIndex {
} }
pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> { pub fn index_entity(&self, e: &Entity) -> tokio::task::JoinHandle<tantivy::Result<()>> {
self.index_entities_batch(std::slice::from_ref(e))
}
pub fn index_entities_batch(&self, entities: &[Entity]) -> tokio::task::JoinHandle<tantivy::Result<()>> {
let writer = Arc::clone(&self.writer); let writer = Arc::clone(&self.writer);
let id_field = self.id_field; let id_field = self.id_field;
let id_val = e.name.clone();
let needs_commit = Arc::clone(&self.needs_commit); let needs_commit = Arc::clone(&self.needs_commit);
let docs: Vec<(String, tantivy::TantivyDocument)> = entities
.iter()
.map(|e| {
let id_val = e.name.clone();
let doc = doc!( let doc = doc!(
self.id_field => e.name.as_str(), self.id_field => e.name.as_str(),
self.title_field => e.name.as_str(), self.title_field => e.name.as_str(),
@@ -101,11 +117,16 @@ impl MemoryIndex {
self.type_field => "entity", self.type_field => "entity",
self.namespace_field => e.namespace.as_str() self.namespace_field => e.namespace.as_str()
); );
(id_val, doc)
})
.collect();
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
let writer = writer.lock().unwrap_or_else(|e| e.into_inner()); let writer_guard = writer.lock().unwrap_or_else(|e| e.into_inner());
writer.delete_term(tantivy::Term::from_field_text(id_field, &id_val)); for (id_val, doc) in docs {
writer.add_document(doc)?; writer_guard.delete_term(tantivy::Term::from_field_text(id_field, &id_val));
writer_guard.add_document(doc)?;
}
needs_commit.store(true, std::sync::atomic::Ordering::SeqCst); needs_commit.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(()) Ok(())
}) })
@@ -172,6 +193,15 @@ impl MemoryIndex {
} }
pub fn commit_if_needed(&self) -> tantivy::Result<()> {
if self.needs_commit.swap(false, std::sync::atomic::Ordering::SeqCst) {
let mut writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
writer.commit()?;
let _ = self.reader.reload();
}
Ok(())
}
pub fn search( pub fn search(
&self, &self,
query: &str, query: &str,
@@ -182,13 +212,28 @@ impl MemoryIndex {
&self.index, &self.index,
vec![self.title_field, self.body_field], vec![self.title_field, self.body_field],
); );
let q = query_parser.parse_query(query)?; let q = query_parser.parse_query(query).or_else(|_| {
let sanitized: String = query
.chars()
.map(|c| match c {
'+' | '-' | '&' | '|' | '!' | '(' | ')' | '{' | '}' | '[' | ']' | '^'
| '"' | '~' | '*' | '?' | ':' | '\\' | '/' => ' ',
_ => c,
})
.collect();
let safe_query = sanitized.trim();
if safe_query.is_empty() {
query_parser.parse_query("*")
} else {
query_parser.parse_query(safe_query)
}
})?;
let top_docs = searcher.search( let top_docs = searcher.search(
&q, &q,
&tantivy::collector::TopDocs::with_limit(50).order_by_score(), &tantivy::collector::TopDocs::with_limit(50).order_by_score(),
)?; )?;
let mut results = Vec::new(); let mut results = Vec::with_capacity(top_docs.len());
for (score, doc_address) in top_docs { for (score, doc_address) in top_docs {
let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?; let retrieved_doc = searcher.doc::<tantivy::TantivyDocument>(doc_address)?;
let id = retrieved_doc let id = retrieved_doc
@@ -273,7 +318,7 @@ impl MemoryIndex {
} }
pub fn add_entity_sync(&self, e: &Entity) { pub fn add_entity_sync(&self, e: &Entity) {
if let Ok(writer) = self.writer.lock() { let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!( let _ = writer.add_document(doc!(
self.id_field => e.name.as_str(), self.id_field => e.name.as_str(),
self.title_field => e.name.as_str(), self.title_field => e.name.as_str(),
@@ -284,19 +329,16 @@ impl MemoryIndex {
self.needs_commit self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst); .store(true, std::sync::atomic::Ordering::SeqCst);
} }
}
pub fn delete_all(&self) { pub fn delete_all(&self) {
if let Ok(writer) = self.writer.lock() { let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.delete_all_documents(); let _ = writer.delete_all_documents();
self.needs_commit self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst); .store(true, std::sync::atomic::Ordering::SeqCst);
} }
}
pub fn add_task_sync(&self, t: &Task) { pub fn add_task_sync(&self, t: &Task) {
// println!("add_task_sync called for task: {}", t.id); let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
if let Ok(writer) = self.writer.lock() {
let _res = writer.add_document(doc!( let _res = writer.add_document(doc!(
self.id_field => t.id.as_str(), self.id_field => t.id.as_str(),
self.title_field => t.title.as_str(), self.title_field => t.title.as_str(),
@@ -304,17 +346,12 @@ impl MemoryIndex {
self.type_field => "task", self.type_field => "task",
self.namespace_field => "global" self.namespace_field => "global"
)); ));
// println!("Writer add_document returned id/result");
self.needs_commit self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst); .store(true, std::sync::atomic::Ordering::SeqCst);
// println!("Needs_commit set to true in add_task_sync");
} else {
// println!("Failed to acquire writer lock in add_task_sync");
}
} }
pub fn add_snippet_sync(&self, s: &Snippet) { pub fn add_snippet_sync(&self, s: &Snippet) {
if let Ok(writer) = self.writer.lock() { let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!( let _ = writer.add_document(doc!(
self.id_field => s.name.as_str(), self.id_field => s.name.as_str(),
self.title_field => s.name.as_str(), self.title_field => s.name.as_str(),
@@ -325,10 +362,9 @@ impl MemoryIndex {
self.needs_commit self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst); .store(true, std::sync::atomic::Ordering::SeqCst);
} }
}
pub fn add_adr_sync(&self, a: &Adr) { pub fn add_adr_sync(&self, a: &Adr) {
if let Ok(writer) = self.writer.lock() { let writer = self.writer.lock().unwrap_or_else(|e| e.into_inner());
let _ = writer.add_document(doc!( let _ = writer.add_document(doc!(
self.id_field => a.id.as_str(), self.id_field => a.id.as_str(),
self.title_field => a.title.as_str(), self.title_field => a.title.as_str(),
@@ -339,7 +375,6 @@ impl MemoryIndex {
self.needs_commit self.needs_commit
.store(true, std::sync::atomic::Ordering::SeqCst); .store(true, std::sync::atomic::Ordering::SeqCst);
} }
}
} }
#[cfg(test)] #[cfg(test)]
@@ -423,12 +458,12 @@ mod tests {
let temp_dir = TempDir::new().unwrap(); let temp_dir = TempDir::new().unwrap();
let index = MemoryIndex::new(temp_dir.path()).unwrap(); let index = MemoryIndex::new(temp_dir.path()).unwrap();
// Malformed lucene query (unclosed parenthesis) // Malformed lucene query (unclosed parenthesis) - sanitized gracefully
let result = index.search("title: (unclosed", None); let result = index.search("title: (unclosed", None);
assert!(result.is_err()); assert!(result.is_ok());
// Another malformed query (unclosed quote) // Another malformed query (unclosed quote) - sanitized gracefully
let result2 = index.search("title: \"unclosed", None); let result2 = index.search("title: \"unclosed", None);
assert!(result2.is_err()); assert!(result2.is_ok());
} }
} }
+76 -44
View File
@@ -4,7 +4,7 @@ use crate::store::Store;
use crate::vector_db::VectorDB; use crate::vector_db::VectorDB;
use std::collections::HashMap; use std::collections::HashMap;
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::{Arc, RwLock}; use std::sync::Arc;
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] #[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
pub struct GenericEvent { pub struct GenericEvent {
@@ -56,7 +56,7 @@ pub struct MemoryState {
pub ttl_notify: Arc<tokio::sync::Notify>, pub ttl_notify: Arc<tokio::sync::Notify>,
pub condense_notify: Arc<tokio::sync::Notify>, pub condense_notify: Arc<tokio::sync::Notify>,
pub graph: Store<KnowledgeGraph>, pub graph: Store<KnowledgeGraph>,
pub search_index: RwLock<MemoryIndex>, pub search_index: tokio::sync::RwLock<MemoryIndex>,
pub vector_db: tokio::sync::RwLock<Option<VectorDB>>, pub vector_db: tokio::sync::RwLock<Option<VectorDB>>,
pub project: ProjectStores, pub project: ProjectStores,
@@ -70,14 +70,37 @@ pub struct MemoryState {
} }
impl MemoryState { impl MemoryState {
pub fn new_in_memory() -> Self {
Self::new(":memory:")
}
pub fn new(base_dir_str: &str) -> Self { pub fn new(base_dir_str: &str) -> Self {
let is_in_memory = base_dir_str == ":memory:";
let base = std::path::PathBuf::from(base_dir_str); let base = std::path::PathBuf::from(base_dir_str);
if !is_in_memory {
if let Err(e) = std::fs::create_dir_all(&base) { if let Err(e) = std::fs::create_dir_all(&base) {
tracing::error!("Failed to create store directory at {:?}: {}", base, e); tracing::error!("Failed to create store directory at {:?}: {}", base, e);
} }
}
let db = crate::db::init_redb(&base); let db = crate::db::init_redb(&base);
let search_index = if is_in_memory {
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
} else {
match crate::search::MemoryIndex::new(&base) {
Ok(idx) => idx,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
}
}
};
let state = Self { let state = Self {
ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()), ollama: Arc::new(crate::ollama::OllamaClient::new_from_env()),
clipboard_watch_mode: tokio::sync::RwLock::new(false), clipboard_watch_mode: tokio::sync::RwLock::new(false),
@@ -87,17 +110,7 @@ impl MemoryState {
condense_notify: Arc::new(tokio::sync::Notify::new()), condense_notify: Arc::new(tokio::sync::Notify::new()),
graph: Store::new("knowledge_graph_master", db.clone()), graph: Store::new("knowledge_graph_master", db.clone()),
base_dir: base.clone(), base_dir: base.clone(),
search_index: RwLock::new(match crate::search::MemoryIndex::new(&base) { search_index: tokio::sync::RwLock::new(search_index),
Ok(idx) => idx,
Err(e) => {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ =
std::fs::write(&log_path, format!("Failed to create MemoryIndex: {}\n", e));
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex")
}
}),
vector_db: tokio::sync::RwLock::new(None), vector_db: tokio::sync::RwLock::new(None),
project: ProjectStores { project: ProjectStores {
@@ -164,11 +177,8 @@ impl MemoryState {
self.graph.modify(update_fn); self.graph.modify(update_fn);
} }
pub fn get_search_index(&self) -> MemoryIndex { pub async fn get_search_index(&self) -> MemoryIndex {
self.search_index self.search_index.read().await.clone()
.read()
.unwrap_or_else(|e| e.into_inner())
.clone()
} }
pub fn search(self: &Arc<Self>) -> SearchService { pub fn search(self: &Arc<Self>) -> SearchService {
@@ -176,8 +186,18 @@ impl MemoryState {
} }
pub async fn rebuild_index(self: &Arc<Self>) { pub async fn rebuild_index(self: &Arc<Self>) {
let idx = self.get_search_index(); let is_in_memory = self.base_dir.to_str() == Some(":memory:");
idx.delete_all(); let new_idx = if is_in_memory {
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
} else {
match crate::search::MemoryIndex::new(&self.base_dir) {
Ok(idx) => idx,
Err(e) => {
tracing::warn!("Failed to create disk MemoryIndex for rebuild ({}), falling back to RAM", e);
crate::search::MemoryIndex::new_in_ram().expect("Failed to create RAM MemoryIndex for rebuild")
}
}
};
let entities: Vec<_> = self let entities: Vec<_> = self
.graph .graph
@@ -192,9 +212,8 @@ impl MemoryState {
tasks.len() tasks.len()
); );
let idx_clone = idx.clone(); let idx_clone = new_idx.clone();
tokio::task::spawn_blocking(move || { tokio::task::spawn_blocking(move || {
// tracing::info!("spawn_blocking started in rebuild_index");
for e in entities { for e in entities {
idx_clone.add_entity_sync(&e); idx_clone.add_entity_sync(&e);
} }
@@ -213,10 +232,8 @@ impl MemoryState {
tracing::error!("Failed to join tantivy index rebuild thread: {}", e); tracing::error!("Failed to join tantivy index rebuild thread: {}", e);
}); });
let _ = idx.commit().await; let _ = new_idx.commit().await;
if let Ok(mut w) = self.search_index.write() { *self.search_index.write().await = new_idx;
*w = idx;
}
self.index_commit_notify.notify_waiters(); self.index_commit_notify.notify_waiters();
} }
@@ -350,7 +367,7 @@ mod tests {
arc_state.rebuild_index().await; arc_state.rebuild_index().await;
// Check search index initialization // Check search index initialization
let idx = arc_state.search_index.read().unwrap(); let idx = arc_state.search_index.read().await;
// Force reload reader to ensure it sees the commit made by rebuild_index // Force reload reader to ensure it sees the commit made by rebuild_index
idx.reader.reload().unwrap(); idx.reader.reload().unwrap();
// tracing::info!( // tracing::info!(
@@ -464,31 +481,46 @@ impl SearchService {
} }
if !vdb_search { if !vdb_search {
let (mut texts_to_embed, mut metadata) = self.state.code.snippets.read_with(|snips| { let mut cached_items = Vec::new();
let mut texts = Vec::with_capacity(snips.len().min(50)); let mut uncached_texts = Vec::new();
let mut meta = Vec::with_capacity(snips.len().min(50)); let mut uncached_meta = Vec::new();
self.state.code.snippets.read_with(|snips| {
for snippet in snips.iter().take(50) { for snippet in snips.iter().take(50) {
texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code)); let title = snippet.name.clone();
meta.push((snippet.name.clone(), "snippet".to_string(), snippet.description.clone())); let desc = snippet.description.clone();
if let Some(ref emb) = snippet.embedding {
cached_items.push((title, "snippet".to_string(), desc, emb.clone()));
} else {
uncached_texts.push(format!("{} {} {}", snippet.name, snippet.description, snippet.code));
uncached_meta.push((title, "snippet".to_string(), desc));
}
} }
(texts, meta)
}); });
self.state.code.sticky.read_with(|sticky| { self.state.code.sticky.read_with(|sticky| {
for note in sticky.iter().take(50) { for note in sticky.iter().take(50) {
texts_to_embed.push(note.content.clone()); let content_preview = note.content.chars().take(200).collect::<String>();
metadata.push(( uncached_texts.push(note.content.clone());
"StickyNote".to_string(), uncached_meta.push(("StickyNote".to_string(), "sticky".to_string(), content_preview));
"sticky".to_string(),
note.content.chars().take(200).collect::<String>(),
));
} }
}); });
if !texts_to_embed.is_empty() for (title, doc_type, body, emb) in cached_items {
&& let Ok(embeddings) = generate_embeddings_async(texts_to_embed).await let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult {
id: title.clone(),
doc_type,
title,
body,
score: sim,
});
}
if !uncached_texts.is_empty()
&& let Ok(embeddings) = generate_embeddings_async(uncached_texts).await
{ {
for (emb, meta) in embeddings.into_iter().zip(metadata) { for (emb, meta) in embeddings.into_iter().zip(uncached_meta) {
let sim = cosine_similarity(&query_emb, &emb); let sim = cosine_similarity(&query_emb, &emb);
results.push(UnifiedSearchResult { results.push(UnifiedSearchResult {
id: meta.0.clone(), id: meta.0.clone(),
@@ -511,13 +543,13 @@ impl SearchService {
Ok(results) Ok(results)
} }
pub fn keyword_search( pub async fn keyword_search(
&self, &self,
query: &str, query: &str,
filter_namespace: Option<&str>, filter_namespace: Option<&str>,
limit: usize, limit: usize,
) -> crate::error::Result<Vec<UnifiedSearchResult>> { ) -> crate::error::Result<Vec<UnifiedSearchResult>> {
let idx = self.state.get_search_index(); let idx = self.state.get_search_index().await;
let matches = idx let matches = idx
.search(query, filter_namespace) .search(query, filter_namespace)
.map_err(|e| crate::error::AppError::Internal(e.to_string()))?; .map_err(|e| crate::error::AppError::Internal(e.to_string()))?;
+138 -42
View File
@@ -8,7 +8,8 @@ pub const STORE_TABLE: TableDefinition<&str, &[u8]> = TableDefinition::new("stor
struct DbWriteTask { struct DbWriteTask {
key: String, key: String,
data: Vec<u8>, data: Vec<u8>,
flushed: Arc<tokio::sync::Notify>, flushed_notifier: Arc<tokio::sync::Notify>,
oneshot_tx: Option<tokio::sync::oneshot::Sender<()>>,
} }
/// Shared centralized write queue actor that handles all database writes serially with micro-batching. /// Shared centralized write queue actor that handles all database writes serially with micro-batching.
@@ -33,7 +34,7 @@ fn get_or_create_queue(db: Arc<Database>) -> DbWriteQueue {
impl DbWriteQueue { impl DbWriteQueue {
pub fn new(db: Arc<Database>) -> Self { pub fn new(db: Arc<Database>) -> Self {
let (tx, mut rx) = tokio::sync::mpsc::channel::<DbWriteTask>(2048); let (tx, mut rx) = tokio::sync::mpsc::channel::<DbWriteTask>(1024);
tokio::spawn(async move { tokio::spawn(async move {
while let Some(first_task) = rx.recv().await { while let Some(first_task) = rx.recv().await {
@@ -69,7 +70,10 @@ impl DbWriteQueue {
// Event-driven notification to all waiting listeners for this micro-batch // Event-driven notification to all waiting listeners for this micro-batch
for task in batch { for task in batch {
task.flushed.notify_waiters(); if let Some(oneshot) = task.oneshot_tx {
let _ = oneshot.send(());
}
task.flushed_notifier.notify_waiters();
} }
}) })
.await; .await;
@@ -79,14 +83,46 @@ impl DbWriteQueue {
Self { tx } Self { tx }
} }
pub fn push(&self, key: String, data: Vec<u8>, flushed: Arc<tokio::sync::Notify>) { pub fn push(
let task = DbWriteTask { key, data, flushed }; &self,
key: String,
data: Vec<u8>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
let task = DbWriteTask {
key,
data,
flushed_notifier,
oneshot_tx: Some(oneshot_tx),
};
if let Err(e) = self.tx.try_send(task) { if let Err(e) = self.tx.try_send(task) {
let task = e.into_inner(); let key = e.into_inner().key;
let tx = self.tx.clone(); tracing::error!("DbWriteQueue channel full or closed; unable to persist key '{}'", key);
tokio::spawn(async move { None
let _ = tx.send(task).await; } else {
}); Some(oneshot_rx)
}
}
pub async fn push_async(
&self,
key: String,
data: Vec<u8>,
flushed_notifier: Arc<tokio::sync::Notify>,
) -> Option<tokio::sync::oneshot::Receiver<()>> {
let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
let task = DbWriteTask {
key,
data,
flushed_notifier,
oneshot_tx: Some(oneshot_tx),
};
if let Err(e) = self.tx.send(task).await {
tracing::error!("DbWriteQueue channel closed; unable to persist key '{}'", e.0.key);
None
} else {
Some(oneshot_rx)
} }
} }
} }
@@ -96,11 +132,12 @@ pub struct Store<T> {
pub flushed: Arc<tokio::sync::Notify>, pub flushed: Arc<tokio::sync::Notify>,
key: String, key: String,
queue: DbWriteQueue, queue: DbWriteQueue,
is_corrupted: bool,
} }
impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static> Store<T> { impl<T: DeserializeOwned + Default + Serialize + Send + Sync + 'static> Store<T> {
pub fn new(key: &str, db: Arc<Database>) -> Self { pub fn new(key: &str, db: Arc<Database>) -> Self {
let initial_data = Self::load_from_db(key, &db); let (initial_data, is_corrupted) = Self::load_from_db(key, &db);
let cache = Arc::new(RwLock::new(initial_data)); let cache = Arc::new(RwLock::new(initial_data));
let flushed = Arc::new(tokio::sync::Notify::new()); let flushed = Arc::new(tokio::sync::Notify::new());
let queue = get_or_create_queue(db); let queue = get_or_create_queue(db);
@@ -110,20 +147,38 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
flushed, flushed,
key: key.to_string(), key: key.to_string(),
queue, queue,
is_corrupted,
} }
} }
fn load_from_db(key: &str, db: &Database) -> T { fn load_from_db(key: &str, db: &Database) -> (T, bool) {
let Ok(read_txn) = db.begin_read() else { let Ok(read_txn) = db.begin_read() else {
return T::default(); tracing::error!("Failed to begin read transaction for key '{}'", key);
return (T::default(), false);
}; };
if let Ok(table) = read_txn.open_table(STORE_TABLE) match read_txn.open_table(STORE_TABLE) {
&& let Ok(Some(value)) = table.get(key) Ok(table) => match table.get(key) {
&& let Ok(parsed) = serde_json::from_slice::<T>(value.value()) Ok(Some(value)) => match serde_json::from_slice::<T>(value.value()) {
{ Ok(parsed) => (parsed, false),
return parsed; Err(e) => {
tracing::error!(
"CRITICAL: Corrupted data for key '{}' in database: {}. Quarantine mode active: state initialized to empty default without overwriting DB key.",
key, e
);
(T::default(), true)
}
},
Ok(None) => (T::default(), false),
Err(e) => {
tracing::error!("Failed to get key '{}' from store table: {}", key, e);
(T::default(), false)
}
},
Err(e) => {
tracing::error!("Failed to open STORE_TABLE for key '{}': {}", key, e);
(T::default(), false)
}
} }
T::default()
} }
pub fn read_with<F, R>(&self, f: F) -> R pub fn read_with<F, R>(&self, f: F) -> R
@@ -134,15 +189,63 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
f(&lock) f(&lock)
} }
pub fn modify<F: FnOnce(&mut T)>(&self, f: F) { pub fn modify<F: FnOnce(&mut T)>(&self, f: F)
let cloned_data = { where
T: Serialize,
{
if self.is_corrupted {
tracing::error!(
"CRITICAL: Refusing to persist changes for corrupted store key '{}' to prevent data loss.",
self.key
);
return;
}
let serialized_res = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner()); let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
f(&mut lock); f(&mut lock);
lock.clone() serde_json::to_vec(&*lock)
}; };
match serde_json::to_vec(&cloned_data) { match serialized_res {
Ok(data) => self.queue.push(self.key.clone(), data, self.flushed.clone()), Ok(data) => {
if self.queue.push(self.key.clone(), data.clone(), self.flushed.clone()).is_none() {
let queue = self.queue.clone();
let key = self.key.clone();
let flushed = self.flushed.clone();
tokio::spawn(async move {
let _ = queue.push_async(key, data, flushed).await;
});
}
}
Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e),
}
}
pub async fn modify_async<F: FnOnce(&mut T)>(&self, f: F)
where
T: Serialize,
{
if self.is_corrupted {
tracing::error!(
"CRITICAL: Refusing to persist changes for corrupted store key '{}' to prevent data loss.",
self.key
);
return;
}
let serialized_res = {
let mut lock = self.cache.write().unwrap_or_else(|e| e.into_inner());
f(&mut lock);
serde_json::to_vec(&*lock)
};
match serialized_res {
Ok(data) => {
if let Some(rx) = self.queue.push_async(self.key.clone(), data, self.flushed.clone()).await {
let _ = rx.await;
}
}
Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e), Err(e) => tracing::error!("Failed to serialize memory store for key '{}': {}", self.key, e),
} }
} }
@@ -151,7 +254,6 @@ impl<T: DeserializeOwned + Default + Serialize + Clone + Send + Sync + 'static>
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use tempfile::NamedTempFile;
#[derive(Serialize, serde::Deserialize, Clone, Default, PartialEq, Debug)] #[derive(Serialize, serde::Deserialize, Clone, Default, PartialEq, Debug)]
struct TestData { struct TestData {
@@ -159,18 +261,21 @@ mod tests {
value: i32, value: i32,
} }
#[tokio::test] fn create_in_memory_test_db() -> Arc<Database> {
async fn test_store_read_write() { let db = Database::builder()
let temp_file = NamedTempFile::new().unwrap(); .create_with_backend(redb::backends::InMemoryBackend::new())
let db = Database::create(temp_file.path()).unwrap(); .unwrap();
let write_txn = db.begin_write().unwrap(); let write_txn = db.begin_write().unwrap();
{ {
write_txn.open_table(STORE_TABLE).unwrap(); write_txn.open_table(STORE_TABLE).unwrap();
} }
write_txn.commit().unwrap(); write_txn.commit().unwrap();
Arc::new(db)
}
let db = Arc::new(db); #[tokio::test]
async fn test_store_read_write() {
let db = create_in_memory_test_db();
let store = Store::<TestData>::new("test_key", db.clone()); let store = Store::<TestData>::new("test_key", db.clone());
assert_eq!(store.read_with(|s| s.clone()), TestData::default()); assert_eq!(store.read_with(|s| s.clone()), TestData::default());
@@ -195,16 +300,7 @@ mod tests {
#[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_store_concurrency() { async fn test_store_concurrency() {
let temp_file = NamedTempFile::new().unwrap(); let db = create_in_memory_test_db();
let db = Database::create(temp_file.path()).unwrap();
let write_txn = db.begin_write().unwrap();
{
write_txn.open_table(STORE_TABLE).unwrap();
}
write_txn.commit().unwrap();
let db = Arc::new(db);
let store = Arc::new(Store::<TestData>::new("concurrent_key", db.clone())); let store = Arc::new(Store::<TestData>::new("concurrent_key", db.clone()));
let mut handles = vec![]; let mut handles = vec![];
+13 -28
View File
@@ -25,10 +25,11 @@ pub struct CreateEntitiesTool {
pub entities: Vec<crate::models::Entity>, pub entities: Vec<crate::models::Entity>,
} }
/// Create new relations between entities in the knowledge graph. Note: relation_type MUST ALWAYS be snake_case (e.g. depends_on, calls, implements). /// Create new relations between entities in the knowledge graph.
/// CRITICAL SCHEMA REQUIREMENT: Each relation item MUST strictly contain fields 'from' (source entity name string), 'to' (target entity name string), and 'relation_type' (snake_case string, e.g., 'depends_on', 'calls', 'implements'). Do NOT use 'source', 'target', or 'relationType'.
#[derive(Debug, Deserialize, Serialize, JsonSchema)] #[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct CreateRelationsTool { pub struct CreateRelationsTool {
/// Array of relations to create. /// Array of relations to create. Each relation item MUST use 'from', 'to', and 'relation_type'.
pub relations: Vec<crate::models::Relation>, pub relations: Vec<crate::models::Relation>,
} }
@@ -81,7 +82,8 @@ pub struct ReadGraphTool {
pub max_tokens: Option<usize>, pub max_tokens: Option<usize>,
} }
/// Search for entities in the knowledge graph by name or type. /// Search specifically for Knowledge Graph entities and nodes by name or type.
/// Note: For searching across tasks, snippets, ADRs, error fixes, and graph entities simultaneously, use 'omni_search' instead.
#[derive(Debug, Deserialize, Serialize, JsonSchema)] #[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct SearchNodesTool { pub struct SearchNodesTool {
/// The search query. /// The search query.
@@ -134,9 +136,13 @@ pub struct VisualizeGraphTool {
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)] #[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
pub enum StickyNoteAction { pub enum StickyNoteAction {
#[serde(alias = "ADD", alias = "Add", alias = "create", alias = "insert")]
Add, Add,
#[serde(alias = "READ", alias = "Read", alias = "list", alias = "get")]
Read, Read,
#[serde(alias = "DELETE", alias = "Delete", alias = "remove", alias = "rm")]
Delete, Delete,
#[serde(alias = "CLEAR", alias = "Clear", alias = "purge")]
Clear, Clear,
} }
@@ -161,31 +167,6 @@ pub struct StickyNotesTool {
pub session_only: Option<bool>, pub session_only: Option<bool>,
} }
/// Add a sticky note for unstructured thoughts or reminders.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct AddStickyNoteTool {
/// The content of the sticky note.
pub content: String,
/// Optional time-to-live in seconds. Automatically expires the note after specified duration.
pub ttl_seconds: Option<u64>,
/// Optional flag to auto-expire the note at the end of the current session.
pub session_only: Option<bool>,
}
/// Read all active sticky notes.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ReadStickyNotesTool {}
/// Delete a specific sticky note by its 1-indexed position.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct DeleteStickyNoteTool {
/// The 1-indexed position of the sticky note to delete.
pub index: usize,
}
/// Clear all sticky notes.
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
pub struct ClearStickyNotesTool {}
/// Condense or summarize an entity's observations to reduce size. /// Condense or summarize an entity's observations to reduce size.
#[derive(Debug, Deserialize, Serialize, JsonSchema)] #[derive(Debug, Deserialize, Serialize, JsonSchema)]
@@ -266,8 +247,11 @@ pub struct TagSnippetTool {
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)] #[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
pub enum SnippetSearchMode { pub enum SnippetSearchMode {
#[serde(alias = "HYBRID", alias = "Hybrid")]
Hybrid, Hybrid,
#[serde(alias = "KEYWORD", alias = "Keyword")]
Keyword, Keyword,
#[serde(alias = "SEMANTIC", alias = "Semantic")]
Semantic, Semantic,
} }
@@ -374,6 +358,7 @@ pub struct LogDecisionTool {
/// The decision made. /// The decision made.
pub decision: String, pub decision: String,
/// The consequence of the decision. /// The consequence of the decision.
#[serde(alias = "consequences")]
pub consequence: String, pub consequence: String,
/// The ID of an older ADR this decision supersedes, if any. /// The ID of an older ADR this decision supersedes, if any.
pub supersedes: Option<String>, pub supersedes: Option<String>,
+39
View File
@@ -14,6 +14,11 @@ pub struct VectorDB {
collection_name: String, collection_name: String,
} }
/// Vector search result containing semantic match metadata.
/// Note: Similarity scores range from 0.0 to 1.0 (Cosine Distance):
/// - Score >= 0.75: High-confidence semantic match
/// - Score 0.50 - 0.74: Moderate semantic relevance
/// - Score < 0.40: Low-confidence noise
#[derive(Debug, Serialize, Deserialize)] #[derive(Debug, Serialize, Deserialize)]
pub struct VectorSearchResult { pub struct VectorSearchResult {
pub id: String, pub id: String,
@@ -98,6 +103,40 @@ impl VectorDB {
Ok(()) Ok(())
} }
pub async fn index_documents_batch(
&self,
docs: Vec<(&str, &str, &str, Vec<f32>)>,
) -> Result<(), Box<dyn Error + Send + Sync>> {
if docs.is_empty() {
return Ok(());
}
let mut points = Vec::with_capacity(docs.len());
for (id, doc_type, text, vector) in docs {
let point_id = match Uuid::parse_str(id) {
Ok(u) => u.to_string(),
Err(_) => {
let namespace = Uuid::NAMESPACE_OID;
Uuid::new_v5(&namespace, id.as_bytes()).to_string()
}
};
let mut payload: std::collections::HashMap<String, serde_json::Value> =
std::collections::HashMap::new();
payload.insert("doc_type".to_string(), serde_json::Value::String(doc_type.to_string()));
payload.insert("text".to_string(), serde_json::Value::String(text.to_string()));
payload.insert("original_id".to_string(), serde_json::Value::String(id.to_string()));
points.push(PointStruct::new(point_id, vector, payload));
}
self.client
.upsert_points(UpsertPointsBuilder::new(&self.collection_name, points))
.await?;
Ok(())
}
pub async fn search( pub async fn search(
&self, &self,
query_vector: Vec<f32>, query_vector: Vec<f32>,
+21 -8
View File
@@ -4,16 +4,18 @@ use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
use tracing::{error, info}; use tracing::{error, info};
pub fn spawn_watcher(_state: Arc<MemoryState>) { pub fn spawn_watcher(state: Arc<MemoryState>) {
let watch_path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")); let watch_path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
info!("Spawning proactive daemon watcher on {:?}", watch_path); info!("Spawning proactive daemon watcher on {:?}", watch_path);
tokio::spawn(async move { tokio::spawn(async move {
let (tx, mut rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(100); let (tx, mut rx) = tokio::sync::mpsc::channel::<notify::Result<Event>>(500);
let mut watcher = match RecommendedWatcher::new( let mut watcher = match RecommendedWatcher::new(
move |res| { move |res| {
let _ = tx.try_send(res); if let Err(e) = tx.try_send(res) {
tracing::warn!("Watcher event dropped due to channel backpressure: {}", e);
}
}, },
Config::default(), Config::default(),
) { ) {
@@ -29,14 +31,26 @@ pub fn spawn_watcher(_state: Arc<MemoryState>) {
return; return;
} }
let mut last_processed: std::collections::HashMap<std::path::PathBuf, std::time::Instant> =
std::collections::HashMap::new();
while let Some(res) = rx.recv().await { while let Some(res) = rx.recv().await {
match res { match res {
Ok(event) => { Ok(event) => {
if event.kind.is_modify() { if event.kind.is_modify() {
let now = std::time::Instant::now();
for path in event.paths { for path in event.paths {
if should_review(&path) { if should_review(&path) {
// 250ms debouncing window per file path
if let Some(last) = last_processed.get(&path) {
if now.duration_since(*last) < std::time::Duration::from_millis(250) {
continue;
}
}
last_processed.insert(path.clone(), now);
info!("Proactive Daemon Hooks: File modified: {:?}", path); info!("Proactive Daemon Hooks: File modified: {:?}", path);
trigger_autonomous_review(&path).await; trigger_autonomous_review(&path, Arc::clone(&state)).await;
} }
} }
} }
@@ -49,7 +63,7 @@ pub fn spawn_watcher(_state: Arc<MemoryState>) {
fn should_review(path: &Path) -> bool { fn should_review(path: &Path) -> bool {
let path_str = path.to_string_lossy(); let path_str = path.to_string_lossy();
if path_str.contains(".git") || path_str.contains("target") { if path_str.contains(".git") || path_str.contains("target") || path_str.contains(".gemini") || path_str.contains("node_modules") {
return false; return false;
} }
@@ -62,10 +76,9 @@ fn should_review(path: &Path) -> bool {
false false
} }
async fn trigger_autonomous_review(path: &Path) { async fn trigger_autonomous_review(path: &Path, state: Arc<MemoryState>) {
info!("Triggering autonomous review for {:?}", path); info!("Triggering autonomous review for {:?}", path);
// Simulate invoking a Neovim Lua diagnostic via mcp-memory-win-nvim daemon state.broadcast_activity("AUTONOMOUS", &format!("Modified: {:?}", path.file_name().unwrap_or_default()));
// MVP: Just write an info log if Neovim isn't reachable.
info!("Autonomous review complete for {:?}", path); info!("Autonomous review complete for {:?}", path);
} }
-823
View File
@@ -1,823 +0,0 @@
mod handlers;
mod mcp;
mod models;
mod search;
mod state;
mod store;
mod tools;
use crate::handlers::MemoryHandler;
use crate::models::*;
use crate::state::MemoryState;
use crate::store::Store;
use std::fs;
use std::path::PathBuf;
use std::sync::{Arc, RwLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use redb::ReadableTable;
use tokio::time::sleep;
use clap::{Parser, Subcommand};
use std::collections::HashMap;
#[derive(Parser)]
#[command(author, version, about = "Antigravity MCP Memory Server", long_about = None)]
struct Cli {
#[command(subcommand)]
command: Option<Commands>,
/// Target URL for the proxy to connect to (e.g., http://127.0.0.1:3000)
#[arg(long)]
target: Option<String>,
/// Run the server as a background daemon process (Windows only)
#[arg(long)]
daemon: bool,
/// Send a shutdown request to the currently running server
#[arg(long)]
exit: bool,
/// Send a shutdown request to the existing server and wait for it to exit
#[arg(long)]
restart: bool,
}
#[derive(Subcommand)]
enum Commands {
/// Manage authorization gates and verification for actions
Gate {
#[command(subcommand)]
subcmd: GateCommands,
},
}
#[derive(Subcommand)]
enum GateCommands {
Set {
#[arg(long)]
action: String,
#[arg(long)]
target: String,
#[arg(long)]
namespace: Option<String>,
#[arg(short = 'p', long = "param")]
params: Vec<String>,
#[arg(long, conflicts_with = "block")]
authorize: bool,
#[arg(long, conflicts_with = "authorize")]
block: bool,
#[arg(long)]
reason: Option<String>,
},
Verify {
#[arg(long)]
action: String,
#[arg(long)]
target: String,
#[arg(long)]
namespace: Option<String>,
#[arg(short = 'p', long = "param")]
params: Vec<String>,
#[arg(long)]
consume: bool,
},
}
async fn reconcile_worker(state: Arc<MemoryState>) {
loop {
sleep(Duration::from_secs(5)).await;
let pattern = format!("{}/delta_*.json", state.base_dir.display());
let has_local = {
let session = state.session_graph.read().unwrap();
!session.entities.is_empty() || !session.relations.is_empty()
};
let has_files = glob::glob(&pattern).map(|p| p.count() > 0).unwrap_or(false);
if has_local || has_files {
state.apply_sync_write(|_master| {}).await;
let state_clone = state.clone();
let _ = tokio::task::spawn_blocking(move || {
state_clone.rebuild_index();
}).await;
}
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
state.ledger.modify(|ledger| {
let seven_days = now.saturating_sub(7 * 24 * 60 * 60);
ledger.retain(|c| c.timestamp >= seven_days);
if ledger.len() > 1000 {
let excess = ledger.len() - 1000;
ledger.drain(0..excess);
}
});
state.sticky.modify(|notes| {
notes.retain(|note| note.timestamp >= now.saturating_sub(24 * 60 * 60));
});
}
}
use axum::{
Json, Router,
extract::{Query, State, ws::{WebSocketUpgrade, WebSocket, Message}},
response::IntoResponse,
routing::{get, post},
};
use futures_util::{SinkExt, StreamExt};
use std::sync::atomic::{AtomicUsize, Ordering};
use tokio::sync::mpsc;
struct AppState {
handler: Arc<MemoryHandler>,
clients: RwLock<HashMap<String, mpsc::Sender<String>>>,
next_id: AtomicUsize,
}
#[derive(serde::Deserialize)]
struct GateVerifyReq {
action: String,
target: String,
namespace: Option<String>,
#[serde(default)]
params: HashMap<String, String>,
#[serde(default)]
consume: bool,
}
#[derive(serde::Deserialize)]
struct GateSetReq {
action: String,
target: String,
namespace: Option<String>,
#[serde(default)]
params: HashMap<String, String>,
authorize: Option<bool>,
block: Option<bool>,
reason: Option<String>,
}
async fn gate_verify_handler(
State(app_state): State<Arc<AppState>>,
Query(q): Query<GateVerifyReq>,
) -> axum::response::Response {
let mut found = None;
let mut to_remove = None;
app_state.handler.state.gates.modify(|gates| {
if let Some(idx) = gates.iter().position(|g| {
g.action == q.action
&& g.target == q.target
&& g.namespace == q.namespace
&& g.params == q.params
}) {
found = Some(gates[idx].clone());
if q.consume {
to_remove = Some(idx);
}
}
if let Some(idx) = to_remove {
gates.remove(idx);
}
});
match found {
Some(record) => {
if record.status == "authorized" {
(axum::http::StatusCode::OK, "Authorized").into_response()
} else {
let msg = if let Some(r) = record.reason {
format!("Action blocked. Reason: {}", r)
} else {
"Action blocked.".to_string()
};
(axum::http::StatusCode::FORBIDDEN, msg).into_response()
}
}
None => {
(axum::http::StatusCode::NOT_FOUND, "Action not yet authorized (no gate record found).").into_response()
}
}
}
async fn gate_set_handler(
State(app_state): State<Arc<AppState>>,
Json(body): Json<GateSetReq>,
) -> axum::response::Response {
let status = if body.block.unwrap_or(false) {
"blocked".to_string()
} else if body.authorize.unwrap_or(false) {
"authorized".to_string()
} else {
"pending".to_string()
};
let record = GateRecord {
id: uuid::Uuid::new_v4().to_string(),
action: body.action.clone(),
target: body.target.clone(),
namespace: body.namespace.clone(),
params: body.params.clone(),
status,
reason: body.reason.clone(),
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
};
app_state.handler.state.gates.modify(|gates| {
gates.retain(|g| !(g.action == record.action && g.target == record.target));
gates.push(record);
});
(axum::http::StatusCode::OK, "Gate state updated.").into_response()
}
fn run_server(state: Arc<MemoryState>) -> Result<(), Box<dyn std::error::Error>> {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async {
tokio::spawn(reconcile_worker(Arc::clone(&state)));
let app_state = Arc::new(AppState {
handler: Arc::new(MemoryHandler { state: Arc::clone(&state) }),
clients: RwLock::new(HashMap::new()),
next_id: AtomicUsize::new(1),
});
let app = Router::new()
.route("/api/version", get(|| async move {
axum::Json(serde_json::json!({
"version": env!("BUILD_DATE"),
"git_hash": option_env!("GIT_HASH").unwrap_or("unknown")
}))
}))
.route("/ws", get(ws_handler))
.route("/health", get(health_handler))
.route("/gate/verify", get(gate_verify_handler))
.route("/gate/set", post(gate_set_handler))
.route(
"/shutdown",
post(|| async move {
std::thread::spawn(|| {
std::thread::sleep(std::time::Duration::from_millis(100));
std::process::exit(0);
});
"Shutting down..."
}),
)
.route(
"/",
get(|| async move { axum::response::Html(include_str!("dashboard.html")) }),
)
.route("/api/graph", get({
let state_clone = app_state.handler.state.clone();
move || async move {
let graph = state_clone.get_full_graph();
axum::Json(graph)
}
}))
.route("/api/tasks/{id}/complete", post({
let state_clone = app_state.handler.state.clone();
move |axum::extract::Path(id): axum::extract::Path<String>| async move {
state_clone.tasks.modify(|tasks| {
for t in tasks.iter_mut() {
if t.id == id {
t.status = "completed".to_string();
break;
}
}
});
axum::Json(serde_json::json!({"status": "success"}))
}
}))
.route("/api/tasks", get({
let state_clone = app_state.handler.state.clone();
move || async move {
let tasks = state_clone.tasks.read();
axum::Json(tasks.clone())
}
}))
.route("/api/search", get({
let state_clone = app_state.handler.state.clone();
move |axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>| async move {
if let Some(q) = params.get("q") {
if let Ok(idx) = state_clone.search_index.read() {
if let Ok(results) = idx.search(q, None) {
let mut formatted_results = Vec::new();
for (type_name, content) in results {
formatted_results.push(serde_json::json!({
"type_name": type_name,
"content": content,
"score": 1.0
}));
}
return axum::Json(serde_json::json!({ "results": formatted_results }));
}
}
}
axum::Json(serde_json::json!({ "results": [] }))
}
}))
.route("/api/stats",
get({
let state_clone = app_state.handler.state.clone();
move || async move {
let (entities, relations) = {
let graph = state_clone.get_full_graph();
(graph.entities.len(), graph.relations.len())
};
let tasks = state_clone.tasks.read().len();
let snippets = state_clone.snippets.read().len();
let tech_debts = state_clone.tech_debts.read().len();
let adrs = state_clone.adrs.read().len();
let ledger = state_clone.ledger.read().len();
let sticky = state_clone.sticky.read().len();
let error_fixes = state_clone.error_fixes.read().len();
let pinned_files = state_clone.pinned_files.read().len();
let session_summaries = state_clone.session_summaries.read().len();
let handoff_memos = state_clone.handoff_memos.read().len();
let env_fingerprints = state_clone.env_fingerprints.read().len();
let env_requirements = state_clone.env_requirements.read().len();
let milestones = state_clone.milestones.read().len();
let environments = state_clone.environments.read().len();
let pr_checklists = state_clone.pr_checklists.read().len();
let gates = state_clone.gates.read().len();
let context_workspaces = state_clone.context_workspaces.read().len();
axum::Json(serde_json::json!({
"entities": entities,
"relations": relations,
"tasks": tasks,
"snippets": snippets,
"tech_debts": tech_debts,
"adrs": adrs,
"ledger": ledger,
"sticky": sticky,
"error_fixes": error_fixes,
"pinned_files": pinned_files,
"session_summaries": session_summaries,
"handoff_memos": handoff_memos,
"env_fingerprints": env_fingerprints,
"env_requirements": env_requirements,
"milestones": milestones,
"environments": environments,
"pr_checklists": pr_checklists,
"gates": gates,
"context_workspaces": context_workspaces
}))
}
}),
)
.with_state(app_state);
let mut retries = 0;
let listener = loop {
match tokio::net::TcpListener::bind("127.0.0.1:3000").await {
Ok(l) => break l,
Err(e) => {
// Check if it's already running and healthy
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
use std::io::{Read, Write};
let _ = stream.write_all(
b"GET /health HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
);
let mut response = String::new();
let _ = stream.read_to_string(&mut response);
if response.contains("200 OK") {
// Already healthy! Just exit cleanly instead of panicking/retrying loop.
std::process::exit(0);
}
}
retries += 1;
if retries > 15 {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_fatal.log");
let _ = std::fs::write(
&log_path,
format!(
"FATAL: Could not bind to 127.0.0.1:3000 after 15 seconds: {}\n",
e
),
);
std::process::exit(1);
}
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
if let Ok(mut file) = std::fs::OpenOptions::new()
.create(true)
.append(true)
.open(&log_path)
{
use std::io::Write;
let _ = writeln!(
file,
"Failed to bind to 127.0.0.1:3000 (attempt {}): {}. Retrying in 1s...",
retries, e
);
}
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
}
}
};
// Background Garbage Collection for old tasks
let state_gc = Arc::clone(&state);
tokio::spawn(async move {
loop {
// Run every 24 hours
tokio::time::sleep(tokio::time::Duration::from_secs(24 * 3600)).await;
let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs();
let fourteen_days = 14 * 24 * 3600;
let cutoff = now.saturating_sub(fourteen_days);
state_gc.tasks.modify(|tasks| {
let initial_len = tasks.len();
tasks.retain(|task| {
if task.status.to_lowercase() == "completed" && task.created_at < cutoff {
false // remove
} else {
true // keep
}
});
if tasks.len() < initial_len {
eprintln!("GC: Removed {} old completed tasks", initial_len - tasks.len());
}
});
}
});
// Git Native Sync Background Task
let state_git = Arc::clone(&state);
tokio::spawn(async move {
let repo_path = std::env::current_dir().unwrap_or_else(|_| ".".into());
let mut last_commit_id = String::new();
loop {
tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
if let Ok(repo) = git2::Repository::discover(&repo_path) {
if let Ok(head) = repo.head() {
if let Ok(commit) = head.peel_to_commit() {
let current_id = commit.id().to_string();
if current_id != last_commit_id && !last_commit_id.is_empty() {
let msg = commit.message().unwrap_or("").to_string();
let branch = head.shorthand().unwrap_or("unknown").to_string();
state_git.ledger.modify(|changes| {
changes.push(crate::models::CodeChange {
git_commit: Some(current_id.clone()),
git_branch: Some(branch),
description: format!("Auto-synced commit: {}", msg.trim()),
timestamp: std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(),
file_path: "".to_string(),
});
});
eprintln!("Git Sync: Logged new commit {}", current_id);
state_git.tasks.modify(|tasks| {
for task in tasks.iter_mut() {
if task.status != "completed" && msg.to_lowercase().contains(&task.title.to_lowercase()) {
task.status = "completed".to_string();
eprintln!("Git Sync: Auto-completed task '{}'", task.title);
}
}
});
}
last_commit_id = current_id;
}
}
}
}
});
eprintln!("MCP Memory Server running on http://127.0.0.1:3000/sse");
if let Err(e) = axum::serve(listener, app).await {
let log_path = dirs::home_dir()
.unwrap_or_default()
.join(".gemini/mcp_memory/daemon_error.log");
let _ = std::fs::write(&log_path, format!("Server crashed: {}\n", e));
}
Ok(())
})
}
async fn ws_handler(
ws: WebSocketUpgrade,
State(state): State<Arc<AppState>>,
Query(query): Query<std::collections::HashMap<String, String>>,
) -> impl axum::response::IntoResponse {
let client_type = query.get("client").cloned().unwrap_or_else(|| "unknown".to_string());
ws.on_upgrade(move |socket| handle_socket(socket, state, client_type))
}
async fn handle_socket(socket: WebSocket, state: Arc<AppState>, client_type: String) {
let session_id = format!("{}", state.next_id.fetch_add(1, Ordering::SeqCst));
let (tx, mut rx) = mpsc::channel::<String>(100);
state.clients.write().unwrap().insert(session_id.clone(), tx.clone());
let (mut sender, mut receiver) = socket.split();
let mut send_task = tokio::spawn(async move {
while let Some(msg) = rx.recv().await {
if sender.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
});
let handler = Arc::clone(&state.handler);
let state_clone = Arc::clone(&state);
let session_id_clone = session_id.clone();
let mut recv_task = tokio::spawn(async move {
while let Some(Ok(Message::Text(text))) = receiver.next().await {
if let Ok(payload) = serde_json::from_str::<serde_json::Value>(&text) {
if client_type == "proxy" {
// Send activity broadcast to UI clients
if let Some(method) = payload.get("method").and_then(|m| m.as_str()) {
if method == "tools/call" {
let name = payload.get("params").and_then(|p| p.get("name")).and_then(|n| n.as_str()).unwrap_or("unknown_tool");
let activity_msg = format!("Agent executed tool: {}", name);
let event = serde_json::json!({
"type": "activity",
"data": activity_msg
});
let clients_map = state_clone.clients.read().unwrap().clone();
for (id, client_tx) in clients_map.iter() {
if id != &session_id_clone {
let _ = client_tx.send(event.to_string()).await;
}
}
}
}
// Process MCP request
if let Some(response) = handler.handle_request(payload).await {
let res_str = serde_json::to_string(&response).unwrap();
let tx_opt = state_clone.clients.read().unwrap().get(&session_id_clone).cloned();
if let Some(client_tx) = tx_opt {
let _ = client_tx.send(res_str).await;
}
}
}
}
}
});
tokio::select! {
_ = (&mut send_task) => recv_task.abort(),
_ = (&mut recv_task) => send_task.abort(),
};
state.clients.write().unwrap().remove(&session_id);
}
async fn health_handler() -> &'static str {
"OK"
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let cli = Cli::parse();
if cli.exit {
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
use std::io::Write;
let _ = stream.write_all(
b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
);
}
println!("Sent shutdown request to server.");
return Ok(());
}
if cli.restart {
if let Ok(mut stream) = std::net::TcpStream::connect("127.0.0.1:3000") {
use std::io::Write;
let _ = stream.write_all(
b"POST /shutdown HTTP/1.1\r\nHost: 127.0.0.1\r\nConnection: close\r\n\r\n",
);
println!("Sent shutdown request to existing server. Waiting for it to exit...");
std::thread::sleep(std::time::Duration::from_millis(1500));
}
return Ok(());
}
#[cfg(target_os = "windows")]
{
use std::os::windows::process::CommandExt;
if !cli.daemon {
// Just spawn the daemon and exit. We no longer act as a proxy.
#[allow(clippy::zombie_processes)]
let _ = std::process::Command::new(std::env::current_exe().unwrap())
.arg("--daemon")
.stdin(std::process::Stdio::null())
.stdout(std::process::Stdio::null())
.stderr(std::process::Stdio::null())
.creation_flags(0x08000000) // CREATE_NO_WINDOW
.spawn()
.expect("Failed to spawn daemon");
return Ok(());
}
}
#[cfg(not(target_os = "windows"))]
{
// Linux no longer executes server logic natively due to workspace split
return Ok(());
}
let base_dir = std::env::var("MCP_MEMORY_STORE_DIR").unwrap_or_else(|_| {
dirs::home_dir()
.map(|mut h| {
h.push(".gemini/mcp_memory");
h.to_string_lossy().into_owned()
})
.unwrap_or_else(|| ".gemini/mcp_memory".into())
});
let base = PathBuf::from(base_dir);
fs::create_dir_all(&base).expect("Failed to create store dir");
let redb_path = base.join("mcp_store.redb");
let db = Arc::new(redb::Database::create(&redb_path).unwrap());
// Ensure table exists and migrate old JSON files
{
let write_txn = db.begin_write().unwrap();
{
let mut table = write_txn.open_table(crate::store::STORE_TABLE).unwrap();
let stores = [
("audit_ledger", "audit_ledger.json"),
("sticky_notes", "sticky_notes.json"),
("tasks", "tasks.json"),
("snippets", "snippets.json"),
("adrs", "adrs.json"),
("preferences", "preferences.json"),
("error_fixes", "error_fixes.json"),
("pinned_files", "pinned_files.json"),
("session_summaries", "session_summaries.json"),
("handoff_memos", "handoff_memos.json"),
("env_fingerprints", "env_fingerprints.json"),
("env_requirements", "env_requirements.json"),
("milestones", "milestones.json"),
("environments", "environments.json"),
("pr_checklists", "pr_checklists.json"),
("tech_debts", "tech_debts.json"),
("gates", "gates.json"),
("context_workspaces", "context_workspaces.json"),
];
for (key, file_name) in stores.iter() {
if table.get(*key).unwrap().is_none() {
let json_path = base.join(file_name);
if json_path.exists() {
if let Ok(data) = fs::read(&json_path) {
if serde_json::from_slice::<serde_json::Value>(&data).is_ok() {
table.insert(*key, data.as_slice()).unwrap();
}
}
}
}
}
}
write_txn.commit().unwrap();
}
let state = Arc::new(MemoryState {
master_path: base.join("knowledge_graph_master.json"),
session_graph: RwLock::new(KnowledgeGraph::default()),
base_dir: base.clone(),
master_cache: RwLock::new((KnowledgeGraph::default(), SystemTime::UNIX_EPOCH)),
search_index: RwLock::new(crate::search::MemoryIndex::new(&base).unwrap()),
ledger: Store::new("audit_ledger", db.clone()),
sticky: Store::new("sticky_notes", db.clone()),
tasks: Store::new("tasks", db.clone()),
snippets: Store::new("snippets", db.clone()),
adrs: Store::new("adrs", db.clone()),
prefs: Store::new("preferences", db.clone()),
error_fixes: Store::new("error_fixes", db.clone()),
pinned_files: Store::new("pinned_files", db.clone()),
session_summaries: Store::new("session_summaries", db.clone()),
handoff_memos: Store::new("handoff_memos", db.clone()),
env_fingerprints: Store::new("env_fingerprints", db.clone()),
env_requirements: Store::new("env_requirements", db.clone()),
milestones: Store::new("milestones", db.clone()),
environments: Store::new("environments", db.clone()),
pr_checklists: Store::new("pr_checklists", db.clone()),
tech_debts: Store::new("tech_debts", db.clone()),
gates: Store::new("gates", db.clone()),
context_workspaces: Store::new("context_workspaces", db.clone()),
});
state.rebuild_index();
if let Some(command) = cli.command {
match command {
Commands::Gate { subcmd } => match subcmd {
GateCommands::Set {
action,
target,
namespace,
params,
authorize,
block,
reason,
} => {
let status = if authorize {
"authorized".to_string()
} else if block {
"blocked".to_string()
} else {
"pending".to_string()
};
let mut param_map = HashMap::new();
for p in params {
if let Some((k, v)) = p.split_once('=') {
param_map.insert(k.to_string(), v.to_string());
}
}
let record = GateRecord {
id: uuid::Uuid::new_v4().to_string(),
action: action.clone(),
target: target.clone(),
namespace,
params: param_map,
status,
reason,
timestamp: SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs(),
};
state.gates.modify(|gates| {
gates.retain(|g| !(g.action == record.action && g.target == record.target));
gates.push(record);
});
println!("Gate state updated.");
std::process::exit(0);
}
GateCommands::Verify {
action,
target,
namespace,
params,
consume,
} => {
let mut param_map = HashMap::new();
for p in params {
if let Some((k, v)) = p.split_once('=') {
param_map.insert(k.to_string(), v.to_string());
}
}
let mut found = None;
let mut to_remove = None;
state.gates.modify(|gates| {
if let Some(idx) = gates.iter().position(|g| {
g.action == action
&& g.target == target
&& g.namespace == namespace
&& g.params == param_map
}) {
found = Some(gates[idx].clone());
if consume {
to_remove = Some(idx);
}
}
if let Some(idx) = to_remove {
gates.remove(idx);
}
});
match found {
Some(record) => {
if record.status == "authorized" {
std::process::exit(0);
} else {
if let Some(r) = record.reason {
eprintln!("❌ Action blocked. Reason: {}", r);
} else {
eprintln!("❌ Action blocked.");
}
std::process::exit(1);
}
}
None => {
eprintln!("❌ Action not yet authorized (no gate record found).");
std::process::exit(2);
}
}
}
},
}
}
run_server(state)
}
-6
View File
@@ -4,7 +4,6 @@ version = "1.0.0"
edition = "2024" edition = "2024"
[dependencies] [dependencies]
#rustls-tls = "0.2"
clap = { version = "4.6.6", features = ["derive"] } clap = { version = "4.6.6", features = ["derive"] }
reqwest = { version = "0.13", default-features = false, features = ["stream"] } reqwest = { version = "0.13", default-features = false, features = ["stream"] }
tokio = { version = "1.53.1", features = ["full"] } tokio = { version = "1.53.1", features = ["full"] }
@@ -18,8 +17,3 @@ serde_json = "1.0.151"
mcp-stdio = { version = "1.0.0", path = "../mcp-stdio" } mcp-stdio = { version = "1.0.0", path = "../mcp-stdio" }
regex = "1.13.1" regex = "1.13.1"
async-channel = "2.5.0" async-channel = "2.5.0"
[dev-dependencies]
serde_json = "1.0.151"