From 031a1f3293174cbe1867a39734b0a6ec8eed8cf9 Mon Sep 17 00:00:00 2001 From: yyhhyyyyyy Date: Wed, 15 Jul 2026 15:35:29 +0800 Subject: [PATCH 1/2] fix(memory): complete provider cancellation --- .../memory-provider-text-cancellation/spec.md | 128 +++++++ .../compactionService.ts | 3 +- src/main/presenter/githubCopilotDeviceFlow.ts | 7 +- src/main/presenter/index.ts | 11 +- .../llmProviderPresenter/aiSdk/runtime.ts | 171 +++++---- .../llmProviderPresenter/baseProvider.ts | 30 +- .../presenter/llmProviderPresenter/index.ts | 13 +- .../providers/acpProvider.ts | 295 +++++++++++---- .../providers/aiSdkProvider.ts | 86 ++++- .../providers/githubCopilotProvider.ts | 30 +- .../providers/ollamaProvider.ts | 20 +- .../providers/voiceAIProvider.ts | 76 ++-- src/main/presenter/memoryProviderBindings.ts | 32 ++ .../types/presenters/core.presenter.d.ts | 3 +- .../presenters/llmprovider.presenter.d.ts | 3 +- test/main/presenter/acpProvider.test.ts | 343 +++++++++++++++++- .../compactionService.test.ts | 8 + .../presenter/githubCopilotDeviceFlow.test.ts | 61 ++++ .../presenter/llmProviderPresenter.test.ts | 23 ++ .../llmProviderPresenter/aiSdkRuntime.test.ts | 340 +++++++++++++++++ .../anthropicProvider.test.ts | 28 +- .../llmProviderPresenter/baseProvider.test.ts | 91 +++++ .../githubCopilotProvider.test.ts | 76 +++- .../ollamaProviderCancellation.test.ts | 42 +++ .../voiceAIProvider.test.ts | 107 ++++++ .../presenter/memoryProviderBindings.test.ts | 54 +++ .../presenter/memoryProviderGateway.test.ts | 83 +++++ 27 files changed, 1958 insertions(+), 206 deletions(-) create mode 100644 docs/issues/memory-provider-text-cancellation/spec.md create mode 100644 src/main/presenter/memoryProviderBindings.ts create mode 100644 test/main/presenter/githubCopilotDeviceFlow.test.ts create mode 100644 test/main/presenter/llmProviderPresenter/ollamaProviderCancellation.test.ts create mode 100644 test/main/presenter/llmProviderPresenter/voiceAIProvider.test.ts create mode 100644 test/main/presenter/memoryProviderBindings.test.ts diff --git a/docs/issues/memory-provider-text-cancellation/spec.md b/docs/issues/memory-provider-text-cancellation/spec.md new file mode 100644 index 0000000000..ea18b910ea --- /dev/null +++ b/docs/issues/memory-provider-text-cancellation/spec.md @@ -0,0 +1,128 @@ +# Memory Provider Text Cancellation + +## Issue + +`MemoryProviderGateway` creates and forwards an `AbortSignal` for every text-generation request, +but the production composition root drops that signal before calling +`ILlmProviderPresenter.generateText`. The presenter and provider-level `generateText` contracts do +not currently accept a caller signal, so an outer Memory deadline or lifecycle abort can return +promptly while the underlying provider request continues running. + +## Impact + +- Cancelled or expired Memory text requests can continue consuming network and provider resources. +- Gateway capacity remains reserved until the uncancelled provider Promise eventually settles. +- Repeated slow requests can exhaust the per-key or global unsettled-request limits even though + callers already observed cancellation. +- Existing generation fences prevent late results from committing, so this is a resource-lifetime + defect rather than a stale-write defect. + +## Root Cause + +- The production Memory provider binding accepts only three arguments for `generateText` and drops + the gateway's fourth `signal` argument. +- `ILlmProviderPresenter.generateText` and `BaseLLMProvider.generateText` have no caller-cancellation + option. +- Provider implementations use model-level timeouts where available, but do not combine them with + the Memory caller signal. +- Gateway unit tests stop at a mocked dependency and therefore do not cover the composition and + transport boundaries where the signal is lost. + +## Review Findings + +The first implementation propagated the signal but exposed additional lifecycle defects: + +- ACP standalone requests share a conversation-scoped session and replace its hooks. Passing a + request-scoped signal into that shared session allows one request to cancel or detach another. +- ACP returns from a caller abort before the original `session/prompt` request settles, so Memory + gateway capacity can be released while the ACP turn is still running. +- GitHub Copilot's primary Device Flow token exchange does not receive the caller signal and can + fall back to a second credential request after cancellation. +- Replacing AI SDK's positional `systemPrompt` argument with an options object silently drops the + prompt for untyped or out-of-tree callers using the old concrete-provider API. +- VoiceAI creates timeout/listener state before validation paths that can throw outside its cleanup + boundary. AI SDK media transports need equivalent cleanup and active-abort coverage. + +## Fix Plan + +- Add an optional trailing `{ signal?: AbortSignal }` option to + `ILlmProviderPresenter.generateText` and forward it into a provider-level options object. +- Add a named AI SDK provider option so `systemPrompt` and `signal` share one extensible internal + contract, while retaining a deprecated runtime-compatible positional `systemPrompt` overload. +- Extract the Memory LLM bindings from the application composition root and test exact signal + forwarding for rate-limit admission, embeddings, dimensions, and text generation. +- Pass the existing compaction signal into `generateText` while retaining its defensive outer abort + race. +- Propagate caller cancellation through AI SDK/Ollama, GitHub Copilot, VoiceAI, and ACP transports. + Combine it with existing model timeouts without changing which source wins or how timeout errors + are classified. +- Preserve the gateway invariant that capacity is released only when the provider Promise settles; + never decrement capacity merely because the outer deadline or abort race completed. +- Serialize ACP standalone turns per reused session with an abort-aware FIFO owner queue. Only the + active owner may install session hooks or cancel a prompt; queued aborts leave the session alone. +- After an ACP prompt is dispatched, send protocol cancellation promptly but retain ownership and + consume late events until the original prompt Promise settles. Cancel pending and late permission + requests without surfacing them to the user. +- Keep the generic provider text options signal-only. Add an AI SDK-specific options type and a + deprecated runtime-compatible overload for the old positional `systemPrompt` argument. +- Propagate cancellation into the GitHub Device Flow token fetch and prevent fallback after abort. +- Extend cleanup boundaries so every timer and caller-signal listener is disposed on validation, + setup, transport, polling, and download failures. + +## Constraints + +- No renderer behavior, route, database, persisted schema, or IPC wire-format changes. +- New API parameters remain optional, so existing callers keep their current behavior. +- Memory cancellation, deadline, and capacity error codes remain unchanged. +- Late provider results and rejections remain absorbed and cannot commit side effects. +- Providers or upstream SDKs that genuinely ignore cancellation may still settle late and retain + capacity until then. +- ACP standalone calls keep the existing per-model session/process reuse. They are serialized + rather than moved to per-request sessions or processes. + +## Tasks + +- [x] Document the issue, constraints, and acceptance criteria. +- [x] Add and forward the presenter/provider `generateText` cancellation options. +- [x] Extract and test production Memory provider bindings. +- [x] Propagate cancellation through every in-tree provider implementation. +- [x] Pass the existing compaction signal into text generation. +- [x] Add composition, transport, timeout-composition, and capacity regression tests. +- [x] Add ACP standalone ownership and prompt-drain semantics. +- [x] Restore AI SDK positional `systemPrompt` compatibility and narrow generic option naming. +- [x] Propagate Device Flow cancellation and close VoiceAI/media cleanup gaps. +- [x] Add review regression tests for ACP, compatibility, Device Flow, cleanup, TTS, and video. +- [x] Run formatting, i18n validation, typecheck, focused and full main tests, and lint. + +## Validation + +- [x] The production Memory text binding forwards the gateway's exact signal. +- [x] A pre-aborted signal prevents a provider transport request from starting. +- [x] Aborting an active request reaches the in-tree transport or ACP cancel operation promptly. +- [x] Caller cancellation and model timeout coexist, with the first source winning. +- [x] Signal-aware provider settlement releases gateway capacity; a provider that ignores abort + continues to retain capacity until settlement. +- [x] Memory deadline/abort classification and late-result side-effect fencing remain unchanged. +- [x] Compaction still returns promptly and never persists a late summary after cancellation. +- [x] Concurrent ACP standalone requests cannot replace or cancel each other's session hooks. +- [x] ACP provider and gateway capacity remain occupied until the dispatched prompt settles. +- [x] Queued ACP requests can abort without opening or mutating the shared session. +- [x] Device Flow token exchange receives the exact request signal and does not fall back after + cancellation. +- [x] Legacy positional AI SDK system prompts and the new options form produce identical prompts. +- [x] Validation and media-route failures leave no request timer or caller-signal listener behind. + +## Verification + +- Focused Vitest: 12 files and 163 tests passed. +- `pnpm run typecheck:node`: passed. +- `pnpm run test:main`: 369 files and 4201 tests passed; 15 files and 196 tests skipped by existing + suite conditions. +- `pnpm run format` and `pnpm run format:check`: passed. +- `pnpm run i18n`: passed with no missing keys or invalid translations. +- `pnpm run lint`: passed, including agent cleanup and architecture guards. +- `git diff --check`: passed. + +## GitHub + +No GitHub issue sync was requested. diff --git a/src/main/presenter/agentRuntimePresenter/compactionService.ts b/src/main/presenter/agentRuntimePresenter/compactionService.ts index 31731f6deb..3bbc876ff7 100644 --- a/src/main/presenter/agentRuntimePresenter/compactionService.ts +++ b/src/main/presenter/agentRuntimePresenter/compactionService.ts @@ -1002,7 +1002,8 @@ export class CompactionService { prompt, model.modelId, 0.2, - this.getSummaryOutputTokens(reserveTokens) + this.getSummaryOutputTokens(reserveTokens), + { signal } ), signal ) diff --git a/src/main/presenter/githubCopilotDeviceFlow.ts b/src/main/presenter/githubCopilotDeviceFlow.ts index a7892f3a65..634895364d 100644 --- a/src/main/presenter/githubCopilotDeviceFlow.ts +++ b/src/main/presenter/githubCopilotDeviceFlow.ts @@ -78,7 +78,8 @@ export class GitHubCopilotDeviceFlow { * 获取 Copilot API token * 使用OAuth token交换Copilot API token */ - public async getCopilotToken(): Promise { + public async getCopilotToken(signal?: AbortSignal): Promise { + signal?.throwIfAborted() if (!this.oauthToken) { throw new Error('No OAuth token available') } @@ -93,7 +94,8 @@ export class GitHubCopilotDeviceFlow { Authorization: `Bearer ${this.oauthToken}`, Accept: 'application/json', 'User-Agent': 'DeepChat/1.0.0' - } + }, + ...(signal ? { signal } : {}) }) if (!response.ok) { @@ -106,6 +108,7 @@ export class GitHubCopilotDeviceFlow { const data = (await response.json()) as { token: string; expires_at: number } return data.token } catch (error) { + signal?.throwIfAborted() console.error('[GitHub Copilot][DeviceFlow] Failed to get Copilot token:', error) throw error } diff --git a/src/main/presenter/index.ts b/src/main/presenter/index.ts index 8124214fb5..93d2e93654 100644 --- a/src/main/presenter/index.ts +++ b/src/main/presenter/index.ts @@ -123,6 +123,7 @@ import { rtkRuntimeService } from '@/agent/shared/process/rtkRuntimeService' import { SessionHistorySearch } from '@/routes/sessions/sessionHistorySearch' import { SessionTranslation } from '@/routes/sessions/sessionTranslation' import { AgentSessionExportService } from './exporter/agentSessionExporter' +import { createMemoryProviderBindings } from './memoryProviderBindings' type MemoryMaintenanceConfigChangeTarget = Pick< MemoryPresenter, @@ -637,15 +638,7 @@ export class Presenter implements IPresenter { .filter( (agentId) => agentRepository.resolveDeepChatAgentConfig(agentId).memoryEnabled === true ), - executeWithRateLimit: (providerId, options) => - this.llmproviderPresenter.executeWithRateLimit(providerId, { signal: options.signal }), - getEmbeddings: (providerId, modelId, texts, signal) => - this.llmproviderPresenter.getEmbeddings(providerId, modelId, texts, signal), - getDimensions: (providerId, modelId, signal) => - this.llmproviderPresenter.getDimensions(providerId, modelId, signal), - generateText: async (providerId, modelId, prompt) => - (await this.llmproviderPresenter.generateText(providerId, prompt, modelId, 0.2)).content ?? - '', + ...createMemoryProviderBindings(this.llmproviderPresenter), createVectorStore: (agentId, embedding, dimensions) => { if (!isSafeAgentId(agentId)) { throw new Error(`[Memory] refusing to open vector store for unsafe agentId: ${agentId}`) diff --git a/src/main/presenter/llmProviderPresenter/aiSdk/runtime.ts b/src/main/presenter/llmProviderPresenter/aiSdk/runtime.ts index c262193c73..74318c8048 100644 --- a/src/main/presenter/llmProviderPresenter/aiSdk/runtime.ts +++ b/src/main/presenter/llmProviderPresenter/aiSdk/runtime.ts @@ -437,6 +437,16 @@ function extractChatAudioContentData(content: unknown): string | undefined { return typeof audioData === 'string' && audioData ? audioData : undefined } +function relayAbortSignal(source: AbortSignal | undefined, target: AbortController): () => void { + if (!source) return () => {} + + const onAbort = () => target.abort(source.reason) + if (source.aborted) onAbort() + else source.addEventListener('abort', onAbort, { once: true }) + + return () => source.removeEventListener('abort', onAbort) +} + /** * Pattern A: calls the standard OpenAI-compatible /audio/speech endpoint. */ @@ -446,7 +456,8 @@ async function executeTtsPatternA( text: string, modelId: string, modelConfig: ModelConfig, - timeout: number | undefined + timeout: number | undefined, + signal?: AbortSignal ): Promise<{ base64: string; mimeType: string }> { const tts = normalizeTtsSettings(modelConfig.tts) const format = tts?.responseFormat ?? 'mp3' @@ -468,10 +479,11 @@ async function executeTtsPatternA( const controller = new AbortController() const timeoutId = timeout ? setTimeout(() => controller.abort(), timeout) : undefined - const proxyUrl = proxyConfig.getProxyUrl() - const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined + const disposeCallerAbort = relayAbortSignal(signal, controller) try { + const proxyUrl = proxyConfig.getProxyUrl() + const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { method: 'POST', headers: { @@ -495,6 +507,7 @@ async function executeTtsPatternA( return { base64, mimeType: ttsFormatToMimeType(format) } } finally { if (timeoutId !== undefined) clearTimeout(timeoutId) + disposeCallerAbort() } } @@ -508,7 +521,8 @@ async function executeTtsPatternB( text: string, modelId: string, modelConfig: ModelConfig, - timeout: number | undefined + timeout: number | undefined, + signal?: AbortSignal ): Promise<{ base64: string; mimeType: string }> { const tts = normalizeTtsSettings(modelConfig.tts) const format = tts?.responseFormat ?? 'wav' @@ -530,10 +544,11 @@ async function executeTtsPatternB( const controller = new AbortController() const timeoutId = timeout ? setTimeout(() => controller.abort(), timeout) : undefined - const proxyUrl = proxyConfig.getProxyUrl() - const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined + const disposeCallerAbort = relayAbortSignal(signal, controller) try { + const proxyUrl = proxyConfig.getProxyUrl() + const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { method: 'POST', headers: { @@ -571,6 +586,7 @@ async function executeTtsPatternB( return { base64: audioData, mimeType: ttsFormatToMimeType(format) } } finally { if (timeoutId !== undefined) clearTimeout(timeoutId) + disposeCallerAbort() } } @@ -580,7 +596,8 @@ async function executeTtsPatternC( text: string, modelId: string, modelConfig: ModelConfig, - timeout: number | undefined + timeout: number | undefined, + signal?: AbortSignal ): Promise<{ base64: string; mimeType: string }> { const tts = normalizeTtsSettings(modelConfig.tts) const baseUrl = resolveGeminiTtsBaseUrl(provider) @@ -611,10 +628,11 @@ async function executeTtsPatternC( const controller = new AbortController() const timeoutId = timeout ? setTimeout(() => controller.abort(), timeout) : undefined - const proxyUrl = proxyConfig.getProxyUrl() - const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined + const disposeCallerAbort = relayAbortSignal(signal, controller) try { + const proxyUrl = proxyConfig.getProxyUrl() + const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { method: 'POST', headers: { @@ -659,6 +677,7 @@ async function executeTtsPatternC( ) } finally { if (timeoutId !== undefined) clearTimeout(timeoutId) + disposeCallerAbort() } } @@ -1010,14 +1029,14 @@ function resolveVideoTaskStatus(response: VideoGenerationTaskResponse | null | u function delayWithAbort(ms: number, signal: AbortSignal): Promise { return new Promise((resolve, reject) => { if (signal.aborted) { - reject(signal.reason instanceof Error ? signal.reason : new Error('Aborted')) + reject(signal.reason) return } const onAbort = () => { clearTimeout(timeoutId) signal.removeEventListener('abort', onAbort) - reject(signal.reason instanceof Error ? signal.reason : new Error('Aborted')) + reject(signal.reason) } const timeoutId = setTimeout(() => { @@ -1035,7 +1054,8 @@ async function executeOpenAICompatibleVideoGeneration( modelId: string, prompt: string, modelConfig: ModelConfig, - timeout: number | undefined + timeout: number | undefined, + signal?: AbortSignal ): Promise<{ base64: string; mimeType: string }> { const normalizedOptions = resolveVideoGenerationRequestOptions( prompt, @@ -1047,54 +1067,56 @@ async function executeOpenAICompatibleVideoGeneration( const body = buildVideoGenerationRequestBody(provider, modelId, prompt, normalizedOptions) const controller = new AbortController() const timeoutId = timeout ? setTimeout(() => controller.abort(), timeout) : undefined - const proxyUrl = proxyConfig.getProxyUrl() - const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined - - const fetchJson = async (url: string, init: RequestInit): Promise => { - const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { - ...init, - headers: { - ...defaultHeaders, - Authorization: `Bearer ${provider.oauthToken || provider.apiKey || ''}`, - ...(init.headers as Record | undefined) - }, - signal: controller.signal - } - if (dispatcher) fetchInit.dispatcher = dispatcher + const disposeCallerAbort = relayAbortSignal(signal, controller) - const response = await fetch(url, fetchInit) - if (!response.ok) { - const errorText = await response.text().catch(() => '') - throw new Error(`Video request failed (${response.status}): ${errorText}`) - } + try { + const proxyUrl = proxyConfig.getProxyUrl() + const dispatcher = proxyUrl ? new ProxyAgent(proxyUrl) : undefined + + const fetchJson = async (url: string, init: RequestInit): Promise => { + const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { + ...init, + headers: { + ...defaultHeaders, + Authorization: `Bearer ${provider.oauthToken || provider.apiKey || ''}`, + ...(init.headers as Record | undefined) + }, + signal: controller.signal + } + if (dispatcher) fetchInit.dispatcher = dispatcher - return (await response.json()) as T - } + const response = await fetch(url, fetchInit) + if (!response.ok) { + const errorText = await response.text().catch(() => '') + throw new Error(`Video request failed (${response.status}): ${errorText}`) + } - const fetchBinary = async (url: string): Promise<{ buffer: ArrayBuffer; mimeType: string }> => { - const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { - method: 'GET', - headers: { - ...defaultHeaders, - Authorization: `Bearer ${provider.oauthToken || provider.apiKey || ''}` - }, - signal: controller.signal + return (await response.json()) as T } - if (dispatcher) fetchInit.dispatcher = dispatcher - const response = await fetch(url, fetchInit) - if (!response.ok) { - const errorText = await response.text().catch(() => '') - throw new Error(`Video content download failed (${response.status}): ${errorText}`) - } + const fetchBinary = async (url: string): Promise<{ buffer: ArrayBuffer; mimeType: string }> => { + const fetchInit: RequestInit & { dispatcher?: ProxyAgent } = { + method: 'GET', + headers: { + ...defaultHeaders, + Authorization: `Bearer ${provider.oauthToken || provider.apiKey || ''}` + }, + signal: controller.signal + } + if (dispatcher) fetchInit.dispatcher = dispatcher - return { - buffer: await response.arrayBuffer(), - mimeType: response.headers.get('content-type')?.split(';')[0]?.trim() || 'video/mp4' + const response = await fetch(url, fetchInit) + if (!response.ok) { + const errorText = await response.text().catch(() => '') + throw new Error(`Video content download failed (${response.status}): ${errorText}`) + } + + return { + buffer: await response.arrayBuffer(), + mimeType: response.headers.get('content-type')?.split(';')[0]?.trim() || 'video/mp4' + } } - } - try { let task = await fetchJson(createUrl, { method: 'POST', headers: { @@ -1138,6 +1160,7 @@ async function executeOpenAICompatibleVideoGeneration( if (timeoutId !== undefined) { clearTimeout(timeoutId) } + disposeCallerAbort() } } @@ -1232,14 +1255,25 @@ function usageToLlmResponse( } } +function combineRequestSignal( + timeout: number | undefined, + callerSignal?: AbortSignal +): AbortSignal | undefined { + const timeoutSignal = timeout ? AbortSignal.timeout(timeout) : undefined + if (timeoutSignal && callerSignal) return AbortSignal.any([callerSignal, timeoutSignal]) + return callerSignal ?? timeoutSignal +} + export async function runAiSdkGenerateText( context: AiSdkRuntimeContext, messages: ChatMessage[], modelId: string, modelConfig: ModelConfig, temperature?: number, - maxTokens?: number + maxTokens?: number, + signal?: AbortSignal ): Promise { + signal?.throwIfAborted() const normalizedModelConfig = normalizeRuntimeModelConfig(context, modelId, modelConfig) const runtime = await buildPromptRuntime(context, messages, modelId, normalizedModelConfig, []) const { shouldSendTemperature, temperature: resolvedTemperature } = resolveRuntimeTemperature( @@ -1265,13 +1299,16 @@ export async function runAiSdkGenerateText( body: requestBody }) + const requestSignal = combineRequestSignal(timeout, signal) + requestSignal?.throwIfAborted() + const result = await generateText({ model: runtime.providerContext.model, ...(runtime.instructions ? { instructions: runtime.instructions } : {}), messages: runtime.messages, allowSystemInMessages: false, providerOptions: runtime.providerOptions as any, - ...(timeout ? { abortSignal: AbortSignal.timeout(timeout) } : {}), + ...(requestSignal ? { abortSignal: requestSignal } : {}), ...(shouldSendTemperature && resolvedTemperature !== undefined ? { temperature: resolvedTemperature } : {}), @@ -1293,8 +1330,10 @@ export async function* runAiSdkCoreStream( modelConfig: ModelConfig, temperature: number, maxTokens: number, - tools: MCPToolDefinition[] + tools: MCPToolDefinition[], + signal?: AbortSignal ): AsyncGenerator { + signal?.throwIfAborted() const normalizedModelConfig = normalizeRuntimeModelConfig(context, modelId, modelConfig) const timeout = resolveRequestTimeout(normalizedModelConfig) @@ -1310,7 +1349,8 @@ export async function* runAiSdkCoreStream( text, modelId, normalizedModelConfig, - timeout + timeout, + signal ) : usePatternB ? await executeTtsPatternB( @@ -1319,7 +1359,8 @@ export async function* runAiSdkCoreStream( text, modelId, normalizedModelConfig, - timeout + timeout, + signal ) : await executeTtsPatternA( context.provider, @@ -1327,7 +1368,8 @@ export async function* runAiSdkCoreStream( text, modelId, normalizedModelConfig, - timeout + timeout, + signal ) const dataUrl = `data:${mimeType};base64,${base64}` @@ -1372,7 +1414,8 @@ export async function* runAiSdkCoreStream( modelId, prompt, normalizedModelConfig, - timeout + timeout, + signal ) yield { @@ -1422,11 +1465,14 @@ export async function* runAiSdkCoreStream( } }) + const requestSignal = combineRequestSignal(timeout, signal) + requestSignal?.throwIfAborted() + const result = await generateImage({ model: providerContext.imageModel, prompt, ...imageGenerationRequestOptions, - ...(timeout ? { abortSignal: AbortSignal.timeout(timeout) } : {}) + ...(requestSignal ? { abortSignal: requestSignal } : {}) }) for (const image of result.images) { @@ -1472,6 +1518,9 @@ export async function* runAiSdkCoreStream( body: requestBody }) + const requestSignal = combineRequestSignal(timeout, signal) + requestSignal?.throwIfAborted() + const result = streamText({ model: runtime.providerContext.model, ...(runtime.instructions ? { instructions: runtime.instructions } : {}), @@ -1479,7 +1528,7 @@ export async function* runAiSdkCoreStream( allowSystemInMessages: false, tools: runtime.tools, providerOptions: runtime.providerOptions as any, - ...(timeout ? { abortSignal: AbortSignal.timeout(timeout) } : {}), + ...(requestSignal ? { abortSignal: requestSignal } : {}), ...(shouldSendTemperature && resolvedTemperature !== undefined ? { temperature: resolvedTemperature } : {}), diff --git a/src/main/presenter/llmProviderPresenter/baseProvider.ts b/src/main/presenter/llmProviderPresenter/baseProvider.ts index 344f36ff86..7445d2cb76 100644 --- a/src/main/presenter/llmProviderPresenter/baseProvider.ts +++ b/src/main/presenter/llmProviderPresenter/baseProvider.ts @@ -24,6 +24,10 @@ export function isAudioTranscriptionNotSupportedError(error: unknown): boolean { return error instanceof Error && error.message === AUDIO_TRANSCRIPTION_NOT_SUPPORTED_ERROR } +export interface ProviderGenerateTextOptions { + signal?: AbortSignal +} + /** * Base LLM Provider Abstract Class * @@ -117,14 +121,24 @@ export abstract class BaseLLMProvider { this.loadCachedModels() } - protected createModelRequestSignal(modelConfig?: Pick | null): { + protected createModelRequestSignal( + modelConfig?: Pick | null, + callerSignal?: AbortSignal + ): { signal?: AbortSignal timeoutMs?: number dispose: () => void } { const timeoutMs = this.resolveModelRequestTimeout(modelConfig) + if (!timeoutMs && !callerSignal) { + return { + dispose: () => {} + } + } + if (!timeoutMs) { return { + signal: callerSignal, dispose: () => {} } } @@ -133,11 +147,20 @@ export abstract class BaseLLMProvider { const timeoutId = setTimeout(() => { controller.abort(this.createModelRequestTimeoutError(timeoutMs)) }, timeoutMs) + const onCallerAbort = () => controller.abort(callerSignal?.reason) + + if (callerSignal) { + if (callerSignal.aborted) onCallerAbort() + else callerSignal.addEventListener('abort', onCallerAbort, { once: true }) + } return { signal: controller.signal, timeoutMs, - dispose: () => clearTimeout(timeoutId) + dispose: () => { + clearTimeout(timeoutId) + callerSignal?.removeEventListener('abort', onCallerAbort) + } } } @@ -705,7 +728,8 @@ ${this.convertToolsToXml(tools)} prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: ProviderGenerateTextOptions ): Promise /** diff --git a/src/main/presenter/llmProviderPresenter/index.ts b/src/main/presenter/llmProviderPresenter/index.ts index f1f3270757..fabbccfbcd 100644 --- a/src/main/presenter/llmProviderPresenter/index.ts +++ b/src/main/presenter/llmProviderPresenter/index.ts @@ -38,7 +38,11 @@ import type { AcpProviderAdminPort } from '../runtimePorts' import { CONFIG_EVENTS, PROVIDER_DB_EVENTS } from '@/events' -import { BaseLLMProvider, isAudioTranscriptionNotSupportedError } from './baseProvider' +import { + BaseLLMProvider, + isAudioTranscriptionNotSupportedError, + type ProviderGenerateTextOptions +} from './baseProvider' import { ProviderConfig, StreamState } from './types' import { RateLimitManager } from './managers/rateLimitManager' import { ProviderInstanceManager } from './managers/providerInstanceManager' @@ -401,10 +405,13 @@ export class LLMProviderPresenter prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: Pick ): Promise { const provider = this.getProviderInstance(providerId) - return provider.generateText(prompt, modelId, temperature, maxTokens) + return options + ? provider.generateText(prompt, modelId, temperature, maxTokens, options) + : provider.generateText(prompt, modelId, temperature, maxTokens) } async generateCompletionStandalone( diff --git a/src/main/presenter/llmProviderPresenter/providers/acpProvider.ts b/src/main/presenter/llmProviderPresenter/providers/acpProvider.ts index 06a389521e..62f2507f43 100644 --- a/src/main/presenter/llmProviderPresenter/providers/acpProvider.ts +++ b/src/main/presenter/llmProviderPresenter/providers/acpProvider.ts @@ -1,7 +1,11 @@ import logger from '@shared/logger' import type * as schema from '@agentclientprotocol/sdk/dist/schema/index.js' import type { ClientSideConnection as ClientSideConnectionType } from '@agentclientprotocol/sdk' -import { BaseLLMProvider, SUMMARY_TITLES_PROMPT } from '../baseProvider' +import { + BaseLLMProvider, + SUMMARY_TITLES_PROMPT, + type ProviderGenerateTextOptions +} from '../baseProvider' import type { AcpConfigState, ChatMessage, @@ -41,6 +45,7 @@ import { nanoid } from 'nanoid' import type { ProviderMcpRuntimePort } from '../runtimePorts' import { resolveAcpAgentAlias } from '@shared/utils/acpAgentAlias' import { toAppSessionId } from '@/agent/shared/agentSessionIds' +import { awaitWithAbort } from '@/lib/awaitWithAbort' type EventQueue = { push: (event: LLMCoreStreamEvent | null) => void @@ -50,6 +55,7 @@ type EventQueue = { type RunPromptOptions = { onPromptSucceeded?: () => void + signal?: AbortSignal } type PermissionRequestContext = { @@ -67,6 +73,20 @@ type PendingPermissionState = { timeoutId: ReturnType } +type StandaloneRequest = { + started: boolean + signal?: AbortSignal + run: () => Promise + resolve: (value: unknown) => void + reject: (reason?: unknown) => void + detachAbort?: () => void +} + +type StandaloneRequestQueue = { + active: boolean + pending: StandaloneRequest[] +} + const ACP_PERMISSION_TIMEOUT_MS = 60_000 type AcpConnectionWithModelSelection = { @@ -127,6 +147,7 @@ export class AcpProvider extends BaseLLMProvider { private readonly promptController: AcpPromptController private readonly messageFormatter = new AcpMessageFormatter() private readonly pendingPermissions = new Map() + private standaloneRequestQueues?: Map constructor( provider: LLM_PROVIDER, @@ -287,7 +308,8 @@ export class AcpProvider extends BaseLLMProvider { messages: ChatMessage[], modelId: string, temperature: number = 0.6, - maxTokens: number = 4096 + maxTokens: number = 4096, + signal?: AbortSignal ): Promise { const modelConfig = this.configPresenter.getModelConfig(modelId, this.provider.id) const { content, reasoning } = await this.collectFromStream( @@ -295,7 +317,8 @@ export class AcpProvider extends BaseLLMProvider { modelId, modelConfig, temperature, - maxTokens + maxTokens, + signal ) return { @@ -317,9 +340,16 @@ export class AcpProvider extends BaseLLMProvider { prompt: string, modelId: string, temperature: number = 0.6, - maxTokens: number = 4096 + maxTokens: number = 4096, + options?: ProviderGenerateTextOptions ): Promise { - return this.completions([{ role: 'user', content: prompt }], modelId, temperature, maxTokens) + return this.completions( + [{ role: 'user', content: prompt }], + modelId, + temperature, + maxTokens, + options?.signal + ) } async *coreStream( @@ -328,8 +358,10 @@ export class AcpProvider extends BaseLLMProvider { modelConfig: ModelConfig, _temperature: number, _maxTokens: number, - _tools: MCPToolDefinition[] + _tools: MCPToolDefinition[], + signal?: AbortSignal ): AsyncGenerator { + signal?.throwIfAborted() const queue = this.createEventQueue() let session: AcpSessionRecord | null = null @@ -352,10 +384,16 @@ export class AcpProvider extends BaseLLMProvider { { onEvents: (events) => events.forEach((event) => queue.push(event)), onPermission: (request) => - this.handlePermissionRequest(queue, request, { - agent, - conversationId: conversationKey - }) + this.handlePermissionRequest( + queue, + request, + { + agent, + conversationId: conversationKey + }, + signal + ), + signal }, workdir ) @@ -370,7 +408,8 @@ export class AcpProvider extends BaseLLMProvider { ? () => { activeSession.systemPromptSent = true } - : undefined + : undefined, + signal }) } } @@ -389,10 +428,12 @@ export class AcpProvider extends BaseLLMProvider { } } finally { if (session) { - try { - await session.connection.cancel({ sessionId: session.sessionId }) - } catch (error) { - console.warn('[ACP] cancel failed:', error) + if (this.promptController.getActiveTurn(session.sessionId)) { + try { + await session.connection.cancel({ sessionId: session.sessionId }) + } catch (error) { + console.warn('[ACP] cancel failed:', error) + } } this.acpRuntime.sessionController.clearMappedSession(session.sessionId) this.clearPendingPermissionsForSession(session.sessionId) @@ -1075,12 +1116,15 @@ export class AcpProvider extends BaseLLMProvider { options: RunPromptOptions = {} ): Promise { const timeoutMs = this.resolveModelRequestTimeout(modelConfig) - let timeoutId: NodeJS.Timeout | null = null + let requestSignal: AbortSignal | undefined + let disposeRequestSignal = () => {} const conversationId = modelConfig.conversationId ?? session.conversationId let turnStarted = false let turnId: string | null = null + let promptRequest: Promise | undefined try { + options.signal?.throwIfAborted() const turn = this.promptController.begin({ sessionId: session.sessionId, conversationId @@ -1120,20 +1164,16 @@ export class AcpProvider extends BaseLLMProvider { body: requestBody }) - const promptRequest = session.connection.prompt({ + const requestCancellation = this.createModelRequestSignal(modelConfig, options.signal) + requestSignal = requestCancellation.signal + disposeRequestSignal = requestCancellation.dispose + requestSignal?.throwIfAborted() + + promptRequest = session.connection.prompt({ sessionId: requestBody.sessionId, prompt: requestBody.prompt }) - const response = await (timeoutMs - ? Promise.race([ - promptRequest, - new Promise((_, reject) => { - timeoutId = setTimeout(() => { - reject(this.createModelRequestTimeoutError(timeoutMs)) - }, timeoutMs) - }) - ]) - : promptRequest) + const response = await awaitWithAbort(promptRequest, requestSignal) options.onPromptSucceeded?.() const responseSummary = { sessionId: session.sessionId, @@ -1161,32 +1201,57 @@ export class AcpProvider extends BaseLLMProvider { } queue.push(createStreamEvent.stop(mapAcpPromptStopReason(response.stopReason))) } catch (error) { - if (timeoutMs && error instanceof Error && error.name === 'AbortError') { - try { - await session.connection.cancel({ sessionId: session.sessionId }) - } catch (cancelError) { - console.warn('[ACP] cancel after timeout failed:', cancelError) + const callerCancelled = + options.signal?.aborted === true && + error === options.signal.reason && + (!requestSignal || requestSignal.reason === options.signal.reason) + const requestCancelled = requestSignal?.aborted === true && error === requestSignal.reason + + if (requestCancelled || callerCancelled) { + disposeRequestSignal() + disposeRequestSignal = () => {} + this.clearPendingPermissionsForSession(session.sessionId) + + if (promptRequest) { + const cancelRequest = Promise.resolve().then(() => + session.connection.cancel({ sessionId: session.sessionId }) + ) + const [cancelResult, promptResult] = await Promise.allSettled([ + cancelRequest, + promptRequest + ]) + if (cancelResult.status === 'rejected') { + console.warn('[ACP] cancel after request abort failed:', cancelResult.reason) + } + if (promptResult.status === 'rejected') { + console.info('[ACP] Prompt settled after cancellation:', promptResult.reason) + } } } if (turnStarted) { - const failedTurn = this.promptController.fail(session.sessionId) - if (failedTurn) { + const settledTurn = callerCancelled + ? this.promptController.cancel(session.sessionId) + : this.promptController.fail(session.sessionId) + if (settledTurn) { await this.persistTurnFinish({ - id: failedTurn.id, - status: 'error', - stopReason: 'error', - completedAt: failedTurn.completedAt ?? Date.now() + id: settledTurn.id, + status: callerCancelled ? 'cancelled' : 'error', + stopReason: callerCancelled ? 'cancelled' : 'error', + completedAt: settledTurn.completedAt ?? Date.now() }) } else if (turnId) { await this.persistTurnFinish({ id: turnId, - status: 'error', - stopReason: 'error', + status: callerCancelled ? 'cancelled' : 'error', + stopReason: callerCancelled ? 'cancelled' : 'error', completedAt: Date.now() }) } } + + if (callerCancelled) return + const message = error instanceof Error ? error.message : typeof error === 'string' ? error : 'Unknown error' console.error(`[ACP] Prompt failed for ACP session ${session.sessionId}:`, error) @@ -1199,9 +1264,7 @@ export class AcpProvider extends BaseLLMProvider { }) queue.push(createStreamEvent.error(`ACP: ${message}`)) } finally { - if (timeoutId) { - clearTimeout(timeoutId) - } + disposeRequestSignal() queue.done() } } @@ -1209,9 +1272,27 @@ export class AcpProvider extends BaseLLMProvider { private async handlePermissionRequest( queue: EventQueue, params: schema.RequestPermissionRequest, - context: PermissionRequestContext + context: PermissionRequestContext, + signal?: AbortSignal ): Promise { + if (signal?.aborted) { + return { outcome: { outcome: 'cancelled' } } + } + const { requestId, promise } = this.registerPendingPermission(params, context) + const onAbort = () => { + this.removePendingPermission(requestId)?.resolve({ outcome: { outcome: 'cancelled' } }) + } + + if (signal) { + signal.addEventListener('abort', onAbort, { once: true }) + if (signal.aborted) onAbort() + } + + if (!this.pendingPermissions.has(requestId)) { + signal?.removeEventListener('abort', onAbort) + return await promise + } const toolLabel = params.toolCall.title ?? params.toolCall.toolCallId queue.push( @@ -1223,7 +1304,11 @@ export class AcpProvider extends BaseLLMProvider { createStreamEvent.permission(this.buildPermissionPayload(params, context, requestId)) ) - return await promise + try { + return await promise + } finally { + signal?.removeEventListener('abort', onAbort) + } } private registerPendingPermission( @@ -1395,32 +1480,118 @@ export class AcpProvider extends BaseLLMProvider { modelId: string, modelConfig: ModelConfig, temperature: number, - maxTokens: number + maxTokens: number, + signal?: AbortSignal ): Promise<{ content: string; reasoning: string }> { const mergedConfig: ModelConfig = { ...modelConfig, temperature: temperature ?? modelConfig.temperature, maxTokens: maxTokens ?? modelConfig.maxTokens } + const conversationKey = mergedConfig.conversationId ?? modelId + + return await this.runStandaloneRequest(conversationKey, signal, async () => { + let content = '' + let reasoning = '' + for await (const chunk of this.coreStream( + messages, + modelId, + mergedConfig, + temperature, + maxTokens, + [], + signal + )) { + logger.info('[ACP] collectFromStream: chunk:', chunk) + if (chunk.type === 'text' && chunk.content) { + content += chunk.content + } else if (chunk.type === 'reasoning' && chunk.reasoning_content) { + reasoning += chunk.reasoning_content + } + } + signal?.throwIfAborted() + return { content, reasoning } + }) + } - let content = '' - let reasoning = '' - for await (const chunk of this.coreStream( - messages, - modelId, - mergedConfig, - temperature, - maxTokens, - [] - )) { - logger.info('[ACP] collectFromStream: chunk:', chunk) - if (chunk.type === 'text' && chunk.content) { - content += chunk.content - } else if (chunk.type === 'reasoning' && chunk.reasoning_content) { - reasoning += chunk.reasoning_content + private runStandaloneRequest( + conversationKey: string, + signal: AbortSignal | undefined, + run: () => Promise + ): Promise { + signal?.throwIfAborted() + const queues = (this.standaloneRequestQueues ??= new Map()) + const queue = queues.get(conversationKey) ?? { active: false, pending: [] } + queues.set(conversationKey, queue) + + return new Promise((resolve, reject) => { + const request: StandaloneRequest = { + started: false, + signal, + run, + resolve: (value) => resolve(value as T), + reject + } + + if (signal) { + const onAbort = () => { + if (request.started) return + const pendingIndex = queue.pending.indexOf(request) + if (pendingIndex < 0) return + queue.pending.splice(pendingIndex, 1) + request.detachAbort?.() + request.reject(signal.reason) + this.deleteStandaloneQueueIfIdle(conversationKey, queue) + } + signal.addEventListener('abort', onAbort, { once: true }) + request.detachAbort = () => signal.removeEventListener('abort', onAbort) + if (signal.aborted) { + request.detachAbort() + reject(signal.reason) + this.deleteStandaloneQueueIfIdle(conversationKey, queue) + return + } + } + + queue.pending.push(request) + this.startNextStandaloneRequest(conversationKey, queue) + }) + } + + private startNextStandaloneRequest(conversationKey: string, queue: StandaloneRequestQueue): void { + if (queue.active) return + + while (queue.pending.length > 0) { + const request = queue.pending.shift()! + request.detachAbort?.() + if (request.signal?.aborted) { + request.reject(request.signal.reason) + continue } + + request.started = true + queue.active = true + void Promise.resolve() + .then(request.run) + .then(request.resolve, request.reject) + .finally(() => { + queue.active = false + this.startNextStandaloneRequest(conversationKey, queue) + }) + return + } + + this.deleteStandaloneQueueIfIdle(conversationKey, queue) + } + + private deleteStandaloneQueueIfIdle( + conversationKey: string, + queue: StandaloneRequestQueue + ): void { + if (queue.active || queue.pending.length > 0) return + if (this.standaloneRequestQueues?.get(conversationKey) === queue) { + this.standaloneRequestQueues.delete(conversationKey) } - return { content, reasoning } } private createEventQueue(): EventQueue { diff --git a/src/main/presenter/llmProviderPresenter/providers/aiSdkProvider.ts b/src/main/presenter/llmProviderPresenter/providers/aiSdkProvider.ts index 0fced707a8..d6af055ff0 100644 --- a/src/main/presenter/llmProviderPresenter/providers/aiSdkProvider.ts +++ b/src/main/presenter/llmProviderPresenter/providers/aiSdkProvider.ts @@ -38,7 +38,11 @@ import { import { BedrockClient, ListFoundationModelsCommand } from '@aws-sdk/client-bedrock' import { fromNodeProviderChain } from '@aws-sdk/credential-providers' import { ProxyAgent } from 'undici' -import { BaseLLMProvider, SUMMARY_TITLES_PROMPT } from '../baseProvider' +import { + BaseLLMProvider, + SUMMARY_TITLES_PROMPT, + type ProviderGenerateTextOptions +} from '../baseProvider' import { runAiSdkCoreStream, runAiSdkDimensions, @@ -105,6 +109,10 @@ type ProviderRequestOptions = { signal?: AbortSignal } +export interface AiSdkGenerateTextOptions extends ProviderGenerateTextOptions { + systemPrompt?: string +} + type AudioTranscriptionResponse = { text?: unknown } @@ -1082,7 +1090,8 @@ export class AiSdkProvider extends BaseLLMProvider { modelId: string, temperature?: number, maxTokens?: number, - modelConfig?: ModelConfig + modelConfig?: ModelConfig, + signal?: AbortSignal ): Promise { if (!this.isInitialized) { throw new Error('Provider not initialized') @@ -1092,6 +1101,17 @@ export class AiSdkProvider extends BaseLLMProvider { } const { context, resolvedModelConfig } = this.buildRuntimeContext(modelId, modelConfig) + if (signal) { + return runAiSdkGenerateText( + context, + messages, + modelId, + resolvedModelConfig, + temperature, + maxTokens, + signal + ) + } return runAiSdkGenerateText( context, messages, @@ -1108,7 +1128,8 @@ export class AiSdkProvider extends BaseLLMProvider { modelConfig: ModelConfig, temperature: number, maxTokens: number, - tools: MCPToolDefinition[] + tools: MCPToolDefinition[], + signal?: AbortSignal ): AsyncGenerator { if (!this.isInitialized) { throw new Error('Provider not initialized') @@ -1118,6 +1139,19 @@ export class AiSdkProvider extends BaseLLMProvider { } const { context, resolvedModelConfig } = this.buildRuntimeContext(modelId, modelConfig) + if (signal) { + yield* runAiSdkCoreStream( + context, + messages, + modelId, + resolvedModelConfig, + temperature, + maxTokens, + tools, + signal + ) + return + } yield* runAiSdkCoreStream( context, messages, @@ -1135,7 +1169,8 @@ export class AiSdkProvider extends BaseLLMProvider { temperature?: number, maxTokens?: number, tools: MCPToolDefinition[] = [], - modelConfig?: ModelConfig + modelConfig?: ModelConfig, + signal?: AbortSignal ): Promise { const response: LLMResponse = { content: '' @@ -1153,7 +1188,8 @@ export class AiSdkProvider extends BaseLLMProvider { resolvedModelConfig, temperature ?? resolvedModelConfig.temperature ?? 0.7, maxTokens ?? resolvedModelConfig.maxTokens ?? 1024, - tools + tools, + signal )) { switch (event.type) { case 'text': @@ -2464,7 +2500,8 @@ export class AiSdkProvider extends BaseLLMProvider { modelId: string, temperature?: number, maxTokens?: number, - systemPrompt?: string + systemPrompt?: string, + signal?: AbortSignal ): Promise { return this.runText( [ @@ -2473,7 +2510,9 @@ export class AiSdkProvider extends BaseLLMProvider { ], modelId, temperature, - maxTokens + maxTokens, + undefined, + signal ) } @@ -2668,24 +2707,53 @@ export class AiSdkProvider extends BaseLLMProvider { } } + public async generateText( + prompt: string, + modelId: string, + temperature?: number, + maxTokens?: number, + options?: AiSdkGenerateTextOptions + ): Promise + /** @deprecated Pass `{ systemPrompt }` as the fifth argument instead. */ public async generateText( prompt: string, modelId: string, temperature?: number, maxTokens?: number, systemPrompt?: string + ): Promise + public async generateText( + prompt: string, + modelId: string, + temperature?: number, + maxTokens?: number, + optionsOrSystemPrompt?: AiSdkGenerateTextOptions | string ): Promise { + const options: AiSdkGenerateTextOptions = + typeof optionsOrSystemPrompt === 'string' + ? { systemPrompt: optionsOrSystemPrompt } + : (optionsOrSystemPrompt ?? {}) const decision = this.resolveRouteDecision(modelId) if (decision.endpointType === 'grok-image' || decision.endpointType === 'image-generation') { return this.collectStreamResponse( [{ role: 'user', content: prompt }], modelId, temperature, - maxTokens + maxTokens, + [], + undefined, + options?.signal ) } - return this.runPromptCompletion(prompt, modelId, temperature, maxTokens, systemPrompt) + return this.runPromptCompletion( + prompt, + modelId, + temperature, + maxTokens, + options?.systemPrompt, + options?.signal + ) } public async suggestions( diff --git a/src/main/presenter/llmProviderPresenter/providers/githubCopilotProvider.ts b/src/main/presenter/llmProviderPresenter/providers/githubCopilotProvider.ts index 9772451e0b..e099defcbf 100644 --- a/src/main/presenter/llmProviderPresenter/providers/githubCopilotProvider.ts +++ b/src/main/presenter/llmProviderPresenter/providers/githubCopilotProvider.ts @@ -9,7 +9,11 @@ import { MCPToolDefinition, IConfigPresenter } from '@shared/presenter' -import { BaseLLMProvider, SUMMARY_TITLES_PROMPT } from '../baseProvider' +import { + BaseLLMProvider, + SUMMARY_TITLES_PROMPT, + type ProviderGenerateTextOptions +} from '../baseProvider' import { HttpsProxyAgent } from 'https-proxy-agent' import { getGlobalGitHubCopilotDeviceFlow, @@ -82,11 +86,13 @@ export class GithubCopilotProvider extends BaseLLMProvider { } private async getCopilotToken(signal?: AbortSignal): Promise { + signal?.throwIfAborted() // 优先使用设备流获取 token if (this.deviceFlow) { try { - return await this.deviceFlow.getCopilotToken() + return await this.deviceFlow.getCopilotToken(signal) } catch (error) { + signal?.throwIfAborted() console.warn( '[GitHub Copilot] Device flow failed, falling back to provider API key:', error @@ -155,6 +161,7 @@ export class GithubCopilotProvider extends BaseLLMProvider { return this.copilotToken } catch (error) { + signal?.throwIfAborted() console.error('[GitHub Copilot] Error getting Copilot token:', error) throw error } @@ -578,9 +585,7 @@ export class GithubCopilotProvider extends BaseLLMProvider { reader.releaseLock() } } catch (error) { - if (signal?.aborted && signal.reason instanceof Error) { - throw signal.reason - } + signal?.throwIfAborted() console.error('GitHub Copilot stream error:', error) throw error } finally { @@ -592,11 +597,12 @@ export class GithubCopilotProvider extends BaseLLMProvider { messages: ChatMessage[], modelId: string, temperature?: number, - _maxTokens?: number + _maxTokens?: number, + callerSignal?: AbortSignal ): Promise { if (!modelId) throw new Error('Model ID is required') const modelConfig = this.configPresenter.getModelConfig(modelId, this.provider.id) - const { signal, dispose } = this.createModelRequestSignal(modelConfig) + const { signal, dispose } = this.createModelRequestSignal(modelConfig, callerSignal) try { const token = await this.getCopilotToken(signal) const formattedMessages = this.formatMessages(messages) @@ -696,9 +702,7 @@ export class GithubCopilotProvider extends BaseLLMProvider { return result } catch (error) { - if (signal?.aborted && signal.reason instanceof Error) { - throw signal.reason - } + signal?.throwIfAborted() console.error('GitHub Copilot completion error:', error) throw error } finally { @@ -730,7 +734,8 @@ export class GithubCopilotProvider extends BaseLLMProvider { prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: ProviderGenerateTextOptions ): Promise { return this.completions( [ @@ -741,7 +746,8 @@ export class GithubCopilotProvider extends BaseLLMProvider { ], modelId, temperature, - maxTokens + maxTokens, + options?.signal ) } diff --git a/src/main/presenter/llmProviderPresenter/providers/ollamaProvider.ts b/src/main/presenter/llmProviderPresenter/providers/ollamaProvider.ts index d15790df0f..205bd51793 100644 --- a/src/main/presenter/llmProviderPresenter/providers/ollamaProvider.ts +++ b/src/main/presenter/llmProviderPresenter/providers/ollamaProvider.ts @@ -13,7 +13,11 @@ import { } from '@shared/presenter' import { ModelType } from '@shared/model' import { DEFAULT_MODEL_CONTEXT_LENGTH, DEFAULT_MODEL_MAX_TOKENS } from '@shared/modelConfigDefaults' -import { BaseLLMProvider, SUMMARY_TITLES_PROMPT } from '../baseProvider' +import { + BaseLLMProvider, + SUMMARY_TITLES_PROMPT, + type ProviderGenerateTextOptions +} from '../baseProvider' import { execFile } from 'node:child_process' import { Ollama, ShowResponse } from 'ollama' import { @@ -535,8 +539,20 @@ export class OllamaProvider extends BaseLLMProvider { prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: ProviderGenerateTextOptions ): Promise { + if (options?.signal) { + return runAiSdkGenerateText( + this.getAiSdkRuntimeContext(), + [{ role: 'user', content: prompt }], + modelId, + this.configPresenter.getModelConfig(modelId, this.provider.id), + temperature, + maxTokens, + options.signal + ) + } return runAiSdkGenerateText( this.getAiSdkRuntimeContext(), [{ role: 'user', content: prompt }], diff --git a/src/main/presenter/llmProviderPresenter/providers/voiceAIProvider.ts b/src/main/presenter/llmProviderPresenter/providers/voiceAIProvider.ts index 087da99c3d..8ed3cc5947 100644 --- a/src/main/presenter/llmProviderPresenter/providers/voiceAIProvider.ts +++ b/src/main/presenter/llmProviderPresenter/providers/voiceAIProvider.ts @@ -10,7 +10,7 @@ import { } from '@shared/presenter' import { DEFAULT_MODEL_CONTEXT_LENGTH, DEFAULT_MODEL_MAX_TOKENS } from '@shared/modelConfigDefaults' import { createStreamEvent } from '@shared/types/core/llm-events' -import { BaseLLMProvider } from '../baseProvider' +import { BaseLLMProvider, type ProviderGenerateTextOptions } from '../baseProvider' import { proxyConfig } from '../../proxyConfig' import { ProxyAgent } from 'undici' @@ -144,7 +144,8 @@ export class VoiceAIProvider extends BaseLLMProvider { prompt: string, modelId: string, temperature?: number, - _maxTokens?: number + _maxTokens?: number, + options?: ProviderGenerateTextOptions ): Promise { if (!prompt) { throw new Error('No prompt provided for Voice.ai TTS') @@ -154,7 +155,8 @@ export class VoiceAIProvider extends BaseLLMProvider { prompt, modelId, temperature, - this.configPresenter.getModelConfig(modelId, this.provider.id) + this.configPresenter.getModelConfig(modelId, this.provider.id), + options?.signal ) return { @@ -420,41 +422,43 @@ export class VoiceAIProvider extends BaseLLMProvider { text: string, modelId: string, temperature?: number, - modelConfig?: ModelConfig + modelConfig?: ModelConfig, + callerSignal?: AbortSignal ): Promise<{ audioBase64: string; mimeType: string }> { - const { signal, dispose } = this.createModelRequestSignal(modelConfig) - const config = this.getTtsConfig() - if (!SUPPORTED_LANGUAGES.has(config.language)) { - throw new Error( - `Unsupported language code: ${config.language}. Supported languages: ${Array.from( - SUPPORTED_LANGUAGES - ).join(', ')}` - ) - } - const voiceId = this.resolveVoiceId(modelId) - const requestBody: Record = { - text, - audio_format: config.audioFormat, - model: config.model, - language: config.language, - temperature: typeof temperature === 'number' ? temperature : config.temperature, - top_p: config.topP - } + const { signal, dispose } = this.createModelRequestSignal(modelConfig, callerSignal) + try { + signal?.throwIfAborted() + const config = this.getTtsConfig() + if (!SUPPORTED_LANGUAGES.has(config.language)) { + throw new Error( + `Unsupported language code: ${config.language}. Supported languages: ${Array.from( + SUPPORTED_LANGUAGES + ).join(', ')}` + ) + } + const voiceId = this.resolveVoiceId(modelId) + const requestBody: Record = { + text, + audio_format: config.audioFormat, + model: config.model, + language: config.language, + temperature: typeof temperature === 'number' ? temperature : config.temperature, + top_p: config.topP + } - if (voiceId) { - requestBody['voice_id'] = voiceId - } + if (voiceId) { + requestBody['voice_id'] = voiceId + } - const headers = this.getAuthHeaders() - if (modelConfig) { - await this.emitRequestTrace(modelConfig, { - endpoint: this.buildUrl('/api/v1/tts/speech'), - headers, - body: requestBody - }) - } + const headers = this.getAuthHeaders() + if (modelConfig) { + await this.emitRequestTrace(modelConfig, { + endpoint: this.buildUrl('/api/v1/tts/speech'), + headers, + body: requestBody + }) + } - try { const response = await fetch(this.buildUrl('/api/v1/tts/speech'), { method: 'POST', headers, @@ -484,9 +488,7 @@ export class VoiceAIProvider extends BaseLLMProvider { const buffer = Buffer.from(await response.arrayBuffer()) return { audioBase64: buffer.toString('base64'), mimeType } } catch (error) { - if (signal?.aborted && signal.reason instanceof Error) { - throw signal.reason - } + signal?.throwIfAborted() throw error } finally { dispose() diff --git a/src/main/presenter/memoryProviderBindings.ts b/src/main/presenter/memoryProviderBindings.ts new file mode 100644 index 0000000000..691b9a0bac --- /dev/null +++ b/src/main/presenter/memoryProviderBindings.ts @@ -0,0 +1,32 @@ +import type { ILlmProviderPresenter } from '@shared/presenter' + +import type { MemoryPresenterDeps } from './memoryPresenter/types' + +type MemoryProviderBindings = Pick< + MemoryPresenterDeps, + 'executeWithRateLimit' | 'getEmbeddings' | 'getDimensions' | 'generateText' +> + +type MemoryLlmProviderPort = Pick< + ILlmProviderPresenter, + 'executeWithRateLimit' | 'getEmbeddings' | 'getDimensions' | 'generateText' +> + +export function createMemoryProviderBindings( + llmProviderPresenter: MemoryLlmProviderPort +): MemoryProviderBindings { + return { + executeWithRateLimit: (providerId, options) => + llmProviderPresenter.executeWithRateLimit(providerId, { signal: options.signal }), + getEmbeddings: (providerId, modelId, texts, signal) => + llmProviderPresenter.getEmbeddings(providerId, modelId, texts, signal), + getDimensions: (providerId, modelId, signal) => + llmProviderPresenter.getDimensions(providerId, modelId, signal), + generateText: async (providerId, modelId, prompt, signal) => + ( + await llmProviderPresenter.generateText(providerId, prompt, modelId, 0.2, undefined, { + signal + }) + ).content ?? '' + } +} diff --git a/src/shared/types/presenters/core.presenter.d.ts b/src/shared/types/presenters/core.presenter.d.ts index 2a61292e15..cf796fc4bb 100644 --- a/src/shared/types/presenters/core.presenter.d.ts +++ b/src/shared/types/presenters/core.presenter.d.ts @@ -1105,7 +1105,8 @@ export interface ILlmProviderPresenter { prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: { signal?: AbortSignal } ): Promise<{ content: string }> stopStream(eventId: string): Promise check(providerId: string, modelId?: string): Promise<{ isOk: boolean; errorMsg: string | null }> diff --git a/src/shared/types/presenters/llmprovider.presenter.d.ts b/src/shared/types/presenters/llmprovider.presenter.d.ts index 1016d89614..3eb7169574 100644 --- a/src/shared/types/presenters/llmprovider.presenter.d.ts +++ b/src/shared/types/presenters/llmprovider.presenter.d.ts @@ -263,7 +263,8 @@ export interface ILlmProviderPresenter { prompt: string, modelId: string, temperature?: number, - maxTokens?: number + maxTokens?: number, + options?: { signal?: AbortSignal } ): Promise<{ content: string }> stopStream(eventId: string): Promise check(providerId: string, modelId?: string): Promise<{ isOk: boolean; errorMsg: string | null }> diff --git a/test/main/presenter/acpProvider.test.ts b/test/main/presenter/acpProvider.test.ts index 3267e85f65..96860cef76 100644 --- a/test/main/presenter/acpProvider.test.ts +++ b/test/main/presenter/acpProvider.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, it, expect, vi } from 'vitest' import { AcpProvider } from '../../../src/main/presenter/llmProviderPresenter/providers/acpProvider' import { AcpSessionController, LEGACY_MODE_CONFIG_ID } from '@/agent/acp/runtime' +import { AcpPromptController } from '@/agent/acp/client' import { eventBus } from '@/eventbus' import type { AcpConfigState } from '../../../src/shared/types/presenters' @@ -72,6 +73,59 @@ describe('AcpProvider runDebugAction error handling', () => { }) } } + const createStandaloneProvider = () => { + let resolveFirstPrompt!: (response: { stopReason: string }) => void + const prompt = vi + .fn() + .mockImplementationOnce( + () => + new Promise<{ stopReason: string }>((resolve) => { + resolveFirstPrompt = resolve + }) + ) + .mockResolvedValue({ stopReason: 'end_turn' }) + const cancel = vi.fn().mockResolvedValue(undefined) + const session = { + sessionId: 'standalone-session', + conversationId: 'agent1', + agentId: 'agent1', + promptCapabilities: {}, + systemPromptSent: false, + connection: { prompt, cancel } + } + const open = vi.fn().mockResolvedValue(session) + const clearMappedSession = vi.fn() + const provider = Object.create(AcpProvider.prototype) as any + provider.provider = { id: 'acp', name: 'ACP' } + provider.configPresenter = { + getModelConfig: vi.fn().mockReturnValue({}), + getAcpEnabled: vi.fn().mockResolvedValue(true), + getAcpAgents: vi.fn().mockResolvedValue([agent]) + } + provider.sessionPersistence = { + getWorkdir: vi.fn().mockResolvedValue('/tmp'), + startTurn: vi.fn().mockResolvedValue(undefined), + finishTurn: vi.fn().mockResolvedValue(undefined) + } + provider.acpRuntime = { + sessionController: { open, clearMappedSession } + } + provider.promptController = new AcpPromptController() + provider.messageFormatter = { + format: vi.fn().mockReturnValue({ blocks: [], includedSystemPrompt: false }) + } + provider.pendingPermissions = new Map() + provider.emitRequestTrace = vi.fn().mockResolvedValue(undefined) + + return { + provider, + prompt, + cancel, + open, + clearMappedSession, + resolveFirstPrompt: (response: { stopReason: string }) => resolveFirstPrompt(response) + } + } const createConfigState = (modelValue = 'gpt-5'): AcpConfigState => ({ source: 'configOptions', options: [ @@ -95,6 +149,92 @@ describe('AcpProvider runDebugAction error handling', () => { ] }) + it('keeps the active standalone owner until its cancelled ACP prompt settles', async () => { + const { provider, prompt, cancel, open, resolveFirstPrompt } = createStandaloneProvider() + const controller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + const first = provider.generateText('first', 'agent1', undefined, undefined, { + signal: controller.signal + }) + await vi.waitFor(() => expect(prompt).toHaveBeenCalledOnce()) + + const second = provider.generateText('second', 'agent1') + controller.abort(reason) + const firstAssertion = expect(first).rejects.toBe(reason) + await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce()) + + expect(open).toHaveBeenCalledTimes(1) + expect(prompt).toHaveBeenCalledTimes(1) + + resolveFirstPrompt({ stopReason: 'cancelled' }) + await firstAssertion + await expect(second).resolves.toEqual({ content: '', reasoning_content: '' }) + + expect(open).toHaveBeenCalledTimes(2) + expect(prompt).toHaveBeenCalledTimes(2) + expect(cancel).toHaveBeenCalledTimes(1) + }) + + it('removes an aborted queued standalone request without opening the shared session', async () => { + const { provider, prompt, cancel, open, resolveFirstPrompt } = createStandaloneProvider() + const first = provider.generateText('first', 'agent1') + await vi.waitFor(() => expect(prompt).toHaveBeenCalledOnce()) + const controller = new AbortController() + const reason = new DOMException('Queued request aborted', 'AbortError') + const queued = provider.generateText('queued', 'agent1', undefined, undefined, { + signal: controller.signal + }) + + controller.abort(reason) + + await expect(queued).rejects.toBe(reason) + expect(open).toHaveBeenCalledTimes(1) + expect(prompt).toHaveBeenCalledTimes(1) + expect(cancel).not.toHaveBeenCalled() + + resolveFirstPrompt({ stopReason: 'end_turn' }) + await expect(first).resolves.toEqual({ content: '', reasoning_content: '' }) + }) + + it('rejects a pre-aborted standalone request before opening an ACP session', async () => { + const { provider, prompt, open } = createStandaloneProvider() + const controller = new AbortController() + const reason = new DOMException('Already aborted', 'AbortError') + controller.abort(reason) + + await expect( + provider.generateText('prompt', 'agent1', undefined, undefined, { + signal: controller.signal + }) + ).rejects.toBe(reason) + + expect(open).not.toHaveBeenCalled() + expect(prompt).not.toHaveBeenCalled() + }) + + it('forwards caller cancellation while opening the standalone ACP session', async () => { + const { provider, prompt, open } = createStandaloneProvider() + let openSignal: AbortSignal | undefined + open.mockReset().mockImplementation((_conversationId, _agent, hooks) => { + openSignal = hooks.signal + return new Promise((_resolve, reject) => { + hooks.signal.addEventListener('abort', () => reject(hooks.signal.reason), { once: true }) + }) + }) + const controller = new AbortController() + const reason = new DOMException('Session open aborted', 'AbortError') + + const generating = provider.generateText('prompt', 'agent1', undefined, undefined, { + signal: controller.signal + }) + await vi.waitFor(() => expect(open).toHaveBeenCalledOnce()) + expect(openSignal).toBe(controller.signal) + controller.abort(reason) + + await expect(generating).rejects.toBe(reason) + expect(prompt).not.toHaveBeenCalled() + }) + it('returns error result when process manager is shutting down', async () => { const provider = Object.create(AcpProvider.prototype) as any provider.configPresenter = { @@ -622,6 +762,41 @@ describe('AcpProvider runDebugAction error handling', () => { } }) + it('cancels pending and late permission requests after caller abort', async () => { + const provider = Object.create(AcpProvider.prototype) as any + provider.provider = { id: 'acp', name: 'ACP' } + provider.pendingPermissions = new Map() + const queue = { push: vi.fn() } + const params = { + sessionId: 'session-1', + toolCall: { + toolCallId: 'tc-terminal', + title: 'Terminal', + kind: 'execute', + rawInput: { command: 'dir' } + }, + options: [] + } + const context = { + conversationId: 'conv-1', + agent: { id: 'agent1', name: 'Claude Agent', command: 'claude' } + } + const controller = new AbortController() + + const pending = provider.handlePermissionRequest(queue, params, context, controller.signal) + await vi.waitFor(() => expect(queue.push).toHaveBeenCalledTimes(2)) + controller.abort(new DOMException('Memory request aborted', 'AbortError')) + + await expect(pending).resolves.toEqual({ outcome: { outcome: 'cancelled' } }) + expect(provider.pendingPermissions.size).toBe(0) + + const lateQueue = { push: vi.fn() } + await expect( + provider.handlePermissionRequest(lateQueue, params, context, controller.signal) + ).resolves.toEqual({ outcome: { outcome: 'cancelled' } }) + expect(lateQueue.push).not.toHaveBeenCalled() + }) + it('clears pending permission timeout when resolving a request', async () => { vi.useFakeTimers() @@ -880,6 +1055,7 @@ describe('AcpProvider runDebugAction error handling', () => { try { const provider = Object.create(AcpProvider.prototype) as any + provider.pendingPermissions = new Map() provider.emitRequestTrace = vi.fn().mockResolvedValue(undefined) provider.promptController = { begin: vi.fn().mockReturnValue({ @@ -901,7 +1077,13 @@ describe('AcpProvider runDebugAction error handling', () => { } const cancel = vi.fn().mockResolvedValue(undefined) - const prompt = vi.fn().mockImplementation(() => new Promise(() => {})) + let resolvePrompt!: (response: { stopReason: string }) => void + const prompt = vi.fn().mockImplementation( + () => + new Promise<{ stopReason: string }>((resolve) => { + resolvePrompt = resolve + }) + ) const queue = { push: vi.fn(), done: vi.fn() @@ -922,6 +1104,10 @@ describe('AcpProvider runDebugAction error handling', () => { ) await vi.advanceTimersByTimeAsync(25) + await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce()) + expect(queue.done).not.toHaveBeenCalled() + + resolvePrompt({ stopReason: 'cancelled' }) await runPrompt expect(cancel).toHaveBeenCalledWith({ sessionId: 'session-timeout' }) @@ -935,6 +1121,161 @@ describe('AcpProvider runDebugAction error handling', () => { } }) + it('waits for prompt settlement when caller cancellation throws synchronously', async () => { + const provider = Object.create(AcpProvider.prototype) as any + provider.pendingPermissions = new Map() + provider.emitRequestTrace = vi.fn().mockResolvedValue(undefined) + provider.promptController = { + begin: vi.fn().mockReturnValue({ + id: 'turn-cancelled', + sessionId: 'session-cancelled', + conversationId: 'conv-cancelled', + userMessageId: null, + startedAt: Date.now() + }), + complete: vi.fn(), + cancel: vi.fn().mockReturnValue({ + id: 'turn-cancelled', + completedAt: Date.now() + }), + fail: vi.fn() + } + provider.sessionPersistence = { + startTurn: vi.fn().mockResolvedValue(undefined), + finishTurn: vi.fn().mockResolvedValue(undefined) + } + + const cancelError = new Error('synchronous cancel failure') + const cancel = vi.fn(() => { + throw cancelError + }) + let resolvePrompt!: (response: { stopReason: string }) => void + const prompt = vi.fn().mockImplementation( + () => + new Promise<{ stopReason: string }>((resolve) => { + resolvePrompt = resolve + }) + ) + const queue = { + push: vi.fn(), + done: vi.fn() + } + const controller = new AbortController() + + const runPrompt = provider['runPrompt']( + { + sessionId: 'session-cancelled', + conversationId: 'conv-cancelled', + connection: { + prompt, + cancel + } + }, + [], + queue, + {}, + { signal: controller.signal } + ) + + await vi.waitFor(() => expect(prompt).toHaveBeenCalledOnce()) + controller.abort(new DOMException('Memory request aborted', 'AbortError')) + await vi.waitFor(() => expect(cancel).toHaveBeenCalledOnce()) + expect(provider.promptController.cancel).not.toHaveBeenCalled() + expect(queue.done).not.toHaveBeenCalled() + + resolvePrompt({ stopReason: 'cancelled' }) + await runPrompt + + expect(cancel).toHaveBeenCalledWith({ sessionId: 'session-cancelled' }) + expect(provider.promptController.cancel).toHaveBeenCalledWith('session-cancelled') + expect(provider.promptController.fail).not.toHaveBeenCalled() + expect(provider.sessionPersistence.finishTurn).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'turn-cancelled', + status: 'cancelled', + stopReason: 'cancelled' + }) + ) + expect(queue.push).not.toHaveBeenCalled() + expect(queue.done).toHaveBeenCalledTimes(1) + }) + + it('keeps a prompt failure as the first result when caller abort follows immediately', async () => { + const provider = Object.create(AcpProvider.prototype) as any + provider.pendingPermissions = new Map() + provider.emitRequestTrace = vi.fn().mockResolvedValue(undefined) + provider.promptController = { + begin: vi.fn().mockReturnValue({ + id: 'turn-prompt-error', + sessionId: 'session-prompt-error', + conversationId: 'conv-prompt-error', + userMessageId: null, + startedAt: Date.now() + }), + complete: vi.fn(), + cancel: vi.fn(), + fail: vi.fn().mockReturnValue({ + id: 'turn-prompt-error', + completedAt: Date.now() + }) + } + provider.sessionPersistence = { + startTurn: vi.fn().mockResolvedValue(undefined), + finishTurn: vi.fn().mockResolvedValue(undefined) + } + + let rejectPrompt!: (reason: Error) => void + const prompt = vi.fn( + () => + new Promise<{ stopReason: string }>((_resolve, reject) => { + rejectPrompt = reject + }) + ) + const cancel = vi.fn().mockResolvedValue(undefined) + const queue = { + push: vi.fn(), + done: vi.fn() + } + const controller = new AbortController() + const promptError = new Error('prompt transport failed') + + const runPrompt = provider['runPrompt']( + { + sessionId: 'session-prompt-error', + conversationId: 'conv-prompt-error', + connection: { + prompt, + cancel + } + }, + [], + queue, + {}, + { signal: controller.signal } + ) + + await vi.waitFor(() => expect(prompt).toHaveBeenCalledOnce()) + rejectPrompt(promptError) + controller.abort(new DOMException('late caller abort', 'AbortError')) + await runPrompt + + expect(cancel).not.toHaveBeenCalled() + expect(provider.promptController.cancel).not.toHaveBeenCalled() + expect(provider.promptController.fail).toHaveBeenCalledWith('session-prompt-error') + expect(provider.sessionPersistence.finishTurn).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'turn-prompt-error', + status: 'error', + stopReason: 'error' + }) + ) + expect(queue.push).toHaveBeenCalledWith({ + type: 'error', + error_message: 'ACP: prompt transport failed' + }) + expect(queue.done).toHaveBeenCalledTimes(1) + }) + it('marks the system prompt as sent only after the ACP prompt succeeds', async () => { const provider = Object.create(AcpProvider.prototype) as any provider.emitRequestTrace = vi.fn().mockResolvedValue(undefined) diff --git a/test/main/presenter/agentRuntimePresenter/compactionService.test.ts b/test/main/presenter/agentRuntimePresenter/compactionService.test.ts index cad9efa2e6..e9a93cf36e 100644 --- a/test/main/presenter/agentRuntimePresenter/compactionService.test.ts +++ b/test/main/presenter/agentRuntimePresenter/compactionService.test.ts @@ -801,6 +801,14 @@ describe('CompactionService', () => { abortController.signal ) await vi.waitFor(() => expect(llmProviderPresenter.generateText).toHaveBeenCalled()) + expect(llmProviderPresenter.generateText).toHaveBeenCalledWith( + 'openai', + expect.any(String), + 'gpt-4o', + 0.2, + 512, + { signal: abortController.signal } + ) abortController.abort() resolveSummary({ content: 'late generated summary' }) diff --git a/test/main/presenter/githubCopilotDeviceFlow.test.ts b/test/main/presenter/githubCopilotDeviceFlow.test.ts new file mode 100644 index 0000000000..8ac1811f6d --- /dev/null +++ b/test/main/presenter/githubCopilotDeviceFlow.test.ts @@ -0,0 +1,61 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { GitHubCopilotDeviceFlow } from '@/presenter/githubCopilotDeviceFlow' + +vi.mock('electron', () => ({ + BrowserWindow: vi.fn(), + shell: { openExternal: vi.fn() } +})) + +vi.mock('@/presenter', () => ({ + presenter: {} +})) + +describe('GitHubCopilotDeviceFlow cancellation', () => { + afterEach(() => { + vi.unstubAllGlobals() + }) + + const createAuthenticatedFlow = () => { + const flow = new GitHubCopilotDeviceFlow({ clientId: 'client', scope: 'read:user' }) + ;(flow as unknown as { oauthToken: string }).oauthToken = 'oauth-token' + return flow + } + + it('forwards the exact caller signal to the Copilot token exchange', async () => { + const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { + return new Promise((_resolve, reject) => { + const signal = options?.signal as AbortSignal | undefined + signal?.addEventListener('abort', () => reject(new Error('transport wrapped abort')), { + once: true + }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + const flow = createAuthenticatedFlow() + const controller = new AbortController() + const reason = { source: 'memory-caller' } + + const token = flow.getCopilotToken(controller.signal) + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + controller.abort(reason) + + await expect(token).rejects.toBe(reason) + expect(fetchMock).toHaveBeenCalledWith( + 'https://api.github.com/copilot_internal/v2/token', + expect.objectContaining({ signal: controller.signal }) + ) + }) + + it('does not start a token exchange for a pre-aborted request', async () => { + const fetchMock = vi.fn() + vi.stubGlobal('fetch', fetchMock) + const flow = createAuthenticatedFlow() + const controller = new AbortController() + const reason = new DOMException('Already cancelled', 'AbortError') + controller.abort(reason) + + await expect(flow.getCopilotToken(controller.signal)).rejects.toBe(reason) + expect(fetchMock).not.toHaveBeenCalled() + }) +}) diff --git a/test/main/presenter/llmProviderPresenter.test.ts b/test/main/presenter/llmProviderPresenter.test.ts index 9bc6923117..d711cca2d5 100644 --- a/test/main/presenter/llmProviderPresenter.test.ts +++ b/test/main/presenter/llmProviderPresenter.test.ts @@ -346,6 +346,29 @@ describe('LLMProviderPresenter Integration Tests', () => { console.log('Completion response:', response.substring(0, 100)) }, 15000) + it('forwards text-generation cancellation options to the provider', async () => { + const provider = llmProviderPresenter.getProviderInstance('mock-openai-api') + const generateTextSpy = vi + .spyOn(provider, 'generateText') + .mockResolvedValue({ content: 'generated' }) + const signal = new AbortController().signal + + await expect( + llmProviderPresenter.generateText( + 'mock-openai-api', + 'prompt', + 'mock-gpt-thinking', + 0.2, + 128, + { signal } + ) + ).resolves.toEqual({ content: 'generated' }) + + expect(generateTextSpy).toHaveBeenCalledWith('prompt', 'mock-gpt-thinking', 0.2, 128, { + signal + }) + }) + it('should generate completion standalone', async () => { const messages: ChatMessage[] = [{ role: 'user', content: '1' }] diff --git a/test/main/presenter/llmProviderPresenter/aiSdkRuntime.test.ts b/test/main/presenter/llmProviderPresenter/aiSdkRuntime.test.ts index ab9dc96c9f..cc4c1d1f92 100644 --- a/test/main/presenter/llmProviderPresenter/aiSdkRuntime.test.ts +++ b/test/main/presenter/llmProviderPresenter/aiSdkRuntime.test.ts @@ -327,6 +327,67 @@ describe('AI SDK runtime', () => { ) }) + it('passes the exact caller signal to generateText when no model timeout is configured', async () => { + const controller = new AbortController() + + await runAiSdkGenerateText( + createTextRuntimeContext(), + [{ role: 'user', content: 'Hello' }], + 'gpt-4', + { apiEndpoint: 'chat' } as any, + undefined, + undefined, + controller.signal + ) + + expect(mockGenerateText).toHaveBeenCalledWith( + expect.objectContaining({ abortSignal: controller.signal }) + ) + }) + + it('combines caller cancellation with the configured model timeout', async () => { + const controller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + + await runAiSdkGenerateText( + createTextRuntimeContext(), + [{ role: 'user', content: 'Hello' }], + 'gpt-4', + { apiEndpoint: 'chat', timeout: 60_000 } as any, + undefined, + undefined, + controller.signal + ) + + const request = mockGenerateText.mock.calls[0]?.[0] as { abortSignal?: AbortSignal } + expect(request.abortSignal).not.toBe(controller.signal) + expect(request.abortSignal?.aborted).toBe(false) + + controller.abort(reason) + + expect(request.abortSignal?.aborted).toBe(true) + expect(request.abortSignal?.reason).toBe(reason) + }) + + it('rejects a pre-aborted text request before invoking the AI SDK', async () => { + const controller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + controller.abort(reason) + + await expect( + runAiSdkGenerateText( + createTextRuntimeContext(), + [{ role: 'user', content: 'Hello' }], + 'gpt-4', + { apiEndpoint: 'chat' } as any, + undefined, + undefined, + controller.signal + ) + ).rejects.toBe(reason) + expect(mockGenerateText).not.toHaveBeenCalled() + }) + it('promotes leading system messages to the top-level instructions option for generateText', async () => { await runAiSdkGenerateText( createTextRuntimeContext(), @@ -549,6 +610,35 @@ describe('AI SDK runtime', () => { ]) }) + it('forwards caller cancellation to the image generation transport', async () => { + const context = { + providerKind: 'openai-compatible', + provider: { + id: 'openai', + apiType: 'openai-compatible' + }, + configPresenter: {}, + defaultHeaders: {}, + shouldUseImageGeneration: () => true + } as any + const signal = new AbortController().signal + + for await (const _event of runAiSdkCoreStream( + context, + [{ role: 'user', content: 'draw a cat' }], + 'gpt-image-2', + { apiEndpoint: 'image' } as any, + 0.7, + 1024, + [], + signal + )) { + // Drain stream. + } + + expect(mockGenerateImage).toHaveBeenCalledWith(expect.objectContaining({ abortSignal: signal })) + }) + it('does not forward gpt-image-2 image options when the config is empty', async () => { const context = { providerKind: 'openai-responses', @@ -762,6 +852,87 @@ describe('AI SDK runtime', () => { expect(request).not.toHaveProperty('providerOptions') }) + it.each([ + { + route: 'OpenAI speech', + provider: { + id: 'openai', + apiType: 'openai-compatible', + baseUrl: 'https://example.com/v1', + apiKey: 'test-key' + }, + modelId: 'tts-1', + modelConfig: { apiEndpoint: 'audio-speech' } + }, + { + route: 'chat audio', + provider: { + id: 'xiaomimimo', + apiType: 'openai-compatible', + baseUrl: 'https://example.com/v1', + apiKey: 'test-key' + }, + modelId: 'mimo-v2.5-tts', + modelConfig: { apiEndpoint: 'chat', tts: { responseFormat: 'wav' } } + }, + { + route: 'Gemini generateContent', + provider: { + id: 'aihubmix', + apiType: 'openai-compatible', + baseUrl: 'https://aihubmix.com/v1', + apiKey: 'test-key' + }, + modelId: 'gemini-2.5-flash-preview-tts', + modelConfig: { apiEndpoint: 'audio-speech' } + } + ])('aborts the $route TTS transport and removes its caller listener', async (scenario) => { + const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { + return new Promise((_resolve, reject) => { + const requestSignal = options?.signal as AbortSignal + if (requestSignal.aborted) { + reject(requestSignal.reason) + return + } + requestSignal.addEventListener('abort', () => reject(requestSignal.reason), { once: true }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + const caller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + const removeEventListener = vi.spyOn(caller.signal, 'removeEventListener') + const context = { + providerKind: 'openai-compatible', + provider: scenario.provider, + configPresenter: {}, + defaultHeaders: {}, + shouldUseTts: () => true + } as any + + const stream = (async () => { + for await (const _event of runAiSdkCoreStream( + context, + [{ role: 'user', content: 'hello' }], + scenario.modelId, + scenario.modelConfig as any, + 0.7, + 1024, + [], + caller.signal + )) { + // Drain stream. + } + })() + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + caller.abort(reason) + + await expect(stream).rejects.toBe(reason) + const transportSignal = (fetchMock.mock.calls[0]?.[1] as RequestInit).signal + expect(transportSignal).not.toBe(caller.signal) + expect(transportSignal?.reason).toBe(reason) + expect(removeEventListener).toHaveBeenCalledWith('abort', expect.any(Function)) + }) + it('uses normal chat streaming for non-TTS MiMo Pro models', async () => { const fetchMock = vi.fn() vi.stubGlobal('fetch', fetchMock) @@ -1125,6 +1296,175 @@ describe('AI SDK runtime', () => { ]) }) + it('aborts the active video creation request and disposes request resources', async () => { + vi.useFakeTimers() + const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { + return new Promise((_resolve, reject) => { + const requestSignal = options?.signal as AbortSignal + if (requestSignal.aborted) { + reject(requestSignal.reason) + return + } + requestSignal.addEventListener('abort', () => reject(requestSignal.reason), { once: true }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + const caller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + const removeEventListener = vi.spyOn(caller.signal, 'removeEventListener') + const context = { + providerKind: 'openai-compatible', + provider: { + id: 'aihubmix', + apiType: 'openai-compatible', + baseUrl: 'https://aihubmix.com/v1', + apiKey: 'test-key' + }, + configPresenter: {}, + defaultHeaders: {}, + shouldUseVideoGeneration: () => true + } as any + + const stream = (async () => { + for await (const _event of runAiSdkCoreStream( + context, + [{ role: 'user', content: 'make a video' }], + 'video-model', + { apiEndpoint: 'video', timeout: 60_000 } as any, + 0.7, + 1024, + [], + caller.signal + )) { + // Drain stream. + } + })() + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + caller.abort(reason) + + await expect(stream).rejects.toBe(reason) + const creationSignal = (fetchMock.mock.calls[0]?.[1] as RequestInit).signal + expect(creationSignal?.reason).toBe(reason) + expect(removeEventListener).toHaveBeenCalledWith('abort', expect.any(Function)) + expect(vi.getTimerCount()).toBe(0) + }) + + it('aborts video generation while waiting to poll without starting another request', async () => { + const fetchMock = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ id: 'task-poll', status: 'submitted' }), { + status: 200, + headers: { 'Content-Type': 'application/json' } + }) + ) + vi.stubGlobal('fetch', fetchMock) + const caller = new AbortController() + const reason = { source: 'memory-caller' } + const context = { + providerKind: 'openai-compatible', + provider: { + id: 'aihubmix', + apiType: 'openai-compatible', + baseUrl: 'https://aihubmix.com/v1', + apiKey: 'test-key' + }, + configPresenter: {}, + defaultHeaders: {}, + shouldUseVideoGeneration: () => true + } as any + + const stream = (async () => { + for await (const _event of runAiSdkCoreStream( + context, + [{ role: 'user', content: 'make a video' }], + 'video-model', + { apiEndpoint: 'video' } as any, + 0.7, + 1024, + [], + caller.signal + )) { + // Drain stream. + } + })() + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + caller.abort(reason) + + await expect(stream).rejects.toBe(reason) + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it('aborts the video content download with the original caller reason', async () => { + vi.useFakeTimers() + const fetchMock = vi + .fn() + .mockResolvedValueOnce( + new Response(JSON.stringify({ id: 'task-download', status: 'submitted' }), { + status: 200, + headers: { 'Content-Type': 'application/json' } + }) + ) + .mockResolvedValueOnce( + new Response( + JSON.stringify({ + id: 'task-download', + status: 'completed', + url: 'https://cdn.example.com/video.mp4' + }), + { + status: 200, + headers: { 'Content-Type': 'application/json' } + } + ) + ) + .mockImplementationOnce((_url: string, options?: RequestInit) => { + return new Promise((_resolve, reject) => { + const requestSignal = options?.signal as AbortSignal + requestSignal.addEventListener('abort', () => reject(requestSignal.reason), { + once: true + }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + const caller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + const context = { + providerKind: 'openai-compatible', + provider: { + id: 'aihubmix', + apiType: 'openai-compatible', + baseUrl: 'https://aihubmix.com/v1', + apiKey: 'test-key' + }, + configPresenter: {}, + defaultHeaders: {}, + shouldUseVideoGeneration: () => true + } as any + + const stream = (async () => { + for await (const _event of runAiSdkCoreStream( + context, + [{ role: 'user', content: 'make a video' }], + 'video-model', + { apiEndpoint: 'video' } as any, + 0.7, + 1024, + [], + caller.signal + )) { + // Drain stream. + } + })() + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + await vi.advanceTimersByTimeAsync(3_000) + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3)) + caller.abort(reason) + + await expect(stream).rejects.toBe(reason) + const downloadSignal = (fetchMock.mock.calls[2]?.[1] as RequestInit).signal + expect(downloadSignal?.reason).toBe(reason) + expect(vi.getTimerCount()).toBe(0) + }) + it('does not inject unsupported Seedance duration from prompt text', async () => { vi.useFakeTimers() const videoBytes = Uint8Array.from([0, 1, 2, 3]) diff --git a/test/main/presenter/llmProviderPresenter/anthropicProvider.test.ts b/test/main/presenter/llmProviderPresenter/anthropicProvider.test.ts index 64618bc868..eaca1689b8 100644 --- a/test/main/presenter/llmProviderPresenter/anthropicProvider.test.ts +++ b/test/main/presenter/llmProviderPresenter/anthropicProvider.test.ts @@ -131,8 +131,12 @@ describe('AiSdkProvider anthropic', () => { it('passes system prompts through the AI SDK text path', async () => { const provider = new AiSdkProvider(createProvider(), createConfigPresenter()) ;(provider as any).isInitialized = true + const signal = new AbortController().signal - await provider.generateText('hi', 'claude-sonnet-4-5-20250929', 0.2, 32, 'Real system prompt') + await provider.generateText('hi', 'claude-sonnet-4-5-20250929', 0.2, 32, { + systemPrompt: 'Real system prompt', + signal + }) expect(mockRunAiSdkGenerateText).toHaveBeenLastCalledWith( expect.objectContaining({ @@ -145,6 +149,28 @@ describe('AiSdkProvider anthropic', () => { 'claude-sonnet-4-5-20250929', expect.any(Object), 0.2, + 32, + signal + ) + }) + + it('preserves the legacy positional system prompt argument', async () => { + const provider = new AiSdkProvider(createProvider(), createConfigPresenter()) + ;(provider as any).isInitialized = true + + await provider.generateText('hi', 'claude-sonnet-4-5-20250929', 0.2, 32, 'Legacy system prompt') + + expect(mockRunAiSdkGenerateText).toHaveBeenLastCalledWith( + expect.objectContaining({ + providerKind: 'anthropic' + }), + [ + { role: 'system', content: 'Legacy system prompt' }, + { role: 'user', content: 'hi' } + ], + 'claude-sonnet-4-5-20250929', + expect.any(Object), + 0.2, 32 ) }) diff --git a/test/main/presenter/llmProviderPresenter/baseProvider.test.ts b/test/main/presenter/llmProviderPresenter/baseProvider.test.ts index 4c3977b490..15554cf672 100644 --- a/test/main/presenter/llmProviderPresenter/baseProvider.test.ts +++ b/test/main/presenter/llmProviderPresenter/baseProvider.test.ts @@ -52,6 +52,10 @@ class TestProvider extends BaseLLMProvider { return this.provider } + public createRequestSignal(timeout: number | undefined, signal?: AbortSignal) { + return this.createModelRequestSignal(timeout ? { timeout } : undefined, signal) + } + public onProxyResolved(): void {} public async check(): Promise<{ isOk: boolean; errorMsg: string | null }> { @@ -275,4 +279,91 @@ describe('BaseLLMProvider tool XML conversion', () => { await expect(provider.fetchModels()).rejects.toThrow('model persistence failed') }) + + it('preserves the caller signal identity when no model timeout is configured', () => { + const provider = new TestProvider(configPresenter) + const caller = new AbortController() + const request = provider.createRequestSignal(undefined, caller.signal) + + expect(request.signal).toBe(caller.signal) + request.dispose() + }) + + it('keeps caller cancellation as the first reason and disposes the model timeout', () => { + vi.useFakeTimers() + try { + const provider = new TestProvider(configPresenter) + const caller = new AbortController() + const reason = new DOMException('caller cancelled', 'AbortError') + const request = provider.createRequestSignal(25, caller.signal) + + caller.abort(reason) + + expect(request.signal?.aborted).toBe(true) + expect(request.signal?.reason).toBe(reason) + request.dispose() + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) + + it('uses an already-aborted caller as the combined request reason', () => { + vi.useFakeTimers() + try { + const provider = new TestProvider(configPresenter) + const caller = new AbortController() + const reason = new DOMException('already cancelled', 'AbortError') + caller.abort(reason) + + const request = provider.createRequestSignal(25, caller.signal) + + expect(request.signal?.aborted).toBe(true) + expect(request.signal?.reason).toBe(reason) + request.dispose() + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) + + it('keeps the model timeout as the first reason when the caller aborts later', async () => { + vi.useFakeTimers() + try { + const provider = new TestProvider(configPresenter) + const caller = new AbortController() + const request = provider.createRequestSignal(25, caller.signal) + + await vi.advanceTimersByTimeAsync(25) + const timeoutReason = request.signal?.reason + caller.abort(new DOMException('late caller cancellation', 'AbortError')) + + expect(timeoutReason).toMatchObject({ + name: 'AbortError', + message: 'Request timed out after 25ms' + }) + expect(request.signal?.reason).toBe(timeoutReason) + request.dispose() + } finally { + vi.useRealTimers() + } + }) + + it('cleans up timer and caller listener when request signal ownership is disposed', async () => { + vi.useFakeTimers() + try { + const provider = new TestProvider(configPresenter) + const caller = new AbortController() + const request = provider.createRequestSignal(25, caller.signal) + + request.dispose() + caller.abort(new DOMException('disposed caller', 'AbortError')) + await vi.advanceTimersByTimeAsync(25) + + expect(request.signal?.aborted).toBe(false) + expect(vi.getTimerCount()).toBe(0) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/test/main/presenter/llmProviderPresenter/githubCopilotProvider.test.ts b/test/main/presenter/llmProviderPresenter/githubCopilotProvider.test.ts index 14aae39c3c..a4665c00f2 100644 --- a/test/main/presenter/llmProviderPresenter/githubCopilotProvider.test.ts +++ b/test/main/presenter/llmProviderPresenter/githubCopilotProvider.test.ts @@ -29,7 +29,9 @@ describe('GithubCopilotProvider request timeout', () => { const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { return new Promise((_, reject) => { const signal = options?.signal as AbortSignal | undefined - signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + signal?.addEventListener('abort', () => reject(new Error('transport wrapped abort')), { + once: true + }) }) }) vi.stubGlobal('fetch', fetchMock) @@ -63,6 +65,78 @@ describe('GithubCopilotProvider request timeout', () => { ) }) + it('forwards caller cancellation through generateText to the fetch request', async () => { + const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { + return new Promise((_, reject) => { + const signal = options?.signal as AbortSignal | undefined + signal?.addEventListener('abort', () => reject(new Error('transport wrapped abort')), { + once: true + }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + + const provider = Object.create(GithubCopilotProvider.prototype) as GithubCopilotProvider & { + provider: { id: string; name: string } + configPresenter: { getModelConfig: ReturnType } + baseApiUrl: string + getCopilotToken: ReturnType + } + provider.provider = { id: 'github-copilot', name: 'GitHub Copilot' } + provider.configPresenter = { + getModelConfig: vi.fn().mockReturnValue({}) + } + provider.baseApiUrl = 'https://api.githubcopilot.com' + provider.getCopilotToken = vi.fn().mockResolvedValue('token') + const controller = new AbortController() + const reason = { source: 'memory-caller' } + + const completion = provider.generateText('hello', 'gpt-5', undefined, undefined, { + signal: controller.signal + }) + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + controller.abort(reason) + + await expect(completion).rejects.toBe(reason) + expect(provider.getCopilotToken).toHaveBeenCalledWith(controller.signal) + expect(fetchMock).toHaveBeenCalledWith( + 'https://api.githubcopilot.com/chat/completions', + expect.objectContaining({ signal: controller.signal }) + ) + }) + + it('does not fall back to the provider API key after Device Flow is aborted', async () => { + const fetchMock = vi.fn() + vi.stubGlobal('fetch', fetchMock) + const deviceFlow = { + getCopilotToken: vi.fn( + (signal?: AbortSignal) => + new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + ) + } + const provider = Object.create(GithubCopilotProvider.prototype) as any + provider.provider = { + id: 'github-copilot', + name: 'GitHub Copilot', + apiKey: 'fallback-oauth-token' + } + provider.deviceFlow = deviceFlow + provider.copilotToken = null + provider.tokenExpiresAt = 0 + const controller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + + const token = provider.getCopilotToken(controller.signal) + await vi.waitFor(() => expect(deviceFlow.getCopilotToken).toHaveBeenCalledOnce()) + controller.abort(reason) + + await expect(token).rejects.toBe(reason) + expect(deviceFlow.getCopilotToken).toHaveBeenCalledWith(controller.signal) + expect(fetchMock).not.toHaveBeenCalled() + }) + it('aborts streamed requests when the model timeout elapses', async () => { vi.useFakeTimers() diff --git a/test/main/presenter/llmProviderPresenter/ollamaProviderCancellation.test.ts b/test/main/presenter/llmProviderPresenter/ollamaProviderCancellation.test.ts new file mode 100644 index 0000000000..e9b11ae967 --- /dev/null +++ b/test/main/presenter/llmProviderPresenter/ollamaProviderCancellation.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it, vi } from 'vitest' + +const { mockRunAiSdkGenerateText } = vi.hoisted(() => ({ + mockRunAiSdkGenerateText: vi.fn().mockResolvedValue({ content: 'generated' }) +})) + +vi.mock('@/presenter/llmProviderPresenter/aiSdk', () => ({ + runAiSdkCoreStream: vi.fn(), + runAiSdkDimensions: vi.fn(), + runAiSdkEmbeddings: vi.fn(), + runAiSdkGenerateText: mockRunAiSdkGenerateText +})) + +vi.mock('ollama', () => ({ + Ollama: class MockOllama {} +})) + +import { OllamaProvider } from '@/presenter/llmProviderPresenter/providers/ollamaProvider' + +describe('OllamaProvider text cancellation', () => { + it('forwards the caller signal to the shared AI SDK runtime', async () => { + const provider = Object.create(OllamaProvider.prototype) as any + provider.provider = { id: 'ollama' } + provider.configPresenter = { + getModelConfig: vi.fn().mockReturnValue({ apiEndpoint: 'chat' }) + } + provider.getAiSdkRuntimeContext = vi.fn().mockReturnValue({ providerKind: 'openai-compatible' }) + const signal = new AbortController().signal + + await provider.generateText('prompt', 'llama3', 0.2, 128, { signal }) + + expect(mockRunAiSdkGenerateText).toHaveBeenCalledWith( + { providerKind: 'openai-compatible' }, + [{ role: 'user', content: 'prompt' }], + 'llama3', + { apiEndpoint: 'chat' }, + 0.2, + 128, + signal + ) + }) +}) diff --git a/test/main/presenter/llmProviderPresenter/voiceAIProvider.test.ts b/test/main/presenter/llmProviderPresenter/voiceAIProvider.test.ts new file mode 100644 index 0000000000..cdf8f6440c --- /dev/null +++ b/test/main/presenter/llmProviderPresenter/voiceAIProvider.test.ts @@ -0,0 +1,107 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' + +import { VoiceAIProvider } from '@/presenter/llmProviderPresenter/providers/voiceAIProvider' + +vi.mock('@/presenter/proxyConfig', () => ({ + proxyConfig: { + getProxyUrl: vi.fn().mockReturnValue(null) + } +})) + +describe('VoiceAIProvider text cancellation', () => { + afterEach(() => { + vi.useRealTimers() + vi.unstubAllGlobals() + }) + + it('forwards caller cancellation to the speech fetch request', async () => { + const fetchMock = vi.fn().mockImplementation((_url: string, options?: RequestInit) => { + return new Promise((_, reject) => { + const signal = options?.signal as AbortSignal | undefined + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + }) + vi.stubGlobal('fetch', fetchMock) + + const provider = Object.create(VoiceAIProvider.prototype) as VoiceAIProvider & { + provider: { id: string; name: string; apiKey: string; baseUrl: string } + configPresenter: { + getModelConfig: ReturnType + getSetting: ReturnType + } + defaultHeaders: Record + } + provider.provider = { + id: 'voiceai', + name: 'VoiceAI', + apiKey: 'test-key', + baseUrl: 'https://voice.example.com' + } + provider.configPresenter = { + getModelConfig: vi.fn().mockReturnValue({}), + getSetting: vi.fn().mockReturnValue(undefined) + } + provider.defaultHeaders = {} + const controller = new AbortController() + const reason = new DOMException('Memory request aborted', 'AbortError') + + const generating = provider.generateText('hello', 'default', undefined, undefined, { + signal: controller.signal + }) + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledOnce()) + controller.abort(reason) + + await expect(generating).rejects.toBe(reason) + expect(fetchMock).toHaveBeenCalledWith( + 'https://voice.example.com/api/v1/tts/speech', + expect.objectContaining({ signal: controller.signal }) + ) + }) + + it.each([ + { + name: 'invalid language', + apiKey: 'test-key', + language: 'invalid-language', + expectedError: 'Unsupported language code' + }, + { + name: 'missing API key', + apiKey: '', + language: undefined, + expectedError: 'API key is required' + } + ])('disposes timeout and caller listeners after $name validation', async (scenario) => { + vi.useFakeTimers() + const fetchMock = vi.fn() + vi.stubGlobal('fetch', fetchMock) + const provider = Object.create(VoiceAIProvider.prototype) as any + provider.provider = { + id: 'voiceai', + name: 'VoiceAI', + apiKey: scenario.apiKey, + baseUrl: 'https://voice.example.com' + } + provider.configPresenter = { + getModelConfig: vi.fn().mockReturnValue({ timeout: 25 }), + getSetting: vi.fn((key: string) => + key === 'voiceAI_language' ? scenario.language : undefined + ) + } + provider.defaultHeaders = {} + const controller = new AbortController() + const addEventListener = vi.spyOn(controller.signal, 'addEventListener') + const removeEventListener = vi.spyOn(controller.signal, 'removeEventListener') + + await expect( + provider.generateText('hello', 'default', undefined, undefined, { + signal: controller.signal + }) + ).rejects.toThrow(scenario.expectedError) + + expect(fetchMock).not.toHaveBeenCalled() + expect(addEventListener).toHaveBeenCalledWith('abort', expect.any(Function), { once: true }) + expect(removeEventListener).toHaveBeenCalledWith('abort', expect.any(Function)) + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/test/main/presenter/memoryProviderBindings.test.ts b/test/main/presenter/memoryProviderBindings.test.ts new file mode 100644 index 0000000000..cad35c53c3 --- /dev/null +++ b/test/main/presenter/memoryProviderBindings.test.ts @@ -0,0 +1,54 @@ +import { describe, expect, it, vi } from 'vitest' + +import { createMemoryProviderBindings } from '@/presenter/memoryProviderBindings' +import type { ILlmProviderPresenter } from '@shared/presenter' + +type MemoryLlmProviderPort = Pick< + ILlmProviderPresenter, + 'executeWithRateLimit' | 'getEmbeddings' | 'getDimensions' | 'generateText' +> + +function createLlmProviderPort(): MemoryLlmProviderPort { + return { + executeWithRateLimit: vi.fn(async () => undefined), + getEmbeddings: vi.fn(async () => [[1, 2, 3]]), + getDimensions: vi.fn(async () => ({ data: { dimensions: 3 } })), + generateText: vi.fn(async () => ({ content: 'generated' })) + } +} + +describe('createMemoryProviderBindings', () => { + it('forwards the exact caller signal across every provider binding', async () => { + const llmProviderPresenter = createLlmProviderPort() + const bindings = createMemoryProviderBindings(llmProviderPresenter) + const signal = new AbortController().signal + + await bindings.executeWithRateLimit('provider', { signal, purpose: 'decision' }) + await bindings.getEmbeddings('provider', 'embedding-model', ['value'], signal) + await bindings.getDimensions('provider', 'embedding-model', signal) + await expect(bindings.generateText('provider', 'text-model', 'prompt', signal)).resolves.toBe( + 'generated' + ) + + expect(llmProviderPresenter.executeWithRateLimit).toHaveBeenCalledWith('provider', { signal }) + expect(llmProviderPresenter.getEmbeddings).toHaveBeenCalledWith( + 'provider', + 'embedding-model', + ['value'], + signal + ) + expect(llmProviderPresenter.getDimensions).toHaveBeenCalledWith( + 'provider', + 'embedding-model', + signal + ) + expect(llmProviderPresenter.generateText).toHaveBeenCalledWith( + 'provider', + 'prompt', + 'text-model', + 0.2, + undefined, + { signal } + ) + }) +}) diff --git a/test/main/presenter/memoryProviderGateway.test.ts b/test/main/presenter/memoryProviderGateway.test.ts index 4b177947bd..8a4dc16bbb 100644 --- a/test/main/presenter/memoryProviderGateway.test.ts +++ b/test/main/presenter/memoryProviderGateway.test.ts @@ -269,6 +269,89 @@ describe('MemoryProviderGateway', () => { await Promise.allSettled([first, second]) }) + it('restores capacity after signal-aware provider calls settle on agent abort', async () => { + const diagnostics = { + recordProviderAdmissionDecision: vi.fn(), + recordProviderRaceEvent: vi.fn() + } + let callCount = 0 + const { gateway } = makeGateway({ + diagnostics, + generateText: vi.fn((_providerId, _modelId, _prompt, signal) => { + callCount += 1 + if (callCount > 2) return Promise.resolve('after-abort') + + return new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + }) + }) + const first = gateway.generateText('agent', 'p', 'm', 'one', 'decision') + const second = gateway.generateText('agent', 'p', 'm', 'two', 'decision') + await vi.waitFor(() => expect(callCount).toBe(2)) + + gateway.abortAgent('agent') + await Promise.allSettled([first, second]) + await vi.waitFor(() => { + expect(diagnostics.recordProviderRaceEvent).toHaveBeenCalledTimes(4) + expect(diagnostics.recordProviderRaceEvent).toHaveBeenCalledWith('aborted') + expect(diagnostics.recordProviderRaceEvent).toHaveBeenCalledWith('lateSettled') + }) + + await expect(gateway.generateText('agent', 'p', 'm', 'after', 'decision')).resolves.toBe( + 'after-abort' + ) + }) + + it('retains capacity after abort until signal-aware provider calls actually settle', async () => { + const releases: Array<() => void> = [] + let callCount = 0 + const diagnostics = { + recordProviderAdmissionDecision: vi.fn(), + recordProviderRaceEvent: vi.fn() + } + const { gateway, deps } = makeGateway({ + diagnostics, + generateText: vi.fn((_providerId, _modelId, _prompt, signal) => { + callCount += 1 + if (callCount > 2) return Promise.resolve('after-settle') + + return new Promise((_resolve, reject) => { + signal?.addEventListener( + 'abort', + () => { + releases.push(() => reject(signal.reason)) + }, + { once: true } + ) + }) + }) + }) + const first = gateway.generateText('agent', 'p', 'm', 'one', 'decision') + const second = gateway.generateText('agent', 'p', 'm', 'two', 'decision') + await vi.waitFor(() => expect(callCount).toBe(2)) + + gateway.abortAgent('agent') + await Promise.allSettled([first, second]) + await vi.waitFor(() => expect(releases).toHaveLength(2)) + + await expect( + gateway.generateText('agent', 'p', 'm', 'still-full', 'decision') + ).rejects.toMatchObject({ name: 'AbortError', code: MEMORY_PROVIDER_CAPACITY_CODE }) + expect(deps.generateText).toHaveBeenCalledTimes(2) + + releases.forEach((release) => release()) + await vi.waitFor(() => { + expect( + diagnostics.recordProviderRaceEvent.mock.calls.filter(([event]) => event === 'lateSettled') + ).toHaveLength(2) + }) + + await expect(gateway.generateText('agent', 'p', 'm', 'after', 'decision')).resolves.toBe( + 'after-settle' + ) + }) + it('caps unsettled provider calls globally', async () => { const { gateway, deps } = makeGateway({ generateText: vi.fn(() => new Promise(() => undefined)) From a02e23f165f05dbc3cf181b2600d909cf9b6c1ad Mon Sep 17 00:00:00 2001 From: yyhhyyyyyy Date: Wed, 15 Jul 2026 15:45:34 +0800 Subject: [PATCH 2/2] test(memory): fix provider binding test scope --- test/main/presenter/memoryProviderBindings.test.ts | 2 +- test/memory-test-scope.json | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/test/main/presenter/memoryProviderBindings.test.ts b/test/main/presenter/memoryProviderBindings.test.ts index cad35c53c3..59ea68345a 100644 --- a/test/main/presenter/memoryProviderBindings.test.ts +++ b/test/main/presenter/memoryProviderBindings.test.ts @@ -12,7 +12,7 @@ function createLlmProviderPort(): MemoryLlmProviderPort { return { executeWithRateLimit: vi.fn(async () => undefined), getEmbeddings: vi.fn(async () => [[1, 2, 3]]), - getDimensions: vi.fn(async () => ({ data: { dimensions: 3 } })), + getDimensions: vi.fn(async () => ({ data: { dimensions: 3, normalized: false } })), generateText: vi.fn(async () => ({ content: 'generated' })) } } diff --git a/test/memory-test-scope.json b/test/memory-test-scope.json index a91e945bcc..024a656054 100644 --- a/test/memory-test-scope.json +++ b/test/memory-test-scope.json @@ -35,6 +35,7 @@ "test/main/presenter/memoryMaintenanceBudget.test.ts", "test/main/presenter/memoryPagination.test.ts", "test/main/presenter/memoryPresenter.test.ts", + "test/main/presenter/memoryProviderBindings.test.ts", "test/main/presenter/memoryProviderGateway.test.ts", "test/main/presenter/memoryRuntimeContext.test.ts", "test/main/presenter/memorySearch.test.ts",