@dbx-tools/appkit-mastra 0.1.58 → 0.1.68

package/src/history.ts

@@ -1,322 +0,0 @@
-/**
- * Thread history loader exposed as a Mastra custom API route.
- *
- * Backed entirely by native Mastra: looks up the active agent by id,
- * asks its `Memory` instance to `recall` a page of `MastraDBMessage`s,
- * and converts the result to AI SDK V5 `UIMessage`s with the official
- * {@link toAISdkV5Messages} helper from `@mastra/ai-sdk/ui`. No direct
- * database reads.
- *
- * The route is registered through {@link historyRoute} as a Mastra
- * `registerApiRoute` so it sits in the same dispatcher pipeline as
- * the standard agent routes. That means the `MastraServer` auth
- * middleware (in `./server.ts`) has already populated `RequestContext` with
- * `MASTRA_THREAD_ID_KEY` and `MASTRA_RESOURCE_ID_KEY` by the time
- * the handler runs - no cookie or user lookups happen here, and the
- * session-cookie logic stays the single source of truth in `server.ts`.
- */
-
-import type {
-  MastraClearHistoryResponse,
-  MastraHistoryResponse,
-  MastraHistoryUIMessage,
-} from "@dbx-tools/appkit-mastra-shared";
-import { commonUtils, logUtils } from "@dbx-tools/shared";
-import { toAISdkV5Messages } from "@mastra/ai-sdk/ui";
-import type { Agent } from "@mastra/core/agent";
-import type { MastraDBMessage } from "@mastra/core/agent/message-list";
-import {
-  MASTRA_RESOURCE_ID_KEY,
-  MASTRA_THREAD_ID_KEY,
-} from "@mastra/core/request-context";
-import type { ContextWithMastra } from "@mastra/core/server";
-import { registerApiRoute } from "@mastra/core/server";
-
-const log = logUtils.logger("mastra/history");
-
-/** Default history page size; matches the Mastra storage default. */
-const DEFAULT_PER_PAGE = 20;
-/** Hard cap so a misbehaving client can't fetch the whole thread at once. */
-const MAX_PER_PAGE = 200;
-
-/** Inputs accepted by {@link loadHistory}. */
-export interface LoadHistoryOptions {
-  agent: Agent;
-  threadId: string;
-  resourceId?: string;
-  page?: number;
-  perPage?: number;
-  /** When true, returns the *oldest* page first (chronological). */
-  ascending?: boolean;
-}
-
-/**
- * Fetch a page of UI-formatted messages for a thread.
- *
- * Uses the agent's resolved `Memory` (`getMemory()`) so per-agent
- * storage namespaces (`mastra_<agentId>` schemas) and any future
- * memory-side filters apply automatically. When the agent has no
- * memory configured the response is a successful empty page so
- * callers don't have to special-case stateless agents.
- *
- * Pagination is descending-by-default: page 0 is the most recent
- * page, page 1 the page before that, etc. The returned `uiMessages`
- * are always re-sorted into chronological order (oldest -> newest)
- * so the client can prepend them above the existing transcript
- * without sorting locally.
- */
-export async function loadHistory(
-  opts: LoadHistoryOptions,
-): Promise<MastraHistoryResponse> {
-  const perPage = clampPerPage(opts.perPage);
-  const page = Math.max(0, Math.trunc(opts.page ?? 0));
-  const memory = await opts.agent.getMemory();
-  if (!memory) {
-    log.debug("recall:no-memory", { agentId: opts.agent.id, threadId: opts.threadId });
-    return { uiMessages: [], page, perPage, total: 0, hasMore: false };
-  }
-  const startedAt = Date.now();
-  const result = await memory.recall({
-    threadId: opts.threadId,
-    ...(opts.resourceId ? { resourceId: opts.resourceId } : {}),
-    page,
-    perPage,
-    orderBy: {
-      field: "createdAt",
-      direction: opts.ascending ? "ASC" : "DESC",
-    },
-  });
-  const chronological = sortChronological(result.messages);
-  const uiMessages = toAISdkV5Messages(
-    chronological,
-  ) as unknown as MastraHistoryUIMessage[];
-  log.debug("recall:done", {
-    agentId: opts.agent.id,
-    threadId: opts.threadId,
-    page,
-    perPage,
-    returned: uiMessages.length,
-    total: result.total,
-    hasMore: result.hasMore,
-    elapsedMs: Date.now() - startedAt,
-  });
-  return {
-    uiMessages,
-    page,
-    perPage,
-    total: result.total,
-    hasMore: result.hasMore,
-  };
-}
-
-/** Inputs accepted by {@link clearHistory}. */
-export interface ClearHistoryOptions {
-  agent: Agent;
-  threadId: string;
-}
-
-/**
- * Wipe every persisted message tied to a thread. Returns the count
- * of messages that were on the thread at delete time so the caller
- * can render a "cleared N messages" affordance without an
- * additional round-trip.
- *
- * Agents without a configured `Memory` resolve to a no-op (count
- * 0), matching {@link loadHistory}'s "stateless agents return an
- * empty page" stance so callers don't have to special-case them.
- * Threads that don't exist yet are also a successful no-op - the
- * operation is idempotent so the UI can fire-and-forget without
- * tracking thread existence.
- */
-export async function clearHistory(
-  opts: ClearHistoryOptions,
-): Promise<{ cleared: number }> {
-  const memory = await opts.agent.getMemory();
-  if (!memory) {
-    log.debug("clear:no-memory", { agentId: opts.agent.id, threadId: opts.threadId });
-    return { cleared: 0 };
-  }
-  // Mastra's `deleteThread` cascades to the message table, so we
-  // can't ask for a count after the fact. Read it pre-delete with a
-  // one-page recall sized to fit common threads in a single round
-  // trip; the value is for telemetry / UI, not correctness.
-  let cleared = 0;
-  try {
-    const probe = await memory.recall({
-      threadId: opts.threadId,
-      page: 0,
-      perPage: 1,
-    });
-    cleared = probe.total;
-  } catch (err) {
-    // A missing-thread error is the happy-path "nothing to count";
-    // every other error is logged but doesn't block the delete.
-    log.debug("clear:probe-failed", {
-      agentId: opts.agent.id,
-      threadId: opts.threadId,
-      error: commonUtils.errorMessage(err),
-    });
-  }
-
-  const startedAt = Date.now();
-  try {
-    await memory.deleteThread(opts.threadId);
-  } catch (err) {
-    // Mastra's `deleteThread` raises when the thread row was never
-    // created (e.g. clearing an empty session). Surface as a soft
-    // warn and treat as success - the user-facing semantic is
-    // "history is now empty" which is already true.
-    log.warn("clear:delete-soft-failed", {
-      agentId: opts.agent.id,
-      threadId: opts.threadId,
-      error: commonUtils.errorMessage(err),
-    });
-  }
-  log.info("clear:done", {
-    agentId: opts.agent.id,
-    threadId: opts.threadId,
-    cleared,
-    elapsedMs: Date.now() - startedAt,
-  });
-  return { cleared };
-}
-
-/** Options accepted by {@link historyRoute}. */
-export type HistoryRouteOptions =
-  | { path: `${string}:agentId${string}`; agent?: never }
-  | { path: string; agent: string };
-
-/**
- * Register the `<path>` Mastra custom API route. Handles two
- * methods on the same mount:
- *
- *   - `GET`: return a page of AI SDK V5 `UIMessage`s for the
- *     caller's current thread ({@link loadHistory}).
- *   - `DELETE`: wipe every persisted message on the caller's
- *     thread ({@link clearHistory}). The session cookie that
- *     anchors the thread id is left alone so the user keeps the
- *     same thread - only the contents go away.
- *
- * Follows the `@mastra/ai-sdk` `chatRoute` convention for agent
- * binding: pass `agent` for a fixed-agent mount, or include
- * `:agentId` in the path for dynamic routing. The plugin registers
- * both `/route/history` (default agent) and `/route/history/:agentId`.
- *
- * The handler reads `threadId` and `resourceId` from `RequestContext`
- * (populated upstream by `MastraServer.registerAuthMiddleware`), so
- * no cookie or user lookups happen here.
- */
-export function historyRoute(options: HistoryRouteOptions) {
-  const { path } = options;
-  const fixedAgent = "agent" in options ? options.agent : undefined;
-  if (!fixedAgent && !path.includes(":agentId")) {
-    throw new Error(
-      "historyRoute path must include `:agentId` or `agent` must be passed explicitly",
-    );
-  }
-  // Tiny resolver shared by GET / DELETE: derive the active agent
-  // and thread id, returning a JSON error response when either is
-  // missing. Keeps both handlers thin and gives them identical
-  // validation behaviour.
-  const resolveContext = (c: ContextWithMastra) => {
-    const mastra = c.get("mastra");
-    const requestContext = c.get("requestContext");
-    const agentId = fixedAgent ?? c.req.param("agentId");
-    if (!agentId) {
-      return { error: c.json({ error: "agentId is required" }, 400) } as const;
-    }
-    const agent = mastra.getAgentById(agentId);
-    if (!agent) {
-      return {
-        error: c.json({ error: `Unknown agent "${agentId}"` }, 404),
-      } as const;
-    }
-    const threadId = requestContext.get(MASTRA_THREAD_ID_KEY) as string | undefined;
-    if (!threadId) {
-      return {
-        error: c.json({ error: "thread id missing from request context" }, 400),
-      } as const;
-    }
-    const resourceId = requestContext.get(MASTRA_RESOURCE_ID_KEY) as string | undefined;
-    return { agentId, agent, threadId, resourceId } as const;
-  };
-
-  return [
-    registerApiRoute(path, {
-      method: "GET",
-      handler: async (c: ContextWithMastra) => {
-        const ctx = resolveContext(c);
-        if ("error" in ctx) return ctx.error;
-        const payload = await loadHistory({
-          agent: ctx.agent,
-          threadId: ctx.threadId,
-          ...(ctx.resourceId ? { resourceId: ctx.resourceId } : {}),
-          page: parseIntParam(c.req.query("page")),
-          perPage: parseIntParam(c.req.query("perPage")),
-        });
-        return c.json(payload);
-      },
-    }),
-    registerApiRoute(path, {
-      method: "DELETE",
-      handler: async (c: ContextWithMastra) => {
-        const ctx = resolveContext(c);
-        if ("error" in ctx) return ctx.error;
-        const { cleared } = await clearHistory({
-          agent: ctx.agent,
-          threadId: ctx.threadId,
-        });
-        const payload: MastraClearHistoryResponse = {
-          ok: true,
-          agentId: ctx.agentId,
-          threadId: ctx.threadId,
-          cleared,
-        };
-        return c.json(payload);
-      },
-    }),
-  ];
-}
-
-/** Coerce / clamp `perPage`; falls back to the page-size default. */
-function clampPerPage(value: number | undefined): number {
-  if (value === undefined || Number.isNaN(value)) return DEFAULT_PER_PAGE;
-  const n = Math.trunc(value);
-  if (n <= 0) return DEFAULT_PER_PAGE;
-  return Math.min(n, MAX_PER_PAGE);
-}
-
-/**
- * Sort messages oldest-first by `createdAt`, falling back to whatever
- * order the storage returned them in. The native `recall` call honors
- * `orderBy` but doesn't guarantee a stable secondary sort, so we
- * normalize here before handing the page to the AI SDK converter.
- */
-function sortChronological(messages: MastraDBMessage[]): MastraDBMessage[] {
-  return [...messages].sort((a, b) => {
-    const ta = toEpoch(a.createdAt);
-    const tb = toEpoch(b.createdAt);
-    return ta - tb;
-  });
-}
-
-function toEpoch(value: unknown): number {
-  if (value instanceof Date) return value.getTime();
-  if (typeof value === "string" || typeof value === "number") {
-    const parsed = new Date(value).getTime();
-    return Number.isNaN(parsed) ? 0 : parsed;
-  }
-  return 0;
-}
-
-/**
- * Coerce a Hono query value into a non-negative integer. Returns
- * `undefined` for empty / non-numeric / negative inputs so
- * {@link loadHistory} can apply its built-in defaults.
- */
-function parseIntParam(value: string | undefined): number | undefined {
-  if (!value) return undefined;
-  const n = Number(value);
-  if (!Number.isFinite(n) || n < 0) return undefined;
-  return Math.trunc(n);
-}

package/src/memory.ts

@@ -1,293 +0,0 @@
-/**
- * Lakebase-backed Mastra memory wiring.
- *
- * Provides a {@link MemoryBuilder} that mints one `Memory` per agent
- * with two independent knobs:
- *
- * - **Storage** (threads / messages via `PostgresStore`): defaults to
- *   **per-agent** namespacing via `schemaName: "mastra_<agentId>"` so
- *   conversation history stays isolated between agents in the same
- *   database. `PostgresStore` auto-creates the schema with
- *   `CREATE SCHEMA IF NOT EXISTS` on init.
- * - **Memory** (semantic recall via `PgVector`): defaults to a single
- *   **shared** instance across every agent. Cross-agent recall on one
- *   index is almost always what users want; opt into per-agent recall
- *   by passing a {@link MastraMemoryConfigOverride} on the agent.
- *
- * Additionally, {@link MemoryBuilder.instanceStorage} returns a
- * **Mastra-instance-level** `PostgresStore` (schema `mastra_instance`)
- * used for workflow snapshots - the persistence layer
- * `agent.resumeStream()` reads from when waking a suspended
- * `requireApproval` tool call. Per-agent stores are not enough for
- * this: workflow runs are scoped to the Mastra instance, not an
- * individual agent's `Memory`.
- *
- * Plugin-level `config.storage` / `config.memory` act as the baseline
- * (auto-defaulted to `true` in `plugin.ts` when the `lakebase` plugin
- * is registered); per-agent settings cascade on top of that.
- */
-
-import { getUsernameWithApiLookup } from "@databricks/appkit";
-import { logUtils } from "@dbx-tools/shared";
-import { fastembed } from "@mastra/fastembed";
-import { Memory } from "@mastra/memory";
-import { PgVector, PostgresStore } from "@mastra/pg";
-import { randomUUID } from "node:crypto";
-import { Pool, type PoolConfig } from "pg";
-
-import type { MastraAgentDefinition, MastraMemoryConfigOverride } from "./agents.js";
-import type { MastraPluginConfig } from "./config.js";
-
-const log = logUtils.logger("mastra/memory");
-
-/**
- * Build a dedicated **service-principal** Lakebase pool for Mastra
- * memory from the lakebase plugin's resolved SP pg config.
- *
- * The plugin's `exports().pool` is a `RoutingPool` that switches to
- * the per-user (OBO) pool whenever a query runs inside an `asUser`
- * scope - exactly the context the mastra plugin establishes around
- * every chat turn. Memory (threads / messages + semantic recall) must
- * instead always act as the app service principal: it owns the
- * auto-created `mastra_*` schemas (a per-user role usually can't
- * `CREATE SCHEMA`) and is shared across users, so it cannot inherit a
- * request's OBO identity.
- *
- * `pgConfig` must be the plugin's `exports().getPgConfig()` evaluated
- * **outside** any `asUser` scope (i.e. during setup), so it carries
- * the SP connection target, OAuth token-refresh `password` callback,
- * and any `lakebase({ pool })` tuning overrides - all of which this
- * pool inherits. See the call site in `plugin.ts`.
- */
-export async function createServicePrincipalPool(pgConfig: PoolConfig): Promise<Pool> {
-  // `getPgConfig()` resolves the SP username synchronously from
-  // `PGUSER` / `DATABRICKS_CLIENT_ID`; fall back to the async API
-  // lookup (e.g. local dev authenticating via PAT) so the pool always
-  // has an identity to connect with.
-  const user = pgConfig.user ?? (await getUsernameWithApiLookup());
-  return new Pool({ ...pgConfig, user });
-}
-
-/** Effective per-knob setting after the plugin/agent cascade. */
-type StorageSetting = MastraAgentDefinition["storage"];
-type MemorySetting = MastraAgentDefinition["memory"];
-
-/**
- * True when any plugin-level or per-agent setting could need the
- * Lakebase pool. Used by `plugin.ts` to gate creation of the
- * service-principal pool and the {@link MemoryBuilder} that consumes
- * it; when false neither is built.
- */
-export function needsLakebase(config: MastraPluginConfig): boolean {
-  if (settingNeedsSharedPool(config.storage)) return true;
-  if (settingNeedsSharedPool(config.memory)) return true;
-  const defs = collectAgentDefinitions(config);
-  return defs.some(
-    (d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory),
-  );
-}
-
-/**
- * Construct a per-agent {@link Memory} factory bound to the supplied
- * service-principal pool (see {@link createServicePrincipalPool}).
- * Caches the shared `PgVector` singleton (built on first need) so each
- * agent build is O(1) after the first.
- */
-export function createMemoryBuilder(
-  config: MastraPluginConfig,
-  servicePrincipalPool: Pool,
-): MemoryBuilder {
-  return new MemoryBuilder(config, servicePrincipalPool);
-}
-
-/**
- * Builds one `Memory` per agent against a shared service-principal
- * Lakebase pool. Per-instance state keeps the shared `PgVector` alive
- * across calls so registering N agents stays cheap.
- */
-export class MemoryBuilder {
-  private sharedVector: PgVector | undefined;
-
-  constructor(
-    private readonly config: MastraPluginConfig,
-    private readonly servicePrincipalPool: Pool,
-  ) {}
-
-  /**
-   * Build a `Memory` for `agentId` after the plugin/agent cascade.
-   * Returns `undefined` when the agent has neither storage nor a
-   * vector store enabled - Mastra accepts a missing `memory` field
-   * and treats the agent as stateless.
-   */
-  /**
-   * Build the Mastra-instance-level storage used for workflow
-   * snapshots. Returns `undefined` when plugin-level `storage` is
-   * disabled, in which case `agent.resumeStream()` (and therefore
-   * the `requireApproval` flow) will not be available.
-   *
-   * The store lives in a dedicated `mastra_instance` schema so it
-   * never collides with per-agent `mastra_<agentId>` namespaces.
-   * Workflow snapshots are not per-agent state; they belong to the
-   * `Mastra` instance that owns the workflow execution.
-   */
-  instanceStorage(): PostgresStore | undefined {
-    const setting = this.config.storage;
-    if (!setting) return undefined;
-    if (typeof setting === "object") {
-      return new PostgresStore(
-        withId(setting, "mastra-store__instance") as ConstructorParameters<
-          typeof PostgresStore
-        >[0],
-      );
-    }
-    return new PostgresStore({
-      id: "mastra-store__instance",
-      schemaName: "mastra_instance",
-      pool: this.servicePrincipalPool,
-    });
-  }
-
-  forAgent(agentId: string, def: MastraAgentDefinition): Memory | undefined {
-    const storageSetting = def.storage ?? this.config.storage;
-    const memorySetting = def.memory ?? this.config.memory;
-
-    const storage = this.buildStorage(agentId, storageSetting);
-    const vector = this.buildVector(memorySetting);
-    if (!storage && !vector) {
-      log.debug("agent:stateless", { agentId });
-      return undefined;
-    }
-
-    log.debug("agent:configured", {
-      agentId,
-      storage: storage !== undefined,
-      vector: vector !== undefined,
-      vectorMode:
-        vector === undefined
-          ? "off"
-          : typeof memorySetting === "object"
-            ? "dedicated"
-            : "shared",
-    });
-
-    return new Memory({
-      ...(storage ? { storage } : {}),
-      ...(vector ? { vector, embedder: fastembed } : {}),
-      options: {
-        lastMessages: 10,
-        ...(vector ? { semanticRecall: { topK: 3, messageRange: 2 } } : {}),
-      },
-    });
-  }
-
-  private buildStorage(
-    agentId: string,
-    setting: StorageSetting,
-  ): PostgresStore | undefined {
-    if (!setting) return undefined;
-    if (typeof setting === "boolean") {
-      return new PostgresStore({
-        id: `mastra-store__${agentId}`,
-        schemaName: `mastra_${agentId}`,
-        pool: this.servicePrincipalPool,
-      });
-    }
-    // Cast: `withId` guarantees `id` is set, but the distributive
-    // Omit + `id?: string` shape doesn't structurally narrow to the
-    // discriminated union members. Runtime shape is identical.
-    return new PostgresStore(
-      withId(setting, `mastra-store__${agentId}`) as ConstructorParameters<
-        typeof PostgresStore
-      >[0],
-    );
-  }
-
-  /**
-   * Resolve the agent's vector store. Cascade:
-   *
-   * - falsy: no vector.
-   * - `boolean` / `undefined-inheriting-true`: return the shared
-   *   singleton (built lazily on first call). All agents that
-   *   default-enable memory write into and recall from one index.
-   * - object: build a dedicated `PgVector` for this agent.
-   */
-  private buildVector(setting: MemorySetting): PgVector | undefined {
-    if (!setting) return undefined;
-    if (typeof setting === "boolean") return this.getSharedVector();
-    return buildPgVector(setting);
-  }
-
-  private getSharedVector(): PgVector {
-    if (!this.sharedVector) {
-      this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
-    }
-    return this.sharedVector;
-  }
-}
-
-/**
- * Build the shared `PgVector` that backs the default
- * `def.memory === true` case across every agent.
- *
- * `PgVector`'s constructor accepts only connection-style configs
- * (`HostConfig` / `ConnectionStringConfig` / `ClientConfig`); there is
- * no `{ pool }` shorthand the way `PostgresStore` has one. Worse, the
- * constructor synchronously kicks off a `cacheWarmupPromise` IIFE that
- * calls `this.pool.connect()` before returning, so we can't cleanly
- * hand it an inert config and patch the pool afterwards.
- *
- * The trick: pass illegal-but-validation-passing placeholders so the
- * warmup's `net.connect()` rejects synchronously with `RangeError`
- * (Node validates `0 <= port < 65536`). The IIFE's `catch {}` swallows
- * it, no DNS lookup or TCP attempt happens, and we then swap
- * `pgVector.pool` to the lakebase pool. Every subsequent `PgVector`
- * method reads `this.pool` at call time, so all real I/O goes through
- * the lakebase pool from then on. The placeholder pool is `.end()`'d
- * so its socket book-keeping is released.
- */
-function buildSharedPgVector(pool: Pool): PgVector {
-  const vector = new PgVector({
-    id: `pg${randomUUID()}`,
-    host: "-1",
-    port: -1,
-    database: "_",
-    user: "_",
-    password: "_",
-  });
-  const placeholder = vector.pool;
-  vector.pool = pool;
-  void placeholder.end().catch(() => undefined);
-  return vector;
-}
-
-/** Per-agent dedicated `PgVector` (rare; opt-in via object override). */
-function buildPgVector(setting: MastraMemoryConfigOverride): PgVector {
-  return new PgVector(
-    withId(setting, `pg-vector__${randomUUID()}`) as ConstructorParameters<
-      typeof PgVector
-    >[0],
-  );
-}
-
-/** True when this setting requires the shared Lakebase pool. */
-function settingNeedsSharedPool(
-  setting: StorageSetting | MemorySetting | undefined,
-): boolean {
-  return setting === true;
-}
-
-/** Walk the three shapes of `config.agents` into a flat list. */
-function collectAgentDefinitions(config: MastraPluginConfig): MastraAgentDefinition[] {
-  const agents = config.agents;
-  if (!agents) return [];
-  if (Array.isArray(agents)) return agents;
-  if (typeof (agents as MastraAgentDefinition).instructions === "string") {
-    return [agents as MastraAgentDefinition];
-  }
-  return Object.values(agents as Record<string, MastraAgentDefinition>);
-}
-
-/** Fill in a default `id` when the caller didn't supply one. */
-function withId<T extends { id?: string }>(value: T, fallback: string): T {
-  return value.id ? value : { ...value, id: fallback };
-}