🧠 pi-tokensaver

April 17, 2026 Β· View on GitHub

🧠 pi-tokensaver

Give your AI coding agent a semantic memory for your codebase

npm version license pi-coding-agent

A Pi coding agent extension that bridges tokensave β€” a Rust-powered local semantic graph engine β€” directly into your AI pair-programming workflow.

Why? β€’ How it works β€’ Install β€’ Usage β€’ Architecture β€’ Configuration


πŸ€” Why?

When you ask an AI agent to explore a codebase, it typically:

  • πŸ” Runs grep, glob, find β€” scanning thousands of lines
  • πŸ“„ Reads entire files just to find one function
  • πŸ’Έ Burns through your token budget on redundant file reads

tokensave fixes this by building a local semantic graph of your code β€” functions, types, call relationships, imports β€” and exposing it as an MCP server with query tools. pi-tokensaver bridges that server into the Pi agent, so your AI gets precision code exploration instead of brute-force file scanning.

Result: Dramatically fewer tokens consumed per codebase query. Your agent finds what it needs in one semantic lookup instead of reading 20 files.


⚑ How it works

β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”         β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
β”‚                      β”‚         β”‚                      β”‚
β”‚   Pi Coding Agent    β”‚  calls  β”‚   pi-tokensaver      β”‚
β”‚   (LLM)              β”œβ”€β”€β”€β”€β”€β”€β”€β”€β–Ίβ”‚   (this extension)   β”‚
β”‚                      β”‚         β”‚                      β”‚
β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜         β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                            β”‚ JSON-RPC over stdio
                                            β–Ό
                                 β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                                 β”‚                      β”‚
                                 β”‚   tokensave serve    β”‚
                                 β”‚   (MCP server)       β”‚
                                 β”‚   queries semantic   β”‚
                                 β”‚   graph locally      β”‚
                                 β”‚                      β”‚
                                 β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜
                                            β”‚ reads
                                            β–Ό
                                 β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”
                                 β”‚   .tokensave/        β”‚
                                 β”‚   (local graph DB)   β”‚
                                 β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜

The extension manages the full lifecycle automatically:

PhaseWhat happens
1. InitRuns tokensave sync to build/update the semantic graph. Ensures .tokensave/ is gitignored.
2. RegistrationSpawns tokensave serve, performs MCP handshake, discovers all available tools, and bridges each one as a Pi tool prefixed with tokensave_.
3. Prompt InjectionInjects a system prompt rule telling the LLM to prefer tokensave tools over raw grep/glob/read.
4. TeardownGracefully kills the tokensave serve process when the Pi session ends.

πŸ“¦ Install

Prerequisites

  1. Pi coding agent β€” the AI agent framework.
  2. tokensave β€” the Rust semantic graph engine:
    cargo install tokensave
    

Install the extension

One command β€” installs globally, works across all your projects automatically:

pi install npm:pi-tokensaver

That's it. Pi auto-discovers the extension from the global install. No per-project config needed.

Uninstall: pi remove npm:pi-tokensaver


πŸš€ Usage

Once installed, it Just Worksβ„’ β€” no manual setup required.

  1. Start a Pi session in any project directory:

    pi
    
  2. The extension automatically:

    • Detects and syncs the semantic graph for your project
    • Spawns the MCP server in the background
    • Registers semantic exploration tools (e.g. tokensave_search, tokensave_related)
    • Tells the LLM to prefer these tools over brute-force file reads
  3. Ask your agent anything β€” it will use the semantic graph for efficient exploration:

    "Where is the authentication middleware applied?" "Which functions call processPayment?" "Show me the type hierarchy for UserEvent"

What if tokensave isn't installed?

The extension fails gracefully β€” you'll see a notification with install instructions, and the Pi session continues normally without tokensave tools.


πŸ—οΈ Architecture

MCP Client

The extension implements a lightweight JSON-RPC 2.0 client over stdio that speaks the Model Context Protocol. It handles:

  • Process lifecycle (spawn, heartbeat, graceful shutdown with SIGTERMβ†’SIGKILL fallback)
  • Message framing (newline-delimited JSON)
  • Request/response correlation with pending promise tracking
  • Clean error propagation from MCP server errors

Tool Bridge

Each MCP tool discovered from tokensave serve is converted into a Pi custom tool:

  • Schema conversion: MCP JSON Schema β†’ TypeBox schema for Pi parameter validation
  • Call forwarding: Pi tool calls are proxied to the MCP server via JSON-RPC
  • Result mapping: MCP response content is flattened into Pi's text-based result format

Type Safety

Written in strict TypeScript with full type coverage for:

  • JSON-RPC message types
  • MCP tool definitions and responses
  • Pi ExtensionAPI surface
  • TypeBox schema construction

βš™οΈ Configuration

No configuration needed β€” the extension auto-detects everything from your project root.

Build from source

git clone https://github.com/xilnick/pi-tokensaver.git
cd pi-tokensaver
npm install
npm run build

πŸ“„ License

MIT Β© contributors


Made with 🧠 for Pi + tokensave

Report a Bug Β· Request a Feature Β· Contribute