Skip to content

Commit b623dc2

Browse files
committed
perf(client): batch large thread sync updates
1 parent c30a6d9 commit b623dc2

3 files changed

Lines changed: 211 additions & 71 deletions

File tree

‎packages/client-runtime/src/state/threads-sync.test.ts‎

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -133,10 +133,12 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
133133
readonly cached?: OrchestrationThread;
134134
readonly httpSnapshot?: Option.Option<OrchestrationThreadDetailSnapshot>;
135135
readonly completionMarker?: boolean;
136+
readonly eventBatchSize?: number;
136137
}) {
137138
const inputs = yield* Queue.unbounded<TestThreadInput>();
138139
const observed = yield* Queue.unbounded<EnvironmentThreadState>();
139140
const latest = yield* Ref.make<EnvironmentThreadState>(EMPTY_ENVIRONMENT_THREAD_STATE);
141+
const statePublicationCount = yield* Ref.make(0);
140142
const retryCount = yield* Ref.make(0);
141143
const subscriptionCount = yield* Ref.make(0);
142144
const loaderCalls = yield* Ref.make(0);
@@ -221,7 +223,9 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
221223
clearVcsRefs: () => Effect.void,
222224
clear: () => Effect.void,
223225
});
224-
const threadState = yield* makeEnvironmentThreadState(THREAD_ID).pipe(
226+
const threadState = yield* makeEnvironmentThreadState(THREAD_ID, {
227+
eventBatchSize: options?.eventBatchSize ?? 1,
228+
}).pipe(
225229
Effect.provideService(EnvironmentSupervisor.EnvironmentSupervisor, supervisor),
226230
Effect.provideService(Persistence.EnvironmentCacheStore, cache),
227231
Effect.provideService(ThreadSnapshotLoader, snapshotLoader),
@@ -232,7 +236,10 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
232236
);
233237
yield* SubscriptionRef.changes(threadState).pipe(
234238
Stream.runForEach((state) =>
235-
Ref.set(latest, state).pipe(Effect.andThen(Queue.offer(observed, state))),
239+
Ref.update(statePublicationCount, (count) => count + 1).pipe(
240+
Effect.andThen(Ref.set(latest, state)),
241+
Effect.andThen(Queue.offer(observed, state)),
242+
),
236243
),
237244
Effect.forkScoped,
238245
);
@@ -241,6 +248,7 @@ const makeHarness = Effect.fn("TestEnvironmentThreads.makeHarness")(function* (o
241248
inputs,
242249
observed,
243250
latest,
251+
statePublicationCount,
244252
retryCount,
245253
subscriptionCount,
246254
loaderCalls,
@@ -347,6 +355,48 @@ describe("EnvironmentThreads", () => {
347355
}),
348356
);
349357

358+
it.effect("applies a live burst in order with one state publication", () =>
359+
Effect.gen(function* () {
360+
const harness = yield* makeHarness({
361+
cached: BASE_THREAD,
362+
completionMarker: true,
363+
eventBatchSize: 64,
364+
});
365+
yield* awaitThreadState(
366+
harness.observed,
367+
(value) => value.status === "synchronizing" && Option.isSome(value.data),
368+
);
369+
const publicationsBeforeBurst = yield* Ref.get(harness.statePublicationCount);
370+
371+
const finalSequence = CACHED_SNAPSHOT_SEQUENCE + 63;
372+
for (let sequence = CACHED_SNAPSHOT_SEQUENCE + 1; sequence <= finalSequence; sequence += 1) {
373+
yield* Queue.offer(
374+
harness.inputs,
375+
titleUpdated(
376+
sequence === finalSequence
377+
? "Final title"
378+
: sequence === CACHED_SNAPSHOT_SEQUENCE + 1
379+
? "First title"
380+
: "Interim title",
381+
sequence,
382+
),
383+
);
384+
}
385+
yield* Queue.offer(harness.inputs, synchronized());
386+
387+
const state = yield* awaitThreadState(
388+
harness.observed,
389+
(value) =>
390+
value.status === "live" &&
391+
Option.isSome(value.data) &&
392+
value.data.value.title === "Final title",
393+
);
394+
395+
expect(Option.getOrThrow(state.data).title).toBe("Final title");
396+
expect(yield* Ref.get(harness.statePublicationCount)).toBe(publicationsBeforeBurst + 1);
397+
}),
398+
);
399+
350400
it.effect("reduces live events and persists the latest thread", () =>
351401
Effect.gen(function* () {
352402
const harness = yield* makeHarness({ cached: BASE_THREAD });

‎packages/client-runtime/src/state/threads.ts‎

Lines changed: 146 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import {
77
type ThreadId as ThreadIdType,
88
} from "@t3tools/contracts";
99
import * as Cause from "effect/Cause";
10+
import * as Duration from "effect/Duration";
1011
import * as Effect from "effect/Effect";
1112
import * as Option from "effect/Option";
1213
import * as Queue from "effect/Queue";
@@ -32,6 +33,111 @@ import {
3233
type EnvironmentThreadStatus,
3334
} from "./threadState.ts";
3435

36+
// Coalesce one render-sized burst so web and mobile derive their large thread
37+
// views once instead of once per streamed event. The size cap bounds catch-up
38+
// memory and prevents a quiet connection from waiting indefinitely.
39+
const THREAD_EVENT_BATCH_WINDOW = Duration.millis(16);
40+
const THREAD_EVENT_BATCH_MAX_SIZE = 64;
41+
42+
interface ThreadStreamBatchReduction {
43+
readonly state: EnvironmentThreadState;
44+
readonly lastSequence: number;
45+
readonly awaitingCompletion: boolean;
46+
readonly threadChanged: boolean;
47+
readonly threadDeleted: boolean;
48+
}
49+
50+
export interface EnvironmentThreadStateOptions {
51+
readonly eventBatchSize?: number;
52+
}
53+
54+
function reduceThreadStreamItems(
55+
currentState: EnvironmentThreadState,
56+
currentSequence: number,
57+
currentAwaitingCompletion: boolean,
58+
items: ReadonlyArray<OrchestrationThreadStreamItem>,
59+
): ThreadStreamBatchReduction {
60+
let state = currentState;
61+
let lastSequence = currentSequence;
62+
let awaitingCompletion = currentAwaitingCompletion;
63+
let thread = Option.getOrNull(currentState.data);
64+
let threadChanged = false;
65+
let threadDeleted = false;
66+
67+
for (const item of items) {
68+
if (item.kind === "synchronized") {
69+
awaitingCompletion = false;
70+
if (thread !== null && state.status !== "deleted") {
71+
state = {
72+
data: state.data,
73+
status: "live",
74+
error: Option.none(),
75+
};
76+
}
77+
continue;
78+
}
79+
80+
if (item.kind === "snapshot") {
81+
lastSequence = item.snapshot.snapshotSequence;
82+
thread = item.snapshot.thread;
83+
threadChanged = true;
84+
threadDeleted = false;
85+
state = {
86+
data: Option.some(thread),
87+
status: awaitingCompletion ? "synchronizing" : "live",
88+
error: Option.none(),
89+
};
90+
continue;
91+
}
92+
93+
if (item.event.sequence <= lastSequence) {
94+
continue;
95+
}
96+
lastSequence = item.event.sequence;
97+
98+
if (thread === null) {
99+
if (item.event.type === "thread.deleted") {
100+
awaitingCompletion = false;
101+
threadDeleted = true;
102+
state = {
103+
data: Option.none(),
104+
status: "deleted",
105+
error: Option.none(),
106+
};
107+
}
108+
continue;
109+
}
110+
111+
const result = applyThreadDetailEvent(thread, item.event);
112+
if (result.kind === "updated") {
113+
thread = result.thread;
114+
threadChanged = true;
115+
state = {
116+
data: Option.some(thread),
117+
status: awaitingCompletion ? "synchronizing" : "live",
118+
error: Option.none(),
119+
};
120+
} else if (result.kind === "deleted") {
121+
awaitingCompletion = false;
122+
thread = null;
123+
threadDeleted = true;
124+
state = {
125+
data: Option.none(),
126+
status: "deleted",
127+
error: Option.none(),
128+
};
129+
}
130+
}
131+
132+
return {
133+
state,
134+
lastSequence,
135+
awaitingCompletion,
136+
threadChanged,
137+
threadDeleted,
138+
};
139+
}
140+
35141
function statusWithoutLiveData(data: Option.Option<OrchestrationThread>): EnvironmentThreadStatus {
36142
return Option.isSome(data) ? "cached" : "empty";
37143
}
@@ -50,12 +156,14 @@ function shouldPersistThread(thread: OrchestrationThread): boolean {
50156

51157
export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make")(function* (
52158
threadId: ThreadIdType,
159+
options?: EnvironmentThreadStateOptions,
53160
) {
54161
const supervisor = yield* EnvironmentSupervisor;
55162
const cache = yield* EnvironmentCacheStore;
56163
const snapshotLoader = yield* ThreadSnapshotLoader;
57164
const wakeups = yield* Effect.serviceOption(ConnectionWakeups.ConnectionWakeups);
58165
const environmentId = supervisor.target.environmentId;
166+
const eventBatchSize = options?.eventBatchSize ?? THREAD_EVENT_BATCH_MAX_SIZE;
59167
const cached = yield* cache.loadThread(environmentId, threadId).pipe(
60168
Effect.catch((error) =>
61169
Effect.logWarning("Could not load cached thread.").pipe(
@@ -141,81 +249,47 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make
141249
),
142250
);
143251

144-
const setThread = Effect.fn("EnvironmentThreadState.setThread")(function* (
145-
thread: OrchestrationThread,
252+
const applyItems = Effect.fn("EnvironmentThreadState.applyItems")(function* (
253+
items: ReadonlyArray<OrchestrationThreadStreamItem>,
146254
) {
147-
const waiting = yield* Ref.get(awaitingCompletion);
148-
yield* SubscriptionRef.set(state, {
149-
data: Option.some(thread),
150-
status: waiting ? "synchronizing" : "live",
151-
error: Option.none(),
152-
});
153-
// Active threads can update many times per second and retain large tool
154-
// payloads. The server remains the source of truth while a turn is active;
155-
// persist once it settles so cache encoding stays off the streaming path.
156-
if (shouldPersistThread(thread)) {
157-
const snapshotSequence = yield* SubscriptionRef.get(lastSequence);
158-
yield* Queue.offer(persistence, { snapshotSequence, thread });
159-
}
160-
});
161-
162-
const setDeleted = Effect.fn("EnvironmentThreadState.setDeleted")(function* () {
163-
yield* Ref.set(awaitingCompletion, false);
164-
yield* SubscriptionRef.set(state, {
165-
data: Option.none(),
166-
status: "deleted",
167-
error: Option.none(),
168-
});
169-
yield* cache.removeThread(environmentId, threadId).pipe(
170-
Effect.catch((error) =>
171-
Effect.logWarning("Could not remove the cached thread.").pipe(
172-
Effect.annotateLogs({
173-
environmentId,
174-
threadId,
175-
error: error.message,
176-
}),
177-
),
178-
),
255+
const currentState = yield* SubscriptionRef.get(state);
256+
const reduction = reduceThreadStreamItems(
257+
currentState,
258+
yield* SubscriptionRef.get(lastSequence),
259+
yield* Ref.get(awaitingCompletion),
260+
items,
179261
);
180-
});
181262

182-
const applyItem = Effect.fn("EnvironmentThreadState.applyItem")(function* (
183-
item: OrchestrationThreadStreamItem,
184-
) {
185-
if (item.kind === "synchronized") {
186-
yield* Ref.set(awaitingCompletion, false);
187-
yield* SubscriptionRef.update(state, (current) =>
188-
Option.isSome(current.data) && current.status !== "deleted"
189-
? { ...current, status: "live" as const, error: Option.none() }
190-
: current,
191-
);
192-
return;
263+
yield* SubscriptionRef.set(lastSequence, reduction.lastSequence);
264+
yield* Ref.set(awaitingCompletion, reduction.awaitingCompletion);
265+
if (reduction.state !== currentState) {
266+
yield* SubscriptionRef.set(state, reduction.state);
193267
}
194268

195-
if (item.kind === "snapshot") {
196-
yield* SubscriptionRef.set(lastSequence, item.snapshot.snapshotSequence);
197-
yield* setThread(item.snapshot.thread);
198-
return;
199-
}
200-
201-
const sequence = yield* SubscriptionRef.get(lastSequence);
202-
if (item.event.sequence <= sequence) {
269+
if (reduction.threadDeleted) {
270+
yield* cache.removeThread(environmentId, threadId).pipe(
271+
Effect.catch((error) =>
272+
Effect.logWarning("Could not remove the cached thread.").pipe(
273+
Effect.annotateLogs({
274+
environmentId,
275+
threadId,
276+
error: error.message,
277+
}),
278+
),
279+
),
280+
);
203281
return;
204282
}
205-
yield* SubscriptionRef.set(lastSequence, item.event.sequence);
206283

207-
const current = yield* SubscriptionRef.get(state);
208-
if (Option.isNone(current.data)) {
209-
if (item.event.type === "thread.deleted") {
210-
yield* setDeleted();
211-
}
212-
return;
213-
}
214-
const result = applyThreadDetailEvent(current.data.value, item.event);
215-
if (result.kind === "updated") {
216-
yield* setThread(result.thread);
217-
} else if (result.kind === "deleted") {
218-
yield* setDeleted();
284+
const thread = Option.getOrNull(reduction.state.data);
285+
// Active threads can update many times per second and retain large tool
286+
// payloads. The server remains the source of truth while a turn is active;
287+
// persist once it settles so cache encoding stays off the streaming path.
288+
if (reduction.threadChanged && thread !== null && shouldPersistThread(thread)) {
289+
yield* Queue.offer(persistence, {
290+
snapshotSequence: reduction.lastSequence,
291+
thread,
292+
});
219293
}
220294
});
221295

@@ -269,7 +343,7 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make
269343
);
270344
const httpSnapshot = yield* snapshotLoader.load(prepared, threadId);
271345
if (Option.isSome(httpSnapshot)) {
272-
yield* applyItem({ kind: "snapshot", snapshot: httpSnapshot.value });
346+
yield* applyItems([{ kind: "snapshot", snapshot: httpSnapshot.value }]);
273347
current = yield* SubscriptionRef.get(state);
274348
}
275349
}
@@ -295,7 +369,10 @@ export const makeEnvironmentThreadState = Effect.fn("EnvironmentThreadState.make
295369
retryExpectedFailureAfter: "250 millis",
296370
resubscribe: foregroundResubscriptions,
297371
},
298-
).pipe(Stream.runForEach(applyItem)),
372+
).pipe(
373+
Stream.groupedWithin(eventBatchSize, THREAD_EVENT_BATCH_WINDOW),
374+
Stream.runForEach(applyItems),
375+
),
299376
);
300377

301378
yield* Effect.addFinalizer(() =>

‎tasks/todo.md‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Large-thread sync optimization
2+
3+
- [x] Add bounded batching for live thread stream items.
4+
- [x] Add regression coverage for ordered, single-publication batches.
5+
- [x] Run focused client-runtime checks and review the diff.
6+
- [ ] Commit the focused sync fix on `perf/large-thread-sync`.
7+
8+
## Review/results
9+
10+
- Client-side batching is implemented in the shared web/mobile runtime.
11+
- Focused sync, reducer, and atom tests pass.
12+
- Targeted typecheck, lint, formatting, and diff checks pass; tsgo reports one unrelated existing suggestion in `src/relay/discovery.ts`.
13+
- Review found no remaining standards issue; server replay/snapshot and per-event reducer costs remain follow-up scope.

0 commit comments

Comments
 (0)