From 243fd0c1a1f0ee0befdeab6328f5a5ba069e5cae Mon Sep 17 00:00:00 2001 From: KevinAndrewDong <2970752254@qq.com> Date: Sun, 25 Oct 2026 10:43:44 +0800 Subject: [PATCH] fix(streaming): preserve array indexes in progressive tool arguments --- core/src/utils/streaming_utils.ts | 63 +++- .../progressive_array_arguments_test.ts | 293 ++++++++++++++++++ 2 files changed, 349 insertions(+), 7 deletions(-) create mode 100644 core/test/models/progressive_array_arguments_test.ts diff --git a/core/src/utils/streaming_utils.ts b/core/src/utils/streaming_utils.ts index 7fce7e14bc..9e82db4eaa 100644 --- a/core/src/utils/streaming_utils.ts +++ b/core/src/utils/streaming_utils.ts @@ -43,6 +43,7 @@ class ProgressiveStrategy implements StreamingStrategy { // For streaming function call arguments private currentFcName?: string; private currentFcArgs: Record = {}; + private readonly arrayCandidates = new WeakSet(); private currentFcId?: string; private currentThoughtSignature?: string | Uint8Array; @@ -66,6 +67,37 @@ class ProgressiveStrategy implements StreamingStrategy { this.currentTextIsThought = undefined; } + private getArgumentPath(jsonPath: string): Array { + // toPathArray strips quotes, so mark array indices before parsing to + // distinguish [0] from the object member ['0']. Leave quoted text alone. + let prefix = '_adkArrayIndex'; + while (jsonPath.includes(prefix)) prefix += '_'; + const indices = new Map(); + const markedPath = jsonPath.replace( + /"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'|\[\s*(0|[1-9]\d*)\s*\]/g, + (match, index: string | undefined) => { + if (index === undefined) return match; + const value = Number(index); + // Other numeric properties cannot be represented as JS array elements. + if (!Number.isSafeInteger(value) || value > 2 ** 32 - 2) return match; + const marker = `${prefix}${indices.size}`; + indices.set(marker, value); + return `['${marker}']`; + }, + ); + const parse = (path: string) => + JSONPath.toPathArray(path).filter( + (part) => part !== '$' && part !== '$[', + ); + const parts = parse(markedPath); + // Preserve existing parsing for expressions where an index was not a + // standalone path segment, rather than leaking an internal marker. + if (parts.filter((part) => indices.has(part)).length !== indices.size) { + return parse(jsonPath); + } + return parts.map((part) => indices.get(part) ?? part); + } + private getValueFromPartialArg( partialArg: PartialArg, jsonPath: string, @@ -82,9 +114,7 @@ class ProgressiveStrategy implements StreamingStrategy { const stringChunk = stringValue; hasValue = true; - const pathParts = JSONPath.toPathArray(jsonPath).filter( - (p) => p !== '$' && p !== '$[', - ); + const pathParts = this.getArgumentPath(jsonPath); let existingValue: unknown = this.currentFcArgs; for (const part of pathParts) { @@ -120,9 +150,7 @@ class ProgressiveStrategy implements StreamingStrategy { } private setValueByJsonPath(jsonPath: string, value: unknown): void { - const pathParts = JSONPath.toPathArray(jsonPath).filter( - (p) => p !== '$' && p !== '$[', - ); + const pathParts = this.getArgumentPath(jsonPath); // Reject model-controlled paths that target prototype-chain properties, // which would otherwise pollute `Object.prototype` (e.g. @@ -136,6 +164,9 @@ class ProgressiveStrategy implements StreamingStrategy { const part = pathParts[i]; if (!Object.prototype.hasOwnProperty.call(current, part)) { current[part] = {}; + if (typeof pathParts[i + 1] === 'number') { + this.arrayCandidates.add(current[part] as object); + } } current = current[part] as Record; } @@ -150,7 +181,25 @@ class ProgressiveStrategy implements StreamingStrategy { const fcPart: Part = { functionCall: { name: this.currentFcName, - args: JSON.parse(JSON.stringify(this.currentFcArgs)), + args: JSON.parse( + JSON.stringify(this.currentFcArgs, (_key, value: unknown) => { + if ( + value !== null && + typeof value === 'object' && + this.arrayCandidates.has(value) + ) { + const keys = Object.keys(value); + // Only complete index sequences become arrays. Accumulating + // as objects preserves sparse and named members like length. + if (keys.every((key, index) => key === String(index))) { + return keys.map( + (key) => (value as Record)[key], + ); + } + } + return value; + }), + ), id: this.currentFcId ?? generateClientFunctionCallId(), } as FunctionCall, }; diff --git a/core/test/models/progressive_array_arguments_test.ts b/core/test/models/progressive_array_arguments_test.ts new file mode 100644 index 0000000000..03b147fa2c --- /dev/null +++ b/core/test/models/progressive_array_arguments_test.ts @@ -0,0 +1,293 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + FinishReason, + FunctionCall, + FunctionDeclaration, + GenerateContentResponse, + PartialArg, + Type, +} from '@google/genai'; +import {describe, expect, it} from 'vitest'; +import {z} from 'zod/v4'; + +import {Context} from '../../src/agents/context.js'; +import {InvocationContext} from '../../src/agents/invocation_context.js'; +import {LlmAgent} from '../../src/agents/llm_agent.js'; +import { + FeatureName, + isFeatureEnabled, + withTemporaryFeatureOverride, +} from '../../src/features/feature_registry.js'; +import {Gemini} from '../../src/models/google_llm.js'; +import {appendTools, LlmRequest} from '../../src/models/llm_request.js'; +import {PluginManager} from '../../src/plugins/plugin_manager.js'; +import {createSession} from '../../src/sessions/session.js'; +import {BaseTool} from '../../src/tools/base_tool.js'; +import {FunctionTool} from '../../src/tools/function_tool.js'; + +function partialCall(partialArgs: PartialArg[], id = 'call'): FunctionCall { + return {id, name: 'collect', partialArgs, willContinue: false}; +} + +async function stream( + calls: FunctionCall[], + progressive = true, + tool?: BaseTool, +) { + const model = new Gemini({ + vertexai: true, + project: 'local-test-project', + location: 'us-central1', + }); + const chunks = calls.map((call) => { + const response = new GenerateContentResponse(); + response.candidates = [ + { + content: {role: 'model', parts: [{functionCall: call}]}, + finishReason: FinishReason.STOP, + }, + ]; + return response; + }); + let sdkCalls = 0; + let declarations: FunctionDeclaration[] | undefined; + Object.defineProperty(model, 'apiClient', { + value: { + vertexai: true, + models: { + async generateContentStream(params: { + config?: { + tools?: Array<{functionDeclarations?: FunctionDeclaration[]}>; + }; + }) { + sdkCalls++; + declarations = params.config?.tools?.[0].functionDeclarations; + return (async function* () { + yield* chunks; + })(); + }, + }, + }, + }); + const request: LlmRequest = { + contents: [{role: 'user', parts: [{text: 'collect'}]}], + toolsDict: {}, + liveConnectConfig: {}, + }; + if (tool) appendTools(request, [tool]); + const feature = FeatureName.PROGRESSIVE_SSE_STREAMING; + const before = isFeatureEnabled(feature); + const result: FunctionCall[] = []; + try { + await withTemporaryFeatureOverride(feature, progressive, async () => { + for await (const response of model.generateContentAsync(request, true)) { + if (!response.partial) { + for (const part of response.content?.parts ?? []) { + if (part.functionCall) result.push(part.functionCall); + } + } + } + }); + } finally { + expect(isFeatureEnabled(feature)).toBe(before); + } + expect(sdkCalls).toBe(1); + return {calls: result, model, declarations}; +} + +describe('Vertex progressive streamed array arguments', () => { + it('reconstructs declared arrays and executes a real FunctionTool once', async () => { + const expected = {items: [{name: 'alpha'}, {name: 'beta'}]}; + let executions = 0; + let received: unknown; + const tool = new FunctionTool({ + name: 'collect', + description: 'Collects names', + parameters: z.object({items: z.array(z.object({name: z.string()}))}), + execute: (args) => { + executions++; + received = args; + return {count: args.items.length}; + }, + }); + const {calls, model, declarations} = await stream( + [ + partialCall([ + {jsonPath: '$.items[0].name', stringValue: 'alpha'}, + {jsonPath: '$.items[1].name', stringValue: 'beta'}, + ]), + ], + true, + tool, + ); + expect(declarations?.[0].parameters?.properties?.['items'].type).toBe( + Type.ARRAY, + ); + expect(calls).toHaveLength(1); + expect(Array.isArray(calls[0].args?.['items'])).toBe(true); + expect(calls[0].args).toEqual(expected); + const context = new Context({ + invocationContext: new InvocationContext({ + invocationId: 'test', + agent: new LlmAgent({name: 'agent', model}), + session: createSession({id: 'session', appName: 'app', userId: 'user'}), + pluginManager: new PluginManager([]), + }), + }); + expect( + await tool.runAsync({args: calls[0].args ?? {}, toolContext: context}), + ).toEqual({count: 2}); + expect(received).toEqual(expected); + expect(executions).toBe(1); + }); + + it.each(["$.items['0'].name", '$.items["0"].name'])( + 'keeps the quoted numeric object member in %s', + async (jsonPath) => { + const {calls} = await stream([ + partialCall([{jsonPath, stringValue: 'alpha'}]), + ]); + expect(calls[0].args).toEqual({items: {'0': {name: 'alpha'}}}); + expect(Array.isArray(calls[0].args?.['items'])).toBe(false); + }, + ); + + it.each([true, false])( + 'preserves complete array args with progressive=%s', + async (progressive) => { + const args = {items: [{name: 'alpha'}]}; + const {calls} = await stream( + [{id: 'call', name: 'collect', args}], + progressive, + ); + expect(calls[0].args).toEqual(args); + }, + ); + + it('reconstructs nested arrays', async () => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.matrix[0][0]', numberValue: 1}, + {jsonPath: '$.matrix[0][1]', numberValue: 2}, + {jsonPath: '$.matrix[1][0]', numberValue: 3}, + ]), + ]); + expect(calls[0].args).toEqual({matrix: [[1, 2], [3]]}); + }); + + it('appends string chunks at the same array element', async () => { + const first = partialCall([{jsonPath: '$.items[0]', stringValue: 'al'}]); + first.willContinue = true; + const last: FunctionCall = { + id: 'call', + willContinue: false, + partialArgs: [{jsonPath: '$.items[0]', stringValue: 'pha'}], + }; + const {calls} = await stream([first, last]); + expect(calls).toHaveLength(1); + expect(calls[0].args).toEqual({items: ['alpha']}); + }); + + it('preserves false, zero, null and empty string in array elements', async () => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.items[0]', boolValue: false}, + {jsonPath: '$.items[1]', numberValue: 0}, + {jsonPath: '$.items[2]', nullValue: 'NULL_VALUE'}, + {jsonPath: '$.items[3]', stringValue: ''}, + ]), + ]); + expect(calls[0].args).toEqual({items: [false, 0, null, '']}); + }); + + it('places consecutive indexes correctly when they arrive in reverse order', async () => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.items[1]', stringValue: 'beta'}, + {jsonPath: '$.items[0]', stringValue: 'alpha'}, + ]), + ]); + expect(calls[0].args).toEqual({items: ['alpha', 'beta']}); + }); + + it('keeps a missing-index path in its previous object form', async () => { + const {calls} = await stream([ + partialCall([{jsonPath: '$.items[1]', stringValue: 'beta'}]), + ]); + expect(calls[0].args).toEqual({items: {'1': 'beta'}}); + }); + + it('preserves an explicitly supplied null slot in a complete array', async () => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.items[0]', nullValue: 'NULL_VALUE'}, + {jsonPath: '$.items[1]', stringValue: 'beta'}, + ]), + ]); + expect(calls[0].args).toEqual({items: [null, 'beta']}); + }); + + it('keeps mixed numeric and named members in their previous object form', async () => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.items[0]', stringValue: 'alpha'}, + {jsonPath: "$.items['label']", stringValue: 'description'}, + ]), + ]); + expect(calls[0].args).toEqual({ + items: {'0': 'alpha', label: 'description'}, + }); + }); + + it.each([0, 'named'])( + 'preserves a length member with value %s', + async (length) => { + const {calls} = await stream([ + partialCall([ + {jsonPath: '$.items[0]', stringValue: 'alpha'}, + typeof length === 'number' + ? {jsonPath: '$.items.length', numberValue: length} + : {jsonPath: '$.items.length', stringValue: length}, + ]), + ]); + expect(calls[0].args).toEqual({items: {'0': 'alpha', length}}); + }, + ); + + it('accepts whitespace around an array index', async () => { + const {calls} = await stream([ + partialCall([{jsonPath: '$.items[ 0 ]', stringValue: 'alpha'}]), + ]); + expect(calls[0].args).toEqual({items: ['alpha']}); + }); + + it('preserves a quoted member containing a period before an array', async () => { + const {calls} = await stream([ + partialCall([{jsonPath: "$['a.b'][0].name", stringValue: 'alpha'}]), + ]); + expect(calls[0].args).toEqual({'a.b': [{name: 'alpha'}]}); + }); + + it('does not confuse an ordinary member with an internal index marker', async () => { + const {calls} = await stream([ + partialCall([{jsonPath: '$._adkArrayIndex0[0]', stringValue: 'alpha'}]), + ]); + expect(calls[0].args).toEqual({_adkArrayIndex0: ['alpha']}); + }); + + it('keeps separately streamed calls and IDs independent', async () => { + const {calls} = await stream([ + partialCall([{jsonPath: '$.items[0]', stringValue: 'alpha'}], 'first'), + partialCall([{jsonPath: '$.items[0]', stringValue: 'beta'}], 'second'), + ]); + expect(calls.map((call) => ({id: call.id, args: call.args}))).toEqual([ + {id: 'first', args: {items: ['alpha']}}, + {id: 'second', args: {items: ['beta']}}, + ]); + }); +});