mirror of
https://github.com/stablyai/orca.git
synced 2026-10-19 16:03:26 +00:00
fix(ssh): fence model admission across reconnect
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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>(() => {}))
|
||||
|
||||
@@ -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 }
|
||||
}
|
||||
@@ -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 () => {
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user