@sixb/core 0.1.1 → 0.1.3

package/src/storage/executions/validation.ts

@@ -0,0 +1,370 @@
+import type {
+  AuthorizablePrincipal,
+  AuthorizationRef,
+  KernelOperation,
+} from "../../execution/types"
+import type { AuthStorage } from "../auth"
+import { ExecutionStorageError } from "./errors"
+import type {
+  CreateExecutionInput,
+  DurableExecutionExecutor,
+  DurableExecutionSource,
+  ExecutionRecord,
+} from "./types"
+
+export interface ExecutionValidationLookup {
+  readonly auth: AuthStorage
+  getExecution(params: {
+    readonly projectId: string
+    readonly id: string
+  }): Promise<ExecutionRecord | null>
+}
+
+export function normalizeExecutionRecord(
+  input: CreateExecutionInput,
+  createdAt: Date = new Date()
+): ExecutionRecord {
+  const record: ExecutionRecord = {
+    id: input.id,
+    projectId: input.projectId,
+    ...(input.requestedBy === undefined ? {} : { requestedBy: structuredClone(input.requestedBy) }),
+    executor: structuredClone(input.executor),
+    source: structuredClone(input.source),
+    correlationId: input.correlationId,
+    authorizationRef: structuredClone(input.authorizationRef),
+    createdAt: new Date(createdAt),
+  }
+
+  assertRecordShape(record)
+  assertExecutorAuthority(record)
+  return record
+}
+
+export async function validateExecutionRecordReferences(
+  record: ExecutionRecord,
+  lookup: ExecutionValidationLookup
+): Promise<void> {
+  await validatePrincipal(record.requestedBy, record.projectId, "Requested-by principal", lookup)
+
+  if (record.authorizationRef.type === "principal") {
+    await validatePrincipal(
+      record.authorizationRef.principal,
+      record.projectId,
+      "Authority principal",
+      lookup
+    )
+    await validateCredential(record, lookup)
+  }
+
+  await validateParent(record, lookup)
+}
+
+export function cloneExecutionRecord(record: ExecutionRecord): ExecutionRecord {
+  return structuredClone(record)
+}
+
+function assertRecordShape(record: ExecutionRecord): void {
+  assertNonBlank(record.id, "Execution id")
+  assertNonBlank(record.projectId, "Execution project id")
+  assertNonBlank(record.correlationId, "Execution correlation id")
+  assertValidDate(record.createdAt, "Execution createdAt")
+
+  if (record.requestedBy) {
+    assertPrincipal(record.requestedBy, "Requested-by principal")
+  }
+
+  assertExecutor(record.executor)
+  assertSource(record.source)
+  assertAuthorizationRef(record.authorizationRef)
+}
+
+function assertExecutor(executor: DurableExecutionExecutor): void {
+  switch (executor.type) {
+    case "request":
+      assertNonBlank(executor.requestId, "Executor request id")
+      return
+    case "primitive":
+      assertTrustedPrimitiveKind(executor.kind, "Executor primitive kind")
+      assertNonBlank(executor.runId, "Executor primitive run id")
+      return
+    case "agent":
+      assertNonBlank(executor.runId, "Executor agent run id")
+      return
+    case "kernel":
+      assertKernelOperation(executor.operation, "Executor kernel operation")
+      return
+    default:
+      invalid(`Unknown execution executor type '${String((executor as { type?: unknown }).type)}'.`)
+  }
+}
+
+function assertSource(source: DurableExecutionSource): void {
+  switch (source.type) {
+    case "http":
+      assertNonBlank(source.requestId, "Execution source request id")
+      return
+    case "webhook":
+      assertNonBlank(source.deliveryId, "Execution source delivery id")
+      return
+    case "schedule":
+    case "event":
+      assertNonBlank(source.eventId, "Execution source event id")
+      return
+    case "datasetVersion":
+      assertNonBlank(source.datasetId, "Execution source dataset id")
+      assertNonBlank(source.versionId, "Execution source dataset version id")
+      return
+    case "execution":
+      assertNonBlank(source.executionId, "Execution source execution id")
+      return
+    default:
+      invalid(
+        `Execution source type '${String((source as { type?: unknown }).type)}' is not durable.`
+      )
+  }
+}
+
+function assertAuthorizationRef(ref: AuthorizationRef): void {
+  switch (ref.type) {
+    case "principal":
+      assertPrincipal(ref.principal, "Authority principal")
+      if (ref.credential) {
+        if (ref.credential.type !== "session" && ref.credential.type !== "accessToken") {
+          invalid(
+            `Unknown authority credential type '${String(
+              (ref.credential as { type?: unknown }).type
+            )}'.`
+          )
+        }
+        assertNonBlank(ref.credential.id, "Authority credential id")
+      }
+      return
+    case "trustedPrimitive":
+      assertTrustedPrimitiveKind(ref.primitive.kind, "Authority primitive kind")
+      assertNonBlank(ref.primitive.id, "Authority primitive id")
+      assertNonBlank(ref.primitive.runId, "Authority primitive run id")
+      return
+    case "kernel":
+      assertKernelOperation(ref.operation, "Authority kernel operation")
+      return
+    case "disabled":
+      return
+    default:
+      invalid(`Unknown authorization reference type '${String((ref as { type?: unknown }).type)}'.`)
+  }
+}
+
+function assertExecutorAuthority(record: ExecutionRecord): void {
+  const { executor, authorizationRef, requestedBy, source } = record
+
+  switch (executor.type) {
+    case "request": {
+      if (source.type !== "http" || source.requestId !== executor.requestId) {
+        invalid("Request executions require a matching HTTP source request id.")
+      }
+      if (authorizationRef.type === "principal") {
+        if (!principalsEqual(requestedBy, authorizationRef.principal)) {
+          invalid("Request authority must match its requested-by principal.")
+        }
+        return
+      }
+      if (authorizationRef.type === "disabled") {
+        if (requestedBy !== undefined) {
+          invalid("Auth-disabled request executions cannot have a requested-by principal.")
+        }
+        return
+      }
+      invalid("Request executions require principal or explicitly disabled authority.")
+      break
+    }
+    case "primitive": {
+      if (authorizationRef.type !== "trustedPrimitive") {
+        invalid("Primitive executions require trusted primitive authority.")
+      }
+      if (
+        authorizationRef.primitive.kind !== executor.kind ||
+        authorizationRef.primitive.runId !== executor.runId
+      ) {
+        invalid("Trusted primitive authority must match the executor kind and run id.")
+      }
+      return
+    }
+    case "agent": {
+      if (
+        authorizationRef.type !== "principal" ||
+        authorizationRef.principal.type !== "serviceAccount"
+      ) {
+        invalid("Agent executions require service-account authority.")
+      }
+      if (authorizationRef.credential !== undefined) {
+        invalid("Agent execution authority cannot carry an external credential.")
+      }
+      return
+    }
+    case "kernel": {
+      if (
+        authorizationRef.type !== "kernel" ||
+        !kernelOperationsEqual(authorizationRef.operation, executor.operation)
+      ) {
+        invalid("Kernel authority must match the executor operation.")
+      }
+      if (requestedBy !== undefined) {
+        invalid("Kernel executions cannot have a requested-by principal.")
+      }
+      return
+    }
+  }
+}
+
+async function validatePrincipal(
+  principal: AuthorizablePrincipal | undefined,
+  projectId: string,
+  label: string,
+  lookup: ExecutionValidationLookup
+): Promise<void> {
+  if (!principal) return
+
+  const record =
+    principal.type === "user"
+      ? await lookup.auth.users.getById({ projectId, id: principal.id })
+      : await lookup.auth.serviceAccounts.getById({ projectId, id: principal.id })
+  if (!record) {
+    throw new ExecutionStorageError(
+      "missing_principal",
+      `[Sixb] ${label} '${principal.id}' does not exist in project '${projectId}'.`
+    )
+  }
+}
+
+async function validateCredential(
+  record: ExecutionRecord,
+  lookup: ExecutionValidationLookup
+): Promise<void> {
+  if (record.authorizationRef.type !== "principal" || !record.authorizationRef.credential) return
+
+  const { principal, credential } = record.authorizationRef
+  if (credential.type === "session") {
+    const session = await lookup.auth.sessions.getById({
+      projectId: record.projectId,
+      id: credential.id,
+    })
+    if (!session) {
+      throw new ExecutionStorageError(
+        "missing_credential",
+        `[Sixb] Authority session '${credential.id}' does not exist in project '${record.projectId}'.`
+      )
+    }
+    if (principal.type !== "user" || session.userId !== principal.id) {
+      throw new ExecutionStorageError(
+        "invalid_credential",
+        `[Sixb] Authority session '${credential.id}' does not belong to principal '${principal.id}'.`
+      )
+    }
+    return
+  }
+
+  const token = await lookup.auth.accessTokens.getById({
+    projectId: record.projectId,
+    id: credential.id,
+  })
+  if (!token) {
+    throw new ExecutionStorageError(
+      "missing_credential",
+      `[Sixb] Authority access token '${credential.id}' does not exist in project '${record.projectId}'.`
+    )
+  }
+  if (token.subjectType !== principal.type || token.subjectId !== principal.id) {
+    throw new ExecutionStorageError(
+      "invalid_credential",
+      `[Sixb] Authority access token '${credential.id}' does not belong to principal '${principal.id}'.`
+    )
+  }
+}
+
+async function validateParent(
+  record: ExecutionRecord,
+  lookup: ExecutionValidationLookup
+): Promise<void> {
+  if (record.source.type !== "execution") return
+  const parentExecutionId = record.source.executionId
+
+  const parent = await lookup.getExecution({
+    projectId: record.projectId,
+    id: parentExecutionId,
+  })
+  if (!parent) {
+    throw new ExecutionStorageError(
+      "missing_parent_execution",
+      `[Sixb] Parent execution '${parentExecutionId}' does not exist in project '${record.projectId}'.`
+    )
+  }
+  if (parent.correlationId !== record.correlationId) {
+    throw new ExecutionStorageError(
+      "invalid_parent_execution",
+      `[Sixb] Child execution '${record.id}' must preserve its parent correlation id.`
+    )
+  }
+  if (!principalsEqual(parent.requestedBy, record.requestedBy)) {
+    throw new ExecutionStorageError(
+      "invalid_parent_execution",
+      `[Sixb] Child execution '${record.id}' must preserve its parent requested-by principal.`
+    )
+  }
+}
+
+function assertPrincipal(principal: AuthorizablePrincipal, label: string): void {
+  if (principal.type !== "user" && principal.type !== "serviceAccount") {
+    invalid(`${label} has unsupported type '${String((principal as { type?: unknown }).type)}'.`)
+  }
+  assertNonBlank(principal.id, `${label} id`)
+}
+
+function assertKernelOperation(operation: KernelOperation, label: string): void {
+  if (operation.type !== "ontology.recover") {
+    invalid(`${label} has unsupported type '${String((operation as { type?: unknown }).type)}'.`)
+  }
+  assertNonBlank(operation.recoveryId, `${label} recovery id`)
+}
+
+function assertTrustedPrimitiveKind(value: string, label: string): void {
+  if (!TRUSTED_PRIMITIVE_KINDS.has(value)) {
+    invalid(`${label} '${value}' is not supported.`)
+  }
+}
+
+function assertNonBlank(value: string, label: string): void {
+  if (typeof value !== "string" || !value.trim()) {
+    invalid(`${label} must be a non-empty string.`)
+  }
+}
+
+function assertValidDate(value: Date, label: string): void {
+  if (!(value instanceof Date) || !Number.isFinite(value.getTime())) {
+    invalid(`${label} must be a valid date.`)
+  }
+}
+
+function principalsEqual(
+  left: AuthorizablePrincipal | undefined,
+  right: AuthorizablePrincipal | undefined
+): boolean {
+  return left?.type === right?.type && left?.id === right?.id
+}
+
+function kernelOperationsEqual(left: KernelOperation, right: KernelOperation): boolean {
+  return left.type === right.type && left.recoveryId === right.recoveryId
+}
+
+function invalid(message: string): never {
+  throw new ExecutionStorageError("invalid_input", `[Sixb] ${message}`)
+}
+
+const TRUSTED_PRIMITIVE_KINDS = new Set([
+  "action",
+  "pipeline",
+  "projection",
+  "rule",
+  "sync",
+  "webhook",
+  "workflow",
+])

package/src/storage/in-memory/index.ts

@@ -1,8 +1,16 @@
 import { AsyncLocalStorage } from "node:async_hooks"
 import { InMemoryActionRunStorage } from "../action-runs"
 import { type AgentStorage, InMemoryAgentStorage } from "../agents"
+import { type AiAccountingAttribution, type AiCostStorage, InMemoryAiCostStorage } from "../ai-cost"
+import { type AiUsageStorage, InMemoryAiUsageStorage } from "../ai-usage"
 import { type AuthStorage, InMemoryAuthStorage } from "../auth"
+import {
+  type ConnectorConnectionStorage,
+  InMemoryConnectorConnectionStorage,
+} from "../connector-connections"
 import { StorageTransactionError } from "../errors"
+import { InMemoryExecutionStorage } from "../executions/in-memory"
+import type { ExecutionStorage } from "../executions/types"
 import { type FileUploadSessionStore, InMemoryFileUploadSessions } from "../file-upload-sessions"
 import type { ObjectStorage } from "../objects"
 import { InMemoryObjectStorage } from "../objects/in-memory"
@@ -25,7 +33,6 @@ import type { TimeseriesStorage } from "../timeseries"
 import { InMemoryTimeseriesStorage } from "../timeseries/store"
 import { createTransactionStorageProxy, throwNestedStorageTransaction } from "../transaction"
 import type { Storage, StorageTransactionOptions } from "../types"
-import { InMemoryWebhookDeliveryStorage, type WebhookDeliveryStorage } from "../webhook-deliveries"
 import { InMemoryWebhookRunStorage, type WebhookRunStorage } from "../webhook-runs"
 import {
   InMemoryWorkflowInterventionStorage,
@@ -43,6 +50,13 @@ import { registerInMemoryStorageTestingAdapter } from "./testing"
  * root capability. Transaction calls reuse the already-active scope and therefore never reacquire
  * its non-reentrant lock.
  */
+export interface InMemoryStorageOptions {
+  readonly connectorConnections?: {
+    /** Storage-authoritative clock override for deterministic development and tests. */
+    readonly now?: () => Date
+  }
+}
+
 export class InMemoryStorage implements Storage {
   readonly objects: ObjectStorage
   readonly timeseries: TimeseriesStorage
@@ -51,31 +65,44 @@ export class InMemoryStorage implements Storage {
   private readonly timeseriesStorage = new InMemoryTimeseriesStorage()
   private readonly ontologyStorage: InMemoryOntologyStorage
   private readonly authStorage = new InMemoryAuthStorage()
-  private readonly agentStorage = new InMemoryAgentStorage()
-  private readonly actionRunStorage = new InMemoryActionRunStorage()
-  private readonly syncRunStorage = new InMemorySyncRunStorage()
-  private readonly pipelineRunStorage = new InMemoryPipelineRunStorage()
-  private readonly projectionRunStorage = new InMemoryProjectionRunStorage()
-  private readonly workflowRunStorage = new InMemoryWorkflowRunStorage()
+  private readonly executionStorage = new InMemoryExecutionStorage(this.authStorage)
+  private readonly agentStorage = new InMemoryAgentStorage(this.executionStorage)
+  private readonly aiUsageStorage = new InMemoryAiUsageStorage(this.executionStorage)
+  private readonly aiCostStorage = new InMemoryAiCostStorage(this.aiUsageStorage, {
+    resolve: (input) => this.resolveAiAccountingAttribution(input),
+  })
+  private readonly actionRunStorage = new InMemoryActionRunStorage(this.executionStorage)
+  private readonly syncRunStorage = new InMemorySyncRunStorage(this.executionStorage)
+  private readonly pipelineRunStorage = new InMemoryPipelineRunStorage(this.executionStorage)
+  private readonly projectionRunStorage = new InMemoryProjectionRunStorage({
+    executions: this.executionStorage,
+  })
+  private readonly workflowRunStorage = new InMemoryWorkflowRunStorage(this.executionStorage)
   private readonly workflowInterventionStorage = new InMemoryWorkflowInterventionStorage()
-  private readonly webhookDeliveryStorage = new InMemoryWebhookDeliveryStorage()
-  private readonly webhookRunStorage = new InMemoryWebhookRunStorage()
+  private readonly webhookRunStorage = new InMemoryWebhookRunStorage(this.executionStorage)
   private readonly rulesStorage = new InMemoryRulesStorage()
   private readonly fileUploadSessionStorage = new InMemoryFileUploadSessions()
+  private readonly connectorConnectionStorage: InMemoryConnectorConnectionStorage
   readonly auth: AuthStorage
+  readonly executions: ExecutionStorage
   readonly agents: AgentStorage
+  readonly aiUsage: AiUsageStorage
+  readonly aiCosts: AiCostStorage
   readonly actionRuns: InMemoryActionRunStorage
   readonly syncRuns: SyncRunStorage
   readonly pipelineRuns: PipelineRunStorage
   readonly projectionRuns: InMemoryProjectionRunStorage
   readonly workflowRuns: WorkflowRunStorage
   readonly workflowInterventions: WorkflowInterventionStorage
-  readonly webhookDeliveries: WebhookDeliveryStorage
   readonly webhookRuns: WebhookRunStorage
   readonly rules: RulesStorage
   readonly fileUploadSessions: FileUploadSessionStore
+  readonly connectorConnections: ConnectorConnectionStorage
 
-  constructor() {
+  constructor(options: InMemoryStorageOptions = {}) {
+    this.connectorConnectionStorage = new InMemoryConnectorConnectionStorage(
+      options.connectorConnections
+    )
     const scope = createStorageOperationScope(
       (run) => this.withStorageOperation(run),
       () => this.assertRootOperationAvailable()
@@ -83,7 +110,10 @@ export class InMemoryStorage implements Storage {
     this.objects = createOperationScopedFacade(this.objectStorage, scope)
     this.timeseries = createOperationScopedFacade(this.timeseriesStorage, scope)
     this.auth = createAuthOperationScope(this.authStorage, scope)
+    this.executions = createOperationScopedFacade(this.executionStorage, scope)
     this.agents = createAgentOperationScope(this.agentStorage, scope)
+    this.aiUsage = createOperationScopedFacade(this.aiUsageStorage, scope)
+    this.aiCosts = createOperationScopedFacade(this.aiCostStorage, scope)
     this.actionRuns = createOperationScopedFacade(this.actionRunStorage, scope)
     this.syncRuns = createOperationScopedFacade(this.syncRunStorage, scope)
     this.pipelineRuns = createOperationScopedFacade(this.pipelineRunStorage, scope)
@@ -93,22 +123,64 @@ export class InMemoryStorage implements Storage {
       this.workflowInterventionStorage,
       scope
     )
-    this.webhookDeliveries = createOperationScopedFacade(this.webhookDeliveryStorage, scope)
     this.webhookRuns = createOperationScopedFacade(this.webhookRunStorage, scope)
     this.rules = createOperationScopedFacade(this.rulesStorage, scope)
     this.fileUploadSessions = createOperationScopedFacade(this.fileUploadSessionStorage, scope)
+    this.connectorConnections = createOperationScopedFacade(this.connectorConnectionStorage, scope)
     this.ontologyStorage = new InMemoryOntologyStorage(this.objectStorage, this.timeseriesStorage, {
       runRootOperation: async (run) => run(),
       getTransactionToken: () => this.getActiveTransactionToken(),
       getMaterializationLifecycle: () => this.getActiveMaterializationLifecycle(),
       assertSourceMaterializationExecution: (input) =>
         this.projectionRunStorage.assertSourceMaterializationExecutionUnlocked(input),
+      executionExists: async (projectId, executionId) =>
+        (await this.executionStorage.getById({ projectId, id: executionId })) !== null,
     })
     this.ontology = createOntologyOperationScope(this.ontologyStorage, scope)
     this.ontologyStorage.registerTestingAlias(this.ontology)
     registerInMemoryStorageTestingAdapter(this, { snapshot: () => this.snapshot() })
   }
 
+  private async resolveAiAccountingAttribution(input: {
+    readonly projectId: string
+    readonly executionId: string
+  }): Promise<AiAccountingAttribution | undefined> {
+    const execution = await this.executionStorage.getById({
+      projectId: input.projectId,
+      id: input.executionId,
+    })
+    if (!execution || execution.executor.type !== "agent") return undefined
+    const directRun = await this.agentStorage.runs.getById({
+      projectId: input.projectId,
+      id: execution.executor.runId,
+    })
+    if (directRun) {
+      return {
+        kind: "agent",
+        agentId: directRun.agentId,
+        agentRunId: directRun.id,
+        threadId: directRun.threadId,
+      }
+    }
+    const agentRun = await this.workflowRunStorage.agentNodes.getByNodeRunId({
+      projectId: input.projectId,
+      nodeRunId: execution.executor.runId,
+    })
+    if (!agentRun) return undefined
+    const node = await this.workflowRunStorage.nodes.getById({
+      projectId: input.projectId,
+      id: agentRun.nodeRunId,
+    })
+    if (!node) return undefined
+    return {
+      kind: "workflowAgent",
+      agentId: agentRun.agentId,
+      nodeRunId: agentRun.nodeRunId,
+      workflowId: node.workflowId,
+      workflowRunId: node.workflowRunId,
+    }
+  }
+
   private readonly transactionScope = new AsyncLocalStorage<object>()
   private readonly activeTransactionTokens = new WeakSet<object>()
   private readonly materializationLifecycles = new WeakMap<
@@ -221,17 +293,20 @@ export class InMemoryStorage implements Storage {
       timeseries: this.timeseriesStorage,
       ontology: this.ontologyStorage,
       auth: this.authStorage,
+      executions: this.executionStorage,
       agents: this.agentStorage,
+      aiUsage: this.aiUsageStorage,
+      aiCosts: this.aiCostStorage,
       actionRuns: this.actionRunStorage,
       syncRuns: this.syncRunStorage,
       pipelineRuns: this.pipelineRunStorage,
       projectionRuns: this.projectionRunStorage,
       workflowRuns: this.workflowRunStorage,
       workflowInterventions: this.workflowInterventionStorage,
-      webhookDeliveries: this.webhookDeliveryStorage,
       webhookRuns: this.webhookRunStorage,
       rules: this.rulesStorage,
       fileUploadSessions: this.fileUploadSessionStorage,
+      connectorConnections: this.connectorConnectionStorage,
       ping: async () => undefined,
       transaction: async <T>(): Promise<T> => {
         throwNestedStorageTransaction()
@@ -245,17 +320,20 @@ export class InMemoryStorage implements Storage {
       timeseries: this.timeseriesStorage.snapshot(),
       ontology: this.ontologyStorage.snapshot(),
       auth: this.authStorage.snapshot(),
+      executions: this.executionStorage.snapshot(),
       agents: this.agentStorage.snapshot(),
+      aiUsage: this.aiUsageStorage.snapshot(),
+      aiCosts: this.aiCostStorage.snapshot(),
       actionRuns: this.actionRunStorage.snapshot(),
       syncRuns: this.syncRunStorage.snapshot(),
       pipelineRuns: this.pipelineRunStorage.snapshot(),
       projectionRuns: this.projectionRunStorage.snapshot(),
       workflowRuns: this.workflowRunStorage.snapshot(),
       workflowInterventions: this.workflowInterventionStorage.snapshot(),
-      webhookDeliveries: this.webhookDeliveryStorage.snapshot(),
       webhookRuns: this.webhookRunStorage.snapshot(),
       rules: this.rulesStorage.snapshot(),
       fileUploadSessions: this.fileUploadSessionStorage.snapshot(),
+      connectorConnections: this.connectorConnectionStorage.snapshot(),
     }
   }
 
@@ -264,17 +342,20 @@ export class InMemoryStorage implements Storage {
     this.timeseriesStorage.restore(snapshot.timeseries)
     this.ontologyStorage.restore(snapshot.ontology)
     this.authStorage.restore(snapshot.auth)
+    this.executionStorage.restore(snapshot.executions)
     this.agentStorage.restore(snapshot.agents)
+    this.aiUsageStorage.restore(snapshot.aiUsage)
+    this.aiCostStorage.restore(snapshot.aiCosts)
     this.actionRunStorage.restore(snapshot.actionRuns)
     this.syncRunStorage.restore(snapshot.syncRuns)
     this.pipelineRunStorage.restore(snapshot.pipelineRuns)
     this.projectionRunStorage.restore(snapshot.projectionRuns)
     this.workflowRunStorage.restore(snapshot.workflowRuns)
     this.workflowInterventionStorage.restore(snapshot.workflowInterventions)
-    this.webhookDeliveryStorage.restore(snapshot.webhookDeliveries)
     this.webhookRunStorage.restore(snapshot.webhookRuns)
     this.rulesStorage.restore(snapshot.rules)
     this.fileUploadSessionStorage.restore(snapshot.fileUploadSessions)
+    this.connectorConnectionStorage.restore(snapshot.connectorConnections)
   }
 }
 
@@ -283,15 +364,18 @@ export interface InMemoryStorageSnapshot {
   readonly timeseries: ReturnType<InMemoryTimeseriesStorage["snapshot"]>
   readonly ontology: ReturnType<InMemoryOntologyStorage["snapshot"]>
   readonly auth: ReturnType<InMemoryAuthStorage["snapshot"]>
+  readonly executions: ReturnType<InMemoryExecutionStorage["snapshot"]>
   readonly agents: ReturnType<InMemoryAgentStorage["snapshot"]>
+  readonly aiUsage: ReturnType<InMemoryAiUsageStorage["snapshot"]>
+  readonly aiCosts: ReturnType<InMemoryAiCostStorage["snapshot"]>
   readonly actionRuns: ReturnType<InMemoryActionRunStorage["snapshot"]>
   readonly syncRuns: ReturnType<InMemorySyncRunStorage["snapshot"]>
   readonly pipelineRuns: ReturnType<InMemoryPipelineRunStorage["snapshot"]>
   readonly projectionRuns: ReturnType<InMemoryProjectionRunStorage["snapshot"]>
   readonly workflowRuns: ReturnType<InMemoryWorkflowRunStorage["snapshot"]>
   readonly workflowInterventions: ReturnType<InMemoryWorkflowInterventionStorage["snapshot"]>
-  readonly webhookDeliveries: ReturnType<InMemoryWebhookDeliveryStorage["snapshot"]>
   readonly webhookRuns: ReturnType<InMemoryWebhookRunStorage["snapshot"]>
   readonly rules: ReturnType<InMemoryRulesStorage["snapshot"]>
   readonly fileUploadSessions: ReturnType<InMemoryFileUploadSessions["snapshot"]>
+  readonly connectorConnections: ReturnType<InMemoryConnectorConnectionStorage["snapshot"]>
 }