@takara-ai/miru-code 1.3.1 → 1.5.0

package/src/mcp/server.ts

@@ -1,13 +1,14 @@
 import * as z from "zod";
 import packageJson from "../../package.json";
-import { attachSearchBenchmark, benchmarkSearchComparison } from "../benchmark/compare.ts";
+import { benchmarkSearchComparison, toAgentBenchmarkSummary } from "../benchmark/compare.ts";
 import {
   appendBenchmarkQuery,
   readBenchmarkRollup,
   recordFromAgentSummary,
   recordFromBenchmark,
 } from "../benchmark/history.ts";
-import { attachLocateBenchmark, benchmarkLocateComparison } from "../benchmark/locate-compare.ts";
+import { benchmarkLocateComparison } from "../benchmark/locate-compare.ts";
+import { appendAgentBenchmark } from "../benchmark/summary.ts";
 import {
   MCP_BENCHMARK_SERVER_INSTRUCTIONS,
   MCP_EXPAND_TOOL_DESCRIPTION,
@@ -32,6 +33,11 @@ import {
   MAX_MCP_TOP_K,
   resolveChunk,
 } from "../utils.ts";
+import {
+  formatExpandResultsText,
+  formatLiteralLocateText,
+  formatResultsText,
+} from "./format-text.ts";
 import { getIndexForRepo, type IndexCache, toolText } from "./index-cache.ts";
 import { MiruMcpServer } from "./runtime.ts";
 
@@ -43,14 +49,11 @@ const REPO_DESCRIPTION =
 const BENCHMARK_LOCAL_ONLY_NOTE =
   "Benchmark comparisons require a local repo path; git URL repos are skipped.";
 
-function withBenchmarkSkipped<T extends Record<string, unknown>>(
-  payload: T,
-): T & { benchmark_skipped: "local_repo_only"; note: string } {
-  return {
-    ...payload,
+function withBenchmarkSkippedNote(body: string): string {
+  return `${body}\n\n${JSON.stringify({
     benchmark_skipped: "local_repo_only",
     note: BENCHMARK_LOCAL_ONLY_NOTE,
-  };
+  })}`;
 }
 
 export function createMcpServer(
@@ -108,20 +111,18 @@ export function createMcpServer(
           });
           const results = comparison.results;
           if (results.length === 0) {
-            return toolText(JSON.stringify({ error: "No results found." }));
+            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(
-            JSON.stringify(
-              attachSearchBenchmark(
-                formatResults(query, results, { repoRoot, snippet: true }),
-                comparison.benchmark,
-              ),
-            ),
+            appendAgentBenchmark(body, toAgentBenchmarkSummary(comparison.benchmark)),
           );
         }
 
@@ -130,13 +131,14 @@ export function createMcpServer(
           results = dedupeResultsByFile(results);
         }
         if (results.length === 0) {
-          return toolText(JSON.stringify({ error: "No results found." }));
+          return toolText("No results found.");
         }
         const payload = formatResults(query, results, { repoRoot, snippet: true });
+        const body = formatResultsText(payload);
         if (benchmark && !repoRoot) {
-          return toolText(JSON.stringify(withBenchmarkSkipped(payload)));
+          return toolText(withBenchmarkSkippedNote(body));
         }
-        return toolText(JSON.stringify(payload));
+        return toolText(body);
       } catch (err) {
         return toolText(err instanceof Error ? err.message : String(err));
       }
@@ -212,18 +214,18 @@ export function createMcpServer(
             } catch {
               // History is best-effort; never fail locate on persist errors.
             }
-            return toolText(
-              JSON.stringify(attachLocateBenchmark(comparison.payload, comparison.benchmark)),
-            );
+            const body = formatLiteralLocateText(comparison.payload);
+            return toolText(appendAgentBenchmark(body, comparison.benchmark));
           }
         }
 
         const result = index.locateLiteral(literal, locateOpts);
         const payload = formatLiteralLocate(result);
+        const body = formatLiteralLocateText(payload);
         if (benchmark && !localRepoRoot(repo)) {
-          return toolText(JSON.stringify(withBenchmarkSkipped(payload)));
+          return toolText(withBenchmarkSkippedNote(body));
         }
-        return toolText(JSON.stringify(payload));
+        return toolText(body);
       } catch (err) {
         return toolText(err instanceof Error ? err.message : String(err));
       }
@@ -277,7 +279,7 @@ export function createMcpServer(
           );
         }
         return toolText(
-          JSON.stringify(
+          formatExpandResultsText(
             formatExpandResults(filePath, anchorLine, anchor, expanded, {
               repoRoot,
               before: beforeCount,
@@ -327,12 +329,10 @@ export function createMcpServer(
         }
         const results = await index.findRelated(chunk, clampMcpTopK(topK));
         if (results.length === 0) {
-          return toolText(
-            JSON.stringify({ error: `No related chunks found for ${filePath}:${anchorLine}.` }),
-          );
+          return toolText(`No related chunks found for ${filePath}:${anchorLine}.`);
         }
         return toolText(
-          JSON.stringify(
+          formatResultsText(
             formatResults(`Chunks related to ${filePath}:${anchorLine}`, results, {
               repoRoot,
               snippet: true,

package/src/prompt.ts

@@ -1,4 +1,5 @@
 import { stdin as input, stdout as output } from "node:process";
+import * as readline from "node:readline/promises";
 
 export interface HiddenPromptState {
   value: string;
@@ -60,6 +61,18 @@ export function applyHiddenPromptChar(
   };
 }
 
+/** Visible-input prompt (not for secrets — use promptHidden for those). */
+export async function promptText(message: string, defaultValue = ""): Promise<string> {
+  const rl = readline.createInterface({ input, output });
+  try {
+    const suffix = defaultValue ? ` (${defaultValue})` : "";
+    const answer = (await rl.question(`${message}${suffix}: `)).trim();
+    return answer || defaultValue;
+  } finally {
+    rl.close();
+  }
+}
+
 export async function promptHidden(message: string): Promise<string> {
   if (!input.isTTY) {
     throw new Error(

package/src/setup.ts

@@ -1,14 +1,22 @@
-import { divider, fail, hint, info, printBrandBanner, success, writeStderr } from "./cli-ui.ts";
+import { divider, fail, hint, info, printBrandBanner, success, writeStdout } from "./cli-ui.ts";
 import {
   clearStoredCredentials,
   loadStoredCredentials,
   readStoredCredentials,
   resolveCredentialsPath,
+  type StoredSageMakerCredentials,
   saveStoredCredentials,
+  saveStoredSageMakerCredentials,
 } from "./credentials.ts";
+import {
+  isSageMakerConfigured,
+  parseSageMakerEndpointArn,
+  type SageMakerEmbeddingConfig,
+  validateSageMakerConnection,
+} from "./embeddings/sagemaker.ts";
 import { validateEmbeddingApiKey } from "./embeddings/validate.ts";
 import { hasTakaraApiKeyInEnv, resolveEmbeddingApiKey } from "./env.ts";
-import { promptHidden } from "./prompt.ts";
+import { promptHidden, promptText } from "./prompt.ts";
 import { Spinner } from "./spinner.ts";
 
 async function promptApiKey(): Promise<string> {
@@ -26,6 +34,9 @@ export interface RunSetupOptions {
   apiKey?: string;
   force?: boolean;
   skipValidation?: boolean;
+  sagemaker?: boolean;
+  sagemakerArn?: string;
+  profile?: string;
 }
 
 export interface RunSetupResult {
@@ -33,7 +44,165 @@ export interface RunSetupResult {
   newlySaved: boolean;
 }
 
+export interface ParsedSetupCliArgs {
+  apiKey?: string;
+  force: boolean;
+  clear: boolean;
+  sagemaker: boolean;
+  sagemakerArn?: string;
+  profile?: string;
+}
+
+export type SetupCliArgError = "clear_with_key" | "sagemaker_with_key";
+
+/** Parse `miru setup` argv (everything after the `setup` command). */
+export function parseSetupCliArgs(rest: string[]): {
+  args: ParsedSetupCliArgs;
+  error?: SetupCliArgError;
+} {
+  let apiKey: string | undefined;
+  let force = false;
+  let clear = false;
+  let sagemaker = false;
+  let sagemakerArn: string | undefined;
+  let profile: string | undefined;
+
+  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];
+    } else if (arg === "--sagemaker") {
+      sagemaker = true;
+    } else if (arg === "--arn" && rest[i + 1]) {
+      sagemakerArn = rest[++i];
+    } else if (arg === "--profile" && rest[i + 1]) {
+      profile = rest[++i];
+    }
+  }
+
+  if (sagemakerArn) {
+    sagemaker = true;
+  }
+
+  const args: ParsedSetupCliArgs = {
+    apiKey,
+    force,
+    clear,
+    sagemaker,
+    sagemakerArn,
+    profile,
+  };
+
+  if (clear && apiKey) {
+    return { args, error: "clear_with_key" };
+  }
+  if (sagemaker && apiKey) {
+    return { args, error: "sagemaker_with_key" };
+  }
+  return { args };
+}
+
+async function promptSageMakerArn(): Promise<string> {
+  while (true) {
+    const arn = await promptText("SageMaker endpoint ARN");
+    if (!arn) {
+      fail("Endpoint ARN cannot be empty.");
+      continue;
+    }
+    try {
+      parseSageMakerEndpointArn(arn);
+      return arn;
+    } catch (err) {
+      fail(err instanceof Error ? err.message : String(err));
+    }
+  }
+}
+
+async function promptAwsProfile(): Promise<string> {
+  while (true) {
+    const profile = await promptText(
+      "AWS profile name (must already exist in ~/.aws)",
+      process.env.AWS_PROFILE || "miru",
+    );
+    if (profile) {
+      return profile;
+    }
+    fail("AWS profile name cannot be empty.");
+  }
+}
+
+/**
+ * Miru only ever inherits AWS credentials — it never creates, writes, or rotates an
+ * AWS profile. Set one up yourself first (e.g. `aws configure --profile miru`).
+ */
+export async function runSageMakerSetup(options: RunSetupOptions = {}): Promise<RunSetupResult> {
+  writeStdout("");
+  printBrandBanner(process.stderr);
+  divider("─", 48, process.stderr);
+  writeStdout("Miru will connect directly to your self-hosted SageMaker embedding endpoint.");
+  hint("Miru only inherits AWS credentials from a profile you've already configured —");
+  hint("it never creates or writes to ~/.aws. Run `aws configure --profile <name>` first.");
+  hint("This replaces any stored Takara API key — only one embedding mode is active at a time.");
+  writeStdout("");
+
+  const arnInput = options.sagemakerArn ?? (await promptSageMakerArn());
+  const parsed = parseSageMakerEndpointArn(arnInput);
+
+  let profile = options.profile;
+  if (!profile) {
+    if (!canPromptForCredentials()) {
+      throw new Error(
+        "An AWS profile name is required. Pass --profile <name>, or run " +
+          "`miru setup --sagemaker` interactively.",
+      );
+    }
+    profile = await promptAwsProfile();
+  }
+
+  process.env.MIRU_SAGEMAKER_ENDPOINT_ARN = arnInput;
+  process.env.AWS_PROFILE = profile;
+
+  if (!options.skipValidation) {
+    const spinner = new Spinner(`Validating SageMaker endpoint (profile "${profile}")`);
+    spinner.start();
+    const config: SageMakerEmbeddingConfig = {
+      endpointName: parsed.endpointName,
+      region: parsed.region,
+      normalize: true,
+      truncate: true,
+      truncationDirection: "Right",
+    };
+    const result = await validateSageMakerConnection(config);
+    if (!result.valid) {
+      spinner.stop();
+      throw new Error(result.message);
+    }
+    spinner.succeed("SageMaker endpoint validated");
+  }
+
+  const stored: StoredSageMakerCredentials = { endpoint_arn: arnInput, profile };
+  const hadTakaraKey = Boolean((await readStoredCredentials())?.takara_api_key);
+  const path = await saveStoredSageMakerCredentials(stored);
+  writeStdout("");
+  success(`Saved SageMaker config to ${path}`);
+  if (hadTakaraKey) {
+    hint("Removed the stored Takara API key — Miru now embeds only via SageMaker.");
+  } else {
+    hint("Miru will embed via this SageMaker endpoint (Takara is not used).");
+  }
+  writeStdout("");
+  return { path, newlySaved: true };
+}
+
 export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupResult> {
+  if (options.sagemaker || options.sagemakerArn) {
+    return runSageMakerSetup(options);
+  }
+
   if (!options.force && hasTakaraApiKeyInEnv()) {
     const path = resolveCredentialsPath();
     const stored = await readStoredCredentials();
@@ -47,20 +216,23 @@ export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupR
 
   if (!options.force) {
     const stored = await readStoredCredentials();
-    if (stored && !options.apiKey) {
+    if (stored?.takara_api_key && !options.apiKey) {
       info(`API key already stored at ${resolveCredentialsPath()}. Use --force to replace.`);
       process.env.TAKARA_API_KEY = stored.takara_api_key;
       return { path: resolveCredentialsPath(), newlySaved: false };
     }
   }
 
-  writeStderr("");
+  writeStdout("");
   // stderr banner: setup runs before stdout may be a TTY (e.g. piped miru search).
   printBrandBanner(process.stderr);
   divider("─", 48, process.stderr);
-  writeStderr("Miru needs a Takara API key for code embeddings.");
+  writeStdout("Miru needs a Takara API key for code embeddings.");
   hint("Get a bearer token from Takara, then enter it below.");
-  writeStderr("");
+  hint(
+    "This replaces any stored SageMaker endpoint — only one embedding mode is active at a time.",
+  );
+  writeStdout("");
 
   const apiKey = options.apiKey ?? (await promptApiKey());
 
@@ -75,22 +247,27 @@ export async function runSetup(options: RunSetupOptions = {}): Promise<RunSetupR
     spinner.succeed("API key validated");
   }
 
+  const hadSageMaker = Boolean((await readStoredCredentials())?.sagemaker);
   const path = await saveStoredCredentials(apiKey);
   process.env.TAKARA_API_KEY = apiKey;
-  writeStderr("");
+  writeStdout("");
   success(`Saved credentials to ${path}`);
-  hint("MCP loads this key from credentials.json automatically.");
-  writeStderr("");
+  if (hadSageMaker) {
+    hint("Removed the stored SageMaker endpoint — Miru now embeds only via Takara.");
+  } else {
+    hint("MCP loads this key from credentials.json automatically.");
+  }
+  writeStdout("");
   return { path, newlySaved: true };
 }
 
 export async function runClearCredentials(): Promise<void> {
   const { cleared, path } = await clearStoredCredentials();
   if (cleared) {
-    success(`Removed stored API key from ${path}`);
+    success(`Removed stored credentials from ${path}`);
     return;
   }
-  info(`No stored API key at ${path}`);
+  info(`No stored credentials at ${path}`);
 }
 
 export function canPromptForCredentials(): boolean {
@@ -98,6 +275,9 @@ export function canPromptForCredentials(): boolean {
 }
 
 export function hasCredentials(): boolean {
+  if (isSageMakerConfigured()) {
+    return true;
+  }
   try {
     resolveEmbeddingApiKey();
     return true;
@@ -122,7 +302,7 @@ export async function ensureCredentials(options?: { interactive?: boolean }): Pr
 
   const wantsPrompt = options?.interactive ?? true;
   if (wantsPrompt && canPromptForCredentials()) {
-    writeStderr("");
+    writeStdout("");
     info("No Takara API key found.");
     hint("Miru needs one for embeddings — enter it below (same as `miru setup`).");
     await runSetup();
@@ -130,10 +310,10 @@ export async function ensureCredentials(options?: { interactive?: boolean }): Pr
     return;
   }
 
-  writeStderr("");
+  writeStdout("");
   // Brand on stderr when credentials are missing in non-interactive mode (stdout may not be a TTY).
   printBrandBanner(process.stderr);
-  writeStderr("");
+  writeStdout("");
 
   throw new Error(
     "Takara API key required. Run `miru setup` in a terminal, or set TAKARA_API_KEY " +