From 40c1812b4b8140768ed6f45af9eb5b576a9fefec Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 17 Apr 2026 15:50:55 -0400 Subject: [PATCH 1/2] refactor: adopt @adcp/client 5.1 state-store helpers in training agent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes #2269. - serializeSession / deserializeSession replaced by structuredSerialize / structuredDeserialize from @adcp/client/server. Disk format is now the SDK canonical form (tagged __adcpType envelopes for Map/Date) instead of our hand-rolled Object.fromEntries + ISO-string approach. - MAX_SESSION_JSON_BYTES manual check in flushDirtySessions removed. PostgresStateStore.put / InMemoryStateStore.put enforce DEFAULT_MAX_DOCUMENT_BYTES (5 MB) at the boundary and throw StateError('PAYLOAD_TOO_LARGE') automatically. - RESERVED_KEYS prototype-pollution filter removed. structuredDeserialize walks tagged envelopes, not Object.entries on untrusted JSONB, and the SDK's validateId / validateCollection protects the store boundary. Migration 411 clears `training_sessions` rows on deploy because the serialization format changed. Training sessions are sandbox state with a 1-hour TTL — losing them is equivalent to a machine restart. Net: server/src/training-agent/state.ts goes from 384 → 356 lines. Tests: 302/302 training-agent + 587/587 tests/ suite pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../adopt-adcp-client-5-1-state-helpers.md | 12 ++ .../411_clear_training_sessions.sql | 10 ++ server/src/training-agent/state.ts | 112 +++++++----------- 3 files changed, 64 insertions(+), 70 deletions(-) create mode 100644 .changeset/adopt-adcp-client-5-1-state-helpers.md create mode 100644 server/src/db/migrations/411_clear_training_sessions.sql diff --git a/.changeset/adopt-adcp-client-5-1-state-helpers.md b/.changeset/adopt-adcp-client-5-1-state-helpers.md new file mode 100644 index 0000000000..3639faefa4 --- /dev/null +++ b/.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. diff --git a/server/src/db/migrations/411_clear_training_sessions.sql b/server/src/db/migrations/411_clear_training_sessions.sql new file mode 100644 index 0000000000..8f46974b92 --- /dev/null +++ b/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'; diff --git a/server/src/training-agent/state.ts b/server/src/training-agent/state.ts index 791d70909a..06834d8340 100644 --- a/server/src/training-agent/state.ts +++ b/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,8 +84,6 @@ export function runWithSessionContext(fn: () => Promise): Promise { 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. * @@ -96,6 +91,11 @@ const MAX_SESSION_JSON_BYTES = 5 * 1024 * 1024; // 5 MB * 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 { 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,47 @@ 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). `lastGetProductsContext.products` is deterministic + * from the catalog so we drop it — `proposals` (session-specific drafts) ride + * along. Returns a JSON-safe Record. */ function serializeSession(session: SessionState): Record { - 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, + lastGetProductsContext: session.lastGetProductsContext?.proposals?.length + ? { proposals: session.lastGetProductsContext.proposals } + : undefined, }; + return structuredSerialize(persisted) as Record; } -const RESERVED_KEYS = new Set(['__proto__', 'constructor', 'prototype']); - -/** Deserialize a stored doc back into a SessionState with Map/Date types. */ +/** Inverse — hydrate a stored doc back into a SessionState. */ function deserializeSession(data: Record): SessionState { - const asMap = (obj: unknown): Map => { - if (!obj || typeof obj !== 'object') return new Map(); - const entries = Object.entries(obj as Record) - .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 & { lastGetProductsContext?: unknown }; + const fresh = createSession(); + // Merge hydrated fields onto a fresh skeleton — missing maps stay empty, + // missing dates stay fresh. Guards against incomplete stored rows. + const asMap = (v: unknown, fallback: Map): Map => + v instanceof Map ? (v as Map) : fallback; return { - mediaBuys: asMap(data.mediaBuys), - creatives: asMap(data.creatives), - signalActivations: asMap(data.signalActivations), - governancePlans: asMap(data.governancePlans), - governanceChecks: asMap(data.governanceChecks), - governanceOutcomes: asMap(data.governanceOutcomes), - propertyLists: asMap(data.propertyLists), - collectionLists: asMap(data.collectionLists), - contentStandards: asMap(data.contentStandards), - rightsGrants: asMap(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['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, }; } From 8505d5e61d52bde727eea7870837b7c138dccdf7 Mon Sep 17 00:00:00 2001 From: Brian O'Kelley Date: Fri, 17 Apr 2026 16:03:58 -0400 Subject: [PATCH 2/2] review: document invariants + pin disk format with explicit test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Applying feedback from code + security review on #2283: - Document the security invariant around RESERVED_KEYS removal: structured- Deserialize DOES walk untrusted JSONB, so prototype pollution is closed at the INPUT boundary (handlers never spread raw request payloads into session state) — not by the SDK. Comment warns maintainers to keep that invariant if adding new handlers. - Document the Map-field maintenance invariant on deserializeSession: each Map in SessionState needs an explicit `asMap(...)` override, or it will hydrate as a raw envelope object. - Lift the "products dropped from persistence" intent onto the call site where the override happens (code-reviewer Should Fix #2). - Add explicit disk-format test that verifies: - Maps serialize as { __adcpType: 'Map', entries: [...] } - Dates serialize as { __adcpType: 'Date', value: ISO } - Round-trip hydrates back to real Map/Date instances Pins the wire format against future SDK regressions. 303/303 training-agent tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) --- server/src/training-agent/state.ts | 27 ++++++++++++---- server/tests/unit/training-agent.test.ts | 41 ++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 6 deletions(-) diff --git a/server/src/training-agent/state.ts b/server/src/training-agent/state.ts index 06834d8340..099792f4c8 100644 --- a/server/src/training-agent/state.ts +++ b/server/src/training-agent/state.ts @@ -175,13 +175,14 @@ function createSession(): SessionState { /** * Serialize a SessionState via the SDK's `structuredSerialize` (tagged - * envelopes for Map/Date). `lastGetProductsContext.products` is deterministic - * from the catalog so we drop it — `proposals` (session-specific drafts) ride - * along. Returns a JSON-safe Record. + * envelopes for Map/Date). Returns a JSON-safe Record. */ function serializeSession(session: SessionState): Record { 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, @@ -189,12 +190,26 @@ function serializeSession(session: SessionState): Record { return structuredSerialize(persisted) as Record; } -/** Inverse — hydrate a stored doc back into a SessionState. */ +/** + * 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): SessionState { const hydrated = structuredDeserialize(data) as Partial & { lastGetProductsContext?: unknown }; const fresh = createSession(); - // Merge hydrated fields onto a fresh skeleton — missing maps stay empty, - // missing dates stay fresh. Guards against incomplete stored rows. const asMap = (v: unknown, fallback: Map): Map => v instanceof Map ? (v as Map) : fallback; return { diff --git a/server/tests/unit/training-agent.test.ts b/server/tests/unit/training-agent.test.ts index 8ad485fb9c..52a69781e3 100644 --- a/server/tests/unit/training-agent.test.ts +++ b/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; + 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');