refactor(training-agent): adopt @adcp/client 5.1 state-store helpers

.changeset/adopt-adcp-client-5-1-state-helpers.md

@@ -0,0 +1,12 @@
+---
+---
+
+Training agent adopts `@adcp/client` 5.1 state-store helpers:
+
+- `serializeSession` / `deserializeSession` replaced by the SDK's `structuredSerialize` / `structuredDeserialize` (tagged `__adcpType` envelopes for Map/Date). Net: ~40 lines of hand-rolled type coercion removed; disk format is now the SDK canonical form.
+- Per-flush 5 MB size check removed. `PostgresStateStore`/`InMemoryStateStore` enforce `DEFAULT_MAX_DOCUMENT_BYTES` at `put()` time and throw `StateError('PAYLOAD_TOO_LARGE')` automatically.
+- Prototype-pollution `RESERVED_KEYS` filter removed. `structuredDeserialize` operates on tagged envelopes, and SDK key validation protects the store boundary.
+
+Migration 411 clears `training_sessions` rows on deploy to avoid a format mismatch between pre- and post-change serialization. Training sessions have a 1-hour TTL and are sandbox state — losing them is equivalent to a machine restart. Other collections in `adcp_state` are unaffected.
+
+Closes #2269.

server/src/db/migrations/411_clear_training_sessions.sql

@@ -0,0 +1,10 @@
+-- One-time clear of training-agent sessions: the serialization format
+-- changed when we adopted @adcp/client 5.1's structuredSerialize helpers
+-- (tagged envelopes for Map/Date instead of hand-rolled Object.fromEntries).
+--
+-- Training sessions are sandbox state with a 1-hour TTL — losing them is
+-- equivalent to a ~1-hour-stale machine restart. Other collections in
+-- adcp_state (if any — today this table is training-sessions only) are
+-- preserved.
+
+DELETE FROM adcp_state WHERE collection = 'training_sessions';

server/src/training-agent/state.ts

@@ -15,16 +15,13 @@
  */
 
 import { AsyncLocalStorage } from 'node:async_hooks';
-import type {
-  SessionState, AccountRef, BrandRef,
-  MediaBuyState, CreativeState, SignalActivationState, GovernancePlanState,
-  GovernanceCheckState, GovernanceOutcomeState, PropertyListState,
-  CollectionListState, ContentStandardsState, RightsGrantState, UsageRecord,
-} from './types.js';
+import type { SessionState, AccountRef, BrandRef } from './types.js';
 import { cleanupExpiredTasks } from '@adcp/client';
 import {
   InMemoryStateStore,
   PostgresStateStore,
+  structuredSerialize,
+  structuredDeserialize,
   type AdcpStateStore,
 } from '@adcp/client/server';
 import { isDatabaseInitialized, getPool } from '../db/client.js';
@@ -87,15 +84,18 @@ export function runWithSessionContext<T>(fn: () => Promise<T>): Promise<T> {
   return requestCtx.run(ctx, fn);
 }
 
-const MAX_SESSION_JSON_BYTES = 5 * 1024 * 1024; // 5 MB
-
 /**
  * Persist sessions that were actually mutated during the current request.
  *
  * Mutation is detected by comparing the current serialized shape to a snapshot
  * taken when the session was first loaded. Read-only accesses do not write,
  * eliminating unnecessary DB traffic on `get_*` / `list_*` tools.
  *
+ * Size enforcement (5 MB default) and key validation live in the SDK's
+ * state store — `store.put` throws `StateError('PAYLOAD_TOO_LARGE')` or
+ * `StateError('INVALID_ID')` automatically. Failures bubble to the MCP
+ * transport layer so operators notice in alert pipelines.
+ *
  * Known limitation: concurrent requests against the same session key use
  * last-writer-wins semantics. Acceptable for the sandbox training agent where
  * storyboards are sequential. Production sellers should use per-entity
@@ -111,17 +111,6 @@ export async function flushDirtySessions(): Promise<void> {
     const currentJson = stableStringify(current);
     const snapshotJson = ctx.snapshots.get(key);
     if (snapshotJson === currentJson) continue;
-    if (currentJson.length > MAX_SESSION_JSON_BYTES) {
-      // Size cap is a defense-in-depth backstop on top of per-collection
-      // mutation limits. If we hit it, something grew unexpectedly —
-      // don't silently drop the write (the caller already got 200 OK).
-      const err = new Error(
-        `Training agent session "${key}" exceeds ${MAX_SESSION_JSON_BYTES} bytes (${currentJson.length}); refusing to persist`,
-      );
-      logger.error({ key, bytes: currentJson.length }, err.message);
-      errors.push({ key, err });
-      continue;
-    }
     try {
       await store.put(SESSIONS_COLLECTION, key, current);
       ctx.snapshots.set(key, currentJson);
@@ -184,64 +173,62 @@ function createSession(): SessionState {
   };
 }
 
-/** Serialize a SessionState to a JSON-safe object for the state store.
- *
- * `lastGetProductsContext.products` is deterministic from the catalog, so we
- * drop it from persistence and let callers recompute on next request.
- * `proposals` (session-specific drafts) are persisted.
+/**
+ * Serialize a SessionState via the SDK's `structuredSerialize` (tagged
+ * envelopes for Map/Date). Returns a JSON-safe Record.
  */
 function serializeSession(session: SessionState): Record<string, unknown> {
-  return {
-    mediaBuys: Object.fromEntries(session.mediaBuys),
-    creatives: Object.fromEntries(session.creatives),
-    signalActivations: Object.fromEntries(session.signalActivations),
-    governancePlans: Object.fromEntries(session.governancePlans),
-    governanceChecks: Object.fromEntries(session.governanceChecks),
-    governanceOutcomes: Object.fromEntries(session.governanceOutcomes),
-    propertyLists: Object.fromEntries(session.propertyLists),
-    collectionLists: Object.fromEntries(session.collectionLists),
-    contentStandards: Object.fromEntries(session.contentStandards),
-    rightsGrants: Object.fromEntries(session.rightsGrants),
-    usageRecords: session.usageRecords,
-    lastGetProductsProposals: session.lastGetProductsContext?.proposals,
-    createdAt: session.createdAt.toISOString(),
-    lastAccessedAt: session.lastAccessedAt.toISOString(),
+  const persisted = {
+    ...session,
+    // `products` is deterministic from the catalog — dropped from persistence
+    // so callers re-derive on the next request. Only `proposals` (session-
+    // specific drafts from refine workflows) ride along.
+    lastGetProductsContext: session.lastGetProductsContext?.proposals?.length
+      ? { proposals: session.lastGetProductsContext.proposals }
+      : undefined,
   };
+  return structuredSerialize(persisted) as Record<string, unknown>;
 }
 
-const RESERVED_KEYS = new Set(['__proto__', 'constructor', 'prototype']);
-
-/** Deserialize a stored doc back into a SessionState with Map/Date types. */
+/**
+ * Hydrate a stored doc back into a SessionState.
+ *
+ * Security invariant (per #2283 security review): `structuredDeserialize`
+ * walks untrusted JSONB via `Object.entries` and can reconstitute a Map
+ * whose entries include a "constructor" or "__proto__" key. Those are safe
+ * on Map lookups (`.get(k)` returns the stored value, not Object.prototype)
+ * — but only as long as handlers never spread raw request payloads into
+ * session state. A handler that does `session.propertyLists.set(id, req)`
+ * verbatim would re-open the vector. Every existing handler builds
+ * PropertyListState/MediaBuyState/etc field-by-field from validated
+ * primitives, not raw spread. Maintainers: keep it that way.
+ *
+ * Map-field safety: each `SessionState` Map gets an explicit `asMap(...)`
+ * override below. If you add a new Map field to `SessionState`, add a
+ * matching line here or it will hydrate as a raw envelope object.
+ */
 function deserializeSession(data: Record<string, unknown>): SessionState {
-  const asMap = <V>(obj: unknown): Map<string, V> => {
-    if (!obj || typeof obj !== 'object') return new Map();
-    const entries = Object.entries(obj as Record<string, V>)
-      .filter(([k]) => !RESERVED_KEYS.has(k));
-    return new Map(entries);
-  };
-  const asDate = (v: unknown): Date => {
-    if (typeof v === 'string') return new Date(v);
-    return new Date();
-  };
+  const hydrated = structuredDeserialize(data) as Partial<SessionState> & { lastGetProductsContext?: unknown };
+  const fresh = createSession();
+  const asMap = <V>(v: unknown, fallback: Map<string, V>): Map<string, V> =>
+    v instanceof Map ? (v as Map<string, V>) : fallback;
   return {
-    mediaBuys: asMap<MediaBuyState>(data.mediaBuys),
-    creatives: asMap<CreativeState>(data.creatives),
-    signalActivations: asMap<SignalActivationState>(data.signalActivations),
-    governancePlans: asMap<GovernancePlanState>(data.governancePlans),
-    governanceChecks: asMap<GovernanceCheckState>(data.governanceChecks),
-    governanceOutcomes: asMap<GovernanceOutcomeState>(data.governanceOutcomes),
-    propertyLists: asMap<PropertyListState>(data.propertyLists),
-    collectionLists: asMap<CollectionListState>(data.collectionLists),
-    contentStandards: asMap<ContentStandardsState>(data.contentStandards),
-    rightsGrants: asMap<RightsGrantState>(data.rightsGrants),
-    usageRecords: Array.isArray(data.usageRecords) ? data.usageRecords as UsageRecord[] : [],
-    // Only proposals persist; products are deterministic from the catalog, so callers
-    // re-derive on the next request via the fallback in the get_products handler.
-    lastGetProductsContext: Array.isArray(data.lastGetProductsProposals) && data.lastGetProductsProposals.length > 0
-      ? { proposals: data.lastGetProductsProposals as NonNullable<SessionState['lastGetProductsContext']>['proposals'] }
-      : undefined,
-    createdAt: asDate(data.createdAt),
-    lastAccessedAt: asDate(data.lastAccessedAt),
+    ...fresh,
+    ...hydrated,
+    mediaBuys: asMap(hydrated.mediaBuys, fresh.mediaBuys),
+    creatives: asMap(hydrated.creatives, fresh.creatives),
+    signalActivations: asMap(hydrated.signalActivations, fresh.signalActivations),
+    governancePlans: asMap(hydrated.governancePlans, fresh.governancePlans),
+    governanceChecks: asMap(hydrated.governanceChecks, fresh.governanceChecks),
+    governanceOutcomes: asMap(hydrated.governanceOutcomes, fresh.governanceOutcomes),
+    propertyLists: asMap(hydrated.propertyLists, fresh.propertyLists),
+    collectionLists: asMap(hydrated.collectionLists, fresh.collectionLists),
+    contentStandards: asMap(hydrated.contentStandards, fresh.contentStandards),
+    rightsGrants: asMap(hydrated.rightsGrants, fresh.rightsGrants),
+    usageRecords: Array.isArray(hydrated.usageRecords) ? hydrated.usageRecords : [],
+    lastGetProductsContext: (hydrated.lastGetProductsContext as SessionState['lastGetProductsContext']) ?? undefined,
+    createdAt: hydrated.createdAt instanceof Date ? hydrated.createdAt : fresh.createdAt,
+    lastAccessedAt: hydrated.lastAccessedAt instanceof Date ? hydrated.lastAccessedAt : fresh.lastAccessedAt,
   };
 }
 

server/tests/unit/training-agent.test.ts

@@ -812,6 +812,47 @@ describe('session state', () => {
       }
     });
 
+    it('disk format uses structuredSerialize envelopes (Maps round-trip losslessly)', async () => {
+      const { setStateStore } = await import('../../src/training-agent/state.js');
+      const { InMemoryStateStore } = await import('@adcp/client/server');
+      const store = new InMemoryStateStore();
+      setStateStore(store);
+      const key = 'format-test';
+      try {
+        const createdAt = new Date('2026-04-17T10:00:00Z');
+        await runWithSessionContext(async () => {
+          const s = await getSession(key);
+          s.mediaBuys.set('mb_abc', { mediaBuyId: 'mb_abc', status: 'active' } as any);
+          s.mediaBuys.set('mb_def', { mediaBuyId: 'mb_def', status: 'paused' } as any);
+          s.createdAt = createdAt;
+          await flushDirtySessions();
+        });
+
+        // Inspect the raw stored doc — must use SDK's tagged-envelope format.
+        const raw = await store.get('training_sessions', key) as Record<string, unknown>;
+        expect(raw).toBeDefined();
+        const mediaBuys = raw.mediaBuys as { __adcpType?: string; entries?: unknown[] };
+        expect(mediaBuys.__adcpType).toBe('Map');
+        expect(Array.isArray(mediaBuys.entries)).toBe(true);
+        expect(mediaBuys.entries!.length).toBe(2);
+        const created = raw.createdAt as { __adcpType?: string; value?: string };
+        expect(created.__adcpType).toBe('Date');
+        expect(created.value).toBe(createdAt.toISOString());
+
+        // Hydrate via getSession and verify both entries come back as real Maps/Dates.
+        await runWithSessionContext(async () => {
+          const s = await getSession(key);
+          expect(s.mediaBuys).toBeInstanceOf(Map);
+          expect(s.mediaBuys.get('mb_abc')).toEqual({ mediaBuyId: 'mb_abc', status: 'active' });
+          expect(s.mediaBuys.get('mb_def')).toEqual({ mediaBuyId: 'mb_def', status: 'paused' });
+          expect(s.createdAt).toBeInstanceOf(Date);
+          expect(s.createdAt.toISOString()).toBe(createdAt.toISOString());
+        });
+      } finally {
+        setStateStore(null);
+      }
+    });
+
     it('dispatcher skips flush when handler throws (real MCP path)', async () => {
       const { setStateStore } = await import('../../src/training-agent/state.js');
       const { InMemoryStateStore } = await import('@adcp/client/server');