@tanstack-ai-browser/afm 0.0.0-stage → 0.2.1

package/src/electron.ts

@@ -0,0 +1,180 @@
+/**
+ * Apple's on-device model in an Electron app on macOS: `serveAfm()` in the
+ * main process runs it, through the package's Node-API addon
+ * (electron/afm.node, built from the same Swift as the Capacitor plugin), and
+ * adds the package's preload (electron/preload.cjs) to the app's pages, which
+ * puts it at `window.afm` for `afmText()`.
+ *
+ * Each frame's sessions are its own: their native ids are the frame's id and
+ * the page's, a frame's calls reach only its sessions, and their events go to
+ * that frame only, while `allow()` still allows its page. A frame's sessions
+ * end when it navigates or closes.
+ *
+ * Electron is required when `serveAfm()` runs, not imported: bundlers of the
+ * app's pages (web, iOS) never see it, and drop this code.
+ */
+import type { App, IpcMain, Session, WebContents, WebFrameMain } from "electron"
+
+const PACKAGE = "@tanstack-ai-browser/afm"
+const CALL = "afm:call"
+const EVENT = "afm:event"
+
+/** The addon (electron/afm.c): AfmBridge's calls, and its events, in JSON. */
+type Addon = {
+  call(method: string, options: string): string
+  setListener(listener: ((json: string) => void) | null): void
+}
+
+type Electron = {
+  app: App
+  ipcMain: IpcMain
+  session: { defaultSession: Session }
+}
+
+/** A native session's frame, and the page's id for it. */
+type Owner = { frame: WebFrameMain; frameId: string; contentsId: number; sessionId: string }
+
+/** The bits of Node used here, typed here: this is a browser package otherwise. */
+type NodeRequire = ((id: string) => unknown) & { resolve(id: string): string }
+type NodeProcess = {
+  platform: string
+  execPath: string
+  getBuiltinModule(id: "node:module"): { createRequire(path: string): NodeRequire }
+  getBuiltinModule(id: "node:path"): { join(...parts: Array<string>): string; dirname(path: string): string }
+}
+
+export type ServeAfmOptions = {
+  /**
+   * The folder with the package's Electron files (afm.node and preload.cjs):
+   * by default the package's own `electron` folder, found from the app's
+   * node_modules. An app that leaves node_modules out of its package copies
+   * the two files in and passes their folder. Inside an asar archive, the
+   * unpacked afm.node beside it is loaded (native code can't load from an archive).
+   */
+  directory?: string
+  /**
+   * Which pages may use the model, by URL. By default the app's own: any page
+   * but an http(s) one on another host than localhost.
+   */
+  allow?: (url: string) => boolean
+}
+
+const appPages = (url: string) => {
+  try {
+    const { protocol, hostname } = new URL(url)
+    if (protocol !== "http:" && protocol !== "https:") return true
+    return hostname === "localhost" || hostname === "127.0.0.1" || hostname === "[::1]"
+  } catch {
+    return false
+  }
+}
+
+const unpacked = (file: string) => file.replace(/app\.asar(?=[\\/])/, "app.asar.unpacked")
+
+let served = false
+
+/**
+ * Runs Apple's on-device model for the app's pages, where `afmText()` finds it.
+ * Call it once in the main process, before opening windows. On a Mac before
+ * macOS 26, or elsewhere, the pages hear the model is unavailable.
+ */
+export function serveAfm(options: ServeAfmOptions = {}): void {
+  if (served) return
+  served = true
+  const node = (globalThis as unknown as { process: NodeProcess }).process
+  const { createRequire } = node.getBuiltinModule("node:module")
+  const path = node.getBuiltinModule("node:path")
+  // Electron answers `require("electron")` from anywhere in the main process.
+  const electron = createRequire(node.execPath)("electron") as Electron
+  const { app } = electron
+  const allow = options.allow ?? appPages
+
+  let directory = options.directory
+  if (!directory) {
+    try {
+      const fromApp = createRequire(path.join(app.getAppPath(), "package.json"))
+      directory = path.join(path.dirname(fromApp.resolve(`${PACKAGE}/package.json`)), "electron")
+    } catch (error) {
+      console.error(`Apple's on-device model is unavailable: ${PACKAGE} isn't in the app's node_modules (pass serveAfm's directory)`, error)
+    }
+  }
+
+  let addon: Addon | undefined
+  if (directory && node.platform === "darwin") {
+    try {
+      addon = createRequire(node.execPath)(unpacked(path.join(directory, "afm.node"))) as Addon
+    } catch (error) {
+      console.error("Apple's on-device model is unavailable: its addon didn't load", error)
+    }
+  }
+
+  // The native sessions, by native id: `${frame id}/${the page's id}`.
+  const owners = new Map<string, Owner>()
+  const frameIdOf = (frame: WebFrameMain) => `${frame.processId}:${frame.frameToken}`
+  const endSessions = (ended: (owner: Owner) => boolean) => {
+    for (const [id, owner] of owners) {
+      if (!ended(owner)) continue
+      owners.delete(id)
+      addon?.call("deleteSession", JSON.stringify({ sessionId: id }))
+    }
+  }
+  // A page's sessions end with its document: a navigation, or its contents closing.
+  const watched = new WeakSet<WebContents>()
+  const watch = (contents: WebContents) => {
+    if (watched.has(contents)) return
+    watched.add(contents)
+    const contentsId = contents.id
+    contents.on("did-start-navigation", ({ isSameDocument, frame }) => {
+      if (isSameDocument || !frame) return
+      const frameId = frameIdOf(frame)
+      endSessions((owner) => owner.frameId === frameId)
+    })
+    contents.once("destroyed", () => endSessions((owner) => owner.contentsId === contentsId))
+  }
+
+  electron.ipcMain.handle(CALL, (event, method: string, args: unknown) => {
+    const frame = event.senderFrame
+    if (!frame || !allow(frame.url)) throw new Error("This page may not use Apple's on-device model")
+    if (!addon) {
+      if (method === "availability") return { available: false, reason: "unsupportedOS" }
+      throw new Error("Apple's on-device model needs macOS 26 on a Mac with Apple silicon")
+    }
+    const options = { ...(typeof args === "object" && args !== null ? args : {}) } as Record<string, unknown>
+    // The frame's own session: another frame's can't be reached.
+    if (typeof options.sessionId === "string") {
+      const frameId = frameIdOf(frame)
+      const id = `${frameId}/${options.sessionId}`
+      if (method === "createSession") {
+        owners.set(id, { frame, frameId, contentsId: event.sender.id, sessionId: options.sessionId })
+        watch(event.sender)
+      } else if (method === "deleteSession") {
+        owners.delete(id)
+      }
+      options.sessionId = id
+    }
+    const reply = JSON.parse(addon.call(method, JSON.stringify(options))) as { result?: unknown; error?: string }
+    if (reply.error) throw new Error(reply.error)
+    return reply.result ?? null
+  })
+  // Each event to its session's frame only, with the page's id, while its page is allowed.
+  addon?.setListener((json) => {
+    const event = JSON.parse(json) as { sessionId?: unknown }
+    const owner = typeof event.sessionId === "string" ? owners.get(event.sessionId) : undefined
+    if (!owner || owner.frame.isDestroyed() || owner.frame.detached || !allow(owner.frame.url)) return
+    owner.frame.send(EVENT, { ...event, sessionId: owner.sessionId })
+  })
+
+  // The preload, in every session's pages: window.afm.
+  if (!directory) return
+  const preload = path.join(directory, "preload.cjs")
+  const preloaded = new WeakSet<Session>()
+  const addPreload = (session: Session) => {
+    if (preloaded.has(session)) return
+    preloaded.add(session)
+    session.registerPreloadScript({ type: "frame", id: "tanstack-ai-browser-afm", filePath: preload })
+  }
+  // Now if the app is ready, so a window opened next has it.
+  if (app.isReady()) addPreload(electron.session.defaultSession)
+  else void app.whenReady().then(() => addPreload(electron.session.defaultSession))
+  app.on("session-created", addPreload)
+}

package/src/index.ts

@@ -0,0 +1,27 @@
+/**
+ * TanStack AI adapter for Apple's on-device Foundation Model, through the
+ * Foundation Models framework, with its native tool calling: a Capacitor
+ * plugin for iOS (`Afm`), Electron's main process on macOS (`serveAfm()`), or
+ * any host implementing `AfmNative`. The model catalog (model-meta.ts) is
+ * exported here too.
+ */
+export { AfmTextAdapter, afmText } from "./adapter.js"
+export type { Availability, AfmConfig, AfmTextModelOptions } from "./adapter.js"
+export { Afm } from "./native.js"
+export { serveAfm } from "./electron.js"
+export type { ServeAfmOptions } from "./electron.js"
+export type {
+  AfmNative,
+  NativeAvailability,
+  NativeEvent,
+  NativeGenerationSchema,
+  NativeMessage,
+  NativeProperty,
+  NativeSchema,
+  NativeTool,
+  NativeUsage,
+  UnavailableReason
+} from "./native.js"
+export { toNativeSchema } from "./schema.js"
+export { AFM_CHAT_MODELS, AFM_MODEL_META } from "./model-meta.js"
+export type { AfmChatModel, AfmModelMeta } from "./model-meta.js"

package/src/messages.ts

@@ -0,0 +1,37 @@
+/**
+ * TanStack AI's messages as a native session's conversation (`NativeMessage`):
+ * text, the assistant's tool calls, and the tools' results. The framework
+ * reads text only here, so other parts are left out.
+ */
+import type { ModelMessage } from "@tanstack/ai"
+
+import type { NativeMessage } from "./native.js"
+
+/** A message's text: its string, or its text parts joined. */
+export function contentText(content: ModelMessage["content"]): string {
+  if (typeof content === "string") return content
+  if (!Array.isArray(content)) return ""
+  return content
+    .map((part) => (part.type === "text" ? part.content : ""))
+    .filter(Boolean)
+    .join("\n")
+}
+
+/** The messages as a session's conversation, each tool result named by its call's tool. */
+export function toNativeMessages(messages: ReadonlyArray<ModelMessage>): Array<NativeMessage> {
+  const toolNames = new Map<string, string>()
+  return messages.map((m): NativeMessage => {
+    if (m.role === "tool") {
+      const id = m.toolCallId ?? ""
+      return { role: "tool", toolCallId: id, name: m.name ?? toolNames.get(id) ?? "", text: contentText(m.content) }
+    }
+    if (m.role === "assistant") {
+      const toolCalls = (m.toolCalls ?? []).map((call) => {
+        toolNames.set(call.id, call.function.name)
+        return { id: call.id, name: call.function.name, arguments: call.function.arguments || "{}" }
+      })
+      return { role: "assistant", text: contentText(m.content), ...(toolCalls.length > 0 && { toolCalls }) }
+    }
+    return { role: "user", text: contentText(m.content) }
+  })
+}

package/src/model-meta.ts

@@ -0,0 +1,37 @@
+/**
+ * The model the framework runs: Apple's system language model
+ * (`SystemLanguageModel.default`), the on-device model of Apple Intelligence.
+ * It calls tools natively, and holds structured answers to a schema. It doesn't
+ * reason, so `chat()` offers no `reasoning` option for it.
+ */
+
+/** What a Foundation Models model supports. */
+export type AfmModelMeta = {
+  /** The model's name here. */
+  name: string
+  supports: {
+    input: ReadonlyArray<"text">
+    output: ReadonlyArray<"text">
+    capabilities: ReadonlyArray<"tools" | "structured-output">
+  }
+}
+
+const SYSTEM_MODEL = {
+  name: "system",
+  supports: {
+    input: ["text"],
+    output: ["text"],
+    capabilities: ["tools", "structured-output"]
+  }
+} as const satisfies AfmModelMeta
+
+/** The models the framework runs: Apple's system language model. */
+export const AFM_CHAT_MODELS = [SYSTEM_MODEL.name] as const
+
+/** A model the framework runs. */
+export type AfmChatModel = (typeof AFM_CHAT_MODELS)[number]
+
+/** Each model's metadata, by name. */
+export const AFM_MODEL_META: Readonly<Record<AfmChatModel, AfmModelMeta>> = {
+  [SYSTEM_MODEL.name]: SYSTEM_MODEL
+}

package/src/native-build.test.ts

@@ -0,0 +1,13 @@
+import assert from "node:assert/strict"
+import { readFileSync } from "node:fs"
+import { test } from "node:test"
+
+import { hashFile, sourcesHash } from "../electron/sources-hash.mjs"
+
+test("the committed Electron addon was built from the current Swift and C sources", () => {
+  assert.equal(
+    readFileSync(hashFile, "utf8").trim(),
+    sourcesHash(),
+    "electron/afm.node is stale: rebuild it (node electron/build.mjs, on a Mac with Xcode) and commit it with electron/afm.node.sha256"
+  )
+})

package/src/native.ts

@@ -0,0 +1,119 @@
+/**
+ * The native side of the adapter: Apple's Foundation Models framework, reached
+ * through the Capacitor plugin on iOS (ios/Sources), or the Electron main
+ * process (`serveAfm()`, whose preload puts it at `window.afm`). Any host can
+ * stand in for them by implementing `AfmNative`.
+ */
+import type { PluginListenerHandle } from "@capacitor/core"
+
+/** Why the model can't run: Apple's reasons, or an OS before 26. */
+export type UnavailableReason =
+  | "deviceNotEligible"
+  | "appleIntelligenceNotEnabled"
+  | "modelNotReady"
+  | "unsupportedOS"
+  | "unavailable"
+
+export type NativeAvailability = {
+  available: boolean
+  reason?: UnavailableReason
+  /** Tokens a session holds (iOS/macOS 26.4 and later). */
+  contextSize?: number
+}
+
+/**
+ * A schema as the framework takes it: JSON Schema read into an ordered form
+ * (`schema.ts`), properties in the order given, and named definitions for
+ * `reference`s.
+ */
+export type NativeSchema =
+  | { kind: "object"; name: string; description?: string; properties: Array<NativeProperty> }
+  | { kind: "array"; items: NativeSchema; minimumElements?: number; maximumElements?: number }
+  | { kind: "string"; choices?: Array<string>; constant?: string }
+  | { kind: "integer" | "number"; minimum?: number; maximum?: number }
+  | { kind: "boolean" }
+  | { kind: "null" }
+  | { kind: "anyOf"; name: string; description?: string; choices: Array<NativeSchema> }
+  | { kind: "reference"; name: string }
+
+export type NativeProperty = { name: string; description?: string; optional: boolean; schema: NativeSchema }
+
+/** A root schema and the definitions its references name. */
+export type NativeGenerationSchema = { root: NativeSchema; definitions: Array<NativeSchema> }
+
+export type NativeTool = { name: string; description: string; parameters: NativeGenerationSchema }
+
+/** A message of the conversation a session starts from. */
+export type NativeMessage =
+  | { role: "user"; text: string }
+  | { role: "assistant"; text: string; toolCalls?: Array<{ id: string; name: string; arguments: string }> }
+  | { role: "tool"; toolCallId: string; name: string; text: string }
+
+/**
+ * A response's tokens, over every pass of the model (a tool call's too).
+ * `cachedTokens` is there only where the framework counts it right: for a
+ * session without tools.
+ */
+export type NativeUsage = { inputTokens: number; cachedTokens?: number; outputTokens: number }
+
+/** What a session sends while it responds, for the `turn` of the `respond` call. */
+export type NativeEvent = { sessionId: string; turn: number } & (
+  | { type: "delta"; text: string }
+  | { type: "toolCalls"; calls: Array<{ id: string; name: string; arguments: string }> }
+  | { type: "done"; text: string; usage?: NativeUsage }
+  | { type: "error"; code: string; message: string }
+)
+
+export interface AfmNative {
+  availability(): Promise<NativeAvailability>
+  createSession(options: {
+    sessionId: string
+    instructions?: string
+    tools: Array<NativeTool>
+    messages: Array<NativeMessage>
+  }): Promise<void>
+  prewarm(options: { sessionId: string }): Promise<void>
+  /** Prompts a session; the response comes as events. A `schema` holds the answer to it (JSON in `done`). */
+  respond(options: {
+    sessionId: string
+    turn: number
+    prompt: string
+    options?: {
+      temperature?: number
+      maximumResponseTokens?: number
+      /** Whether the model may call the session's tools, must, or mustn't (iOS/macOS 27). Left out, the framework decides. */
+      toolCallingMode?: "allowed" | "required" | "disallowed"
+    }
+    schema?: NativeGenerationSchema
+  }): Promise<void>
+  /** A tool's result: the session's waiting call returns it, and the response goes on. */
+  submitToolResult(options: { sessionId: string; callId: string; output: string }): Promise<void>
+  /** Stops a session's response. */
+  cancel(options: { sessionId: string }): Promise<void>
+  deleteSession(options: { sessionId: string }): Promise<void>
+  addListener(eventName: "event", listener: (event: NativeEvent) => void): Promise<PluginListenerHandle>
+}
+
+let loading: Promise<{ plugin: AfmNative }> | undefined
+/**
+ * The plugin, loading Capacitor on first use: pages elsewhere never load it.
+ * It's wrapped: a Capacitor plugin answers every property, `then` too, so a
+ * promise of the plugin itself would wait on it forever.
+ */
+const capacitor = () =>
+  (loading ??= import("@capacitor/core").then(({ registerPlugin }) => ({ plugin: registerPlugin<AfmNative>("Afm") })))
+
+/**
+ * The Capacitor plugin (iOS), loading Capacitor on its first call. Elsewhere
+ * its calls reject: Capacitor has no web version of it.
+ */
+export const Afm: AfmNative = {
+  availability: async () => (await capacitor()).plugin.availability(),
+  createSession: async (options) => (await capacitor()).plugin.createSession(options),
+  prewarm: async (options) => (await capacitor()).plugin.prewarm(options),
+  respond: async (options) => (await capacitor()).plugin.respond(options),
+  submitToolResult: async (options) => (await capacitor()).plugin.submitToolResult(options),
+  cancel: async (options) => (await capacitor()).plugin.cancel(options),
+  deleteSession: async (options) => (await capacitor()).plugin.deleteSession(options),
+  addListener: async (name, listener) => (await capacitor()).plugin.addListener(name, listener)
+}

package/src/schema.test.ts

@@ -0,0 +1,146 @@
+import assert from "node:assert/strict"
+import { test } from "node:test"
+
+import { toNativeSchema } from "./schema.js"
+
+test("an object keeps its properties' order; required, nullable and descriptions carry over", () => {
+  const schema = toNativeSchema(
+    {
+      type: "object",
+      description: "Where",
+      properties: {
+        city: { type: "string", description: "The city" },
+        units: { type: "string", enum: ["celsius", "fahrenheit"] },
+        days: { type: ["integer", "null"], minimum: 1, maximum: 7 },
+        hourly: { type: "boolean" }
+      },
+      required: ["city", "days"]
+    },
+    "get_weather"
+  )
+  assert.deepEqual(schema, {
+    root: {
+      kind: "object",
+      name: "get_weather",
+      description: "Where",
+      properties: [
+        { name: "city", description: "The city", optional: false, schema: { kind: "string" } },
+        { name: "units", optional: true, schema: { kind: "string", choices: ["celsius", "fahrenheit"] } },
+        { name: "days", optional: true, schema: { kind: "integer", minimum: 1, maximum: 7 } },
+        { name: "hourly", optional: true, schema: { kind: "boolean" } }
+      ]
+    },
+    definitions: []
+  })
+})
+
+test("nested objects and arrays are named by their path", () => {
+  const { root } = toNativeSchema(
+    {
+      type: "object",
+      properties: {
+        stops: { type: "array", minItems: 1, maxItems: 3, items: { type: "object", properties: { name: { type: "string" } } } }
+      }
+    },
+    "plan-trip"
+  )
+  assert.deepEqual(root, {
+    kind: "object",
+    name: "plan_trip",
+    properties: [
+      {
+        name: "stops",
+        optional: true,
+        schema: {
+          kind: "array",
+          minimumElements: 1,
+          maximumElements: 3,
+          items: { kind: "object", name: "plan_trip_stops_item", properties: [{ name: "name", optional: true, schema: { kind: "string" } }] }
+        }
+      }
+    ]
+  })
+})
+
+test("anyOf becomes a choice (null left out); one non-null choice is the schema itself; const is a constant", () => {
+  const { root } = toNativeSchema(
+    {
+      type: "object",
+      properties: {
+        target: { anyOf: [{ type: "string" }, { type: "number" }] },
+        note: { anyOf: [{ type: "string" }, { type: "null" }] },
+        kind: { const: "search" }
+      }
+    },
+    "tool"
+  )
+  assert.equal(root.kind, "object")
+  const [target, note, kind] = root.kind === "object" ? root.properties : []
+  assert.deepEqual(target.schema, { kind: "anyOf", name: "tool_target", choices: [{ kind: "string" }, { kind: "number" }] })
+  assert.deepEqual(note, { name: "note", optional: true, schema: { kind: "string" } })
+  assert.deepEqual(kind.schema, { kind: "string", constant: "search" })
+})
+
+test("$refs to object definitions are references; to plain values, written in place", () => {
+  const schema = toNativeSchema(
+    {
+      type: "object",
+      properties: { home: { $ref: "#/$defs/Place" }, mood: { $ref: "#/$defs/Mood" } },
+      $defs: {
+        Place: { type: "object", properties: { city: { type: "string" } } },
+        Mood: { type: "string", enum: ["happy", "sad"] }
+      }
+    },
+    "profile"
+  )
+  assert.deepEqual(schema.root.kind === "object" && schema.root.properties.map((p) => p.schema), [
+    { kind: "reference", name: "Place" },
+    { kind: "string", choices: ["happy", "sad"] }
+  ])
+  assert.deepEqual(schema.definitions, [
+    { kind: "object", name: "Place", properties: [{ name: "city", optional: true, schema: { kind: "string" } }] }
+  ])
+})
+
+test("enums hold to their exact members: strings as choices, integers (with strings) as a choice of exact values; what the framework can't hold is refused", () => {
+  const { root } = toNativeSchema(
+    {
+      type: "object",
+      properties: {
+        unit: { type: "string", enum: ["c", "f"] },
+        stars: { type: "integer", enum: [1, 10], description: "A rating." },
+        size: { enum: ["auto", 8] },
+        on: { type: "boolean", enum: [true, false] }
+      },
+      required: ["unit", "stars", "size", "on"]
+    },
+    "t"
+  )
+  assert.equal(root.kind, "object")
+  const property = (name: string) => (root.kind === "object" ? root.properties.find((p) => p.name === name) : undefined)
+  assert.deepEqual(property("unit")?.schema, { kind: "string", choices: ["c", "f"] })
+  assert.deepEqual(property("stars")?.schema, {
+    kind: "anyOf",
+    name: "t_stars",
+    description: "A rating.",
+    choices: [
+      { kind: "integer", minimum: 1, maximum: 1 },
+      { kind: "integer", minimum: 10, maximum: 10 }
+    ]
+  })
+  assert.equal(property("stars")?.description, "A rating.")
+  assert.deepEqual(property("size")?.schema, {
+    kind: "anyOf",
+    name: "t_size",
+    choices: [
+      { kind: "string", constant: "auto" },
+      { kind: "integer", minimum: 8, maximum: 8 }
+    ]
+  })
+  assert.deepEqual(property("on")?.schema, { kind: "boolean" })
+  // One integer is that integer.
+  assert.deepEqual(toNativeSchema({ type: "integer", enum: [3] }, "n").root, { kind: "integer", minimum: 3, maximum: 3 })
+  assert.throws(() => toNativeSchema({ type: "boolean", enum: [true] }, "flag"), /enum of only true isn't supported/)
+  assert.throws(() => toNativeSchema({ type: "number", enum: [0.5, 1.5] }, "ratio"), /decimal numbers isn't supported/)
+  assert.throws(() => toNativeSchema({ enum: ["a", true] }, "mixed"), /mixing string and boolean/)
+})