afm-bridge 0.1.0

package/src/session.ts

@@ -0,0 +1,290 @@
+import { AbortError, AfmError, ClientClosedError, SessionLostError, SessionNotFoundError } from './errors.js'
+import type { RequestOptions } from './rpc.js'
+import { type SchemaInput, type StructuredSchema, toStructuredSchema } from './schema.js'
+import type {
+  GenerateResult,
+  GenerationOptions,
+  ModelId,
+  Prompt,
+  RespondOptions,
+  SessionInfo,
+  StreamChunk,
+  StructuredOptions,
+  StructuredResult,
+  StructuredStreamChunk,
+  Usage,
+} from './types.js'
+
+/** What a session needs from the client. Internal. */
+export interface SessionHost {
+  request<T>(epoch: number, method: string, params: unknown, options?: RequestOptions): Promise<T>
+  listen(epoch: number, requestId: number, listener: (params: DeltaParams) => void): () => void
+  disposed(epoch: number, sessionId: string): void
+}
+
+export interface DeltaParams {
+  requestId: number
+  delta?: string
+  reset?: boolean
+  text?: string
+  partial?: unknown
+}
+
+interface RespondResult {
+  text: string
+  content?: unknown
+  usage?: Usage
+}
+
+type AnyOptions = RespondOptions & { schema?: SchemaInput<unknown> }
+
+/**
+ * A conversation with the model. The sidecar keeps its history, so every call sees
+ * the previous prompts and responses. Only one `respond`/`stream` may run at a time.
+ */
+export class Session {
+  private disposed = false
+
+  /** @internal Use `AppleFoundationModels.createSession()`. */
+  constructor(
+    private readonly host: SessionHost,
+    private readonly epoch: number,
+    readonly id: string,
+    /** Context window in tokens. */
+    readonly contextSize: number,
+    readonly model: ModelId,
+  ) {}
+
+  /**
+   * Generates a complete response. Returns the text, or with `schema` the parsed and
+   * validated value.
+   */
+  respond(prompt: Prompt, options?: RespondOptions & { schema?: undefined }): Promise<string>
+  respond<T>(prompt: Prompt, options: StructuredOptions<T>): Promise<T>
+  async respond(prompt: Prompt, options: AnyOptions = {}): Promise<unknown> {
+    const result = await this.generate(prompt, options as StructuredOptions<unknown>)
+    return options.schema ? (result as StructuredResult<unknown>).content : result.text
+  }
+
+  /** Like `respond`, but also returns token usage (macOS 27+). */
+  generate(prompt: Prompt, options?: RespondOptions & { schema?: undefined }): Promise<GenerateResult>
+  generate<T>(prompt: Prompt, options: StructuredOptions<T>): Promise<StructuredResult<T>>
+  async generate(
+    prompt: Prompt,
+    options: AnyOptions = {},
+  ): Promise<GenerateResult | StructuredResult<unknown>> {
+    const { signal, timeoutMs } = options
+    const schema = options.schema ? toStructuredSchema(options.schema) : undefined
+    const result = await this.call<RespondResult>(
+      'session/respond',
+      this.respondParams(prompt, options, schema),
+      {
+        signal,
+        timeoutMs,
+      },
+    )
+    if (!schema) return withUsage({ text: result.text }, result.usage)
+    return withUsage({ content: parseContent(schema, result), text: result.text }, result.usage)
+  }
+
+  /**
+   * Streams the response. Text: each chunk carries the new `delta` and the full `text`
+   * so far. With `schema`: each chunk carries the `partial` value so far. The last chunk
+   * has `done: true` with the complete result. Breaking out of the loop cancels the generation.
+   */
+  stream(
+    prompt: Prompt,
+    options?: RespondOptions & { schema?: undefined },
+  ): AsyncGenerator<StreamChunk, void, undefined>
+  stream<T>(
+    prompt: Prompt,
+    options: StructuredOptions<T>,
+  ): AsyncGenerator<StructuredStreamChunk<T>, void, undefined>
+  async *stream(
+    prompt: Prompt,
+    options: AnyOptions = {},
+  ): AsyncGenerator<StreamChunk | StructuredStreamChunk<unknown>, void, undefined> {
+    const { signal, timeoutMs } = options
+    if (signal?.aborted) throw new AbortError()
+    this.assertUsable()
+    const schema = options.schema ? toStructuredSchema(options.schema) : undefined
+
+    const controller = new AbortController()
+    const onAbort = () => controller.abort()
+    signal?.addEventListener('abort', onAbort, { once: true })
+
+    const queue = new ChunkQueue()
+    let text = ''
+    let unsubscribe = () => {}
+    const request = this.host.request<RespondResult>(
+      this.epoch,
+      'session/respond',
+      { ...this.respondParams(prompt, options, schema), stream: true },
+      {
+        signal: controller.signal,
+        timeoutMs,
+        onStart: (id) => {
+          unsubscribe = this.host.listen(this.epoch, id, (p) => {
+            if (schema) {
+              if (p.partial !== undefined) queue.push({ done: false, partial: p.partial })
+            } else if (p.reset) {
+              text = p.text ?? ''
+              queue.push({ done: false, delta: '', text, reset: true })
+            } else if (p.delta) {
+              text += p.delta
+              queue.push({ done: false, delta: p.delta, text })
+            }
+          })
+        },
+      },
+    )
+    request.then(
+      (result) => queue.end(result),
+      (error: unknown) => queue.fail(error),
+    )
+
+    let finished = false
+    try {
+      while (true) {
+        const item = await queue.next()
+        if (item.kind === 'chunk') {
+          yield item.chunk
+          continue
+        }
+        finished = true
+        const { result } = item
+        yield schema
+          ? withUsage(
+              { done: true as const, content: parseContent(schema, result), text: result.text },
+              result.usage,
+            )
+          : withUsage({ done: true as const, delta: '' as const, text: result.text }, result.usage)
+        return
+      }
+    } finally {
+      unsubscribe()
+      signal?.removeEventListener('abort', onAbort)
+      if (!finished) {
+        // The consumer stopped early or an error was thrown: make sure the
+        // generation is cancelled and the session is free before we return.
+        controller.abort()
+        await request.catch(() => {})
+      }
+    }
+  }
+
+  /** Loads model resources ahead of time to reduce the latency of the first response. */
+  async prewarm(promptPrefix?: string): Promise<void> {
+    await this.call('session/prewarm', { sessionId: this.id, ...(promptPrefix ? { promptPrefix } : {}) })
+  }
+
+  /** Context usage of the session. */
+  async info(): Promise<SessionInfo> {
+    return this.call<SessionInfo>('session/info', { sessionId: this.id })
+  }
+
+  /** Frees the session in the sidecar and cancels a running generation. Safe to call twice. */
+  async dispose(): Promise<void> {
+    if (this.disposed) return
+    this.disposed = true
+    this.host.disposed(this.epoch, this.id)
+    try {
+      await this.host.request(this.epoch, 'session/dispose', { sessionId: this.id })
+    } catch (error) {
+      // Nothing to free when the process is already gone.
+      if (error instanceof SessionLostError || error instanceof ClientClosedError) return
+      throw error
+    }
+  }
+
+  private call<T>(method: string, params: unknown, options?: RequestOptions): Promise<T> {
+    try {
+      this.assertUsable()
+    } catch (error) {
+      return Promise.reject(error)
+    }
+    return this.host.request<T>(this.epoch, method, params, options)
+  }
+
+  private assertUsable() {
+    if (this.disposed) {
+      throw new SessionNotFoundError(`Session ${this.id} was disposed`, {
+        type: 'sessionNotFound',
+        data: { sessionId: this.id },
+      })
+    }
+  }
+
+  private respondParams(prompt: Prompt, options: AnyOptions, schema: StructuredSchema<unknown> | undefined) {
+    const { temperature, maximumResponseTokens, sampling, toolCalling, reasoningLevel } = options
+    const generation: GenerationOptions = {
+      temperature,
+      maximumResponseTokens,
+      sampling,
+      toolCalling,
+      reasoningLevel,
+    }
+    const defined = Object.fromEntries(Object.entries(generation).filter(([, v]) => v !== undefined))
+    return {
+      sessionId: this.id,
+      prompt,
+      ...(Object.keys(defined).length > 0 ? { options: defined } : {}),
+      ...(schema ? { schema: schema.jsonSchema } : {}),
+      ...(schema && options.includeSchemaInPrompt !== undefined
+        ? { includeSchemaInPrompt: options.includeSchemaInPrompt }
+        : {}),
+    }
+  }
+}
+
+function parseContent(schema: StructuredSchema<unknown>, result: RespondResult): unknown {
+  return schema.parse(result.content !== undefined ? result.content : JSON.parse(result.text))
+}
+
+function withUsage<T extends object>(value: T, usage: Usage | undefined): T & { usage?: Usage } {
+  return usage ? { ...value, usage } : value
+}
+
+type QueueItem =
+  | { kind: 'chunk'; chunk: StreamChunk | StructuredStreamChunk<unknown> }
+  | { kind: 'end'; result: RespondResult }
+
+/** Single-consumer queue that delivers buffered chunks before an error. */
+class ChunkQueue {
+  private readonly items: QueueItem[] = []
+  private error: { value: unknown } | undefined
+  private waiter: (() => void) | undefined
+
+  push(chunk: StreamChunk | StructuredStreamChunk<unknown>) {
+    this.items.push({ kind: 'chunk', chunk })
+    this.wake()
+  }
+
+  end(result: RespondResult) {
+    this.items.push({ kind: 'end', result })
+    this.wake()
+  }
+
+  fail(error: unknown) {
+    this.error = { value: error }
+    this.wake()
+  }
+
+  async next(): Promise<QueueItem> {
+    while (true) {
+      const item = this.items.shift()
+      if (item) return item
+      if (this.error)
+        throw this.error.value instanceof Error ? this.error.value : new AfmError(String(this.error.value))
+      await new Promise<void>((resolve) => {
+        this.waiter = resolve
+      })
+    }
+  }
+
+  private wake() {
+    const waiter = this.waiter
+    this.waiter = undefined
+    waiter?.()
+  }
+}

package/src/tauri/index.ts

@@ -0,0 +1,100 @@
+import { type Child, Command } from '@tauri-apps/plugin-shell'
+import { TransportError } from '../errors.js'
+import type { Connection, Transport, TransportHandlers } from '../transport.js'
+
+export interface TauriSidecarTransportOptions {
+  /**
+   * Sidecar name as listed in `tauri.conf.json > bundle > externalBin` (without the
+   * target-triple suffix). Default `binaries/afm-bridge-server`.
+   */
+  program?: string
+  args?: string[]
+  /** Extra environment variables for the sidecar. */
+  env?: Record<string, string>
+  /** Sidecar log level (stderr). Default `warn`. */
+  logLevel?: 'error' | 'warn' | 'info' | 'debug'
+  /** How long `close()` waits for the sidecar to exit after `shutdown` before killing it. Default 2000 ms. */
+  shutdownTimeoutMs?: number
+  /**
+   * Kill the sidecar when the page goes away (reload, dev-server HMR, navigation).
+   * The Rust side keeps the process alive across reloads otherwise. Default true.
+   */
+  killOnPageHide?: boolean
+}
+
+/**
+ * Runs the sidecar through Tauri's shell plugin, straight from the webview.
+ * Requires `@tauri-apps/plugin-shell` and the capabilities described in the README.
+ */
+export class TauriSidecarTransport implements Transport {
+  constructor(private readonly options: TauriSidecarTransportOptions = {}) {}
+
+  async connect(handlers: TransportHandlers): Promise<Connection> {
+    const program = this.options.program ?? 'binaries/afm-bridge-server'
+    const command = Command.sidecar(program, this.options.args ?? [], {
+      env: { AFM_BRIDGE_LOG: this.options.logLevel ?? 'warn', ...this.options.env },
+    })
+
+    let exited = false
+    let resolveExit = () => {}
+    const exit = new Promise<void>((resolve) => {
+      resolveExit = resolve
+    })
+    command.on('close', ({ code, signal }) => {
+      exited = true
+      removePageHide()
+      handlers.onClose({ code, signal: signal === null ? null : String(signal) })
+      resolveExit()
+    })
+    command.on('error', (error) => handlers.onLog?.(`[tauri shell] ${error}`))
+    // The shell plugin emits one event per line; tolerate a trailing newline either way.
+    command.stdout.on('data', (data) => forEachLine(data, handlers.onMessage))
+    command.stderr.on('data', (data) => forEachLine(data, (line) => handlers.onLog?.(line)))
+
+    let child: Child
+    try {
+      child = await command.spawn()
+    } catch (error) {
+      throw new TransportError(
+        `Could not start sidecar '${program}': ${error}. Check bundle.externalBin in tauri.conf.json, the binary name suffix (-aarch64-apple-darwin) and the shell:allow-spawn capability.`,
+        { type: 'transportError', cause: error },
+      )
+    }
+
+    const onPageHide = () => {
+      if (!exited) void child.kill()
+    }
+    const removePageHide = () => {
+      if (typeof window !== 'undefined') window.removeEventListener('pagehide', onPageHide)
+    }
+    if ((this.options.killOnPageHide ?? true) && typeof window !== 'undefined') {
+      window.addEventListener('pagehide', onPageHide)
+    }
+
+    const shutdownTimeoutMs = this.options.shutdownTimeoutMs ?? 2000
+    return {
+      send(line) {
+        if (exited) return
+        child
+          .write(`${line}\n`)
+          .catch((error: unknown) => handlers.onLog?.(`[tauri shell] write failed: ${error}`))
+      },
+      async close() {
+        removePageHide()
+        if (exited) return
+        // The client sends `shutdown` first, so the sidecar is normally already exiting.
+        const timedOut = await Promise.race([
+          exit.then(() => false),
+          new Promise<boolean>((resolve) => setTimeout(() => resolve(true), shutdownTimeoutMs)),
+        ])
+        if (timedOut && !exited) await child.kill()
+      },
+    }
+  }
+}
+
+function forEachLine(data: string, onLine: (line: string) => void) {
+  for (const line of data.split(/\r?\n/)) {
+    if (line.length > 0) onLine(line)
+  }
+}

package/src/tool.ts

@@ -0,0 +1,80 @@
+import { AfmError } from './errors.js'
+import { type SchemaInput, toStructuredSchema } from './schema.js'
+
+export interface ToolContext {
+  /** Aborted when the generation is cancelled or the tool times out. */
+  signal: AbortSignal
+  sessionId: string
+}
+
+/**
+ * A function the model may call. Arguments are generated by the model to match
+ * `parameters`, validated, and passed to `execute`. The return value goes back to the
+ * model: strings as they are, anything else as JSON.
+ */
+export interface Tool<Args = unknown> {
+  /** `^[a-zA-Z_][a-zA-Z0-9_]{0,63}$` */
+  name: string
+  /** Tells the model when to use the tool. */
+  description: string
+  /** Object schema of the arguments (JSON Schema or `zodSchema()`). */
+  parameters: SchemaInput<Args>
+  includeSchemaInInstructions?: boolean
+  /** Default 60 000 ms. */
+  timeoutMs?: number
+  execute(args: Args, context: ToolContext): unknown
+}
+
+/** Identity helper that infers the argument type from `parameters`. */
+export function tool<Args>(definition: Tool<Args>): Tool<Args> {
+  return definition
+}
+
+/** Wire format for `session/create`. */
+export function toolDefinition(t: Tool<unknown>) {
+  return {
+    name: t.name,
+    description: t.description,
+    parameters: toStructuredSchema(t.parameters).jsonSchema,
+    ...(t.includeSchemaInInstructions === undefined
+      ? {}
+      : { includeSchemaInInstructions: t.includeSchemaInInstructions }),
+  }
+}
+
+const DEFAULT_TOOL_TIMEOUT_MS = 60_000
+
+/** Runs a tool for a `tool/call` request. Errors become the error reply to the sidecar. */
+export async function runTool(
+  t: Tool<unknown>,
+  rawArguments: unknown,
+  context: { signal: AbortSignal; sessionId: string },
+): Promise<{ output: string }> {
+  const args = toStructuredSchema(t.parameters).parse(rawArguments)
+  const timeoutMs = t.timeoutMs ?? DEFAULT_TOOL_TIMEOUT_MS
+  const timeout = AbortSignal.timeout(timeoutMs)
+  const signal = AbortSignal.any([context.signal, timeout])
+
+  // Settle as soon as the signal aborts, even if `execute` ignores it.
+  let onAbort = () => {}
+  const aborted = new Promise<never>((_, reject) => {
+    onAbort = () =>
+      reject(
+        timeout.aborted
+          ? new AfmError(`Tool '${t.name}' timed out after ${timeoutMs} ms`, { type: 'toolTimeout' })
+          : new AfmError(`Tool '${t.name}' was cancelled`, { type: 'cancelled' }),
+      )
+    if (signal.aborted) onAbort()
+    else signal.addEventListener('abort', onAbort, { once: true })
+  })
+  try {
+    const result = await Promise.race([
+      Promise.resolve().then(() => t.execute(args, { signal, sessionId: context.sessionId })),
+      aborted,
+    ])
+    return { output: typeof result === 'string' ? result : JSON.stringify(result ?? null) }
+  } finally {
+    signal.removeEventListener('abort', onAbort)
+    aborted.catch(() => {})
+  }
+}

package/src/transport.ts

@@ -0,0 +1,60 @@
+/**
+ * A transport starts the sidecar and moves NDJSON lines in both directions.
+ * Implement it to run the sidecar somewhere else (for example behind Electron IPC).
+ */
+export interface Transport {
+  /** Starts a new sidecar process. Called again after a crash to restart it. */
+  connect(handlers: TransportHandlers): Promise<Connection>
+}
+
+export interface TransportHandlers {
+  /** One complete protocol line from the sidecar's stdout, without the trailing newline. */
+  onMessage(line: string): void
+  /** The connection ended (process exited). Called exactly once. */
+  onClose(info: CloseInfo): void
+  /** One line of the sidecar's stderr log. */
+  onLog?(line: string): void
+}
+
+export interface CloseInfo {
+  code: number | null
+  signal: string | null
+  error?: Error
+}
+
+export interface Connection {
+  /** Sends one protocol line. The transport appends the newline. */
+  send(line: string): void
+  /** Ends the connection gracefully (close stdin, wait, then kill). Resolves when the process is gone. */
+  close(): Promise<void>
+  /**
+   * Hint from the client: `true` while requests are in flight. Node transports use it to
+   * keep the event loop alive only while needed, so an idle sidecar never blocks process exit.
+   */
+  setActive?(active: boolean): void
+}
+
+/** Splits a stream of text chunks into lines. Handles `\n` and `\r\n`. */
+export function createLineSplitter(onLine: (line: string) => void): {
+  push(chunk: string): void
+  flush(): void
+} {
+  let buffer = ''
+  return {
+    push(chunk) {
+      buffer += chunk
+      let index = buffer.indexOf('\n')
+      while (index !== -1) {
+        const line = buffer.slice(0, index).replace(/\r$/, '')
+        buffer = buffer.slice(index + 1)
+        if (line.length > 0) onLine(line)
+        index = buffer.indexOf('\n')
+      }
+    },
+    flush() {
+      const rest = buffer.replace(/\r$/, '')
+      buffer = ''
+      if (rest.length > 0) onLine(rest)
+    },
+  }
+}

package/src/types.ts

@@ -0,0 +1,142 @@
+import type { DeepPartial, SchemaInput } from './schema.js'
+import type { Tool } from './tool.js'
+
+/** Which model to use. `private-cloud` is reserved; see PROTOCOL.md section 10. */
+export type ModelId = 'on-device' | 'private-cloud'
+
+export type UnavailableReason =
+  /** The Mac cannot run Apple Intelligence (Apple silicon required). */
+  | 'deviceNotEligible'
+  /** Apple Intelligence is switched off in System Settings. */
+  | 'appleIntelligenceNotEnabled'
+  /** Model assets are still downloading or being prepared. Try again later. */
+  | 'modelNotReady'
+  /** Private Cloud Compute: the system is not ready yet. */
+  | 'systemNotReady'
+  /** macOS is too old (afm-bridge needs macOS 26 or newer). */
+  | 'unsupportedOS'
+  /** Not macOS on Apple silicon (detected by the client, the sidecar never starts). */
+  | 'unsupportedPlatform'
+  /** The model needs an entitlement afm-bridge does not have. */
+  | 'requiresEntitlement'
+  | 'unknown'
+
+export type Capability = 'vision' | 'guidedGeneration' | 'reasoning' | 'toolCalling'
+
+export type ModelStatus =
+  | {
+      available: true
+      /** Context window in tokens (instructions + history + prompt + response). */
+      contextSize: number
+      /** Model variant name, macOS 27+. */
+      variant?: string
+      /** Empty on macOS 26, which does not report capabilities. */
+      capabilities: Capability[]
+      /** BCP 47 codes, for example `en`, `en-GB`, `de`. */
+      supportedLanguages: string[]
+    }
+  | { available: false; reason: UnavailableReason; message: string }
+
+export type Sampling =
+  | { mode: 'greedy' }
+  | { mode: 'topK'; k: number; seed?: number }
+  | { mode: 'topP'; threshold: number; seed?: number }
+
+export interface GenerationOptions {
+  temperature?: number
+  maximumResponseTokens?: number
+  sampling?: Sampling
+  /** macOS 27+. */
+  toolCalling?: 'allowed' | 'required' | 'disallowed'
+  /** macOS 27+, only for models with the `reasoning` capability. */
+  reasoningLevel?: 'light' | 'moderate' | 'deep'
+}
+
+export type PromptPart =
+  | { type: 'text'; text: string }
+  | { type: 'image'; path: string }
+  | { type: 'image'; data: string; mimeType?: string }
+
+export type Prompt = string | PromptPart[]
+
+export interface Usage {
+  inputTokens: number
+  cachedInputTokens: number
+  outputTokens: number
+  reasoningTokens: number
+}
+
+export interface ServerInfo {
+  protocolVersion: number
+  server: { name: string; version: string }
+  os: { version: string }
+  models: Record<ModelId, ModelStatus>
+}
+
+export interface SessionOptions {
+  model?: ModelId
+  /** System instructions that apply to the whole conversation. */
+  instructions?: string
+  guardrails?: 'default' | 'permissiveContentTransformations'
+  useCase?: 'general' | 'contentTagging'
+  /** Functions the model may call while responding. */
+  // biome-ignore lint/suspicious/noExplicitAny: tools with different argument types
+  tools?: Array<Tool<any>>
+  signal?: AbortSignal
+}
+
+export interface CallOptions {
+  /** Cancels the request. The promise rejects with `AbortError`. */
+  signal?: AbortSignal
+  /** Client-side timeout; rejects with `RequestTimeoutError` and cancels the request. */
+  timeoutMs?: number
+}
+
+export type RespondOptions = GenerationOptions &
+  CallOptions & {
+    /** Whether the schema is also described in the prompt (default true). Only with `schema`. */
+    includeSchemaInPrompt?: boolean
+  }
+
+/** Options that request structured output: the response is parsed and validated as `T`. */
+export type StructuredOptions<T> = RespondOptions & { schema: SchemaInput<T> }
+
+export interface StructuredResult<T> {
+  /** Parsed and validated value. */
+  content: T
+  /** The same value as JSON text. */
+  text: string
+  usage?: Usage
+}
+
+export type StructuredStreamChunk<T> =
+  | { done: false; partial: DeepPartial<T> }
+  | { done: true; content: T; text: string; usage?: Usage }
+
+export interface GenerateResult {
+  text: string
+  /** Tokens used by this request (macOS 27+). */
+  usage?: Usage
+}
+
+export type StreamChunk =
+  | {
+      done: false
+      /** Text appended since the previous chunk. Empty when `reset` is true. */
+      delta: string
+      /** Full text so far. */
+      text: string
+      /** The model rewrote earlier output; replace what you showed with `text`. */
+      reset?: boolean
+    }
+  | { done: true; delta: ''; text: string; usage?: Usage }
+
+export interface SessionInfo {
+  model: ModelId
+  contextSize: number
+  isResponding: boolean
+  /** Tokens currently occupied by the conversation history (macOS 26.4+). */
+  historyTokenCount?: number
+  /** Cumulative usage of the session (macOS 27+). */
+  usage?: Usage
+}