@seanhogg/builderforce-memory 2026.6.18

package/src/runtime/SSMRuntime.ts

@@ -0,0 +1,309 @@
+/**
+ * SSMRuntime – core runtime class for SSM.js.
+ *
+ * Wraps a MambaSession and adds:
+ *   - Hybrid inference routing (SSM ↔ transformer bridge)
+ *   - Unified generate / stream API
+ *   - Perplexity-aware routing (InferenceRouter auto mode)
+ *   - Thin save/load pass-throughs for MemoryStore integration
+ *
+ * GPU lifecycle is owned entirely by this class via the underlying session.
+ */
+
+import {
+    MambaSession,
+    type MambaSessionOptions,
+    type CompleteOptions,
+    type AdaptOptions,
+    type AdaptResult,
+    type SaveOptions,
+    type LoadOptions,
+    type CreateCallbacks,
+} from '../session/index.js';
+
+import { InferenceRouter, type RoutingStrategy, type RoutingAuditEntry } from '../router/InferenceRouter.js';
+import type { DistillationLog } from '../distillation/DistillationEngine.js';
+import type { TransformerBridge, BridgeGenerateOptions } from '../bridges/TransformerBridge.js';
+import { SSMError } from '../errors/SSMError.js';
+
+// ── Public types ──────────────────────────────────────────────────────────────
+
+export type GenerateOptions = CompleteOptions & {
+    /**
+     * Options forwarded to the transformer bridge when the router selects it.
+     * Ignored when SSM is selected.
+     */
+    bridgeOpts?: BridgeGenerateOptions;
+
+    /**
+     * Stable system / context prefix, kept separate from the volatile `input`.
+     *
+     * Splitting it out is what makes transformer-side prompt caching possible:
+     * on the SSM path it is prepended to `input` (reproducing the single-string
+     * trained prompt format, so SSM behaviour is unchanged), while on the
+     * transformer path it is sent as the cacheable `system` field rather than
+     * being concatenated into the user message. A bridge can then cache it
+     * across turns; concatenating it into the per-turn input could never cache.
+     */
+    system?: string;
+};
+
+export interface SSMRuntimeOptions {
+    /**
+     * Options forwarded verbatim to MambaSession.create().
+     * Controls model size, SSM variant, tokenizer, and checkpoint URL.
+     */
+    session              : MambaSessionOptions;
+
+    /**
+     * Optional transformer bridge.  When absent, all requests go to the SSM
+     * and distillation is unavailable.
+     */
+    bridge?              : TransformerBridge;
+
+    /**
+     * Routing strategy when a bridge is present.
+     * Default: 'auto'
+     */
+    routingStrategy?     : RoutingStrategy;
+
+    /**
+     * Character length above which auto-routing prefers the transformer.
+     * Default: 1200
+     */
+    longInputThreshold?  : number;
+
+    /**
+     * SSM perplexity above which auto-routing prefers the transformer.
+     * Default: 80
+     */
+    perplexityThreshold? : number;
+
+    /**
+     * MambaSession.create() progress callbacks.
+     */
+    callbacks?           : CreateCallbacks;
+}
+
+// ── SSMRuntime ────────────────────────────────────────────────────────────────
+
+export class SSMRuntime {
+    private readonly _session  : MambaSession;
+    private readonly _bridge   : TransformerBridge | undefined;
+    private readonly _router   : InferenceRouter;
+    private _destroyed = false;
+
+    private constructor(
+        session : MambaSession,
+        bridge  : TransformerBridge | undefined,
+        router  : InferenceRouter,
+    ) {
+        this._session = session;
+        this._bridge  = bridge;
+        this._router  = router;
+    }
+
+    // ── Factory ───────────────────────────────────────────────────────────────
+
+    /**
+     * Creates and initialises a new SSMRuntime.
+     *
+     * Delegates to MambaSession.create() — can throw SessionError for
+     * GPU / tokenizer failures.
+     */
+    static async create(opts: SSMRuntimeOptions): Promise<SSMRuntime> {
+        const session = await MambaSession.create(opts.session, opts.callbacks);
+
+        const router = new InferenceRouter({
+            strategy            : opts.routingStrategy,
+            longInputThreshold  : opts.longInputThreshold,
+            perplexityThreshold : opts.perplexityThreshold,
+            hasBridge           : !!opts.bridge,
+            // Pass perplexity probe as callback — avoids circular import
+            perplexityProbe     : opts.bridge
+                ? (text) => session.evaluate(text)
+                : undefined,
+        });
+
+        return new SSMRuntime(session, opts.bridge, router);
+    }
+
+    // ── Inference ─────────────────────────────────────────────────────────────
+
+    /**
+     * Generates a full response for the given input.
+     * Routes to SSM or transformer bridge based on the configured strategy.
+     */
+    async generate(input: string, opts: GenerateOptions = {}): Promise<string> {
+        this._checkAlive();
+
+        const decision = await this._router.route(input);
+
+        if (decision.target === 'transformer' && this._bridge) {
+            return this._bridge.generate(input, this._bridgeOptsWithSystem(opts));
+        }
+
+        // SSM path — extract only CompleteOptions fields
+        const { bridgeOpts: _b, system, ...completeOpts } = opts;
+        return this._session.complete(this._ssmInput(input, system), completeOpts);
+    }
+
+    /**
+     * Streaming token generation.
+     *
+     * Routing note: streaming always uses the SSM path for consistent
+     * latency characteristics.  Use `generate()` if transformer streaming
+     * is needed (bridgeOpts.stream, handled by the bridge's own `stream()` method).
+     */
+    async *stream(input: string, opts: GenerateOptions = {}): AsyncIterable<string> {
+        this._checkAlive();
+        const { bridgeOpts: _b, system, ...completeOpts } = opts;
+        yield* this._session.completeStream(this._ssmInput(input, system), completeOpts);
+    }
+
+    /**
+     * Hybrid streaming: routes to SSM or transformer bridge stream.
+     * Falls back to `generate()` for bridges that don't support streaming.
+     */
+    async *streamHybrid(input: string, opts: GenerateOptions = {}): AsyncIterable<string> {
+        this._checkAlive();
+
+        const decision = await this._router.route(input);
+
+        if (decision.target === 'transformer' && this._bridge) {
+            const bridgeOpts = this._bridgeOptsWithSystem(opts);
+            if (this._bridge.supportsStreaming && this._bridge.stream) {
+                yield* this._bridge.stream(input, bridgeOpts);
+            } else {
+                yield await this._bridge.generate(input, bridgeOpts);
+            }
+            return;
+        }
+
+        const { bridgeOpts: _b, system, ...completeOpts } = opts;
+        yield* this._session.completeStream(this._ssmInput(input, system), completeOpts);
+    }
+
+    // ── Prompt-prefix helpers ─────────────────────────────────────────────────
+
+    /**
+     * Combines the stable `system` prefix with the volatile `input` for the SSM
+     * path, reproducing the single-string format the model was trained on.
+     */
+    private _ssmInput(input: string, system: string | undefined): string {
+        return system ? `${system}\n${input}` : input;
+    }
+
+    /**
+     * Maps the top-level `system` prefix onto the bridge's `systemPrompt` so the
+     * transformer receives it as a cacheable `system` field. An explicit
+     * `bridgeOpts.systemPrompt` wins if the caller set one.
+     */
+    private _bridgeOptsWithSystem(opts: GenerateOptions): BridgeGenerateOptions {
+        return {
+            ...opts.bridgeOpts,
+            systemPrompt: opts.bridgeOpts?.systemPrompt ?? opts.system,
+        };
+    }
+
+    // ── Adaptation ────────────────────────────────────────────────────────────
+
+    /**
+     * Fine-tunes the SSM on the provided text.
+     * Pass-through to MambaSession.adapt().
+     */
+    async adapt(data: string, opts?: AdaptOptions): Promise<AdaptResult> {
+        this._checkAlive();
+        return this._session.adapt(data, opts);
+    }
+
+    /**
+     * Evaluates SSM perplexity on the provided text.
+     * Exposed for power users and InferenceRouter's perplexity probe.
+     */
+    async evaluate(text: string): Promise<number> {
+        this._checkAlive();
+        return this._session.evaluate(text);
+    }
+
+    /**
+     * Returns an L2-normalised SSM embedding for `text`.
+     * Pass-through to MambaSession.embed(). Consumed by MemoryStore.recallSimilar
+     * for SSM-based semantic search.
+     */
+    async embed(text: string): Promise<Float32Array> {
+        this._checkAlive();
+        return this._session.embed(text);
+    }
+
+    // ── Persistence ───────────────────────────────────────────────────────────
+
+    /**
+     * Saves SSM weights. Pass-through to MambaSession.save().
+     * Accepts a `key` option used by MemoryStore to save under a custom key.
+     */
+    async save(opts?: SaveOptions & { key?: string }): Promise<void> {
+        this._checkAlive();
+        await this._session.save(opts);
+    }
+
+    /**
+     * Loads SSM weights. Pass-through to MambaSession.load().
+     * Returns false when no checkpoint exists.
+     */
+    async load(opts?: LoadOptions & { key?: string }): Promise<boolean> {
+        this._checkAlive();
+        return this._session.load(opts) as Promise<boolean>;
+    }
+
+    // ── Accessors ─────────────────────────────────────────────────────────────
+
+    /**
+     * Returns a copy of the in-memory routing audit log.
+     * Delegates to the internal InferenceRouter.
+     */
+    getRoutingAuditLog(): RoutingAuditEntry[] {
+        return this._router.getAuditLog();
+    }
+
+    /**
+     * Returns the distillation log if a DistillationEngine has been attached.
+     * Returns an empty array when no distillation engine is available.
+     *
+     * Note: DistillationEngine is a separate object — attach it by calling
+     *   `new DistillationEngine(runtime, bridge)` and keeping a reference.
+     *   This method is a stub for future inline engine support.
+     */
+    getDistillationLog(): DistillationLog[] {
+        return [];
+    }
+
+    /** The transformer bridge attached to this runtime, if any. */
+    get bridge(): TransformerBridge | undefined { return this._bridge; }
+
+    /** Whether this runtime has been destroyed. */
+    get destroyed(): boolean { return this._destroyed; }
+
+    /**
+     * Escape hatch to the underlying MambaSession internals.
+     * Use sparingly — the SSMRuntime API covers the common cases.
+     */
+    get internals() { return this._session.internals; }
+
+    // ── Lifecycle ─────────────────────────────────────────────────────────────
+
+    /**
+     * Destroys the underlying MambaSession and releases all GPU resources.
+     * After calling this, all methods throw SSMError('RUNTIME_DESTROYED').
+     */
+    destroy(): void {
+        this._session.destroy();
+        this._destroyed = true;
+    }
+
+    private _checkAlive(): void {
+        if (this._destroyed) {
+            throw new SSMError('RUNTIME_DESTROYED', 'SSMRuntime has been destroyed.');
+        }
+    }
+}

package/src/runtime/index.ts

@@ -0,0 +1,2 @@
+export { SSMRuntime } from './SSMRuntime.js';
+export type { SSMRuntimeOptions, GenerateOptions } from './SSMRuntime.js';

package/src/session/errors.ts

@@ -0,0 +1,24 @@
+/**
+ * errors.ts – Typed error class for the MambaSession layer.
+ */
+
+export type SessionErrorCode =
+    | 'GPU_UNAVAILABLE'          // navigator.gpu not present or adapter request failed
+    | 'TOKENIZER_LOAD_FAILED'    // vocab.json or merges.txt could not be fetched/parsed
+    | 'CHECKPOINT_FETCH_FAILED'  // checkpoint URL returned non-OK response after retries
+    | 'CHECKPOINT_INVALID'       // loadWeights threw (bad magic, version, or size mismatch)
+    | 'INPUT_TOO_SHORT'          // adapt() input encodes to fewer than 2 tokens
+    | 'STORAGE_UNAVAILABLE'      // IndexedDB or File System Access API not available
+    | 'SESSION_DESTROYED'        // method called after destroy()
+    | 'UNKNOWN';                 // unexpected error (original in .cause)
+
+export class SessionError extends Error {
+    constructor(
+        public readonly code: SessionErrorCode,
+        message: string,
+        public readonly cause?: unknown,
+    ) {
+        super(message);
+        this.name = 'SessionError';
+    }
+}

package/src/session/index.ts

@@ -0,0 +1,25 @@
+/**
+ * session/index.ts – barrel export for the MambaSession layer.
+ */
+
+export { MambaSession }              from './session.js';
+export { SessionError }              from './errors.js';
+export { MODEL_PRESETS, resolveLayerSchedule, resolveModelConfig } from './presets.js';
+
+export type { SessionErrorCode }     from './errors.js';
+export type { LayerSchedulePreset }  from './presets.js';
+export type {
+    MambaSessionOptions,
+    CompleteOptions,
+    AdaptOptions,
+    AdaptResult,
+    SaveOptions,
+    LoadOptions,
+    StorageTarget,
+    CreateProgressEvent,
+    CreateStage,
+    CreateCallbacks,
+    SessionInternals,
+    GpuMode,
+} from './session.js';
+export type { Tokenizer } from './tokenizer.js';

package/src/session/persistence.ts

@@ -0,0 +1,142 @@
+/**
+ * persistence.ts – Storage helpers for the MambaSession layer.
+ *
+ * Supports three storage targets:
+ *  - IndexedDB  (browser default)
+ *  - Download   (Blob URL trigger)
+ *  - File System Access API
+ */
+
+import { SessionError } from './errors.js';
+
+const DB_NAME    = 'ssmjs-session';
+const DB_VERSION = 1;
+const STORE_NAME = 'checkpoints';
+
+// ── IndexedDB helpers ─────────────────────────────────────────────────────────
+
+function openDB(idb?: IDBFactory): Promise<IDBDatabase> {
+    return new Promise((resolve, reject) => {
+        const factory = idb ?? (typeof indexedDB !== 'undefined' ? indexedDB : undefined);
+        if (!factory) {
+            reject(new SessionError(
+                'STORAGE_UNAVAILABLE',
+                'IndexedDB is not available in this environment. Pass an idbFactory option (e.g. from fake-indexeddb) for Node.js support.',
+            ));
+            return;
+        }
+
+        const req = factory.open(DB_NAME, DB_VERSION);
+
+        req.onupgradeneeded = (e) => {
+            const db = (e.target as IDBOpenDBRequest).result;
+            if (!db.objectStoreNames.contains(STORE_NAME)) {
+                db.createObjectStore(STORE_NAME);
+            }
+        };
+
+        req.onsuccess = () => resolve(req.result);
+        req.onerror   = ()  => reject(new SessionError(
+            'STORAGE_UNAVAILABLE',
+            `Failed to open IndexedDB database "${DB_NAME}": ${req.error?.message ?? 'unknown error'}`,
+            req.error,
+        ));
+    });
+}
+
+export async function saveToIndexedDB(key: string, buffer: ArrayBuffer, idb?: IDBFactory): Promise<void> {
+    const db = await openDB(idb);
+    return new Promise((resolve, reject) => {
+        const tx    = db.transaction(STORE_NAME, 'readwrite');
+        const store = tx.objectStore(STORE_NAME);
+        const req   = store.put(buffer, key);
+
+        req.onsuccess = () => resolve();
+        req.onerror   = () => reject(new SessionError(
+            'STORAGE_UNAVAILABLE',
+            `Failed to write checkpoint to IndexedDB (key="${key}"): ${req.error?.message ?? 'unknown error'}`,
+            req.error,
+        ));
+        tx.oncomplete = () => db.close();
+    });
+}
+
+export async function loadFromIndexedDB(key: string, idb?: IDBFactory): Promise<ArrayBuffer | undefined> {
+    const db = await openDB(idb);
+    return new Promise((resolve, reject) => {
+        const tx    = db.transaction(STORE_NAME, 'readonly');
+        const store = tx.objectStore(STORE_NAME);
+        const req   = store.get(key);
+
+        req.onsuccess = () => {
+            db.close();
+            resolve(req.result as ArrayBuffer | undefined);
+        };
+        req.onerror = () => reject(new SessionError(
+            'STORAGE_UNAVAILABLE',
+            `Failed to read checkpoint from IndexedDB (key="${key}"): ${req.error?.message ?? 'unknown error'}`,
+            req.error,
+        ));
+    });
+}
+
+/** Milliseconds a download Blob URL is kept alive before being revoked. */
+const DOWNLOAD_URL_TTL_MS = 10_000;
+
+// ── Download helper ───────────────────────────────────────────────────────────
+
+export async function triggerDownload(filename: string, buffer: ArrayBuffer): Promise<void> {
+    const blob = new Blob([buffer], { type: 'application/octet-stream' });
+    const url  = URL.createObjectURL(blob);
+
+    const anchor       = document.createElement('a');
+    anchor.href        = url;
+    anchor.download    = filename;
+    anchor.style.display = 'none';
+
+    document.body.appendChild(anchor);
+    anchor.click();
+    document.body.removeChild(anchor);
+
+    // Release the object URL after the TTL to allow the download to start
+    setTimeout(() => URL.revokeObjectURL(url), DOWNLOAD_URL_TTL_MS);
+}
+
+// ── File System Access API helpers ────────────────────────────────────────────
+
+export async function saveViaFileSystemAPI(filename: string, buffer: ArrayBuffer): Promise<void> {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const win = window as any;
+    if (typeof win.showSaveFilePicker !== 'function') {
+        throw new SessionError(
+            'STORAGE_UNAVAILABLE',
+            'File System Access API (showSaveFilePicker) is not available in this browser.',
+        );
+    }
+
+    const handle = await win.showSaveFilePicker({
+        suggestedName: filename,
+        types: [{ description: 'MambaSession Checkpoint', accept: { 'application/octet-stream': ['.bin'] } }],
+    });
+    const writable = await handle.createWritable();
+    await writable.write(buffer);
+    await writable.close();
+}
+
+export async function loadViaFileSystemAPI(): Promise<ArrayBuffer> {
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const win = window as any;
+    if (typeof win.showOpenFilePicker !== 'function') {
+        throw new SessionError(
+            'STORAGE_UNAVAILABLE',
+            'File System Access API (showOpenFilePicker) is not available in this browser.',
+        );
+    }
+
+    const [handle] = await win.showOpenFilePicker({
+        types: [{ description: 'MambaSession Checkpoint', accept: { 'application/octet-stream': ['.bin'] } }],
+        multiple: false,
+    });
+    const file = await handle.getFile();
+    return file.arrayBuffer();
+}

package/src/session/presets.ts

@@ -0,0 +1,122 @@
+/**
+ * presets.ts – Model size presets and config resolver for the MambaSession layer.
+ */
+
+import type { HybridMambaModelConfig, LayerSpec, LayerType } from '@seanhogg/builderforce-memory-engine';
+import type { MambaSessionOptions } from './session.js';
+
+// ── Model size presets ────────────────────────────────────────────────────────
+
+/**
+ * Pre-defined model size presets.
+ * nHeads is used by Mamba-2/3 and Attention layers; ignored for Mamba-1.
+ */
+export const MODEL_PRESETS: Record<string, Partial<HybridMambaModelConfig>> = {
+    nano  : { dModel: 128, numLayers:  4, dState: 16, dConv: 4, expand: 2, nHeads:  4 },
+    small : { dModel: 256, numLayers:  6, dState: 16, dConv: 4, expand: 2, nHeads:  8 },
+    medium: { dModel: 512, numLayers:  8, dState: 16, dConv: 4, expand: 2, nHeads:  8 },
+    large : { dModel: 768, numLayers: 12, dState: 16, dConv: 4, expand: 2, nHeads: 12 },
+};
+
+const DEFAULT_PRESET = 'nano';
+
+// ── Layer schedule presets ────────────────────────────────────────────────────
+
+export type LayerSchedulePreset = 'jamba' | 'zamba';
+
+/**
+ * Resolve a layer schedule from a preset name, explicit array, or undefined.
+ *
+ *   'jamba' — Jamba-style: every 4th layer (index 3, 7, 11…) is attention, rest mamba2
+ *   'zamba' — Zamba-style: every 6th layer (index 5, 11…) is attention, rest mamba3
+ */
+export function resolveLayerSchedule(
+    schedule   : LayerSpec[] | LayerSchedulePreset | undefined,
+    numLayers  : number,
+    defaultType: LayerType,
+): LayerSpec[] {
+    if (!schedule) {
+        return Array.from({ length: numLayers }, () => ({ type: defaultType }));
+    }
+
+    if (schedule === 'jamba') {
+        return Array.from({ length: numLayers }, (_, i) => ({
+            type: (i % 4 === 3 ? 'attention' : 'mamba2') as LayerType,
+        }));
+    }
+
+    if (schedule === 'zamba') {
+        return Array.from({ length: numLayers }, (_, i) => ({
+            type: (i % 6 === 5 ? 'attention' : 'mamba3') as LayerType,
+        }));
+    }
+
+    return schedule;
+}
+
+// ── Config resolution ─────────────────────────────────────────────────────────
+
+/**
+ * Resolves a fully-populated HybridMambaModelConfig from session options and
+ * the actual tokenizer vocab size.
+ *
+ * Resolution order:
+ *  1. Preset fields (default: 'nano')
+ *  2. modelConfig overrides (only applied when modelSize === 'custom')
+ *  3. vocabSize from the tokenizer
+ *  4. mambaVersion → default layer type for schedule resolution
+ *  5. layerSchedule → per-layer type array (preset string or explicit array)
+ */
+export function resolveModelConfig(
+    options  : MambaSessionOptions,
+    vocabSize: number,
+): HybridMambaModelConfig {
+    const presetName = options.modelSize === 'custom' || options.modelSize == null
+        ? DEFAULT_PRESET
+        : options.modelSize;
+
+    const preset = MODEL_PRESETS[presetName] ?? MODEL_PRESETS[DEFAULT_PRESET]!;
+
+    const overrides: Partial<HybridMambaModelConfig> =
+        options.modelSize === 'custom' && options.modelConfig
+            ? options.modelConfig
+            : {};
+
+    const dModel     = overrides.dModel     ?? preset.dModel     ?? 128;
+    const numLayers  = overrides.numLayers  ?? preset.numLayers  ?? 4;
+    const dState     = overrides.dState     ?? preset.dState     ?? 16;
+    const dConv      = overrides.dConv      ?? preset.dConv      ?? 4;
+    const expand     = overrides.expand     ?? preset.expand     ?? 2;
+    const nHeads     = overrides.nHeads     ?? preset.nHeads     ?? 4;
+    const nGroups    = overrides.nGroups    ?? preset.nGroups    ?? 1;
+    const chunkLen   = overrides.chunkLen   ?? preset.chunkLen   ?? 256;
+    const mimoGroup  = overrides.mimoGroup  ?? preset.mimoGroup  ?? 1;
+    const eosId      = overrides.eosId      ?? preset.eosId      ?? -1;
+
+    // Validate nHeads divides dModel for multi-head blocks
+    if (dModel % nHeads !== 0) {
+        throw new Error(
+            `resolveModelConfig: dModel (${dModel}) must be divisible by nHeads (${nHeads}).`
+        );
+    }
+
+    // Layer schedule
+    const defaultType: LayerType = options.mambaVersion ?? 'mamba1';
+    const layers = resolveLayerSchedule(options.layerSchedule, numLayers, defaultType);
+
+    return {
+        vocabSize,
+        dModel,
+        numLayers,
+        dState,
+        dConv,
+        expand,
+        nHeads,
+        nGroups,
+        chunkLen,
+        mimoGroup,
+        eosId,
+        layers,
+        seed: options.seed,
+    };
+}