@takara-ai/miru-code 1.4.0 → 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.4.0",
+  "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/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 };
 }