fix(sdk): preserve session stream ordering

This commit is contained in:
Kit Langton
2026-06-25 23:37:44 -04:00
parent 605fb0b6c5
commit 5d27fac711
7 changed files with 228 additions and 88 deletions
+60 -43
View File
@@ -72,7 +72,11 @@ export interface Interface {
readonly after?: number
readonly live: (event: Payload) => boolean
}) => Effect.Effect<
{ readonly replay: ReadonlyArray<Payload>; readonly updates: Stream.Stream<Payload> },
{
readonly replay: ReadonlyArray<Payload>
readonly updates: Stream.Stream<Payload>
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<Payload>(),
durable: new Map<string, Set<PubSub.PubSub<void>>>(),
durable: new Map<string, Set<PubSub.PubSub<number>>>(),
typed: new Map<string, PubSub.PubSub<Payload>>(),
}
const projectors = new Map<string, Subscriber[]>()
// TODO: Bind durable projectors to exact type+version before supporting incompatible historical payloads.
const listeners = new Array<Subscriber>()
const listeners = new Set<Subscriber>()
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<void>(1)
const wake = yield* PubSub.sliding<number>(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<ReadonlyArray<Payload>>) => {
let sequence = after
return (through?: number) =>
through !== undefined && through <= sequence
? Effect.succeed<ReadonlyArray<Payload>>([])
: 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<Payload> =>
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<Signal, Cause.Done>(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<Unsubscribe> =>
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 = <D extends Definition>(definition: D, projector: Subscriber<D>): Effect.Effect<void> =>
+18 -27
View File
@@ -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<string>(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),
),
),
),
)
}),
),
+4 -9
View File
@@ -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<ReadonlySet<SessionSchema.ID>>
/** Observes foreground ownership with an authoritative initial snapshot. */
readonly activity: (
sessionID: SessionSchema.ID,
) => Effect.Effect<
{ readonly snapshot: Effect.Effect<boolean>; readonly changes: Stream.Stream<boolean> },
never,
Scope.Scope
>
readonly activity: (sessionID: SessionSchema.ID) => Effect.Effect<SessionRunCoordinator.Activity, never, Scope.Scope>
/** Starts execution while idle or joins the active execution. */
readonly resume: (sessionID: SessionSchema.ID) => Effect.Effect<void, SessionRunner.RunError>
/** 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,
+11 -9
View File
@@ -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<boolean>
readonly changes: Stream.Stream<boolean>
/** Installs a synchronous observer and snapshots current ownership atomically. */
readonly attach: (observer: (active: boolean) => void) => Effect.Effect<boolean>
}
/** Serializes execution for each key while allowing different keys to run concurrently. */
@@ -119,9 +118,9 @@ export const make = <Key, E>(options: {
const activity = (key: Key) =>
Effect.gen(function* () {
const queue = yield* Queue.dropping<boolean, Cause.Done>(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 = <Key, E>(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)
}),
}
})
+51
View File
@@ -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
+78
View File
@@ -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<void>()
const release = yield* Deferred.make<void>()
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
+6
View File
@@ -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,