ashlr-plugin 1.36.4

package/servers/_token-measure.ts

@@ -0,0 +1,160 @@
+/**
+ * _token-measure.ts — Measured token counting via Anthropic API.
+ *
+ * Two modes:
+ *   estimateTokens(text)  — synchronous chars/4 heuristic (always available)
+ *   measureTokens(text)   — async Anthropic API call returning true input_tokens
+ *
+ * Design constraints:
+ *   - NEVER blocks a hot path. All API calls are fire-and-forget from callers.
+ *   - Returns null on missing API key, timeout, or any failure.
+ *   - SHA-256 LRU cache persisted at ~/.ashlr/token-measure-cache.json (max 500).
+ *   - 5s AbortSignal timeout per call.
+ */
+
+import { createHash } from "crypto";
+import { existsSync, readFileSync, writeFileSync, mkdirSync } from "fs";
+import { homedir } from "os";
+import { join } from "path";
+import { estimateTokensFromString } from "@ashlr/core-efficiency";
+
+export type MeasurementMode = "estimate" | "measured";
+
+// ---------------------------------------------------------------------------
+// Estimate (synchronous, always available)
+// ---------------------------------------------------------------------------
+
+/**
+ * Estimate token count using the chars/4 heuristic. Wraps the shared
+ * estimateTokensFromString from @ashlr/core-efficiency for a single import
+ * surface.
+ */
+export function estimateTokens(text: string): number {
+  return estimateTokensFromString(text);
+}
+
+// ---------------------------------------------------------------------------
+// Cache
+// ---------------------------------------------------------------------------
+
+interface CacheEntry {
+  tokens: number;
+  /** ISO timestamp — used for eviction ordering */
+  at: string;
+}
+
+interface CacheShape {
+  /** sha256(text) → entry */
+  entries: Record<string, CacheEntry>;
+}
+
+const MAX_CACHE_ENTRIES = 500;
+
+function cachePath(): string {
+  return join(process.env.HOME ?? homedir(), ".ashlr", "token-measure-cache.json");
+}
+
+function loadCache(): CacheShape {
+  try {
+    const p = cachePath();
+    if (!existsSync(p)) return { entries: {} };
+    const raw = readFileSync(p, "utf-8");
+    const parsed = JSON.parse(raw) as Partial<CacheShape>;
+    if (!parsed.entries || typeof parsed.entries !== "object") return { entries: {} };
+    return { entries: parsed.entries };
+  } catch {
+    return { entries: {} };
+  }
+}
+
+function saveCache(cache: CacheShape): void {
+  try {
+    const p = cachePath();
+    mkdirSync(join(p, ".."), { recursive: true });
+    writeFileSync(p, JSON.stringify(cache), "utf-8");
+  } catch {
+    // best-effort
+  }
+}
+
+function cacheKey(text: string): string {
+  return createHash("sha256").update(text).digest("hex");
+}
+
+function evictIfNeeded(cache: CacheShape): void {
+  const keys = Object.keys(cache.entries);
+  if (keys.length <= MAX_CACHE_ENTRIES) return;
+  // Sort by timestamp ascending, drop oldest to get under limit
+  const sorted = keys.sort((a, b) => {
+    const ta = cache.entries[a]?.at ?? "";
+    const tb = cache.entries[b]?.at ?? "";
+    return ta < tb ? -1 : ta > tb ? 1 : 0;
+  });
+  const toRemove = sorted.slice(0, keys.length - MAX_CACHE_ENTRIES);
+  for (const k of toRemove) {
+    delete cache.entries[k];
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Measure (async, requires ANTHROPIC_API_KEY)
+// ---------------------------------------------------------------------------
+
+/**
+ * Measure true token count via Anthropic API.
+ *
+ * Returns null when:
+ *   - ANTHROPIC_API_KEY is not set
+ *   - Request times out (5s)
+ *   - Any network or API error
+ *   - Response shape is unexpected
+ *
+ * Results are cached by SHA-256(text) to avoid redundant API calls.
+ * NEVER throws — callers must be able to fire-and-forget.
+ */
+export async function measureTokens(text: string): Promise<number | null> {
+  const apiKey = process.env.ANTHROPIC_API_KEY;
+  if (!apiKey) return null;
+
+  const key = cacheKey(text);
+  const cache = loadCache();
+
+  // Cache hit
+  const cached = cache.entries[key];
+  if (cached && typeof cached.tokens === "number") {
+    return cached.tokens;
+  }
+
+  try {
+    const res = await fetch("https://api.anthropic.com/v1/messages", {
+      method: "POST",
+      signal: AbortSignal.timeout(5_000),
+      headers: {
+        "x-api-key": apiKey,
+        "anthropic-version": "2023-06-01",
+        "content-type": "application/json",
+      },
+      body: JSON.stringify({
+        model: "claude-haiku-4-5",
+        max_tokens: 1,
+        messages: [{ role: "user", content: [{ type: "text", text }] }],
+      }),
+    });
+
+    if (!res.ok) return null;
+
+    const json = (await res.json()) as { usage?: { input_tokens?: number } };
+    const tokens = json?.usage?.input_tokens;
+    if (typeof tokens !== "number" || !Number.isFinite(tokens)) return null;
+
+    // Store in cache
+    cache.entries[key] = { tokens, at: new Date().toISOString() };
+    evictIfNeeded(cache);
+    saveCache(cache);
+
+    return tokens;
+  } catch {
+    // timeout, network error, JSON parse error — all return null
+    return null;
+  }
+}

package/servers/_tool-base.ts

@@ -0,0 +1,248 @@
+/**
+ * Shared MCP server scaffold.
+ *
+ * Tools migrate from per-server boilerplate (stdio setup, ListTools/CallTool
+ * request handlers, lookup-by-name switch) to this handler-style registration
+ * so they can be composed into a single router process (see `_router.ts`) or
+ * continue running as standalone MCP servers via `runStandalone()`.
+ *
+ * Adoption is incremental. Existing servers keep working unchanged; they can
+ * opt-in by switching their `new Server(...)` + `setRequestHandler(...)`
+ * boilerplate to `registerTool({...})` + `runStandalone(name)`.
+ *
+ * Cross-track contract: all four tracks (router, context-db, ast, compression)
+ * target this interface so handlers can be hoisted into `_router.ts` without
+ * per-track refactors.
+ */
+
+import { Server } from "@modelcontextprotocol/sdk/server/index.js";
+import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import {
+  CallToolRequestSchema,
+  ListToolsRequestSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+
+import { logEvent } from "./_events";
+import { writeCrashDump } from "./_crash-dump";
+import { openContextDb, type ContextDb } from "./_embedding-cache";
+import { recordTelemetryEvent, isTelemetryEnabled } from "./_telemetry";
+import { detectMcpHost, type McpHost } from "./_mcp-host";
+
+/** Protocol-native guidance applies to every MCP host without host-specific hooks. */
+export const MCP_EFFICIENCY_INSTRUCTIONS = [
+  "Use Ashlr tools for scoped, token-efficient repository retrieval: orient/tree for structure, grep for relevant symbols, then read the required files or ranges.",
+  "Compact or summarized results are navigation aids. Retrieve omitted source before editing or verifying behavior; never treat omitted content as absent.",
+  "Preserve task requirements, AGENTS.md instructions, authorization boundaries, citations, exact error details and verification evidence. Expand retrieval when any of these are missing.",
+  "Use diff-format edits and focused tests. Savings estimates apply only to the tool output actually shortened; model quality, task success and token savings require separate evidence.",
+].join("\n");
+
+export interface ToolCallContext {
+  /** CLAUDE_SESSION_ID (or ASHLR_SESSION_ID override) when present. */
+  sessionId?: string;
+  /** Raw process env — handlers should prefer this over globals for testability. */
+  env: NodeJS.ProcessEnv;
+  /** Which MCP host invoked this tool. Defaults to "claude-code" for backwards compat. */
+  host?: McpHost;
+  /**
+   * When true, the handler should skip savings accounting (recordSaving).
+   * Set by ashlr__pipe on intermediate ctx calls so that only the pipe
+   * itself records ONE aggregate saving — not each sub-call individually.
+   * Handlers that call recordSaving unconditionally may ignore this field;
+   * the savings-server and stats helpers check it in the pipe's call path.
+   */
+  _noAccounting?: boolean;
+}
+
+export interface ToolTextContent {
+  type: "text";
+  text: string;
+}
+
+export interface ToolResult {
+  content: ToolTextContent[];
+  isError?: boolean;
+}
+
+export interface ToolHandler {
+  name: string;
+  description: string;
+  inputSchema: unknown;
+  handler: (args: Record<string, unknown>, ctx: ToolCallContext) => Promise<ToolResult>;
+}
+
+const registry = new Map<string, ToolHandler>();
+
+/**
+ * Register a tool handler. Safe to call multiple times per-name (last
+ * registration wins — useful for tests and for track overrides during the
+ * router rollout).
+ */
+export function registerTool(tool: ToolHandler): void {
+  registry.set(tool.name, tool);
+}
+
+/** Return all currently registered handlers. */
+export function listTools(): ToolHandler[] {
+  return Array.from(registry.values());
+}
+
+/** Lookup a registered handler by name. */
+export function getTool(name: string): ToolHandler | undefined {
+  return registry.get(name);
+}
+
+/** Test helper: drop all registrations. Do not call from production code. */
+export function __resetRegistryForTests(): void {
+  registry.clear();
+}
+
+/**
+ * Test helper: return a snapshot of the current registry so a caller can
+ * restore it after a scoped reset. Pair with `__restoreRegistryForTests`
+ * to keep the crash-isolation suite from leaking a wiped registry into
+ * tests that depend on the production handler set being present.
+ */
+export function __snapshotRegistryForTests(): ReadonlyMap<string, ToolHandler> {
+  return new Map(registry);
+}
+
+/** Test helper: overwrite the current registry with a previously snapshotted map. */
+export function __restoreRegistryForTests(snap: ReadonlyMap<string, ToolHandler>): void {
+  registry.clear();
+  for (const [k, v] of snap) registry.set(k, v);
+}
+
+// ---------------------------------------------------------------------------
+// Shared process-wide resources
+// ---------------------------------------------------------------------------
+//
+// Handlers registered on a single router process share the expensive
+// resources (SQLite handles, LLM clients) via the getters below. Before
+// router consolidation, every server opened its own context.db handle —
+// fine standalone, wasteful once N handlers live in one process. These
+// getters are lazy so importing `_tool-base` stays cheap for tests that
+// don't touch embedding infra.
+
+let _sharedCtxDb: ContextDb | null = null;
+
+/**
+ * Process-wide embedding cache handle. Migrated handlers should call this
+ * instead of `openContextDb()` directly so the router keeps one SQLite
+ * handle per process. Honors `ASHLR_CONTEXT_DB_DISABLE=1` (returns a
+ * no-op stub) since the underlying factory already does.
+ *
+ * Concurrency note: this check-then-set is safe only while `openContextDb()`
+ * is synchronous (which it is today — bun:sqlite opens on the main thread).
+ * If the factory ever becomes async, two concurrent awaits will each observe
+ * `null` and open two handles. At that point, replace the lazy init with a
+ * `Promise<ContextDb>` memo so the first caller owns the open.
+ */
+export function getEmbeddingCache(): ContextDb {
+  if (!_sharedCtxDb) _sharedCtxDb = openContextDb();
+  return _sharedCtxDb;
+}
+
+/**
+ * Build an `isError: true` ToolResult for a caught exception with a
+ * tool-specific prefix. Handlers that want to surface expected errors as
+ * first-class results (rather than letting the dispatch catch emit a
+ * `tool_crashed` event) use this to keep the boilerplate in one place.
+ */
+export function toErrorResult(prefix: string, err: unknown): ToolResult {
+  const message = err instanceof Error ? err.message : String(err);
+  return {
+    content: [{ type: "text", text: `${prefix}: ${message}` }],
+    isError: true,
+  };
+}
+
+/**
+ * Run the current tool registry as a standalone MCP server on stdio. Used by
+ * per-server entry points during the v1.12 → v1.13 migration so stale
+ * `plugin.json` files keep working. After router migration completes, most
+ * callers will go through `_router.ts` instead.
+ */
+export async function runStandalone(
+  serverName: string,
+  serverVersion: string = "1.0.0",
+): Promise<void> {
+  const server = new Server(
+    { name: serverName, version: serverVersion },
+    { capabilities: { tools: {} }, instructions: MCP_EFFICIENCY_INSTRUCTIONS },
+  );
+
+  server.setRequestHandler(ListToolsRequestSchema, async () => ({
+    tools: listTools().map((t) => ({
+      name: t.name,
+      description: t.description,
+      inputSchema: t.inputSchema,
+    })),
+  }));
+
+  // Matches the inline pattern every per-server uses today: TS infers the
+  // object-literal return as a valid variant of the SDK's ServerResult union.
+  server.setRequestHandler(CallToolRequestSchema, async (req) => {
+    const tool = getTool(req.params.name);
+    if (!tool) {
+      return {
+        content: [{ type: "text" as const, text: `Unknown tool: ${req.params.name}` }],
+        isError: true,
+      };
+    }
+    const ctx: ToolCallContext = {
+      sessionId:
+        process.env.CLAUDE_SESSION_ID || process.env.ASHLR_SESSION_ID || undefined,
+      env: process.env,
+      host: detectMcpHost(),
+    };
+    try {
+      // SDK's ServerResult is a union that includes a task-based variant; a
+      // typed ToolResult does not widen to the union from a variable, so we
+      // cast at the call site (same effective shape, glob-server-style).
+      const dispatchStart = Date.now();
+      const result = (await tool.handler(req.params.arguments ?? {}, ctx)) as unknown;
+      const dispatchMs = Date.now() - dispatchStart;
+
+      // 1.1 — Emit tool_name on every MCP tool dispatch (opt-in, privacy-safe).
+      // tool.name is the MCP tool name (e.g. "ashlr__grep"). No args/paths logged.
+      if (isTelemetryEnabled()) {
+        recordTelemetryEvent({
+          kind: "tool_call",
+          tool: tool.name,
+          tool_name: tool.name,
+          rawBytes: 0,       // byte accounting happens inside each handler
+          compactBytes: 0,
+          fellBack: false,
+          providerUsed: "local",
+          durationMs: dispatchMs,
+        });
+      }
+
+      return result as { content: unknown[] };
+    } catch (err) {
+      // Per-handler crash isolation: one handler's throw must not take the
+      // whole router down. Emit a tool_crashed event for observability, write
+      // a redacted on-disk crash dump for post-mortem, then return a
+      // structured error response.
+      const msg = err instanceof Error ? err.message : String(err);
+      const stack = err instanceof Error ? err.stack : undefined;
+      await logEvent("tool_crashed", {
+        tool: tool.name,
+        reason: msg,
+        extra: stack ? { stack: stack.split("\n").slice(0, 5).join("\n") } : undefined,
+      }).catch(() => undefined);
+      await writeCrashDump({
+        tool: tool.name,
+        args: req.params.arguments ?? {},
+        error: err,
+      });
+      return {
+        content: [{ type: "text" as const, text: `[ashlr:${tool.name}] handler crashed: ${msg}` }],
+        isError: true,
+      };
+    }
+  });
+
+  const transport = new StdioServerTransport();
+  await server.connect(transport);
+}