diff --git a/core/src/agents/active_streaming_tool.ts b/core/src/agents/active_streaming_tool.ts index b073a7dd..6734d88f 100644 --- a/core/src/agents/active_streaming_tool.ts +++ b/core/src/agents/active_streaming_tool.ts @@ -4,13 +4,14 @@ * SPDX-License-Identifier: Apache-2.0 */ +import {Task} from '../utils/task.js'; import {LiveRequestQueue} from './live_request_queue.js'; /** * The parameters for creating an ActiveStreamingTool. */ export interface ActiveStreamingToolParams { - task?: Promise; + task?: Task; stream?: LiveRequestQueue; } @@ -20,10 +21,8 @@ export interface ActiveStreamingToolParams { export class ActiveStreamingTool { /** * The active task of this streaming tool. - * TODO: Replace 'Promise' with a proper Task type if available in this - * env. */ - task?: Promise; + task?: Task; /** * The active (input) streams of this streaming tool. diff --git a/core/src/common.ts b/core/src/common.ts index 409ef94f..f9a09650 100644 --- a/core/src/common.ts +++ b/core/src/common.ts @@ -253,6 +253,7 @@ export {LogLevel, getLogger, setLogLevel, setLogger} from './utils/logger.js'; export type {Logger} from './utils/logger.js'; export {isGemini2OrAbove, isGemini3xFlashLive} from './utils/model_name.js'; export {zodObjectToSchema} from './utils/simple_zod_to_json.js'; +export {Task} from './utils/task.js'; export {GoogleLLMVariant} from './utils/variant_utils.js'; export {version} from './version.js'; @@ -272,6 +273,11 @@ export {LoadSkillTool} from './tools/skill/load_skill_tool.js'; export {SearchSkillsTool} from './tools/skill/search_skills_tool.js'; export {SkillToolset} from './tools/skill/skill_toolset.js'; +export * from './artifacts/base_artifact_service.js'; +export * from './features/feature_registry.js'; +export * from './memory/base_memory_service.js'; +export * from './sessions/base_session_service.js'; +export * from './tools/base_tool.js'; export {OpenApiSpecParser} from './tools/openapi_tool/openapi_spec_parser/openapi_spec_parser.js'; export type { OperationEndpoint, @@ -286,9 +292,3 @@ export { RestApiTool, createRestApiTool, } from './tools/openapi_tool/rest_api_tool.js'; - -export * from './artifacts/base_artifact_service.js'; -export * from './features/feature_registry.js'; -export * from './memory/base_memory_service.js'; -export * from './sessions/base_session_service.js'; -export * from './tools/base_tool.js'; diff --git a/core/src/utils/task.ts b/core/src/utils/task.ts new file mode 100644 index 00000000..b07cf036 --- /dev/null +++ b/core/src/utils/task.ts @@ -0,0 +1,36 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Represents a runtime task wrapping a promise, allowing status check and cancellation. + */ +export class Task { + private isDone = false; + + constructor( + readonly promise: Promise, + private readonly cancelFn?: () => void, + ) { + const markDone = () => { + this.isDone = true; + }; + promise.then(markDone, markDone); + } + + /** + * Cancels the task execution. + */ + cancel(): void { + this.cancelFn?.(); + } + + /** + * Returns true if the task has completed (either resolved or rejected). + */ + done(): boolean { + return this.isDone; + } +} diff --git a/core/test/agents/active_streaming_tool_test.ts b/core/test/agents/active_streaming_tool_test.ts new file mode 100644 index 00000000..69c6f842 --- /dev/null +++ b/core/test/agents/active_streaming_tool_test.ts @@ -0,0 +1,23 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {ActiveStreamingTool, Task} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('ActiveStreamingTool', () => { + it('should construct with default parameters', () => { + const tool = new ActiveStreamingTool(); + expect(tool.task).toBeUndefined(); + expect(tool.stream).toBeUndefined(); + }); + + it('should store task when constructed with one', () => { + const promise = Promise.resolve(); + const task = new Task(promise); + const tool = new ActiveStreamingTool({task}); + expect(tool.task).toBe(task); + }); +}); diff --git a/core/test/utils/task_test.ts b/core/test/utils/task_test.ts new file mode 100644 index 00000000..3fee2117 --- /dev/null +++ b/core/test/utils/task_test.ts @@ -0,0 +1,61 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {Task} from '@google/adk'; +import {describe, expect, it, vi} from 'vitest'; + +describe('task utils', () => { + describe('Task class', () => { + it('should initialize with done() returning false', () => { + const promise = new Promise(() => {}); + const task = new Task(promise); + expect(task.done()).toBe(false); + }); + + it('should set done() to true when promise resolves', async () => { + const promise = Promise.resolve(); + const task = new Task(promise); + + expect(task.done()).toBe(false); + + await promise; + + expect(task.done()).toBe(true); + }); + + it('should set done() to true when promise rejects', async () => { + const promise = Promise.reject(new Error('test error')); + const task = new Task(promise); + + expect(task.done()).toBe(false); + + try { + await promise; + } catch (_) { + // expected + } + + expect(task.done()).toBe(true); + }); + + it('should call cancelFn when cancel is called', () => { + const cancelFn = vi.fn(); + const promise = new Promise(() => {}); + const task = new Task(promise, cancelFn); + + task.cancel(); + + expect(cancelFn).toHaveBeenCalledTimes(1); + }); + + it('should not throw if cancel is called but no cancelFn is provided', () => { + const promise = new Promise(() => {}); + const task = new Task(promise); + + expect(() => task.cancel()).not.toThrow(); + }); + }); +}); diff --git a/tests/e2e/streaming/task_e2e_test.ts b/tests/e2e/streaming/task_e2e_test.ts new file mode 100644 index 00000000..0bf41f6d --- /dev/null +++ b/tests/e2e/streaming/task_e2e_test.ts @@ -0,0 +1,52 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {ActiveStreamingTool, LiveRequestQueue, Task} from '@google/adk'; +import {describe, expect, it} from 'vitest'; + +describe('ActiveStreamingTool E2E Simulation', () => { + it('should manage a simulated streaming tool task', async () => { + const queue = new LiveRequestQueue(); + let cancelled = false; + + // Simulate a background task + const promise = new Promise((resolve) => { + const interval = setInterval(() => { + if (cancelled) { + clearInterval(interval); + resolve(); + return; + } + // Push some dummy data to the queue + queue.sendContent({parts: [{text: 'data'}]}); + }, 10); + }); + + const task = new Task(promise, () => { + cancelled = true; + }); + + const activeTool = new ActiveStreamingTool({task, stream: queue}); + + expect(activeTool.task).toBe(task); + expect(activeTool.stream).toBe(queue); + expect(activeTool.task?.done()).toBe(false); + + // Let it run a bit to generate some data + await new Promise((resolve) => setTimeout(resolve, 50)); + + // Verify queue got some data + const nextItem = await queue.get(); + expect(nextItem).toBeDefined(); + expect(nextItem?.content?.parts?.[0]?.text).toBe('data'); + + // Cancel it + activeTool.task?.cancel(); + await promise; + + expect(activeTool.task?.done()).toBe(true); + }); +});