import type { DatabaseSync } from "node:sqlite"; import { PROVIDER_EXECUTION_MODE, 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; workspace_scope: string; } export interface CreateProviderRequestInput { orgId: string; provider: ProviderRequestOwner; model: string; workspaceId: string | null; 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 ProviderBackgroundAcceptedInput { orgId: string; requestId: string; rawResponseJson: string; externalResponseId: string; } export interface ProviderBackgroundProgressInput { orgId: string; requestId: string; rawResponseJson: string; externalResponseId: string; state: typeof PROVIDER_REQUEST_STATE.PENDING | typeof PROVIDER_REQUEST_STATE.STREAMING; } 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; rawResponseJson?: string | null; } 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 interface ProviderBackgroundRequestRow { id: string; org_id: string; provider_id: string; provider_key: string; model: string; state: string; external_response_id: string; input_char_count: number; execution_mode: ProviderExecutionMode; } 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, workspace_id, 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.workspaceId, input.model, input.promptSha256, PROVIDER_REQUEST_STATE.PENDING, input.inputCharCount, input.executionMode, policyDecision(input.provider, input.workspaceId), 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, ); } markBackgroundAccepted(input: ProviderBackgroundAcceptedInput): void { this.db.prepare(` update provider_requests set state = ?, raw_response_json = ?, external_response_id = ? where id = ? and org_id = ? `).run( PROVIDER_REQUEST_STATE.PENDING, input.rawResponseJson, input.externalResponseId, input.requestId, input.orgId, ); } markBackgroundProgress(input: ProviderBackgroundProgressInput): void { 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); } 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 = ?, 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); } 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; } listRecoverableBackgroundRequests(orgId?: string): ProviderBackgroundRequestRow[] { const sql = ` select id, org_id, provider_id, provider_key, model, state, external_response_id, input_char_count, execution_mode from provider_requests where execution_mode = ? and state in (?, ?) and external_response_id is not null ${orgId ? "and org_id = ?" : ""} order by created_at asc `; return orgId ? this.db.prepare(sql).all(PROVIDER_EXECUTION_MODE.BACKGROUND, PROVIDER_REQUEST_STATE.PENDING, PROVIDER_REQUEST_STATE.STREAMING, orgId) as unknown as ProviderBackgroundRequestRow[] : this.db.prepare(sql).all(PROVIDER_EXECUTION_MODE.BACKGROUND, PROVIDER_REQUEST_STATE.PENDING, PROVIDER_REQUEST_STATE.STREAMING) as unknown as ProviderBackgroundRequestRow[]; } loadBackgroundRequest(orgId: string, requestId: string): ProviderBackgroundRequestRow { const request = this.db.prepare(` select id, org_id, provider_id, provider_key, model, state, external_response_id, input_char_count, execution_mode from provider_requests where id = ? and org_id = ? and execution_mode = ? and external_response_id is not null `).get(requestId, orgId, PROVIDER_EXECUTION_MODE.BACKGROUND) as ProviderBackgroundRequestRow | undefined; if (!request) throw new Error("Provider background 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); } } function policyDecision(provider: ProviderRequestOwner, workspaceId: string | null): string { if (provider.workspace_scope === "workspace") return `allowed:${provider.data_policy}:workspace:${workspaceId}`; return `allowed:${provider.data_policy}:organization`; }