@tensor-cad/mcp 0.1.2 → 0.1.4

package/index.js

@@ -310,7 +310,12 @@ Content-Length: 0\r
         setTimeout(() => socket.destroy(), 50).unref();
         return;
       }
-      this.wss.handleUpgrade(req, socket, head, (ws) => this.attach(ws));
+      this.wss.handleUpgrade(req, socket, head, (ws) => {
+        this.attach(ws).catch((e) => {
+          this.log(`tensorcad bridge: could not greet a client: ${e.message}`);
+          this.send(ws, { type: "error", message: e.message });
+        });
+      });
     });
   }
   get port() {
@@ -384,14 +389,14 @@ Content-Length: 0\r
     }
     return;
   }
-  attach(ws) {
+  async attach(ws) {
     this.send(ws, {
       type: "hello",
       protocol: BRIDGE_PROTOCOL,
       server: this.options.name,
       version: this.options.version,
       root: this.options.root,
-      designs: this.options.store.list()
+      designs: await this.options.store.list()
     });
     ws.on("message", (raw) => {
       let message;
@@ -401,23 +406,21 @@ Content-Length: 0\r
         this.send(ws, { type: "error", message: `not JSON: ${e.message}` });
         return;
       }
-      try {
-        this.handle(ws, message);
-      } catch (e) {
+      this.handle(ws, message).catch((e) => {
         this.send(ws, { type: "error", message: e.message, about: message?.type });
-      }
+      });
     });
   }
-  handle(ws, message) {
+  async handle(ws, message) {
     switch (message?.type) {
       case "publish": {
         const doc = asDocument(message.doc);
-        const record = this.during(ws, () => this.options.store.adopt(doc));
+        const record = await this.during(ws, () => this.options.store.adopt(doc));
         this.send(ws, designMessage(record, "published"));
         return;
       }
       case "attach": {
-        const record = this.options.store.get(message.design_id);
+        const record = await this.options.store.get(message.design_id);
         this.send(ws, designMessage(record, "requested"));
         return;
       }
@@ -425,12 +428,12 @@ Content-Length: 0\r
       case "replace": {
         const write = message.type === "ops" ? () => this.options.store.apply(message.design_id, parseOps(message.ops), message.revision) : () => this.options.store.replace(message.design_id, asDocument(message.doc), message.revision);
         try {
-          const { record } = this.during(ws, write);
+          const { record } = await this.during(ws, write);
           this.send(ws, designMessage(record, message.type === "ops" ? "applied" : "replaced"));
         } catch (e) {
           if (e instanceof RevisionConflictError) {
             this.send(ws, { type: "error", message: e.message, about: message.type });
-            this.send(ws, designMessage(this.options.store.get(message.design_id), "requested"));
+            this.send(ws, designMessage(await this.options.store.get(message.design_id), "requested"));
             return;
           }
           throw e;
@@ -460,13 +463,11 @@ Content-Length: 0\r
       }
     }
   }
-  during(ws, work) {
+  async during(ws, work) {
     this.acting = ws;
-    try {
-      return work();
-    } finally {
-      this.acting = undefined;
-    }
+    const running = work();
+    this.acting = undefined;
+    return await running;
   }
   send(ws, message) {
     if (ws.readyState !== ws.OPEN)
@@ -549,7 +550,7 @@ import { McpServer as McpServer2 } from "@modelcontextprotocol/server";
 // packages/mcp/package.json
 var package_default = {
   name: "@tensor-cad/mcp",
-  version: "0.1.2",
+  version: "0.1.4",
   description: "Model Context Protocol server for TensorCAD: design, validate, analyze and generate LLM architectures from an agent",
   mcpName: "io.github.filip-pajalic/tensorcad",
   type: "module",
@@ -824,7 +825,7 @@ class FileStore {
       }
     }
   }
-  list() {
+  async list() {
     return [...this.entries.values()].map((e) => summaryOf(e.record)).sort((a, b) => b.updated_at.localeCompare(a.updated_at));
   }
   async listFiles() {
@@ -850,13 +851,13 @@ class FileStore {
     await walk(this.root, this.depth);
     return out.sort();
   }
-  get(id) {
+  async get(id) {
     const entry = this.entries.get(id);
     if (!entry)
       throw new UnknownDesignError(id, [...this.entries.keys()]);
     return entry.record;
   }
-  create(options) {
+  async create(options) {
     let doc;
     let source;
     if (options.preset) {
@@ -873,7 +874,7 @@ class FileStore {
       doc.meta.name = options.name;
     return this.register(doc, source, undefined, true);
   }
-  adopt(doc) {
+  async adopt(doc) {
     return this.register(doc, "derived", undefined, true);
   }
   async open(path) {
@@ -915,7 +916,7 @@ class FileStore {
     this.emit({ kind: "registered", record });
     return record;
   }
-  apply(id, ops, expectedRevision) {
+  async apply(id, ops, expectedRevision) {
     const entry = this.entry(id);
     const { record } = entry;
     if (expectedRevision !== undefined && expectedRevision !== record.revision) {
@@ -933,7 +934,7 @@ class FileStore {
     this.emit({ kind: "applied", record, ops });
     return { record, applied, previousRevision };
   }
-  replace(id, doc, expectedRevision) {
+  async replace(id, doc, expectedRevision) {
     const entry = this.entry(id);
     const { record } = entry;
     if (expectedRevision !== undefined && expectedRevision !== record.revision) {
@@ -950,7 +951,7 @@ class FileStore {
     return { record, applied: ["replaced the document"], previousRevision };
   }
   async save(id, path) {
-    const record = this.get(id);
+    const record = await this.get(id);
     const target = path ? this.resolvePath(path) : record.path ?? join2(this.root, `${slug(record.name)}.tensorcad.json`);
     const text = `${JSON.stringify(record.doc, null, 2)}
 `;
@@ -962,7 +963,7 @@ class FileStore {
     this.emit({ kind: "saved", record });
     return { record, path: target, bytes: Buffer.byteLength(text, "utf8") };
   }
-  checkpoint(id, label) {
+  async checkpoint(id, label) {
     const entry = this.entry(id);
     const info = {
       checkpoint_id: `ckpt_${this.nextCheckpoint++}`,
@@ -973,11 +974,11 @@ class FileStore {
     entry.checkpoints.set(info.checkpoint_id, { ...info, doc: structuredClone(entry.record.doc) });
     return info;
   }
-  checkpoints(id) {
+  async checkpoints(id) {
     const entry = this.entry(id);
     return [...entry.checkpoints.values()].map(({ doc: _doc, ...info }) => info);
   }
-  restore(id, checkpointId) {
+  async restore(id, checkpointId) {
     const entry = this.entry(id);
     const { record } = entry;
     let doc;
@@ -1034,6 +1035,25 @@ function slug(name) {
   return name.replace(/[^a-zA-Z0-9._-]+/g, "-").replace(/^-+|-+$/g, "") || "design";
 }
 
+// packages/mcp/src/store/disk-artifacts.ts
+import { mkdir as mkdir3, writeFile as writeFile3 } from "node:fs/promises";
+import { dirname as dirname3, isAbsolute as isAbsolute2, join as join3, resolve as resolve2 } from "node:path";
+
+class DiskSink {
+  root;
+  label = "disk";
+  constructor(root) {
+    this.root = root;
+  }
+  async write(prefix, path, contents) {
+    const base = isAbsolute2(prefix) ? prefix : resolve2(this.root, prefix);
+    const target = join3(base, path);
+    await mkdir3(dirname3(target), { recursive: true });
+    await writeFile3(target, contents, "utf8");
+    return { location: target };
+  }
+}
+
 // packages/mcp/src/prompts.ts
 import * as z2 from "zod";
 import { HARDWARE, PRESET_NAMES as PRESET_NAMES2 } from "@tensor-cad/engine/node";
@@ -1500,6 +1520,7 @@ function catalogEntry(def) {
     type: def.type,
     kind: def.kind,
     category: def.category,
+    name: def.docs.name ?? def.type,
     summary: def.docs.summary ?? "",
     refs: def.docs.refs ?? [],
     params: def.paramOrder.map((name) => catalogParam(name, def.params[name])),
@@ -1538,7 +1559,10 @@ function allCatalogEntries() {
 }
 function catalogText(entries) {
   return entries.map((e) => {
-    const lines = [`${e.type}  (${e.kind}/${e.category})`, `  ${e.summary}`];
+    const lines = [
+      e.name && e.name !== e.type ? `${e.type}  — ${e.name}  (${e.kind}/${e.category})` : `${e.type}  (${e.kind}/${e.category})`,
+      `  ${e.summary}`
+    ];
     if (e.formula)
       lines.push(`  formula: ${e.formula}`);
     if (e.params.length > 0) {
@@ -1696,8 +1720,8 @@ function registerResources(server, store) {
     title: "Design rule report",
     description: "Every finding for a design: shape errors, memory fit, kernel constraints, Chinchilla sanity.",
     mimeType: JSON_MIME
-  }, (uri, { id }) => {
-    const record = store.get(String(id));
+  }, async (uri, { id }) => {
+    const record = await store.get(String(id));
     const report = validate(record.doc);
     return json(uri, {
       design_id: record.design_id,
@@ -1712,14 +1736,14 @@ function registerResources(server, store) {
     title: "Design analysis",
     description: "Parameters, FLOPs, KV cache, memory, throughput and cost at the document's own defaults.",
     mimeType: JSON_MIME
-  }, (uri, { id }) => {
-    const record = store.get(String(id));
+  }, async (uri, { id }) => {
+    const record = await store.get(String(id));
     const result = analyze(record.doc);
     return json(uri, { design_id: record.design_id, revision: record.revision, ...analysisJson(result) });
   });
   server.registerResource("design", new ResourceTemplate("tensorcad://designs/{id}", {
-    list: () => ({
-      resources: store.list().map((d) => ({
+    list: async () => ({
+      resources: (await store.list()).map((d) => ({
         uri: `tensorcad://designs/${d.design_id}`,
         name: d.name,
         title: `${d.name} (revision ${d.revision})`,
@@ -1732,8 +1756,8 @@ function registerResources(server, store) {
     title: "Design document",
     description: "The literal .tensorcad.json document, with a compact outline beside it.",
     mimeType: JSON_MIME
-  }, (uri, { id }) => {
-    const record = store.get(String(id));
+  }, async (uri, { id }) => {
+    const record = await store.get(String(id));
     return json(uri, {
       design_id: record.design_id,
       revision: record.revision,
@@ -1747,7 +1771,7 @@ function registerResources(server, store) {
     title: "Block catalog",
     description: "Every block type with its parameter schema, port patterns and documentation. " + "Primitives carry the formulas; composites expand into primitives; repeat is the only container.",
     mimeType: JSON_MIME
-  }, (uri) => {
+  }, async (uri) => {
     const blocks = allCatalogEntries();
     return json(uri, {
       count: blocks.length,
@@ -1756,7 +1780,7 @@ function registerResources(server, store) {
     });
   });
   server.registerResource("catalog-block", new ResourceTemplate("tensorcad://catalog/{type}", {
-    list: () => ({
+    list: async () => ({
       resources: Object.keys(CATALOG).sort().map((type) => ({
         uri: `tensorcad://catalog/${type}`,
         name: type,
@@ -1791,12 +1815,10 @@ function registerResources(server, store) {
   }));
 }
 function completeDesignId(store) {
-  return (value) => store.list().map((d) => d.design_id).filter((id) => id.startsWith(value));
+  return async (value) => (await store.list()).map((d) => d.design_id).filter((id) => id.startsWith(value));
 }
 
 // packages/mcp/src/tools.ts
-import { mkdir as mkdir3, writeFile as writeFile3 } from "node:fs/promises";
-import { dirname as dirname3, isAbsolute as isAbsolute2, join as join3, resolve as resolve2 } from "node:path";
 import * as z3 from "zod";
 import { formatBytes as formatBytes2, formatCount as formatCount2 } from "@tensor-cad/engine";
 import {
@@ -1873,7 +1895,7 @@ function toAnalysisOptions(input) {
 var READ = { readOnlyHint: true, idempotentHint: true, openWorldHint: false };
 var WRITE = { readOnlyHint: false, idempotentHint: false, openWorldHint: false };
 var DESTRUCTIVE = { readOnlyHint: false, idempotentHint: false, destructiveHint: true, openWorldHint: false };
-function registerTools(server, store) {
+function registerTools(server, store, artifacts) {
   server.registerTool("tensorcad_list_designs", {
     title: "List designs",
     description: "List the designs this server has open, the built-in reference architectures you can start from, " + "and the .tensorcad.json files it can see on disk. Start here when you do not already hold a design_id.",
@@ -1892,7 +1914,7 @@ function registerTools(server, store) {
     }),
     annotations: { ...READ, title: "List designs" }
   }, async ({ include_files }) => guard(async () => {
-    const designs = store.list();
+    const designs = await store.list();
     const presets = PRESET_NAMES3.map((name) => {
       const doc = getPreset2(name);
       const p = { name };
@@ -1933,8 +1955,8 @@ ${designs.map((d) => `  ${d.design_id}  ${d.name}  rev ${d.revision}${d.dirty ?
       outline: Outline
     }),
     annotations: { ...WRITE, title: "New design" }
-  }, async (args) => guard(() => {
-    const record = store.create(args);
+  }, async (args) => guard(async () => {
+    const record = await store.create(args);
     const outline = outlineOf(record.doc);
     return ok(`${record.design_id} (revision ${record.revision})
 
@@ -2019,8 +2041,8 @@ ${formatCount2(report.analysis.params.total)} parameters, ` + `${report.counts.e
       document: z3.record(z3.string(), z3.unknown()).optional().describe("The literal design document.")
     }),
     annotations: { ...READ, title: "Get design" }
-  }, async ({ design_id, format }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, format }) => guard(async () => {
+    const record = await store.get(design_id);
     const outline = outlineOf(record.doc);
     const mode = format ?? "outline";
     const base = {
@@ -2074,8 +2096,8 @@ ${outlineText(outline)}`, {
       children: z3.array(z3.string())
     }),
     annotations: { ...READ, title: "Get block" }
-  }, async ({ design_id, path }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, path }) => guard(async () => {
+    const record = await store.get(design_id);
     const detail = blockDetail(record.doc, path);
     return ok(blockText(detail), {
       design_id: record.design_id,
@@ -2099,6 +2121,7 @@ ${outlineText(outline)}`, {
         type: z3.string(),
         kind: z3.string(),
         category: z3.string(),
+        name: z3.string().describe("What a drawing calls this block."),
         summary: z3.string(),
         formula: z3.string().optional(),
         refs: z3.array(z3.string()),
@@ -2115,7 +2138,7 @@ ${outlineText(outline)}`, {
       }))
     }),
     annotations: { ...READ, title: "Search catalog" }
-  }, async ({ query, category, kind, limit }) => guard(() => {
+  }, async ({ query, category, kind, limit }) => guard(async () => {
     const all = allCatalogEntries();
     const q = query?.toLowerCase();
     const matched = all.filter((e) => {
@@ -2125,7 +2148,7 @@ ${outlineText(outline)}`, {
         return false;
       if (!q)
         return true;
-      const hay = `${e.type} ${e.category} ${e.summary} ${e.formula ?? ""}`.toLowerCase();
+      const hay = `${e.type} ${e.name} ${e.category} ${e.summary} ${e.formula ?? ""}`.toLowerCase();
       return hay.includes(q);
     });
     const blocks = matched.slice(0, limit ?? 20);
@@ -2159,10 +2182,10 @@ ${catalogText(blocks)}`, {
       validation: ValidationSummary
     }),
     annotations: { ...WRITE, title: "Apply edits" }
-  }, async ({ design_id, expected_revision, ops }) => guard(() => {
-    const before = store.get(design_id);
+  }, async ({ design_id, expected_revision, ops }) => guard(async () => {
+    const before = await store.get(design_id);
     const paramsBefore = outlineOf(before.doc).params_total;
-    const outcome = store.apply(design_id, ops, expected_revision);
+    const outcome = await store.apply(design_id, ops, expected_revision);
     const report = validate2(outcome.record.doc);
     const total = report.analysis.params.total;
     const delta = total - paramsBefore;
@@ -2205,8 +2228,8 @@ ${catalogText(blocks)}`, {
       params_total: z3.number()
     }),
     annotations: { ...READ, title: "Validate design" }
-  }, async ({ design_id, severity, ...rest }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, severity, ...rest }) => guard(async () => {
+    const record = await store.get(design_id);
     const report = validate2(record.doc, toAnalysisOptions(rest));
     const rank = { error: 0, warning: 1, info: 2 };
     const findings = findingsJson(report).filter((f) => severity === undefined || rank[f.severity] <= rank[severity]);
@@ -2233,8 +2256,8 @@ ${catalogText(blocks)}`, {
     inputSchema: z3.object({ design_id: DESIGN_ID, ...analysisOptionsShape }),
     outputSchema: AnalysisOutput,
     annotations: { ...READ, title: "Analyze design" }
-  }, async ({ design_id, ...rest }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, ...rest }) => guard(async () => {
+    const record = await store.get(design_id);
     const result = analyze2(record.doc, toAnalysisOptions(rest));
     return ok(analysisText(result), {
       design_id: record.design_id,
@@ -2267,30 +2290,36 @@ ${catalogText(blocks)}`, {
     }),
     annotations: { readOnlyHint: false, idempotentHint: true, openWorldHint: false, title: "Generate code" }
   }, async ({ design_id, class_name, include_smoke_test, out_dir }) => guard(async () => {
-    const record = store.get(design_id);
+    const record = await store.get(design_id);
     const generated = generateTorch(record.doc, {
       ...class_name ? { className: class_name } : {},
       includeSmokeTest: include_smoke_test ?? false
     });
-    const root = out_dir ? isAbsolute2(out_dir) ? out_dir : resolve2(process.cwd(), out_dir) : undefined;
     const files = [];
     for (const file of generated.files) {
-      const bytes = Buffer.byteLength(file.contents, "utf8");
+      const bytes = new TextEncoder().encode(file.contents).length;
       const lines = file.contents.split(`
 `).length;
-      if (root) {
-        const target = join3(root, file.path);
-        await mkdir3(dirname3(target), { recursive: true });
-        await writeFile3(target, file.contents, "utf8");
-        files.push({ path: file.path, bytes, lines, written_to: target });
-      } else {
+      if (out_dir === undefined) {
         files.push({ path: file.path, bytes, lines, contents: file.contents });
+        continue;
       }
+      const written = await artifacts.write(out_dir, file.path, file.contents);
+      files.push({
+        path: file.path,
+        bytes,
+        lines,
+        written_to: written.location,
+        ...written.url ? { url: written.url } : {}
+      });
     }
-    const text = root ? [`wrote ${files.length} file(s) to ${root}`, ...files.map((f) => `  ${f.path}  ${f.bytes} bytes`)].join(`
-`) : generated.files.map((f) => `# ${f.path}
+    const text = out_dir === undefined ? generated.files.map((f) => `# ${f.path}
 ${f.contents}`).join(`
 
+`) : [
+      `wrote ${files.length} file(s) to ${artifacts.label}`,
+      ...files.map((f) => `  ${f.path}  ${f.bytes} bytes  ${f.url ?? f.written_to}`)
+    ].join(`
 `);
     return ok(generated.warnings.length > 0 ? `${text}
 
@@ -2299,8 +2328,8 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
 `)}` : text, {
       design_id: record.design_id,
       revision: record.revision,
-      wrote: Boolean(root),
-      ...root ? { out_dir: root } : {},
+      wrote: out_dir !== undefined,
+      ...out_dir !== undefined ? { out_dir } : {},
       files,
       warnings: generated.warnings
     });
@@ -2326,12 +2355,12 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       }))
     }),
     annotations: { ...WRITE, title: "Checkpoint design" }
-  }, async ({ design_id, label }) => guard(() => {
-    const info = store.checkpoint(design_id, label);
+  }, async ({ design_id, label }) => guard(async () => {
+    const info = await store.checkpoint(design_id, label);
     return ok(`${info.checkpoint_id} at revision ${info.revision}: ${info.label}`, {
       design_id,
       ...info,
-      checkpoints: store.checkpoints(design_id)
+      checkpoints: await store.checkpoints(design_id)
     });
   }));
   server.registerTool("tensorcad_restore", {
@@ -2350,8 +2379,8 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       validation: ValidationSummary
     }),
     annotations: { ...DESTRUCTIVE, title: "Restore design" }
-  }, async ({ design_id, checkpoint_id }) => guard(() => {
-    const { record, restoredFrom } = store.restore(design_id, checkpoint_id);
+  }, async ({ design_id, checkpoint_id }) => guard(async () => {
+    const { record, restoredFrom } = await store.restore(design_id, checkpoint_id);
     const report = validate2(record.doc);
     return ok(`${record.design_id} restored from ${restoredFrom}; now revision ${record.revision}, ` + `${formatCount2(report.analysis.params.total)} parameters`, {
       design_id: record.design_id,
@@ -2394,8 +2423,8 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       })
     }),
     annotations: { ...READ, title: "Explain a block" }
-  }, async ({ design_id, path, ...rest }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, path, ...rest }) => guard(async () => {
+    const record = await store.get(design_id);
     const e = explain(record.doc, path, toAnalysisOptions(rest));
     const lines = [
       `${path}  ${e.type} (${e.kind})`,
@@ -2456,8 +2485,8 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       notes: z3.array(z3.string())
     }),
     annotations: { readOnlyHint: false, idempotentHint: true, openWorldHint: false, title: "Scale a design" }
-  }, async ({ design_id, target_params, target_basis, vocab, tie_head, keep_depth }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, target_params, target_basis, vocab, tie_head, keep_depth }) => guard(async () => {
+    const record = await store.get(design_id);
     const result = scaleDesign(record.doc, {
       targetParams: target_params,
       ...target_basis ? { targetBasis: target_basis } : {},
@@ -2465,7 +2494,7 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       ...tie_head !== undefined ? { tieHead: tie_head } : {},
       ...keep_depth !== undefined ? { keepDepth: keep_depth } : {}
     });
-    const saved = store.adopt(result.doc);
+    const saved = await store.adopt(result.doc);
     const changes = Object.entries(result.changes).map(([symbol, c]) => ({
       symbol,
       from: c.from,
@@ -2518,28 +2547,32 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       notes: z3.array(z3.string())
     }),
     annotations: { readOnlyHint: false, idempotentHint: true, openWorldHint: false, title: "Build a width ladder" }
-  }, async ({ design_id, widths, base_width }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, widths, base_width }) => guard(async () => {
+    const record = await store.get(design_id);
     const ladder = mupLadder(record.doc, {
       ...widths ? { widths } : {},
       ...base_width !== undefined ? { baseWidth: base_width } : {}
     });
-    const rungs = ladder.rungs.map((rung) => ({
-      design_id: store.adopt(rung.doc).design_id,
-      width: rung.width,
-      multiplier: rung.multiplier,
-      heads: rung.heads,
-      params: rung.params,
-      base: rung.base,
-      scaling: rung.scaling.map((s) => ({
-        class: s.class,
-        init_std: s.initStd,
-        adam_lr: s.adamLr,
-        paths: s.paths,
-        why: s.why
-      })),
-      notes: rung.notes
-    }));
+    const rungs = [];
+    for (const rung of ladder.rungs) {
+      const adopted = await store.adopt(rung.doc);
+      rungs.push({
+        design_id: adopted.design_id,
+        width: rung.width,
+        multiplier: rung.multiplier,
+        heads: rung.heads,
+        params: rung.params,
+        base: rung.base,
+        scaling: rung.scaling.map((s) => ({
+          class: s.class,
+          init_std: s.initStd,
+          adam_lr: s.adamLr,
+          paths: s.paths,
+          why: s.why
+        })),
+        notes: rung.notes
+      });
+    }
     const text = [
       `${record.doc.meta.name} laddered by ${ladder.widthSymbol}, tuned at ${ladder.baseWidth}, ` + `heads of ${ladder.headDim} throughout`,
       ...rungs.map((r) => `  ${r.base ? "base " : "     "}${r.width} wide, ${r.heads} heads, ${formatCount2(r.params)}` + `  ${r.scaling.filter((s) => s.class !== "input").map((s) => `${s.class} init x${s.init_std.toPrecision(4)} rate x${s.adam_lr.toPrecision(4)}`).join(", ")}`),
@@ -2588,8 +2621,8 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       notes: z3.array(z3.string())
     }),
     annotations: { ...READ, title: "Plan a cluster" }
-  }, async ({ design_id, gpus, gpus_per_node, headroom, limit, ...rest }) => guard(() => {
-    const record = store.get(design_id);
+  }, async ({ design_id, gpus, gpus_per_node, headroom, limit, ...rest }) => guard(async () => {
+    const record = await store.get(design_id);
     const result = planCluster(record.doc, toAnalysisOptions(rest), {
       gpus,
       ...gpus_per_node !== undefined ? { gpusPerNode: gpus_per_node } : {},
@@ -2665,9 +2698,9 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       }))
     }),
     annotations: { ...READ, title: "Compare two designs" }
-  }, async ({ a, b, ...rest }) => guard(() => {
-    const left = store.get(a);
-    const right = store.get(b);
+  }, async ({ a, b, ...rest }) => guard(async () => {
+    const left = await store.get(a);
+    const right = await store.get(b);
     const d = diffDesigns(left.doc, right.doc, toAnalysisOptions(rest));
     const brief = (v) => {
       if (v === undefined || v === null)
@@ -2733,9 +2766,9 @@ ${generated.warnings.map((w) => `  ${w}`).join(`
       warnings: z3.array(z3.string())
     }),
     annotations: { readOnlyHint: false, idempotentHint: true, openWorldHint: false, title: "Import a config" }
-  }, async ({ config, name }) => guard(() => {
+  }, async ({ config, name }) => guard(async () => {
     const result = importHfConfig(config, name);
-    const record = store.adopt(result.doc);
+    const record = await store.adopt(result.doc);
     const total = analyze2(result.doc).params.total;
     const text = [
       `${result.doc.meta.name}: ${formatCount2(total)} parameters`,
@@ -2798,13 +2831,14 @@ var INSTRUCTIONS = [
 ].join(`
 `);
 function createServer2(options = {}) {
-  const { store: given, bridge, ...storeOptions } = options;
+  const { store: given, bridge, artifacts: givenArtifacts, ...storeOptions } = options;
   const store = given ?? new FileStore(storeOptions);
+  const artifacts = givenArtifacts ?? new DiskSink(storeOptions.root ?? process.cwd());
   const server = new McpServer2({ name: SERVER_NAME, version: SERVER_VERSION }, {
     capabilities: { tools: {}, resources: {}, prompts: {}, completions: {} },
     instructions: INSTRUCTIONS
   });
-  registerTools(server, store);
+  registerTools(server, store, artifacts);
   registerResources(server, store);
   registerPrompts(server);
   if (bridge)
@@ -2863,11 +2897,14 @@ async function startBridge(store, root) {
   return bridge;
 }
 export {
+  DiskSink,
   FileStore,
   PROMPT_NAMES,
+  RevisionConflictError,
   SERVER_NAME,
   SERVER_VERSION,
   TOOL_NAMES,
+  UnknownDesignError,
   applyOps,
   createServer2 as createServer,
   serve

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@tensor-cad/mcp",
-  "version": "0.1.2",
+  "version": "0.1.4",
   "description": "Model Context Protocol server for TensorCAD: design, validate, analyze and generate LLM architectures from an agent",
   "mcpName": "io.github.filip-pajalic/tensorcad",
   "type": "module",
@@ -25,7 +25,7 @@
   },
   "dependencies": {
     "@modelcontextprotocol/server": "^2.0.0",
-    "@tensor-cad/engine": "0.1.2",
+    "@tensor-cad/engine": "0.1.4",
     "ws": "^8.21.3",
     "zod": "^4.2.0"
   },