diff --git a/apps/server/src/core/ai-chat/ai-chat.service.ts b/apps/server/src/core/ai-chat/ai-chat.service.ts index 14826bd6..f9050185 100644 --- a/apps/server/src/core/ai-chat/ai-chat.service.ts +++ b/apps/server/src/core/ai-chat/ai-chat.service.ts @@ -1238,6 +1238,13 @@ export class AiChatService implements OnModuleInit { system, messages, tools, + // Pin the AI SDK per-request retry budget explicitly instead of relying + // on its default (which is also 2). Connection arithmetic per turn: + // (1 + maxRetries=2) × (1 + AI_STREAM_PRE_RESPONSE_RETRIES) network + // connects worst-case — the two retry layers compose, so making the SDK + // side explicit keeps that ceiling visible and pinned against SDK-default + // drift. + maxRetries: 2, // No maxOutputTokens cap on the agent: tool-call arguments (e.g. a full // page body for the write tools) are emitted as OUTPUT tokens, so a fixed // cap would truncate complex tool calls mid-argument. Let the model use its diff --git a/apps/server/src/core/ai-chat/public-share-chat.service.ts b/apps/server/src/core/ai-chat/public-share-chat.service.ts index 236067eb..0c2d21ac 100644 --- a/apps/server/src/core/ai-chat/public-share-chat.service.ts +++ b/apps/server/src/core/ai-chat/public-share-chat.service.ts @@ -307,6 +307,10 @@ export class PublicShareChatService { system, messages: modelMessages, tools, + // Pin the AI SDK per-request retry budget explicitly (matches the SDK + // default of 2). Connection arithmetic: (1 + maxRetries) × (1 + + // AI_STREAM_PRE_RESPONSE_RETRIES) worst-case connects per turn. + maxRetries: 2, // Bound the agent loop for anonymous callers. stopWhen: stepCountIs(5), // Cap per-request output so one anonymous call cannot run up the provider diff --git a/apps/server/src/integrations/ai/ai.service.provider-fetch.spec.ts b/apps/server/src/integrations/ai/ai.service.provider-fetch.spec.ts new file mode 100644 index 00000000..558aca5b --- /dev/null +++ b/apps/server/src/integrations/ai/ai.service.provider-fetch.spec.ts @@ -0,0 +1,86 @@ +// `.provider` alone cannot prove the gemini/ollama chat factories were built +// with the instrumented streaming fetch — a regression dropping it (which drops +// them back to the global undici fetch: no keep-alive recycle, no reset retries, +// unbounded silence timeout; incident classes #140/#175/#310) would still pass. +// So mock the factories and assert the exact fetch argument. jest.mock is +// module-scoped, hence a dedicated file. + +const mockGeminiModel = { provider: 'google.generative-ai', modelId: 'm' }; +const mockOllamaModel = { provider: 'ollama.chat', modelId: 'm' }; + +// jest allows `mock`-prefixed vars inside a jest.mock factory. +const mockCreateGoogle = jest.fn((_settings: unknown) => () => mockGeminiModel); +const mockCreateOllama = jest.fn((_settings: unknown) => () => mockOllamaModel); + +jest.mock('@ai-sdk/google', () => ({ + createGoogleGenerativeAI: (settings: unknown) => mockCreateGoogle(settings), +})); +jest.mock('ai-sdk-ollama', () => ({ + createOllama: (settings: unknown) => mockCreateOllama(settings), +})); + +import { AiService } from './ai.service'; + +describe('AiService.getChatModel provider transport fetch (gemini/ollama)', () => { + function serviceWith(cfg: Record) { + const aiSettings = { + resolve: jest.fn().mockResolvedValue(cfg), + }; + return new AiService( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + aiSettings as any, + { find: jest.fn() } as never, + { decryptSecret: jest.fn() } as never, + ); + } + + beforeEach(() => { + mockCreateGoogle.mockClear(); + mockCreateOllama.mockClear(); + }); + + it('builds the gemini chat model with the instrumented streaming fetch', async () => { + await serviceWith({ + driver: 'gemini', + chatModel: 'gemini-2.5-pro', + apiKey: 'the-key', + }).getChatModel('ws-1'); + expect(mockCreateGoogle).toHaveBeenCalledTimes(1); + expect(mockCreateGoogle).toHaveBeenCalledWith( + expect.objectContaining({ + apiKey: 'the-key', + fetch: expect.any(Function), + }), + ); + }); + + it('builds the ollama chat model with the instrumented streaming fetch', async () => { + await serviceWith({ + driver: 'ollama', + chatModel: 'llama3', + baseUrl: 'http://localhost:11434/api', + }).getChatModel('ws-1'); + expect(mockCreateOllama).toHaveBeenCalledTimes(1); + expect(mockCreateOllama).toHaveBeenCalledWith( + expect.objectContaining({ + baseURL: 'http://localhost:11434/api', + fetch: expect.any(Function), + }), + ); + }); + + it('reuses ONE service-lifetime fetch instance across both providers', async () => { + const svc = serviceWith({ + driver: 'gemini', + chatModel: 'gemini-2.5-pro', + apiKey: 'k', + }); + await svc.getChatModel('ws-1'); + const geminiFetch = mockCreateGoogle.mock.calls[0][0] as { fetch: unknown }; + // Same instance on a second call — the fetch is held for the service + // lifetime to reuse the streaming dispatcher's connection pool. + await svc.getChatModel('ws-1'); + const geminiFetch2 = mockCreateGoogle.mock.calls[1][0] as { fetch: unknown }; + expect(geminiFetch.fetch).toBe(geminiFetch2.fetch); + }); +}); diff --git a/apps/server/src/integrations/ai/ai.service.ts b/apps/server/src/integrations/ai/ai.service.ts index 16aa6997..2f999e72 100644 --- a/apps/server/src/integrations/ai/ai.service.ts +++ b/apps/server/src/integrations/ai/ai.service.ts @@ -190,10 +190,22 @@ export class AiService { }).chat(chatModel); } case 'gemini': - return createGoogleGenerativeAI({ apiKey })(chatModel); + // Route gemini through the same instrumented streaming fetch as openai + // (finite silence timeouts + keep-alive recycling + pre-response + // connection-reset retry). Without it the provider ran on the global + // undici fetch — no keep-alive recycle, no reset retries, default + // (unbounded silence) timeout — so incident classes #140/#175/#310 were + // reproducible for gemini too. + return createGoogleGenerativeAI({ + apiKey, + fetch: this.aiProviderFetch, + })(chatModel); case 'ollama': - // Ollama needs no API key. - return createOllama({ baseURL: baseUrl })(chatModel); + // Ollama needs no API key. Same transport hardening as above (#140/#175/#310). + return createOllama({ + baseURL: baseUrl, + fetch: this.aiProviderFetch, + })(chatModel); default: throw new AiNotConfiguredException(); }