@highstate/backend 0.9.16 → 0.9.19

package/src/business/instance-state.ts

@@ -1,292 +1,679 @@
-import type { HotStateManager } from "../hotstate"
-import type { LockManager } from "../lock"
-import type { PubSubManager } from "../pubsub"
-import type { StateManager } from "../state"
+import type {
+  ComponentKind,
+  UnitPage,
+  UnitTerminal,
+  UnitTrigger,
+  UnitWorker,
+} from "@highstate/contract"
 import type { Logger } from "pino"
-import type { RunnerBackend } from "../runner"
 import type { ArtifactService } from "../artifact"
-import { randomBytes } from "node:crypto"
-import { parseInstanceId } from "@highstate/contract"
-import { unique } from "remeda"
+import type { SecretService, UnitExtraService, WorkerService } from "../business"
+import type { PubSubManager } from "../pubsub"
+import type { RunnerBackend } from "../runner"
+import { type InstanceId, parseInstanceId } from "@highstate/contract"
+import { isNonNullish, omit } from "remeda"
+import {
+  type DatabaseManager,
+  DbNull,
+  type InstanceOperationStateCreateManyInput,
+  type InstanceOperationStateUpdateInput,
+  type InstanceStateInclude,
+  type InstanceStateUpdateInput,
+  type ProjectTransaction,
+} from "../database"
 import {
-  applyStatePatch,
-  isStableOperationStatus,
-  isStateEmpty,
-  type InstanceCustomStatus,
+  forSchema,
+  type InstanceCustomStatusInput,
+  InstanceLockedError,
+  type InstanceOperationState,
   type InstanceState,
-  type InstanceStatePatch,
-  type Operation,
+  InstanceStateNotFoundError,
+  ProjectNotFoundError,
+  projectOutputSchema,
+  waitAll,
 } from "../shared"
 
-export class InstanceStateService {
-  constructor(
-    private readonly stateManager: StateManager,
-    private readonly hotStateManager: HotStateManager,
-    private readonly pubsubManager: PubSubManager,
-    private readonly lockManager: LockManager,
-    private readonly runnerBackend: RunnerBackend,
-    private readonly artifactManager: ArtifactService,
-    private readonly logger: Logger,
-  ) {}
-
+export type GetProjectInstancesOptions = {
   /**
-   * Gets the current states of all instances in a project.
+   * Whether to include evaluation states in the result.
    *
-   * @param projectId The ID of the project for which to retrieve instance states.
+   * By default, this is false.
    */
-  async getInstanceStates(projectId: string): Promise<InstanceState[]> {
-    const entries = await this.hotStateManager.hgetall(["instance-states", projectId])
+  includeEvaluationState?: boolean
 
-    return entries.map(([, value]) => value)
-  }
+  /**
+   * Include last operation state in the result.
+   *
+   * By default, this is false.
+   */
+  includeLastOperationState?: boolean
 
-  async getInstanceState(projectId: string, instanceId: string): Promise<InstanceState | null> {
-    return await this.hotStateManager.hget(["instance-states", projectId], instanceId)
-  }
+  /**
+   * Include the instance ID of the parent instance in the result.
+   *
+   * By default, this is false.
+   */
+  includeParentInstanceId?: boolean
 
-  async updateInstanceStates(
-    projectId: string,
-    states: InstanceState[],
-    persistent = false,
-  ): Promise<InstanceState[]> {
-    const batch = this.stateManager.batch()
-    const promises: Promise<void>[] = []
+  /**
+   * Whether to include `terminalIds` and `pageIds` in the result.
+   *
+   * By default, this is false.
+   */
+  includeExtra?: boolean
 
-    for (const state of states) {
-      promises.push(this.hotStateManager.hset(["instance-states", projectId], state.id, state))
+  /**
+   * Whether to load custom statuses for each instance.
+   *
+   * By default, this is false.
+   */
+  loadCustomStatuses?: boolean
+}
 
-      promises.push(
-        this.pubsubManager.publish(["instance-state", projectId], { type: "updated", state }),
-      )
+export type ForgetInstanceStateOptions = {
+  /**
+   * Whether to delete terminals and their sessions associated with the instance.
+   *
+   * If `false`, the terminals will be marked as deleted and no new sessions will be allowed to be created,
+   * but existing sessions will remain to provide history and logs.
+   *
+   * By default, this is false.
+   */
+  clearTerminalData?: boolean
 
-      if (persistent) {
-        promises.push(this.stateManager.getInstanceStateRepository(projectId).putItem(state, batch))
-      }
-    }
+  /**
+   * Whether to delete the secrets associated with the instance.
+   *
+   * By default, this is false.
+   */
+  deleteSecrets?: boolean
+}
 
-    await Promise.allSettled(promises)
-    await batch.write()
+export type InstanceStatePatch = Pick<
+  Partial<InstanceState>,
+  | "status"
+  | "message"
+  | "statusFields"
+  | "parentId"
+  | "lastOperationState"
+  | "inputHash"
+  | "outputHash"
+  | "dependencyOutputHash"
+  | "model"
+  | "resolvedInputs"
+  | "currentResourceCount"
+>
+
+export type UpdateOperationStateOptions = {
+  /**
+   * The operation state to update.
+   */
+  operationState: InstanceOperationStateUpdateInput
 
-    this.logger.debug({ projectId, states }, `instance states updated for project "%s"`, projectId)
+  /**
+   * Instance state patch to update or function to compute patch from current state.
+   */
+  instanceState?: InstanceStatePatch
 
-    return states
+  /**
+   * Unit-specific extra data to update.
+   */
+  unitExtra?: {
+    /**
+     * Unit pages to update.
+     */
+    pages: UnitPage[]
+
+    /**
+     * Unit terminals to update.
+     */
+    terminals: UnitTerminal[]
+
+    /**
+     * Unit triggers to update.
+     */
+    triggers: UnitTrigger[]
+
+    /**
+     * Unit workers to update.
+     */
+    workers: UnitWorker[]
+
+    /**
+     * Unit secrets to update.
+     */
+    secrets: Record<string, unknown>
   }
+}
 
-  async patchInstanceState(projectId: string, patch: InstanceStatePatch): Promise<InstanceState> {
-    return await this.lockManager.acquire(["instance-state", projectId, patch.id], async () => {
-      const state = await this.hotStateManager.hget(["instance-states", projectId], patch.id)
+export function includeForInstanceState(
+  options: GetProjectInstancesOptions = {},
+): InstanceStateInclude {
+  return {
+    secrets: {
+      select: { name: true },
+    },
 
-      if (!state) {
-        throw new Error(`Instance state with ID "${patch.id}" not found in project "${projectId}"`)
-      }
+    parent: options.includeParentInstanceId ? { select: { instanceId: true } } : undefined,
 
-      const patchedState = await this.applyStatePatch(
-        projectId,
-        state,
-        patch,
-        patch.operationStatus !== undefined,
-      )
+    operationStates: options.includeLastOperationState
+      ? { take: 1, orderBy: { startedAt: "desc" } }
+      : undefined,
 
-      this.logger.debug(
-        { projectId, instanceId: patch.id, patch },
-        `instance state updated for instance "%s" in project "%s"`,
-        patch.id,
-        projectId,
-      )
+    evaluationState: options.includeEvaluationState,
 
-      return patchedState
-    })
+    terminals: options.includeExtra ? { select: { id: true } } : undefined,
+    pages: options.includeExtra ? { select: { id: true } } : undefined,
+
+    customStatuses: options.loadCustomStatuses
+      ? {
+          orderBy: [{ order: "asc" }, { createdAt: "asc" }],
+        }
+      : undefined,
   }
+}
 
-  async patchOperationInstanceState(
-    projectId: string,
-    operation: Operation,
-    patch: InstanceStatePatch,
-  ): Promise<InstanceState> {
-    return await this.lockManager.acquire(["instance-state", projectId, patch.id], async () => {
-      const state = await this.hotStateManager.hget(["instance-states", projectId], patch.id)
+export function mapInstanceStateResult(
+  instance: InstanceState & {
+    secrets: { name: string | null }[]
+    parent?: { instanceId: InstanceId } | null
+    operationStates?: InstanceOperationState[]
+    terminals?: { id: string }[]
+    pages?: { id: string }[]
+    customStatuses: InstanceState["customStatuses"]
+  },
+): InstanceState {
+  return {
+    ...omit(instance, [
+      "secrets",
+      "parent",
+      "operationStates",
+      "terminals",
+      "pages",
+      "customStatuses",
+    ]),
+
+    secretNames: instance.secrets.map(secret => secret.name).filter(isNonNullish),
+    parentInstanceId: instance.parent ? (instance.parent?.instanceId ?? null) : undefined,
+    lastOperationState: instance.operationStates?.[0],
+    evaluationState: instance.evaluationState ?? undefined,
+    terminalIds: instance.terminals ? instance.terminals.map(terminal => terminal.id) : undefined,
+    pageIds: instance.pages ? instance.pages.map(page => page.id) : undefined,
+    customStatuses: instance.customStatuses ?? undefined,
+  }
+}
 
-      const patchedState = await this.applyStatePatch(
-        projectId,
-        state,
-        patch,
-        operation.type !== "preview",
-      )
+export class InstanceStateService {
+  constructor(
+    private readonly database: DatabaseManager,
+    private readonly pubsubManager: PubSubManager,
+    private readonly runnerBackend: RunnerBackend,
+    private readonly workerService: WorkerService,
+    private readonly artifactService: ArtifactService,
+    private readonly unitExtraService: UnitExtraService,
+    private readonly secretService: SecretService,
+    private readonly logger: Logger,
+  ) {}
 
-      this.logger.debug(
-        { projectId, instanceId: patch.id, patch },
-        `instance state updated for instance "%s" in project "%s"`,
-        patch.id,
-        projectId,
-      )
+  /**
+   * Gets the aggregates of all instances in the project.
+   *
+   * @param projectId The ID of the project for which to retrieve instances.
+   * @param options Options to customize the retrieval of instances.
+   */
+  async getInstanceStates(
+    projectId: string,
+    options: GetProjectInstancesOptions = {},
+  ): Promise<InstanceState[]> {
+    const database = await this.database.forProject(projectId)
 
-      return patchedState
+    // load instance states from the database
+    const queryResult = await database.instanceState.findMany({
+      include: includeForInstanceState(options),
     })
+
+    // aggregate the results from the database
+    return queryResult.map(mapInstanceStateResult)
   }
 
-  async updateStateSecretNames(
+  /**
+   * Marks an instance state as deleted and cleans up associated resources.
+   *
+   * This operation:
+   * - only allows deletion of instances with operation state null or "destroyed";
+   * - prevents deletion of instances that have active locks;
+   * - marks the instance state as "deleted" (never actually removes the record);
+   * - handles terminals: deletes data if requested, otherwise marks as "unavailable";
+   * - handles secrets: deletes if requested, otherwise ignores them;
+   * - deletes artifact references for the instance;
+   * - recursively handles child instances using parentId relationship;
+   * - performs worker cleanup and synchronization;
+   * - performs artifact garbage collection.
+   *
+   * @param projectId The ID of the project containing the instance.
+   * @param instanceId The ID of the instance whose state is to be marked as deleted.
+   * @param options Configuration options for terminal and secret handling.
+   */
+  async forgetInstanceState(
     projectId: string,
-    instanceId: string,
-    secretNamesToAdd: string[],
-    secretNamesToRemove: string[],
-    invalidateState = true,
+    instanceId: InstanceId,
+    { deleteSecrets = false, clearTerminalData = false }: ForgetInstanceStateOptions = {},
   ): Promise<void> {
-    await this.lockManager.acquire(["instance-state", projectId, instanceId], async () => {
-      const state = await this.hotStateManager.hget(["instance-states", projectId], instanceId)
+    const database = await this.database.forProject(projectId)
+    const project = await this.database.backend.project.findUnique({
+      where: { id: projectId },
+      select: forSchema(projectOutputSchema),
+    })
 
-      const secretNameSet = new Set(state?.secretNames ?? [])
-      for (const name of secretNamesToAdd) {
-        secretNameSet.add(name)
-      }
+    if (!project) {
+      throw new ProjectNotFoundError(projectId)
+    }
 
-      for (const name of secretNamesToRemove) {
-        secretNameSet.delete(name)
-      }
+    // collect instances to process cleanup after transaction
+    const unitInstancesToCleanup: { id: string; instanceId: InstanceId }[] = []
+    const updatedStateIds: string[] = []
 
-      const patch: InstanceStatePatch = {
-        id: instanceId,
-        secretNames: Array.from(secretNameSet),
+    await database.$transaction(async tx => {
+      const state = await tx.instanceState.findUnique({
+        where: { instanceId },
+        select: { id: true, kind: true, instanceId: true, lock: { select: { stateId: true } } },
+      })
+
+      if (!state) {
+        throw new InstanceStateNotFoundError(projectId, instanceId)
       }
 
-      if (state?.operationStatus && invalidateState) {
-        // invalidate the input hash nonce if the operation status is present
-        // if not, the invalidation makes no sense
-        patch.operationStatus = {
-          inputHashNonce: randomBytes(4).readInt32LE(),
-        }
+      if (state.lock) {
+        throw new InstanceLockedError(projectId, instanceId)
       }
 
-      await this.applyStatePatch(projectId, state, patch, true)
+      await this.processInstanceDeletion(
+        tx,
+        projectId,
+        state,
+        { deleteSecrets, clearTerminalData },
+        unitInstancesToCleanup,
+        updatedStateIds,
+      )
     })
+
+    // publish state events for all updated instances
+    for (const updatedStateId of updatedStateIds) {
+      void this.pubsubManager.publish(["instance-state", projectId], {
+        type: "patched",
+        stateId: updatedStateId,
+        patch: {
+          status: "undeployed",
+          statusFields: null,
+          inputHash: null,
+          outputHash: null,
+          dependencyOutputHash: null,
+          message: null,
+          currentResourceCount: null,
+          model: null,
+          resolvedInputs: null,
+          secretNames: deleteSecrets ? [] : undefined,
+          terminalIds: clearTerminalData ? [] : undefined,
+          pageIds: [],
+          customStatuses: [],
+          triggerIds: [],
+          evaluationState: null,
+        },
+      })
+    }
+
+    // process side effects
+    try {
+      await waitAll([
+        this.workerService.cleanupWorkerUsageAndSync(projectId),
+        this.artifactService.collectGarbage(projectId),
+        ...unitInstancesToCleanup.map(async ({ id, instanceId }) => {
+          const [instanceType, instanceName] = parseInstanceId(instanceId)
+
+          await this.runnerBackend.deleteState({
+            projectId: project.id,
+            stateId: id,
+            libraryId: project.libraryId,
+            instanceName,
+            instanceType,
+          })
+        }),
+      ])
+    } catch (error) {
+      this.logger.warn(
+        { error, projectId, instanceId },
+        "failed to perform side effects after forgetting instance state",
+      )
+    }
   }
 
   /**
-   * Completely deletes the instance state from the system including all related resources.
-   *
-   * @param projectId The ID of the project from which to delete the instance state.
-   * @param instanceId The ID of the instance whose state is to be deleted.
+   * Processes the deletion of an instance within a transaction.
+   * Handles validation, deletion logic, and recursive child deletion.
    */
-  async deleteInstanceState(projectId: string, instanceId: string): Promise<void> {
-    await this.lockManager.acquire(["instance-state", projectId, instanceId], async () => {
-      const [instanceType, instanceName] = parseInstanceId(instanceId)
-      const state = await this.getInstanceState(projectId, instanceId)
-
-      const batch = this.stateManager.batch()
-
-      await Promise.allSettled([
-        // remove from the state
-        this.stateManager.getInstanceStateRepository(projectId).delete(instanceId, batch),
-
-        // clear related resources
-        this.stateManager.getInstanceLockRepository(projectId).delete(instanceId, batch),
-
-        this.stateManager
-          .getPageRepository(projectId)
-          .deleteMany(Object.values(state?.extra?.pageIds ?? {}), batch),
+  private async processInstanceDeletion(
+    tx: ProjectTransaction,
+    projectId: string,
+    state: { id: string; kind: ComponentKind; instanceId: InstanceId },
+    options: ForgetInstanceStateOptions,
+    unitInstancesToCleanup: { id: string; instanceId: InstanceId }[],
+    updatedStateIds: string[],
+  ): Promise<void> {
+    const { deleteSecrets = false, clearTerminalData = false } = options
+
+    // always mark instance state as undeployed (never actually delete the record)
+    await tx.instanceState.update({
+      where: { id: state.id },
+      data: {
+        status: "undeployed",
+        statusFields: DbNull,
+        inputHash: null,
+        outputHash: null,
+        dependencyOutputHash: null,
+        message: null,
+        currentResourceCount: null,
+        model: DbNull,
+        resolvedInputs: DbNull,
+      },
+    })
 
-        this.stateManager
-          .getTriggerRepository(projectId)
-          .deleteMany(Object.values(state?.extra?.triggerIds ?? {}), batch),
+    // handle terminals
+    if (clearTerminalData) {
+      await tx.terminal.deleteMany({ where: { stateId: state.id } })
+    } else {
+      await tx.terminal.updateMany({
+        where: { stateId: state.id },
+        data: { status: "unavailable" },
+      })
+    }
 
-        this.stateManager
-          .getTerminalRepository(projectId)
-          .deleteMany(Object.values(state?.extra?.terminalIds ?? {}), batch),
-      ])
+    // handle secrets
+    if (deleteSecrets) {
+      await tx.secret.deleteMany({
+        where: { stateId: state.id },
+      })
+    }
 
-      await Promise.allSettled([
-        // commit the state batch
-        batch.write(),
+    // delete custom statuses for this instance
+    await tx.instanceCustomStatus.deleteMany({
+      where: { stateId: state.id },
+    })
 
-        // update hotstate
-        this.hotStateManager.hdel(["instance-states", projectId], instanceId),
+    // delete other related resources
+    await tx.page.deleteMany({ where: { stateId: state.id } })
+    await tx.trigger.deleteMany({ where: { stateId: state.id } })
+
+    // remove artifact references for this instance
+    await tx.instanceState.update({
+      where: { id: state.id },
+      data: {
+        artifacts: {
+          set: [],
+        },
+      },
+    })
 
-        // publish the instance state deletion event
-        this.pubsubManager.publish(["instance-state", projectId], { type: "deleted", instanceId }),
+    // collect unit instances for Pulumi cleanup (to be done outside transaction)
+    if (state.kind === "unit") {
+      unitInstancesToCleanup.push({ id: state.id, instanceId: state.instanceId })
+    }
 
-        // remove all artifact references
-        this.artifactManager.removeUsages(
+    // track this instance as updated
+    updatedStateIds.push(state.id)
+
+    this.logger.info({ projectId }, `marked state "%s" as deleted`, state.id)
+
+    // recursively handle child instances using parentId
+    if (state.kind === "composite") {
+      const childStates = await tx.instanceState.findMany({
+        where: {
+          parentId: state.id,
+          status: { not: "undeployed" }, // don't process undeployed children
+        },
+        select: { id: true, kind: true, instanceId: true },
+      })
+
+      // recursively delete child states (within the same transaction)
+      for (const child of childStates) {
+        await this.processInstanceDeletion(
+          tx,
           projectId,
-          unique([
-            ...Object.values(state?.extra?.ownedArtifactIds ?? []).flat(),
-            ...(state?.extra?.usedArtifactIds ?? []),
-          ]),
-          [{ type: "instance", instanceId }],
-        ),
-
-        // clear pulumi state
-        this.runnerBackend.deleteState({ projectId, instanceName, instanceType }),
-      ])
-    })
+          child,
+          options,
+          unitInstancesToCleanup,
+          updatedStateIds,
+        )
+      }
+    }
   }
 
+  /**
+   * Replaces or adds a custom status for an instance in a project.
+   *
+   * @param projectId The ID of the project containing the instance.
+   * @param serviceAccoundtId The ID of the service account owning the instance.
+   * @param stateId The ID of the instance state to update.
+   * @param status The custom status to replace or add.
+   */
   async updateCustomStatus(
     projectId: string,
-    instanceId: string,
-    status: InstanceCustomStatus,
+    stateId: string,
+    serviceAccountId: string,
+    status: InstanceCustomStatusInput,
   ): Promise<void> {
-    await this.lockManager.acquire(["instance-state", projectId, instanceId], async () => {
-      const state = await this.hotStateManager.hget(["instance-states", projectId], instanceId)
-      if (!state) {
-        throw new Error(
-          `Instance state with ID "${instanceId}" not found in project "${projectId}"`,
-        )
-      }
+    const database = await this.database.forProject(projectId)
+
+    const customStatuses = await database.$transaction(async tx => {
+      await tx.instanceCustomStatus.upsert({
+        where: {
+          stateId_serviceAccountId_name: {
+            stateId,
+            serviceAccountId,
+            name: status.name,
+          },
+        },
+        create: {
+          stateId: stateId,
+          serviceAccountId,
+          name: status.name,
+          meta: status.meta,
+          value: status.value,
+          message: status.message ?? null,
+          order: status.order ?? 50,
+        },
+        update: {
+          meta: status.meta,
+          value: status.value,
+          message: status.message ?? null,
+          order: status.order ?? 50,
+        },
+      })
+
+      return await tx.instanceCustomStatus.findMany({
+        where: { stateId, serviceAccountId },
+        orderBy: [{ order: "asc" }, { createdAt: "asc" }],
+      })
+    })
 
-      const existingStatus = state.extra?.customStatuses?.find(s => s.name === status.name)
-      if (existingStatus) {
-        Object.assign(existingStatus, status)
-      } else {
-        state.extra ??= {}
-        state.extra.customStatuses ??= []
-        state.extra.customStatuses?.push(status)
-      }
+    void this.pubsubManager.publish(["instance-state", projectId], {
+      type: "patched",
+      stateId,
+      patch: { customStatuses },
+    })
+  }
+
+  /**
+   * Removes a custom status from an instance in a project.
+   *
+   * @param projectId The ID of the project containing the instance.
+   * @param stateId The ID of the instance state to update.
+   * @param serviceAccountId The ID of the service account owning the instance.
+   * @param statusName The name of the custom status to remove.
+   */
+  async removeCustomStatus(
+    projectId: string,
+    stateId: string,
+    serviceAccountId: string,
+    statusName: string,
+  ): Promise<void> {
+    const database = await this.database.forProject(projectId)
+
+    const customStatuses = await database.$transaction(async tx => {
+      await tx.instanceCustomStatus.deleteMany({
+        where: {
+          stateId,
+          serviceAccountId,
+          name: statusName,
+        },
+      })
+
+      return await tx.instanceCustomStatus.findMany({
+        where: { stateId, serviceAccountId },
+        orderBy: [{ order: "asc" }, { createdAt: "asc" }],
+      })
+    })
 
-      await this.applyStatePatch(projectId, state, { id: instanceId, extra: state.extra })
+    void this.pubsubManager.publish(["instance-state", projectId], {
+      type: "patched",
+      stateId,
+      patch: { customStatuses },
     })
   }
 
-  private async applyStatePatch(
+  /**
+   * Creates the provided operation states.
+   * Also updates the instance state if provided.
+   *
+   * @param projectId The ID of the project containing the instances.
+   * @param operationStates The tuples of operation state data to create and instance state patch to apply.
+   */
+  async createOperationStates(
     projectId: string,
-    state: InstanceState | null,
-    patch: InstanceStatePatch,
-    persistent = true,
-  ): Promise<InstanceState> {
-    const patchedState = applyStatePatch(state, patch)
-
-    const promises: Promise<void>[] = [
-      this.hotStateManager.hset(["instance-states", projectId], patchedState.id, patchedState),
-    ]
-
-    const shouldPersist =
-      persistent &&
-      (!patchedState.operationStatus ||
-        isStableOperationStatus(patchedState.operationStatus?.status))
-
-    const repo = this.stateManager.getInstanceStateRepository(projectId)
-
-    if (isStateEmpty(patchedState)) {
-      promises.push(
-        this.pubsubManager.publish(["instance-state", projectId], {
-          type: "deleted",
-          instanceId: patchedState.id,
+    operationStates: [
+      opState: InstanceOperationStateCreateManyInput,
+      instanceState: InstanceStatePatch,
+    ][],
+  ): Promise<Partial<InstanceState>[]> {
+    const database = await this.database.forProject(projectId)
+
+    const patches = await database.$transaction(async tx => {
+      return await Promise.all(
+        operationStates.map(async ([opState, instanceState]) => {
+          const operationState = await tx.instanceOperationState.create({
+            data: opState,
+          })
+
+          const state = await tx.instanceState.update({
+            where: { id: opState.stateId },
+            data: instanceState as InstanceStateUpdateInput,
+          })
+
+          return { ...state, lastOperationState: operationState }
         }),
       )
+    })
+
+    // publish patches after transaction
+    for (const patch of patches) {
+      void this.pubsubManager.publish(["instance-state", projectId], {
+        type: "patched",
+        stateId: patch.id,
+        patch,
+      })
+    }
+
+    return patches
+  }
+
+  /**
+   * Updates the operation state and instance state for a specific instance.
+   *
+   * @param projectId The ID of the project containing the instance.
+   * @param stateId The ID of the instance state to update.
+   * @param operationId The ID of the operation (required if operationState is provided).
+   * @param options Update options containing operation state, instance state, and unit-specific data.
+   * @return The instance state patch containing the updated operation state and instance state fields.
+   */
+  async updateOperationState(
+    projectId: string,
+    stateId: string,
+    operationId: string,
+    options: UpdateOperationStateOptions,
+  ): Promise<Partial<InstanceState>> {
+    const { operationState, instanceState, unitExtra } = options
+    const database = await this.database.forProject(projectId)
+
+    const result = await database.$transaction(async tx => {
+      let unitExtraData = null
+
+      // update operation state
+      const updatedOperationState = await tx.instanceOperationState.update({
+        where: {
+          operationId_stateId: {
+            operationId,
+            stateId,
+          },
+        },
+        data: operationState,
+      })
+
+      const project = await this.database.backend.project.findUnique({
+        where: { id: projectId },
+        select: { libraryId: true },
+      })
+
+      if (!project) {
+        throw new ProjectNotFoundError(projectId)
+      }
 
-      if (shouldPersist) {
-        promises.push(repo.delete(patchedState.id))
+      // update unit-specific data if provided
+      if (unitExtra) {
+        const [pageIds, terminalIds, triggerIds, secretNames] = await Promise.all([
+          this.unitExtraService.processUnitPages(tx, stateId, unitExtra.pages),
+          this.unitExtraService.processUnitTerminals(tx, stateId, unitExtra.terminals),
+          this.unitExtraService.processUnitTriggers(tx, stateId, unitExtra.triggers),
+          this.secretService.updateInstanceSecretsCore(
+            tx,
+            project.libraryId,
+            stateId,
+            unitExtra.secrets,
+          ),
+          this.workerService.updateUnitRegistrations(tx, projectId, stateId, unitExtra.workers),
+        ])
+
+        unitExtraData = { pageIds, terminalIds, triggerIds, secretNames }
       }
-    } else {
-      promises.push(
-        this.pubsubManager.publish(["instance-state", projectId], { type: "patched", patch }),
-      )
 
-      if (shouldPersist) {
-        promises.push(repo.putItem(patchedState))
+      // update instance state if provided
+      if (instanceState) {
+        await tx.instanceState.update({
+          where: { id: stateId },
+          data: instanceState as InstanceStateUpdateInput,
+        })
       }
+
+      return { updatedOperationState, unitExtraData }
+    })
+
+    // build patch combining operation state, instance state, and unit extra data
+    const patch: Partial<InstanceState> = {
+      ...instanceState,
+      ...result.unitExtraData,
+      lastOperationState: result.updatedOperationState,
     }
 
-    await Promise.allSettled(promises)
+    // emit the patch after transaction
+    void this.pubsubManager.publish(["instance-state", projectId], {
+      type: "patched",
+      stateId,
+      patch,
+    })
+
+    this.logger.debug(
+      { projectId, stateId, operationId, options },
+      "updated operation state for instance",
+    )
 
-    return patchedState
+    return patch
   }
 }