diff --git a/src/internal-agents/run-stream.test.ts b/src/internal-agents/run-stream.test.ts index fab6d10b62..cf6c089a03 100644 --- a/src/internal-agents/run-stream.test.ts +++ b/src/internal-agents/run-stream.test.ts @@ -21,8 +21,22 @@ import type { import { registerSkill } from "#veryfront/skill/registry.ts"; import type { RemoteToolSource, Tool } from "#veryfront/tool"; import { __resetLoggerConfigForTests, type LogEntry } from "#veryfront/utils/logger/logger.ts"; +import type { AgentRunEventSink } from "#veryfront/runtime/model-call-context.ts"; +import { getActiveRunEventSinks } from "#veryfront/runtime/run-event-sink-context.ts"; import { AgentRunSessionManager } from "./session-manager.ts"; -import { buildMergedTools, createRuntimeAgentStreamResponse } from "./run-stream.ts"; +import { + buildMergedTools, + createRuntimeAgentStreamResponse, + MODEL_CALL_CONTEXT_SSE_EVENT_NAME, +} from "./run-stream.ts"; + +function parseSseFrames(body: string): Array<{ event: string; data: unknown }> { + return body.split("\n\n").flatMap((frame) => { + const event = /^event: (.+)$/m.exec(frame)?.[1]; + const data = /^data: (.+)$/m.exec(frame)?.[1]; + return event && data ? [{ event, data: JSON.parse(data) as unknown }] : []; + }); +} class RecordingSpan implements Span { readonly attributes: Record = {}; @@ -2561,4 +2575,118 @@ describe("internal-agents/run-stream", () => { ); assertEquals(debugEntry?.component, "internal-agent-run-stream"); }); + describe("model call context", () => { + const modelCallContextEvent = { + type: "AGENT_RUN_MODEL_CALL_CONTEXT", + messages: [{ role: "system", content: "test system prompt" }], + tools: [{ type: "function", name: "granted_tool", inputSchema: {} }], + }; + + function contextAgent(): Agent { + return { + id: "context-agent", + config: { id: "context-agent", model: "anthropic/claude-opus-4-6", system: "test" }, + } as unknown as Agent; + } + + function contextRunInput(runId: string) { + return { + agentId: "context-agent", + threadId: crypto.randomUUID(), + runId, + messages: [], + tools: [], + context: [], + } as Parameters[0]; + } + + it("streams a context emitted while the runtime stream is created", async () => { + let sinkDuringCreate: AgentRunEventSink | undefined; + + const response = await createRuntimeAgentStreamResponse( + contextRunInput("run_context_setup"), + contextAgent(), + { + sessionManager: new AgentRunSessionManager(), + createRuntime: () => ({ + stream: async () => { + // The real runtime dispatches its first model call here, so the + // sink has to already be scoped by the time stream() runs. + sinkDuringCreate = getActiveRunEventSinks().mandatory; + await sinkDuringCreate?.(modelCallContextEvent as never); + return new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + }, + }), + }, + ); + + const frames = parseSseFrames(await response.text()); + assertEquals(Boolean(sinkDuringCreate), true); + assertEquals( + frames + .filter((frame) => frame.event === MODEL_CALL_CONTEXT_SSE_EVENT_NAME) + .map((frame) => frame.data), + [modelCallContextEvent], + ); + }); + + it("keeps the context ahead of the step it describes", async () => { + const response = await createRuntimeAgentStreamResponse( + contextRunInput("run_context_order"), + contextAgent(), + { + sessionManager: new AgentRunSessionManager(), + createRuntime: () => ({ + stream: async () => { + await getActiveRunEventSinks().mandatory?.(modelCallContextEvent as never); + return new ReadableStream({ + start(controller) { + controller.close(); + }, + }); + }, + }), + }, + ); + + const names = parseSseFrames(await response.text()).map((frame) => frame.event); + assertEquals(names[0], "RunStarted"); + assertEquals(names[1], MODEL_CALL_CONTEXT_SSE_EVENT_NAME); + }); + + it("streams a context emitted for a later step while the client reads", async () => { + let sinkDuringConsume: AgentRunEventSink | undefined; + + const response = await createRuntimeAgentStreamResponse( + contextRunInput("run_context_step_two"), + contextAgent(), + { + sessionManager: new AgentRunSessionManager(), + createRuntime: () => ({ + stream: async () => + new ReadableStream({ + // Multi-step runs dispatch later model calls as the stream is + // pulled, long after stream() returned. + async pull(controller) { + sinkDuringConsume = getActiveRunEventSinks().mandatory; + await sinkDuringConsume?.(modelCallContextEvent as never); + controller.close(); + }, + }), + }), + }, + ); + + const frames = parseSseFrames(await response.text()); + assertEquals(Boolean(sinkDuringConsume), true); + assertEquals( + frames.filter((frame) => frame.event === MODEL_CALL_CONTEXT_SSE_EVENT_NAME).length, + 1, + ); + }); + }); }); diff --git a/src/internal-agents/run-stream.ts b/src/internal-agents/run-stream.ts index 9cbb2427a1..6376069e76 100644 --- a/src/internal-agents/run-stream.ts +++ b/src/internal-agents/run-stream.ts @@ -51,6 +51,8 @@ import { mapRuntimeEventToAgUi, parseSseJsonEvents, } from "./ag-ui-sse.ts"; +import type { AgentRunEvent, AgentRunEventSink } from "#veryfront/runtime/model-call-context.ts"; +import { runWithMandatoryRunEventSink } from "#veryfront/runtime/run-event-sink-context.ts"; import { AgentRunCancelledError, type AgentRunSessionManager } from "./session-manager.ts"; import { composeInternalAgentRunSystemPrompt } from "./run-system-prompt.ts"; import type { RuntimeRunAgentInput } from "./schema.ts"; @@ -64,6 +66,12 @@ const INTERNAL_AGENT_RUNTIME_HEARTBEAT_INTERVAL_MS = 25_000; const INTERNAL_AGENT_RUNTIME_HEARTBEAT_FRAME = new TextEncoder().encode( ": internal-agent-runtime-heartbeat\n\n", ); +/** + * SSE frame name carrying AGENT_RUN_MODEL_CALL_CONTEXT to veryfront-api. Not an + * AG-UI event: veryfront-api persists it under its own event type rather than + * folding it into the run's public event sequence. + */ +export const MODEL_CALL_CONTEXT_SSE_EVENT_NAME = "AgentRunModelCallContext"; type RuntimeFilteredAgent = Agent & { config: Agent["config"] & { @@ -638,6 +646,31 @@ function compactRuntimeMessagesForStream( ) as Message[]; } +/** + * Relays run events produced by the runtime into the run's SSE stream. + * + * The first model call is dispatched while `runtime.stream()` is still being + * awaited, before there is a controller to enqueue into, so events raised + * before `attach` are buffered and replayed once the stream opens. + */ +function createModelCallContextRelay(): { + sink: AgentRunEventSink; + attach: (emit: (event: AgentRunEvent) => void) => void; +} { + const buffered: AgentRunEvent[] = []; + let emit: ((event: AgentRunEvent) => void) | undefined; + return { + sink: (event) => { + if (emit) emit(event); + else buffered.push(event); + }, + attach: (next) => { + emit = next; + for (const event of buffered.splice(0)) next(event); + }, + }; +} + export async function createRuntimeAgentStreamResponse( input: RuntimeRunAgentInput, agent: Agent, @@ -659,6 +692,7 @@ export async function createRuntimeAgentStreamResponse( let completedResponse: AgentResponse | null = null; let runtimeStream: ReadableStream; let closeSandbox = createIdempotentAsyncCleanup(); + const modelCallContextRelay = createModelCallContextRelay(); try { const forwardedAllowedRemoteToolNames = getAllowedRemoteToolNames(input.forwardedProps); const sourceAllowedRemoteToolNames = getAgentAllowedRemoteToolNames(agent); @@ -788,27 +822,34 @@ export async function createRuntimeAgentStreamResponse( runtimeToolNames.length, ); const maxOutputTokens = getForwardedMaxOutputTokens(input.forwardedProps); - const candidateRuntimeStream = await runtime.stream( - runtimeMessages, - { - threadId: input.threadId, - runId: input.runId, - ...(deps.projectAgentSandbox?.authToken - ? { authToken: deps.projectAgentSandbox.authToken } - : {}), - ...(input.parentRunId ? { parentRunId: input.parentRunId } : {}), - ...(input.state !== undefined ? { state: input.state } : {}), - context: input.context, - forwardedProps: input.forwardedProps, - }, - { - onFinish: (response) => { - completedResponse = response; - }, - }, - undefined, - maxOutputTokens, - abortSignal, + // Scoped here because the runtime dispatches the run's first model call + // before stream() resolves. Later steps inherit this scope through the + // stream they are pumped from. + const candidateRuntimeStream = await runWithMandatoryRunEventSink( + modelCallContextRelay.sink, + () => + runtime.stream( + runtimeMessages, + { + threadId: input.threadId, + runId: input.runId, + ...(deps.projectAgentSandbox?.authToken + ? { authToken: deps.projectAgentSandbox.authToken } + : {}), + ...(input.parentRunId ? { parentRunId: input.parentRunId } : {}), + ...(input.state !== undefined ? { state: input.state } : {}), + context: input.context, + forwardedProps: input.forwardedProps, + }, + { + onFinish: (response) => { + completedResponse = response; + }, + }, + undefined, + maxOutputTokens, + abortSignal, + ), ); if (candidateRuntimeStream.locked) { throw new TypeError("Internal agent runtime returned a locked stream"); @@ -960,6 +1001,14 @@ export async function createRuntimeAgentStreamResponse( threadId: input.threadId, agentId: agent.id, }); + // Replays whatever the first model call already produced, then + // forwards later steps as they happen. RunStarted stays first. + modelCallContextRelay.attach((event) => + enqueueIfAttached( + MODEL_CALL_CONTEXT_SSE_EVENT_NAME, + event as unknown as Record, + ) + ); heartbeatTimer = setInterval( enqueueHeartbeatIfAttached, INTERNAL_AGENT_RUNTIME_HEARTBEAT_INTERVAL_MS, @@ -968,7 +1017,13 @@ export async function createRuntimeAgentStreamResponse( while (true) { throwIfAborted(); - const { done, value } = await reader.read(); + // A runtime that dispatches later model calls from its pull() + // rather than from a continuation of stream() needs the sink in + // scope on the read that triggers the pull. + const { done, value } = await runWithMandatoryRunEventSink( + modelCallContextRelay.sink, + () => reader.read(), + ); throwIfAborted(); if (done) {