diff --git a/package.json b/package.json index f8fcc5e0c..192d9be7e 100644 --- a/package.json +++ b/package.json @@ -66,7 +66,7 @@ "api-extractor:prod:node": "api-extractor run -c api-extractor.node.json --verbose", "api-extractor:prod:web": "api-extractor run -c api-extractor.web.json --verbose", "api-extractor:prod:tokenizer-node": "api-extractor run -c api-extractor.tokenizer-node.json --verbose", - "unit-test": "tsc && cp src/cross/sentencepiece/sentencepiece_model.pb.js dist/src/cross/sentencepiece/ && jasmine dist/test/unit/**/*_test.js dist/test/unit/**/**/*_test.js dist/test/unit/*_test.js", + "unit-test": "tsc && cp src/cross/sentencepiece/sentencepiece_model.pb.js dist/src/cross/sentencepiece/ && jasmine dist/test/unit/**/*_test.js dist/test/unit/**/**/*_test.js dist/test/unit/*_t[...]", "system-test": "tsc && jasmine dist/test/system/**/*_test.js", "test-server-tests": "tsc && GOOGLE_CLOUD_PROJECT=googcloudproj GOOGLE_CLOUD_LOCATION=googcloudloc jasmine dist/test/system/node/*_test.js -- --test-server", "test-server-tests:record": "tsc && jasmine --fail-fast dist/test/system/node/*_test.js -- --test-server --record", @@ -77,7 +77,7 @@ "lint": "eslint '**/*.ts'", "lint-fix": "eslint --fix '**/*.ts'", "coverage-report": "./test/generate_report.sh", - "generate-proto": "pbjs -t static-module -w es6 -o src/cross/sentencepiece/sentencepiece_model.pb.js src/cross/sentencepiece/sentencepiece_model.proto && pbts -o src/cross/sentencepiece/sentencepiece_model.pb.d.ts src/cross/sentencepiece/sentencepiece_model.pb.js && sed -i.bak 's/import \\* as \\$protobuf from \"protobufjs\\/minimal\"/import \\$protobuf from \"protobufjs\\/minimal.js\"/' src/cross/sentencepiece/sentencepiece_model.pb.js && rm src/cross/sentencepiece/sentencepiece_model.pb.js.bak" + "generate-proto": "pbjs -t static-module -w es6 -o src/cross/sentencepiece/sentencepiece_model.pb.js src/cross/sentencepiece/sentencepiece_model.proto && pbts -o src/cross/sentencepiece/senten[...]" }, "engines": { "node": ">=20.0.0" @@ -135,15 +135,14 @@ "typedoc": "^0.27.0", "typescript": "~5.4.0", "typescript-eslint": "8.24.1", - "undici": "^7.16.0", - "undici-types": "^7.16.0", "zod": "^3.25.0", "zod-to-json-schema": "^3.25.0" }, "dependencies": { "google-auth-library": "^10.3.0", "protobufjs": "^7.5.4", - "ws": "^8.18.0" + "ws": "^8.18.0", + "undici": "^7.16.0" }, "peerDependencies": { "@modelcontextprotocol/sdk": "^1.24.0" @@ -163,4 +162,4 @@ "homepage": "https://github.com/googleapis/js-genai#readme", "author": "", "license": "Apache-2.0" -} +} \ No newline at end of file diff --git a/src/_api_client.ts b/src/_api_client.ts index 83b4d0a85..02fa3c85d 100644 --- a/src/_api_client.ts +++ b/src/_api_client.ts @@ -12,6 +12,7 @@ import {uploadToFileSearchStoreConfigToMldev} from './converters/_filesearchstor import {ApiError} from './errors.js'; import {GeminiNextGenAPIClientAdapter} from './interactions/client-adapter.js'; import * as types from './types.js'; +import type { RequestInit as UndiciRequestInit } from 'undici'; const CONTENT_TYPE_HEADER = 'Content-Type'; const SERVER_TIMEOUT_HEADER = 'X-Server-Timeout'; @@ -462,6 +463,17 @@ export class ApiClient implements GeminiNextGenAPIClientAdapter { // https://nodejs.org/api/timers.html#timeoutunref timeoutHandle.unref(); } + if (typeof process !== 'undefined' && process.versions?.node) { + try { + const { Agent } = await import('undici'); + (requestInit as UndiciRequestInit).dispatcher = new Agent({ + headersTimeout: httpOptions.timeout, + bodyTimeout: httpOptions.timeout, + }); + } catch { + // Ignore errors, undici might not be available. + } + } } if (abortSignal) { abortSignal.addEventListener('abort', () => { diff --git a/test/unit/api_client_test.ts b/test/unit/api_client_test.ts index 8803c4dcf..0698c77f9 100644 --- a/test/unit/api_client_test.ts +++ b/test/unit/api_client_test.ts @@ -14,6 +14,7 @@ import {CrossDownloader} from '../../src/cross/_cross_downloader.js'; import {CrossUploader} from '../../src/cross/_cross_uploader.js'; import * as types from '../../src/types.js'; import {FakeAuth} from '../_fake_auth.js'; +import {Agent, type RequestInit as UndiciRequestInit } from 'undici'; const fetchOkOptions = { status: 200, @@ -809,6 +810,29 @@ describe('ApiClient', () => { // @ts-expect-error TS2532: Object is possibly 'undefined'. expect(fetchArgs[0][1].signal.aborted).toBeTrue(); }); + it('should set dispatcher with timeouts in Node.js', async () => { + const client = new ApiClient({ + auth: new FakeAuth('test-api-key'), + apiKey: 'test-api-key', + httpOptions: {timeout: 1000}, + uploader: new CrossUploader(), + downloader: new CrossDownloader(), + }); + const fetchSpy = spyOn(global, 'fetch').and.returnValue( + Promise.resolve( + new Response( + JSON.stringify(mockGenerateContentResponse), + fetchOkOptions, + ), + ), + ); + + await client.request({path: 'test-path', httpMethod: 'POST'}); + const fetchArgs = fetchSpy.calls.first().args; + const requestInit = fetchArgs[1] as UndiciRequestInit; + expect(requestInit.dispatcher).toBeDefined(); + expect(requestInit.dispatcher).toBeInstanceOf(Agent); + }); it('should apply requestHttpOptions when provided', async () => { const client = new ApiClient({ auth: new FakeAuth('test-api-key'),