feat(relay): port the #16741 T1 seam (work drain, publication drain, release gate)

src/relay/relay-work-admission.test.ts

@@ -0,0 +1,223 @@
+import { describe, expect, it, vi } from 'vitest'
+import type { RequestContext } from './dispatcher-contract'
+import { RelayWorkAdmission } from './relay-work-admission'
+import { SKILL_SSH_RELAY_CANCEL_UPLOAD_METHOD } from '../shared/skill-ssh-relay-contract'
+
+function context(): RequestContext {
+  return { clientId: 1, isStale: () => false }
+}
+
+function pendingOperation() {
+  let resolve!: () => void
+  const promise = new Promise<void>((settle) => {
+    resolve = settle
+  })
+  return { promise, resolve }
+}
+
+async function nextTurn(): Promise<void> {
+  await new Promise<void>((resolve) => setImmediate(resolve))
+}
+
+describe('RelayWorkAdmission', () => {
+  it('starts admitted operations synchronously and preserves their result', async () => {
+    const admission = new RelayWorkAdmission()
+    const operation = vi.fn(() => 'started')
+    const result = admission.run('fs.writeFile', context(), operation)
+    expect(operation).toHaveBeenCalledOnce()
+    await expect(result).resolves.toBe('started')
+  })
+
+  it('closes admission synchronously and drains existing work despite abort and disconnect', async () => {
+    const admission = new RelayWorkAdmission()
+    const controller = new AbortController()
+    let disconnected = false
+    const mutation = pendingOperation()
+    const running = admission.run(
+      'fs.writeFile',
+      { ...context(), signal: controller.signal, isStale: () => disconnected },
+      () => mutation.promise
+    )
+    const drained = vi.fn()
+    const drain = admission.beginDrain().then(drained)
+    const rejectedOperation = vi.fn()
+    await expect(admission.run('fs.writeFile', context(), rejectedOperation)).rejects.toThrow(
+      'relay_work_admission_closed'
+    )
+    expect(rejectedOperation).not.toHaveBeenCalled()
+    controller.abort()
+    disconnected = true
+    await nextTurn()
+    expect(drained).not.toHaveBeenCalled()
+    mutation.resolve()
+    await running
+    await drain
+    expect(drained).toHaveBeenCalledOnce()
+  })
+
+  it.each([
+    'relay.status',
+    'fs.unwatchAndWait',
+    'agent.cancelExec',
+    SKILL_SSH_RELAY_CANCEL_UPLOAD_METHOD
+  ])('admits cleanup request %s while draining', async (method) => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const running = admission.run('git.diff', context(), () => mutation.promise)
+    const drain = admission.beginDrain()
+    await expect(
+      admission.run(method, context(), () => {
+        mutation.resolve()
+        return 'acknowledged'
+      })
+    ).resolves.toBe('acknowledged')
+    await running
+    await drain
+  })
+
+  it.each([
+    'rpc.cancel',
+    'git.responseAck',
+    'git.cancelResponseStream',
+    'fs.streamAck',
+    'fs.cancelStream',
+    'fs.unwatch',
+    'pty.ackData',
+    'pty.setDeliveryPaused'
+  ])('admits cleanup notification %s while draining', async (method) => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const running = admission.run('git.diff', context(), () => mutation.promise)
+    const drain = admission.beginDrain()
+    const acknowledge = vi.fn(() => mutation.resolve())
+    admission.runNotification(method, context(), acknowledge)
+    expect(acknowledge).toHaveBeenCalledOnce()
+    await running
+    await drain
+    await expect(admission.run('git.diff', context(), () => {})).rejects.toThrow(
+      'relay_work_admission_closed'
+    )
+  })
+
+  it('excludes the exact reset request without excluding another request from the same client', async () => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const running = admission.run('fs.writeFile', context(), () => mutation.promise)
+    const resetContext = context()
+    const drained = vi.fn()
+    const reset = admission.run('relay.reset', resetContext, async () => {
+      await admission.beginDrain(resetContext)
+      drained()
+    })
+    await nextTurn()
+    expect(drained).not.toHaveBeenCalled()
+    mutation.resolve()
+    await running
+    await reset
+    expect(drained).toHaveBeenCalledOnce()
+  })
+
+  it('rejects a forged context without closing admission', async () => {
+    const admission = new RelayWorkAdmission()
+    const actualContext = context()
+    await admission.run('relay.reset', actualContext, async () => {
+      await expect(admission.beginDrain({ ...actualContext })).rejects.toThrow(
+        'relay_work_drain_context_not_active'
+      )
+      await expect(admission.run('fs.writeFile', context(), () => 'open')).resolves.toBe('open')
+    })
+  })
+
+  it('rejects an already-settled request context without closing admission', async () => {
+    const admission = new RelayWorkAdmission()
+    const oldContext = context()
+    await admission.run('relay.reset', oldContext, () => {})
+    await expect(admission.beginDrain(oldContext)).rejects.toThrow(
+      'relay_work_drain_context_not_active'
+    )
+    await expect(admission.run('fs.writeFile', context(), () => 'open')).resolves.toBe('open')
+  })
+
+  it('waits for reset work when the host lifecycle supplies no exclusion', async () => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const reset = admission.run('relay.reset', context(), () => mutation.promise)
+    const drained = vi.fn()
+    const drain = admission.beginDrain().then(drained)
+    await nextTurn()
+    expect(drained).not.toHaveBeenCalled()
+    mutation.resolve()
+    await reset
+    await drain
+    expect(drained).toHaveBeenCalledOnce()
+  })
+
+  it('drops new mutation notifications without invoking their handlers during drain', async () => {
+    const admission = new RelayWorkAdmission()
+    await admission.beginDrain()
+    const mutate = vi.fn()
+    admission.runNotification('pty.write', context(), mutate)
+    expect(mutate).not.toHaveBeenCalled()
+    await expect(admission.run('pty.ackData', context(), mutate)).rejects.toThrow(
+      'relay_work_admission_closed'
+    )
+  })
+
+  it('tracks asynchronous notifications through settlement while draining', async () => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const operation = vi.fn(async () => mutation.promise)
+    admission.runNotification('pty.write', context(), operation)
+    expect(operation).toHaveBeenCalledOnce()
+    const drained = vi.fn()
+    const drain = admission.beginDrain().then(drained)
+    await nextTurn()
+    expect(drained).not.toHaveBeenCalled()
+    mutation.resolve()
+    await drain
+    expect(drained).toHaveBeenCalledOnce()
+  })
+
+  it('releases tracking after synchronous notification failure', async () => {
+    const admission = new RelayWorkAdmission()
+    expect(() =>
+      admission.runNotification('pty.write', context(), () => {
+        throw new Error('notification_failed')
+      })
+    ).toThrow('notification_failed')
+    await admission.beginDrain()
+    const mutate = vi.fn()
+    admission.runNotification('pty.write', context(), mutate)
+    expect(mutate).not.toHaveBeenCalled()
+  })
+
+  it('settles concurrent drains after an admitted operation fails without reopening', async () => {
+    const admission = new RelayWorkAdmission()
+    const mutation = pendingOperation()
+    const running = admission.run('fs.writeFile', context(), async () => {
+      await mutation.promise
+      throw new Error('write_failed')
+    })
+    const failure = expect(running).rejects.toThrow('write_failed')
+    const firstDrain = admission.beginDrain()
+    const secondDrain = admission.beginDrain()
+    mutation.resolve()
+    await failure
+    await Promise.all([firstDrain, secondDrain])
+    await expect(admission.run('fs.writeFile', context(), () => {})).rejects.toThrow(
+      'relay_work_admission_closed'
+    )
+  })
+
+  it('turns synchronous operation failures into rejected promises and releases tracking', async () => {
+    const admission = new RelayWorkAdmission()
+    const result = admission.run('fs.writeFile', context(), () => {
+      throw new Error('write_failed')
+    })
+    await expect(result).rejects.toThrow('write_failed')
+    await admission.beginDrain()
+    await expect(admission.run('fs.writeFile', context(), () => {})).rejects.toThrow(
+      'relay_work_admission_closed'
+    )
+  })
+})

src/relay/relay-work-admission.ts

@@ -0,0 +1,84 @@
+import type { RequestContext } from './dispatcher-contract'
+import { allowsRelayWorkDuringDrain } from '../shared/relay-work-drain-contract'
+
+/** Handler drain is not proof of physical process exit or detached producer cleanup. */
+export class RelayWorkAdmission {
+  private draining = false
+  private readonly active = new Map<object, { context: RequestContext; done: Promise<void> }>()
+
+  allows(method: string, notification = false): boolean {
+    return !this.draining || allowsRelayWorkDuringDrain(method, notification)
+  }
+
+  async run<T>(
+    method: string,
+    context: RequestContext,
+    operation: () => T | Promise<T>
+  ): Promise<T> {
+    if (!this.allows(method)) {
+      throw new Error('relay_work_admission_closed')
+    }
+    return this.track(context, operation)
+  }
+
+  runNotification(method: string, context: RequestContext, operation: () => void): void {
+    if (!this.allows(method, true)) {
+      return
+    }
+    const finish = this.trackEntry(context)
+    try {
+      const result = operation()
+      void Promise.resolve(result).then(finish, (error) => {
+        finish()
+        process.stderr.write(`[relay] Notification handler failed: ${String(error)}\n`)
+      })
+    } catch (error) {
+      finish()
+      throw error
+    }
+  }
+
+  assertActiveContext(context: RequestContext): void {
+    if (![...this.active.values()].some((entry) => entry.context === context)) {
+      throw new Error('relay_work_drain_context_not_active')
+    }
+  }
+
+  /** Only the actual initiating request context may be excluded from its own drain. */
+  async beginDrain(exclude?: RequestContext): Promise<void> {
+    if (exclude) {
+      this.assertActiveContext(exclude)
+    }
+    this.draining = true
+    for (;;) {
+      const pending = [...this.active.values()].filter((entry) => entry.context !== exclude)
+      if (pending.length === 0) {
+        return
+      }
+      await Promise.all(pending.map((entry) => entry.done))
+    }
+  }
+
+  private async track<T>(context: RequestContext, operation: () => T | Promise<T>): Promise<T> {
+    const finish = this.trackEntry(context)
+    try {
+      return await operation()
+    } finally {
+      finish()
+    }
+  }
+
+  // Why: no Promise.withResolvers — the legacy relay bundle still targets Node 18 hosts.
+  private trackEntry(context: RequestContext): () => void {
+    const key = {}
+    let resolve!: () => void
+    const done = new Promise<void>((settle) => {
+      resolve = settle
+    })
+    this.active.set(key, { context, done })
+    return () => {
+      this.active.delete(key)
+      resolve()
+    }
+  }
+}

src/shared/pty-ownership-transfer-release-gate.test.ts

@@ -0,0 +1,29 @@
+import { describe, expect, it } from 'vitest'
+import {
+  isPtyOwnershipTransferMutationEnabled,
+  PTY_OWNERSHIP_TRANSFER_CANARY_ENV
+} from './pty-ownership-transfer-release-gate'
+
+describe('PTY ownership-transfer release gate', () => {
+  it('is closed unless the exact canary value is present', () => {
+    expect(isPtyOwnershipTransferMutationEnabled({})).toBe(false)
+    expect(
+      isPtyOwnershipTransferMutationEnabled({
+        [PTY_OWNERSHIP_TRANSFER_CANARY_ENV]: 'true'
+      })
+    ).toBe(false)
+    expect(
+      isPtyOwnershipTransferMutationEnabled({
+        [PTY_OWNERSHIP_TRANSFER_CANARY_ENV]: '1 '
+      })
+    ).toBe(false)
+  })
+
+  it('opens only for an explicit canary opt-in', () => {
+    expect(
+      isPtyOwnershipTransferMutationEnabled({
+        [PTY_OWNERSHIP_TRANSFER_CANARY_ENV]: '1'
+      })
+    ).toBe(true)
+  })
+})

src/shared/pty-ownership-transfer-release-gate.ts

@@ -0,0 +1,12 @@
+/**
+ * The live PTY ownership-transfer path is a release canary until routed,
+ * cross-platform recovery evidence is complete. Keep the opt-in exact and
+ * process-scoped so ordinary profiles cannot inherit it accidentally.
+ */
+export const PTY_OWNERSHIP_TRANSFER_CANARY_ENV = 'ORCA_ENABLE_PTY_OWNERSHIP_TRANSFER_MUTATION'
+
+export function isPtyOwnershipTransferMutationEnabled(
+  env: Record<string, string | undefined> = process.env
+): boolean {
+  return env[PTY_OWNERSHIP_TRANSFER_CANARY_ENV] === '1'
+}

src/shared/relay-work-drain-contract.ts

@@ -0,0 +1,24 @@
+import { SKILL_SSH_RELAY_CANCEL_UPLOAD_METHOD } from './skill-ssh-relay-contract'
+
+// Why: a drain must still admit the requests and notifications that let in-flight work finish or cancel.
+// Owner-reset (T3) and network-tunnel (T4) methods join these sets when their contracts land.
+const drainRequests = new Set<string>([
+  'relay.status',
+  'fs.unwatchAndWait',
+  'agent.cancelExec',
+  SKILL_SSH_RELAY_CANCEL_UPLOAD_METHOD
+])
+const drainNotifications = new Set<string>([
+  'rpc.cancel',
+  'git.responseAck',
+  'git.cancelResponseStream',
+  'fs.streamAck',
+  'fs.cancelStream',
+  'fs.unwatch',
+  'pty.ackData',
+  'pty.setDeliveryPaused'
+])
+
+export function allowsRelayWorkDuringDrain(method: string, notification = false): boolean {
+  return (notification ? drainNotifications : drainRequests).has(method)
+}