orca/src/shared/remote-runtime-shared-control-state.ts

import type { RemoteRuntimeClientError } from './remote-runtime-client-error'
import { releaseRemoteRuntimePreparedRequest } from './remote-runtime-prepared-request-admission'
import { remoteRuntimeUnavailableError } from './remote-runtime-request-frames'
import type { RuntimeRpcResponse } from './runtime-rpc-envelope'
import type {
  RemoteRuntimeSharedConnectionDiagnostics,
  SharedControlConnectionState,
  SharedControlLogicalSubscription,
  SharedControlPendingRequest,
  SharedControlReadyWaiter
} from './remote-runtime-shared-control-types'
import { getSubscriptionId, isEndResult } from './remote-runtime-shared-control-protocol'
import { withReconnectJitter } from './reconnect-jitter'
import { tagRuntimeSubscriptionReplayResponse } from './runtime-subscription-replay'

export function buildSharedControlDiagnostics(args: {
  state: SharedControlConnectionState
  reconnecting: boolean
  pendingRequestCount: number
  subscriptionCount: number
  reconnectAttempt: number
  lastConnectedAt: number | null
  lastClose: { code: number; reason: string } | null
  lastError: string | null
}): RemoteRuntimeSharedConnectionDiagnostics {
  return {
    state: args.reconnecting ? 'reconnecting' : args.state,
    pendingRequestCount: args.pendingRequestCount,
    subscriptionCount: args.subscriptionCount,
    reconnectAttempt: args.reconnectAttempt,
    lastConnectedAt: args.lastConnectedAt,
    lastClose: args.lastClose,
    lastError: args.lastError
  }
}

export function rejectSharedControlPendingRequest(
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>,
  requestId: string,
  error: Error
): void {
  const pending = pendingRequests.get(requestId)
  if (!pending) {
    return
  }
  pendingRequests.delete(requestId)
  clearTimeout(pending.timeout)
  releaseRemoteRuntimePreparedRequest(pending)
  pending.reject(error)
}

export function resolveSharedControlPendingResponse(
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>,
  requestId: string,
  response: RuntimeRpcResponse<unknown>
): void {
  const pending = pendingRequests.get(requestId)
  if (!pending) {
    return
  }
  pendingRequests.delete(requestId)
  clearTimeout(pending.timeout)
  releaseRemoteRuntimePreparedRequest(pending)
  pending.resolve(response)
}

export function refreshSharedControlPendingRequestTimeouts(
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>
): void {
  for (const pending of pendingRequests.values()) {
    // Why: only long-poll requests opted into keepalive refresh; refreshing an
    // ordinary short RPC would keep a genuinely-stuck server call alive forever.
    if (!pending.refreshTimeoutOnKeepalive) {
      continue
    }
    const timeout = pending.timeout as ReturnType<typeof setTimeout> & { refresh?: () => void }
    timeout.refresh?.()
  }
}

export function rejectAllSharedControlPendingRequests(
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>,
  error?: Error
): void {
  const closeError = error ?? remoteRuntimeUnavailableError()
  for (const [requestId, pending] of pendingRequests) {
    clearTimeout(pending.timeout)
    pendingRequests.delete(requestId)
    releaseRemoteRuntimePreparedRequest(pending)
    pending.reject(closeError)
  }
}

export function markSharedControlSubscriptionsUnsent(
  subscriptions: Map<string, SharedControlLogicalSubscription<unknown>>
): void {
  for (const subscription of subscriptions.values()) {
    subscription.sent = false
  }
}

export function finishSharedControlSubscription(
  subscriptions: Map<string, SharedControlLogicalSubscription<unknown>>,
  subscription: SharedControlLogicalSubscription<unknown>,
  notifyClose: boolean,
  error?: RemoteRuntimeClientError
): void {
  if (subscription.closed) {
    return
  }
  subscription.closed = true
  subscriptions.delete(subscription.requestId)
  if (error) {
    subscription.callbacks.onError(error)
  }
  if (notifyClose) {
    subscription.callbacks.onClose?.()
  }
}

export function resolveSharedControlReadyWaiters(waiters: SharedControlReadyWaiter[]): void {
  for (const waiter of waiters.splice(0)) {
    waiter.resolve()
  }
}

export function rejectSharedControlReadyWaiters(
  waiters: SharedControlReadyWaiter[],
  error: Error
): void {
  for (const waiter of waiters.splice(0)) {
    waiter.reject(error)
  }
}

export function handleSharedControlSubscriptionResponse(
  subscriptions: Map<string, SharedControlLogicalSubscription<unknown>>,
  subscription: SharedControlLogicalSubscription<unknown>,
  response: RuntimeRpcResponse<unknown>
): void {
  // The replay window ends on either success or error. An error created no
  // remote subscription, so a later close must finish locally instead of
  // waiting forever for an id that will never arrive.
  subscription.awaitingResubscribe = false
  if (!response.ok) {
    subscription.sent = false
  }
  if (response.ok) {
    const subscriptionId = getSubscriptionId(response.result)
    if (subscriptionId) {
      subscription.remoteSubscriptionId = subscriptionId
    }
  }
  let delivered = response
  if (subscription.pendingReplayTag) {
    subscription.pendingReplayTag = false
    if (response.ok) {
      delivered = tagRuntimeSubscriptionReplayResponse(response)
    }
  }
  subscription.callbacks.onResponse(delivered)
  if (response.ok && isEndResult(response.result)) {
    finishSharedControlSubscription(subscriptions, subscription, false)
  }
}

export function closeSharedControlSocketState(args: {
  readyWaiters: SharedControlReadyWaiter[]
  pendingRequests: Map<string, SharedControlPendingRequest<unknown>>
  subscriptions: Map<string, SharedControlLogicalSubscription<unknown>>
  socketCleanup: (() => void) | null
  ws: { close: () => void } | null
  error?: Error
  preserveReadyWaitersAndPendingRequests?: boolean
}): void {
  if (!args.preserveReadyWaitersAndPendingRequests) {
    rejectSharedControlReadyWaiters(
      args.readyWaiters,
      args.error ?? remoteRuntimeUnavailableError()
    )
    rejectAllSharedControlPendingRequests(args.pendingRequests, args.error)
  }
  markSharedControlSubscriptionsUnsent(args.subscriptions)
  try {
    args.socketCleanup?.()
    args.ws?.close()
  } catch {
    // Best-effort cleanup of remote runtime control socket.
  }
}

export function scheduleSharedControlReconnect(args: {
  current: ReturnType<typeof setTimeout> | null
  intentionallyClosed: boolean
  reconnectAttempt: number
  delaysMs: readonly number[]
  open: () => void
}): { timer: ReturnType<typeof setTimeout> | null; reconnectAttempt: number } {
  if (args.current || args.intentionallyClosed) {
    return { timer: args.current, reconnectAttempt: args.reconnectAttempt }
  }
  const delay = withReconnectJitter(
    args.delaysMs[Math.min(args.reconnectAttempt, args.delaysMs.length - 1)]
  )
  const timer = setTimeout(args.open, delay)
  if (typeof timer.unref === 'function') {
    timer.unref()
  }
  return { timer, reconnectAttempt: args.reconnectAttempt + 1 }
}