From 5d27fac7113da92f03c1e044fe67bfa690e250a3 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Thu, 25 Jun 2026 23:37:44 -0400 Subject: [PATCH] fix(sdk): preserve session stream ordering --- packages/core/src/event.ts | 103 +++++++++++-------- packages/core/src/session.ts | 45 ++++---- packages/core/src/session/execution.ts | 13 +-- packages/core/src/session/run-coordinator.ts | 20 ++-- packages/core/test/event.test.ts | 51 +++++++++ packages/core/test/session-prompt.test.ts | 78 ++++++++++++++ packages/schema/src/session-event.ts | 6 ++ 7 files changed, 228 insertions(+), 88 deletions(-) diff --git a/packages/core/src/event.ts b/packages/core/src/event.ts index 1a1c08c3c9..d089627f77 100644 --- a/packages/core/src/event.ts +++ b/packages/core/src/event.ts @@ -72,7 +72,11 @@ export interface Interface { readonly after?: number readonly live: (event: Payload) => boolean }) => Effect.Effect< - { readonly replay: ReadonlyArray; readonly updates: Stream.Stream }, + { + readonly replay: ReadonlyArray + readonly updates: Stream.Stream + readonly offer: (event: Payload, position?: "after" | "before") => boolean + }, never, Scope.Scope > @@ -103,12 +107,12 @@ export const layerWith = (options?: LayerOptions) => Effect.gen(function* () { const pubsub = { all: yield* PubSub.unbounded(), - durable: new Map>>(), + durable: new Map>>(), typed: new Map>(), } const projectors = new Map() // TODO: Bind durable projectors to exact type+version before supporting incompatible historical payloads. - const listeners = new Array() + const listeners = new Set() const { db } = yield* Database.Service const getOrCreate = (definition: Definition) => @@ -284,7 +288,7 @@ export const layerWith = (options?: LayerOptions) => if (committed) { yield* Effect.forEach( pubsub.durable.get(committed.aggregateID) ?? [], - (wake) => PubSub.publish(wake, undefined), + (wake) => PubSub.publish(wake, committed.seq), { discard: true }, ) } @@ -513,7 +517,7 @@ export const layerWith = (options?: LayerOptions) => const subscribeDurable = (aggregateID: string) => Effect.gen(function* () { - const wake = yield* PubSub.sliding(1) + const wake = yield* PubSub.sliding(1) const subscription = yield* PubSub.subscribe(wake) yield* Effect.acquireRelease( Effect.sync(() => { @@ -531,70 +535,83 @@ export const layerWith = (options?: LayerOptions) => return subscription }) + const aggregateDrain = ( + aggregateID: string, + after: number, + ): ((through?: number) => Effect.Effect>) => { + let sequence = after + return (through?: number) => + through !== undefined && through <= sequence + ? Effect.succeed>([]) + : Effect.suspend(() => readAfter(aggregateID, sequence, through)).pipe( + Effect.tap((events) => + Effect.sync(() => { + sequence = events.at(-1)?.durable?.seq ?? sequence + }), + ), + ) + } + const durable = (input: { readonly aggregateID: string; readonly after?: number }): Stream.Stream => Stream.unwrap( Effect.gen(function* () { const wakes = yield* subscribeDurable(input.aggregateID) - let sequence = input.after ?? -1 - const read = Effect.suspend(() => readAfter(input.aggregateID, sequence)).pipe( - Effect.tap((events) => - Effect.sync(() => { - sequence = events.at(-1)?.durable?.seq ?? sequence - }), - ), - ) - const historical = yield* read - const live = Stream.fromSubscription(wakes).pipe( - Stream.mapEffect(() => read), - Stream.flattenIterable, - ) + const drain = aggregateDrain(input.aggregateID, input.after ?? -1) + const historical = yield* drain() + const live = Stream.fromSubscription(wakes).pipe(Stream.mapEffect(drain), Stream.flattenIterable) return Stream.concat(Stream.fromIterable(historical), live) }), ) const observeAggregate: Interface["observeAggregate"] = (input) => Effect.gen(function* () { - type Signal = { readonly _tag: "durable" } | { readonly _tag: "live"; readonly event: Payload } + type Signal = + | { readonly _tag: "durable" } + | { readonly _tag: "transient"; readonly event: Payload; readonly position: "after" | "before" } const signals = yield* Queue.dropping(256) + const wakes = yield* subscribeDurable(input.aggregateID) + let durableQueued = false + let durableThrough = -1 + const offer = (event: Payload, position: "after" | "before" = "after") => { + const offered = Queue.offerUnsafe(signals, { _tag: "transient", event, position }) + if (!offered) Queue.endUnsafe(signals) + return offered + } const unsubscribe = yield* listen((event) => Effect.sync(() => { - if (event.durable?.aggregateID === input.aggregateID) { - // Durable payloads are only wakeups; a full queue already contains work that will drain the database. - Queue.offerUnsafe(signals, { _tag: "durable" }) - return - } - if (!input.live(event)) return - if (!Queue.offerUnsafe(signals, { _tag: "live", event })) Queue.endUnsafe(signals) + if (input.live(event)) offer(event) }), ) yield* Effect.addFinalizer(() => unsubscribe.pipe(Effect.andThen(Queue.shutdown(signals)))) + yield* Stream.runForEach(Stream.fromSubscription(wakes), (sequence) => + Effect.sync(() => { + durableThrough = Math.max(durableThrough, sequence) + if (durableQueued) return + durableQueued = Queue.offerUnsafe(signals, { _tag: "durable" }) + }), + ).pipe(Effect.forkScoped) const cutoff = yield* latestSequence(db, input.aggregateID) const replay = yield* readAfter(input.aggregateID, input.after ?? -1, cutoff) - let sequence = cutoff - const drain = Effect.suspend(() => readAfter(input.aggregateID, sequence)).pipe( - Effect.tap((events) => - Effect.sync(() => { - sequence = events.at(-1)?.durable?.seq ?? sequence - }), - ), - ) + const drain = aggregateDrain(input.aggregateID, cutoff) const updates = Stream.fromQueue(signals).pipe( - Stream.mapEffect((signal) => - drain.pipe(Effect.map((events) => (signal._tag === "live" ? [...events, signal.event] : events))), - ), + Stream.mapEffect((signal) => { + if (signal._tag === "durable") { + durableQueued = false + return drain(durableThrough) + } + if (signal.position === "before") return Effect.succeed([signal.event]) + return drain().pipe(Effect.map((events) => [...events, signal.event])) + }), Stream.flattenIterable, ) - return { replay, updates } + return { replay, updates, offer } }) const listen = (listener: Subscriber): Effect.Effect => Effect.sync(() => { - listeners.push(listener) - return Effect.sync(() => { - const index = listeners.indexOf(listener) - if (index >= 0) listeners.splice(index, 1) - }) + listeners.add(listener) + return Effect.sync(() => listeners.delete(listener)).pipe(Effect.asVoid) }) const project = (definition: D, projector: Subscriber): Effect.Effect => diff --git a/packages/core/src/session.ts b/packages/core/src/session.ts index d422dca323..49e083de16 100644 --- a/packages/core/src/session.ts +++ b/packages/core/src/session.ts @@ -186,8 +186,11 @@ export const layer = Layer.unwrap( const store = yield* SessionStore.Service const locations = yield* LocationServiceMap const decodeMessage = Schema.decodeUnknownEffect(SessionMessage.Message) - const isDurableSessionEvent = Schema.is(SessionEvent.Durable) const sessionEventTypes = new Set(SessionEvent.Definitions.map((definition) => definition.type)) + const isSessionEvent = (event: EventV2.Payload): event is SessionEvent.Event => + sessionEventTypes.has(event.type) + const isDurableSessionEvent = (event: EventV2.Payload): event is SessionEvent.DurableEvent => + event.durable !== undefined && isSessionEvent(event) const decode = (row: typeof SessionMessageTable.$inferSelect) => decodeMessage({ ...row.data, id: row.id, type: row.type }).pipe( Effect.mapError( @@ -351,36 +354,24 @@ export const layer = Layer.unwrap( aggregateID: input.sessionID, after: input.after, live: (event) => - event.durable === undefined && - sessionEventTypes.has(event.type) && - typeof event.data === "object" && - event.data !== null && - "sessionID" in event.data && - event.data.sessionID === input.sessionID, + event.durable === undefined && isSessionEvent(event) && event.data.sessionID === input.sessionID, }) - const initialActivity: SessionEvent.Activity = { - id: EventV2.ID.create(), - type: SessionEvent.Activity.type, - data: { sessionID: input.sessionID, active: yield* activity.snapshot }, - } - const replay = observed.replay.filter((event): event is SessionEvent.DurableEvent => - isDurableSessionEvent(event), - ) - const updates = observed.updates.pipe( - Stream.filter((event): event is SessionEvent.Event => Schema.is(SessionEvent.All)(event)), - ) - const activityChanges = activity.changes.pipe( - Stream.map( - (active): SessionEvent.Activity => ({ - id: EventV2.ID.create(), - type: SessionEvent.Activity.type, - data: { sessionID: input.sessionID, active }, - }), + const initialActivity = SessionEvent.makeActivity( + input.sessionID, + yield* activity.attach((active) => + observed.offer(SessionEvent.makeActivity(input.sessionID, active), active ? "before" : "after"), ), ) - return Stream.fromIterable(replay).pipe( + return Stream.fromIterable(observed.replay.filter(isDurableSessionEvent)).pipe( Stream.concat(Stream.make(initialActivity)), - Stream.concat(Stream.merge(updates, activityChanges)), + Stream.concat( + observed.updates.pipe( + Stream.filter( + (event): event is SessionEvent.StreamEvent => + event.type === SessionEvent.Activity.type || isSessionEvent(event), + ), + ), + ), ) }), ), diff --git a/packages/core/src/session/execution.ts b/packages/core/src/session/execution.ts index 13eb8de860..1873215922 100644 --- a/packages/core/src/session/execution.ts +++ b/packages/core/src/session/execution.ts @@ -1,20 +1,15 @@ export * as SessionExecution from "./execution" -import { Context, Effect, Layer, Scope, Stream } from "effect" +import { Context, Effect, Layer, Scope } from "effect" import { SessionRunner } from "./runner/index" +import { SessionRunCoordinator } from "./run-coordinator" import { SessionSchema } from "./schema" export interface Interface { /** Snapshots active execution owned by this process. */ readonly active: Effect.Effect> /** Observes foreground ownership with an authoritative initial snapshot. */ - readonly activity: ( - sessionID: SessionSchema.ID, - ) => Effect.Effect< - { readonly snapshot: Effect.Effect; readonly changes: Stream.Stream }, - never, - Scope.Scope - > + readonly activity: (sessionID: SessionSchema.ID) => Effect.Effect /** Starts execution while idle or joins the active execution. */ readonly resume: (sessionID: SessionSchema.ID) => Effect.Effect /** Registers newly recorded work. Repeated wakeups may coalesce. */ @@ -31,7 +26,7 @@ export const noopLayer = Layer.succeed( Service, Service.of({ active: Effect.succeed(new Set()), - activity: () => Effect.succeed({ snapshot: Effect.succeed(false), changes: Stream.never }), + activity: () => Effect.succeed({ attach: () => Effect.succeed(false) }), resume: () => Effect.void, wake: () => Effect.void, interrupt: () => Effect.void, diff --git a/packages/core/src/session/run-coordinator.ts b/packages/core/src/session/run-coordinator.ts index 00738147eb..13b65ba53d 100644 --- a/packages/core/src/session/run-coordinator.ts +++ b/packages/core/src/session/run-coordinator.ts @@ -1,11 +1,10 @@ export * as SessionRunCoordinator from "./run-coordinator" -import { Cause, Deferred, Effect, Exit, Fiber, FiberSet, Queue, Scope, Stream } from "effect" +import { Deferred, Effect, Exit, Fiber, FiberSet, Scope } from "effect" export interface Activity { - /** Discards earlier transitions and snapshots current ownership. */ - readonly snapshot: Effect.Effect - readonly changes: Stream.Stream + /** Installs a synchronous observer and snapshots current ownership atomically. */ + readonly attach: (observer: (active: boolean) => void) => Effect.Effect } /** Serializes execution for each key while allowing different keys to run concurrently. */ @@ -119,9 +118,9 @@ export const make = (options: { const activity = (key: Key) => Effect.gen(function* () { - const queue = yield* Queue.dropping(256) + let attached: ((active: boolean) => void) | undefined const observer = (value: boolean) => { - if (!Queue.offerUnsafe(queue, value)) Queue.endUnsafe(queue) + attached?.(value) } yield* Effect.acquireRelease( Effect.sync(() => { @@ -134,11 +133,14 @@ export const make = (options: { const observers = activityObservers.get(key) observers?.delete(observer) if (observers?.size === 0) activityObservers.delete(key) - }).pipe(Effect.andThen(Queue.shutdown(queue))), + }), ) return { - snapshot: Queue.clear(queue).pipe(Effect.andThen(Effect.sync(() => active.has(key)))), - changes: Stream.fromQueue(queue), + attach: (next: (active: boolean) => void) => + Effect.sync(() => { + attached = next + return active.has(key) + }), } }) diff --git a/packages/core/test/event.test.ts b/packages/core/test/event.test.ts index 0c6280fe54..3273491923 100644 --- a/packages/core/test/event.test.ts +++ b/packages/core/test/event.test.ts @@ -450,6 +450,30 @@ describe("EventV2", () => { }), ) + it.effect("observes replay commits that are not republished", () => + Effect.gen(function* () { + const events = yield* EventV2.Service + const aggregateID = Session.ID.create() + const observed = yield* events.observeAggregate({ aggregateID, live: () => false }) + const update = yield* observed.updates.pipe(Stream.take(1), Stream.runCollect, Effect.forkScoped) + + yield* events.replay( + { + id: EventV2.ID.create(), + aggregateID, + seq: 0, + type: EventV2.versionedType(DurableMessage.type, 1), + data: durableData(aggregateID, "replayed"), + }, + { publish: false }, + ) + + expect(Array.from(yield* Fiber.join(update)).map((event) => event.data)).toEqual([ + durableData(aggregateID, "replayed"), + ]) + }), + ) + it.effect("drains causally preceding durable rows before a live payload", () => Effect.gen(function* () { const events = yield* EventV2.Service @@ -488,6 +512,33 @@ describe("EventV2", () => { }), ) + it.effect("coalesces durable notifications without repeated empty reads", () => + Effect.gen(function* () { + let reads = 0 + const eventLayer = EventV2.layerWith({ + beforeAggregateRead: () => + Effect.sync(() => { + reads++ + }), + }).pipe(Layer.provide(Database.defaultLayer)) + + yield* Effect.gen(function* () { + const events = yield* EventV2.Service + const aggregateID = Session.ID.create() + const count = 20 + const observed = yield* events.observeAggregate({ aggregateID, live: () => false }) + + for (let index = 0; index < count; index++) { + yield* events.publish(DurableMessage, durableData(aggregateID, String(index))) + } + const updates = yield* observed.updates.pipe(Stream.take(count), Stream.runCollect) + + expect(updates).toHaveLength(count) + expect(reads).toBe(2) + }).pipe(Effect.provide(Layer.mergeAll(Database.defaultLayer, eventLayer))) + }), + ) + it.effect("coalesces durable aggregate wakes while draining every committed event", () => Effect.gen(function* () { const events = yield* EventV2.Service diff --git a/packages/core/test/session-prompt.test.ts b/packages/core/test/session-prompt.test.ts index 1fc1301321..3e921b8a15 100644 --- a/packages/core/test/session-prompt.test.ts +++ b/packages/core/test/session-prompt.test.ts @@ -257,6 +257,84 @@ describe("SessionV2.prompt", () => { }), ) + it.effect("orders activity around the durable rows owned by a run", () => + Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + const events = yield* EventV2.Service + const assistantMessageID = SessionMessage.ID.create() + const streamed = yield* session.events({ sessionID }).pipe(Stream.take(5), Stream.runCollect, Effect.forkScoped) + yield* Effect.yieldNow + + yield* setActivity(true) + yield* events.publish(SessionEvent.Text.Started, { + sessionID, + assistantMessageID, + textID: "text-activity", + timestamp: yield* DateTime.now, + }) + yield* events.publish(SessionEvent.Text.Ended, { + sessionID, + assistantMessageID, + textID: "text-activity", + text: "done", + timestamp: yield* DateTime.now, + }) + yield* setActivity(false) + + expect(Array.from(yield* Fiber.join(streamed)).map((event) => [event.type, event.data])).toEqual([ + ["session.activity", { sessionID, active: false }], + ["session.activity", { sessionID, active: true }], + ["session.next.text.started", expect.objectContaining({ sessionID })], + ["session.next.text.ended", expect.objectContaining({ sessionID })], + ["session.activity", { sessionID, active: false }], + ]) + }), + ) + + it.effect("terminates the Session stream when live event buffering saturates", () => + Effect.gen(function* () { + yield* setup + const session = yield* SessionV2.Service + const events = yield* EventV2.Service + const initial = yield* Deferred.make() + const release = yield* Deferred.make() + let seenInitial = false + const streamed = yield* session.events({ sessionID }).pipe( + Stream.mapEffect((event) => { + if (!seenInitial) { + seenInitial = true + return Deferred.succeed(initial, undefined).pipe(Effect.as(event)) + } + return Deferred.await(release).pipe(Effect.as(event)) + }), + Stream.runCollect, + Effect.forkScoped, + ) + yield* Deferred.await(initial) + + const timestamp = yield* DateTime.now + const assistantMessageID = SessionMessage.ID.create() + yield* Effect.forEach( + Array.from({ length: 1_000 }, (_, index) => index), + (index) => + events.publish(SessionEvent.Text.Delta, { + sessionID, + assistantMessageID, + textID: "text-saturation", + delta: String(index), + timestamp, + }), + { concurrency: "unbounded", discard: true }, + ) + yield* Deferred.succeed(release, undefined) + + const result = Array.from(yield* Fiber.join(streamed)) + expect(result[0]?.type).toBe("session.activity") + expect(result.length).toBeLessThan(1_001) + }), + ) + it.effect("resumes through a recorded message without appending another prompt", () => Effect.gen(function* () { yield* setup diff --git a/packages/schema/src/session-event.ts b/packages/schema/src/session-event.ts index 7667b79378..cb5f359c7a 100644 --- a/packages/schema/src/session-event.ts +++ b/packages/schema/src/session-event.ts @@ -60,6 +60,12 @@ export const Activity = Event.define({ }) export type Activity = typeof Activity.Type +export const makeActivity = (sessionID: SessionID, active: boolean): Activity => ({ + id: Event.ID.create(), + type: Activity.type, + data: { sessionID, active }, +}) + export const AgentSwitched = Event.define({ type: "session.next.agent.switched", ...options,