fix(ssh): fence model admission across reconnect

This commit is contained in:
OrcaWin
2026-07-28 21:04:25 -07:00
parent 37c81ff42d
commit 66c3f4347c
26 changed files with 960 additions and 165 deletions
+21 -3
View File
@@ -3089,14 +3089,15 @@
"https://github.com/stablyai/orca/pull/7002",
"https://github.com/stablyai/orca/pull/7054"
],
"invariant": "High-volume terminal output must stay bounded across daemon socket writes, SSH relay writer/decoder/source retention/reconnect, main model and projection admission, runtime metadata, detectors, and tail checks, main-to-renderer in-flight bytes, renderer scheduler queues, and hidden-output restore without starving focused input. Negotiated SSH spans remain exact and contiguous through activation, exit, cancellation, recovery, desktop projection, and required remote replacement; provisional activation data cannot project, reconnect may release it only into the exact attempt's private recovery quarantine after contract validation, and exit seals that quarantine against later same-token frames until ordered admission or exact cancellation proof. Pending renderer projection IDs are capped and compacted across split remainders, exit preparation owns its renderer fence through finalization, exit-time cancellation transfers published projections before proof commit, generation close fences late proof while draining exact projection waiters, closed-generation identity compacts without weakening stale rejection, canceled source-delivery retirement retains at most one ordered token per PTY, and pause/resume targets only the exact provider generation. Additional subscribers cannot stall the owner, while the required legacy primary retains backpressure.",
"oracle": "The current executable slice injects main-process renderer backlog pressure, then asserts unchanged path-provenance history reuse for pathless output, ordinary terminal-wait detection without a joined-tail allocation and with full prompt-family continuity, per-PTY and total pending-output caps, a 1,024-ID projection cap with split-remainder compaction, preserved sequenced-tail metadata, active-pending protection ahead of background trimming, exact provider-generation pause ownership, renderer exit/data ordering, and ACK-gated in-flight bounds. The joined source-intake/multiplex oracle blocks renderer or headless serialization, admits one snapshot-covered span and one trailing span, requires replacement reservation only after the authoritative sequence is known, advances upstream ACK eligibility exactly to the covered boundary after SnapshotEnd, and delivers then ACKs the trailing span through the ordinary live path without cancellation. Source-range ledger and multiplex seams additionally require cumulative byte ACK 40 then 100 for one 100-byte frame to release exact byte credit without early source settlement, admit a contiguous higher-generation recovered token while the prior token remains unsettled, and reject stale stream or source generations without releasing in-flight byte credit. Replacement reservation rolls back every earlier span on a later failure, commit requires the exact frozen transfer states, and an authoritatively reclaimed covered span rejects commit without local trim authority while rollback removes the reservation idempotently. The exit-deadline oracle publishes a source-backed desktop projection, advances a fake clock to cancellation, and requires one projection transfer, proof commit, preparation, and final exit with no provider close, retained obligation, duplicate finalization, or process-lifetime cancellation tombstone; a second controlled proof remains pending across generation close and must never publish final exit, while both paths release the renderer preparation lease exactly once. The intake additionally closes 2,048 sequential provider generations into one exact range, rejects stale events across that range, and preserves an out-of-order live gap until its own close. Direct-SSH contracts additionally prove activation-response settlement before the first source frame, claim-gated projection under same-decoder-turn response/data delivery, exact early-ACK reservation through send settlement or same-token retry, recovery completion visible before lease-held frames transfers those frames only into the private quarantine and admits no output until exact checkpoint-to-recoveryEnd coverage, invalid-checkpoint restore retention through response settlement followed by fresh-token retry and one live source frame, isolated saturated-subscriber eviction with healthy owner liveness, preserved gate-off primary backpressure, one-write ownership across write(false)/drain, control selection ahead of queued ordinary PTY input with FIFO lanes and bounded fairness, liveness rebasing, failed-exit retention through late ACK or exact recovery, response-settlement cancellation authority, token-scoped exit cancellation, exact checkpoint-to-recoveryEnd continuity, stale-owner one-shot retry, no physical PTY teardown on recovery failure, exact provider-generation closure and once-only cleanup when recovery cancellation publication or proof rejects, 16 MiB-plus-1 MiB decoder caps, charged per-PTY/session retention budgets, and eight-wide isolated reattach. The live Electron perf oracle adds hidden-output floods, renderer scheduler queue depth, dropped-output-zero normal scenarios, and active key latency budgets before promotion.",
"invariant": "High-volume terminal output must stay bounded across daemon socket writes, SSH relay writer/decoder/source retention/reconnect, main model and projection admission, runtime metadata, detectors, and tail checks, main-to-renderer in-flight bytes, renderer scheduler queues, and hidden-output restore without starving focused input. Negotiated SSH spans remain exact and contiguous through activation, exit, cancellation, recovery, desktop projection, and required remote replacement; provisional activation data cannot project, reconnect may release it only into the exact attempt's private recovery quarantine after contract validation, and exit seals that quarantine against later same-token frames until ordered admission or exact cancellation proof. Reconnect checkpoint capture freezes exact provider-generation/PTy admission, cancels queued old work, and waits for the running raw completion before exposing its accepted checkpoint; timeout or failure detaches only that PTy's old model, makes the checkpoint unavailable, and releases the bounded fence, while overlapping reconnects preserve any earlier outstanding fence. Pending renderer projection IDs are capped and compacted across split remainders, exit preparation owns its renderer fence through finalization, exit-time cancellation transfers published projections before proof commit, generation close fences late proof while draining exact projection waiters, closed-generation identity compacts without weakening stale rejection, canceled source-delivery retirement retains at most one ordered token per PTY, and pause/resume targets only the exact provider generation. Additional subscribers cannot stall the owner, while the required legacy primary retains backpressure.",
"oracle": "The current executable slice injects main-process renderer backlog pressure, then asserts unchanged path-provenance history reuse for pathless output, ordinary terminal-wait detection without a joined-tail allocation and with full prompt-family continuity, per-PTY and total pending-output caps, a 1,024-ID projection cap with split-remainder compaction, preserved sequenced-tail metadata, active-pending protection ahead of background trimming, exact provider-generation pause ownership, renderer exit/data ordering, and ACK-gated in-flight bounds. The joined source-intake/multiplex oracle blocks renderer or headless serialization, admits one snapshot-covered span and one trailing span, requires replacement reservation only after the authoritative sequence is known, advances upstream ACK eligibility exactly to the covered boundary after SnapshotEnd, and delivers then ACKs the trailing span through the ordinary live path without cancellation. Source-range ledger and multiplex seams additionally require cumulative byte ACK 40 then 100 for one 100-byte frame to release exact byte credit without early source settlement, admit a contiguous higher-generation recovered token while the prior token remains unsettled, and reject stale stream or source generations without releasing in-flight byte credit. Replacement reservation rolls back every earlier span on a later failure, commit requires the exact frozen transfer states, and an authoritatively reclaimed covered span rejects commit without local trim authority while rollback removes the reservation idempotently. The exit-deadline oracle publishes a source-backed desktop projection, advances a fake clock to cancellation, and requires one projection transfer, proof commit, preparation, and final exit with no provider close, retained obligation, duplicate finalization, or process-lifetime cancellation tombstone; a second controlled proof remains pending across generation close and must never publish final exit, while both paths release the renderer preparation lease exactly once. The model-migration oracle accepts span A, blocks span B's raw emulator completion, begins disconnect migration, and proves no checkpoint or reconnect attach can remain at A; releasing B advances the exact checkpoint to B and projects B once. Its fake-clock arm reaches the 10-second deadline, requires checkpoint-unavailable, one exact model reset, zero retained admission charge and timers, and no effect from late raw settlement. The intake additionally closes 2,048 sequential provider generations into one exact range, rejects stale events across that range, and preserves an out-of-order live gap until its own close. Direct-SSH contracts additionally prove activation-response settlement before the first source frame, claim-gated projection under same-decoder-turn response/data delivery, exact early-ACK reservation through send settlement or same-token retry, recovery completion visible before lease-held frames transfers those frames only into the private quarantine and admits no output until exact checkpoint-to-recoveryEnd coverage, invalid-checkpoint restore retention through response settlement followed by fresh-token retry and one live source frame, isolated saturated-subscriber eviction with healthy owner liveness, preserved gate-off primary backpressure, one-write ownership across write(false)/drain, control selection ahead of queued ordinary PTY input with FIFO lanes and bounded fairness, liveness rebasing, failed-exit retention through late ACK or exact recovery, response-settlement cancellation authority, token-scoped exit cancellation, exact checkpoint-to-recoveryEnd continuity, stale-owner one-shot retry, no physical PTY teardown on recovery failure, exact provider-generation closure and once-only cleanup when recovery cancellation publication or proof rejects, 16 MiB-plus-1 MiB decoder caps, charged per-PTY/session retention budgets, and eight-wide isolated reattach. The live Electron perf oracle adds hidden-output floods, renderer scheduler queue depth, dropped-output-zero normal scenarios, and active key latency budgets before promotion.",
"commands": [
"pnpm exec vitest run --config config/vitest.config.ts src/main/ipc/pty.test.ts",
"pnpm exec vitest run --config config/vitest.config.ts src/main/ipc/pty-pending-projection-admissions.test.ts src/main/ipc/ssh-pty-legacy-projection.test.ts src/main/ipc/ssh-pty-model-admission.test.ts --reporter=dot",
"pnpm exec vitest run --config config/vitest.config.ts src/main/runtime/orca-runtime-path-candidate-history.test.ts",
"pnpm exec vitest run --config config/vitest.config.ts src/main/runtime/orca-runtime.test.ts src/main/runtime/orca-runtime-tail-wait-memo.test.ts",
"pnpm exec vitest run --config config/vitest.config.ts src/main/providers/ssh-pty-provider.test.ts src/main/providers/ssh-pty-source-delivery-ledger.test.ts src/main/ssh/ssh-pty-retired-source-deliveries.test.ts src/main/ssh/ssh-relay-session.test.ts src/main/ssh/ssh-relay-session-data-delivery.test.ts src/main/ssh/ssh-relay-session-recovery-races.test.ts src/main/ssh/ssh-relay-session-incarnation.test.ts src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts src/main/ssh/ssh-relay-session-terminal-error.test.ts src/main/ssh/ssh-pty-recovery-retention-budget.test.ts src/main/ssh/relay-protocol-backpressure.test.ts src/relay/protocol-backpressure.test.ts src/relay/pty-source-credit-ledger.test.ts src/relay/pty-source-credit-scheduler.test.ts src/relay/relay-pty-source-publication.test.ts src/relay/ssh-pty-source-credit-adapter.test.ts --reporter=dot",
"pnpm exec vitest run --config config/vitest.config.ts src/main/ipc/ssh-pty-model-admission.test.ts src/main/ipc/ssh-pty-output-model-migration.test.ts src/main/ssh/ssh-relay-session-model-migration.test.ts --reporter=dot",
"pnpm exec vitest run --config config/vitest.config.ts src/relay/git-response-stream-ownership.test.ts src/relay/pty-handler.test.ts --reporter=dot",
"pnpm exec vitest run --config config/vitest.config.ts src/main/providers/ssh-pty-notification-routing.test.ts src/main/providers/ssh-pty-provider-agent-session-create-operation.test.ts src/main/providers/ssh-pty-provider-exit-race.test.ts --reporter=dot",
"pnpm exec vitest run --config config/vitest.config.ts src/relay/relay-pty-source-recovery-interleavings.test.ts src/relay/relay-pty-source-recovery-completion.test.ts src/relay/relay-pty-source-restore-retry.test.ts --reporter=dot",
@@ -3126,6 +3127,7 @@
"src/main/providers/ssh-pty-provider-exit-race.test.ts",
"src/main/ssh/ssh-relay-session.test.ts",
"src/main/ssh/ssh-relay-session-data-delivery.test.ts",
"src/main/ssh/ssh-relay-session-model-migration.test.ts",
"src/main/ssh/ssh-relay-session-recovery-races.test.ts",
"src/main/ssh/ssh-relay-session-incarnation.test.ts",
"src/main/ssh/ssh-relay-session-reconnect-incarnation.test.ts",
@@ -3148,6 +3150,7 @@
"src/shared/pty-consumer-session.test.ts",
"src/main/ipc/ssh-pty-model-admission.test.ts",
"src/main/ipc/ssh-pty-output-intake.test.ts",
"src/main/ipc/ssh-pty-output-model-migration.test.ts",
"src/main/ipc/ssh-pty-output-generation-guard.test.ts",
"src/main/ipc/ssh-pty-output-exit-deadline.test.ts",
"src/main/ipc/ssh-pty-remote-source-range-consumers.test.ts",
@@ -3251,6 +3254,13 @@
"each transient PTY failure is retried once without tearing down successful provider registration"
]
},
{
"file": "src/main/ssh/ssh-relay-session-model-migration.test.ts",
"assertions": [
"overlapping reconnect attaches remain blocked behind the old-generation per-PTY model fence and receive only the post-completion exact checkpoint",
"stale-owner retry retains the old-generation fence through raw completion, then requests checkpoint-unavailable restore"
]
},
{
"file": "src/main/ssh/ssh-pty-recovery-retention-budget.test.ts",
"assertions": [
@@ -3339,7 +3349,8 @@
{
"file": "src/main/ipc/ssh-pty-model-admission.test.ts",
"assertions": [
"disposal resumes every paused provider generation exactly once, including pause-only rejected entries"
"disposal resumes every paused provider generation exactly once, including pause-only rejected entries",
"migration cancels queued old-generation work while retaining the one running raw completion"
]
},
{
@@ -3352,6 +3363,13 @@
"2,048 monotonic provider closes compact into one exact stale-generation range while out-of-order live gaps remain admissible"
]
},
{
"file": "src/main/ipc/ssh-pty-output-model-migration.test.ts",
"assertions": [
"migration waits for blocked span B before exporting its checkpoint and projects B exactly once",
"migration timeout returns checkpoint-unavailable, resets one exact PTY model, and releases admission charge and timers before late raw settlement"
]
},
{
"file": "src/main/ipc/ssh-pty-output-generation-guard.test.ts",
"assertions": [
@@ -815,14 +815,21 @@ after main's synchronous install but before authority commit is bounded by the
provisional lease and is not visible to runtime, desktop, remote consumers, or
`livePtyIds`. Main drops unknown-token notifications before any side effect.
Before issuing attach or any same-client replacement, main serializes that
`(providerGeneration, ptyIncarnation)` through a migration fence: freeze new
old-token intake, cancel queued-but-unstarted model entries, await the one
in-flight emulator callback, then record the last completed receipt. Every old
callback also checks its captured token/generation before mutating checkpoint
state. If the bounded migration deadline expires or the callback fails, reset
that model generation and request `checkpointUnavailable`; recovery reports
`restoreRequired` instead of replaying a guessed gap.
Before issuing attach or any same-client replacement, main serializes each
`(providerGeneration, appPtyId)` through a migration fence: freeze new
old-generation intake, cancel queued-but-unstarted model entries, await the one
in-flight emulator callback, then record the last completed receipt. A later
overlapping reconnect preserves and awaits an earlier fence for the same PTY;
an empty newer generation cannot expose the earlier checkpoint. Before the
deadline, every old callback remains owned by its admission and generation
until raw completion settles. Stale-owner retry invalidates the checkpoint
value but retains the same fence, then requests restore after it settles. If
the 10-second migration deadline expires or
the callback fails, main cancels only that admission, detaches that PTY's old headless model,
prefers provider/renderer snapshot restore, and requests
`checkpointUnavailable`; recovery reports `restoreRequired` instead of
replaying a guessed gap. Late raw completion stays on the detached model and
cannot advance the canceled source checkpoint.
Attach sends `sourceRecovery: { status: 'checkpoint', clientGeneration,
ownerGeneration, ptyIncarnation, deliveryToken, acceptedSourceEndSu }`. A
@@ -2027,6 +2034,9 @@ transition:
frames, and prune cached mappings whose authoritative spans were reclaimed;
- reject and remove exit barriers plus cancellation/preparation ownership on
provider-generation close so late proofs cannot publish final exit;
- close each per-generation model-migration gate after its settlement, remove
its per-PTY fence and timer exactly once, and retain any earlier outstanding
per-PTY fence across an overlapping reconnect;
- replace the latest canceled-token record per PTY on ordered cancellation and
clear it on exact next activation, PTY exit, or provider teardown;
- release shared spans no longer referenced;
@@ -2352,6 +2362,10 @@ boundaries even though they ship in one PR.
- Adds a bounded source-range index beside the current
`RecentPtyOutputBuffer`, preserving its append performance and legacy
output.
- `SshPtyOutputModelMigration` freezes exact generation/PTy admission,
`SshRelaySession` awaits the returned per-PTY Promise before constructing
`sourceRecovery`, and `OrcaRuntimeService` detaches only the failed PTY's old
headless model before snapshot-backed restore.
- Reconnect V1 remains under the startup gate; its late-ACK, timeout,
supersession, and generation-close deterministic oracles pass.
@@ -2619,10 +2633,6 @@ and are then cleaned up idempotently.
`recovery pty.data -> pty.recoveryComplete`; prove both lifecycle fences
remain behind their token's quarantined data, while a same-token cancellation
proof bypasses only after atomically canceling that prefix;
- reconnect and supersede with queued plus in-flight emulator work; prove the
migration fence either checkpoints after the final guarded receipt or times
out to a model-generation reset/restore, with no late mutation or duplicate
gap ingestion;
- reject recovery cancellation publication and local proof application while
the attempt is current; prove the exact provider generation, mux, provider,
publishers, transferred activation, and registration close once, while
+1
View File
@@ -3228,6 +3228,7 @@ export function registerPtyHandlers(
provider as (IPtyProvider & { closeOutputIntake?: (reason: string) => void }) | undefined
)?.closeOutputIntake?.(reason)
},
resetModelForMigration: (_generation, id) => runtime?.resetPtyModelAfterMigrationFailure(id),
onGenerationClosed: (providerGeneration) => {
for (const id of pendingData.keys()) {
const pending = pendingData.get(id)
@@ -15,6 +15,7 @@ export type SshPtyModelAdmissionDebugSnapshot = Readonly<{
pressureFrames: number
pressureBytes: number
pausedPtys: number
migratingPtys: number
}>
export type SshPtyModelAdmissionOptions = {
@@ -0,0 +1,56 @@
import type { SshPtyModelAdmissionKey } from './ssh-pty-model-admission-contract'
import type { AdmissionCharge, AdmissionEntry, PtyUsage } from './ssh-pty-model-admission-entry'
import { admissionError, admissionKeyId } from './ssh-pty-model-admission-entry'
import type { SshPtyModelAdmissionPressure } from './ssh-pty-model-admission-pressure'
export function beginSshPtyModelAdmissionMigration(args: {
key: SshPtyModelAdmissionKey
migratingPtys: Set<string>
pressure: SshPtyModelAdmissionPressure
usageByPty: Map<string, PtyUsage>
release: (key: SshPtyModelAdmissionKey, charge: AdmissionCharge) => void
cleanup: (id: string, usage: PtyUsage) => void
}): void {
const id = admissionKeyId(args.key)
if (args.migratingPtys.has(id)) {
return
}
args.migratingPtys.add(id)
const error = admissionError('ssh_model_migration_queued_canceled')
args.pressure.cancelQueuedPty(args.key, error)
const usage = args.usageByPty.get(id)
if (!usage) {
return
}
const queued = usage.queued
usage.queued = []
for (const entry of queued) {
cancelQueuedEntry(entry, error, args.release)
}
args.cleanup(id, usage)
}
export function closeSshPtyModelAdmissionMigrations(
migratingPtys: Set<string>,
providerGeneration: number
): void {
const prefix = `${providerGeneration}\0`
for (const id of migratingPtys) {
if (id.startsWith(prefix)) {
migratingPtys.delete(id)
}
}
}
function cancelQueuedEntry(
entry: AdmissionEntry,
error: Error,
release: (key: SshPtyModelAdmissionKey, charge: AdmissionCharge) => void
): void {
if (entry.state === 'settled') {
return
}
entry.state = 'settled'
release(entry.key, entry.charge)
entry.reject(error)
}
@@ -83,6 +83,11 @@ export class SshPtyModelAdmissionPressure {
}
cancelPty(key: SshPtyModelAdmissionKey, error: Error, cancelReserved: () => void): void {
this.cancelQueuedPty(key, error)
cancelReserved()
}
cancelQueuedPty(key: SshPtyModelAdmissionKey, error: Error): void {
const id = admissionKeyId(key)
const canceled: AdmissionEntry[] = []
for (let index = this.entries.length - 1; index >= 0; index--) {
@@ -91,7 +96,6 @@ export class SshPtyModelAdmissionPressure {
}
}
this.rejectPressureEntries(canceled, error)
cancelReserved()
const paused = this.pausedKeys.get(id)
if (paused) {
this.pausedKeys.delete(id)
@@ -0,0 +1,18 @@
import type { SshPtyModelAdmissionDebugSnapshot } from './ssh-pty-model-admission-contract'
import type { SshPtyModelAdmissionPressure } from './ssh-pty-model-admission-pressure'
export function sshPtyModelAdmissionSnapshot(
sourceUnits: number,
bytes: number,
pressure: SshPtyModelAdmissionPressure,
migratingPtys: ReadonlySet<string>
): SshPtyModelAdmissionDebugSnapshot {
return {
sourceUnits,
bytes,
pressureFrames: pressure.frameCount,
pressureBytes: pressure.bytes,
pausedPtys: pressure.pausedPtyCount,
migratingPtys: migratingPtys.size
}
}
@@ -19,6 +19,28 @@ function accept(admission: SshPtyModelAdmission, completion: Promise<void>) {
}
describe('SshPtyModelAdmission', () => {
it('freezes migration while retaining the running raw completion', async () => {
const runningCompletion = deferred()
const admission = new SshPtyModelAdmission()
const running = accept(admission, runningCompletion.promise)
const queued = accept(admission, Promise.resolve())
admission.beginMigration({ ptyId: 'pty-1', providerGeneration: 7 })
await expect(queued).rejects.toThrow('ssh_model_migration_queued_canceled')
expect(admission.getDebugSnapshot()).toMatchObject({ sourceUnits: 4 })
await expect(
admission.accept({ ptyId: 'pty-1', providerGeneration: 7 }, 'late', 4, () => ({
sequence: 8,
completion: Promise.resolve()
}))
).rejects.toThrow('ssh_model_admission_migrating')
runningCompletion.resolve()
await expect(running).resolves.toMatchObject({ sequence: 4 })
expect(admission.getDebugSnapshot()).toMatchObject({ sourceUnits: 0, bytes: 0 })
})
it('cancels a never-settling running entry when its generation closes', async () => {
const admission = new SshPtyModelAdmission()
const receipt = accept(admission, new Promise<void>(() => {}))
+34 -13
View File
@@ -1,5 +1,4 @@
import type {
SshPtyModelAdmissionDebugSnapshot,
SshPtyModelAdmissionKey,
SshPtyModelAdmissionOptions,
SshPtyModelAdmissionReceipt
@@ -15,11 +14,13 @@ import {
type AdmissionEntry,
type PtyUsage
} from './ssh-pty-model-admission-entry'
import { resolveSshPtyModelAdmissionLimits } from './ssh-pty-model-admission-limits'
import {
resolveSshPtyModelAdmissionLimits,
type SshPtyModelAdmissionLimits
} from './ssh-pty-model-admission-limits'
beginSshPtyModelAdmissionMigration,
closeSshPtyModelAdmissionMigrations
} from './ssh-pty-model-admission-migration'
import { SshPtyModelAdmissionPressure } from './ssh-pty-model-admission-pressure'
import { sshPtyModelAdmissionSnapshot } from './ssh-pty-model-admission-snapshot'
export type {
SshPtyModelAdmissionKey,
@@ -28,12 +29,13 @@ export type {
} from './ssh-pty-model-admission-contract'
export class SshPtyModelAdmission {
private readonly limits: SshPtyModelAdmissionLimits
private readonly limits: ReturnType<typeof resolveSshPtyModelAdmissionLimits>
private readonly closeProvider: (providerGeneration: number, reason: string) => void
private readonly usageByPty = new Map<string, PtyUsage>()
private readonly pressure: SshPtyModelAdmissionPressure
private readonly idleWaiters = new Map<string, Set<() => void>>()
private readonly closingGenerations = new Set<number>()
private readonly migratingPtys = new Set<string>()
private globalSourceUnits = 0
private globalBytes = 0
private disposed = false
@@ -60,6 +62,9 @@ export class SshPtyModelAdmission {
if (this.closingGenerations.has(key.providerGeneration)) {
return Promise.reject(admissionError('ssh_model_admission_generation_closed'))
}
if (this.migratingPtys.has(admissionKeyId(key))) {
return Promise.reject(admissionError('ssh_model_admission_migrating'))
}
const charge = { sourceUnits, bytes: retainedBytes(data) }
return new Promise((resolve, reject) => {
const entry: AdmissionEntry = {
@@ -94,6 +99,7 @@ export class SshPtyModelAdmission {
release: (key, charge) => this.release(key, charge)
})
)
closeSshPtyModelAdmissionMigrations(this.migratingPtys, providerGeneration)
const generationPrefix = `${providerGeneration}\0`
for (const id of this.idleWaiters.keys()) {
if (id.startsWith(generationPrefix)) {
@@ -102,6 +108,19 @@ export class SshPtyModelAdmission {
}
}
beginMigration(key: SshPtyModelAdmissionKey): void {
beginSshPtyModelAdmissionMigration({
key,
migratingPtys: this.migratingPtys,
pressure: this.pressure,
usageByPty: this.usageByPty,
release: (entryKey, charge) => this.release(entryKey, charge),
cleanup: (id, usage) => this.cleanupUsage(id, usage)
})
const id = admissionKeyId(key)
resolveAdmissionIdleWaiters(this.usageByPty, this.pressure.values, this.idleWaiters, id)
}
cancelPty(key: SshPtyModelAdmissionKey, reason: string): void {
const id = admissionKeyId(key)
const error = admissionError(reason)
@@ -164,17 +183,19 @@ export class SshPtyModelAdmission {
}
}
getDebugSnapshot(): SshPtyModelAdmissionDebugSnapshot {
return {
sourceUnits: this.globalSourceUnits,
bytes: this.globalBytes,
pressureFrames: this.pressure.frameCount,
pressureBytes: this.pressure.bytes,
pausedPtys: this.pressure.pausedPtyCount
}
getDebugSnapshot() {
return sshPtyModelAdmissionSnapshot(
this.globalSourceUnits,
this.globalBytes,
this.pressure,
this.migratingPtys
)
}
private canReserve(key: SshPtyModelAdmissionKey, charge: AdmissionCharge): boolean {
if (this.migratingPtys.has(admissionKeyId(key))) {
return false
}
return canReserveAdmission({
key,
charge,
@@ -58,6 +58,7 @@ export type SshPtyOutputIntakeDependencies = {
pauseProvider?: (providerGeneration: number, id: string) => boolean
resumeProvider?: (providerGeneration: number, id: string) => void
closeProvider?: (providerGeneration: number, reason: string) => void
resetModelForMigration?: (providerGeneration: number, id: string) => void
onGenerationClosed?: (providerGeneration: number, reason: string) => void
publishSourceAck?: (
providerGeneration: number,
@@ -8,6 +8,7 @@ import type {
} from './ssh-pty-output-intake-contract'
import type { SshPtyAcceptedSourceCheckpoint } from './ssh-pty-output-source-obligations'
import type { PtySourceCreditAckBatch } from '../../shared/pty-source-credit-contract'
import type { SshPtyOutputGenerationMigration } from './ssh-pty-output-model-migration'
let installedIntake: SshPtyOutputIntake | null = null
let nextProviderGeneration = 1
@@ -61,6 +62,17 @@ export function getSshPtyAcceptedSourceCheckpoints(
return installedIntake?.getAcceptedSourceCheckpoints(providerGeneration) ?? []
}
export function beginSshPtyOutputGenerationMigration(
providerGeneration: number
): SshPtyOutputGenerationMigration {
return (
installedIntake?.beginGenerationMigration(providerGeneration) ?? {
byPty: new Map(),
completion: Promise.resolve()
}
)
}
export function installSshPtySourceAckPublisher(
providerGeneration: number,
publish: SshPtySourceAckPublisher
@@ -0,0 +1,62 @@
import { vi } from 'vitest'
import {
SshPtyOutputIntake,
type SshPtyOutputDataEvent,
type SshPtyOutputIntakeDependencies
} from './ssh-pty-output-intake'
export function sshPtyOutputEvent(
overrides: Partial<SshPtyOutputDataEvent> = {}
): SshPtyOutputDataEvent {
return {
id: 'pty-1',
data: 'aaaa',
providerGeneration: 1,
ptyIncarnation: 'incarnation-1',
rawLength: 4,
transformed: false,
...overrides
}
}
export function createSshPtyOutputIntakeHarness(
overrides: Partial<SshPtyOutputIntakeDependencies> = {},
options: ConstructorParameters<typeof SshPtyOutputIntake>[1] = {}
) {
let sequence = 0
const completions: ReturnType<typeof deferred>[] = []
const order: string[] = []
const dependencies: SshPtyOutputIntakeDependencies = {
getModelSequence: () => sequence,
acceptModel: (input) => {
order.push(`model:${input.data}`)
sequence += input.rawLength
const completion = deferred()
completions.push(completion)
return { sequence, completion: completion.promise }
},
project: (input) => order.push(`project:${input.data}`),
prepareExit: vi.fn(),
finalizeExit: () => order.push('exit'),
pauseProvider: vi.fn(() => true),
resumeProvider: vi.fn(),
closeProvider: vi.fn(),
...overrides
}
return {
intake: new SshPtyOutputIntake(dependencies, options),
dependencies,
completions,
order
}
}
function deferred() {
let resolve!: () => void
let reject!: (error: Error) => void
const promise = new Promise<void>((promiseResolve, promiseReject) => {
resolve = promiseResolve
reject = promiseReject
})
return { promise, resolve, reject }
}
+3 -58
View File
@@ -1,64 +1,9 @@
import { describe, expect, it, vi } from 'vitest'
import type { LegacySshProjectionSemantics } from './ssh-pty-legacy-projection'
import {
SshPtyOutputIntake,
type SshPtyOutputDataEvent,
type SshPtyOutputIntakeDependencies
} from './ssh-pty-output-intake'
function deferred() {
let resolve!: () => void
let reject!: (error: Error) => void
const promise = new Promise<void>((promiseResolve, promiseReject) => {
resolve = promiseResolve
reject = promiseReject
})
return { promise, resolve, reject }
}
function event(overrides: Partial<SshPtyOutputDataEvent> = {}): SshPtyOutputDataEvent {
return {
id: 'pty-1',
data: 'aaaa',
providerGeneration: 1,
ptyIncarnation: 'incarnation-1',
rawLength: 4,
transformed: false,
...overrides
}
}
function createHarness(
overrides: Partial<SshPtyOutputIntakeDependencies> = {},
options: ConstructorParameters<typeof SshPtyOutputIntake>[1] = {}
) {
let sequence = 0
const completions: ReturnType<typeof deferred>[] = []
const order: string[] = []
const dependencies: SshPtyOutputIntakeDependencies = {
getModelSequence: () => sequence,
acceptModel: (input) => {
order.push(`model:${input.data}`)
sequence += input.rawLength
const completion = deferred()
completions.push(completion)
return { sequence, completion: completion.promise }
},
project: (input) => order.push(`project:${input.data}`),
prepareExit: vi.fn(),
finalizeExit: () => order.push('exit'),
pauseProvider: vi.fn(() => true),
resumeProvider: vi.fn(),
closeProvider: vi.fn(),
...overrides
}
return {
intake: new SshPtyOutputIntake(dependencies, options),
dependencies,
completions,
order
}
}
createSshPtyOutputIntakeHarness as createHarness,
sshPtyOutputEvent as event
} from './ssh-pty-output-intake-test-harness'
describe('SshPtyOutputIntake', () => {
it('plateaus at the model and pressure budgets, then resumes below low water', async () => {
+88 -64
View File
@@ -6,6 +6,11 @@ import { SshPtyModelAdmission } from './ssh-pty-model-admission'
import { SshPtyOutputExitDeadline } from './ssh-pty-output-exit-deadline'
import { settleSshPtyOutputExit } from './ssh-pty-output-exit'
import { SshPtyOutputGenerationGuard } from './ssh-pty-output-generation-guard'
import {
SshPtyOutputModelMigration,
type SshPtyOutputGenerationMigration,
type SshPtyTrackedModelAdmission
} from './ssh-pty-output-model-migration'
import {
SshPtyOutputSourceObligations,
type SshPtyOutputSourceReservation
@@ -32,6 +37,7 @@ export class SshPtyOutputIntake {
private readonly sourceObligations: SshPtyOutputSourceObligations
private readonly generationGuard = new SshPtyOutputGenerationGuard(() => this.disposed)
private readonly admission: SshPtyModelAdmission
private readonly modelMigration: SshPtyOutputModelMigration
private readonly exitDeadline: SshPtyOutputExitDeadline
private disposed = false
@@ -53,6 +59,12 @@ export class SshPtyOutputIntake {
closeProvider: (providerGeneration, reason) =>
this.dependencies.closeProvider?.(providerGeneration, reason)
})
this.modelMigration = new SshPtyOutputModelMigration(
this.admission,
this.sourceObligations,
(providerGeneration, ptyId) =>
this.dependencies.resetModelForMigration?.(providerGeneration, ptyId)
)
this.exitDeadline = new SshPtyOutputExitDeadline({
admission: this.admission,
projections: this.projections,
@@ -71,71 +83,69 @@ export class SshPtyOutputIntake {
}
let projection: LegacySshProjectionSemantics | undefined
let sourceReservation: SshPtyOutputSourceReservation | undefined
const receipt = this.admission.accept(
{ ptyId: event.id, providerGeneration: event.providerGeneration },
event.data,
event.rawLength,
() => {
const expectedSequence = this.dependencies.getModelSequence(event.id) + event.rawLength
const reservation = this.projections.reserve({
ptyId: event.id,
providerGeneration: event.providerGeneration,
ptyIncarnation: event.ptyIncarnation,
data: event.data,
sequenceEnd: expectedSequence,
rawLength: event.rawLength,
transformed: event.transformed,
source: event.source
})
try {
if (reservation.semantics.desktopSpan) {
sourceReservation = this.sourceObligations.reserve(
event,
reservation.semantics.desktopSpan
)
}
} catch (error) {
this.projections.rollback(reservation)
throw error
const key = { ptyId: event.id, providerGeneration: event.providerGeneration }
const tracked: SshPtyTrackedModelAdmission = { key, started: false }
const receipt = this.admission.accept(key, event.data, event.rawLength, () => {
tracked.started = true
const expectedSequence = this.dependencies.getModelSequence(event.id) + event.rawLength
const reservation = this.projections.reserve({
ptyId: event.id,
providerGeneration: event.providerGeneration,
ptyIncarnation: event.ptyIncarnation,
data: event.data,
sequenceEnd: expectedSequence,
rawLength: event.rawLength,
transformed: event.transformed,
source: event.source
})
try {
if (reservation.semantics.desktopSpan) {
sourceReservation = this.sourceObligations.reserve(
event,
reservation.semantics.desktopSpan
)
}
try {
projection = this.projections.commit(reservation)
if (sourceReservation) {
this.sourceObligations.commit(
sourceReservation,
event.id,
projection.identity.sequenceEnd
)
}
} catch (error) {
if (sourceReservation) {
this.sourceObligations.rollback(sourceReservation)
}
if (!this.projections.rollbackCommitted(reservation)) {
this.projections.rollback(reservation)
}
throw error
}
let model: { sequence: number; completion: Promise<void> }
try {
model = this.dependencies.acceptModel(event, projection)
} catch (error) {
if (sourceReservation) {
this.sourceObligations.rollback(sourceReservation)
}
this.projections.rollbackCommitted(reservation)
throw error
}
try {
this.dependencies.project(event, projection)
} catch {
const id = projection.identity.projectionSemanticsId
this.projections.transfer([id], 'projection-admission-failed')
}
return model
} catch (error) {
this.projections.rollback(reservation)
throw error
}
)
return receipt.then(
try {
projection = this.projections.commit(reservation)
if (sourceReservation) {
this.sourceObligations.commit(
sourceReservation,
event.id,
projection.identity.sequenceEnd
)
}
} catch (error) {
if (sourceReservation) {
this.sourceObligations.rollback(sourceReservation)
}
if (!this.projections.rollbackCommitted(reservation)) {
this.projections.rollback(reservation)
}
throw error
}
let model: { sequence: number; completion: Promise<void> }
try {
model = this.dependencies.acceptModel(event, projection)
} catch (error) {
if (sourceReservation) {
this.sourceObligations.rollback(sourceReservation)
}
this.projections.rollbackCommitted(reservation)
throw error
}
try {
this.dependencies.project(event, projection)
} catch {
const id = projection.identity.projectionSemanticsId
this.projections.transfer([id], 'projection-admission-failed')
}
return model
})
const completion = receipt.then(
(modelReceipt) => {
if (!projection) {
throw outputIntakeError('ssh_projection_receipt_missing')
@@ -152,12 +162,19 @@ export class SshPtyOutputIntake {
'model-admission-failed'
)
}
if ((error as { code?: unknown }).code !== 'ssh_exit_delivery_canceled') {
const code = (error as { code?: unknown }).code
if (
code !== 'ssh_exit_delivery_canceled' &&
!(typeof code === 'string' && code.startsWith('ssh_model_migration_'))
) {
this.dependencies.closeProvider?.(event.providerGeneration, 'model-admission-failed')
}
throw error
}
)
tracked.completion = completion
this.modelMigration.track(tracked)
return completion
}
async acceptExit(event: SshPtyOutputExitEvent): Promise<void> {
@@ -246,6 +263,13 @@ export class SshPtyOutputIntake {
return this.sourceObligations.acceptedCheckpoints(providerGeneration)
}
beginGenerationMigration(
providerGeneration: number,
timeoutMs?: number
): SshPtyOutputGenerationMigration {
return this.modelMigration.beginGeneration(providerGeneration, timeoutMs)
}
applySourceCancellationProof(
event: SshPtyOutputExitEvent,
proof: Readonly<{ sentEndSu: number; creditedEndSu: number }>
@@ -0,0 +1,96 @@
import { describe, expect, it, vi } from 'vitest'
import {
createSshPtyOutputIntakeHarness as createHarness,
sshPtyOutputEvent as event
} from './ssh-pty-output-intake-test-harness'
describe('SshPtyOutputModelMigration', () => {
it('fences a running source span before exporting its migration checkpoint', async () => {
const harness = createHarness()
const first = harness.intake.acceptData(
event({
data: 'aaaa',
source: {
spanId: 'span-a',
clientGeneration: 3,
ownerGeneration: 4,
deliveryToken: 'delivery-token-1',
sourceStartSu: 0,
sourceEndSu: 4
}
})
)
harness.completions[0]!.resolve()
await first
const second = harness.intake.acceptData(
event({
data: 'bbbb',
source: {
spanId: 'span-b',
clientGeneration: 3,
ownerGeneration: 4,
deliveryToken: 'delivery-token-1',
sourceStartSu: 4,
sourceEndSu: 8
}
})
)
const migration = harness.intake.beginGenerationMigration(1)
const result = migration.byPty.get('pty-1')
expect(result).toBeDefined()
await expect(Promise.race([result, Promise.resolve('pending')])).resolves.toBe('pending')
expect(harness.intake.getAcceptedSourceCheckpoints(1)[0]?.acceptedSourceEndSu).toBe(4)
harness.completions[1]!.resolve()
await expect(second).resolves.toMatchObject({ sequence: 8 })
await expect(result).resolves.toMatchObject({
status: 'settled',
checkpoint: { acceptedSourceEndSu: 8 }
})
expect(harness.order.filter((entry) => entry === 'project:bbbb')).toHaveLength(1)
await migration.completion
})
it('times out one migration, resets its model, and releases retained admission once', async () => {
vi.useFakeTimers()
try {
const resetModelForMigration = vi.fn()
const harness = createHarness({ resetModelForMigration })
const receipt = harness.intake.acceptData(
event({
source: {
spanId: 'span-b',
clientGeneration: 3,
ownerGeneration: 4,
deliveryToken: 'delivery-token-1',
sourceStartSu: 0,
sourceEndSu: 4
}
})
)
const migration = harness.intake.beginGenerationMigration(1, 10_000)
const result = migration.byPty.get('pty-1')
await vi.advanceTimersByTimeAsync(10_000)
await expect(result).resolves.toMatchObject({ status: 'checkpoint-unavailable' })
await expect(receipt).rejects.toThrow('ssh_model_migration_timeout')
expect(resetModelForMigration).toHaveBeenCalledOnce()
expect(resetModelForMigration).toHaveBeenCalledWith(1, 'pty-1')
expect(harness.intake.getDebugSnapshot().model).toMatchObject({
sourceUnits: 0,
bytes: 0,
migratingPtys: 1
})
expect(vi.getTimerCount()).toBe(0)
harness.completions[0]!.resolve()
await Promise.resolve()
expect(resetModelForMigration).toHaveBeenCalledOnce()
harness.intake.closeGeneration(1, 'connection_lost')
expect(harness.intake.getDebugSnapshot().model.migratingPtys).toBe(0)
} finally {
vi.useRealTimers()
}
})
})
@@ -0,0 +1,134 @@
import type {
SshPtyAcceptedSourceCheckpoint,
SshPtyOutputSourceObligations
} from './ssh-pty-output-source-obligations'
import type { SshPtyModelAdmissionKey } from './ssh-pty-model-admission-contract'
import type { SshPtyModelAdmission } from './ssh-pty-model-admission'
export const SSH_PTY_MODEL_MIGRATION_TIMEOUT_MS = 10_000
export type SshPtyOutputMigrationResult =
| Readonly<{ status: 'settled'; checkpoint: SshPtyAcceptedSourceCheckpoint }>
| Readonly<{ status: 'checkpoint-unavailable'; reason: 'completion-failed' | 'timeout' }>
export type SshPtyOutputGenerationMigration = Readonly<{
byPty: ReadonlyMap<string, Promise<SshPtyOutputMigrationResult>>
completion: Promise<void>
}>
export type SshPtyTrackedModelAdmission = {
readonly key: SshPtyModelAdmissionKey
started: boolean
completion?: Promise<unknown>
}
export class SshPtyOutputModelMigration {
private readonly pendingByPty = new Map<string, Set<SshPtyTrackedModelAdmission>>()
constructor(
private readonly admission: SshPtyModelAdmission,
private readonly sourceObligations: SshPtyOutputSourceObligations,
private readonly resetModel: (providerGeneration: number, ptyId: string) => void
) {}
track(record: SshPtyTrackedModelAdmission): void {
const id = migrationKey(record.key)
const records = this.pendingByPty.get(id) ?? new Set<SshPtyTrackedModelAdmission>()
records.add(record)
this.pendingByPty.set(id, records)
record.completion?.then(
() => this.remove(id, record),
() => this.remove(id, record)
)
}
beginGeneration(
providerGeneration: number,
timeoutMs = SSH_PTY_MODEL_MIGRATION_TIMEOUT_MS
): SshPtyOutputGenerationMigration {
const checkpoints = this.sourceObligations.acceptedCheckpoints(providerGeneration)
const keys = checkpoints.map((checkpoint) => ({
ptyId: checkpoint.id,
providerGeneration
}))
for (const key of keys) {
this.admission.beginMigration(key)
}
const byPty = new Map<string, Promise<SshPtyOutputMigrationResult>>()
for (const key of keys) {
byPty.set(key.ptyId, this.settlePty(key, timeoutMs))
}
const completion = Promise.allSettled(byPty.values()).then(() => {})
return Object.freeze({ byPty, completion })
}
private async settlePty(
key: SshPtyModelAdmissionKey,
timeoutMs: number
): Promise<SshPtyOutputMigrationResult> {
const running = Array.from(this.pendingByPty.get(migrationKey(key)) ?? []).find(
(record) => record.started
)
if (!running?.completion) {
return this.settledCheckpoint(key)
}
let timer: ReturnType<typeof setTimeout> | undefined
const timeout = new Promise<'timeout'>((resolve) => {
timer = setTimeout(() => resolve('timeout'), normalizedTimeout(timeoutMs))
timer.unref?.()
})
try {
const outcome = await Promise.race([
running.completion.then(
() => 'settled' as const,
() => 'failed' as const
),
timeout
])
if (outcome === 'settled') {
return this.settledCheckpoint(key)
}
const reason = outcome === 'timeout' ? 'timeout' : 'completion-failed'
this.resetModel(key.providerGeneration, key.ptyId)
this.admission.cancelPty(key, `ssh_model_migration_${reason}`)
if (outcome === 'timeout') {
await running.completion.catch(() => {})
}
return Object.freeze({ status: 'checkpoint-unavailable', reason })
} finally {
if (timer) {
clearTimeout(timer)
}
}
}
private settledCheckpoint(key: SshPtyModelAdmissionKey): SshPtyOutputMigrationResult {
const checkpoint = this.sourceObligations.acceptedCheckpoint(key)
if (!checkpoint) {
this.resetModel(key.providerGeneration, key.ptyId)
return Object.freeze({
status: 'checkpoint-unavailable',
reason: 'completion-failed'
})
}
return Object.freeze({ status: 'settled', checkpoint })
}
private remove(id: string, record: SshPtyTrackedModelAdmission): void {
const records = this.pendingByPty.get(id)
records?.delete(record)
if (records?.size === 0) {
this.pendingByPty.delete(id)
}
}
}
function migrationKey(key: SshPtyModelAdmissionKey): string {
return `${key.providerGeneration}\0${key.ptyId}`
}
function normalizedTimeout(timeoutMs: number): number {
return Number.isFinite(timeoutMs) && timeoutMs >= 0
? Math.floor(timeoutMs)
: SSH_PTY_MODEL_MIGRATION_TIMEOUT_MS
}
@@ -226,6 +226,17 @@ export class SshPtyOutputSourceObligations {
return Object.freeze(checkpoints)
}
acceptedCheckpoint(key: {
ptyId: string
providerGeneration: number
}): SshPtyAcceptedSourceCheckpoint | null {
return (
this.acceptedCheckpoints(key.providerGeneration).find(
(checkpoint) => checkpoint.id === key.ptyId
) ?? null
)
}
dispose(): void {
this.coordinator.dispose()
}
+4
View File
@@ -113,6 +113,10 @@ vi.mock('./ssh-pty-output-intake-registry', () => ({
let generation = 0
return () => ++generation
})(),
beginSshPtyOutputGenerationMigration: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
applySshPtySourceCancellationProof: vi.fn().mockReturnValue(false),
applySshPtySourceRecoveryCancellationProof: vi.fn().mockReturnValue(false),
closeSshPtyOutputGeneration: vi.fn(),
+5
View File
@@ -7727,6 +7727,11 @@ export class OrcaRuntimeService {
}
}
resetPtyModelAfterMigrationFailure(ptyId: string): void {
this.providerSnapshotPreferredPtys.add(ptyId)
this.disposeHeadlessTerminal(ptyId)
}
/**
* Handles incoming data from a PTY process, running agent detection,
* updating terminal tail buffers, and triggering foreground agent refreshes.
@@ -15,6 +15,7 @@ const {
pauseAdapterMock,
muxDisposeMock,
attachForReconnectMock,
beginMigrationMock,
ptyDataHandlerRef
} = vi.hoisted(() => ({
acceptOutputDataMock: vi.fn().mockResolvedValue(undefined),
@@ -25,6 +26,10 @@ const {
pauseAdapterMock: vi.fn(),
muxDisposeMock: vi.fn(),
attachForReconnectMock: vi.fn().mockResolvedValue({}),
beginMigrationMock: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
ptyDataHandlerRef: { current: undefined as undefined | ((payload: unknown) => void) }
}))
@@ -36,6 +41,7 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: acceptOutputDataMock,
acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined),
allocateSshPtyProviderGeneration: vi.fn(() => 23),
beginSshPtyOutputGenerationMigration: beginMigrationMock,
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
installSshPtySourceAckPublisher: vi.fn(() => () => {}),
@@ -20,6 +20,10 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined),
acceptSshPtyOutputExit: acceptOutputExitMock,
allocateSshPtyProviderGeneration: vi.fn(() => 31),
beginSshPtyOutputGenerationMigration: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
applySshPtySourceRecoveryCancellationProof: vi.fn(() => true),
@@ -0,0 +1,269 @@
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { SshRelaySession } from './ssh-relay-session'
import {
createMismatchedOwnerRecoveryError,
createMockDeps,
mockDeploySuccess
} from './ssh-relay-session-test-fixtures'
type SettledMigration = {
status: 'settled'
checkpoint: {
id: string
providerGeneration: number
clientGeneration: number
ownerGeneration: number
ptyIncarnation: string
deliveryToken: string
acceptedSourceEndSu: number
}
}
const { attachForReconnectMock, beginMigrationMock, muxRequestMock, openConsumerSessionMock } =
vi.hoisted(() => ({
attachForReconnectMock: vi.fn().mockResolvedValue({}),
beginMigrationMock: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
muxRequestMock: vi.fn(),
openConsumerSessionMock: vi.fn()
}))
vi.mock('./ssh-relay-deploy', () => ({ deployAndLaunchRelay: vi.fn() }))
vi.mock('./ssh-pty-consumer-session', () => ({
openSshPtyConsumerSession: openConsumerSessionMock
}))
vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined),
acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined),
allocateSshPtyProviderGeneration: vi.fn(() => 23),
beginSshPtyOutputGenerationMigration: beginMigrationMock,
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
installSshPtySourceAckPublisher: vi.fn(() => () => {}),
installSshPtySourceCancellationPublisher: vi.fn(() => () => {}),
applySshPtySourceCancellationProof: vi.fn(),
applySshPtySourceRecoveryCancellationProof: vi.fn()
}))
vi.mock('./ssh-channel-multiplexer', () => ({
SshChannelMultiplexer: class MockSshChannelMultiplexer {
notify = vi.fn()
notifyWithSettlement = vi.fn()
request = muxRequestMock
onNotification = vi.fn().mockReturnValue(() => {})
onNotificationByMethod = vi.fn().mockReturnValue(() => {})
onRequest = vi.fn().mockReturnValue(() => {})
onDispose = vi.fn().mockReturnValue(() => {})
dispose = vi.fn()
isDisposed = vi.fn().mockReturnValue(false)
}
}))
vi.mock('../providers/ssh-pty-provider', () => ({
SshPtyProvider: class MockSshPtyProvider {
onData = vi.fn().mockReturnValue(() => {})
onReplay = vi.fn().mockReturnValue(() => {})
onExit = vi.fn().mockReturnValue(() => {})
attach = vi.fn().mockResolvedValue(undefined)
attachForReconnect = attachForReconnectMock
setPtyDeliveryPauseAdapter = vi.fn()
dispose = vi.fn()
}
}))
vi.mock('../providers/ssh-filesystem-provider', () => ({
SshFilesystemProvider: class MockSshFilesystemProvider {
dispose = vi.fn()
}
}))
vi.mock('../providers/ssh-git-provider', () => ({
SshGitProvider: class MockSshGitProvider {}
}))
vi.mock('../ipc/pty', () => ({
registerSshPtyProvider: vi.fn(),
unregisterSshPtyProvider: vi.fn(),
getSshPtyProvider: vi.fn().mockReturnValue({ dispose: vi.fn() }),
getPtyIdsForConnection: vi.fn().mockReturnValue([]),
clearPtyOwnershipForConnection: vi.fn(),
clearProviderPtyState: vi.fn(),
deletePtyOwnership: vi.fn(),
restorePtyIncarnation: vi.fn(),
setPtyOwnership: vi.fn()
}))
vi.mock('../providers/ssh-filesystem-dispatch', () => ({
registerSshFilesystemProvider: vi.fn(),
unregisterSshFilesystemProvider: vi.fn(),
getSshFilesystemProvider: vi.fn().mockReturnValue({ dispose: vi.fn() })
}))
vi.mock('../providers/ssh-git-dispatch', () => ({
registerSshGitProvider: vi.fn(),
unregisterSshGitProvider: vi.fn()
}))
const { getPtyIdsForConnection, getSshPtyProvider, registerSshPtyProvider } =
await import('../ipc/pty')
const { getSshPtyAcceptedSourceCheckpoints } = await import('../ipc/ssh-pty-output-intake-registry')
function pendingMigration() {
let resolve!: (result: SettledMigration) => void
const result = new Promise<SettledMigration>((promiseResolve) => {
resolve = promiseResolve
})
return { result, resolve }
}
function settledMigration(appPtyId: string, acceptedSourceEndSu: number): SettledMigration {
return {
status: 'settled',
checkpoint: {
id: appPtyId,
providerGeneration: 23,
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'old-token',
acceptedSourceEndSu
}
}
}
describe('SshRelaySession model migration', () => {
beforeEach(() => {
vi.clearAllMocks()
vi.mocked(getPtyIdsForConnection).mockReturnValue([])
vi.mocked(getSshPtyAcceptedSourceCheckpoints).mockReturnValue([])
openConsumerSessionMock.mockImplementation(async (_mux, options) => ({
mode: 'negotiated',
clientInstanceId: options.clientInstanceId,
clientGeneration: 1,
ownerGeneration: 1,
ownerLease: 'owner-lease-1',
...(options.outputFlowControl
? {
outputFlowControl: { version: 1, windowSu: options.outputFlowControl.requestedWindowSu }
}
: {})
}))
muxRequestMock.mockResolvedValue([])
mockDeploySuccess()
})
it('keeps overlapping reconnects behind the old-generation model migration', async () => {
const targetId = 'migration-fence'
const appPtyId = `ssh:${targetId}@@pty-1`
const migration = pendingMigration()
beginMigrationMock.mockReturnValueOnce({
byPty: new Map([[appPtyId, migration.result]]),
completion: migration.result.then(() => {})
})
vi.mocked(getSshPtyAcceptedSourceCheckpoints).mockReturnValue([
{
id: appPtyId,
providerGeneration: 23,
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'old-token',
acceptedSourceEndSu: 4
}
])
const deps = createMockDeps()
const session = new SshRelaySession(
targetId,
deps.getMainWindow,
deps.mockStore,
deps.mockPortForward,
undefined,
undefined,
() => true
)
await session.establish(deps.mockConn)
vi.mocked(getPtyIdsForConnection).mockReturnValue([appPtyId])
vi.mocked(getSshPtyProvider).mockImplementation(
() => vi.mocked(registerSshPtyProvider).mock.calls.at(-1)?.[1]
)
attachForReconnectMock.mockResolvedValue({
incarnationId: 'incarnation-1',
sourceRecovery: { status: 'restoreRequired', reason: 'relayRestoreRequired' }
})
const firstReconnect = session.reconnect(deps.mockConn)
await vi.waitFor(() => expect(registerSshPtyProvider).toHaveBeenCalledTimes(2))
const secondReconnect = session.reconnect(deps.mockConn)
await vi.waitFor(() => expect(beginMigrationMock).toHaveBeenCalledTimes(2))
expect(attachForReconnectMock).not.toHaveBeenCalled()
migration.resolve(settledMigration(appPtyId, 8))
await Promise.all([firstReconnect, secondReconnect])
expect(attachForReconnectMock).toHaveBeenCalledWith(
'pty-1',
undefined,
Object.freeze({
status: 'checkpoint',
clientGeneration: 1,
ownerGeneration: 1,
ptyIncarnation: 'incarnation-1',
deliveryToken: 'old-token',
acceptedSourceEndSu: 8
})
)
})
it('waits for a stale-owner migration before requesting restore', async () => {
const targetId = 'migration-stale-owner'
const appPtyId = `ssh:${targetId}@@pty-1`
const migration = pendingMigration()
beginMigrationMock.mockReturnValueOnce({
byPty: new Map([[appPtyId, migration.result]]),
completion: migration.result.then(() => {})
})
vi.mocked(getSshPtyAcceptedSourceCheckpoints).mockReturnValue([
settledMigration(appPtyId, 4).checkpoint
])
const deps = createMockDeps()
const session = new SshRelaySession(
targetId,
deps.getMainWindow,
deps.mockStore,
deps.mockPortForward,
undefined,
undefined,
() => true
)
await session.establish(deps.mockConn)
vi.mocked(getPtyIdsForConnection).mockReturnValue([appPtyId])
vi.mocked(getSshPtyProvider).mockImplementation(
() => vi.mocked(registerSshPtyProvider).mock.calls.at(-1)?.[1]
)
openConsumerSessionMock
.mockRejectedValueOnce(createMismatchedOwnerRecoveryError())
.mockImplementationOnce(async (_mux, options) => ({
mode: 'negotiated',
clientInstanceId: options.clientInstanceId,
clientGeneration: 2,
ownerGeneration: 2,
ownerLease: 'fresh-owner-lease',
outputFlowControl: { version: 1, windowSu: 256 * 1024 }
}))
attachForReconnectMock.mockResolvedValue({
incarnationId: 'incarnation-1',
sourceRecovery: { status: 'restoreRequired', reason: 'checkpointUnavailable' }
})
const reconnect = session.reconnect(deps.mockConn)
await vi.waitFor(() => expect(openConsumerSessionMock).toHaveBeenCalledTimes(3))
expect(attachForReconnectMock).not.toHaveBeenCalled()
migration.resolve(settledMigration(appPtyId, 8))
await reconnect
expect(attachForReconnectMock).toHaveBeenCalledWith(
'pty-1',
undefined,
Object.freeze({ status: 'checkpointUnavailable' })
)
})
})
@@ -20,6 +20,10 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: vi.fn().mockResolvedValue(undefined),
acceptSshPtyOutputExit: acceptOutputExitMock,
allocateSshPtyProviderGeneration: vi.fn(() => 17),
beginSshPtyOutputGenerationMigration: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
applySshPtySourceRecoveryCancellationProof: vi.fn(() => true),
@@ -38,6 +38,10 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: acceptOutputDataMock,
acceptSshPtyOutputExit: acceptOutputExitMock,
allocateSshPtyProviderGeneration: vi.fn(() => 23),
beginSshPtyOutputGenerationMigration: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
installSshPtySourceAckPublisher: vi.fn(() => sourceAckCleanupMock),
+4
View File
@@ -29,6 +29,10 @@ vi.mock('../ipc/ssh-pty-output-intake-registry', () => ({
acceptSshPtyOutputData: acceptOutputDataMock,
acceptSshPtyOutputExit: vi.fn().mockResolvedValue(undefined),
allocateSshPtyProviderGeneration: vi.fn(() => 41),
beginSshPtyOutputGenerationMigration: vi.fn(() => ({
byPty: new Map(),
completion: Promise.resolve()
})),
closeSshPtyOutputGeneration: vi.fn(),
getSshPtyAcceptedSourceCheckpoints: vi.fn(() => []),
applySshPtySourceRecoveryCancellationProof: vi.fn(() => true),
+73 -14
View File
@@ -45,12 +45,14 @@ import {
allocateSshPtyProviderGeneration,
applySshPtySourceCancellationProof,
applySshPtySourceRecoveryCancellationProof,
beginSshPtyOutputGenerationMigration,
closeSshPtyOutputGeneration,
getSshPtyAcceptedSourceCheckpoints,
installSshPtySourceAckPublisher,
installSshPtySourceCancellationPublisher
} from '../ipc/ssh-pty-output-intake-registry'
import type { SshPtyAcceptedSourceCheckpoint } from '../ipc/ssh-pty-output-source-obligations'
import type { SshPtyOutputMigrationResult } from '../ipc/ssh-pty-output-model-migration'
import {
registerSshFilesystemProvider,
unregisterSshFilesystemProvider,
@@ -223,6 +225,7 @@ type PtyConsumerRecovery = {
serverBuildId?: string
owner?: SshPtyConsumerOwnerState
checkpointsByAppPtyId: Map<string, SshPtyAcceptedSourceCheckpoint>
modelMigrationsByAppPtyId: Map<string, Promise<SshPtyOutputMigrationResult>>
}
const ptyConsumerRecoveryByTarget = new Map<string, PtyConsumerRecovery>()
@@ -236,7 +239,8 @@ function ptyConsumerRecoveryForTarget(targetId: string): PtyConsumerRecovery {
const created = {
clientInstanceId: randomUUID(),
detached: false,
checkpointsByAppPtyId: new Map<string, SshPtyAcceptedSourceCheckpoint>()
checkpointsByAppPtyId: new Map<string, SshPtyAcceptedSourceCheckpoint>(),
modelMigrationsByAppPtyId: new Map<string, Promise<SshPtyOutputMigrationResult>>()
}
ptyConsumerRecoveryByTarget.set(targetId, created)
return created
@@ -878,6 +882,17 @@ export class SshRelaySession {
if (recovery) {
delete recovery.owner
recovery.checkpointsByAppPtyId.clear()
for (const [ptyId, migration] of recovery.modelMigrationsByAppPtyId) {
recovery.modelMigrationsByAppPtyId.set(
ptyId,
migration.then(() =>
Object.freeze({
status: 'checkpoint-unavailable' as const,
reason: 'completion-failed' as const
})
)
)
}
}
this.ptyConsumerSessionState = null
return openSshPtyConsumerSession(mux, options)
@@ -896,7 +911,10 @@ export class SshRelaySession {
serverBuildId,
owner,
checkpointsByAppPtyId:
previous?.checkpointsByAppPtyId ?? new Map<string, SshPtyAcceptedSourceCheckpoint>()
previous?.checkpointsByAppPtyId ?? new Map<string, SshPtyAcceptedSourceCheckpoint>(),
modelMigrationsByAppPtyId:
previous?.modelMigrationsByAppPtyId ??
new Map<string, Promise<SshPtyOutputMigrationResult>>()
})
}
@@ -1151,17 +1169,12 @@ export class SshRelaySession {
}
this.ptyRecoveryNotificationCleanups = []
if (this.activePtyProviderGeneration !== null) {
const providerGeneration = this.activePtyProviderGeneration
if (reason === 'connection_lost' && this.negotiatedPtyConsumerOwner()?.outputFlowControl) {
const recovery = ptyConsumerRecoveryByTarget.get(this.targetId)
if (recovery) {
for (const checkpoint of getSshPtyAcceptedSourceCheckpoints(
this.activePtyProviderGeneration
)) {
recovery.checkpointsByAppPtyId.set(checkpoint.id, checkpoint)
}
}
this.beginPtyModelMigration(providerGeneration, outputGenerationReason)
} else {
closeSshPtyOutputGeneration(providerGeneration, outputGenerationReason)
}
closeSshPtyOutputGeneration(this.activePtyProviderGeneration, outputGenerationReason)
this.activePtyProviderGeneration = null
}
this.sourceAckPublisherCleanup?.()
@@ -1728,7 +1741,7 @@ export class SshRelaySession {
let sourceActivationLease: SshPtyAttachResult['sourceActivationLease']
let recoveryActivationLease: SshPtyRecoveryActivationLease | undefined
try {
const recoveryRequest = this.sourceRecoveryRequest(appPtyId)
const recoveryRequest = await this.sourceRecoveryRequest(appPtyId)
const attachResult = await this.attachPtyWithRetry(
ptyProvider,
ptyId,
@@ -2025,11 +2038,24 @@ export class SshRelaySession {
}
}
private sourceRecoveryRequest(appPtyId: string): PtySourceRecoveryRequest | undefined {
private async sourceRecoveryRequest(
appPtyId: string
): Promise<PtySourceRecoveryRequest | undefined> {
if (!this.negotiatedPtyConsumerOwner()?.outputFlowControl) {
return undefined
}
const checkpoints = ptyConsumerRecoveryByTarget.get(this.targetId)?.checkpointsByAppPtyId
const recovery = ptyConsumerRecoveryByTarget.get(this.targetId)
const migration = recovery?.modelMigrationsByAppPtyId.get(appPtyId)
if (migration) {
const outcome = await migration
if (recovery?.modelMigrationsByAppPtyId.get(appPtyId) === migration) {
recovery.modelMigrationsByAppPtyId.delete(appPtyId)
}
if (outcome.status !== 'settled') {
return Object.freeze({ status: 'checkpointUnavailable' })
}
}
const checkpoints = recovery?.checkpointsByAppPtyId
const relayPtyId = toRelaySshPtyId(this.targetId, appPtyId)
const checkpoint = checkpoints?.get(appPtyId) ?? checkpoints?.get(relayPtyId)
if (!checkpoint) {
@@ -2045,6 +2071,39 @@ export class SshRelaySession {
})
}
private beginPtyModelMigration(providerGeneration: number, closeReason: string): void {
const recovery = ptyConsumerRecoveryByTarget.get(this.targetId)
if (!recovery) {
closeSshPtyOutputGeneration(providerGeneration, closeReason)
return
}
for (const checkpoint of getSshPtyAcceptedSourceCheckpoints(providerGeneration)) {
recovery.checkpointsByAppPtyId.set(checkpoint.id, checkpoint)
}
const migration = beginSshPtyOutputGenerationMigration(providerGeneration)
for (const [ptyId, result] of migration.byPty) {
const previous = recovery.modelMigrationsByAppPtyId.get(ptyId)
const fence = previous ? previous.then(() => result) : result
recovery.modelMigrationsByAppPtyId.set(ptyId, fence)
void fence.then((outcome) => {
const current = ptyConsumerRecoveryByTarget.get(this.targetId)
if (current?.modelMigrationsByAppPtyId.get(ptyId) !== fence) {
return
}
if (outcome.status === 'settled') {
current.checkpointsByAppPtyId.set(ptyId, outcome.checkpoint)
} else {
current.checkpointsByAppPtyId.delete(ptyId)
current.checkpointsByAppPtyId.delete(toRelaySshPtyId(this.targetId, ptyId))
}
current.modelMigrationsByAppPtyId.delete(ptyId)
})
}
void migration.completion.then(() => {
closeSshPtyOutputGeneration(providerGeneration, closeReason)
})
}
private async finishSourceRecovery(
relayPtyId: string,
appPtyId: string,