This commit is contained in:
Harivansh Rathi 2026-03-09 10:46:43 -07:00
parent 43c6b56dfa
commit 6b2a639fb6
3 changed files with 282 additions and 45 deletions

View file

@ -913,6 +913,14 @@ export class AgentSession {
await this._memoryDisposePromise;
}
private async _awaitAgentEventProcessing(): Promise<void> {
try {
await this._agentEventQueue;
} catch {
// Agent event failures are surfaced through normal listener paths.
}
}
private _enqueueMemoryPromotion(messages: AgentMessage[]): void {
this._memoryWriteQueue = this._memoryWriteQueue
.catch(() => undefined)
@ -1161,6 +1169,7 @@ export class AgentSession {
await this.agent.prompt(messages);
await this.waitForRetry();
await this._awaitAgentEventProcessing();
}
/**
@ -1369,6 +1378,8 @@ export class AgentSession {
}
} else if (options?.triggerTurn) {
await this.agent.prompt(appMessage);
await this.waitForRetry();
await this._awaitAgentEventProcessing();
} else {
this.agent.appendMessage(appMessage);
this.sessionManager.appendCustomMessageEntry(

View file

@ -2,6 +2,43 @@ import { randomUUID } from "node:crypto";
import type { ServerResponse } from "node:http";
import type { AgentSessionEvent } from "../agent-session.js";
type TextStreamState = {
started: boolean;
ended: boolean;
streamedText: string;
};
function isTextContent(
value: unknown,
): value is { type: "text"; text: string } {
return (
typeof value === "object" &&
value !== null &&
"type" in value &&
"text" in value &&
(value as { type: unknown }).type === "text" &&
typeof (value as { text: unknown }).text === "string"
);
}
function getAssistantTextParts(
event: Extract<AgentSessionEvent, { type: "message_end" }>,
): Array<{ contentIndex: number; text: string }> {
if (event.message.role !== "assistant" || !Array.isArray(event.message.content)) {
return [];
}
const textParts: Array<{ contentIndex: number; text: string }> = [];
for (const [contentIndex, content] of event.message.content.entries()) {
if (!isTextContent(content) || content.text.length === 0) {
continue;
}
textParts.push({ contentIndex, text: content.text });
}
return textParts;
}
/**
* Write a single Vercel AI SDK v5+ SSE chunk to the response.
* Format: `data: <JSON>\n\n`
@ -63,6 +100,76 @@ export function createVercelStreamListener(
// so these guards only need to bound the stream to that prompt's event span.
let active = false;
const msgId = messageId ?? randomUUID();
let textStates = new Map<number, TextStreamState>();
const getTextState = (contentIndex: number): TextStreamState => {
const existing = textStates.get(contentIndex);
if (existing) {
return existing;
}
const initial: TextStreamState = {
started: false,
ended: false,
streamedText: "",
};
textStates.set(contentIndex, initial);
return initial;
};
const emitTextStart = (contentIndex: number): void => {
const state = getTextState(contentIndex);
if (state.started) {
return;
}
writeChunk(response, {
type: "text-start",
id: `text_${contentIndex}`,
});
state.started = true;
};
const emitTextDelta = (contentIndex: number, delta: string): void => {
if (delta.length === 0) {
return;
}
emitTextStart(contentIndex);
writeChunk(response, {
type: "text-delta",
id: `text_${contentIndex}`,
delta,
});
getTextState(contentIndex).streamedText += delta;
};
const emitTextEnd = (contentIndex: number): void => {
const state = getTextState(contentIndex);
if (state.ended) {
return;
}
emitTextStart(contentIndex);
writeChunk(response, {
type: "text-end",
id: `text_${contentIndex}`,
});
state.ended = true;
};
const flushAssistantMessageText = (
event: Extract<AgentSessionEvent, { type: "message_end" }>,
): void => {
for (const part of getAssistantTextParts(event)) {
const state = getTextState(part.contentIndex);
if (!part.text.startsWith(state.streamedText)) {
continue;
}
const suffix = part.text.slice(state.streamedText.length);
if (suffix.length > 0) {
emitTextDelta(part.contentIndex, suffix);
}
emitTextEnd(part.contentIndex);
}
};
return (event: AgentSessionEvent) => {
if (response.writableEnded) return;
@ -85,6 +192,7 @@ export function createVercelStreamListener(
switch (event.type) {
case "turn_start":
textStates = new Map();
writeChunk(response, { type: "start-step" });
return;
@ -92,23 +200,13 @@ export function createVercelStreamListener(
const inner = event.assistantMessageEvent;
switch (inner.type) {
case "text_start":
writeChunk(response, {
type: "text-start",
id: `text_${inner.contentIndex}`,
});
emitTextStart(inner.contentIndex);
return;
case "text_delta":
writeChunk(response, {
type: "text-delta",
id: `text_${inner.contentIndex}`,
delta: inner.delta,
});
emitTextDelta(inner.contentIndex, inner.delta);
return;
case "text_end":
writeChunk(response, {
type: "text-end",
id: `text_${inner.contentIndex}`,
});
emitTextEnd(inner.contentIndex);
return;
case "toolcall_start": {
const content = inner.partial.content[inner.contentIndex];
@ -163,6 +261,12 @@ export function createVercelStreamListener(
return;
}
case "message_end":
if (event.message.role === "assistant") {
flushAssistantMessageText(event);
}
return;
case "turn_end":
writeChunk(response, { type: "finish-step" });
return;