@takara-ai/miru-code 1.5.0 → 1.5.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@takara-ai/miru-code",
-  "version": "1.5.0",
+  "version": "1.5.3",
   "description": "Miru (見る) — hybrid code search for agents. Bun, TypeScript, Takara embeddings.",
   "type": "module",
   "license": "MIT",

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/credentials.ts

@@ -61,6 +61,19 @@ export async function readStoredCredentials(): Promise<StoredCredentials | null>
   }
 }
 
+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;
+}
+
 function hydrateSageMakerEnv(sagemaker: StoredSageMakerCredentials): void {
   process.env.MIRU_SAGEMAKER_ENDPOINT_ARN = sagemaker.endpoint_arn;
   if (sagemaker.profile && !process.env.AWS_PROFILE) {
@@ -68,81 +81,81 @@ function hydrateSageMakerEnv(sagemaker: StoredSageMakerCredentials): void {
   }
 }
 
-/** Hydrate TAKARA_API_KEY / SageMaker env vars from the credentials file when env is unset. */
-export async function loadStoredCredentials(): Promise<boolean> {
-  normalizeTakaraApiKeyEnv();
-  const needsTakara = !hasTakaraApiKeyInEnv();
-  const needsSageMaker = !process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
-  if (!needsTakara && !needsSageMaker) {
-    return false;
-  }
-
+/** Drop the other backend from env so setup validation cannot see a stale mode. */
+export async function beginModeSwitch(to: "takara" | "sagemaker"): Promise<void> {
   const stored = await readStoredCredentials();
-  if (!stored) {
-    return false;
+  if (to === "takara") {
+    clearSageMakerEnv(stored?.sagemaker);
+  } else {
+    clearTakaraEnv();
   }
-
-  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> {
+async function writeExclusive(payload: StoredCredentials): 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,
-  };
   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 (previousSageMaker) {
-    if (process.env.MIRU_SAGEMAKER_ENDPOINT_ARN === previousSageMaker.endpoint_arn) {
-      delete process.env.MIRU_SAGEMAKER_ENDPOINT_ARN;
+  return path;
+}
+
+/** Apply credentials.json as the sole mode — never leave Takara and SageMaker both active. */
+export async function loadStoredCredentials(): Promise<boolean> {
+  normalizeTakaraApiKeyEnv();
+  const stored = await readStoredCredentials();
+  if (!stored) {
+    return false;
+  }
+
+  if (stored.takara_api_key) {
+    let changed = false;
+    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) {
+    let changed = false;
+    if (!process.env.MIRU_SAGEMAKER_ENDPOINT_ARN?.trim()) {
+      hydrateSageMakerEnv(stored.sagemaker);
+      changed = true;
     }
-    if (previousSageMaker.profile && process.env.AWS_PROFILE === previousSageMaker.profile) {
-      delete process.env.AWS_PROFILE;
+    if (hasTakaraApiKeyInEnv()) {
+      clearTakaraEnv();
+      changed = true;
     }
+    return changed;
   }
+  return false;
+}
+
+export async function saveStoredCredentials(apiKey: string): Promise<string> {
+  const previous = await readStoredCredentials();
+  const path = await writeExclusive({ version: CREDENTIALS_VERSION, takara_api_key: apiKey });
+  process.env.TAKARA_API_KEY = apiKey;
+  clearSageMakerEnv(previous?.sagemaker);
   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;
+  const path = await writeExclusive({ version: CREDENTIALS_VERSION, sagemaker });
+  clearTakaraEnv();
+  process.env.MIRU_SAGEMAKER_ENDPOINT_ARN = sagemaker.endpoint_arn;
+  if (sagemaker.profile) {
+    process.env.AWS_PROFILE = sagemaker.profile;
   }
   return path;
 }
@@ -157,16 +170,10 @@ export async function clearStoredCredentials(): Promise<{ cleared: boolean; path
   await Bun.file(path).delete();
 
   if (stored?.takara_api_key && process.env.TAKARA_API_KEY === stored.takara_api_key) {
-    delete process.env.TAKARA_API_KEY;
+    clearTakaraEnv();
   }
-  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;
-    }
+  if (stored?.sagemaker) {
+    clearSageMakerEnv(stored.sagemaker);
   }
 
   return { cleared: true, path };

package/src/embeddings/validate.ts

@@ -3,7 +3,7 @@ import {
   embeddingDimensions,
   resolveEmbeddingBaseUrl,
   resolveEmbeddingDimensions,
-  resolveEmbeddingModel,
+  resolveTakaraEmbeddingModel,
 } from "./openai.ts";
 
 export interface ValidateApiKeyResult {
@@ -19,7 +19,7 @@ export async function validateEmbeddingApiKey(options: {
   dimensions?: number;
 }): Promise<ValidateApiKeyResult> {
   const baseUrl = (options.baseUrl ?? resolveEmbeddingBaseUrl()).replace(/\/$/, "");
-  const model = options.model ?? resolveEmbeddingModel();
+  const model = options.model ?? resolveTakaraEmbeddingModel();
   const dimensions = options.dimensions ?? resolveEmbeddingDimensions(model);
   const endpoint = `${baseUrl}/embeddings`;
 

package/src/git.ts

@@ -20,22 +20,17 @@ export async function cloneGitRepository(url: string, ref?: string | null): Prom
       stdout: "pipe",
       stderr: "pipe",
       stdin: "ignore",
+      timeout: timeoutSec * 1000,
     });
   } catch {
     await rm(dir, { recursive: true, force: true });
     throw new Error("git is not installed or not on PATH");
   }
 
-  let timedOut = false;
-  const timer = setTimeout(() => {
-    timedOut = true;
-    proc.kill();
-  }, timeoutSec * 1000);
-
   const code = await proc.exited;
-  clearTimeout(timer);
 
-  if (timedOut) {
+  // Prefer signalCode over killed — Bun may set killed=true after a normal exit.
+  if (proc.signalCode != null) {
     await rm(dir, { recursive: true, force: true });
     throw new Error(`git clone timed out for ${url} (limit: ${timeoutSec}s)`);
   }

package/src/help.ts

@@ -202,6 +202,10 @@ export function printCommandHelp(command: string): void {
       writeStdout("to a Grep baseline (rg search + read of the top matched file),");
       writeStdout("not bare chunk content and not the agent's full tool chain.");
       writeStdout("Local paths only — git URL repos skip comparison with benchmark_skipped.");
+      writeStdout(
+        "Cap each rg/grep spawn with MIRU_BENCHMARK_SEARCH_TIMEOUT (seconds, default 10).",
+      );
+      writeStdout("On timeout Miru still returns results with benchmark_skipped=grep_timeout.");
       writeStdout("");
       return;
     case "init":

package/src/mcp/server.ts

@@ -1,6 +1,7 @@
 import * as z from "zod";
 import packageJson from "../../package.json";
 import { benchmarkSearchComparison, toAgentBenchmarkSummary } from "../benchmark/compare.ts";
+import { withGrepTimeoutFallback } from "../benchmark/grep.ts";
 import {
   appendBenchmarkQuery,
   readBenchmarkRollup,
@@ -46,16 +47,32 @@ const REPO_DESCRIPTION =
   "Pass the project root for local workspaces. " +
   "The index is built on the first tool call and cached for the session.";
 
-const BENCHMARK_LOCAL_ONLY_NOTE =
-  "Benchmark comparisons require a local repo path; git URL repos are skipped.";
+const BENCHMARK_SKIP_NOTES = {
+  local_repo_only: "Benchmark comparisons require a local repo path; git URL repos are skipped.",
+  grep_timeout:
+    "Benchmark Grep baseline timed out; Miru results are still returned. " +
+    "Raise MIRU_BENCHMARK_SEARCH_TIMEOUT (seconds) if needed.",
+} as const;
 
-function withBenchmarkSkippedNote(body: string): string {
+type BenchmarkSkipReason = keyof typeof BENCHMARK_SKIP_NOTES;
+
+function withBenchmarkSkippedNote(body: string, reason: BenchmarkSkipReason): string {
   return `${body}\n\n${JSON.stringify({
-    benchmark_skipped: "local_repo_only",
-    note: BENCHMARK_LOCAL_ONLY_NOTE,
+    benchmark_skipped: reason,
+    note: BENCHMARK_SKIP_NOTES[reason],
   })}`;
 }
 
+async function persistBenchmarkQuery(
+  ...args: Parameters<typeof appendBenchmarkQuery>
+): Promise<void> {
+  try {
+    await appendBenchmarkQuery(...args);
+  } catch {
+    // History is best-effort; never fail the tool on persist errors.
+  }
+}
+
 export function createMcpServer(
   cache: IndexCache,
   options?: { benchmark?: boolean },
@@ -101,29 +118,33 @@ export function createMcpServer(
         const repoRoot = localRepoRoot(repo);
         const k = clampMcpTopK(topK);
 
+        let skip: BenchmarkSkipReason | undefined;
         if (benchmark && repoRoot) {
-          const comparison = await benchmarkSearchComparison({
-            query,
-            repoPath: repoRoot,
-            index,
-            topK: k,
-            dedupeByFile: dedupeByFile !== false,
-          });
-          const results = comparison.results;
-          if (results.length === 0) {
-            return toolText("No results found.");
-          }
-          try {
-            await appendBenchmarkQuery(recordFromBenchmark(repoRoot, comparison.benchmark));
-          } catch {
-            // History is best-effort; never fail the search on persist errors.
-          }
-          const body = formatResultsText(
-            formatResults(query, results, { repoRoot, snippet: true }),
-          );
-          return toolText(
-            appendAgentBenchmark(body, toAgentBenchmarkSummary(comparison.benchmark)),
+          const comparison = await withGrepTimeoutFallback(() =>
+            benchmarkSearchComparison({
+              query,
+              repoPath: repoRoot,
+              index,
+              topK: k,
+              dedupeByFile: dedupeByFile !== false,
+            }),
           );
+          if (comparison) {
+            const results = comparison.results;
+            if (results.length === 0) {
+              return toolText("No results found.");
+            }
+            await persistBenchmarkQuery(recordFromBenchmark(repoRoot, comparison.benchmark));
+            const body = formatResultsText(
+              formatResults(query, results, { repoRoot, snippet: true }),
+            );
+            return toolText(
+              appendAgentBenchmark(body, toAgentBenchmarkSummary(comparison.benchmark)),
+            );
+          }
+          skip = "grep_timeout";
+        } else if (benchmark && !repoRoot) {
+          skip = "local_repo_only";
         }
 
         let results = await index.search({ query, topK: k });
@@ -135,10 +156,7 @@ export function createMcpServer(
         }
         const payload = formatResults(query, results, { repoRoot, snippet: true });
         const body = formatResultsText(payload);
-        if (benchmark && !repoRoot) {
-          return toolText(withBenchmarkSkippedNote(body));
-        }
-        return toolText(body);
+        return toolText(skip ? withBenchmarkSkippedNote(body, skip) : body);
       } catch (err) {
         return toolText(err instanceof Error ? err.message : String(err));
       }
@@ -200,32 +218,31 @@ export function createMcpServer(
       try {
         const index = await getIndexForRepo(repo, cache);
 
-        if (benchmark) {
-          const repoRoot = localRepoRoot(repo);
-          if (repoRoot && typeof literal === "string") {
-            const comparison = await benchmarkLocateComparison({
+        let skip: BenchmarkSkipReason | undefined;
+        const repoRoot = localRepoRoot(repo);
+        if (benchmark && repoRoot && typeof literal === "string") {
+          const comparison = await withGrepTimeoutFallback(() =>
+            benchmarkLocateComparison({
               literal,
               repoPath: repoRoot,
               index,
               locate: locateOpts,
-            });
-            try {
-              await appendBenchmarkQuery(recordFromAgentSummary(repoRoot, comparison.benchmark));
-            } catch {
-              // History is best-effort; never fail locate on persist errors.
-            }
+            }),
+          );
+          if (comparison) {
+            await persistBenchmarkQuery(recordFromAgentSummary(repoRoot, comparison.benchmark));
             const body = formatLiteralLocateText(comparison.payload);
             return toolText(appendAgentBenchmark(body, comparison.benchmark));
           }
+          skip = "grep_timeout";
+        } else if (benchmark && !repoRoot) {
+          skip = "local_repo_only";
         }
 
         const result = index.locateLiteral(literal, locateOpts);
         const payload = formatLiteralLocate(result);
         const body = formatLiteralLocateText(payload);
-        if (benchmark && !localRepoRoot(repo)) {
-          return toolText(withBenchmarkSkippedNote(body));
-        }
-        return toolText(body);
+        return toolText(skip ? withBenchmarkSkippedNote(body, skip) : body);
       } catch (err) {
         return toolText(err instanceof Error ? err.message : String(err));
       }

package/src/setup.ts

@@ -1,5 +1,6 @@
 import { divider, fail, hint, info, printBrandBanner, success, writeStdout } from "./cli-ui.ts";
 import {
+  beginModeSwitch,
   clearStoredCredentials,
   loadStoredCredentials,
   readStoredCredentials,
@@ -163,6 +164,7 @@ export async function runSageMakerSetup(options: RunSetupOptions = {}): Promise<
     profile = await promptAwsProfile();
   }
 
+  await beginModeSwitch("sagemaker");
   process.env.MIRU_SAGEMAKER_ENDPOINT_ARN = arnInput;
   process.env.AWS_PROFILE = profile;
 
@@ -206,7 +208,15 @@ export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupR
   if (!options.force && hasTakaraApiKeyInEnv()) {
     const path = resolveCredentialsPath();
     const stored = await readStoredCredentials();
-    if (stored) {
+    if (stored?.sagemaker) {
+      const saved = await saveStoredCredentials(resolveEmbeddingApiKey());
+      writeStdout("");
+      success(`Saved credentials to ${saved}`);
+      hint("Removed the stored SageMaker endpoint — Miru now embeds only via Takara.");
+      writeStdout("");
+      return { path: saved, newlySaved: true };
+    }
+    if (stored?.takara_api_key) {
       info(`API key already configured (env + ${path}). Use --force to replace stored key.`);
       return { path, newlySaved: false };
     }
@@ -236,6 +246,7 @@ export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupR
 
   const apiKey = options.apiKey ?? (await promptApiKey());
 
+  await beginModeSwitch("takara");
   if (!options.skipValidation) {
     const spinner = new Spinner("Validating API key");
     spinner.start();
@@ -249,7 +260,6 @@ export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupR
 
   const hadSageMaker = Boolean((await readStoredCredentials())?.sagemaker);
   const path = await saveStoredCredentials(apiKey);
-  process.env.TAKARA_API_KEY = apiKey;
   writeStdout("");
   success(`Saved credentials to ${path}`);
   if (hadSageMaker) {