@takara-ai/miru-code 1.4.0 → 1.5.2

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.4.0",
+  "version": "1.5.2",
   "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/grep.ts

@@ -1,6 +1,7 @@
 /** Native grep baseline for Miru benchmark comparisons. */
 
 import { relative } from "node:path";
+import { envOptionalInt } from "../env.ts";
 import { countTokens } from "../token-count.ts";
 
 const STOPWORDS = new Set(
@@ -12,6 +13,9 @@ const STOPWORDS = new Set(
 export const GREP_LINES_PER_FILE = 3;
 export const GREP_CONTEXT = 2;
 
+/** Default cap for each rg/grep/findstr spawn in benchmark mode (seconds). */
+const DEFAULT_BENCHMARK_SEARCH_TIMEOUT_SEC = 10;
+
 /** Shared `-g` excludes for ripgrep baselines (search + locate). */
 export const RG_EXCLUDE_ARGS = [
   "-g",
@@ -24,6 +28,44 @@ export const RG_EXCLUDE_ARGS = [
   "!tokenizer/**",
 ] as const;
 
+export class BenchmarkSearchTimeoutError extends Error {
+  constructor(timeoutSec: number) {
+    super(`Benchmark search timed out after ${timeoutSec}s`);
+    this.name = "BenchmarkSearchTimeoutError";
+  }
+}
+
+/** `Bun.spawn` for benchmark baselines, killed after MIRU_BENCHMARK_SEARCH_TIMEOUT. */
+export async function spawnBenchmarkSearch(args: string[], cwd?: string): Promise<string> {
+  const timeoutSec =
+    envOptionalInt(["MIRU_BENCHMARK_SEARCH_TIMEOUT"], 1) ?? DEFAULT_BENCHMARK_SEARCH_TIMEOUT_SEC;
+  const proc = Bun.spawn(args, {
+    stdout: "pipe",
+    stderr: "pipe",
+    cwd,
+    timeout: timeoutSec * 1000,
+  });
+  const text = await new Response(proc.stdout).text();
+  await proc.exited;
+  // Prefer signalCode over killed — Bun may set killed=true after a normal exit.
+  if (proc.signalCode != null) {
+    throw new BenchmarkSearchTimeoutError(timeoutSec);
+  }
+  return text;
+}
+
+/** Run a Grep baseline; return null on timeout so Miru results can still be returned. */
+export async function withGrepTimeoutFallback<T>(run: () => Promise<T>): Promise<T | null> {
+  try {
+    return await run();
+  } catch (err) {
+    if (err instanceof BenchmarkSearchTimeoutError) {
+      return null;
+    }
+    throw err;
+  }
+}
+
 export interface GrepFileHit {
   file: string;
   matchCount: number;
@@ -172,34 +214,31 @@ export async function grepSearch(
 }
 
 async function rgRankedMatches(repoRoot: string, pattern: string, topK: number) {
-  const countProc = Bun.spawn(
-    ["rg", "-i", "--count-matches", pattern, repoRoot, ...RG_EXCLUDE_ARGS],
-    { stdout: "pipe", stderr: "pipe" },
-  );
-  const countText = await new Response(countProc.stdout).text();
-  await countProc.exited;
+  const countText = await spawnBenchmarkSearch([
+    "rg",
+    "-i",
+    "--count-matches",
+    pattern,
+    repoRoot,
+    ...RG_EXCLUDE_ARGS,
+  ]);
   return parseCountMatches(repoRoot, countText, topK);
 }
 
 async function grepRankedMatches(repoRoot: string, pattern: string, topK: number) {
-  const countProc = Bun.spawn(
-    [
-      "grep",
-      "-R",
-      "-I",
-      "-i",
-      "-E",
-      "-o",
-      "--exclude-dir=node_modules",
-      "--exclude-dir=.git",
-      "--exclude-dir=tokenizer",
-      pattern,
-      repoRoot,
-    ],
-    { stdout: "pipe", stderr: "pipe" },
-  );
-  const countText = await new Response(countProc.stdout).text();
-  await countProc.exited;
+  const countText = await spawnBenchmarkSearch([
+    "grep",
+    "-R",
+    "-I",
+    "-i",
+    "-E",
+    "-o",
+    "--exclude-dir=node_modules",
+    "--exclude-dir=.git",
+    "--exclude-dir=tokenizer",
+    pattern,
+    repoRoot,
+  ]);
   const counts = new Map<string, number>();
   for (const line of countText.split("\n")) {
     if (!line) {
@@ -240,13 +279,10 @@ function splitSearchToolLine(line: string): { path: string; rest: string } | nul
 async function findstrRankedMatches(repoRoot: string, keywords: string[], topK: number) {
   const counts = new Map<string, number>();
   for (const keyword of keywords) {
-    const proc = Bun.spawn(["findstr", "/S", "/N", "/I", "/P", `/C:${keyword}`, "*"], {
-      stdout: "pipe",
-      stderr: "pipe",
-      cwd: repoRoot,
-    });
-    const text = await new Response(proc.stdout).text();
-    await proc.exited;
+    const text = await spawnBenchmarkSearch(
+      ["findstr", "/S", "/N", "/I", "/P", `/C:${keyword}`, "*"],
+      repoRoot,
+    );
     for (const line of text.split("\n")) {
       const parsed = parsePathLinePrefix(line);
       if (!parsed?.path) {
@@ -302,51 +338,42 @@ function parseCountMatches(repoRoot: string, countText: string, topK: number) {
 }
 
 async function rgFilePreview(absPath: string, pattern: string): Promise<string> {
-  const proc = Bun.spawn(
-    [
-      "rg",
-      "-i",
-      "-n",
-      "-C",
-      String(GREP_CONTEXT),
-      "-m",
-      String(GREP_LINES_PER_FILE),
-      pattern,
-      absPath,
-    ],
-    { stdout: "pipe", stderr: "pipe" },
-  );
-  const output = await new Response(proc.stdout).text();
-  await proc.exited;
-  return output;
+  return spawnBenchmarkSearch([
+    "rg",
+    "-i",
+    "-n",
+    "-C",
+    String(GREP_CONTEXT),
+    "-m",
+    String(GREP_LINES_PER_FILE),
+    pattern,
+    absPath,
+  ]);
 }
 
 async function grepFilePreview(absPath: string, pattern: string): Promise<string> {
-  const proc = Bun.spawn(
-    [
-      "grep",
-      "-I",
-      "-i",
-      "-n",
-      "-E",
-      "-C",
-      String(GREP_CONTEXT),
-      "-m",
-      String(GREP_LINES_PER_FILE),
-      pattern,
-      absPath,
-    ],
-    { stdout: "pipe", stderr: "pipe" },
-  );
-  const output = await new Response(proc.stdout).text();
-  await proc.exited;
-  return output;
+  return spawnBenchmarkSearch([
+    "grep",
+    "-I",
+    "-i",
+    "-n",
+    "-E",
+    "-C",
+    String(GREP_CONTEXT),
+    "-m",
+    String(GREP_LINES_PER_FILE),
+    pattern,
+    absPath,
+  ]);
 }
 
 async function findstrFilePreview(absPath: string, keywords: string[]): Promise<string> {
-  const args = ["findstr", "/N", "/I", "/P", ...keywords.map((k) => `/C:${k}`), absPath];
-  const proc = Bun.spawn(args, { stdout: "pipe", stderr: "pipe" });
-  const output = await new Response(proc.stdout).text();
-  await proc.exited;
-  return output;
+  return spawnBenchmarkSearch([
+    "findstr",
+    "/N",
+    "/I",
+    "/P",
+    ...keywords.map((k) => `/C:${k}`),
+    absPath,
+  ]);
 }

package/src/benchmark/rg-literal.ts

@@ -2,7 +2,7 @@
 
 import { relative } from "node:path";
 import { countTokens } from "../token-count.ts";
-import { RG_EXCLUDE_ARGS, selectBenchmarkSearchTool } from "./grep.ts";
+import { RG_EXCLUDE_ARGS, selectBenchmarkSearchTool, spawnBenchmarkSearch } from "./grep.ts";
 
 export interface RgLiteralOutput {
   text: string;
@@ -76,16 +76,10 @@ export async function rgLiteralOutput(
     throw new Error("No search tool found in PATH (tried rg, grep, and findstr on Windows)");
   }
   const start = performance.now();
-  const proc = Bun.spawn(
+  const text = await spawnBenchmarkSearch(
     buildLiteralArgs(tool, repoRoot, literals, context, maxCount, !!options.ignoreCase),
-    {
-      stdout: "pipe",
-      stderr: "pipe",
-      ...(tool === "findstr" ? { cwd: repoRoot } : {}),
-    },
+    tool === "findstr" ? repoRoot : undefined,
   );
-  const text = await new Response(proc.stdout).text();
-  await proc.exited;
   const latency_ms = performance.now() - start;
   const stats = parseRgLiteralStats(text, repoRoot);
   return { text, tokens: countTokens(text), latency_ms, ...stats };

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,65 @@ 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;
+  }
+}
+
+function clearSageMakerEnv(sagemaker?: StoredSageMakerCredentials): void {
+  delete process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
+  delete process.env.MIRU_SAGEMAKER_ENDPOINT_NAME;
+  delete process.env.MIRU_SAGEMAKER_REGION;
+  if (sagemaker?.profile && process.env.AWS_PROFILE === sagemaker.profile) {
+    delete process.env.AWS_PROFILE;
+  }
+}
+
+function clearTakaraEnv(): void {
+  delete process.env.TAKARA_API_KEY;
+}
+
+/** Apply credentials.json as the sole mode — never leave Takara and SageMaker both active. */
 export async function loadStoredCredentials(): Promise<boolean> {
   normalizeTakaraApiKeyEnv();
-  if (hasTakaraApiKeyInEnv()) {
-    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 (stored.takara_api_key) {
+    if (!hasTakaraApiKeyInEnv()) {
+      process.env.TAKARA_API_KEY = stored.takara_api_key;
+      changed = true;
+    }
+    if (process.env.MIRU_SAGEMAKER_ENDPOINT_ARN || process.env.MIRU_SAGEMAKER_ENDPOINT_NAME) {
+      clearSageMakerEnv(stored.sagemaker);
+      changed = true;
+    }
+    return changed;
+  }
+
+  if (stored.sagemaker) {
+    if (!process.env.MIRU_SAGEMAKER_ENDPOINT_ARN?.trim()) {
+      hydrateSageMakerEnv(stored.sagemaker);
+      changed = true;
+    }
+    if (hasTakaraApiKeyInEnv()) {
+      clearTakaraEnv();
+      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 previousSageMaker = (await readStoredCredentials())?.sagemaker;
   const payload: StoredCredentials = {
     version: CREDENTIALS_VERSION,
     takara_api_key: apiKey,
@@ -83,6 +130,27 @@ export async function saveStoredCredentials(apiKey: string): Promise<string> {
   } catch {
     // Windows may not support Unix mode bits on all filesystems.
   }
+  clearSageMakerEnv(previousSageMaker);
+  return path;
+}
+
+export async function saveStoredSageMakerCredentials(
+  sagemaker: StoredSageMakerCredentials,
+): Promise<string> {
+  const dir = resolveCredentialsDir();
+  const path = resolveCredentialsPath();
+  await mkdir(dir, { recursive: true, mode: 0o700 });
+  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.
+  }
+  clearTakaraEnv();
   return path;
 }
 
@@ -95,8 +163,11 @@ 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) {
-    delete process.env.TAKARA_API_KEY;
+  if (stored?.takara_api_key && process.env.TAKARA_API_KEY === stored.takara_api_key) {
+    clearTakaraEnv();
+  }
+  if (stored?.sagemaker) {
+    clearSageMakerEnv(stored.sagemaker);
   }
 
   return { cleared: true, path };