From 1f4e02d83dea14a1c7550361d9a0a1be56e2f87c Mon Sep 17 00:00:00 2001 From: Varun Nuthalapati Date: Tue, 14 Apr 2026 20:41:16 -0700 Subject: [PATCH] feat: add intermediate flag to Event for multi-step agent streaming Add an optional Event.intermediate boolean field that is set to true on all events emitted during agent steps that produce function calls (i.e. steps that are followed by tool execution and more LLM calls). This allows streaming consumers to distinguish between model text produced as part of a reasoning/tool-calling step versus the final response text returned to the user. Motivation: issue #261. Aligned with adk-python behaviour. --- core/src/agents/llm_agent.ts | 13 +- core/src/events/event.ts | 11 ++ core/test/agents/event_intermediate_test.ts | 159 ++++++++++++++++++++ 3 files changed, 181 insertions(+), 2 deletions(-) create mode 100644 core/test/agents/event_intermediate_test.ts diff --git a/core/src/agents/llm_agent.ts b/core/src/agents/llm_agent.ts index 30aba287..493e030c 100644 --- a/core/src/agents/llm_agent.ts +++ b/core/src/agents/llm_agent.ts @@ -667,6 +667,7 @@ export class LlmAgent extends BaseAgent { ): AsyncGenerator { while (true) { let lastEvent: Event | undefined = undefined; + const stepEvents: Event[] = []; for await (const event of this.runOneStepAsync(context)) { if (context.abortSignal?.aborted) { return; @@ -674,14 +675,22 @@ export class LlmAgent extends BaseAgent { lastEvent = event; this.maybeSaveOutputToState(event); + stepEvents.push(event); + } + + const isFinal = !lastEvent || isFinalResponse(lastEvent); + for (const event of stepEvents) { + if (!isFinal) { + event.intermediate = true; + } yield event; } - if (!lastEvent || isFinalResponse(lastEvent)) { + if (isFinal) { break; } - if (lastEvent.partial) { + if (lastEvent?.partial) { logger.warn('The last event is partial, which is not expected.'); break; } diff --git a/core/src/events/event.ts b/core/src/events/event.ts index a25c28be..feb54050 100644 --- a/core/src/events/event.ts +++ b/core/src/events/event.ts @@ -62,6 +62,17 @@ export interface Event extends LlmResponse { * The timestamp of the event. */ timestamp: number; + + /** + * Whether this event is from an intermediate agent step that precedes a + * function call. When `true`, the event was produced by the model during a + * reasoning step that was followed by tool invocations and is not the final + * response to the user. + * + * Streaming consumers can use this flag to suppress intermediate model text + * (e.g. "thinking" output) from the user-visible response. + */ + intermediate?: boolean; } /** diff --git a/core/test/agents/event_intermediate_test.ts b/core/test/agents/event_intermediate_test.ts new file mode 100644 index 00000000..e9f321f7 --- /dev/null +++ b/core/test/agents/event_intermediate_test.ts @@ -0,0 +1,159 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + BaseLlm, + BaseLlmConnection, + createEvent, + Event, + FunctionTool, + InMemorySessionService, + LlmAgent, + LlmRequest, + LlmResponse, + Runner, +} from '@google/adk'; +import {beforeEach, describe, expect, it} from 'vitest'; + +// --------------------------------------------------------------------------- +// Unit tests — Event.intermediate field +// --------------------------------------------------------------------------- + +describe('Event.intermediate field', () => { + it('is undefined by default in createEvent', () => { + const event = createEvent(); + expect(event.intermediate).toBeUndefined(); + }); + + it('is preserved when set via createEvent', () => { + const event = createEvent({intermediate: true}); + expect(event.intermediate).toBe(true); + }); + + it('can be set to false', () => { + const event = createEvent({intermediate: false}); + expect(event.intermediate).toBe(false); + }); +}); + +// --------------------------------------------------------------------------- +// Integration tests — LlmAgent marks intermediate events +// --------------------------------------------------------------------------- + +/** + * A mock LLM that returns responses from a queue in order. + */ +class QueuedMockLlm extends BaseLlm { + private readonly queue: LlmResponse[]; + + constructor(responses: LlmResponse[]) { + super({model: 'mock-llm'}); + this.queue = [...responses]; + } + + async *generateContentAsync( + _request: LlmRequest, + ): AsyncGenerator { + const response = this.queue.shift(); + if (response) { + yield response; + } + } + + async connect(_llmRequest: LlmRequest): Promise { + throw new Error('connect() not supported in QueuedMockLlm'); + } +} + +const APP_NAME = 'test-app'; +const USER_ID = 'test-user'; + +describe('LlmAgent.runAsyncImpl intermediate flag', () => { + let sessionService: InMemorySessionService; + + beforeEach(() => { + sessionService = new InMemorySessionService(); + }); + + async function runAgent(agent: LlmAgent): Promise { + const runner = new Runner({appName: APP_NAME, agent, sessionService}); + const session = await sessionService.createSession({ + appName: APP_NAME, + userId: USER_ID, + }); + + const events: Event[] = []; + for await (const event of runner.runAsync({ + userId: session.userId, + sessionId: session.id, + newMessage: {role: 'user', parts: [{text: 'go'}]}, + })) { + events.push(event); + } + return events; + } + + it('does not mark events as intermediate when agent responds in one step', async () => { + const agent = new LlmAgent({ + name: 'single-step-agent', + model: new QueuedMockLlm([ + {content: {role: 'model', parts: [{text: 'Hello!'}]}}, + ]), + }); + + const events = await runAgent(agent); + + const modelEvents = events.filter((e) => e.author === 'single-step-agent'); + expect(modelEvents.length).toBeGreaterThan(0); + for (const event of modelEvents) { + expect(event.intermediate).toBeUndefined(); + } + }); + + it('marks events as intermediate when the step contains a function call', async () => { + const noopTool = new FunctionTool({ + name: 'noop', + description: 'Does nothing.', + execute: async () => 'ok', + }); + + const agent = new LlmAgent({ + name: 'multi-step-agent', + model: new QueuedMockLlm([ + // Step 1: function call → not a final response → intermediate + { + content: { + role: 'model', + parts: [{functionCall: {name: 'noop', args: {}}}], + }, + }, + // Step 2: text → final response → not intermediate + {content: {role: 'model', parts: [{text: 'Done!'}]}}, + ]), + tools: [noopTool], + }); + + const events = await runAgent(agent); + + // Events from step 1 (function call + function response) are intermediate + const step1Events = events.filter( + (e) => + e.content?.parts?.some((p) => p.functionCall) || + e.content?.parts?.some((p) => p.functionResponse), + ); + expect(step1Events.length).toBeGreaterThan(0); + for (const event of step1Events) { + expect(event.intermediate).toBe(true); + } + + // The final text response is not intermediate + const finalEvent = events.find((e) => + e.content?.parts?.some((p) => p.text === 'Done!'), + ); + expect(finalEvent).toBeDefined(); + expect(finalEvent!.intermediate).toBeUndefined(); + }); +});