diff --git a/apps/hub/src/mailbox-persist.ts b/apps/hub/src/mailbox-persist.ts index e33bc1e46..89f0e835d 100644 --- a/apps/hub/src/mailbox-persist.ts +++ b/apps/hub/src/mailbox-persist.ts @@ -25,11 +25,12 @@ export function createHubPersistMailWithSessionEnsure Promise, ): (args: MailboxPersistArgs) => Promise { return async (args) => { - const sender = await resolveRoutableAddress(db, args.senderAddress); + let sender = await resolveRoutableAddress(db, args.senderAddress); try { const eventCollectors = eventCollectorsRef.current; if (sender !== undefined && eventCollectors !== undefined) { await ensureRunSession({ db, eventCollectors, runId: sender.id }); + sender = await resolveRoutableAddress(db, args.senderAddress); } } catch (err) { reportError(err, { @@ -37,10 +38,9 @@ export function createHubPersistMailWithSessionEnsure { + let db: DB; + + const tenantId = generateId("tenant"); + const userPrincipalId = generateId("principal"); + const runPrincipalId = generateId("principal"); + const definitionId = generateId("workflowDefinition"); + const runId = generateId("workflowRun"); + const domain = `first-reply-${tenantId}.localhost`; + const sessionId = generateId("session"); + const unanchoredRunId = generateId("workflowRun"); + const unanchoredSessionId = generateId("session"); + const unanchoredAddress = `${unanchoredRunId}@${domain}`; + const closedRunId = generateId("workflowRun"); + const closedSessionId = generateId("session"); + const closedAddress = `${closedRunId}@${domain}`; + const doneRunId = generateId("workflowRun"); + const doneSessionId = generateId("session"); + const doneAddress = `${doneRunId}@${domain}`; + const donePrincipalId = generateId("principal"); + const closedPrincipalId = generateId("principal"); + const runIds = [runId, unanchoredRunId, closedRunId, doneRunId]; + const sessionIds = [sessionId, unanchoredSessionId, closedSessionId, doneSessionId]; + const address = `${runId}@${domain}`; + + beforeAll(async () => { + if (databaseUrl === undefined) return; + db = createDB(dbConfigFromUrl(databaseUrl)); + + await db.db.insert(schema.tenant).values({ + id: tenantId, + name: "First Reply Test Tenant", + slug: `first-reply-${tenantId.replaceAll("_", "-")}`, + domain, + parentId: null, + config: null, + }); + await db.db.insert(schema.principal).values([ + { id: userPrincipalId, tenantId, kind: "user", refId: "usr_test", status: "active" }, + { id: runPrincipalId, tenantId, kind: "workflow", refId: runId, status: "active" }, + { id: donePrincipalId, tenantId, kind: "workflow", refId: doneRunId, status: "active" }, + { id: closedPrincipalId, tenantId, kind: "workflow", refId: closedRunId, status: "active" }, + ]); + await db.db.insert(schema.workflowDefinition).values({ + id: definitionId, + tenantId, + creatorPrincipalId: userPrincipalId, + assetId: null, + name: "first-reply-test-definition", + }); + const seedRun = async ( + id: string, + session: string, + principalId: string | null, + runAddress: string, + ) => { + await db.db.insert(schema.workflowRun).values({ + id, + definitionId, + anchorRunId: id, + tenantId, + principalId, + address: runAddress, + status: "running", + }); + await db.db.insert(schema.workflowRunLaunchSpec).values({ + anchorRunId: id, + sessionId: session, + deploymentDomain: domain, + sourceAuthorityPrincipalId: userPrincipalId, + frozenApprovalBundle: {}, + sourceOfferingIds: [], + defaultSourceOfferingId: "off_test", + deployContent: { systemPrompt: "" }, + }); + }; + await seedRun(runId, sessionId, runPrincipalId, address); + await seedRun(unanchoredRunId, unanchoredSessionId, null, unanchoredAddress); + await seedRun(doneRunId, doneSessionId, donePrincipalId, doneAddress); + await seedRun(closedRunId, closedSessionId, closedPrincipalId, closedAddress); + }); + + afterAll(async () => { + if (databaseUrl === undefined) return; + await db.db.delete(schema.inferenceTurn).where(eq(schema.inferenceTurn.tenantId, tenantId)); + await db.db.delete(schema.agentSession).where(inArray(schema.agentSession.id, sessionIds)); + await db.db + .delete(schema.workflowRunLaunchSpec) + .where(inArray(schema.workflowRunLaunchSpec.anchorRunId, runIds)); + await db.db.delete(schema.workflowRun).where(inArray(schema.workflowRun.id, runIds)); + await db.db + .delete(schema.workflowDefinition) + .where(eq(schema.workflowDefinition.id, definitionId)); + await db.db + .delete(schema.principal) + .where( + inArray(schema.principal.id, [ + runPrincipalId, + closedPrincipalId, + donePrincipalId, + userPrincipalId, + ]), + ); + await db.db.delete(schema.tenant).where(eq(schema.tenant.id, tenantId)); + }); + + async function waitFor(done: () => boolean | Promise): Promise { + for (let i = 0; i < 50 && !(await done()); i++) await Bun.sleep(20); + } + const start = (seq: number) => ({ type: "inference.start", seq, data: { model: "m" } }) as const; + const turnsFor = (session: string) => + db.db.query.inferenceTurn.findMany({ where: eq(schema.inferenceTurn.sessionId, session) }); + + test("persists the run's first turn although no collector existed yet", async () => { + const lazy = withLazyRunCollector(createEventCollectorRegistry({ db: db.db }), db.db); + + lazy.dispatch(address, start(1)); + await Bun.sleep(300); + + expect(lazy.has(address)).toBe(true); + expect(await turnsFor(sessionId)).toHaveLength(1); + }); + + test("creates no collector until the run has a session, then retries on the next event", async () => { + const lazy = withLazyRunCollector(createEventCollectorRegistry({ db: db.db }), db.db); + + lazy.dispatch(unanchoredAddress, start(1)); + await Bun.sleep(300); + expect(lazy.has(unanchoredAddress)).toBe(false); + + const [principal] = await db.db + .insert(schema.principal) + .values({ + id: generateId("principal"), + tenantId, + kind: "workflow", + refId: unanchoredRunId, + status: "active", + }) + .returning(); + await db.db + .update(schema.workflowRun) + .set({ principalId: principal!.id }) + .where(eq(schema.workflowRun.id, unanchoredRunId)); + + lazy.dispatch(unanchoredAddress, start(2)); + await Bun.sleep(300); + expect(lazy.has(unanchoredAddress)).toBe(true); + expect(await turnsFor(unanchoredSessionId)).toHaveLength(1); + + await db.db.delete(schema.inferenceTurn).where(eq(schema.inferenceTurn.tenantId, tenantId)); + await db.db.delete(schema.agentSession).where(eq(schema.agentSession.id, unanchoredSessionId)); + await db.db + .update(schema.workflowRun) + .set({ principalId: null }) + .where(eq(schema.workflowRun.id, unanchoredRunId)); + await db.db.delete(schema.principal).where(eq(schema.principal.id, principal!.id)); + }); + + test("recreates the collector after an abandon so a reconnected run keeps persisting", async () => { + const lazy = withLazyRunCollector(createEventCollectorRegistry({ db: db.db }), db.db); + + lazy.dispatch(closedAddress, start(1)); + await waitFor(() => lazy.has(closedAddress)); + lazy.abandon(closedAddress); + lazy.dispatch(closedAddress, start(2)); + await waitFor(() => lazy.has(closedAddress)); + + expect(lazy.has(closedAddress)).toBe(true); + await waitFor(async () => (await turnsFor(closedSessionId)).length === 2); + expect(await turnsFor(closedSessionId)).toHaveLength(2); + }); + + test("keeps the reply of a run whose done follows its start immediately", async () => { + const lazy = withLazyRunCollector(createEventCollectorRegistry({ db: db.db }), db.db); + + lazy.dispatch(doneAddress, start(1)); + lazy.dispatch(doneAddress, { type: "reactor.done", seq: 2, data: {} }); + await waitFor(async () => (await turnsFor(doneSessionId)).length === 1); + + expect(await turnsFor(doneSessionId)).toHaveLength(1); + }); + + test("does not recreate the collector for mail after the run is done", async () => { + const lazy = withLazyRunCollector(createEventCollectorRegistry({ db: db.db }), db.db); + + lazy.dispatch(doneAddress, start(3)); + lazy.dispatch(doneAddress, { type: "reactor.done", seq: 4, data: {} }); + await waitFor(async () => (await turnsFor(doneSessionId)).length === 2); + await Bun.sleep(100); + await ensureRunSession({ db: db.db, eventCollectors: lazy, runId: doneRunId }); + + expect(lazy.has(doneAddress)).toBe(false); + }); + + test("persists the run's first outbound mail although its session is created mid-call", async () => { + await db.db.delete(schema.agentSession).where(eq(schema.agentSession.id, sessionId)); + const upstreamCalls: unknown[] = []; + const persist = createHubPersistMailWithSessionEnsure( + db.db, + { current: { create: () => {}, has: () => true } }, + async (args) => { + upstreamCalls.push(args); + return []; + }, + ); + + await persist({ senderAddress: address } as Parameters[0]); + + expect(upstreamCalls).toHaveLength(1); + }); +}); diff --git a/packages/workflows/src/index.ts b/packages/workflows/src/index.ts index dd9897493..4f4245053 100644 --- a/packages/workflows/src/index.ts +++ b/packages/workflows/src/index.ts @@ -13,4 +13,8 @@ export { type WorkflowCatalogEntry, type TriggerFieldsValidation, } from "./catalog"; -export { ensureRunSession, type EventCollectorPort } from "./launch/agent-session"; +export { + ensureRunSession, + withLazyRunCollector, + type EventCollectorPort, +} from "./launch/agent-session"; diff --git a/packages/workflows/src/launch/agent-session.ts b/packages/workflows/src/launch/agent-session.ts index f1bcf8bf9..539af23fe 100644 --- a/packages/workflows/src/launch/agent-session.ts +++ b/packages/workflows/src/launch/agent-session.ts @@ -2,8 +2,14 @@ // turn. See docs/agent-session-provisioning.md. import { eq } from "drizzle-orm"; import type { DB } from "@intx/db"; -import { agentSession, workflowRun, workflowRunLaunchSpec } from "@intx/db/schema"; +import { + agentSession, + isLiveWorkflowRunStatus, + workflowRun, + workflowRunLaunchSpec, +} from "@intx/db/schema"; import type { EventCollectorRegistry } from "@intx/hub-sessions"; +import { reportError } from "@corbits/error-sink"; /** The one port every launcher threads the same wrapped * `EventCollectorRegistry` through, never a second registry construction. */ @@ -23,6 +29,15 @@ export async function ensureRunSession(params: { if (runRow === undefined) { return null; } + return ensureSessionForRun(db, eventCollectors, runRow); +} + +async function ensureSessionForRun( + db: DB["db"], + eventCollectors: Pick, + runRow: typeof workflowRun.$inferSelect, +): Promise { + const runId = runRow.id; const launchSpecRow = await db.query.workflowRunLaunchSpec.findFirst({ where: eq(workflowRunLaunchSpec.anchorRunId, runId), @@ -58,10 +73,73 @@ export async function ensureRunSession(params: { .set({ principalId: runRow.principalId, updatedAt: new Date() }) .where(eq(agentSession.id, sessionId)); } - } - if (runRow.address !== null && !eventCollectors.has(runRow.address)) { - eventCollectors.create(runRow.address, runRow.tenantId, sessionId, runRow.id); + if (runRow.address !== null && !eventCollectors.has(runRow.address)) { + eventCollectors.create(runRow.address, runRow.tenantId, sessionId, runRow.id); + } } return sessionId; } + +const MAX_TRACKED_ADDRESSES = 1000; + +function track(set: Set, address: string): void { + if (set.size >= MAX_TRACKED_ADDRESSES) { + const oldest = set.values().next().value; + if (oldest !== undefined) set.delete(oldest); + } + set.add(address); +} + +/** Creates a run's event collector on its first inference event instead of + * its first outbound mail, so the first reply's parts are not dropped. Events + * for an address queue behind its in-flight ensure to keep their order. Once + * a run delivered its terminal event, `create` ignores its address, so no path + * (lazy dispatch or mail) can resurrect one. A disconnect's abandon is not + * terminal: the next event recreates the collector. */ +export function withLazyRunCollector< + R extends Pick, +>(registry: R, db: DB["db"]): R { + const pending = new Map>(); + const closed = new Set(); + + const create: R["create"] = (address, ...rest) => { + if (!closed.has(address)) registry.create(address, ...rest); + }; + + async function ensureForAddress(address: string): Promise { + if (closed.has(address)) return; + try { + const runRow = await db.query.workflowRun.findFirst({ + where: eq(workflowRun.address, address), + }); + if (runRow === undefined || !isLiveWorkflowRunStatus(runRow.status)) return; + await ensureSessionForRun(db, { has: registry.has, create }, runRow); + } catch (err) { + reportError(err, { operation: "workflows.lazyRunCollector.ensure", extra: { address } }); + } + } + + const dispatch: R["dispatch"] = (address, event) => { + const terminal = + event.type === "reactor.done" || (event.type === "reactor.error" && event.data.fatal); + const deliver = () => { + registry.dispatch(address, event); + if (terminal) track(closed, address); + }; + const inFlight = pending.get(address); + if (inFlight === undefined && (registry.has(address) || closed.has(address))) { + deliver(); + return; + } + const next = (inFlight ?? ensureForAddress(address)).then(deliver).catch((err: unknown) => { + reportError(err, { operation: "workflows.lazyRunCollector.dispatch", extra: { address } }); + }); + pending.set(address, next); + void next.finally(() => { + if (pending.get(address) === next) pending.delete(address); + }); + }; + + return { ...registry, create, dispatch }; +}