diff --git a/src/main/services/provider-background-service.ts b/src/main/services/provider-background-service.ts index 56c5b7a3..c495fc88 100644 --- a/src/main/services/provider-background-service.ts +++ b/src/main/services/provider-background-service.ts @@ -152,7 +152,7 @@ export class ProviderBackgroundService { if (result.status === "completed") { const responseText = requireBackgroundText(result.text, provider.provider_key); this.requestStore.clearPollFailures(request.id); - this.requestStore.markSucceeded({ + const finished = this.requestStore.markSucceeded({ orgId: request.org_id, requestId: request.id, responseText, @@ -161,6 +161,7 @@ export class ProviderBackgroundService { externalResponseId: result.responseId, completedAt: new Date().toISOString(), }); + if (!finished) return; this.responseObserver.recordSuccess({ orgId: request.org_id, provider, diff --git a/src/main/services/provider-request-store.ts b/src/main/services/provider-request-store.ts index e91b1ef7..310a0d63 100644 --- a/src/main/services/provider-request-store.ts +++ b/src/main/services/provider-request-store.ts @@ -124,12 +124,13 @@ export class ProviderRequestStore { return requestId; } - markStreaming(input: ProviderRequestProgressInput): void { - this.db.prepare(` + /** Terminal states (succeeded, failed, cancelled) are final: these transitions return false once one is reached. */ + markStreaming(input: ProviderRequestProgressInput): boolean { + return this.db.prepare(` update provider_requests set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, stream_event_count = ?, external_response_id = coalesce(?, external_response_id) - where id = ? and org_id = ? + where id = ? and org_id = ? and state not in ('succeeded', 'failed', 'cancelled') `).run( PROVIDER_REQUEST_STATE.STREAMING, input.responseText, @@ -139,21 +140,21 @@ export class ProviderRequestStore { input.externalResponseId, input.requestId, input.orgId, - ); + ).changes === 1; } - markBackgroundAccepted(input: ProviderBackgroundAcceptedInput): void { - this.db.prepare(` + markBackgroundAccepted(input: ProviderBackgroundAcceptedInput): boolean { + return this.db.prepare(` update provider_requests set state = ?, raw_response_json = ?, external_response_id = ? - where id = ? and org_id = ? + where id = ? and org_id = ? and state not in ('succeeded', 'failed', 'cancelled') `).run( PROVIDER_REQUEST_STATE.PENDING, input.rawResponseJson, input.externalResponseId, input.requestId, input.orgId, - ); + ).changes === 1; } /** Counts a failed status check on a background request and returns how many happened in a row. */ @@ -166,20 +167,20 @@ export class ProviderRequestStore { this.db.prepare("update provider_requests set poll_failures = 0, last_poll_error = null where id = ?").run(requestId); } - markBackgroundProgress(input: ProviderBackgroundProgressInput): void { - this.db.prepare(` + markBackgroundProgress(input: ProviderBackgroundProgressInput): boolean { + return this.db.prepare(` update provider_requests set state = ?, raw_response_json = ?, external_response_id = ? - where id = ? and org_id = ? - `).run(input.state, input.rawResponseJson, input.externalResponseId, input.requestId, input.orgId); + where id = ? and org_id = ? and state not in ('succeeded', 'failed', 'cancelled') + `).run(input.state, input.rawResponseJson, input.externalResponseId, input.requestId, input.orgId).changes === 1; } - markSucceeded(input: ProviderRequestResultInput): void { - this.db.prepare(` + markSucceeded(input: ProviderRequestResultInput): boolean { + return this.db.prepare(` update provider_requests set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, external_response_id = coalesce(?, external_response_id), completed_at = ? - where id = ? and org_id = ? + where id = ? and org_id = ? and state not in ('succeeded', 'failed', 'cancelled') `).run( PROVIDER_REQUEST_STATE.SUCCEEDED, input.responseText, @@ -189,15 +190,15 @@ export class ProviderRequestStore { input.completedAt, input.requestId, input.orgId, - ); + ).changes === 1; } - markFailed(input: ProviderRequestFailureInput): void { - this.db.prepare(` + markFailed(input: ProviderRequestFailureInput): boolean { + return this.db.prepare(` update provider_requests set state = ?, error_message = ?, raw_response_json = coalesce(?, raw_response_json), completed_at = ? - where id = ? and org_id = ? - `).run(PROVIDER_REQUEST_STATE.FAILED, input.errorMessage, input.rawResponseJson ?? null, input.completedAt, input.requestId, input.orgId); + where id = ? and org_id = ? and state not in ('succeeded', 'failed', 'cancelled') + `).run(PROVIDER_REQUEST_STATE.FAILED, input.errorMessage, input.rawResponseJson ?? null, input.completedAt, input.requestId, input.orgId).changes === 1; } markCancelled(input: ProviderRequestCancellationInput): void { diff --git a/src/main/services/provider-response-observer.ts b/src/main/services/provider-response-observer.ts index 768252df..56c8a493 100644 --- a/src/main/services/provider-response-observer.ts +++ b/src/main/services/provider-response-observer.ts @@ -69,13 +69,15 @@ export class ProviderResponseObserver { recordFailure(input: ProviderResponseFailureInput): void { const message = input.error instanceof Error ? input.error.message : String(input.error); this.recordObservedRateLimit(input.orgId, input.provider, input.provider.model, providerRateLimitStateFromError(input.error)); - this.requestStore.markFailed({ + const failed = this.requestStore.markFailed({ orgId: input.orgId, requestId: input.requestId, errorMessage: message, completedAt: new Date().toISOString(), rawResponseJson: input.rawResponseJson, }); + // A late error on a request already cancelled or finished changes nothing. + if (!failed) return; this.requestStore.recordUsage(input.orgId, input.provider.id, input.inputCharCount, 0); this.recordProviderHealth(input.provider.id, input.orgId, HEALTH_STATE_FAILED, message); this.operations.recordSignal({ diff --git a/src/main/services/provider-response-runner.ts b/src/main/services/provider-response-runner.ts index 2703afdd..00eb9756 100644 --- a/src/main/services/provider-response-runner.ts +++ b/src/main/services/provider-response-runner.ts @@ -187,7 +187,7 @@ export class ProviderResponseRunner { if (event.textDelta) textParts.push(event.textDelta); const responseText = textParts.join(""); const streamEventCount = rawEvents.length; - this.deps.requestStore.markStreaming({ + const stillOpen = this.deps.requestStore.markStreaming({ orgId: input.orgId, requestId, responseText, @@ -196,6 +196,10 @@ export class ProviderResponseRunner { streamEventCount, externalResponseId: providerResponseId, }); + if (!stillOpen) { + console.info(`[provider-stream] request ${requestId} was cancelled; ignoring the rest of the stream`); + return requestId; + } onEvent({ requestId, providerId: execution.provider.id, @@ -210,7 +214,7 @@ export class ProviderResponseRunner { }); } const responseText = requireStreamText(textParts); - this.deps.requestStore.markSucceeded({ + const finished = this.deps.requestStore.markSucceeded({ orgId: input.orgId, requestId, responseText, @@ -219,6 +223,7 @@ export class ProviderResponseRunner { externalResponseId: providerResponseId, completedAt: new Date().toISOString(), }); + if (!finished) return requestId; this.deps.responseObserver.recordSuccess({ orgId: input.orgId, provider: execution.provider, diff --git a/tests/unit/provider-service-streaming.test.ts b/tests/unit/provider-service-streaming.test.ts index 48f913c4..9ac42e0d 100644 --- a/tests/unit/provider-service-streaming.test.ts +++ b/tests/unit/provider-service-streaming.test.ts @@ -182,6 +182,52 @@ describe("provider service streaming", () => { }); }); +describe("cancelled streams", () => { + test("a cancelled stream stays cancelled when later events and the final result arrive", async () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + const emitted: string[] = []; + let cancel = async () => {}; + async function* liveEvents(): AsyncIterable { + yield event("text_delta", "Hel", "delta-1", "resp_live_2"); + await cancel(); + yield event("text_delta", "lo", "delta-2", "resp_live_2"); + yield event("message_stop", null, "completed", "resp_live_2"); + } + const adapters: ProviderAdapterRegistry = { + ...providerAdapters, + openai: { + ...providerAdapters.openai, + parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), + checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }), + streamResponse: async () => ({ events: liveEvents(), rateLimitState: null }), + cancelResponse: async ({ responseId }) => ({ responseId, rawResponseJson: JSON.stringify({ id: responseId, status: "cancelled" }), rateLimitState: null }), + }, + }; + const providers = new ProviderService(connection.db, adapters, testVault); + let orgId = ""; + let openAiProviderId = ""; + connection.transaction(() => { + orgId = createOrganization(organization); + openAiProviderId = providerId(connection, orgId, "openai"); + providers.saveCredential({ orgId, providerId: openAiProviderId, secret: "sk-live-provider", model: "gpt-5.4" }); + }); + cancel = async () => { + const live = connection.db.prepare("select id from provider_requests where state = 'streaming'").get() as { id: string }; + await providers.cancelResponse({ orgId, requestId: live.id }); + }; + const requestId = await providers.streamResponse({ + orgId, providerId: openAiProviderId, prompt: "Draft", instructions: null, executionMode: PROVIDER_EXECUTION_MODE.FOREGROUND, + }, (streamEvent) => emitted.push(streamEvent.responseText)); + + expect(connection.db.prepare("select state, response_text from provider_requests where id = ?").get(requestId)) + .toEqual({ state: PROVIDER_REQUEST_STATE.CANCELLED, response_text: "Hel" }); + expect(emitted).toEqual(["Hel"]); + expect(connection.db.prepare("select count(*) as n from provider_usage_windows where org_id = ? and output_char_count > 0").get(orgId)).toEqual({ n: 0 }); + connection.close(); + }); +}); + async function* streamEvents(events: ProviderStreamEvent[]): AsyncIterable { for (const item of events) yield item; }