From 08364c25716f3654953019874c2e6165bc84b697 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Sat, 23 May 2026 20:20:14 -0400 Subject: [PATCH] Wire provider streaming request lifecycle --- AgentCompany-New/src/main/app/ipc-handlers.ts | 20 ++ .../023-provider-request-runtime.ts | 12 + AgentCompany-New/src/main/database/schema.ts | 2 + .../main/services/provider-request-store.ts | 181 ++++++++++ .../services/provider-response-observer.ts | 112 ++++++ .../main/services/provider-response-utils.ts | 16 + .../src/main/services/provider-service.ts | 337 +++++++++++------- .../main/services/provider-sign-in-window.ts | 43 +++ .../src/main/services/provider-streaming.ts | 10 +- .../src/main/services/sdlc-service.ts | 2 +- AgentCompany-New/src/preload/index.ts | 2 + AgentCompany-New/src/shared/commands.ts | 8 + AgentCompany-New/src/shared/desktop-api.ts | 11 +- AgentCompany-New/src/shared/domain.ts | 25 +- AgentCompany-New/src/shared/ipc.ts | 2 + .../src/shared/provider-request-domain.ts | 46 +++ .../tests/unit/domain-services.test.ts | 2 +- .../unit/provider-service-streaming.test.ts | 196 ++++++++++ 18 files changed, 863 insertions(+), 164 deletions(-) create mode 100644 AgentCompany-New/src/main/database/migrations/023-provider-request-runtime.ts create mode 100644 AgentCompany-New/src/main/services/provider-request-store.ts create mode 100644 AgentCompany-New/src/main/services/provider-response-observer.ts create mode 100644 AgentCompany-New/src/main/services/provider-response-utils.ts create mode 100644 AgentCompany-New/src/main/services/provider-sign-in-window.ts create mode 100644 AgentCompany-New/src/shared/provider-request-domain.ts create mode 100644 AgentCompany-New/tests/unit/provider-service-streaming.test.ts diff --git a/AgentCompany-New/src/main/app/ipc-handlers.ts b/AgentCompany-New/src/main/app/ipc-handlers.ts index 0912d9c7..aa86d6ee 100644 --- a/AgentCompany-New/src/main/app/ipc-handlers.ts +++ b/AgentCompany-New/src/main/app/ipc-handlers.ts @@ -134,6 +134,26 @@ export function registerIpcHandlers(deps: IpcDeps): void { return fail("provider_response_failed", error instanceof Error ? error.message : String(error)); } }); + ipcMain.handle(IPC_CHANNELS.streamProviderResponse, async (_event, input) => { + try { + await providers.streamResponse(input, (streamEvent) => { + deps.mainWindow.webContents.send(IPC_CHANNELS.rendererEvent, { type: "provider-response-stream", event: streamEvent }); + }); + deps.mainWindow.webContents.send(IPC_CHANNELS.rendererEvent, { type: "snapshot-changed" }); + return ok(snapshot.load()); + } catch (error) { + return fail("provider_stream_failed", error instanceof Error ? error.message : String(error)); + } + }); + ipcMain.handle(IPC_CHANNELS.cancelProviderResponse, async (_event, input) => { + try { + await providers.cancelResponse(input); + deps.mainWindow.webContents.send(IPC_CHANNELS.rendererEvent, { type: "snapshot-changed" }); + return ok(snapshot.load()); + } catch (error) { + return fail("provider_cancel_failed", error instanceof Error ? error.message : String(error)); + } + }); ipcMain.handle(IPC_CHANNELS.importSkillDirectory, (_event, input) => mutate(() => skills.importDirectory(input))); ipcMain.handle(IPC_CHANNELS.bindSkill, (_event, input) => mutate(() => skills.bind(input))); ipcMain.handle(IPC_CHANNELS.resolveAgentSkills, (_event, input) => mutate(() => skills.resolveRuntimeSnapshot(input))); diff --git a/AgentCompany-New/src/main/database/migrations/023-provider-request-runtime.ts b/AgentCompany-New/src/main/database/migrations/023-provider-request-runtime.ts new file mode 100644 index 00000000..b9eafe8a --- /dev/null +++ b/AgentCompany-New/src/main/database/migrations/023-provider-request-runtime.ts @@ -0,0 +1,12 @@ +export const providerRequestRuntimeSchema = ` + alter table provider_requests add column external_response_id text; + alter table provider_requests add column stream_event_count integer not null default 0; + alter table provider_requests add column cancelled_at text; + + create index if not exists idx_provider_requests_external_response + on provider_requests(provider_id, external_response_id) + where external_response_id is not null; + + insert into schema_migrations(version, name, applied_at) + values (23, 'provider_request_runtime', datetime('now')); +`; diff --git a/AgentCompany-New/src/main/database/schema.ts b/AgentCompany-New/src/main/database/schema.ts index 229ab8fa..e0e12b71 100644 --- a/AgentCompany-New/src/main/database/schema.ts +++ b/AgentCompany-New/src/main/database/schema.ts @@ -21,6 +21,7 @@ import { auditReplaySchema } from "./migrations/019-audit-replay"; import { rollbackSnapshotsSchema } from "./migrations/020-rollback-snapshots"; import { providerAuthBrokerSchema } from "./migrations/021-provider-auth-broker"; import { providerRateLimitObservationSchema } from "./migrations/022-provider-rate-limit-observations"; +import { providerRequestRuntimeSchema } from "./migrations/023-provider-request-runtime"; interface Migration { version: number; @@ -51,6 +52,7 @@ const MIGRATIONS: Migration[] = [ { version: 20, name: "rollback_snapshots", sql: rollbackSnapshotsSchema }, { version: 21, name: "provider_auth_broker", sql: providerAuthBrokerSchema }, { version: 22, name: "provider_rate_limit_observations", sql: providerRateLimitObservationSchema }, + { version: 23, name: "provider_request_runtime", sql: providerRequestRuntimeSchema }, ]; const SCHEMA_VERSION = MIGRATIONS[MIGRATIONS.length - 1].version; diff --git a/AgentCompany-New/src/main/services/provider-request-store.ts b/AgentCompany-New/src/main/services/provider-request-store.ts new file mode 100644 index 00000000..8269f4a1 --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-request-store.ts @@ -0,0 +1,181 @@ +import type { DatabaseSync } from "node:sqlite"; +import { PROVIDER_REQUEST_STATE, type ProviderExecutionMode } from "../../shared/domain"; +import { createId } from "../../shared/ids"; +import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; + +export interface ProviderRequestOwner { + id: string; + provider_key: string; + data_policy: string; +} + +export interface CreateProviderRequestInput { + orgId: string; + provider: ProviderRequestOwner; + model: string; + promptSha256: string; + inputCharCount: number; + executionMode: ProviderExecutionMode; + createdAt: string; +} + +export interface ProviderRequestProgressInput { + orgId: string; + requestId: string; + responseText: string; + rawResponseJson: string; + outputCharCount: number; + streamEventCount: number; + externalResponseId: string | null; +} + +export interface ProviderRequestResultInput { + orgId: string; + requestId: string; + responseText: string; + rawResponseJson: string; + outputCharCount: number; + externalResponseId: string | null; + completedAt: string; +} + +export interface ProviderRequestFailureInput { + orgId: string; + requestId: string; + errorMessage: string; + completedAt: string; +} + +export interface ProviderRequestCancellationInput { + orgId: string; + requestId: string; + rawResponseJson: string; + externalResponseId: string; + cancelledAt: string; +} + +export interface ProviderCancellationRow { + id: string; + org_id: string; + provider_id: string; + provider_key: string; + model: string; + state: string; + external_response_id: string | null; +} + +export class ProviderRequestStore { + constructor(private readonly db: DatabaseSync) {} + + create(input: CreateProviderRequestInput): string { + const requestId = createId("preq"); + this.db.prepare(` + insert into provider_requests( + id, org_id, provider_id, provider_key, model, prompt_sha256, + response_text, raw_response_json, state, error_message, input_char_count, + output_char_count, execution_mode, policy_decision, external_response_id, + stream_event_count, cancelled_at, created_at, completed_at + ) + values (?, ?, ?, ?, ?, ?, null, null, ?, null, ?, 0, ?, ?, null, 0, null, ?, null) + `).run( + requestId, + input.orgId, + input.provider.id, + input.provider.provider_key, + input.model, + input.promptSha256, + PROVIDER_REQUEST_STATE.PENDING, + input.inputCharCount, + input.executionMode, + `allowed:${input.provider.data_policy}`, + input.createdAt, + ); + return requestId; + } + + markStreaming(input: ProviderRequestProgressInput): void { + 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 = ? + `).run( + PROVIDER_REQUEST_STATE.STREAMING, + input.responseText, + input.rawResponseJson, + input.outputCharCount, + input.streamEventCount, + input.externalResponseId, + input.requestId, + input.orgId, + ); + } + + markSucceeded(input: ProviderRequestResultInput): void { + 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 = ? + `).run( + PROVIDER_REQUEST_STATE.SUCCEEDED, + input.responseText, + input.rawResponseJson, + input.outputCharCount, + input.externalResponseId, + input.completedAt, + input.requestId, + input.orgId, + ); + } + + markFailed(input: ProviderRequestFailureInput): void { + this.db.prepare(` + update provider_requests + set state = ?, error_message = ?, completed_at = ? + where id = ? and org_id = ? + `).run(PROVIDER_REQUEST_STATE.FAILED, input.errorMessage, input.completedAt, input.requestId, input.orgId); + } + + markCancelled(input: ProviderRequestCancellationInput): void { + this.db.prepare(` + update provider_requests + set state = ?, raw_response_json = ?, external_response_id = ?, cancelled_at = ?, completed_at = ?, error_message = null + where id = ? and org_id = ? + `).run( + PROVIDER_REQUEST_STATE.CANCELLED, + input.rawResponseJson, + input.externalResponseId, + input.cancelledAt, + input.cancelledAt, + input.requestId, + input.orgId, + ); + } + + loadCancellationRow(orgId: string, requestId: string): ProviderCancellationRow { + const request = this.db.prepare(` + select id, org_id, provider_id, provider_key, model, state, external_response_id + from provider_requests + where id = ? and org_id = ? + `).get(requestId, orgId) as ProviderCancellationRow | undefined; + if (!request) throw new Error("Provider request was not found."); + return request; + } + + recordUsage(orgId: string, providerId: string, inputCharCount: number, outputCharCount: number): void { + const window = currentProviderUsageWindow(); + const now = new Date().toISOString(); + this.db.prepare(` + insert into provider_usage_windows( + id, org_id, provider_id, window_start, window_end, request_count, + input_char_count, output_char_count, created_at, updated_at + ) values (?, ?, ?, ?, ?, 1, ?, ?, ?, ?) + on conflict(provider_id, window_start) do update set + request_count = request_count + 1, + input_char_count = input_char_count + excluded.input_char_count, + output_char_count = output_char_count + excluded.output_char_count, + updated_at = excluded.updated_at + `).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now); + } +} diff --git a/AgentCompany-New/src/main/services/provider-response-observer.ts b/AgentCompany-New/src/main/services/provider-response-observer.ts new file mode 100644 index 00000000..2e6b3567 --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-response-observer.ts @@ -0,0 +1,112 @@ +import type { DatabaseSync } from "node:sqlite"; +import { SLO_SIGNAL_KEY, type ProviderExecutionMode } from "../../shared/domain"; +import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; +import { AuditService } from "./audit-service"; +import { OperationsService } from "./operations-service"; +import { providerRateLimitStateFromError } from "./provider-adapters"; +import { ProviderRateLimitRegistry } from "./provider-rate-limit-registry"; +import { ProviderRequestStore } from "./provider-request-store"; + +const HEALTH_STATE_HEALTHY = "healthy"; +const HEALTH_STATE_FAILED = "failed"; + +export interface ProviderResponseObserverProvider { + id: string; + provider_key: string; + model: string | null; +} + +export interface ProviderResponseSuccessInput { + orgId: string; + provider: ProviderResponseObserverProvider; + requestId: string; + executionMode: ProviderExecutionMode; + startedAt: number; + action: "response_created" | "response_streamed"; +} + +export interface ProviderResponseFailureInput { + orgId: string; + provider: ProviderResponseObserverProvider; + requestId: string; + executionMode: ProviderExecutionMode; + startedAt: number; + inputCharCount: number; + error: unknown; +} + +export class ProviderResponseObserver { + constructor( + private readonly db: DatabaseSync, + private readonly audit: AuditService, + private readonly operations: OperationsService, + private readonly requestStore: ProviderRequestStore, + private readonly rateLimitRegistry: ProviderRateLimitRegistry, + ) {} + + recordSuccess(input: ProviderResponseSuccessInput): void { + this.recordProviderHealth(input.provider.id, input.orgId, HEALTH_STATE_HEALTHY, null); + this.audit.record({ + orgId: input.orgId, + entityType: "provider_request", + entityId: input.requestId, + actorType: "system", + actorId: null, + action: input.action, + payload: { providerKey: input.provider.provider_key, model: input.provider.model }, + }); + this.operations.recordSignal({ + orgId: input.orgId, + signalKey: SLO_SIGNAL_KEY.PROVIDER_REQUEST_SUCCESS, + source: "provider_gateway", + success: true, + durationMs: Date.now() - input.startedAt, + detail: { requestId: input.requestId, providerKey: input.provider.provider_key, executionMode: input.executionMode }, + }); + } + + 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({ + orgId: input.orgId, + requestId: input.requestId, + errorMessage: message, + completedAt: new Date().toISOString(), + }); + 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({ + orgId: input.orgId, + signalKey: SLO_SIGNAL_KEY.PROVIDER_REQUEST_SUCCESS, + source: "provider_gateway", + success: false, + durationMs: Date.now() - input.startedAt, + detail: { requestId: input.requestId, providerKey: input.provider.provider_key, executionMode: input.executionMode }, + }); + } + + recordObservedRateLimit( + orgId: string, + provider: ProviderResponseObserverProvider, + model: string | null, + state: ProviderObservedRateLimitState | null, + ): void { + if (!state) return; + this.rateLimitRegistry.recordObservation({ + orgId, + providerId: provider.id, + providerKey: provider.provider_key, + model, + state, + }); + } + + private recordProviderHealth(providerId: string, orgId: string, healthState: string, error: string | null): void { + this.db.prepare(` + update providers + set health_state = ?, last_error = ?, last_checked_at = ? + where id = ? and org_id = ? + `).run(healthState, error, new Date().toISOString(), providerId, orgId); + } +} diff --git a/AgentCompany-New/src/main/services/provider-response-utils.ts b/AgentCompany-New/src/main/services/provider-response-utils.ts new file mode 100644 index 00000000..a5a40366 --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-response-utils.ts @@ -0,0 +1,16 @@ +export function readProviderResponseId(rawResponseJson: string): string | null { + const parsed = JSON.parse(rawResponseJson) as unknown; + if (!isRecord(parsed)) return null; + const id = parsed.id; + return typeof id === "string" && id.trim() ? id.trim() : null; +} + +export function requireStreamText(parts: string[]): string { + const text = parts.join(""); + if (!text.trim()) throw new Error("Provider stream returned no text content."); + return text; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/AgentCompany-New/src/main/services/provider-service.ts b/AgentCompany-New/src/main/services/provider-service.ts index 278844ef..0a7e5708 100644 --- a/AgentCompany-New/src/main/services/provider-service.ts +++ b/AgentCompany-New/src/main/services/provider-service.ts @@ -1,17 +1,20 @@ import type { DatabaseSync } from "node:sqlite"; import { createHash } from "node:crypto"; -import { BrowserWindow, session } from "electron"; -import type { CreateProviderResponseInput, RevokeProviderCredentialInput, SaveProviderCredentialInput, UpdateProviderPolicyInput } from "../../shared/commands"; -import { FEATURE_FLAG_KEY, PROVIDER_AUTH_TYPE, PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, SLO_SIGNAL_KEY, type ProviderAuthType, type ProviderExecutionMode } from "../../shared/domain"; +import type { CancelProviderResponseInput, CreateProviderResponseInput, RevokeProviderCredentialInput, SaveProviderCredentialInput, StreamProviderResponseInput, UpdateProviderPolicyInput } from "../../shared/commands"; +import { FEATURE_FLAG_KEY, PROVIDER_AUTH_TYPE, PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE, type ProviderAuthType, type ProviderExecutionMode } from "../../shared/domain"; +import type { ProviderStreamRendererEvent } from "../../shared/provider-request-domain"; import { PROVIDER_CATALOG } from "../../shared/provider-catalog"; import { currentProviderUsageWindow } from "../../shared/provider-runtime-domain"; -import { createId } from "../../shared/ids"; import { AuditService } from "./audit-service"; import { AuthBrokerService, type ProviderSecretVault } from "./auth-broker-service"; import { OperationsService } from "./operations-service"; import { SecretVault } from "./secret-vault"; import { providerAdapters, providerRateLimitStateFromError, type ProviderAdapterRegistry } from "./provider-adapters"; import { ProviderRateLimitRegistry } from "./provider-rate-limit-registry"; +import { ProviderResponseObserver } from "./provider-response-observer"; +import { readProviderResponseId, requireStreamText } from "./provider-response-utils"; +import { ProviderRequestStore } from "./provider-request-store"; +import { launchProviderSignInWindow } from "./provider-sign-in-window"; import type { ProviderObservedRateLimitState } from "../../shared/provider-runtime-domain"; const AUTH_STATE_CONFIGURED = "configured"; @@ -26,6 +29,8 @@ export class ProviderService { private readonly authBroker: AuthBrokerService; private readonly operations: OperationsService; private readonly rateLimitRegistry: ProviderRateLimitRegistry; + private readonly responseObserver: ProviderResponseObserver; + private readonly requestStore: ProviderRequestStore; constructor( private readonly db: DatabaseSync, @@ -36,6 +41,8 @@ export class ProviderService { this.authBroker = new AuthBrokerService(db, vault); this.operations = new OperationsService(db); this.rateLimitRegistry = new ProviderRateLimitRegistry(db); + this.requestStore = new ProviderRequestStore(db); + this.responseObserver = new ProviderResponseObserver(db, this.audit, this.operations, this.requestStore, this.rateLimitRegistry); } saveCredential(input: SaveProviderCredentialInput): void { @@ -135,30 +142,9 @@ export class ProviderService { where id = ? and org_id = ? `).get(providerId, orgId) as { id: string; provider_key: string; label: string; auth_type: string } | undefined; if (!provider) throw new Error("Provider was not found."); - const catalogEntry = PROVIDER_CATALOG.find((entry) => entry.key === provider.provider_key); - if (!catalogEntry) throw new Error(`Provider catalog entry missing for ${provider.provider_key}.`); + const catalogEntry = this.requireCatalogEntry(provider.provider_key); if (!catalogEntry.signInUrl) throw new Error(`${provider.label} does not support browser sign-in.`); - - const partition = `persist:agentcompany-provider-${provider.id}`; - const providerSession = session.fromPartition(partition); - providerSession.setPermissionRequestHandler((_webContents, _permission, callback) => callback(false)); - providerSession.setPermissionCheckHandler(() => false); - - const win = new BrowserWindow({ - width: 1100, - height: 760, - title: provider.label, - webPreferences: { - session: providerSession, - nodeIntegration: false, - contextIsolation: true, - sandbox: true, - webSecurity: true, - }, - }); - win.webContents.setWindowOpenHandler(() => ({ action: "deny" })); - this.attachNavigationPolicy(win, catalogEntry.signInUrl); - win.loadURL(catalogEntry.signInUrl); + launchProviderSignInWindow({ providerId: provider.id, label: provider.label, signInUrl: catalogEntry.signInUrl }); this.audit.record({ orgId, @@ -205,7 +191,165 @@ export class ProviderService { }); } - async createResponse(input: CreateProviderResponseInput): Promise { + async createResponse(input: CreateProviderResponseInput): Promise { + const execution = this.prepareExecution(input); + if (!execution.adapter.createResponse) throw new Error(`Provider response execution is unavailable for ${execution.provider.provider_key}.`); + const requestId = this.createRequestRecord(execution); + const startedAt = Date.now(); + try { + const authProfile = execution.adapter.parseAuthProfile(execution.credential.secret); + const result = await execution.adapter.createResponse({ + authProfile, + model: execution.provider.model, + prompt: execution.prompt, + instructions: execution.instructions, + }); + this.recordObservedRateLimit(input.orgId, execution.provider, execution.provider.model, result.rateLimitState); + this.requestStore.markSucceeded({ + orgId: input.orgId, + requestId, + responseText: result.text, + rawResponseJson: result.rawResponseJson, + outputCharCount: result.text.length, + externalResponseId: readProviderResponseId(result.rawResponseJson), + completedAt: new Date().toISOString(), + }); + this.responseObserver.recordSuccess({ + orgId: input.orgId, + provider: execution.provider, + requestId, + executionMode: execution.executionMode, + startedAt, + action: "response_created", + }); + this.requestStore.recordUsage(input.orgId, execution.provider.id, execution.prompt.length, result.text.length); + return requestId; + } catch (error) { + this.responseObserver.recordFailure({ + orgId: input.orgId, + provider: execution.provider, + requestId, + executionMode: execution.executionMode, + startedAt, + inputCharCount: execution.prompt.length, + error, + }); + throw error; + } + } + + async streamResponse(input: StreamProviderResponseInput, onEvent: (event: ProviderStreamRendererEvent) => void): Promise { + const execution = this.prepareExecution(input); + const requestId = this.createRequestRecord(execution); + const startedAt = Date.now(); + const rawEvents: string[] = []; + const textParts: string[] = []; + let providerResponseId: string | null = null; + try { + const authProfile = execution.adapter.parseAuthProfile(execution.credential.secret); + const result = await execution.adapter.streamResponse({ + authProfile, + model: execution.provider.model, + prompt: execution.prompt, + instructions: execution.instructions, + }); + this.recordObservedRateLimit(input.orgId, execution.provider, execution.provider.model, result.rateLimitState); + for await (const event of result.events) { + if (event.type === "error") throw new Error(event.textDelta ?? "Provider stream failed."); + rawEvents.push(event.rawEventJson); + if (event.providerResponseId) providerResponseId = event.providerResponseId; + if (event.textDelta) textParts.push(event.textDelta); + const responseText = textParts.join(""); + const streamEventCount = rawEvents.length; + this.requestStore.markStreaming({ + orgId: input.orgId, + requestId, + responseText, + rawResponseJson: JSON.stringify(rawEvents), + outputCharCount: responseText.length, + streamEventCount, + externalResponseId: providerResponseId, + }); + onEvent({ + requestId, + providerId: execution.provider.id, + providerKey: execution.provider.provider_key, + model: execution.provider.model, + streamEventType: event.type, + textDelta: event.textDelta, + responseText, + providerResponseId, + rawEventJson: event.rawEventJson, + streamEventCount, + }); + } + const responseText = requireStreamText(textParts); + this.requestStore.markSucceeded({ + orgId: input.orgId, + requestId, + responseText, + rawResponseJson: JSON.stringify(rawEvents), + outputCharCount: responseText.length, + externalResponseId: providerResponseId, + completedAt: new Date().toISOString(), + }); + this.responseObserver.recordSuccess({ + orgId: input.orgId, + provider: execution.provider, + requestId, + executionMode: execution.executionMode, + startedAt, + action: "response_streamed", + }); + this.requestStore.recordUsage(input.orgId, execution.provider.id, execution.prompt.length, responseText.length); + return requestId; + } catch (error) { + this.responseObserver.recordFailure({ + orgId: input.orgId, + provider: execution.provider, + requestId, + executionMode: execution.executionMode, + startedAt, + inputCharCount: execution.prompt.length, + error, + }); + throw error; + } + } + + async cancelResponse(input: CancelProviderResponseInput): Promise { + const request = this.requestStore.loadCancellationRow(input.orgId, input.requestId); + if (request.state !== PROVIDER_REQUEST_STATE.PENDING && request.state !== PROVIDER_REQUEST_STATE.STREAMING) { + throw new Error(`Provider request cannot be cancelled from state ${request.state}.`); + } + const providerResponseId = input.providerResponseId?.trim() || request.external_response_id; + if (!providerResponseId) throw new Error("Provider response id is required."); + const provider = this.loadConfiguredProvider(input.orgId, request.provider_id); + const adapter = this.requireAdapter(provider.provider_key, provider.auth_type); + const credential = this.loadCredential(input.orgId, provider); + const authProfile = adapter.parseAuthProfile(credential.secret); + const result = await adapter.cancelResponse({ authProfile, responseId: providerResponseId }); + this.recordObservedRateLimit(input.orgId, provider, provider.model, result.rateLimitState); + this.requestStore.markCancelled({ + orgId: input.orgId, + requestId: input.requestId, + rawResponseJson: result.rawResponseJson, + externalResponseId: providerResponseId, + cancelledAt: new Date().toISOString(), + }); + this.recordProviderHealth(provider.id, input.orgId, HEALTH_STATE_HEALTHY, null); + this.audit.record({ + orgId: input.orgId, + entityType: "provider_request", + entityId: input.requestId, + actorType: "system", + actorId: null, + action: "response_cancelled", + payload: { providerKey: provider.provider_key, model: provider.model }, + }); + } + + private prepareExecution(input: CreateProviderResponseInput | StreamProviderResponseInput): PreparedProviderExecution { const prompt = input.prompt.trim(); if (!prompt) throw new Error("Prompt is required."); const provider = this.loadConfiguredProvider(input.orgId, input.providerId); @@ -215,87 +359,27 @@ export class ProviderService { throw new Error("Provider response execution requires an API key adapter."); } if (!provider.model) throw new Error("Provider model is required."); - const adapter = this.requireAdapter(provider.provider_key, provider.auth_type); - if (!adapter.createResponse) throw new Error(`Provider response execution is unavailable for ${provider.provider_key}.`); - const credential = this.loadCredential(input.orgId, provider); - - const requestId = createId("preq"); - const now = new Date().toISOString(); - const startedAt = Date.now(); - this.db.prepare(` - insert into provider_requests( - id, org_id, provider_id, provider_key, model, prompt_sha256, - response_text, raw_response_json, state, error_message, input_char_count, - output_char_count, execution_mode, policy_decision, created_at, completed_at - ) - values (?, ?, ?, ?, ?, ?, null, null, ?, null, ?, 0, ?, ?, ?, null) - `).run( - requestId, - input.orgId, - provider.id, - provider.provider_key, - provider.model, - hashPrompt(prompt), - PROVIDER_REQUEST_STATE.PENDING, - prompt.length, + return { + orgId: input.orgId, + prompt, + instructions: input.instructions?.trim() || null, executionMode, - `allowed:${provider.data_policy}`, - now, - ); + provider: { ...provider, model: provider.model }, + adapter: this.requireAdapter(provider.provider_key, provider.auth_type), + credential: this.loadCredential(input.orgId, provider), + }; + } - try { - const authProfile = adapter.parseAuthProfile(credential.secret); - const result = await adapter.createResponse({ - authProfile, - model: provider.model, - prompt, - instructions: input.instructions?.trim() || null, - }); - this.recordObservedRateLimit(input.orgId, provider, provider.model, result.rateLimitState); - this.db.prepare(` - update provider_requests - set state = ?, response_text = ?, raw_response_json = ?, output_char_count = ?, completed_at = ? - where id = ? and org_id = ? - `).run(PROVIDER_REQUEST_STATE.SUCCEEDED, result.text, result.rawResponseJson, result.text.length, new Date().toISOString(), requestId, input.orgId); - this.recordUsage(input.orgId, provider.id, prompt.length, result.text.length); - this.recordProviderHealth(provider.id, input.orgId, HEALTH_STATE_HEALTHY, null); - this.audit.record({ - orgId: input.orgId, - entityType: "provider_request", - entityId: requestId, - actorType: "system", - actorId: null, - action: "response_created", - payload: { providerKey: provider.provider_key, model: provider.model }, - }); - this.operations.recordSignal({ - orgId: input.orgId, - signalKey: SLO_SIGNAL_KEY.PROVIDER_REQUEST_SUCCESS, - source: "provider_gateway", - success: true, - durationMs: Date.now() - startedAt, - detail: { requestId, providerKey: provider.provider_key, executionMode }, - }); - } catch (error) { - const message = error instanceof Error ? error.message : String(error); - this.recordObservedRateLimit(input.orgId, provider, provider.model, providerRateLimitStateFromError(error)); - this.db.prepare(` - update provider_requests - set state = ?, error_message = ?, completed_at = ? - where id = ? and org_id = ? - `).run(PROVIDER_REQUEST_STATE.FAILED, message, new Date().toISOString(), requestId, input.orgId); - this.recordUsage(input.orgId, provider.id, prompt.length, 0); - this.recordProviderHealth(provider.id, input.orgId, HEALTH_STATE_FAILED, message); - this.operations.recordSignal({ - orgId: input.orgId, - signalKey: SLO_SIGNAL_KEY.PROVIDER_REQUEST_SUCCESS, - source: "provider_gateway", - success: false, - durationMs: Date.now() - startedAt, - detail: { requestId, providerKey: provider.provider_key, executionMode }, - }); - throw error; - } + private createRequestRecord(execution: PreparedProviderExecution): string { + return this.requestStore.create({ + orgId: execution.orgId, + provider: execution.provider, + model: execution.provider.model, + promptSha256: hashPrompt(execution.prompt), + inputCharCount: execution.prompt.length, + executionMode: execution.executionMode, + createdAt: new Date().toISOString(), + }); } private enforcePolicy(orgId: string, provider: ProviderRow, executionMode: ProviderExecutionMode): void { @@ -357,19 +441,6 @@ export class ProviderService { return entry; } - private attachNavigationPolicy(win: BrowserWindow, signInUrl: string): void { - const allowedHost = new URL(signInUrl).hostname.toLowerCase(); - const allow = (url: string) => { - const host = new URL(url).hostname.toLowerCase(); - return host === allowedHost || host.endsWith(`.${allowedHost}`); - }; - const blockOffHost = (event: Electron.Event, url: string) => { - if (!allow(url)) event.preventDefault(); - }; - win.webContents.on("will-navigate", blockOffHost); - win.webContents.on("will-redirect", blockOffHost); - } - private recordProviderHealth(providerId: string, orgId: string, healthState: string, error: string | null): void { this.db.prepare(` update providers @@ -378,22 +449,6 @@ export class ProviderService { `).run(healthState, error, new Date().toISOString(), providerId, orgId); } - private recordUsage(orgId: string, providerId: string, inputCharCount: number, outputCharCount: number): void { - const window = currentProviderUsageWindow(); - const now = new Date().toISOString(); - this.db.prepare(` - insert into provider_usage_windows( - id, org_id, provider_id, window_start, window_end, request_count, - input_char_count, output_char_count, created_at, updated_at - ) values (?, ?, ?, ?, ?, 1, ?, ?, ?, ?) - on conflict(provider_id, window_start) do update set - request_count = request_count + 1, - input_char_count = input_char_count + excluded.input_char_count, - output_char_count = output_char_count + excluded.output_char_count, - updated_at = excluded.updated_at - `).run(createId("puw"), orgId, providerId, window.start, window.end, inputCharCount, outputCharCount, now, now); - } - private recordObservedRateLimit( orgId: string, provider: Pick, @@ -411,6 +466,16 @@ export class ProviderService { } } +interface PreparedProviderExecution { + orgId: string; + prompt: string; + instructions: string | null; + executionMode: ProviderExecutionMode; + provider: ProviderRow & { model: string }; + adapter: ProviderAdapterRegistry[string]; + credential: { secret: string }; +} + interface ProviderRow { id: string; provider_key: string; diff --git a/AgentCompany-New/src/main/services/provider-sign-in-window.ts b/AgentCompany-New/src/main/services/provider-sign-in-window.ts new file mode 100644 index 00000000..5a2ccfff --- /dev/null +++ b/AgentCompany-New/src/main/services/provider-sign-in-window.ts @@ -0,0 +1,43 @@ +import { BrowserWindow, session } from "electron"; + +export interface LaunchProviderSignInWindowInput { + providerId: string; + label: string; + signInUrl: string; +} + +export function launchProviderSignInWindow(input: LaunchProviderSignInWindowInput): void { + const partition = `persist:agentcompany-provider-${input.providerId}`; + const providerSession = session.fromPartition(partition); + providerSession.setPermissionRequestHandler((_webContents, _permission, callback) => callback(false)); + providerSession.setPermissionCheckHandler(() => false); + + const win = new BrowserWindow({ + width: 1100, + height: 760, + title: input.label, + webPreferences: { + session: providerSession, + nodeIntegration: false, + contextIsolation: true, + sandbox: true, + webSecurity: true, + }, + }); + win.webContents.setWindowOpenHandler(() => ({ action: "deny" })); + attachNavigationPolicy(win, input.signInUrl); + win.loadURL(input.signInUrl); +} + +function attachNavigationPolicy(win: BrowserWindow, signInUrl: string): void { + const allowedHost = new URL(signInUrl).hostname.toLowerCase(); + const allow = (url: string) => { + const host = new URL(url).hostname.toLowerCase(); + return host === allowedHost || host.endsWith(`.${allowedHost}`); + }; + const blockOffHost = (event: Electron.Event, url: string) => { + if (!allow(url)) event.preventDefault(); + }; + win.webContents.on("will-navigate", blockOffHost); + win.webContents.on("will-redirect", blockOffHost); +} diff --git a/AgentCompany-New/src/main/services/provider-streaming.ts b/AgentCompany-New/src/main/services/provider-streaming.ts index 0922429a..414a70f7 100644 --- a/AgentCompany-New/src/main/services/provider-streaming.ts +++ b/AgentCompany-New/src/main/services/provider-streaming.ts @@ -5,6 +5,7 @@ import { isRecord } from "./provider-http"; export interface ProviderStreamEvent { type: "message_start" | "text_delta" | "message_stop" | "error"; textDelta: string | null; + providerResponseId: string | null; rawEventJson: string; } @@ -119,7 +120,7 @@ export function cancelUnsupported(providerLabel: string): () => Promise= 22, value: schemaVersion }, + { key: "schema_version", passed: schemaVersion >= 23, value: schemaVersion }, ]; } diff --git a/AgentCompany-New/src/preload/index.ts b/AgentCompany-New/src/preload/index.ts index edbfeef1..fb40560f 100644 --- a/AgentCompany-New/src/preload/index.ts +++ b/AgentCompany-New/src/preload/index.ts @@ -29,6 +29,8 @@ const api: DesktopApi = { updateProviderPolicy: (input) => invoke(IPC_CHANNELS.updateProviderPolicy, input), openExternalUrl: (url) => invoke(IPC_CHANNELS.openExternalUrl, url), createProviderResponse: (input) => invoke(IPC_CHANNELS.createProviderResponse, input), + streamProviderResponse: (input) => invoke(IPC_CHANNELS.streamProviderResponse, input), + cancelProviderResponse: (input) => invoke(IPC_CHANNELS.cancelProviderResponse, input), importSkillDirectory: (input) => invoke(IPC_CHANNELS.importSkillDirectory, input), bindSkill: (input) => invoke(IPC_CHANNELS.bindSkill, input), resolveAgentSkills: (input) => invoke(IPC_CHANNELS.resolveAgentSkills, input), diff --git a/AgentCompany-New/src/shared/commands.ts b/AgentCompany-New/src/shared/commands.ts index 89b442fe..bf0dd669 100644 --- a/AgentCompany-New/src/shared/commands.ts +++ b/AgentCompany-New/src/shared/commands.ts @@ -173,6 +173,14 @@ export interface CreateProviderResponseInput { executionMode?: ProviderExecutionMode; } +export type StreamProviderResponseInput = CreateProviderResponseInput; + +export interface CancelProviderResponseInput { + orgId: string; + requestId: string; + providerResponseId?: string; +} + export interface ImportSkillDirectoryInput { orgId: string; directoryPath: string; diff --git a/AgentCompany-New/src/shared/desktop-api.ts b/AgentCompany-New/src/shared/desktop-api.ts index 9d7e9989..20443959 100644 --- a/AgentCompany-New/src/shared/desktop-api.ts +++ b/AgentCompany-New/src/shared/desktop-api.ts @@ -4,6 +4,7 @@ import type { BindSkillInput, BootstrapOnboardingInput, BootstrapOnboardingSummary, + CancelProviderResponseInput, CreateAgentInput, CreateApprovalInput, CreateAuditReplayInput, @@ -43,14 +44,16 @@ import type { SetRollbackSnapshotActiveInput, SetRetentionFreezeStateInput, SetScheduledPostStateInput, + StreamProviderResponseInput, SyncEmployeeIdentityInput, UpdateProviderPolicyInput, UpdateTaskStatusInput, } from "./commands"; +import type { ProviderStreamRendererEvent } from "./provider-request-domain"; -export interface RendererEvent { - type: "snapshot-changed"; -} +export type RendererEvent = + | { type: "snapshot-changed" } + | { type: "provider-response-stream"; event: ProviderStreamRendererEvent }; export interface DesktopApi { loadSnapshot(): Promise>; @@ -75,6 +78,8 @@ export interface DesktopApi { updateProviderPolicy(input: UpdateProviderPolicyInput): Promise>; openExternalUrl(url: string): Promise>; createProviderResponse(input: CreateProviderResponseInput): Promise>; + streamProviderResponse(input: StreamProviderResponseInput): Promise>; + cancelProviderResponse(input: CancelProviderResponseInput): Promise>; importSkillDirectory(input: ImportSkillDirectoryInput): Promise>; bindSkill(input: BindSkillInput): Promise>; resolveAgentSkills(input: ResolveAgentSkillsInput): Promise>; diff --git a/AgentCompany-New/src/shared/domain.ts b/AgentCompany-New/src/shared/domain.ts index 23f74768..0d955985 100644 --- a/AgentCompany-New/src/shared/domain.ts +++ b/AgentCompany-New/src/shared/domain.ts @@ -5,7 +5,8 @@ import type { AccessDecisionRecord, EmployeeIdentityRecord, RoleGrantRecord, Ser import type { ApprovalAssignmentRecord, ApprovalRoutingRuleRecord, HiringPolicyRecord, OrgDepartmentRecord } from "./org-governance-domain"; import type { ChaosDrillRecord, ChaosDrillRunRecord, FeatureFlagRecord, OperationalEventRecord, SloTargetRecord } from "./operations-domain"; import type { ProviderAuthProfileRecord } from "./provider-auth-domain"; -import type { ProviderDataPolicy, ProviderExecutionMode } from "./provider-policy-domain"; +import type { ProviderDataPolicy } from "./provider-policy-domain"; +import type { ProviderRequestRecord } from "./provider-request-domain"; import type { ProviderRateLimitObservationRecord, ProviderRateLimitStateRecord } from "./provider-runtime-domain"; import type { RollbackSnapshotRecord } from "./rollback-domain"; import type { EvalCaseRecord, EvalRunRecord, EvalSuiteRecord, PrBotGateRecord, SdlcReviewAgentRecord, SdlcReviewRunRecord } from "./sdlc-domain"; @@ -17,6 +18,7 @@ export { ROUTE_TO_TYPE, type ApprovalAssignmentRecord, type ApprovalRoutingRuleR export { CHAOS_DRILL_SCENARIO, CHAOS_DRILL_STATE, FEATURE_FLAG_KEY, SLO_SIGNAL_KEY, type ChaosDrillRecord, type ChaosDrillRunRecord, type ChaosDrillScenario, type ChaosDrillState, type FeatureFlagKey, type FeatureFlagRecord, type OperationalEventRecord, type SloSignalKey, type SloTargetRecord } from "./operations-domain"; export type { ProviderAuthProfileRecord } from "./provider-auth-domain"; export { PROVIDER_DATA_POLICY, PROVIDER_EXECUTION_MODE, type ProviderDataPolicy, type ProviderExecutionMode } from "./provider-policy-domain"; +export { PROVIDER_REQUEST_STATE, type ProviderRequestRecord, type ProviderRequestState, type ProviderStreamRendererEvent } from "./provider-request-domain"; export type { ProviderRateLimitObservationRecord, ProviderRateLimitStateRecord } from "./provider-runtime-domain"; export type { RollbackSnapshotRecord } from "./rollback-domain"; export { EVAL_SUITE_KEY, PR_BOT_GATE_KEY, SDLC_REVIEW_AGENT_KEY, SDLC_REVIEW_TARGET_TYPE, SDLC_RUN_STATE, type EvalCaseRecord, type EvalRunRecord, type EvalSuiteKey, type EvalSuiteRecord, type PrBotGateKey, type PrBotGateRecord, type SdlcReviewAgentKey, type SdlcReviewAgentRecord, type SdlcReviewRunRecord, type SdlcReviewTargetType, type SdlcRunState } from "./sdlc-domain"; @@ -27,7 +29,6 @@ export const RISK_LEVEL = { LOW: "low", MEDIUM: "medium", HIGH: "high" } as cons export const TASK_PRIORITY = { LOW: "low", NORMAL: "normal", HIGH: "high" } as const; export const PROVIDER_AUTH_TYPE = { API_KEY: "api_key", OAUTH_PKCE: "oauth_pkce", ACCOUNT_SUBSCRIPTION: "account_subscription", CLOUD: "cloud" } as const; export const SOCIAL_PLATFORM = { X: "x", LINKEDIN: "linkedin", YOUTUBE: "youtube", INSTAGRAM: "instagram" } as const; -export const PROVIDER_REQUEST_STATE = { PENDING: "pending", SUCCEEDED: "succeeded", FAILED: "failed" } as const; export const FILE_SCAN_STATE = { APPROVED: "approved", BLOCKED: "blocked" } as const; export type ValueOf = T[keyof T]; @@ -37,7 +38,6 @@ export type RiskLevel = ValueOf; export type TaskPriority = ValueOf; export type ProviderAuthType = ValueOf; export type SocialPlatform = ValueOf; -export type ProviderRequestState = ValueOf; export type FileScanState = ValueOf; export interface OrganizationRecord { @@ -323,25 +323,6 @@ export interface ProviderCatalogEntry { models: ProviderCatalogModel[]; } -export interface ProviderRequestRecord { - id: string; - orgId: string; - providerId: string; - providerKey: string; - model: string; - promptSha256: string; - responseText: string | null; - rawResponseJson: string | null; - state: ProviderRequestState; - errorMessage: string | null; - inputCharCount: number; - outputCharCount: number; - executionMode: ProviderExecutionMode; - policyDecision: string; - createdAt: string; - completedAt: string | null; -} - export interface SkillDefinitionRecord { id: string; orgId: string; diff --git a/AgentCompany-New/src/shared/ipc.ts b/AgentCompany-New/src/shared/ipc.ts index 01a6624e..31835afb 100644 --- a/AgentCompany-New/src/shared/ipc.ts +++ b/AgentCompany-New/src/shared/ipc.ts @@ -21,6 +21,8 @@ export const IPC_CHANNELS = { updateProviderPolicy: "agentcompany:update-provider-policy", openExternalUrl: "agentcompany:open-external-url", createProviderResponse: "agentcompany:create-provider-response", + streamProviderResponse: "agentcompany:stream-provider-response", + cancelProviderResponse: "agentcompany:cancel-provider-response", importSkillDirectory: "agentcompany:import-skill-directory", bindSkill: "agentcompany:bind-skill", resolveAgentSkills: "agentcompany:resolve-agent-skills", diff --git a/AgentCompany-New/src/shared/provider-request-domain.ts b/AgentCompany-New/src/shared/provider-request-domain.ts new file mode 100644 index 00000000..61eeefeb --- /dev/null +++ b/AgentCompany-New/src/shared/provider-request-domain.ts @@ -0,0 +1,46 @@ +import type { ProviderExecutionMode } from "./provider-policy-domain"; + +export const PROVIDER_REQUEST_STATE = { + PENDING: "pending", + STREAMING: "streaming", + SUCCEEDED: "succeeded", + FAILED: "failed", + CANCELLED: "cancelled", +} as const; + +export type ProviderRequestState = typeof PROVIDER_REQUEST_STATE[keyof typeof PROVIDER_REQUEST_STATE]; + +export interface ProviderRequestRecord { + id: string; + orgId: string; + providerId: string; + providerKey: string; + model: string; + promptSha256: string; + responseText: string | null; + rawResponseJson: string | null; + state: ProviderRequestState; + errorMessage: string | null; + inputCharCount: number; + outputCharCount: number; + executionMode: ProviderExecutionMode; + policyDecision: string; + externalResponseId: string | null; + streamEventCount: number; + cancelledAt: string | null; + createdAt: string; + completedAt: string | null; +} + +export interface ProviderStreamRendererEvent { + requestId: string; + providerId: string; + providerKey: string; + model: string; + streamEventType: "message_start" | "text_delta" | "message_stop" | "error"; + textDelta: string | null; + responseText: string; + providerResponseId: string | null; + rawEventJson: string; + streamEventCount: number; +} diff --git a/AgentCompany-New/tests/unit/domain-services.test.ts b/AgentCompany-New/tests/unit/domain-services.test.ts index 9ff8d5f9..44be863f 100644 --- a/AgentCompany-New/tests/unit/domain-services.test.ts +++ b/AgentCompany-New/tests/unit/domain-services.test.ts @@ -23,7 +23,7 @@ describe("domain services", () => { const connection = createConnection(); const rows = connection.db.prepare("select version from schema_migrations order by version asc").all() as Array<{ version: number }>; - expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]); + expect(rows.map((row) => row.version)).toEqual([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23]); connection.close(); }); diff --git a/AgentCompany-New/tests/unit/provider-service-streaming.test.ts b/AgentCompany-New/tests/unit/provider-service-streaming.test.ts new file mode 100644 index 00000000..f2ac91d9 --- /dev/null +++ b/AgentCompany-New/tests/unit/provider-service-streaming.test.ts @@ -0,0 +1,196 @@ +import { randomUUID } from "node:crypto"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { describe, expect, test } from "vitest"; +import { DatabaseConnection } from "../../src/main/database/connection"; +import { OrganizationService } from "../../src/main/services/organization-service"; +import { providerAdapters, type ProviderAdapterRegistry } from "../../src/main/services/provider-adapters"; +import { ProviderService } from "../../src/main/services/provider-service"; +import type { ProviderStreamEvent } from "../../src/main/services/provider-streaming"; +import { PROVIDER_EXECUTION_MODE, PROVIDER_REQUEST_STATE } from "../../src/shared/domain"; +import type { ProviderStreamRendererEvent } from "../../src/shared/provider-request-domain"; + +function createConnection(): DatabaseConnection { + return new DatabaseConnection(join(tmpdir(), `agentcompany-new-provider-stream-${randomUUID()}.db`)); +} + +function createOrganization(organization: OrganizationService): string { + return organization.createOrganization({ + organizationName: "Acme Operations", + employeeName: "Dana Lee", + mission: "Coordinate agents with auditable approvals.", + goal: "Ship a verified control plane.", + }); +} + +function providerId(connection: DatabaseConnection, orgId: string, key: string): string { + return (connection.db.prepare("select id from providers where org_id = ? and provider_key = ?") + .get(orgId, key) as { id: string }).id; +} + +const testVault = { + encryptSecret: (secret: string) => Buffer.from(secret, "utf8"), + decryptSecret: (ciphertext: Buffer) => ciphertext.toString("utf8"), +}; + +describe("provider service streaming", () => { + test("persists provider stream progress and completion", async () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + const emitted: ProviderStreamRendererEvent[] = []; + 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: streamEvents([ + event("message_start", null, "created", "resp_live_1"), + event("text_delta", "Hel", "delta-1", "resp_live_1"), + event("text_delta", "lo", "delta-2", "resp_live_1"), + event("message_stop", null, "completed", "resp_live_1"), + ]), + rateLimitState: { + requestLimit: 120, + remainingRequests: 118, + resetAt: "2026-05-23T12:00:00.000Z", + retryAfterMs: null, + rawStateJson: JSON.stringify({ "x-ratelimit-remaining-requests": "118" }), + }, + }), + }, + }; + 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" }); + }); + const requestId = await providers.streamResponse({ + orgId, + providerId: openAiProviderId, + prompt: "Draft board memo", + instructions: null, + executionMode: PROVIDER_EXECUTION_MODE.FOREGROUND, + }, (streamEvent) => emitted.push(streamEvent)); + + const request = connection.db.prepare(` + select state, response_text, output_char_count, external_response_id, stream_event_count, raw_response_json + from provider_requests + where id = ? + `).get(requestId) as { + state: string; + response_text: string; + output_char_count: number; + external_response_id: string; + stream_event_count: number; + raw_response_json: string; + }; + const usage = connection.db.prepare(` + select request_count, input_char_count, output_char_count + from provider_usage_windows + where org_id = ? and provider_id = ? + `).get(orgId, openAiProviderId); + + expect(request).toMatchObject({ + state: PROVIDER_REQUEST_STATE.SUCCEEDED, + response_text: "Hello", + output_char_count: 5, + external_response_id: "resp_live_1", + stream_event_count: 4, + }); + expect(JSON.parse(request.raw_response_json)).toEqual(["created", "delta-1", "delta-2", "completed"]); + expect(usage).toEqual({ request_count: 1, input_char_count: 16, output_char_count: 5 }); + expect(emitted.map((item) => [item.streamEventType, item.textDelta, item.responseText])).toEqual([ + ["message_start", null, ""], + ["text_delta", "Hel", "Hel"], + ["text_delta", "lo", "Hello"], + ["message_stop", null, "Hello"], + ]); + expect(connection.db.prepare("select remaining_requests from provider_rate_limit_observations").get()).toEqual({ remaining_requests: 118 }); + connection.close(); + }); + + test("cancels a provider request through the adapter response id", async () => { + const connection = createConnection(); + const organization = new OrganizationService(connection.db); + let cancelledResponseId = ""; + const adapters: ProviderAdapterRegistry = { + ...providerAdapters, + openai: { + ...providerAdapters.openai, + parseAuthProfile: (secret) => ({ type: "api_key", apiKey: secret.trim() }), + checkHealth: async () => ({ checkedAt: new Date().toISOString(), rateLimitState: null }), + cancelResponse: async ({ responseId }) => { + cancelledResponseId = responseId; + return { responseId, rawResponseJson: JSON.stringify({ id: responseId, status: "cancelled" }), rateLimitState: null }; + }, + }, + }; + const providers = new ProviderService(connection.db, adapters, testVault); + let orgId = ""; + let openAiProviderId = ""; + const requestId = `preq_${randomUUID()}`; + + connection.transaction(() => { + orgId = createOrganization(organization); + openAiProviderId = providerId(connection, orgId, "openai"); + providers.saveCredential({ orgId, providerId: openAiProviderId, secret: "sk-live-provider", model: "gpt-5.4" }); + connection.db.prepare(` + insert into provider_requests( + id, org_id, provider_id, provider_key, model, prompt_sha256, + response_text, raw_response_json, state, error_message, input_char_count, + output_char_count, execution_mode, policy_decision, external_response_id, + stream_event_count, cancelled_at, created_at, completed_at + ) values (?, ?, ?, 'openai', 'gpt-5.4', ?, 'partial', '[]', ?, null, 6, 7, ?, 'allowed:standard', ?, 2, null, ?, null) + `).run( + requestId, + orgId, + openAiProviderId, + "prompt-sha", + PROVIDER_REQUEST_STATE.STREAMING, + PROVIDER_EXECUTION_MODE.FOREGROUND, + "resp_live_cancel", + new Date().toISOString(), + ); + }); + await providers.cancelResponse({ orgId, requestId }); + + const request = connection.db.prepare(` + select state, external_response_id, raw_response_json, cancelled_at, completed_at + from provider_requests + where id = ? + `).get(requestId) as { + state: string; + external_response_id: string; + raw_response_json: string; + cancelled_at: string; + completed_at: string; + }; + + expect(cancelledResponseId).toBe("resp_live_cancel"); + expect(request.state).toBe(PROVIDER_REQUEST_STATE.CANCELLED); + expect(request.external_response_id).toBe("resp_live_cancel"); + expect(JSON.parse(request.raw_response_json)).toEqual({ id: "resp_live_cancel", status: "cancelled" }); + expect(request.cancelled_at).toHaveLength(24); + expect(request.completed_at).toBe(request.cancelled_at); + connection.close(); + }); +}); + +async function* streamEvents(events: ProviderStreamEvent[]): AsyncIterable { + for (const item of events) yield item; +} + +function event( + type: ProviderStreamEvent["type"], + textDelta: string | null, + rawEventJson: string, + providerResponseId: string | null, +): ProviderStreamEvent { + return { type, textDelta, rawEventJson, providerResponseId }; +}