@takara-ai/miru-code 1.3.1 → 1.5.0

package/README.md

@@ -238,6 +238,8 @@ Miru sends **file contents** to the [Takara inference API](https://takara.ai) wh
 
 If you index proprietary code, make sure that sending snippets to Takara's endpoint fits your security and compliance requirements. `MIRU_WORKSPACE_ROOT` is an opt-in boundary for MCP local `repo` paths only, and restricts indexing to a single workspace directory when set.
 
+Enterprise self-hosted embeddings (no Takara egress): see [docs/self-hosted-sagemaker.md](docs/self-hosted-sagemaker.md).
+
 ## Benchmark mode
 
 Optional measurement of how many tokens Miru saves versus a simple Grep workflow (ripgrep + reading the top matched file). Useful when evaluating Miru; leave it off day-to-day. Only local repo paths are compared — git URL repos skip the comparison and return `benchmark_skipped: "local_repo_only"`.

package/docs/self-hosted-sagemaker.md

@@ -0,0 +1,116 @@
+# Self-hosted embeddings (AWS SageMaker)
+
+Run Miru against your own SageMaker embedding endpoint instead of Takara’s hosted API (`infer.takara.ai`).
+
+The embedding model is sold on AWS Marketplace:
+
+**[Takara DS1 Miru (int8) — AWS Marketplace](https://aws.amazon.com/marketplace/pp/prodview-la5olwsicr6ue)**
+
+## Overview
+
+
+| Step | Who       | What                                                             |
+| ---- | --------- | ---------------------------------------------------------------- |
+| 1    | AWS admin | Subscribe on Marketplace and deploy a SageMaker endpoint         |
+| 2    | AWS admin | Create an invoke-only IAM user / profile (runbook below)         |
+| 3    | Developer | `miru setup --sagemaker` — Miru validates and saves the endpoint |
+
+
+Miru **never** creates IAM users or writes `~/.aws`. It only inherits a profile you already have and checks that it can call the endpoint.
+
+## One mode at a time
+
+Miru stores **either** a Takara API key **or** a SageMaker endpoint — never both.
+
+- `miru setup --sagemaker …` **removes** any stored Takara API key
+- `miru setup` (Takara) **removes** any stored SageMaker endpoint
+
+You do not need `--clear` when switching; the other mode is purged automatically.
+
+## 1. Deploy the Marketplace endpoint
+
+1. Open the [Marketplace listing](https://aws.amazon.com/marketplace/pp/prodview-la5olwsicr6ue)
+2. Subscribe and deploy the model package as a SageMaker endpoint in your account
+3. Copy the endpoint ARN — it looks like:
+
+```text
+arn:aws:sagemaker:<region>:<account-id>:endpoint/<name>
+```
+
+
+
+## 2. Create an invoke-only AWS profile (admin runbook)
+
+You need an AWS identity that can call `sagemaker:InvokeEndpoint` on that ARN. The simplest path for a laptop / CI user is a **least-privilege IAM user** scoped to that one endpoint, installed as a named profile.
+
+From a checkout of this repo, with the [AWS CLI](https://docs.aws.amazon.com/cli/) available and **admin** credentials in your current shell:
+
+```bash
+bun run sagemaker:create-invoke-user -- \
+  --endpoint-arn arn:aws:sagemaker:<region>:<account-id>:endpoint/<name> \
+  --profile miru
+```
+
+That script (`scripts/create-sagemaker-invoke-user.ts`):
+
+1. Creates an IAM user tagged for this purpose (or reuses one with the same name)
+2. Attaches an inline policy allowing **only** `sagemaker:InvokeEndpoint` on that ARN
+3. Creates an access key and writes it into `~/.aws` as the named profile (`miru` by default)
+
+Optional flags: `--user-name <name>`, `--profile <name>` (alias: `--arn` for the endpoint).
+
+Prefer SSO / IAM roles in production if your org requires it — the runbook script is the easy start with long-lived keys. Rotate or delete keys when someone leaves.
+
+Smoke-test invoke (optional):
+
+```bash
+bun run sagemaker:auth-check -- --arn arn:aws:sagemaker:<region>:<account-id>:endpoint/<name>
+```
+
+
+
+## 3. Point Miru at the endpoint
+
+```bash
+miru setup --sagemaker \
+  --arn arn:aws:sagemaker:<region>:<account-id>:endpoint/<name> \
+  --profile miru
+```
+
+Or run `miru setup --sagemaker` and answer the prompts.
+
+**What setup does**
+
+1. Parses the endpoint ARN + AWS profile name
+2. Invokes the endpoint once to confirm auth and that it returns embeddings
+3. Saves the SageMaker config and **deletes any stored Takara API key**
+
+After that, indexing and search embed via SageMaker — no egress to `infer.takara.ai`.
+
+If `TAKARA_API_KEY` is still set in your shell or `.env.local`, remove it there too so it cannot conflict later.
+
+## Switch back to Takara
+
+```bash
+miru setup
+```
+
+Enter your Takara API key. Setup validates it, saves it, and **deletes any stored SageMaker endpoint**.
+
+Also unset any `MIRU_SAGEMAKER_*` variables if you set them in the environment or `.env.local`.
+
+## Environment alternatives
+
+Prefer `miru setup --sagemaker` so auth is checked automatically. You can also set:
+
+
+| Variable                                                 | Purpose                                                   |
+| -------------------------------------------------------- | --------------------------------------------------------- |
+| `MIRU_SAGEMAKER_ENDPOINT_ARN`                            | `arn:aws:sagemaker:<region>:<account-id>:endpoint/<name>` |
+| `MIRU_SAGEMAKER_ENDPOINT_NAME` + `MIRU_SAGEMAKER_REGION` | Alternative to the ARN                                    |
+| `AWS_PROFILE` / `AWS_ACCESS_KEY_ID` / …                  | Standard AWS credential chain                             |
+
+
+Optional TEI-style knobs: `MIRU_SAGEMAKER_NORMALIZE`, `MIRU_SAGEMAKER_TRUNCATE`, `MIRU_SAGEMAKER_TRUNCATION_DIRECTION`, `MIRU_SAGEMAKER_PROMPT_NAME`.
+
+See `miru help setup` and `miru -h` for the same flags in the CLI.

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@takara-ai/miru-code",
-  "version": "1.3.1",
+  "version": "1.5.0",
   "description": "Miru (見る) — hybrid code search for agents. Bun, TypeScript, Takara embeddings.",
   "type": "module",
   "license": "MIT",
@@ -38,6 +38,7 @@
     "src",
     "grammars",
     "tokenizer",
+    "docs",
     "README.md",
     "LICENSE"
   ],
@@ -61,9 +62,13 @@
     "benchmark:literal": "bun run scripts/benchmark-literal-tokens.ts",
     "benchmark:quality": "bun run scripts/benchmark-quality.ts",
     "benchmark:quantization": "bun run scripts/benchmark-quantization.ts",
-    "benchmark:tokenizer-ab": "bun run scripts/tokenizer-ab.ts"
+    "benchmark:tokenizer-ab": "bun run scripts/tokenizer-ab.ts",
+    "sagemaker:auth-check": "bun run scripts/sagemaker-auth-check.ts",
+    "sagemaker:create-invoke-user": "bun run scripts/create-sagemaker-invoke-user.ts",
+    "benchmark:embedding-latency": "bun run scripts/embedding-latency-ab.ts"
   },
   "dependencies": {
+    "@aws-sdk/client-sagemaker-runtime": "^3.1090.0",
     "ignore": "^7.0.5",
     "web-tree-sitter": "^0.26.9",
     "zod": "^4.4.3"

package/src/benchmark/compare.ts

@@ -1,8 +1,9 @@
 import { readFile } from "node:fs/promises";
 import { join } from "node:path";
 import { relativePathFromRoot } from "../index/incremental.ts";
+import { formatExpandResultsText, formatResultsText } from "../mcp/format-text.ts";
 import type { MiruIndex } from "../miru-index.ts";
-import { applySnippetsToResults, estimateResultTokens } from "../snippet.ts";
+import { applySnippetsToResults } from "../snippet.ts";
 import { countTokens, tokenCountMethod, tokenizerJsonPath } from "../token-count.ts";
 import type { SearchResult } from "../types.ts";
 import {
@@ -10,6 +11,8 @@ import {
   DEFAULT_EXPAND_BEFORE,
   dedupeResultsByFile,
   expandChunksAtLine,
+  formatExpandResults,
+  formatResults,
 } from "../utils.ts";
 import { type GrepFileHit, grepSearch } from "./grep.ts";
 import { agentBenchmarkFromTokens, attachAgentBenchmark, tokenSavingsPct } from "./summary.ts";
@@ -83,6 +86,7 @@ async function readLineWindowTokens(
   }
 }
 
+/** Tokens of the MCP `expand` text body (headers + chunk content), matching what agents receive. */
 function miruExpandTokens(
   index: MiruIndex,
   repoPath: string,
@@ -94,7 +98,7 @@ function miruExpandTokens(
   }
   const [{ meta }] = applySnippetsToResults([top], query);
   const line = meta.truncated ? meta.anchor_line : top.chunk.start_line;
-  const { chunks } = expandChunksAtLine(
+  const { anchor, chunks } = expandChunksAtLine(
     index.chunks,
     top.chunk.file_path,
     line,
@@ -102,8 +106,18 @@ function miruExpandTokens(
     DEFAULT_EXPAND_BEFORE,
     DEFAULT_EXPAND_AFTER,
   );
+  if (!anchor) {
+    return { tokens: 0, lineSpan: 0 };
+  }
+  const text = formatExpandResultsText(
+    formatExpandResults(top.chunk.file_path, line, anchor, chunks, {
+      repoRoot: repoPath,
+      before: DEFAULT_EXPAND_BEFORE,
+      after: DEFAULT_EXPAND_AFTER,
+    }),
+  );
   return {
-    tokens: chunks.reduce((sum, chunk) => sum + countTokens(chunk.content), 0),
+    tokens: countTokens(text),
     lineSpan: expandLineSpan(chunks),
   };
 }
@@ -164,8 +178,10 @@ export async function benchmarkSearchComparison(options: {
     const started = performance.now();
     const raw = await index.search({ query, topK, rerank: true });
     const results = (dedupeByFile ? dedupeResultsByFile(raw) : raw).slice(0, topK);
-    const snippetResults = applySnippetsToResults(results, query).map((entry) => entry.result);
-    const searchTokens = estimateResultTokens(snippetResults);
+    // Count the MCP tool text body (path headers + snippets), not bare chunk content.
+    const searchTokens = countTokens(
+      formatResultsText(formatResults(query, results, { repoRoot: repoPath, snippet: true })),
+    );
     const expand = miruExpandTokens(index, repoPath, results[0], query);
     return {
       results,
@@ -264,6 +280,7 @@ export function toAgentBenchmarkSummary(block: SearchBenchmarkBlock): AgentBench
     block.accuracy.rank1_match,
   );
   summary.save_pct = block.efficiency.token_savings_pct;
+  summary.search_tok = block.miru.search_tokens;
   if (block.accuracy.miru_only.length > 0) {
     summary.miru_only = block.accuracy.miru_only.slice(0, MAX_AGENT_MIRU_ONLY);
   }

package/src/benchmark/locate-compare.ts

@@ -1,10 +1,11 @@
 /**
- * Compare Miru `locate` agent payload tokens vs the costlier native Grep path.
+ * Compare Miru `locate` MCP response tokens vs the costlier native Grep path.
  *
- * Baseline mirrors Cursor/Claude Grep-style output: fixed-string matches with
- * ±GREP_CONTEXT lines (not slim `rg -n`). That is the expensive path agents
- * take without `locate`, so savings measure real workflow cost — not an
- * apples-to-apples payload contest against optimally-minimal ripgrep.
+ * Miru side counts the agent-facing text body (`formatLiteralLocateText`), not
+ * the intermediate JSON payload. Baseline mirrors Cursor/Claude Grep-style
+ * output: fixed-string matches with ±GREP_CONTEXT lines (not slim `rg -n`).
+ * That is the expensive path agents take without `locate`, so savings measure
+ * real workflow cost — not an apples-to-apples contest against minimal ripgrep.
  */
 
 import {
@@ -13,6 +14,7 @@ import {
   type LiteralLocateOptions,
   type LiteralLocateResult,
 } from "../literal.ts";
+import { formatLiteralLocateText } from "../mcp/format-text.ts";
 import type { MiruIndex } from "../miru-index.ts";
 import { countTokens } from "../token-count.ts";
 import { GREP_CONTEXT } from "./grep.ts";
@@ -44,7 +46,7 @@ export async function benchmarkLocateComparison(options: {
   const result = options.index.locateLiteral(options.literal, locateOpts);
   const miruMs = performance.now() - miruStart;
   const payload = formatLiteralLocate(result);
-  const miruTok = countTokens(JSON.stringify(payload));
+  const miruTok = countTokens(formatLiteralLocateText(payload));
 
   // Unbounded agent-style Grep dump (±context) — the costlier path without locate.
   // Never use less context than the caller actually requested via context_lines,

package/src/benchmark/summary.ts

@@ -20,10 +20,19 @@ export function agentBenchmarkFromTokens(
   };
 }
 
-/** Attach compact `benchmark` to an MCP tool payload. */
+/** Attach compact `benchmark` to an MCP tool payload (JSON path / tests). */
 export function attachAgentBenchmark<T extends Record<string, unknown>>(
   payload: T,
   benchmark: AgentBenchmarkSummary,
 ): T & { benchmark: AgentBenchmarkSummary } {
   return { ...payload, benchmark };
 }
+
+/**
+ * Append a compact trailing JSON object so plaintext MCP bodies can still expose
+ * `benchmark` without wrapping the result in JSON scaffolding.
+ * `miru_tok` counts the body only — never this trailer.
+ */
+export function appendAgentBenchmark(body: string, benchmark: AgentBenchmarkSummary): string {
+  return `${body}\n\n${JSON.stringify({ benchmark })}`;
+}

package/src/benchmark/types.ts

@@ -42,7 +42,12 @@ export interface SearchBenchmarkBlock {
   };
 }
 
-/** Compact per-search stats attached to MCP `search` responses (agent-facing). */
+/**
+ * Compact per-search stats appended to MCP `search` / `locate` responses.
+ * `miru_tok` is the agent-facing MCP text body for that tool (search also adds
+ * one expand body for workflow vs grep+read). Not JSON scaffolding, not the
+ * trailing `{"benchmark":...}` line.
+ */
 export interface AgentBenchmarkSummary {
   save_pct: number;
   miru_tok: number;
@@ -51,6 +56,11 @@ export interface AgentBenchmarkSummary {
   rank1: boolean;
   /** Present only when non-empty; capped at 3 paths. */
   miru_only?: string[];
+  /**
+   * Search only: tokens of the search MCP text body alone (before expand).
+   * `miru_tok` = search_tok + expand body when present.
+   */
+  search_tok?: number;
 }
 
 /** Compact cumulative rollup for MCP `read_benchmark` (agent-facing). */

package/src/cli.ts

@@ -38,6 +38,7 @@ import { MiruIndex } from "./miru-index.ts";
 import {
   canPromptForCredentials,
   ensureCredentials,
+  parseSetupCliArgs,
   runClearCredentials,
   runSetup,
 } from "./setup.ts";
@@ -471,30 +472,28 @@ async function runCli(argv: string[]): Promise<void> {
   }
 
   if (command === "setup") {
-    let apiKey: string | undefined;
-    let force = false;
-    let clear = false;
-    for (let i = 0; i < rest.length; i++) {
-      const arg = rest[i];
-      if (arg === "--force") {
-        force = true;
-      } else if (arg === "--clear") {
-        clear = true;
-      } else if ((arg === "--key" || arg === "-k") && rest[i + 1]) {
-        apiKey = rest[++i];
-      }
+    const { args, error } = parseSetupCliArgs(rest);
+    if (error === "clear_with_key") {
+      fail("miru setup --clear cannot be combined with --key.");
+      process.exit(1);
     }
-    if (clear) {
-      if (apiKey) {
-        fail("miru setup --clear cannot be combined with --key.");
-        process.exit(1);
-      }
+    if (error === "sagemaker_with_key") {
+      fail("miru setup --sagemaker cannot be combined with --key.");
+      process.exit(1);
+    }
+    if (args.clear) {
       await runClearCredentials();
       return;
     }
-    const { newlySaved } = await runSetup({ apiKey, force });
+    const { newlySaved } = await runSetup({
+      apiKey: args.apiKey,
+      force: args.force,
+      sagemaker: args.sagemaker,
+      sagemakerArn: args.sagemakerArn,
+      profile: args.profile,
+    });
     if (newlySaved) {
-      const offerInstall = canPromptForCredentials() && !apiKey && !force;
+      const offerInstall = canPromptForCredentials() && !args.apiKey && !args.force;
       if (offerInstall) {
         const install = await promptConfirm("Configure Miru in your coding agent now?");
         if (install) {

package/src/credentials.ts

@@ -5,9 +5,15 @@ import { hasTakaraApiKeyInEnv, normalizeTakaraApiKeyEnv } from "./env.ts";
 const CREDENTIALS_FILENAME = "credentials.json";
 const CREDENTIALS_VERSION = 1;
 
+export interface StoredSageMakerCredentials {
+  endpoint_arn: string;
+  profile?: string;
+}
+
 interface StoredCredentials {
   version: number;
-  takara_api_key: string;
+  takara_api_key?: string;
+  sagemaker?: StoredSageMakerCredentials;
 }
 
 export function resolveCredentialsDir(): string {
@@ -46,7 +52,7 @@ export async function readStoredCredentials(): Promise<StoredCredentials | null>
   }
   try {
     const parsed = JSON.parse(await Bun.file(path).text()) as StoredCredentials;
-    if (parsed.version !== CREDENTIALS_VERSION || !parsed.takara_api_key) {
+    if (parsed.version !== CREDENTIALS_VERSION || (!parsed.takara_api_key && !parsed.sagemaker)) {
       return null;
     }
     return parsed;
@@ -55,24 +61,46 @@ export async function readStoredCredentials(): Promise<StoredCredentials | null>
   }
 }
 
-/** Set TAKARA_API_KEY from the user credentials file when env is unset. */
+function hydrateSageMakerEnv(sagemaker: StoredSageMakerCredentials): void {
+  process.env.MIRU_SAGEMAKER_ENDPOINT_ARN = sagemaker.endpoint_arn;
+  if (sagemaker.profile && !process.env.AWS_PROFILE) {
+    process.env.AWS_PROFILE = sagemaker.profile;
+  }
+}
+
+/** Hydrate TAKARA_API_KEY / SageMaker env vars from the credentials file when env is unset. */
 export async function loadStoredCredentials(): Promise<boolean> {
   normalizeTakaraApiKeyEnv();
-  if (hasTakaraApiKeyInEnv()) {
+  const needsTakara = !hasTakaraApiKeyInEnv();
+  const needsSageMaker = !process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
+  if (!needsTakara && !needsSageMaker) {
     return false;
   }
+
   const stored = await readStoredCredentials();
   if (!stored) {
     return false;
   }
-  process.env.TAKARA_API_KEY = stored.takara_api_key;
-  return true;
+
+  let changed = false;
+  if (needsTakara && stored.takara_api_key) {
+    process.env.TAKARA_API_KEY = stored.takara_api_key;
+    changed = true;
+  }
+  if (needsSageMaker && stored.sagemaker) {
+    hydrateSageMakerEnv(stored.sagemaker);
+    changed = true;
+  }
+  return changed;
 }
 
 export async function saveStoredCredentials(apiKey: string): Promise<string> {
   const dir = resolveCredentialsDir();
   const path = resolveCredentialsPath();
   await mkdir(dir, { recursive: true, mode: 0o700 });
+  const existing = await readStoredCredentials();
+  // Takara and SageMaker modes are mutually exclusive — saving a key drops any stored endpoint.
+  const previousSageMaker = existing?.sagemaker;
   const payload: StoredCredentials = {
     version: CREDENTIALS_VERSION,
     takara_api_key: apiKey,
@@ -83,6 +111,39 @@ export async function saveStoredCredentials(apiKey: string): Promise<string> {
   } catch {
     // Windows may not support Unix mode bits on all filesystems.
   }
+  if (previousSageMaker) {
+    if (process.env.MIRU_SAGEMAKER_ENDPOINT_ARN === previousSageMaker.endpoint_arn) {
+      delete process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
+    }
+    if (previousSageMaker.profile && process.env.AWS_PROFILE === previousSageMaker.profile) {
+      delete process.env.AWS_PROFILE;
+    }
+  }
+  return path;
+}
+
+export async function saveStoredSageMakerCredentials(
+  sagemaker: StoredSageMakerCredentials,
+): Promise<string> {
+  const dir = resolveCredentialsDir();
+  const path = resolveCredentialsPath();
+  await mkdir(dir, { recursive: true, mode: 0o700 });
+  const existing = await readStoredCredentials();
+  // Takara and SageMaker modes are mutually exclusive — saving an endpoint drops any stored key.
+  const previousTakaraKey = existing?.takara_api_key;
+  const payload: StoredCredentials = {
+    version: CREDENTIALS_VERSION,
+    sagemaker,
+  };
+  await Bun.write(path, `${JSON.stringify(payload, null, 2)}\n`, { mode: 0o600 });
+  try {
+    await chmod(path, 0o600);
+  } catch {
+    // Windows may not support Unix mode bits on all filesystems.
+  }
+  if (previousTakaraKey && process.env.TAKARA_API_KEY === previousTakaraKey) {
+    delete process.env.TAKARA_API_KEY;
+  }
   return path;
 }
 
@@ -95,9 +156,18 @@ export async function clearStoredCredentials(): Promise<{ cleared: boolean; path
   const stored = await readStoredCredentials();
   await Bun.file(path).delete();
 
-  if (stored && process.env.TAKARA_API_KEY === stored.takara_api_key) {
+  if (stored?.takara_api_key && process.env.TAKARA_API_KEY === stored.takara_api_key) {
     delete process.env.TAKARA_API_KEY;
   }
+  const sagemaker = stored?.sagemaker;
+  if (sagemaker) {
+    if (process.env.MIRU_SAGEMAKER_ENDPOINT_ARN === sagemaker.endpoint_arn) {
+      delete process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
+    }
+    if (sagemaker.profile && process.env.AWS_PROFILE === sagemaker.profile) {
+      delete process.env.AWS_PROFILE;
+    }
+  }
 
   return { cleared: true, path };
 }