@pai/storage-mongo

Package: @pai/storage-mongo

Persists PAI thread runtime state in MongoDB. This provider is intended for shared server deployments that need durable threads, run lifecycle records, queued turns, and active-run leases across processes.

Quick start

import { createAgentRuntime } from "@pai/core";
import { createMongoThreadStore } from "@pai/storage-mongo";

const storage = createMongoThreadStore({
  url: process.env.MONGO_URL!,
});

await storage.initialize();

const runtime = createAgentRuntime({
  agent,
  storage,
  scopeKey: (identity) => identity.workspaceId,
});

initialize() creates the package-owned MongoDB client connection, creates required indexes, and verifies transaction support. The store also initializes lazily on first use, but calling it during server startup surfaces connection/index/topology errors before accepting traffic.

Split writes are committed with MongoDB transactions, so the connection must target a replica set or sharded cluster. For local development, start MongoDB as a single-node replica set rather than a standalone mongod.

Options

type MongoThreadStoreOptions = {
  /** MongoDB connection string. If dbName is omitted, the URI database is used. */
  url: string;

  /** Database name. Defaults to the URI database, then "pai". */
  dbName?: string;

  /** Thread collection name. Defaults to "pai_threads". */
  collectionName?: string;

  /** Test-only clock override. Production uses MongoDB server time. */
  now?: () => Date;

  /**
   * Best-effort operational warning callback.
   *
   * Used for recoverable storage integrity issues, such as skipping an invalid
   * thread document during listThreads().
   */
  onWarning?: (warning: MongoThreadStoreWarning) => void | Promise<void>;
};

When now is omitted, observations, commits, and lease expiry use time reported by the MongoDB primary rather than the application process clock. The override exists only for deterministic tests and should not be configured in production.

onWarning lets server code route recoverable Mongo storage issues into the application logger or observability system. For example, listThreads() skips thread documents whose stored identity is unsafe to trust so one bad row cannot fail the whole list operation:

const storage = createMongoThreadStore({
  url: process.env.MONGO_URL!,
  onWarning: (warning) => {
    logger.warn({ warning }, "PAI Mongo storage warning");
  },
});

Warning callbacks are best-effort. If the callback throws or rejects, the store ignores that failure so logging cannot break storage reads.

Data model

PAI stores runtime records in split MongoDB collections keyed by { scopeKey, threadId }.

With the default collectionName: "pai_threads" the provider uses:

When collectionName is customized, related collection names are derived from that base. For example, collectionName: "agent_threads" uses agent_threads_messages, agent_threads_pending_messages, and agent_threads_runs.

The thread collection does not embed transcript messages, queued items, or run records. This keeps root listing and metadata updates bounded as transcripts grow. readThread() queries only the selected child families and returns them from the same committed ThreadVersion.

Documents are flat: the logical record is stored at the top level with no nested payload wrapper, so scopeKey, threadId, and each child id appear exactly once. Alongside them sit the provider's own keys, underscore-prefixed so they can never collide with a protocol field:

• _id — a meaningless ObjectId assigned by MongoDB. Logical identity is not derived from it; the scope_thread_*_unique indexes below are what enforce uniqueness.
• _schemaVersion — currently 1. Underscore-prefixed because ThreadRecord already owns a version field meaning the CAS token.
• _order — a provider-maintained dense integer on child documents only.

Timestamps are stored as native BSON Date at every depth, and converted back to ISO strings when a record is read. Opaque JSON — application metadata and privateMetadata, tool input/output, content providerMetadata, and the separate tool callProviderMetadata/resultProviderMetadata bags — is never converted, so an ISO-looking string stored there round-trips byte-identical. The strict message parser rejects unrecognized part types; there is no durable unknown-part payload.

The thread root contains required strict-JSON privateMetadata and metadata objects. Private metadata is server-only and root-local, so metadata lookup and cursor updates do not scan transcript child documents.

Each pai_runs document contains both required strict-JSON metadata objects: server-only privateMetadata and server-written, client-readable metadata. A document missing a required field is invalid current data; the schema number alone does not make an older development document compatible with the current record shape.

The provider creates these indexes:

// thread collection
{ scopeKey: 1, threadId: 1 } // unique
{ scopeKey: 1, updatedAt: -1, threadId: 1 }
{ scopeKey: 1, listVisibility: 1, updatedAt: -1, threadId: 1 }
{ scopeKey: 1, agent: 1, listVisibility: 1, updatedAt: -1, threadId: 1 }
{ scopeKey: 1, agent: 1, listVisibility: 1, parentThreadId: 1, updatedAt: -1, threadId: 1 }
{ scopeKey: 1, agent: 1, listVisibility: 1, category: 1, updatedAt: -1, threadId: 1 }

// message collection
{ scopeKey: 1, threadId: 1, messageId: 1 } // unique
{ scopeKey: 1, threadId: 1, _order: 1 } // unique
{ scopeKey: 1, threadId: 1, "visibility.transcript": 1, _order: 1 }
{ scopeKey: 1, threadId: 1, "visibility.context": 1, _order: 1 }
{ scopeKey: 1, threadId: 1, "association.kind": 1, "association.runId": 1, _order: 1 }
{ scopeKey: 1, threadId: 1, status: 1, _order: 1 }

// queued item collection
{ scopeKey: 1, threadId: 1, queueItemId: 1 } // unique
{ scopeKey: 1, threadId: 1, _order: 1 } // unique
{ scopeKey: 1, threadId: 1, runId: 1, queueItemId: 1 }

// run collection
{ scopeKey: 1, threadId: 1, runId: 1 } // unique
{ scopeKey: 1, threadId: 1, _order: 1 } // unique
{ scopeKey: 1, threadId: 1, status: 1, runId: 1 }
{ scopeKey: 1, threadId: 1, status: 1, "admission.queueItemId": 1, runId: 1 }

agent and listVisibility lead the thread filter keys because the ordinary client list constrains both on every call — they are sub-partitions rather than optional filters. Neither agent nor includeHidden is reachable from a remote client: the HTTP receiver enumerates its query parameters rather than forwarding them.

The visibility predicate is always present, never omitted. Listing visible threads sends listVisibility: "visible"; includeHidden: true sends { $in: ["visible", "hidden"] } rather than dropping the key. Both forms keep the key bound, which is what preserves the trailing updatedAt ordering — an unconstrained key ahead of the sort columns does not, and neither does a negation like $ne: "hidden", which cannot bound an index at all. MongoDB serves the $in form by scanning each value and combining them with a streaming SORT_MERGE.

The two cross-agent indexes serve the trusted list, which is the only surface that reads across agents or includes hidden threads. It needs both: a middle key the query leaves unconstrained cannot deliver updatedAt order. Trusted lists that also filter by parentThreadId or category fall back to residual filtering — an operator surface, not a user hot path.

Safety model

Delete commits are scoped to one { scopeKey, threadId }. They transactionally remove only the matching root and provider-owned child records. The provider does not call dropDatabase(), drop(), or unscoped deleteMany().

Writes compare the complete ThreadVersion and commit root, child-family, and lease changes in one MongoDB transaction. A stale version returns { type: "conflict", currentVersion } without applying a partial mutation. Deleting and recreating a key produces a new threadInstanceId.

The provider stores strict JSON and validates schema-versioned documents on read. listThreads() skips an invalid root and reports it through onWarning; targeted reads reject invalid persisted data instead of silently coercing it. There is no legacy-document migration or compatibility read path.

Mongo page tokens encode the complete original query and continuation position. They are opaque to callers, scope-bound, and distinct from ThreadVersion.

Conformance

createMongoThreadStore passes the full @pai/storage/test conformance suite. The package's tests run against a MongoDB 7 in-memory replica set and include split collection layout, schema validation, scoped deletion, incarnation-safe compare-and-swap writes, concurrent commits, provider-clock leases, and Mongo-specific pagination.

import { randomUUID } from "node:crypto";
import { createThreadStoreConformanceSuite } from "@pai/storage/test";
import { createMongoThreadStore } from "@pai/storage-mongo";

createThreadStoreConformanceSuite({
  name: "createMongoThreadStore",
  createStore: async () => {
    const store = createMongoThreadStore({
      url: mongoUrl,
      dbName: "pai_test",
      collectionName: `pai_threads_${randomUUID()}`,
    });
    await store.initialize();
    return store;
  },
});