History

Spill oversized buffered assistant deltas without losing final text

- cap buffered assistant text memory and spill older chunks as delta events - finalize buffered messages by emitting retained remainder before completion when spilled - add regression coverage for oversized buffered deltas and tighten type-safe test assertions

Julius Marminge committed 8 months ago

Commit cc29cee1 parent 32a17b93 files · +137 −17Browse files

Unified · Split

3 files
  1. apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts +53 −0
  2. apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts +77 −12
  3. apps/server/src/orchestration/decider.projectScripts.test.ts +7 −5
apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts+53−0
@@ -402,6 +402,59 @@ describe("ProviderRuntimeIngestion", () => { · 401 unmodified lines
402402 expect(finalMessage?.streaming).toBe(false);
403403 });
404404
405+ it("spills oversized buffered deltas and still finalizes full assistant text", async () => {
406+ const harness = await createHarness();
407+ const now = new Date().toISOString();
408+ const oversizedText = "x".repeat(40_000);
409+
410+ harness.emit({
411+ type: "turn.started",
412+ eventId: asEventId("evt-turn-started-buffer-spill"),
413+ provider: "codex",
414+ sessionId: asSessionId("sess-1"),
415+ createdAt: now,
416+ turnId: asProviderTurnId("turn-buffer-spill"),
417+ });
418+ await waitForThread(
419+ harness.engine,
420+ (thread) =>
421+ thread.session?.status === "running" &&
422+ thread.session?.activeTurnId === "turn-buffer-spill",
423+ );
424+
425+ harness.emit({
426+ type: "message.delta",
427+ eventId: asEventId("evt-message-delta-buffer-spill"),
428+ provider: "codex",
429+ sessionId: asSessionId("sess-1"),
430+ createdAt: now,
431+ turnId: asProviderTurnId("turn-buffer-spill"),
432+ itemId: asItemId("item-buffer-spill"),
433+ delta: oversizedText,
434+ });
435+ harness.emit({
436+ type: "message.completed",
437+ eventId: asEventId("evt-message-completed-buffer-spill"),
438+ provider: "codex",
439+ sessionId: asSessionId("sess-1"),
440+ createdAt: now,
441+ turnId: asProviderTurnId("turn-buffer-spill"),
442+ itemId: asItemId("item-buffer-spill"),
443+ });
444+
445+ const thread = await waitForThread(
446+ harness.engine,
447+ (entry) =>
448+ entry.messages.some(
449+ (message) => message.id === "assistant:item-buffer-spill" && !message.streaming,
450+ ),
451+ );
452+ const message = thread.messages.find((entry) => entry.id === "assistant:item-buffer-spill");
453+ expect(message?.text.length).toBe(oversizedText.length);
454+ expect(message?.text).toBe(oversizedText);
455+ expect(message?.streaming).toBe(false);
456+ });
457+
405458 it("does not duplicate assistant completion when message.completed is followed by turn.completed", async () => {
406459 const harness = await createHarness();
407460 const now = new Date().toISOString();
apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts+77−12
@@ -28,6 +28,10 @@ const TURN_MESSAGE_IDS_BY_TURN_CACHE_CAPACITY = 10_000; · 27 unmodified lines
2828 const TURN_MESSAGE_IDS_BY_TURN_TTL = Duration.minutes(120);
2929 const BUFFERED_MESSAGE_TEXT_BY_MESSAGE_ID_CACHE_CAPACITY = 20_000;
3030 const BUFFERED_MESSAGE_TEXT_BY_MESSAGE_ID_TTL = Duration.minutes(120);
31+const BUFFERED_ASSISTANT_MESSAGE_SPILLED_CACHE_CAPACITY = 20_000;
32+const BUFFERED_ASSISTANT_MESSAGE_SPILLED_TTL = Duration.minutes(120);
33+const MAX_BUFFERED_ASSISTANT_CHARS = 24_000;
34+const BUFFERED_ASSISTANT_KEEP_CHARS = 8_000;
3135
3236 type TurnStartRequestedDomainEvent = Extract<
3337 OrchestrationEvent,
@@ -169,6 +173,11 @@ const make = Effect.gen(function* () { · 135 unmodified lines
169173 timeToLive: BUFFERED_MESSAGE_TEXT_BY_MESSAGE_ID_TTL,
170174 lookup: () => Effect.succeed(""),
171175 });
176+ const bufferedAssistantMessageSpilledByMessageId = yield* Cache.make<MessageId, true>({
177+ capacity: BUFFERED_ASSISTANT_MESSAGE_SPILLED_CACHE_CAPACITY,
178+ timeToLive: BUFFERED_ASSISTANT_MESSAGE_SPILLED_TTL,
179+ lookup: () => Effect.succeed(true),
180+ });
172181
173182 const rememberAssistantMessageId = (
174183 sessionId: ProviderSessionId,
@@ -226,14 +235,23 @@ const make = Effect.gen(function* () { · 51 unmodified lines
226235 const appendBufferedAssistantText = (messageId: MessageId, delta: string) =>
227236 Cache.getOption(bufferedAssistantTextByMessageId, messageId).pipe(
228237 Effect.flatMap((existingText) =>
229− Cache.set(
230− bufferedAssistantTextByMessageId,
231− messageId,
232− Option.match(existingText, {
238+ Effect.gen(function* () {
239+ const nextText = Option.match(existingText, {
233240 onNone: () => delta,
234241 onSome: (text) => `${text}${delta}`,
235− }),
236− ),
242+ });
243+ if (nextText.length <= MAX_BUFFERED_ASSISTANT_CHARS) {
244+ yield* Cache.set(bufferedAssistantTextByMessageId, messageId, nextText);
245+ return "";
246+ }
247+
248+ // Safety valve: spill oldest text as a persisted assistant delta to cap memory.
249+ const spillChars = Math.max(1, nextText.length - BUFFERED_ASSISTANT_KEEP_CHARS);
250+ const spillChunk = nextText.slice(0, spillChars);
251+ const retainedTail = nextText.slice(spillChars);
252+ yield* Cache.set(bufferedAssistantTextByMessageId, messageId, retainedTail);
253+ return spillChunk;
254+ }),
237255 ),
238256 );
239257
@@ -249,18 +267,51 @@ const make = Effect.gen(function* () { · 9 unmodified lines
249267 const clearBufferedAssistantText = (messageId: MessageId) =>
250268 Cache.invalidate(bufferedAssistantTextByMessageId, messageId);
251269
252− const clearAssistantMessageState = (messageId: MessageId) => clearBufferedAssistantText(messageId);
270+ const markBufferedAssistantMessageSpilled = (messageId: MessageId) =>
271+ Cache.set(bufferedAssistantMessageSpilledByMessageId, messageId, true);
272+
273+ const takeBufferedAssistantMessageSpilled = (messageId: MessageId) =>
274+ Cache.getOption(bufferedAssistantMessageSpilledByMessageId, messageId).pipe(
275+ Effect.flatMap((spilled) =>
276+ Cache.invalidate(bufferedAssistantMessageSpilledByMessageId, messageId).pipe(
277+ Effect.as(Option.isSome(spilled)),
278+ ),
279+ ),
280+ );
281+
282+ const clearBufferedAssistantMessageSpilled = (messageId: MessageId) =>
283+ Cache.invalidate(bufferedAssistantMessageSpilledByMessageId, messageId);
284+
285+ const clearAssistantMessageState = (messageId: MessageId) =>
286+ Effect.all(
287+ [clearBufferedAssistantText(messageId), clearBufferedAssistantMessageSpilled(messageId)],
288+ { concurrency: 1 },
289+ ).pipe(Effect.asVoid);
253290
254− const dispatchAssistantCompletion = (input: {
291+ const finalizeAssistantMessage = (input: {
255292 event: ProviderRuntimeEvent;
256293 threadId: ThreadId;
257294 messageId: MessageId;
258295 turnId?: TurnId;
259296 createdAt: string;
260297 commandTag: string;
298+ spillRemainderCommandTag: string;
261299 }) =>
262300 Effect.gen(function* () {
263301 const text = yield* takeBufferedAssistantText(input.messageId);
302+ const spilled = yield* takeBufferedAssistantMessageSpilled(input.messageId);
303+
304+ if (spilled && text.length > 0) {
305+ yield* orchestrationEngine.dispatch({
306+ type: "thread.message.assistant.delta",
307+ commandId: providerCommandId(input.event, input.spillRemainderCommandTag),
308+ threadId: input.threadId,
309+ messageId: input.messageId,
310+ delta: text,
311+ ...(input.turnId ? { turnId: input.turnId } : {}),
312+ createdAt: input.createdAt,
313+ });
314+ }
264315
265316 yield* orchestrationEngine.dispatch({
266317 type: "thread.message.assistant.complete",
@@ -268,7 +319,7 @@ const make = Effect.gen(function* () { · 1 unmodified line
268319 threadId: input.threadId,
269320 messageId: input.messageId,
270321 ...(input.turnId ? { turnId: input.turnId } : {}),
271− text,
322+ text: spilled ? "" : text,
272323 createdAt: input.createdAt,
273324 });
274325 yield* clearAssistantMessageState(input.messageId);
@@ -369,7 +420,19 @@ const make = Effect.gen(function* () { · 94 unmodified lines
369420
370421 const assistantDeliveryMode = yield* Ref.get(assistantDeliveryModeRef);
371422 if (assistantDeliveryMode === "buffered") {
372− yield* appendBufferedAssistantText(assistantMessageId, event.delta);
423+ const spillChunk = yield* appendBufferedAssistantText(assistantMessageId, event.delta);
424+ if (spillChunk.length > 0) {
425+ yield* markBufferedAssistantMessageSpilled(assistantMessageId);
426+ yield* orchestrationEngine.dispatch({
427+ type: "thread.message.assistant.delta",
428+ commandId: providerCommandId(event, "assistant-delta-buffer-spill"),
429+ threadId: thread.id,
430+ messageId: assistantMessageId,
431+ delta: spillChunk,
432+ ...(turnId ? { turnId } : {}),
433+ createdAt: now,
434+ });
435+ }
373436 } else {
374437 yield* orchestrationEngine.dispatch({
375438 type: "thread.message.assistant.delta",
@@ -390,13 +453,14 @@ const make = Effect.gen(function* () { · 14 unmodified lines
390453 yield* rememberAssistantMessageId(event.sessionId, turnId, assistantMessageId);
391454 }
392455
393− yield* dispatchAssistantCompletion({
456+ yield* finalizeAssistantMessage({
394457 event,
395458 threadId: thread.id,
396459 messageId: assistantMessageId,
397460 ...(turnId ? { turnId } : {}),
398461 createdAt: now,
399462 commandTag: "assistant-complete",
463+ spillRemainderCommandTag: "assistant-delta-buffer-remainder",
400464 });
401465
402466 if (turnId) {
@@ -412,13 +476,14 @@ const make = Effect.gen(function* () { · 9 unmodified lines
412476 assistantMessageIds,
413477 (assistantMessageId) =>
414478 Effect.gen(function* () {
415− yield* dispatchAssistantCompletion({
479+ yield* finalizeAssistantMessage({
416480 event,
417481 threadId: thread.id,
418482 messageId: assistantMessageId,
419483 turnId,
420484 createdAt: now,
421485 commandTag: "assistant-complete-finalize",
486+ spillRemainderCommandTag: "assistant-delta-buffer-remainder-finalize",
422487 });
423488 }),
424489 { concurrency: 1 },
apps/server/src/orchestration/decider.projectScripts.test.ts+7−5
@@ -162,10 +162,12 @@ describe("decider project scripts", () => { · 161 unmodified lines
162162 const events = Array.isArray(result) ? result : [result];
163163 expect(events).toHaveLength(2);
164164 expect(events[0]?.type).toBe("thread.message-sent");
165− expect(events[1]?.type).toBe("thread.turn-start-requested");
166− expect(events[1]?.causationEventId).toBe(events[0]?.eventId ?? null);
167− expect((events[1]?.payload as { assistantDeliveryMode?: string }).assistantDeliveryMode).toBe(
168− "buffered",
169− );
165+ const turnStartEvent = events[1];
166+ expect(turnStartEvent?.type).toBe("thread.turn-start-requested");
167+ expect(turnStartEvent?.causationEventId).toBe(events[0]?.eventId ?? null);
168+ if (turnStartEvent?.type !== "thread.turn-start-requested") {
169+ return;
170+ }
171+ expect(turnStartEvent.payload.assistantDeliveryMode).toBe("buffered");
170172 });
171173 });