AgentCompany/src/main/services/provider-request-store.ts

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`;
}