feat(tokenless): add MCP tokenless_retrieve server#1285
Conversation
30e7d49 to
bec0bc8
Compare
Forrest-ly
left a comment
There was a problem hiding this comment.
Code Review: feat(tokenless): add MCP tokenless_retrieve server
This PR adds a minimal MCP stdio server (tokenless mcp serve) that exposes tokenless_retrieve so an MCP-connected agent can recover stashed payloads on demand. The implementation is hand-rolled JSON-RPC (no SDK dependency), deliberately minimal in scope (one tool),
and reuses the existing open_stash_store_or_err + extract_hash infrastructure from #1270. The test coverage (6 tests) covers tool listing, dispatch errors, hash validation, round-trip retrieve, marker-line input, and missing-payload errors.
Findings
- 🟠 SQLite store opened per-request — wasteful for a long-running server
src/tokenless/crates/tokenless-cli/src/mcp.rs — retrieve() function
retrieve() calls open_stash_store_or_err(None) on every tools/call request. SqliteStore::new() runs Connection::open, three PRAGMAs (journal_mode=WAL, busy_timeout=5000, synchronous=NORMAL), CREATE TABLE IF NOT EXISTS, and CREATE INDEX IF NOT EXISTS — five SQL
statements before any actual retrieve query. For a long-running MCP server handling many retrieval requests in a session, this creates and tears down a SQLite connection per request.
SqliteStore is designed for shared long-lived use (wraps Connection in Mutex with poison-recovery), so creating one per request defeats its design.
Suggested fix: Open the store once at the top of serve() (or lazily on first tools/call) and pass it into retrieve() / retrieve_from_store():
pub fn serve() -> Result<(), (String, i32)> {
let store = open_stash_store_or_err(None).ok(); // fail-open at startup
// ... loop ...
"tools/call" => ok(id, handle_tool_call(params, &store)),
// ...
}
- 🟠 Malformed JSON silently dropped — violates JSON-RPC 2.0
src/tokenless/crates/tokenless-cli/src/mcp.rs:32-34
Err(_) => {
continue; // Skip unparseable lines
}
JSON-RPC 2.0 Section 5 requires a Parse Error response: {"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"Parse error"}}. If a client sends syntactically broken JSON that it expects a reply to, it will hang indefinitely waiting for a response that never
comes.
The comment says "some clients prefix LSP-style headers" — but those lines are typically Content-Length: N\r\n which would fail JSON parsing. A better approach: check if the line looks like a header (starts with a letter, contains :) before attempting JSON parse, and
only send -32700 for lines that look like JSON but fail to parse.
Suggested fix: Send a -32700 error on parse failure:
Err(_) => {
let _ = writeln!(out, "{}", err(Value::Null, -32700, "Parse error"));
let _ = out.flush();
continue;
}
- 🟡 Missing ping handler — MCP clients may mark server as unhealthy
src/tokenless/crates/tokenless-cli/src/mcp.rs:40-57
MCP 2024-11-05 defines ping as a standard utility method that either side can send to verify liveness, expecting an empty result {}. The current code falls through to the other arm and returns -32601 method not found: ping. MCP clients (like Claude Desktop) that send
periodic health-check pings would receive a protocol error, potentially marking the server as unhealthy or disconnecting.
Suggested fix: Add "ping" => ok(id, json!({})) to the method match.
- 🟡 No input size limit on stdin — unbounded memory allocation
src/tokenless/crates/tokenless-cli/src/mcp.rs:20
BufRead::read_line appends to the String buffer without any size limit. A malicious or buggy client could send a single JSON line containing gigabytes of data, causing OOM. The main CLI already guards against this with MAX_INPUT_BYTES (64 MiB) via Read::take. The MCP
server has no equivalent protection.
Suggested fix: Check line.len() after each read_line and reject oversized messages, or use a capped reader.
- 🟡 tools/call and tools/list accepted before initialize handshake
src/tokenless/crates/tokenless-cli/src/mcp.rs:40-57
MCP requires the initialize → notifications/initialized handshake before other methods are available. The server handles tools/list and tools/call at any time, even before initialize. A well-behaved client will always initialize first, but the server should track
initialization state and return -32002 ("Server not initialized") for premature requests.
Suggested fix: Add a let mut initialized = false; flag, set it on initialize, and gate tools/list and tools/call behind it.
- 🔵 Duplicate extract_hash call — validation result discarded then recomputed
src/tokenless/crates/tokenless-cli/src/mcp.rs
retrieve() calls extract_hash(hash) to validate the format (checking .is_none()), discards the result, then retrieve_from_store() calls extract_hash(hash) again to actually extract the key. The string scanning work is done twice per request.
Suggested fix: Extract once in retrieve(), pass the resolved key directly to retrieve_from_store():
let key = match extract_hash(hash) {
Some(h) => h,
None if is_valid_hash(hash) => hash,
None => return tool_error("invalid stash hash: ..."),
};
// pass key directly, no re-extraction needed
- 🔵 Error messages disclose internal filesystem paths to the LLM
src/tokenless/crates/tokenless-cli/src/mcp.rs — retrieve() function
Error strings like "stash unavailable: cannot open stash db at /home/user/.tokenless/stash.db: " are returned as MCP tool results, which the LLM sees and may echo to the end user or include in shared transcripts. This leaks the username, home directory
structure, and database implementation details.
Consider wrapping errors in generic messages for the tool result and logging the detailed error to stderr.
Notes
- Code shared with #1270 (stash integration, Retrieve CLI, open_stash_store_or_err, etc.) was already reviewed in the #1270 review and is not re-reviewed here.
- SQL injection is not a concern — SqliteStore::retrieve uses rusqlite::params![] parameterized queries, and both backends normalize keys to lowercase.
- The is_valid_hash import flagged by one reviewer is provided by the stacked #1270 commit included in this branch.
- Review methodology: 4 independent finder angles (line-by-line scan, MCP protocol compliance, efficiency/architecture, security) followed by dedup and cross-validation.
Expose the stash retrieval path over MCP so an MCP-connected agent can recover truncated payloads on demand — the MCP analogue of the `tokenless retrieve` CLI (closes the stash feature's MCP gap vs Headroom CCR's headroom_retrieve). - New `tokenless mcp serve` subcommand: a minimal stdio JSON-RPC server (hand-rolled, no MCP SDK dep — keeps the zero-runtime-dependency core path). Implements initialize / tools/list / tools/call per MCP 2024-11-05; notifications get no response; unknown methods return -32601. - Exposes one tool, `tokenless_retrieve` (hash -> payload text). Accepts a bare 24-hex hash or a line containing a `<<tokenless:KEY>>` marker (extracted via tokenless_ccr::extract_hash), mirroring the CLI. - Reuses open_stash_store + the passwd-rooted stash db trust model from the Retrieve CLI — no duplication of path validation. - Fail model: per-request tool errors return MCP tool results (isError: true) so the client keeps talking; only stdin/stdout I/O failure terminates the server. 6 unit tests (tool schema, dispatch, missing-arg, round-trip, marker line, not-found) + end-to-end stdio smoke verified (compress writes stash in one process, MCP tools/call retrieves in another). Stacked on alibaba#1270 (compressor integration + Retrieve CLI); rebase to a single commit once that lands. Signed-off-by: Shile Zhang <shile.zhang@linux.alibaba.com>
bec0bc8 to
64e3f7b
Compare
|
Thanks @Forrest-ly — all three findings addressed in 1. 🟠 SQLite store opened per-request — Fixed. The store is opened once at the top of 2. 🟠 Malformed JSON silently dropped — Fixed. Per JSON-RPC 2.0 §5, a line that looks like JSON (starts with 3. 🟡 Missing ping handler — Fixed. Added Tests updated for the new cc @samchu-zsl |
Description
Expose the stash retrieval path over MCP so an MCP-connected agent can recover truncated payloads on demand — the MCP analogue of the
tokenless retrieveCLI. Closes the stash feature's MCP gap vs Headroom CCR'sheadroom_retrieve(the only retrieval path before this PR was the CLI).This is the next priority item from the stash-feature completion list (part 1 = crate #1255 merged; part 2 = #1270 compressor integration + CLI; part 3 = #1271 tests + docs; this PR = MCP retrieve).
tokenless mcp serve— minimal stdio MCP server:initialize/tools/list/tools/callper MCP 2024-11-05; notifications get no response; unknown methods return-32601.tokenless_retrieve(hash→ payload text). Accepts a bare 24-hex hash or a line containing a<<tokenless:KEY>>marker (extracted viatokenless_ccr::extract_hash), mirroring the CLI.open_stash_store+ the passwd-rooted stash db trust model from theRetrieveCLI — no duplication of path validation.isError: true) so the client keeps talking; only stdin/stdout I/O failure terminates the server.Related Issue
no-issue: internal tokenless stash feature completion (MCP retrieve).
Type of Change
Scope
tokenless(tokenless)Checklist
descriptionfields are the MCP-facing docs)tokenless:cargo clippy --workspace --all-targets -- -D warningsandcargo fmt --all --checkpassTesting
End-to-end (separate processes — CLI compress writes stash, MCP
tools/callretrieves):Additional Notes
open_stash_store/extract_hashCLI infra, so this branch contains its commit. Merge #1270 first, then I rebase this to a single commit.Scope is deliberately minimal: only
tokenless_retrieveis exposed (the stash retrieval tool). The broader 4-tool MCP server from the planning doc (tokenless_compress/tokenless_stats/tokenless_env_check) is a separate effort — those are already CLI commands and exposing them via MCP is not stash-related.🤖 Generated with Claude Code