🧠 agent-brain

Plug-and-pack graph memory for local AI agents backed by embedded Kùzu graph engine.

agent-brain gives local AI coding agents structured, long-term memory of entire codebases. It parses your source code using native Tree-sitter grammars and Python AST into a high-performance Kùzu graph (CompanyBrain), exposes search and multi-hop dependency analysis via the official Model Context Protocol (MCP), and securely encrypts and synchronizes the memory state across environments using Fernet AES encryption.

🚀 Key Features

• Multi-Language Tree-sitter AST (8 Languages Supported): Out-of-the-box parsing for Python, Rust, JavaScript, TypeScript, Go, Scala, SQL, and PySpark using native Tree-sitter grammars (zero regex heuristics) with full syntax-error resilience.
• Lexical Scope & Collision Prevention: Resolves calls within local classes, methods, and explicit file imports, eliminating spiderweb graph collisions on common function names (connect, init, execute).
• Parallel Multi-Core Processing: Concurrent worker threads parse hundreds of files simultaneously across CPU cores, dramatically accelerating ingestion speed.
• Zero-Config Embedded Kùzu Graph Database: Uses Kùzu (kuzu)—a vectorized, columnar C++ embedded graph database—to persist knowledge locally at ./.agent-memory/brain.db with graph CompanyBrain. Runs natively on Windows, macOS, and Linux with zero background daemons or Docker.
• Structural Graph Extractor: Extracts :CodeFile, :Class/Struct/Table, and :Function/Method/Procedure nodes alongside hierarchical :CONTAINS and multi-hop :CALLS relationships with graceful syntax error tolerance.
• Official Model Context Protocol (MCP) Server: Connects directly to AI agent platforms (Cursor, Claude Desktop, Antigravity) over stdio to provide:
  • search_code_graph(symbol_name): Immediate callers/callees and symbol definition metadata.
  • get_dependency_impact(function_name): Multi-hop graph traversal ([:CALLS*1..5]) to evaluate blast radius and downstream dependents.
• Encrypted Synchronization & Hydration: Compresses memory into in-memory zip streams and encrypts with AES (Fernet) to produce memory-seed.enc. Decrypts and auto-hydrates on startup with Zip Slip protection.
• Unified CLI: Single command-line interface (agent-brain, alias agent-memory) with commands for indexing (--workers), exporting, pulling, registering with AI agents, benchmarking, and running the MCP agent server.

🌐 Supported Languages & Parser Backends

🏛️ Architecture & Graph Schema

graph TD
    CF[":CodeFile (path, filename)"]
    CL[":Class (id, name, qualified_name, docstring)"]
    FN[":Function (id, name, qualified_name, is_method, is_async)"]
    CALLEE[":Function / :Class (callee)"]

    CF -->|:CONTAINS| CL
    CF -->|:CONTAINS| FN
    CL -->|:CONTAINS| FN
    FN -->|:CALLS| CALLEE

📦 Installation

1. Clone & Install Dependencies

git clone https://github.com/your-org/agent-brain.git
cd agent-brain

# Install requirements
pip install -r requirements.txt

# Or install as an editable package with the global 'agent-brain' command
pip install -e .

2. Environment Note

• Cross-Platform Out-of-the-Box:
  • Runs directly on native Windows, macOS, and Linux via embedded Kùzu (kuzu>=0.8.0).
  • No Docker, WSL, Redis, or external server daemons required.
  • Vectorized openCypher query engine with Compressed Sparse Row (CSR) storage representation.

Step 1: Initialize Codebase Memory

Scan any codebase directory to extract AST symbols and populate CompanyBrain:

# Using global CLI command:
agent-brain init --path ./core

# Or directly with Python:
python cli.py init --path ./core

Output:

[INFO] Scanning codebase in parallel at: /path/to/project
  * Files indexed:   21
  * Languages:       python: 21
  * Classes/Structs: 29
  * Functions/Defs:  160
  * Function calls:  1531

[INGEST] Ingesting AST nodes & relationships into graph 'CompanyBrain'...
[OK] Graph successfully populated!
  * :CodeFile nodes:       21
  * :Class nodes:          29
  * :Function nodes:       160
  * :CONTAINS edges:       189
  * :CALLS edges:          1164
  * Database file:         ./.agent-memory/brain.db

Step 2: Auto-Connect or Configure AI Agents (Antigravity, Claude, Cursor)

Automated Registration (Recommended):

Connect the memory MCP server to your local AI coding agents with a single command:

# Auto-detect all installed AI agent platforms and register the MCP server:
agent-brain register

# Or register specifically with Antigravity:
agent-brain register --client antigravity

# Or index and connect in a single step:
agent-brain init --path ./my-project --connect antigravity

Manual Configuration (Alternative):

If you prefer adding the configuration manually to your agent's config file:

• Antigravity (~/.gemini/config/mcp_config.json):
• Claude Desktop (%APPDATA%\Claude\claude_desktop_config.json on Windows, ~/Library/Application Support/Claude/... on macOS):
• Cursor (.cursor/mcp.json):

{
  "mcpServers": {
    "agent-brain": {
      "command": "agent-brain",
      "args": ["mcp"],
      "env": {
        "AGENT_MEMORY_DB_PATH": "C:/path/to/project/.agent-memory/brain.db"
      }
    }
  }
}

Step 3: Run the MCP Server

Start the Model Context Protocol (MCP) server over stdio for local AI agents:

agent-brain mcp

# Or start while ensuring target agent registration:
agent-brain mcp --client antigravity

Note: When configured in Antigravity or Claude Desktop, the AI client automatically launches and manages this server process in the background.

Available MCP Tools

Step 4: Export & Encrypt Memory Archive

Compress and encrypt ./.agent-memory/ into memory-seed.enc:

# 1. Generate an encryption key (if not already set)
agent-brain generate-key

# 2. Set the environment variable
export AGENT_MEMORY_SECRET_KEY='your-base64-secret-key'     # Linux / macOS / WSL
$env:AGENT_MEMORY_SECRET_KEY='your-base64-secret-key'      # Windows PowerShell

# 3. Export encrypted archive
agent-brain export

Output:

🔒 Compressing and encrypting './.agent-memory' -> 'memory-seed.enc'...
✅ Memory exported and encrypted successfully!
  • Encrypted archive: /path/to/memory-seed.enc
  • Archive size:      1.42 KB
  • Encryption method: AES-128-CBC / Fernet HMAC-SHA256

Step 5: Pull & Decrypt on Another Machine

Restore and hydrate the database from the encrypted seed:

agent-brain pull

Output:

🔓 Decrypting and restoring 'memory-seed.enc' -> './.agent-memory'...
✅ Database restored successfully!
  • Restored location: /path/to/.agent-memory
  • KùzuDB file:  ./.agent-memory/brain.db (Ready)

🐍 Python SDK Usage

You can also use agent-brain directly in your Python applications and custom agents:

from core.database import Database, execute_query
from extractors.ast_parser import parse_codebase
from sync.s3_sync import hydrate_database_on_startup

# 1. Auto-hydrate database on startup if seed exists
hydrate_database_on_startup()

# 2. Parse and index codebase
summary = parse_codebase("./my_service")

# 3. Query graph with openCypher
with Database() as db:
    result = db.execute_query(
        "MATCH (fn:Function)-[:CALLS]->(target:Function {name: 'connect'}) RETURN fn.name"
    )
    for record in result.result_set:
        print(f"Caller: {record[0]}")

📊 Benchmark: Traditional RAG vs. Kùzu Graph

The benchmark automatically detects if an AI agent MCP client is configured:

• When an MCP client is found: Compares Traditional Raw-File RAG vs. Kùzu Graph Traversal on the target codebase.
• When no MCP client is found: Benchmarks multi-language AST file indexing latency, parsed symbol counts, and token weight.

# Auto-detects configured agents and benchmarks the default './core' directory:
python benchmark.py
# Or via CLI:
agent-brain benchmark

# Benchmark any custom codebase folder and function symbol:
agent-brain benchmark --path "C:/path/to/custom/codebase" --symbol "my_function"

Results (Tested on Real-World Codebase x-algo, 72 Files)

🧪 Test Suite

The project includes an automated test suite with 50 passing unit tests executing in ~1.5 seconds with 100% test isolation:

python -m unittest discover tests

Test Coverage Summary:

• [tests/test_ast_parser.py]: Tree-sitter AST extraction across 7 languages, Python AST visitor, scope qualification, call graph resolution, and syntax error recovery.
• [tests/test_database.py]: Kùzu initialization, strongly-typed schema DDL, in-memory execution, multi-threaded concurrency lock, query error handling, and clean teardown.
• [tests/test_mcp_server.py]: MCP server registration, caller/callee lookup via search_code_graph, and multi-hop impact analysis via get_dependency_impact.
• [tests/test_integrations.py]: Cross-platform auto-registration for Antigravity, Claude Desktop, and Cursor; safe JSON atomic updates.
• [tests/test_benchmark.py]: Indexing throughput, 3-tier comparative benchmarking (Raw-file vs Semantic chunk vs Graph), and deduplication metrics.
• [tests/test_antigravity_token_tracker.py]: Ground-truth token tracking across Antigravity IDE and Desktop environments, tool telemetry, and cache hit metrics.
• [tests/test_s3_sync.py]: In-memory streaming zip, Fernet AES roundtrip encryption, Zip Slip path traversal guard, and startup hydration.
• [tests/test_cli.py]: Click CLI commands (init, register, export, pull, benchmark, mcp, generate-key).

⚖️ Engineering Design Decisions & Trade-offs

📂 Project Structure

agent-brain/
├── cli.py                     # Unified Click CLI (init, register, benchmark, export, pull, mcp)
├── benchmark.py               # Token efficiency & indexing performance benchmark
├── core/
│   ├── __init__.py
│   ├── database.py            # Embedded Kùzu graph manager
│   └── integrations.py        # Antigravity, Claude, & Cursor auto-registration
├── extractors/
│   ├── __init__.py
│   └── ast_parser.py          # Parallel multi-language AST parser (8 languages)
├── tools/
│   ├── __init__.py
│   └── mcp_server.py          # Official MCP SDK server (search_code_graph, get_dependency_impact)
├── sync/
│   ├── __init__.py
│   └── s3_sync.py             # AES Fernet encryption, in-memory zip, startup hydration
├── tests/
│   ├── test_cli.py
│   ├── test_database.py
│   ├── test_ast_parser.py
│   ├── test_mcp_server.py
│   ├── test_s3_sync.py
│   ├── test_integrations.py
│   └── test_benchmark.py
├── pyproject.toml
├── requirements.txt
└── README.md