@dbx-tools/appkit-mastra 0.1.5 → 0.1.10

package/src/memory.ts

@@ -5,7 +5,7 @@
  * with two independent knobs:
  *
  * - **Storage** (threads / messages via `PostgresStore`): defaults to
- *   **per-agent** namespacing via `schemaName: "mastra_<agentId>"` so
+ *   **per-agent** namespacing via {@link agentStorageSchemaName} so
  *   conversation history stays isolated between agents in the same
  *   database. `PostgresStore` auto-creates the schema with
  *   `CREATE SCHEMA IF NOT EXISTS` on init.
@@ -14,29 +14,61 @@
  *   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 { lakebase } from "@databricks/appkit";
-import { pluginUtils } from "@dbx-tools/appkit-shared";
+import { getUsernameWithApiLookup } from "@databricks/appkit";
 import { fastembed } from "@mastra/fastembed";
 import { Memory } from "@mastra/memory";
 import { PgVector, PostgresStore } from "@mastra/pg";
 import { randomUUID } from "node:crypto";
-import type { Pool } from "pg";
+import { Pool, type PoolConfig } from "pg";
+
+import type { MastraAgentDefinition, MastraMemoryConfigOverride } from "./agents";
+import type { MastraPluginConfig } from "./config";
+import { agentStorageSchemaName } from "./storage-schema";
+import { summaryModel, TITLE_INSTRUCTIONS } from "./summarize";
+import { log } from "@dbx-tools/shared-core";
 
-import type {
-  MastraAgentDefinition,
-  MastraMemoryConfigOverride,
-} from "./agents.js";
-import type { MastraPluginConfig } from "./config.js";
+const logger = log.logger("mastra/memory");
 
-/** Pool handle returned by the AppKit `lakebase` plugin `exports().pool`. */
-export type LakebasePool = ReturnType<
-  InstanceType<ReturnType<typeof lakebase>["plugin"]>["exports"]
->["pool"];
+/**
+ * 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"];
@@ -44,59 +76,69 @@ type MemorySetting = MastraAgentDefinition["memory"];
 
 /**
  * True when any plugin-level or per-agent setting could need the
- * Lakebase pool. Used by `plugin.ts` to gate pool acquisition; the
- * builder also acquires lazily so missed cases still fail with a
- * clear lakebase-not-registered error.
+ * 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),
-  );
+  return defs.some((d) => settingNeedsSharedPool(d.storage) || settingNeedsSharedPool(d.memory));
 }
 
 /**
- * Look up the `lakebase` plugin and return its managed `pg.Pool`.
- * Throws when the sibling plugin is not registered; enabling
- * `storage` / `memory` without lakebase is a wiring bug, not a runtime
- * condition we can recover from.
- */
-export function resolveLakebasePool(
-  context: pluginUtils.PluginContextLike | undefined,
-  caller: MastraPluginConfig,
-): LakebasePool {
-  return pluginUtils.require(context, lakebase, caller).exports().pool;
-}
-
-/**
- * Construct a per-agent {@link Memory} factory. Caches the shared
- * `PgVector` singleton (built on first need) and the lazily-resolved
- * Lakebase pool so each agent build is O(1) after the first.
+ * 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,
-  context: pluginUtils.PluginContextLike | undefined,
+  servicePrincipalPool: Pool,
 ): MemoryBuilder {
-  return new MemoryBuilder(config, context);
+  return new MemoryBuilder(config, servicePrincipalPool);
 }
 
 /**
- * Builds one `Memory` per agent. Per-instance state keeps the shared
- * `PgVector` and the resolved Lakebase pool alive across calls so
- * registering N agents stays cheap.
+ * 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;
-  private pool: LakebasePool | undefined;
 
   constructor(
     private readonly config: MastraPluginConfig,
-    private readonly context: pluginUtils.PluginContextLike | undefined,
+    private readonly servicePrincipalPool: Pool,
   ) {}
 
+  /**
+   * 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 {@link agentStorageSchemaName} 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,
+    });
+  }
+
   /**
    * Build a `Memory` for `agentId` after the plugin/agent cascade.
    * Returns `undefined` when the agent has neither storage nor a
@@ -109,39 +151,57 @@ export class MemoryBuilder {
 
     const storage = this.buildStorage(agentId, storageSetting);
     const vector = this.buildVector(memorySetting);
-    if (!storage && !vector) return undefined;
+    if (!storage && !vector) {
+      logger.debug("agent:stateless", { agentId });
+      return undefined;
+    }
+
+    logger.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 } }
+        ...(vector ? { semanticRecall: { topK: 3, messageRange: 2 } } : {}),
+        // Auto-name each thread from its opening turn so the
+        // conversation list the UI renders shows meaningful titles
+        // instead of raw ids. Titling runs on the small / fast chat
+        // tier (see `summarize.ts`) rather than the agent's primary
+        // model, so naming a thread never spends the heavyweight model.
+        // Only meaningful when storage is on; harmless otherwise.
+        ...(storage
+          ? {
+              generateTitle: {
+                model: summaryModel(this.config),
+                instructions: TITLE_INSTRUCTIONS,
+              },
+            }
           : {}),
       },
     });
   }
 
-  private buildStorage(
-    agentId: string,
-    setting: StorageSetting,
-  ): PostgresStore | undefined {
+  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.requirePool() as Pool,
+        schemaName: agentStorageSchemaName(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],
+      withId(setting, `mastra-store__${agentId}`) as ConstructorParameters<typeof PostgresStore>[0],
     );
   }
 
@@ -162,17 +222,10 @@ export class MemoryBuilder {
 
   private getSharedVector(): PgVector {
     if (!this.sharedVector) {
-      this.sharedVector = buildSharedPgVector(this.requirePool());
+      this.sharedVector = buildSharedPgVector(this.servicePrincipalPool);
     }
     return this.sharedVector;
   }
-
-  private requirePool(): LakebasePool {
-    if (!this.pool) {
-      this.pool = resolveLakebasePool(this.context, this.config);
-    }
-    return this.pool;
-  }
 }
 
 /**
@@ -195,9 +248,17 @@ export class MemoryBuilder {
  * the lakebase pool from then on. The placeholder pool is `.end()`'d
  * so its socket book-keeping is released.
  */
-function buildSharedPgVector(pool: LakebasePool): PgVector {
+function buildSharedPgVector(pool: Pool): PgVector {
   const vector = new PgVector({
     id: `pg${randomUUID()}`,
+    // Keep the recall index out of `public`: on a Lakebase database the app
+    // service principal has no CREATE on `public` (PG15+ locks it down), so a
+    // default-schema PgVector fails on CREATE INDEX with "permission denied for
+    // schema public". PgVector runs CREATE SCHEMA IF NOT EXISTS for a named
+    // schema, so the SP creates and owns `mastra_instance` (the same
+    // instance-level schema the workflow-snapshot store uses) and the index
+    // lands there. Table names don't collide with the storage tables.
+    schemaName: "mastra_instance",
     host: "-1",
     port: -1,
     database: "_",
@@ -205,7 +266,7 @@ function buildSharedPgVector(pool: LakebasePool): PgVector {
     password: "_",
   });
   const placeholder = vector.pool;
-  vector.pool = pool as Pool;
+  vector.pool = pool;
   void placeholder.end().catch(() => undefined);
   return vector;
 }
@@ -213,23 +274,17 @@ function buildSharedPgVector(pool: LakebasePool): PgVector {
 /** 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],
+    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 {
+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[] {
+function collectAgentDefinitions(config: MastraPluginConfig): MastraAgentDefinition[] {
   const agents = config.agents;
   if (!agents) return [];
   if (Array.isArray(agents)) return agents;

package/src/mlflow.ts

@@ -0,0 +1,149 @@
+/**
+ * MLflow user-feedback logging: detect whether MLflow tracing is wired
+ * for this deployment, and log a thumbs / comment as a trace
+ * *assessment* via the Databricks MLflow REST API.
+ *
+ * Feedback attaches to a trace, and the plugin's spans reach MLflow
+ * through the same OTel pipeline as every other AppKit span (see
+ * `observability.ts`). MLflow derives its trace id from the OpenTelemetry
+ * trace id (`tr-<hex(otelTraceId)>`), so the server stamps the active
+ * trace id on each turn's response and the client sends it back here.
+ *
+ * There is no MLflow JS SDK, so this posts to the assessments REST
+ * endpoint directly using the OBO-scoped workspace client (the feedback
+ * is thus attributed to the signed-in user). Trace export is
+ * asynchronous, so the just-finished trace may not exist in MLflow yet
+ * when the user reacts; the log call retries briefly on "not found"
+ * before giving up softly.
+ */
+
+import { feedback } from "@dbx-tools/shared-mastra";
+import { databricksFetch, readResponseJson, readResponseText } from "./rest";
+import { async, error, log } from "@dbx-tools/shared-core";
+import { appkit } from "@dbx-tools/appkit";
+
+const logger = log.logger("mastra/mlflow");
+
+/** Workspace client carried on an AppKit execution context. */
+type WorkspaceClient = appkit.WorkspaceClientLike;
+
+/** Assessments REST path for a trace. `3.0` is the current MLflow API version. */
+const assessmentsPath = (traceId: string): string =>
+  `/api/3.0/mlflow/traces/${encodeURIComponent(traceId)}/assessments`;
+
+/** Number of times to retry a "trace not found" response before giving up. */
+const NOT_FOUND_RETRIES = 3;
+/** Base backoff between "trace not found" retries, in ms (grows linearly). */
+const NOT_FOUND_BACKOFF_MS = 1200;
+
+/**
+ * Whether MLflow feedback logging is available for this deployment.
+ *
+ * Enabled when an OTLP exporter endpoint is configured (traces are
+ * actually shipped somewhere) AND an MLflow experiment is named - the
+ * two signals that the OTLP backend is MLflow and traces will
+ * materialize there. Both are standard env vars, so no plugin config is
+ * required; a deployment opts in simply by wiring MLflow tracing.
+ */
+export function mlflowEnabled(): boolean {
+  const hasExporter = Boolean(
+    process.env.OTEL_EXPORTER_OTLP_ENDPOINT?.trim() ||
+    process.env.OTEL_EXPORTER_OTLP_TRACES_ENDPOINT?.trim(),
+  );
+  const hasExperiment = Boolean(
+    process.env.MLFLOW_EXPERIMENT_ID?.trim() || process.env.MLFLOW_EXPERIMENT_NAME?.trim(),
+  );
+  return hasExporter && hasExperiment;
+}
+
+/**
+ * Resolve whether user feedback is enabled from the plugin's optional
+ * `config.feedback` override: the explicit boolean wins, otherwise fall
+ * back to auto-detecting MLflow tracing ({@link mlflowEnabled}). Shared
+ * by the plugin's client-config gate and the server's trace-id header so
+ * the two never disagree.
+ */
+export function resolveFeedbackEnabled(explicit: boolean | undefined): boolean {
+  return explicit ?? mlflowEnabled();
+}
+
+/** Parameters for {@link logFeedback}. */
+export interface LogFeedbackParams {
+  /** MLflow trace id the assessment attaches to (`tr-<hex>`). */
+  traceId: string;
+  /** Assessment name; defaults per whether a value or a comment is sent. */
+  name?: string;
+  /** Thumbs / rating / label value. Omit for a comment-only submission. */
+  value?: boolean | number | string;
+  /** Freeform comment: the rationale alongside a value, or the value itself when none. */
+  comment?: string;
+  /** Identity the feedback is attributed to (user email / id). */
+  sourceId?: string;
+}
+
+/**
+ * Log a HUMAN feedback assessment to a trace. Returns the created
+ * assessment id on success, or `undefined` when the trace can't be
+ * found (even after retrying for export lag) or the request otherwise
+ * fails - callers surface that as a soft "not recorded" rather than an
+ * error, keeping the chat usable.
+ */
+export async function logFeedback(
+  client: WorkspaceClient,
+  params: LogFeedbackParams,
+): Promise<string | undefined> {
+  // A comment with no thumbs value is logged as text feedback; a value
+  // (with an optional comment as the rationale) is the thumbs path.
+  const hasValue = params.value !== undefined;
+  const name =
+    params.name?.trim() ||
+    (hasValue ? feedback.DEFAULT_FEEDBACK_NAME : feedback.DEFAULT_COMMENT_NAME);
+  const value = hasValue ? params.value : params.comment;
+  const assessment: Record<string, unknown> = {
+    trace_id: params.traceId,
+    assessment_name: name,
+    source: {
+      source_type: "HUMAN",
+      source_id: params.sourceId?.trim() || "user",
+    },
+    feedback: { value },
+    ...(hasValue && params.comment?.trim() ? { rationale: params.comment } : {}),
+  };
+  const body = { assessment };
+
+  for (let attempt = 0; attempt <= NOT_FOUND_RETRIES; attempt++) {
+    let res: Response;
+    try {
+      res = await databricksFetch(client, assessmentsPath(params.traceId), {
+        method: "POST",
+        body,
+      });
+    } catch (err) {
+      logger.warn("feedback request failed", {
+        traceId: params.traceId,
+        error: error.errorMessage(err),
+      });
+      return undefined;
+    }
+    if (res.ok) {
+      const parsed = await readResponseJson(res);
+      const assessmentId =
+        (parsed as { assessment?: { assessment_id?: unknown } })?.assessment?.assessment_id ??
+        (parsed as { assessment_id?: unknown })?.assessment_id;
+      return typeof assessmentId === "string" ? assessmentId : "";
+    }
+    // Trace export is async; a fresh trace may not exist yet. Retry a
+    // few times with a short backoff before giving up softly.
+    if (res.status === 404 && attempt < NOT_FOUND_RETRIES) {
+      await async.sleep(NOT_FOUND_BACKOFF_MS * (attempt + 1));
+      continue;
+    }
+    logger.warn("feedback not recorded", {
+      traceId: params.traceId,
+      status: res.status,
+      body: await readResponseText(res),
+    });
+    return undefined;
+  }
+  return undefined;
+}