@takara-ai/miru-code 1.9.0 → 1.10.0

package/src/embeddings/sagemaker.ts

@@ -26,6 +26,11 @@ function loadAwsCredentialProvider(): Promise<AwsCredentialProviderModule> {
   return awsCredentialProviderPromise;
 }
 
+/** Warm the lazily loaded AWS SDK modules before the first SageMaker request. */
+export async function preloadSageMakerSdk(): Promise<void> {
+  await Promise.all([loadSageMakerRuntime(), loadAwsCredentialProvider()]);
+}
+
 /** arn:aws:sagemaker:<region>:<account-id>:endpoint/<endpoint-name> (also covers aws-cn/aws-us-gov). */
 const ENDPOINT_ARN_PATTERN = /^arn:aws[a-z0-9-]*:sagemaker:([a-z0-9-]+):(\d{12}):endpoint\/(.+)$/;
 
@@ -344,7 +349,13 @@ function buildRequestBody(
   return body;
 }
 
-export function createSageMakerClient(config: SageMakerEmbeddingConfig): EmbeddingClient {
+export function createSageMakerClient(
+  config: SageMakerEmbeddingConfig,
+  loaders: {
+    loadRuntime?: () => Promise<SageMakerRuntimeModule>;
+    loadCredentials?: () => Promise<AwsCredentialProviderModule>;
+  } = {},
+): EmbeddingClient {
   let client: InstanceType<SageMakerRuntimeModule["SageMakerRuntimeClient"]> | null = null;
 
   return {
@@ -354,7 +365,10 @@ export function createSageMakerClient(config: SageMakerEmbeddingConfig): Embeddi
       dimensions?: number,
     ): Promise<EmbeddingResponse> {
       const [{ InvokeEndpointCommand, SageMakerRuntimeClient }, { defaultProvider }] =
-        await Promise.all([loadSageMakerRuntime(), loadAwsCredentialProvider()]);
+        await Promise.all([
+          (loaders.loadRuntime ?? loadSageMakerRuntime)(),
+          (loaders.loadCredentials ?? loadAwsCredentialProvider)(),
+        ]);
       if (!client) {
         const profile = config.profile?.trim() || process.env.AWS_PROFILE?.trim();
         const sdkCredentials = defaultProvider({ profile });
@@ -414,9 +428,10 @@ export interface ValidateSageMakerResult {
 /** Real InvokeEndpoint round trip, used by `miru setup --sagemaker` to confirm auth + shape. */
 export async function validateSageMakerConnection(
   config: SageMakerEmbeddingConfig,
+  createClient: typeof createSageMakerClient = createSageMakerClient,
 ): Promise<ValidateSageMakerResult> {
   try {
-    const client = createSageMakerClient(config);
+    const client = createClient(config);
     const response = await client.createEmbeddings("miru setup validation", "sagemaker-setup");
     const embedding = response.data[0]?.embedding;
     if (!embedding) {

package/src/embeddings/validate.ts

@@ -12,12 +12,15 @@ export interface ValidateApiKeyResult {
   message: string;
 }
 
-export async function validateEmbeddingApiKey(options: {
-  apiKey: string;
-  baseUrl?: string;
-  model?: string;
-  dimensions?: number;
-}): Promise<ValidateApiKeyResult> {
+export async function validateEmbeddingApiKey(
+  options: {
+    apiKey: string;
+    baseUrl?: string;
+    model?: string;
+    dimensions?: number;
+  },
+  fetchImpl: (input: string | URL | Request, init?: RequestInit) => Promise<Response> = fetch,
+): Promise<ValidateApiKeyResult> {
   const baseUrl = (options.baseUrl ?? resolveEmbeddingBaseUrl()).replace(/\/$/, "");
   const model = options.model ?? resolveTakaraEmbeddingModel();
   const dimensions = options.dimensions ?? resolveEmbeddingDimensions(model);
@@ -33,7 +36,7 @@ export async function validateEmbeddingApiKey(options: {
 
   let response: Response;
   try {
-    response = await fetch(endpoint, {
+    response = await fetchImpl(endpoint, {
       method: "POST",
       headers: {
         Authorization: `Bearer ${options.apiKey}`,

package/src/git.ts

@@ -5,7 +5,11 @@ import { envOptionalInt } from "./env.ts";
 
 const DEFAULT_CLONE_TIMEOUT_SEC = 60;
 
-export async function cloneGitRepository(url: string, ref?: string | null): Promise<string> {
+export async function cloneGitRepository(
+  url: string,
+  ref?: string | null,
+  spawn: typeof Bun.spawn = Bun.spawn,
+): Promise<string> {
   const timeoutSec = envOptionalInt(["MIRU_CLONE_TIMEOUT"], 1) ?? DEFAULT_CLONE_TIMEOUT_SEC;
   const dir = await mkdtemp(join(tmpdir(), "miru-git-"));
   const args = ["clone", "--depth", "1"];
@@ -16,7 +20,7 @@ export async function cloneGitRepository(url: string, ref?: string | null): Prom
 
   let proc: ReturnType<typeof Bun.spawn>;
   try {
-    proc = Bun.spawn(["git", ...args], {
+    proc = spawn(["git", ...args], {
       stdout: "pipe",
       stderr: "pipe",
       stdin: "ignore",

package/src/help.ts

@@ -28,7 +28,7 @@ export function printMainHelp(): void {
   header();
 
   section("Usage");
-  writeStdout("  miru                         Start MCP server (stdio)");
+  writeStdout("  miru                         Show help in a terminal");
   writeStdout("  miru <command> [options]");
   divider();
 
@@ -37,6 +37,7 @@ export function printMainHelp(): void {
   commandRow("locate", "Exact substring location in the index");
   commandRow("expand", "Adjacent chunks in the same file as a hit");
   commandRow("find-related", "Find chunks related to a file:line");
+  commandRow("mcp", "Start MCP server (stdio)");
   commandRow("setup", "Authenticate with Takara");
   commandRow("install", "Configure miru across coding agents");
   commandRow("uninstall", "Remove miru agent configuration");
@@ -177,7 +178,7 @@ export function printCommandHelp(command: string): void {
     case "install":
       commandHeader("install", "Interactive global agent setup.");
       writeStdout("Configures MCP server, instructions, sub-agent files, and optional");
-      writeStdout("experimental integrations (search hooks, Caveman, STE writing) under");
+      writeStdout("experimental integrations (Caveman, STE writing) under");
       writeStdout("your user config (~/.claude, ~/.cursor, etc.).");
       writeStdout("");
       writeStdout(
@@ -243,9 +244,9 @@ export function printCommandHelp(command: string): void {
       writeStdout("");
       return;
     case "mcp":
-      commandHeader("mcp", "Stdio MCP server (default with no subcommand).");
+      commandHeader("mcp", "Stdio MCP server.");
       section("Usage");
-      writeStdout("  miru [--ref BRANCH] [--content TYPE ...] [--benchmark]");
+      writeStdout("  miru mcp [--ref BRANCH] [--content TYPE ...] [--benchmark]");
       writeStdout("  Default content: code config");
       writeStdout("");
       writeStdout("Indexes on the first search/expand/find_related tool call (repo argument).");

package/src/index/bm25.ts

@@ -42,12 +42,21 @@ function runScoreWorker(job: Bm25ScoreJob): Promise<Bm25ScoreResult> {
 }
 
 export class BM25Index {
-  private docFreq = new Map<string, number>();
-  private docLengths: number[] = [];
-  private postings = new Map<string, PostingsList>();
-  private avgDocLength = 0;
-  private numDocs = 0;
-  private totalLen = 0;
+  private docFreq: Map<string, number>;
+  private docLengths: number[];
+  private postings: Map<string, PostingsList>;
+  private avgDocLength: number;
+  private numDocs: number;
+  private totalLen: number;
+
+  constructor() {
+    this.docFreq = new Map();
+    this.docLengths = [];
+    this.postings = new Map();
+    this.avgDocLength = 0;
+    this.numDocs = 0;
+    this.totalLen = 0;
+  }
 
   /** Append one tokenized document. Safe to call incrementally while indexing. */
   addDocument(doc: string[]): number {
@@ -73,6 +82,57 @@ export class BM25Index {
     return docIndex;
   }
 
+  get size(): number {
+    return this.numDocs;
+  }
+
+  /**
+   * A new index with the given documents removed and later doc indices shifted down,
+   * equal to re-indexing the surviving docs in order. `this` is left untouched, so
+   * searches already holding it never see a half-updated index.
+   */
+  withoutDocuments(removed: ReadonlySet<number>): BM25Index {
+    const next = new BM25Index();
+    const remap = new Int32Array(this.numDocs);
+    let kept = 0;
+    let firstRemoved = this.numDocs;
+    for (let doc = 0; doc < this.numDocs; doc++) {
+      if (removed.has(doc)) {
+        remap[doc] = -1;
+        if (firstRemoved === this.numDocs) {
+          firstRemoved = doc;
+        }
+      } else {
+        remap[doc] = kept++;
+        const length = this.docLengths[doc] ?? 0;
+        next.docLengths.push(length);
+        next.totalLen += length;
+      }
+    }
+    next.numDocs = kept;
+    next.avgDocLength = kept === 0 ? 0 : next.totalLen / kept;
+
+    for (const [term, list] of this.postings) {
+      // Lists ascend by doc index, so one ending before the first removal is untouched.
+      // Copy the array (addDocument appends to it) but share the immutable tuples.
+      const untouched = (list.at(-1)?.[0] ?? 0) < firstRemoved;
+      const nextList: PostingsList = untouched ? list.slice() : [];
+      if (!untouched) {
+        for (const entry of list) {
+          const mapped = remap[entry[0]] ?? -1;
+          if (mapped !== -1) {
+            nextList.push(mapped === entry[0] ? entry : [mapped, entry[1]]);
+          }
+        }
+      }
+      if (nextList.length > 0) {
+        next.postings.set(term, nextList);
+        next.docFreq.set(term, nextList.length);
+      }
+    }
+    return next;
+  }
+
   index(tokenizedDocs: string[][]): void {
     this.numDocs = 0;
     this.docFreq.clear();

package/src/index/create.ts

@@ -136,9 +136,6 @@ export async function createIndexFromPath(
       }
 
       const take = Math.min(embedBatchSize, pendingWindows.length);
-      if (!force && take < embedBatchSize) {
-        break;
-      }
 
       const batch = pendingWindows.splice(0, take);
       emittedApiBatches++;
@@ -158,9 +155,10 @@ export async function createIndexFromPath(
       })();
       embedPromises.push(promise);
       embedInFlight.add(promise);
-      promise.finally(() => {
-        embedInFlight.delete(promise);
-      });
+      promise.then(
+        () => embedInFlight.delete(promise),
+        () => embedInFlight.delete(promise),
+      );
     }
   };
 
@@ -209,9 +207,7 @@ export async function createIndexFromPath(
   await maybeScheduleEmbed(true);
   await Promise.all(embedPromises);
   // Finish any BM25 docs not yet processed during idle gaps.
-  while (bm25Cursor < emittedChunks.length) {
-    progressBm25(50);
-  }
+  progressBm25(Number.POSITIVE_INFINITY);
   const embedsDoneAt = performance.now();
 
   const chunks = emittedChunks;

package/src/index/files.ts

@@ -52,6 +52,9 @@ const EXTENSION_TO_LANGUAGE: Record<string, string> = {
   ".zig": "zig",
   ".vue": "vue",
   ".svelte": "svelte",
+  ".astro": "astro",
+  ".erb": "embeddedtemplate",
+  ".ejs": "embeddedtemplate",
   ".md": "markdown",
   ".mdx": "markdown",
   ".rst": "rst",
@@ -64,6 +67,7 @@ const EXTENSION_TO_LANGUAGE: Record<string, string> = {
   ".cfg": "ini",
   ".xml": "xml",
   ".html": "html",
+  ".htm": "html",
   ".css": "css",
   ".scss": "scss",
   ".less": "less",
@@ -74,7 +78,7 @@ const CODE_EXTENSIONS = new Set(
   Object.entries(EXTENSION_TO_LANGUAGE)
     .filter(
       ([, lang]) =>
-        !["markdown", "rst", "text", "json", "yaml", "toml", "ini", "xml", "html"].includes(lang),
+        !["markdown", "rst", "text", "json", "yaml", "toml", "ini", "xml"].includes(lang),
     )
     .map(([ext]) => ext),
 );

package/src/index/incremental.ts

@@ -5,7 +5,7 @@ import type { Chunk, ContentType } from "../types.ts";
 import type { BM25Index } from "./bm25.ts";
 import { detectLanguage, getExtensions, getFileStatus, readFileText } from "./files.ts";
 import type { SemanticIndex } from "./semantic-index.ts";
-import { buildBm25FromChunks } from "./sparse.ts";
+import { buildBm25FromChunks, patchBm25 } from "./sparse.ts";
 import { buildSemanticIndex } from "./vector-storage.ts";
 import { vectorAt } from "./vectors.ts";
 
@@ -40,9 +40,34 @@ function isIndexableRelativePath(relativePath: string, extensions: Set<string>):
   return extensions.has(ext);
 }
 
+/** Vector for every text: reused from `known`, otherwise embedded once (duplicates share a call). */
+async function vectorsByText(
+  embeddings: EmbeddingBackend,
+  texts: readonly string[],
+  known: ReadonlyMap<string, Float32Array>,
+): Promise<Map<string, Float32Array>> {
+  const result = new Map(known);
+  const missing = [...new Set(texts)].filter((text) => !result.has(text));
+  if (missing.length === 0) {
+    return result;
+  }
+  const embedded = await embeddings.embedDocuments(missing);
+  if (embedded.length !== missing.length) {
+    throw new Error(
+      `Vector count ${embedded.length} does not match chunk count ${missing.length}.`,
+    );
+  }
+  for (const [i, text] of missing.entries()) {
+    result.set(text, embedded[i] as Float32Array);
+  }
+  return result;
+}
+
 /**
  * Replace chunks for the given repo-relative paths only: remove old chunks,
- * embed new ones, rebuild BM25 + semantic indexes from the merged vector set.
+ * embed new ones, rebuild the semantic index from the merged vector set, and update
+ * BM25 by dropping the old docs and indexing only the new chunks. Pass `bm25` (the
+ * index matching `chunks`) to enable that; without it BM25 is rebuilt from scratch.
  */
 export async function applyIncrementalFileChanges(options: {
   root: string;
@@ -50,6 +75,7 @@ export async function applyIncrementalFileChanges(options: {
   embeddings: EmbeddingBackend;
   chunks: Chunk[];
   semanticIndex: SemanticIndex;
+  bm25?: BM25Index;
   relativePaths: readonly string[];
 }): Promise<{ chunks: Chunk[]; bm25: BM25Index; semantic: SemanticIndex }> {
   const root = resolve(options.root);
@@ -60,16 +86,24 @@ export async function applyIncrementalFileChanges(options: {
       .filter((p) => p.length > 0 && isIndexableRelativePath(p, extensions)),
   );
 
+  // BM25 doc i is chunk i; an index of another size can't be patched safely.
+  const baseBm25 = options.bm25?.size === options.chunks.length ? options.bm25 : undefined;
+
   if (targets.size === 0) {
     return {
       chunks: options.chunks,
-      bm25: buildBm25FromChunks(options.chunks),
+      bm25: baseBm25 ?? buildBm25FromChunks(options.chunks),
       semantic: options.semanticIndex,
     };
   }
 
   const keptChunks: Chunk[] = [];
   const keptVectors: Float32Array[] = [];
+  const removedDocs = new Set<number>();
+  // An embedding is a pure function of chunk text, so a re-chunked file can reuse
+  // the vector of any old chunk with byte-identical content (keyed on the exact
+  // string, never a hash, so a collision can't leave a stale vector).
+  const reusableVectors = new Map<string, Float32Array>();
 
   for (let i = 0; i < options.chunks.length; i++) {
     const chunk = options.chunks[i];
@@ -78,6 +112,10 @@ export async function applyIncrementalFileChanges(options: {
     }
     const rel = normalizeRelativePath(chunk.file_path);
     if (targets.has(rel)) {
+      removedDocs.add(i);
+      if (!reusableVectors.has(chunk.content)) {
+        reusableVectors.set(chunk.content, vectorAt(options.semanticIndex, i));
+      }
       continue;
     }
     keptChunks.push(chunk);
@@ -89,10 +127,12 @@ export async function applyIncrementalFileChanges(options: {
     addedChunks.push(...(await chunksForFile(root, rel)));
   }
 
-  let addedVectors: Float32Array[] = [];
-  if (addedChunks.length > 0) {
-    addedVectors = await options.embeddings.embedDocuments(addedChunks.map((c) => c.content));
-  }
+  const vectorFor = await vectorsByText(
+    options.embeddings,
+    addedChunks.map((c) => c.content),
+    reusableVectors,
+  );
+  const addedVectors = addedChunks.map((c) => vectorFor.get(c.content) as Float32Array);
 
   const chunks = [...keptChunks, ...addedChunks];
   const vectors = [...keptVectors, ...addedVectors];
@@ -107,7 +147,7 @@ export async function applyIncrementalFileChanges(options: {
 
   return {
     chunks,
-    bm25: buildBm25FromChunks(chunks),
+    bm25: baseBm25 ? patchBm25(baseBm25, removedDocs, addedChunks) : buildBm25FromChunks(chunks),
     semantic: buildSemanticIndex(vectors),
   };
 }

package/src/index/sparse.ts

@@ -30,10 +30,22 @@ export function addChunkToBm25(bm25: BM25Index, chunk: Chunk): void {
   bm25.addDocument(tokenize(enrichForBm25(chunk)));
 }
 
-export function buildBm25FromChunks(chunks: readonly Chunk[]): BM25Index {
-  const bm25 = new BM25Index();
+function appendChunksToBm25(bm25: BM25Index, chunks: readonly Chunk[]): BM25Index {
   for (const chunk of chunks) {
     addChunkToBm25(bm25, chunk);
   }
   return bm25;
 }
+
+export function buildBm25FromChunks(chunks: readonly Chunk[]): BM25Index {
+  return appendChunksToBm25(new BM25Index(), chunks);
+}
+
+/** Drop docs `removedDocs` from `base` and append `added`; `base` is left untouched. */
+export function patchBm25(
+  base: BM25Index,
+  removedDocs: ReadonlySet<number>,
+  added: readonly Chunk[],
+): BM25Index {
+  return appendChunksToBm25(base.withoutDocuments(removedDocs), added);
+}