@dbx-tools/appkit-memory 0.0.0-stage → 0.9.132

package/src/graphiti/runtime.ts

@@ -0,0 +1,261 @@
+/**
+ * Node Graphiti runtime and process supervision.
+ *
+ * @module
+ */
+
+import { delimiter, dirname } from "node:path";
+import { exec } from "@dbx-tools/core";
+import type { ChildProcessResult } from "@dbx-tools/core/exec";
+import {
+  asyncUtils,
+  environmentUtils,
+  functionUtils,
+  options as sharedOptions,
+} from "@dbx-tools/shared-core";
+import { coerce, rcompare, type SemVer } from "semver";
+
+import {
+  graphitiOptionsEnvironment,
+  resolveGraphitiOptions,
+  type GraphitiOptions,
+  type ResolvedGraphitiOptions,
+} from "./options.ts";
+
+/** Context passed to one Graphiti sidecar readiness probe. */
+export interface GraphitiHealthCheckContext {
+  readonly process: ChildProcessResult;
+  readonly signal: AbortSignal;
+  readonly attempt: number;
+}
+
+/** HTTP probe used to decide that the Graphiti sidecar has started listening. */
+export type GraphitiHealthCheck = (
+  context: GraphitiHealthCheckContext,
+) => boolean | PromiseLike<boolean>;
+
+/** Async creation result that preserves the awaitable child handle. */
+export interface GraphitiChildProcess {
+  child: ChildProcessResult;
+}
+
+/** Shared Graphiti options plus the Node-only sidecar readiness probe. */
+export type GraphitiRuntimeOptions = GraphitiOptions & {
+  /** Readiness probe for {@link createGraphitiChildProcess}. Defaults to `/healthcheck`. */
+  healthCheck?: GraphitiHealthCheck;
+  /** Cancels sidecar startup without changing the serialized Graphiti configuration. */
+  signal?: AbortSignal;
+};
+
+const HEALTH_CHECK_INTERVAL_MS = 200;
+const PYTHON_CANDIDATES = ["python", "python3"] as const;
+const GRAPHITI_IMPORT_PROBE = "import dbx_tools.graphiti";
+
+interface PythonCandidate {
+  command: string;
+  version: SemVer;
+}
+
+const graphitiPython = functionUtils.memoize(async (): Promise<string | undefined> => {
+  const available = (
+    await Promise.all(PYTHON_CANDIDATES.map((command) => _probePython(command)))
+  ).filter((candidate): candidate is PythonCandidate => candidate !== undefined);
+  available.sort((left, right) => rcompare(left.version, right.version));
+  return available[0]?.command;
+});
+
+function graphitiOptions(options: GraphitiRuntimeOptions): GraphitiOptions {
+  const { healthCheck: _healthCheck, signal: _signal, ...graphiti } = options;
+  return graphiti;
+}
+
+/** Milliseconds left in a startup budget after `startedAt`. */
+export function remainingTimeoutMs(startedAt: number, timeoutMs: number, now = Date.now()): number {
+  return Math.max(0, timeoutMs - (now - startedAt));
+}
+
+function _runtimeEnvironment(options: ResolvedGraphitiOptions): NodeJS.ProcessEnv {
+  if (options.listen.scheme !== "tcp") throw new Error("Graphiti requires a TCP listener");
+  const nodeBin = process.env.DBX_TOOLS_NODE_BIN?.trim();
+  const path = nodeBin
+    ? [dirname(nodeBin), process.env.PATH].filter(Boolean).join(delimiter)
+    : process.env.PATH;
+  return {
+    ...process.env,
+    ...(path ? { PATH: path } : {}),
+    ...graphitiOptionsEnvironment(options),
+  };
+}
+
+/** Return whether the Graphiti `/healthcheck` endpoint responds successfully. */
+export async function graphitiHealthCheck(
+  options: Pick<ResolvedGraphitiOptions, "listen" | "bearer">,
+  signal: AbortSignal,
+): Promise<boolean> {
+  try {
+    return (
+      await fetch(graphitiHttpUrl(options, "/healthcheck"), {
+        signal,
+        headers: graphitiRequestHeaders(options),
+      })
+    ).ok;
+  } catch (error) {
+    if (signal.aborted) throw error;
+    return false;
+  }
+}
+
+/** Poll `/healthcheck` until it succeeds or `timeoutMs` elapses. */
+export async function waitForGraphitiHealth(
+  options: Pick<ResolvedGraphitiOptions, "listen" | "bearer">,
+  timeoutMs: number,
+  signal?: AbortSignal,
+): Promise<void> {
+  const budget = Math.max(1, timeoutMs);
+  for await (const healthy of asyncUtils.poll(
+    ({ signal: pollSignal }) => graphitiHealthCheck(options, pollSignal),
+    {
+      intervalMs: HEALTH_CHECK_INTERVAL_MS,
+      timeoutMs: budget,
+      signal,
+      predicate: (ready) => !ready,
+    },
+  )) {
+    if (healthy) return;
+  }
+  throw new Error("Graphiti sidecar healthcheck did not succeed");
+}
+
+/** Build authorization headers for one optionally secured Graphiti runtime. */
+export function graphitiRequestHeaders(
+  options: Pick<ResolvedGraphitiOptions, "bearer">,
+): Readonly<Record<string, string>> {
+  return options.bearer ? { authorization: `Bearer ${options.bearer}` } : {};
+}
+
+/** Executable and arguments used to launch Graphiti. */
+export interface GraphitiPythonCommand {
+  command: string;
+  args: string[];
+}
+
+/** Resolve configured or app-installed Python, otherwise a uv-provisioned runtime. */
+export async function resolveGraphitiPythonCommand(
+  options: { dev?: boolean } = {},
+  ...args: string[]
+): Promise<GraphitiPythonCommand> {
+  const python = process.env.PYTHON?.trim();
+  if (python) return { command: python, args };
+  if (environmentUtils.isDatabricksAppEnv()) {
+    const installed = await graphitiPython();
+    if (installed) return { command: installed, args };
+  }
+  const extra = options.dev ? "[dev]" : "";
+  const packageArgs = [
+    "--with",
+    `dbx-tools-graphiti${extra}==${(await import("../../index.ts")).PACKAGE_VERSION}`,
+  ];
+  return {
+    command: "uv",
+    args: ["run", "--no-project", ...packageArgs, "--python", "python3", "python", ...args],
+  };
+}
+
+async function _probePython(command: string): Promise<PythonCandidate | undefined> {
+  const versionResult = await exec.spawn(command, ["-V"], {
+    check: false,
+    stdin: "ignore",
+    stdout: "capture",
+    stderr: "capture",
+    trim: true,
+  });
+  if (versionResult.exitCode !== 0) return undefined;
+  const version = coerce(`${versionResult.stdout}\n${versionResult.stderr}`);
+  if (!version) return undefined;
+
+  const importResult = await exec.spawn(command, ["-c", GRAPHITI_IMPORT_PROBE], {
+    check: false,
+    stdin: "ignore",
+    stdout: "ignore",
+    stderr: "ignore",
+  });
+  return importResult.exitCode === 0 ? { command, version } : undefined;
+}
+
+/** Build one HTTP URL for a resolved Graphiti listener. */
+export function graphitiHttpUrl(
+  options: Pick<ResolvedGraphitiOptions, "listen">,
+  path: string,
+): string {
+  return `${sharedOptions
+    .formatListenAddress(options.listen)
+    .replace(/^tcp:/, "http:")}${path.startsWith("/") ? path : `/${path}`}`;
+}
+
+/** Build one readiness-gated Graphiti sidecar from the shared typed options. */
+export async function createGraphitiChildProcess(
+  options: GraphitiRuntimeOptions = {},
+): Promise<GraphitiChildProcess> {
+  const healthCheck = options.healthCheck;
+  const resolved = resolveGraphitiOptions(graphitiOptions(options));
+  const python = await resolveGraphitiPythonCommand(
+    { dev: !resolved.databaseUrl },
+    "-m",
+    "dbx_tools.graphiti",
+  );
+  const child = exec.spawn(python.command, python.args, {
+    env: _runtimeEnvironment(resolved),
+    forwardTerminationSignals: false,
+  });
+  const controller = new AbortController();
+  const startupSignal = options.signal
+    ? AbortSignal.any([controller.signal, options.signal])
+    : controller.signal;
+  const probe = healthCheck ?? (({ signal }) => graphitiHealthCheck(resolved, signal));
+  try {
+    await Promise.race([
+      waitForGraphitiChild(child, probe, resolved.startupTimeoutMs, startupSignal),
+      child.then((result) => {
+        throw new Error(
+          `Graphiti child ${child.pid ?? "<unknown>"} exited before becoming ready ` +
+            `(exit ${result.exitCode})`,
+        );
+      }),
+    ]);
+    return { child };
+  } catch (error) {
+    await child.shutdown();
+    throw error;
+  } finally {
+    controller.abort();
+  }
+}
+
+async function waitForGraphitiChild(
+  child: ChildProcessResult,
+  healthCheck: GraphitiHealthCheck,
+  timeoutMs: number,
+  signal: AbortSignal,
+): Promise<void> {
+  for await (const healthy of asyncUtils.poll(
+    ({ attempt, signal: pollSignal }) =>
+      healthCheck({ process: child, signal: pollSignal, attempt }),
+    {
+      intervalMs: HEALTH_CHECK_INTERVAL_MS,
+      timeoutMs,
+      signal,
+      predicate: (ready) => !ready,
+    },
+  )) {
+    if (healthy) return;
+  }
+}
+
+/** Run the supervised stack until the child exits. */
+export async function runGraphiti(options: GraphitiRuntimeOptions = {}): Promise<void> {
+  const { child } = await createGraphitiChildProcess(options);
+  const outcome = await child;
+  if (outcome.exitCode !== 0) {
+    throw new Error(`Graphiti child exited with status ${outcome.exitCode}`);
+  }
+}

package/src/plugin.ts

@@ -0,0 +1,254 @@
+/**
+ * AppKit memory plugin with Databricks-first tools and lazy Graphiti fallback.
+ *
+ * @module
+ */
+
+import { getExecutionContext, toPlugin, type PluginManifest } from "@databricks/appkit";
+import {
+  defineTool,
+  executeFromRegistry,
+  type ToolkitEntry,
+  type ToolkitOptions,
+  type ToolRegistry,
+} from "@databricks/appkit/beta";
+import {
+  appkit as dbxAppkit,
+  toolkitEntries,
+  ToolRegistryPlugin,
+  type AgentToolExecutionContext,
+} from "@dbx-tools/appkit";
+import { errorUtils, log } from "@dbx-tools/shared-core";
+
+import {
+  RecallMemoryInputSchema,
+  RememberMemoryInputSchema,
+  type GraphitiMemoryBackend,
+  type MemoryTurn,
+  type RecallMemoryInput,
+  type RecallMemoryResult,
+  type RememberMemoryInput,
+  type RememberMemoryResult,
+} from "./_contracts.ts";
+import {
+  MEMORY_CONFIG_SCHEMA,
+  resolveMemoryPluginConfig,
+  type MemoryPluginConfig,
+  type ResolvedMemoryPluginConfig,
+} from "./config.ts";
+import { DatabricksMemoryRuntime } from "./databricks/_runtime.ts";
+
+/** AppKit memory plugin with Databricks managed memory and Graphiti fallback. */
+export class MemoryPlugin extends ToolRegistryPlugin<MemoryPluginConfig> {
+  static phase = "deferred" as const;
+  static manifest: PluginManifest<"memory"> = {
+    name: "memory",
+    displayName: "Memory",
+    description: "Publishes Databricks-first durable memory tools with a lazy Graphiti fallback.",
+    stability: "beta",
+    resources: { required: [], optional: [] },
+    config: { schema: MEMORY_CONFIG_SCHEMA },
+  };
+
+  private readonly logger = log.logger(this);
+  private resolved?: ResolvedMemoryPluginConfig;
+  private databricks?: DatabricksMemoryRuntime;
+  private graphiti?: GraphitiMemoryBackend;
+  private graphitiStarting?: Promise<GraphitiMemoryBackend>;
+  private readonly graphitiStartupAbort = new AbortController();
+  private closed = false;
+
+  protected get toolRegistry(): ToolRegistry {
+    return this.registry();
+  }
+
+  override async setup(): Promise<void> {
+    const resolved = resolveMemoryPluginConfig(this.config);
+    if (resolved.saveTurns === true && resolved.graphiti === false) {
+      throw new Error("AppKit memory saveTurns requires the Graphiti fallback");
+    }
+    this.resolved = resolved;
+    if (resolved.databricks) {
+      this.databricks = new DatabricksMemoryRuntime(
+        getExecutionContext().client,
+        resolved.databricks,
+      );
+    }
+
+    let mastra = false;
+    const context = this.context;
+    const shouldAttachTurnSaver =
+      resolved.saveTurns !== false &&
+      resolved.graphiti !== false &&
+      context?.hasPlugin("mastra") === true;
+    if (shouldAttachTurnSaver && context) {
+      const { attachMastraTurnSaver } = await import("./_mastra.ts");
+      mastra = await attachMastraTurnSaver(context, (turn, signal) => this.saveTurn(turn, signal));
+    }
+    if (resolved.saveTurns === true && !mastra) {
+      throw new Error("AppKit memory saveTurns requires a registered Mastra plugin");
+    }
+    this.logger.info("ready", {
+      databricks: resolved.databricks !== false,
+      graphiti: resolved.graphiti !== false ? "lazy" : false,
+      mastraTurnSaver: mastra,
+    });
+  }
+
+  async toolkit(options?: ToolkitOptions): Promise<Record<string, ToolkitEntry>> {
+    return toolkitEntries.entries("memory", this.getAgentTools(), {
+      ...options,
+      prefix: options?.prefix ?? "",
+    });
+  }
+
+  override executeAgentTool(
+    name: string,
+    args: unknown,
+    signal?: AbortSignal,
+    context?: AgentToolExecutionContext,
+  ): Promise<unknown> {
+    return executeFromRegistry(this.registry(context?.resourceId), name, args, signal);
+  }
+
+  override abortActiveOperations(): void {
+    super.abortActiveOperations();
+    this.graphitiStartupAbort.abort(new Error("AppKit memory is shutting down"));
+  }
+
+  async shutdown(): Promise<void> {
+    this.closed = true;
+    this.graphitiStartupAbort.abort(new Error("AppKit memory is shutting down"));
+    const starting = this.graphitiStarting;
+    const backend = this.graphiti;
+    this.graphitiStarting = undefined;
+    this.graphiti = undefined;
+    await starting?.catch(() => undefined);
+    await backend?.shutdown();
+  }
+
+  private registry(resourceId?: string): ToolRegistry {
+    return {
+      remember: defineTool({
+        description:
+          "Save a durable user preference, fact, decision, or instruction for future conversations.",
+        schema: RememberMemoryInputSchema,
+        annotations: { effect: "write", requiresUserContext: true },
+        autoInheritable: false,
+        execute: (input, signal) => this.remember(this.actorId(resourceId), input, signal),
+      }),
+      recall_memory: defineTool({
+        description:
+          "Search durable actor-scoped memory for facts and preferences relevant to the current task.",
+        schema: RecallMemoryInputSchema,
+        annotations: { effect: "read", requiresUserContext: true },
+        autoInheritable: false,
+        execute: (input, signal) => this.recall(this.actorId(resourceId), input, signal),
+      }),
+    };
+  }
+
+  private async remember(
+    actorId: string,
+    input: RememberMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RememberMemoryResult> {
+    const databricks = await this.databricksBackend(actorId, signal);
+    if (databricks) {
+      try {
+        return await databricks.remember(actorId, input, signal);
+      } catch (error) {
+        if (!this.shouldFallbackFromDatabricks(error)) throw error;
+      }
+    }
+    return (await this.graphitiBackend()).remember(actorId, input, signal);
+  }
+
+  private async recall(
+    actorId: string,
+    input: RecallMemoryInput,
+    signal?: AbortSignal,
+  ): Promise<RecallMemoryResult> {
+    const parsed = RecallMemoryInputSchema.parse(input);
+    const databricks = await this.databricksBackend(actorId, signal);
+    if (databricks) {
+      try {
+        return await databricks.recall(actorId, parsed, signal);
+      } catch (error) {
+        if (!this.shouldFallbackFromDatabricks(error)) throw error;
+      }
+    }
+    return (await this.graphitiBackend()).recall(actorId, parsed, signal);
+  }
+
+  private async saveTurn(turn: MemoryTurn, signal?: AbortSignal): Promise<void> {
+    if (await this.databricksBackend(turn.actorId, signal)) return;
+    await (await this.graphitiBackend()).saveTurn(turn, signal);
+  }
+
+  private async databricksBackend(
+    actorId: string,
+    signal?: AbortSignal,
+  ): Promise<DatabricksMemoryRuntime | undefined> {
+    const backend = this.databricks;
+    if (!backend) return undefined;
+    if (await backend.available(actorId, signal)) return backend;
+    const configured = this.requireConfig().databricks;
+    if (configured && configured.mode === "required") {
+      throw new Error("Required Databricks managed memory is unavailable for the current actor");
+    }
+    return undefined;
+  }
+
+  private shouldFallbackFromDatabricks(error: unknown): boolean {
+    const configured = this.requireConfig().databricks;
+    return (
+      configured !== false &&
+      configured.mode === "auto" &&
+      errorUtils.errorContext(error).hasStatusCode(403, 404)
+    );
+  }
+
+  private async graphitiBackend(): Promise<GraphitiMemoryBackend> {
+    if (this.closed) throw new Error("AppKit memory is closed");
+    if (this.graphiti?.running) return this.graphiti;
+    if (this.graphiti) {
+      await this.graphiti.shutdown();
+      this.graphiti = undefined;
+    }
+    const options = this.requireConfig().graphiti;
+    if (!options) throw new Error("Graphiti fallback is disabled");
+    this.graphitiStarting ??= import("./graphiti/_backend.ts")
+      .then(({ startGraphitiBackend }) =>
+        startGraphitiBackend(options, this.graphitiStartupAbort.signal),
+      )
+      .then((backend) => {
+        if (this.closed) {
+          return backend.shutdown().then(() => {
+            throw new Error("AppKit memory closed during Graphiti startup");
+          });
+        }
+        this.graphiti = backend;
+        return backend;
+      })
+      .finally(() => {
+        this.graphitiStarting = undefined;
+      });
+    return this.graphitiStarting;
+  }
+
+  private actorId(resourceId?: string): string {
+    if (resourceId) return resourceId;
+    const context = dbxAppkit.tryGetExecutionContext();
+    if (!context) throw new Error("Memory tools require an AppKit execution context");
+    return "userId" in context ? context.userId : context.serviceUserId;
+  }
+
+  private requireConfig(): ResolvedMemoryPluginConfig {
+    if (!this.resolved) throw new Error("AppKit memory is not initialized");
+    return this.resolved;
+  }
+}
+
+/** AppKit registration factory for {@link MemoryPlugin}. */
+export const memory = toPlugin(MemoryPlugin);