goodbuddy/src/main/agent/runtime-controller.ts

import type {
  AgentQuestionAnswer,
  AgentRuntimeStatus,
  RuntimeConversationCompactInput,
  RuntimeNativeSnapshot
} from '../../shared/contracts'
import type {
  AgentExecutionRequest,
  AgentRuntime,
  RuntimeAuthorizer,
  RuntimeConversationCompactOutcome,
  RuntimeEvent
} from './runtime'
import type {
  DesktopDiagnosticFailureObserver
} from '../desktop-diagnostics'

type RuntimeSlot = {
  runtime: AgentRuntime
  activeRequests: number
  retiring: boolean
  drainable: boolean
  drainStarted?: Promise<void>
  drainError?: Error
  disposal?: Promise<void>
  disposeStarted?: Promise<void>
  resolveDisposal?: () => void
}

export const MAX_DRAINING_RUNTIME_GENERATIONS = 2

export type RuntimeReplacementStatus = {
  drainingGenerations: number
  failedDrainingGenerations: number
  pendingReplacements: number
  maximumDrainingGenerations: number
  saturated: boolean
}

export class RuntimeReplacementCapacityError extends Error {
  readonly code = 'runtime-drain-capacity'

  constructor(
    readonly occupiedGenerations: number,
    readonly maximumGenerations: number
  ) {
    super(
      `Runtime replacement is unavailable while previous generations are draining (${occupiedGenerations}/${maximumGenerations})`
    )
    this.name = 'RuntimeReplacementCapacityError'
  }
}

export interface DrainableAgentRuntime extends AgentRuntime {
  /** Stop accepting new sessions while existing session ownership remains. */
  beginDrain(): Promise<void>
  /** Resolve only after all owned sessions are released or interrupted. */
  waitForDrain(): Promise<void>
  /** Explicit bounded-shutdown escalation; never used by replacement. */
  forceShutdown(): Promise<void>
}

function isDrainableRuntime(
  runtime: AgentRuntime
): runtime is DrainableAgentRuntime {
  const candidate = runtime as Partial<DrainableAgentRuntime>
  return (
    typeof candidate.beginDrain === 'function' &&
    typeof candidate.waitForDrain === 'function' &&
    typeof candidate.forceShutdown === 'function'
  )
}

export class AgentRuntimeController implements AgentRuntime {
  private current: RuntimeSlot
  private readonly retired = new Set<RuntimeSlot>()
  private replacementQueue: Promise<void> = Promise.resolve()
  private pendingReplacements = 0
  private closing = false
  private retirement?: Promise<void>
  private disposal?: Promise<void>
  private readonly ownedConversationIds = new Set<string>()
  private readonly questionOwners = new Map<string, { slot: RuntimeSlot; requestId: string }>()

  constructor(
    runtime: AgentRuntime,
    private readonly shutdownGraceMs = 2_000,
    private readonly observeFailure?:
      DesktopDiagnosticFailureObserver
  ) {
    this.current = {
      runtime,
      activeRequests: 0,
      retiring: false,
      drainable: isDrainableRuntime(runtime)
    }
  }

  get requiresToolApproval(): boolean {
    return this.current.runtime.requiresToolApproval
  }

  get runtimeId(): AgentRuntimeStatus['id'] | undefined {
    return this.current.runtime.runtimeId
  }

  get supportsToolExecution(): boolean {
    return this.current.runtime.supportsToolExecution
  }

  get supportsScopedDataTools(): boolean {
    return this.current.runtime.supportsScopedDataTools !== false
  }

  get consumesTrustedInstructions(): boolean {
    return this.current.runtime.consumesTrustedInstructions === true
  }

  get capability(): AgentRuntime['capability'] {
    return this.current.runtime.capability
  }

  get replacementStatus(): RuntimeReplacementStatus {
    const drainingGenerations = [...this.retired].filter(
      (slot) => slot.drainable
    ).length
    const failedDrainingGenerations = [...this.retired].filter(
      (slot) =>
        slot.drainable &&
        slot.drainError !== undefined
    ).length
    const occupiedGenerations =
      drainingGenerations + this.pendingReplacements
    return {
      drainingGenerations,
      failedDrainingGenerations,
      pendingReplacements: this.pendingReplacements,
      maximumDrainingGenerations: MAX_DRAINING_RUNTIME_GENERATIONS,
      saturated:
        occupiedGenerations >= MAX_DRAINING_RUNTIME_GENERATIONS
    }
  }

  get activeRequestCount(): number {
    return [...new Set([this.current, ...this.retired])].reduce(
      (total, slot) => total + slot.activeRequests,
      0
    )
  }

  get ownedConversationCount(): number {
    return this.ownedConversationIds.size
  }

  get sharedProcess(): boolean {
    return this.current.runtime.sharedProcess === true
  }

  get canRetire(): boolean {
    return (
      !this.closing &&
      this.activeRequestCount === 0 &&
      (this.sharedProcess || this.ownedConversationIds.size === 0)
    )
  }

  replace(next: AgentRuntime): Promise<void> {
    if (this.closing) {
      return this.disposeRejectedRuntime(next).then(() => {
        throw new Error('Agent Runtime 正在关闭')
      })
    }
    const status = this.replacementStatus
    const occupiedGenerations =
      status.drainingGenerations + status.pendingReplacements
    if (occupiedGenerations >= MAX_DRAINING_RUNTIME_GENERATIONS) {
      return this.disposeRejectedRuntime(next).then(() => {
        throw new RuntimeReplacementCapacityError(
          occupiedGenerations,
          MAX_DRAINING_RUNTIME_GENERATIONS
        )
      })
    }
    this.pendingReplacements += 1
    const operation = this.replacementQueue
      .then(() => this.performReplace(next))
      .finally(() => {
        this.pendingReplacements -= 1
      })
    this.replacementQueue = operation.catch(() => undefined)
    return operation
  }

  private async performReplace(next: AgentRuntime): Promise<void> {
    const previous = this.current
    this.current = {
      runtime: next,
      activeRequests: 0,
      retiring: false,
      drainable: isDrainableRuntime(next)
    }
    const disposal = this.retireSlot(previous)
    if (previous.drainable) {
      await previous.drainStarted
      return
    }
    await waitWithin(disposal, this.shutdownGraceMs)
  }

  private async disposeRejectedRuntime(runtime: AgentRuntime): Promise<void> {
    await waitWithin(
      runtime.dispose().catch(() => undefined),
      this.shutdownGraceMs
    )
  }

  async getStatus(): Promise<AgentRuntimeStatus> {
    return this.probeStatus(
      'status',
      (runtime) => runtime.getStatus()
    )
  }

  async testConnection(): Promise<AgentRuntimeStatus> {
    return this.probeStatus(
      'connection-test',
      (runtime) =>
        runtime.testConnection?.() ?? runtime.getStatus()
    )
  }

  async getNativeSnapshot(): Promise<RuntimeNativeSnapshot> {
    return this.invoke('native-snapshot', (runtime) => {
      if (!runtime.getNativeSnapshot) {
        throw new Error('当前 Runtime 不支持原生能力清单')
      }
      return runtime.getNativeSnapshot()
    })
  }

  async compactConversation(
    request: RuntimeConversationCompactInput,
    signal: AbortSignal
  ): Promise<RuntimeConversationCompactOutcome> {
    this.ownedConversationIds.add(request.conversationId)
    return this.invoke(
      'compact',
      (runtime) => {
        if (!runtime.compactConversation) {
          throw new Error('当前 Runtime 不支持手动压缩')
        }
        return runtime.compactConversation(request, signal)
      },
      signal
    )
  }

  private async probeStatus(
    stage: string,
    operation: (runtime: AgentRuntime) => Promise<AgentRuntimeStatus>
  ): Promise<AgentRuntimeStatus> {
    const status = await this.invoke(stage, operation)
    return {
      ...status,
      supportsToolExecution: this.current.runtime.supportsToolExecution
    }
  }

  private async invoke<T>(
    stage: string,
    operation: (runtime: AgentRuntime) => Promise<T>,
    signal?: AbortSignal
  ): Promise<T> {
    if (this.closing) {
      throw new Error('Agent Runtime 正在关闭')
    }
    const slot = this.current
    slot.activeRequests += 1
    try {
      const status = await operation(slot.runtime)
      if (slot !== this.current && !(slot.runtime.sharedProcess && stage === 'compact')) {
        throw new Error('Runtime 已切换，请重试')
      }
      return status
    } catch (error) {
      if (
        slot === this.current &&
        !signal?.aborted &&
        !isAbortError(error)
      ) {
        this.reportFailure(stage, 'runtime.operation.failed', error)
      }
      throw error
    } finally {
      slot.activeRequests -= 1
      if (
        slot.retiring &&
        !slot.drainable &&
        slot.activeRequests === 0
      ) {
        await this.disposeSlot(slot)
      }
    }
  }

  async *run(
    request: AgentExecutionRequest,
    signal: AbortSignal,
    authorize?: RuntimeAuthorizer
  ): AsyncGenerator<RuntimeEvent, void, void> {
    if (this.closing) {
      throw new Error('Agent Runtime 正在关闭')
    }
    const slot = this.current
    this.ownedConversationIds.add(request.conversationId)
    let toolDenied = false
    const effectiveAuthorize: RuntimeAuthorizer = authorize
      ? async (approval, approvalSignal) => {
          const decision = await authorize(approval, approvalSignal)
          toolDenied = decision === 'deny'
          return decision
        }
      : async () => 'once'
    slot.activeRequests += 1
    try {
      try {
        for await (const event of slot.runtime.run(
          request,
          signal,
          effectiveAuthorize
        )) {
          if (slot !== this.current && !slot.drainable && !slot.runtime.sharedProcess) {
            throw new Error('Runtime 已切换，当前请求已中断')
          }
          toolDenied = false
          if (event.type === 'question') {
            this.questionOwners.set(event.questionId, { slot, requestId: request.requestId })
          }
          yield event
        }
        if (slot !== this.current && !slot.drainable && !slot.runtime.sharedProcess) {
          throw new Error('Runtime 已切换，当前请求已中断')
        }
      } catch (error) {
        if (
          !signal.aborted &&
          !toolDenied &&
          !isAbortError(error) &&
          (slot === this.current || slot.drainable)
        ) {
          this.reportFailure('run', 'runtime.run.failed', error)
        }
        throw error
      }
    } finally {
      for (const [questionId, owner] of this.questionOwners) {
        if (owner.requestId === request.requestId) this.questionOwners.delete(questionId)
      }
      slot.activeRequests -= 1
      if (
        slot.retiring &&
        !slot.drainable &&
        slot.activeRequests === 0
      ) {
        await this.disposeSlot(slot)
      }
    }
  }

  private reportFailure(
    stage: string,
    code: string,
    error: unknown
  ): void {
    try {
      this.observeFailure?.({
        component: 'runtime',
        stage,
        code,
        error
      })
    } catch {
      // Diagnostics must not alter Runtime behavior.
    }
  }

  async releaseConversation(conversationId: string): Promise<void> {
    const results = await Promise.allSettled([
      this.current.runtime.releaseConversation?.(conversationId),
      ...[...this.retired].map((slot) =>
        slot.runtime.releaseConversation?.(conversationId)
      )
    ])
    this.ownedConversationIds.delete(conversationId)
    const failures = results.flatMap((result) =>
      result.status === 'rejected' ? [result.reason] : []
    )
    if (failures.length > 0) {
      throw new AggregateError(
        failures,
        'One or more Agent Runtimes failed to release the conversation'
      )
    }
  }

  async respondToQuestion(
    questionId: string,
    answers?: AgentQuestionAnswer[]
  ): Promise<void> {
    const owner = this.questionOwners.get(questionId)
    if (this.closing && !owner) {
      throw new Error('Agent Runtime 正在关闭')
    }
    const runtime = (owner?.slot ?? this.current).runtime
    if (!runtime.respondToQuestion) {
      throw new Error('当前 Runtime 不支持回答交互式问题')
    }
    await runtime.respondToQuestion(questionId, answers)
    this.questionOwners.delete(questionId)
  }

  private retireSlot(slot: RuntimeSlot): Promise<void> {
    slot.retiring = true
    this.retired.add(slot)
    if (!slot.disposal) {
      slot.disposal = new Promise((resolve) => {
        slot.resolveDisposal = resolve
      })
    }
    if (slot.drainable && !slot.drainStarted) {
      const runtime = slot.runtime as DrainableAgentRuntime
      slot.drainStarted = runtime.beginDrain()
      void slot.drainStarted
        .then(() => runtime.waitForDrain())
        .then(() => this.disposeSlot(slot))
        .catch((error: unknown) => {
          slot.drainError =
            error instanceof Error ? error : new Error(String(error))
        })
    }
    if (!slot.drainable && slot.activeRequests === 0) {
      void this.disposeSlot(slot)
    }
    return slot.disposal
  }

  private async disposeSlot(
    slot: RuntimeSlot,
    timeoutMs?: number,
    abandonOnTimeout = false
  ): Promise<void> {
    if (!slot.resolveDisposal) {
      return
    }
    slot.disposeStarted ??= Promise.resolve().then(() =>
      slot.runtime.dispose()
    )
    let completed = true
    try {
      if (timeoutMs === undefined) {
        await slot.disposeStarted
      } else {
        completed = await waitWithin(
          slot.disposeStarted,
          timeoutMs
        )
      }
    } catch {
      completed = true
    }
    if (!completed && !abandonOnTimeout) {
      return
    }
    const resolve = slot.resolveDisposal
    if (resolve === undefined) {
      return
    }
    slot.resolveDisposal = undefined
    this.retired.delete(slot)
    resolve()
  }

  retire(): Promise<void> {
    if (this.retirement !== undefined) {
      return this.retirement
    }
    this.closing = true
    const operation = this.replacementQueue.then(async () => {
      this.retireSlot(this.current)
      await Promise.all(
        [...this.retired].map(
          (slot) => slot.disposal ?? Promise.resolve()
        )
      )
    })
    this.retirement = operation
    return operation
  }

  dispose(): Promise<void> {
    if (this.disposal !== undefined) {
      return this.disposal
    }
    this.closing = true
    const operation = this.replacementQueue.then(async () => {
      const slot = this.current
      this.retireSlot(slot)
      const slots = [...this.retired]
      const disposals = slots.map(
        (candidate) => candidate.disposal ?? Promise.resolve()
      )
      await waitWithin(
        Promise.all(disposals).then(() => undefined),
        this.shutdownGraceMs
      )
      await Promise.all(
        slots.map(async (candidate) => {
          if (!candidate.resolveDisposal) {
            return
          }
          if (candidate.drainable) {
            await waitWithin(
              (
                candidate.runtime as DrainableAgentRuntime
              ).forceShutdown().catch(() => undefined),
              this.shutdownGraceMs
            )
          }
          await this.disposeSlot(
            candidate,
            this.shutdownGraceMs,
            true
          )
        })
      )
      await Promise.all(disposals)
    })
    this.disposal = operation
    return operation
  }

  detachForApplicationExit(): Promise<void> {
    if (this.disposal !== undefined) {
      return this.disposal
    }
    this.closing = true
    const operation = this.replacementQueue.then(async () => {
      const slots = [...new Set([this.current, ...this.retired])]
      const detachable = slots.filter(
        (slot) =>
          typeof slot.runtime.detachForApplicationExit === 'function'
      )
      const ordinary = slots.filter(
        (slot) =>
          typeof slot.runtime.detachForApplicationExit !== 'function'
      )

      await Promise.allSettled(
        detachable.map(async (slot) => {
          await slot.runtime.detachForApplicationExit?.()
          this.abandonSlot(slot)
        })
      )

      for (const slot of ordinary) {
        this.retireSlot(slot)
      }
      const disposals = ordinary.map(
        (slot) => slot.disposal ?? Promise.resolve()
      )
      await waitWithin(
        Promise.all(disposals).then(() => undefined),
        this.shutdownGraceMs
      )
      await Promise.all(
        ordinary.map(async (slot) => {
          if (!slot.resolveDisposal) {
            return
          }
          if (slot.drainable) {
            await waitWithin(
              (
                slot.runtime as DrainableAgentRuntime
              ).forceShutdown().catch(() => undefined),
              this.shutdownGraceMs
            )
          }
          await this.disposeSlot(
            slot,
            this.shutdownGraceMs,
            true
          )
        })
      )
      await Promise.all(disposals)
    })
    this.disposal = operation
    return operation
  }

  private abandonSlot(slot: RuntimeSlot): void {
    const resolve = slot.resolveDisposal
    slot.resolveDisposal = undefined
    this.retired.delete(slot)
    resolve?.()
  }
}

function isAbortError(error: unknown): boolean {
  if (typeof error !== 'object' || error === null) {
    return false
  }
  try {
    return 'name' in error && error.name === 'AbortError'
  } catch {
    return false
  }
}

async function waitWithin(
  promise: Promise<void>,
  timeoutMs: number
): Promise<boolean> {
  let timer: ReturnType<typeof setTimeout> | undefined
  try {
    return await Promise.race([
      promise.then(() => true),
      new Promise<false>((resolve) => {
        timer = setTimeout(() => resolve(false), timeoutMs)
        timer.unref?.()
      })
    ])
  } finally {
    if (timer !== undefined) {
      clearTimeout(timer)
    }
  }
}