From 3b83f6e63b87f5753ccfb05b66a99a45c7e6826f Mon Sep 17 00:00:00 2001 From: Priya Singh Date: Mon, 10 Aug 2026 03:40:49 +0000 Subject: [PATCH] fix(runtime): update shell session wire protocol --- src/runtime/shell/__tests__/protocol.test.ts | 8 - src/runtime/shell/__tests__/session.test.ts | 206 ++++++------------ src/runtime/shell/config.ts | 10 +- src/runtime/shell/protocol.ts | 5 - src/runtime/shell/session.ts | 214 ++----------------- 5 files changed, 86 insertions(+), 357 deletions(-) diff --git a/src/runtime/shell/__tests__/protocol.test.ts b/src/runtime/shell/__tests__/protocol.test.ts index f3b6f34..0124adb 100644 --- a/src/runtime/shell/__tests__/protocol.test.ts +++ b/src/runtime/shell/__tests__/protocol.test.ts @@ -160,11 +160,3 @@ describe('ShellFramer.encodeHeartbeat', () => { expect(frame).toEqual(Buffer.from([ShellChannel.HEARTBEAT])) }) }) - -describe('ShellFramer.encodeClose', () => { - it('encodes close as single byte', () => { - const framer = new ShellFramer() - const frame = framer.encodeClose() - expect(frame).toEqual(Buffer.from([ShellChannel.CLOSE])) - }) -}) diff --git a/src/runtime/shell/__tests__/session.test.ts b/src/runtime/shell/__tests__/session.test.ts index f1a0150..540e04a 100644 --- a/src/runtime/shell/__tests__/session.test.ts +++ b/src/runtime/shell/__tests__/session.test.ts @@ -150,30 +150,15 @@ describe('ShellSession: construction', () => { it('initial attributes are correct', () => { const session = new ShellSession(makeOpts([])) - expect(session.reconnected).toBe(false) expect(session.kicked).toBe(false) - expect(session.bytesDropped).toBe(0) expect(session.exitCode).toBeNull() }) }) describe('ShellSession: connect', () => { - it('reads shellId and reconnected=false from STATUS confirmation frame', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('srv-shell', false)])) - await session.connect() - expect(session.shellId).toBe('srv-shell') - expect(session.reconnected).toBe(false) - }) - - it('sets reconnected=true from STATUS frame', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s', true)])) - await session.connect() - expect(session.reconnected).toBe(true) - }) - it('reads shellId and sessionId from 101 upgrade headers', async () => { const session = new ShellSession( - makeOpts([confirmationFrame('hdr-shell')], { + makeOpts([], { headers: { 'x-amzn-bedrock-agentcore-shell-id': 'hdr-shell', 'x-amzn-bedrock-agentcore-runtime-session-id': 'hdr-session', @@ -184,49 +169,11 @@ describe('ShellSession: connect', () => { expect(session.shellId).toBe('hdr-shell') expect(session.sessionId).toBe('hdr-session') }) - - it('stashes non-STATUS frames received before confirmation', async () => { - const session = new ShellSession(makeOpts([stdoutFrame('early'), confirmationFrame('s')])) - await session.connect() - const frames = [] - for await (const f of session) { - frames.push(f) - if (frames.length === 1) break - } - expect(frames[0]!.channel).toBe(ShellChannel.STDOUT) - expect(frames[0]!.text).toBe('early') - }) - - it('proceeds with warning when STATUS confirmation times out', async () => { - const warnSpy = vi.spyOn(console, 'warn') - - let capturedWs: MockWs | null = null - const wsFactory = (_url: string): WebSocket => { - const ws = makeMockWs() - capturedWs = ws - return ws as unknown as WebSocket - } - const connectFn: ConnectFn = vi.fn(async (_shellId, _sessionId) => { - process.nextTick(() => { - capturedWs!.emit('upgrade', { headers: {} }) - capturedWs!.emit('open') - // No frames emitted — metadata timeout fires after DEFAULT_METADATA_TIMEOUT ms - }) - return { url: 'wss://test.local/runtimes/x/ws/shells', headers: {} } - }) - - const session = new ShellSession({ connectFn, _wsFactory: wsFactory, logger: console }) - // AbortSignal.timeout() uses real timers — no fake timers needed, just await with real timeout - await session.connect() - expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('imed out')) - }, 15_000) }) describe('ShellSession: iterator', () => { it('yields STDOUT frames in order', async () => { - const session = new ShellSession( - makeOpts([confirmationFrame('s'), stdoutFrame('hello'), stdoutFrame('world'), closeFrame()]) - ) + const session = new ShellSession(makeOpts([stdoutFrame('hello'), stdoutFrame('world'), closeFrame()])) await session.connect() const texts: string[] = [] for await (const f of session) { @@ -236,9 +183,7 @@ describe('ShellSession: iterator', () => { }) it('swallows HEARTBEAT frames — never yields to caller', async () => { - const session = new ShellSession( - makeOpts([confirmationFrame('s'), heartbeatFrame(), stdoutFrame('after-hb'), closeFrame()]) - ) + const session = new ShellSession(makeOpts([heartbeatFrame(), stdoutFrame('after-hb'), closeFrame()])) await session.connect() const channels: ShellChannel[] = [] for await (const f of session) channels.push(f.channel) @@ -247,7 +192,7 @@ describe('ShellSession: iterator', () => { }) it('stops on CLOSE frame', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), stdoutFrame('x'), closeFrame()])) + const session = new ShellSession(makeOpts([stdoutFrame('x'), closeFrame()])) await session.connect() let count = 0 for await (const _ of session) count++ @@ -255,7 +200,7 @@ describe('ShellSession: iterator', () => { }) it('stops on WebSocket close 1000', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), stdoutFrame('x')], { closeCode: 1000 })) + const session = new ShellSession(makeOpts([stdoutFrame('x')], { closeCode: 1000 })) await session.connect() let count = 0 for await (const _ of session) count++ @@ -263,7 +208,7 @@ describe('ShellSession: iterator', () => { }) it('sets kicked=true on close code 4000', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')], { closeCode: 4000 })) + const session = new ShellSession(makeOpts([], { closeCode: 4000 })) await session.connect() for await (const _ of session) { /* drain */ @@ -273,9 +218,7 @@ describe('ShellSession: iterator', () => { it('stops on close code 1003 without reconnecting and warns', async () => { const warnSpy = vi.spyOn(console, 'warn') - // reconnectConfig is present so that if 1003 incorrectly fell through to the - // reconnect path it would trigger a retry; we assert it does not. - const opts = makeOpts([confirmationFrame('s'), stdoutFrame('x')], { closeCode: 1003 }) + const opts = makeOpts([stdoutFrame('x')], { closeCode: 1003 }) const session = new ShellSession({ ...opts, reconnectConfig: { maxRetries: 1, baseDelay: 0 }, logger: console }) await session.connect() let count = 0 @@ -283,12 +226,11 @@ describe('ShellSession: iterator', () => { expect(count).toBe(1) expect((session as unknown as { _state: { status: string } })._state.status).not.toBe('open') expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining('1003')) - // connectFn called exactly once — reconnect loop was not entered expect((opts.connectFn as ReturnType).mock.calls.length).toBe(1) }) it('stops cleanly when close() is called mid-iteration', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), stdoutFrame('x')])) + const session = new ShellSession(makeOpts([stdoutFrame('x')])) await session.connect() let count = 0 for await (const _ of session) { @@ -302,7 +244,7 @@ describe('ShellSession: iterator', () => { describe('ShellSession: exitCode', () => { it('is null during STDOUT frames, set after termination STATUS', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), stdoutFrame('x'), exitFrame(0)])) + const session = new ShellSession(makeOpts([stdoutFrame('x'), exitFrame(0)])) await session.connect() let midLoopCode: number | null | undefined for await (const f of session) { @@ -313,7 +255,7 @@ describe('ShellSession: exitCode', () => { }) it('exitCode=0 for clean exit (status=Success)', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), exitFrame(0)])) + const session = new ShellSession(makeOpts([exitFrame(0)])) await session.connect() for await (const _ of session) { /* drain */ @@ -322,7 +264,7 @@ describe('ShellSession: exitCode', () => { }) it('exitCode reflects non-zero exit', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), exitFrame(42)])) + const session = new ShellSession(makeOpts([exitFrame(42)])) await session.connect() for await (const _ of session) { /* drain */ @@ -339,70 +281,32 @@ describe('ShellSession: exitCode', () => { reason: 'InternalError', code: 500, }) - const session = new ShellSession(makeOpts([confirmationFrame('s'), errFrame])) + const session = new ShellSession(makeOpts([errFrame])) await session.connect() for await (const _ of session) { /* drain */ } expect(session.exitCode).toBeNull() }) - - it('exitCode set via pending-frames drain path', async () => { - // STDOUT arrives before STATUS confirmation → gets stashed; exitCode set when STATUS drained - const session = new ShellSession(makeOpts([stdoutFrame('stashed'), confirmationFrame('s'), exitFrame(5)])) - await session.connect() - const channels: ShellChannel[] = [] - for await (const f of session) channels.push(f.channel) - expect(session.exitCode).toBe(5) - expect(channels[0]).toBe(ShellChannel.STDOUT) // pending frame drained first - }) -}) - -describe('ShellSession: bytesDropped', () => { - it('set from second confirmation frame with bytesDropped field', async () => { - const secondConf = statusFrame({ - kind: 'Status', - apiVersion: 'v1', - metadata: { shellId: 's', reconnected: true, bytesDropped: 512 }, - status: 'Success', - }) - const session = new ShellSession( - makeOpts([confirmationFrame('s'), stdoutFrame('output'), secondConf, closeFrame()]) - ) - await session.connect() - const channels: ShellChannel[] = [] - for await (const f of session) channels.push(f.channel) - expect(session.bytesDropped).toBe(512) - expect(channels).toEqual([ShellChannel.STDOUT]) // second confirmation swallowed - }) - - it('stays 0 when no ring-buffer overflow', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s'), closeFrame()])) - await session.connect() - for await (const _ of session) { - /* drain */ - } - expect(session.bytesDropped).toBe(0) - }) }) describe('ShellSession: close()', () => { it('transitions to closed state after close()', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')])) + const session = new ShellSession(makeOpts([])) await session.connect() await session.close() expect((session as unknown as { _state: { status: string } })._state.status).toBe('closed') }) it('is idempotent — calling twice does not throw', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')])) + const session = new ShellSession(makeOpts([])) await session.connect() await session.close() await expect(session.close()).resolves.toBeUndefined() }) it('send() throws after close()', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')])) + const session = new ShellSession(makeOpts([])) await session.connect() await session.close() await expect(session.send('hello\n')).rejects.toThrow() @@ -413,7 +317,7 @@ describe('ShellSession: disconnect handling', () => { // Build a session whose socket we can drive frame-by-frame, with full control over // when frames/close arrive — so we can drop the connection while NOT iterating. function manualOpts( - opts: { reconnect?: boolean; onReconnect?: (reconnected: boolean) => void | Promise } = {} + opts: { reconnect?: boolean; onReconnect?: () => void | Promise } = {} ): ShellSessionOptions & { sockets: MockWs[] confirm: () => void @@ -425,13 +329,10 @@ describe('ShellSession: disconnect handling', () => { return ws as unknown as WebSocket } const connectFn: ConnectFn = vi.fn(async () => { - // Read the socket lazily inside nextTick — wsFactory runs AFTER connectFn returns. process.nextTick(() => { const ws = sockets[sockets.length - 1]! - const isReconnect = sockets.length > 1 ws.emit('upgrade', { headers: {} }) ws.emit('open') - process.nextTick(() => ws.emit('message', confirmationFrame('s', isReconnect))) }) return { url: 'wss://test.local/runtimes/x/ws/shells', headers: {} } }) @@ -676,7 +577,7 @@ describe('ShellSession: disconnect handling', () => { describe('ShellSession: keepalive', () => { it('calls ws.ping() on the interval and stops after close()', async () => { vi.useFakeTimers() - const opts = makeOpts([confirmationFrame('s')]) + const opts = makeOpts([]) const session = new ShellSession({ ...opts, keepaliveIntervalMs: 1000 }) await session.connect() @@ -701,7 +602,7 @@ describe('ShellSession: keepalive', () => { it('keepalive disabled when keepaliveIntervalMs=0', async () => { vi.useFakeTimers() - const opts = makeOpts([confirmationFrame('s')]) + const opts = makeOpts([]) const session = new ShellSession({ ...opts, keepaliveIntervalMs: 0 }) await session.connect() @@ -714,7 +615,7 @@ describe('ShellSession: keepalive', () => { it('keepalive timer is stopped when iterator finishes naturally', async () => { vi.useFakeTimers() - const opts = makeOpts([confirmationFrame('s')], { closeCode: 1000 }) + const opts = makeOpts([], { closeCode: 1000 }) const session = new ShellSession({ ...opts, keepaliveIntervalMs: 500 }) await session.connect() @@ -733,14 +634,14 @@ describe('ShellSession: keepalive', () => { describe('ShellSession: connect() state guards', () => { it('throws when called on an already-closed session', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')])) + const session = new ShellSession(makeOpts([])) await session.connect() await session.close() await expect(session.connect()).rejects.toThrow('closed') }) it('throws when called while already connecting/open', async () => { - const session = new ShellSession(makeOpts([confirmationFrame('s')])) + const session = new ShellSession(makeOpts([])) await session.connect() await expect(session.connect()).rejects.toThrow('idle') }) @@ -762,21 +663,6 @@ describe('ShellSession: connect() state guards', () => { }) }) -describe('ShellSession: HEARTBEAT in pendingFrames', () => { - it('does not yield a HEARTBEAT that arrived before STATUS confirmation', async () => { - // HEARTBEAT arrives before the confirmation STATUS — it ends up in pendingFrames. - // The pendingFrames drain must filter it out just like the main loop does. - const session = new ShellSession( - makeOpts([heartbeatFrame(), confirmationFrame('s'), stdoutFrame('hi'), closeFrame()]) - ) - await session.connect() - const channels: ShellChannel[] = [] - for await (const f of session) channels.push(f.channel) - expect(channels).not.toContain(ShellChannel.HEARTBEAT) - expect(channels).toContain(ShellChannel.STDOUT) - }) -}) - describe('ShellSession: reconnectWindow null = unlimited', () => { it('does not expire and reconnects when reconnectWindow is null', async () => { let callCount = 0 @@ -795,14 +681,11 @@ describe('ShellSession: reconnectWindow null = unlimited', () => { ws.emit('upgrade', { headers: {} }) ws.emit('open') process.nextTick(() => { - ws.emit('message', confirmationFrame('s')) - process.nextTick(() => { - if (thisCallCount === 1) { - ws.emit('close', 1006, Buffer.from('')) // abnormal → triggers reconnect - } else { - ws.emit('message', closeFrame()) // clean close on second call - } - }) + if (thisCallCount === 1) { + ws.emit('close', 1006, Buffer.from('')) // abnormal → triggers reconnect + } else { + ws.emit('message', closeFrame()) // clean close on second call + } }) }) return { url: 'wss://test.local/runtimes/x/ws/shells', headers: {} } @@ -821,3 +704,40 @@ describe('ShellSession: reconnectWindow null = unlimited', () => { expect(callCount).toBeGreaterThanOrEqual(2) }) }) + +describe('ShellSession: confirmation frame swallowed during iteration', () => { + it('swallows a confirmation STATUS frame and does not yield it to caller', async () => { + const session = new ShellSession(makeOpts([confirmationFrame('s'), stdoutFrame('after-conf'), closeFrame()])) + await session.connect() + const channels: ShellChannel[] = [] + for await (const f of session) channels.push(f.channel) + // Only STDOUT should be yielded — confirmation frame is swallowed + expect(channels).toEqual([ShellChannel.STDOUT]) + }) +}) + +describe('ShellSession: connect ready immediately', () => { + it('connect resolves without waiting for any inbound frames', async () => { + // No frames at all — connect should still succeed immediately + let capturedWs: MockWs | null = null + const wsFactory = (_url: string): WebSocket => { + const ws = makeMockWs() + capturedWs = ws + return ws as unknown as WebSocket + } + const connectFn: ConnectFn = vi.fn(async () => { + process.nextTick(() => { + capturedWs!.emit('upgrade', { + headers: { 'x-amzn-bedrock-agentcore-shell-id': 'instant-shell' }, + }) + capturedWs!.emit('open') + // No frames emitted at all — connect should still complete + }) + return { url: 'wss://test.local/runtimes/x/ws/shells', headers: {} } + }) + + const session = new ShellSession({ connectFn, _wsFactory: wsFactory }) + await session.connect() + expect(session.shellId).toBe('instant-shell') + }) +}) diff --git a/src/runtime/shell/config.ts b/src/runtime/shell/config.ts index ae86fb6..8b59457 100644 --- a/src/runtime/shell/config.ts +++ b/src/runtime/shell/config.ts @@ -32,7 +32,6 @@ export const noopLogger: Logger = { export const DEFAULT_MAX_RETRIES = 5 export const DEFAULT_BASE_DELAY = 1000 // ms export const DEFAULT_MAX_DELAY = 15000 // ms -export const DEFAULT_METADATA_TIMEOUT = 10_000 // ms export const DEFAULT_RECONNECT_WINDOW = 900_000 // ms — ~15 min, matches server-side KARP idle timeout export const DEFAULT_OUTER_LOOP_DELAY = 30_000 // ms export const DEFAULT_KEEPALIVE_INTERVAL = 30_000 // ms — KARP idle timeout is ~60s; ping every 30s @@ -49,9 +48,8 @@ export const DEFAULT_KEEPALIVE_INTERVAL = 30_000 // ms — KARP idle timeout is * const config: ReconnectConfig = { * maxRetries: 5, * reconnectWindow: null, // unlimited - * onReconnect: async (reconnected) => { - * if (reconnected) console.log('Reattached — buffered output will follow') - * else console.log('New shell started') + * onReconnect: async () => { + * console.log('Reconnected to shell') * } * } * const shell = await client.openShell(runtimeArn, { reconnectConfig: config }) @@ -86,8 +84,6 @@ export interface ReconnectConfig { /** * Optional callback invoked after each successful reconnect. - * Receives `reconnected: true` when the existing PTY was reattached; - * `false` when a fresh shell was started. */ - onReconnect?: (reconnected: boolean) => void | Promise + onReconnect?: () => void | Promise } diff --git a/src/runtime/shell/protocol.ts b/src/runtime/shell/protocol.ts index aa328f2..3d9012b 100644 --- a/src/runtime/shell/protocol.ts +++ b/src/runtime/shell/protocol.ts @@ -117,9 +117,4 @@ export class ShellFramer { encodeHeartbeat(): Buffer { return Buffer.from([ShellChannel.HEARTBEAT]) } - - /** Encode a graceful-shutdown CLOSE frame (empty payload). */ - encodeClose(): Buffer { - return Buffer.from([ShellChannel.CLOSE]) - } } diff --git a/src/runtime/shell/session.ts b/src/runtime/shell/session.ts index 2b0166a..29ad96e 100644 --- a/src/runtime/shell/session.ts +++ b/src/runtime/shell/session.ts @@ -1,8 +1,8 @@ /** * ShellSession — async-iterable interactive PTY WebSocket session. * - * Connects on `connect()`, reads the initial STATUS confirmation frame, and exposes - * typed `send()` / `resize()` / `[Symbol.asyncIterator]()` / `close()`. + * Connects on `connect()` and exposes typed `send()` / `resize()` / + * `[Symbol.asyncIterator]()` / `close()`. * * When `reconnectConfig` is provided, transparently reconnects on unexpected disconnects * using the same `shellId` so the shell's working directory, environment, background jobs, @@ -11,10 +11,9 @@ * Reconnect restores the *connection* on its own (it is driven by the socket close event, * not by your read loop, and `send()`/`resize()` wait for it). However, on reattach the * server replays the buffered output as inbound frames — to receive that replay (and to see - * `bytesDropped` updated and `exitCode` set) you must be consuming the session with - * `for await (const frame of shell)`. A write-only caller that never iterates stays - * connected across drops but will not observe the replayed output. Keep a `for await` loop - * running for the life of the session. + * `exitCode` set) you must be consuming the session with `for await (const frame of shell)`. + * A write-only caller that never iterates stays connected across drops but will not observe + * the replayed output. Keep a `for await` loop running for the life of the session. * * @example * ```typescript @@ -43,7 +42,6 @@ import { DEFAULT_KEEPALIVE_INTERVAL, DEFAULT_MAX_DELAY, DEFAULT_MAX_RETRIES, - DEFAULT_METADATA_TIMEOUT, DEFAULT_OUTER_LOOP_DELAY, DEFAULT_RECONNECT_WINDOW, noopLogger, @@ -99,7 +97,7 @@ export interface ShellSessionOptions { // // Transitions: // idle → connecting connect() called -// connecting → open upgrade + metadata handshake succeeded +// connecting → open upgrade succeeded // connecting → closed close() called during upgrade // connecting → (throws) upgrade failed; caller retries or surfaces error // open → reconnecting WS dropped; _iterate catches the error @@ -120,7 +118,6 @@ type SessionState = ws: WebSocket messageIterator: AsyncIterableIterator keepaliveTimer: ReturnType | null - pendingFrames: ShellFrame[] } | { status: 'reconnecting' } | { status: 'closed' } @@ -131,22 +128,16 @@ type SessionState = * Read-only observable attributes (updated by the session as events arrive): * - `shellId` — Server-confirmed shell identifier. Preserve to reconnect to the same PTY. * - `sessionId` — Runtime session ID routing to the VM. - * - `reconnected` — True when the most recent connect reattached an existing PTY. * - `kicked` — True when another client connected with the same shellId (close 4000). * Check this after the `for await` loop exits to distinguish a kick from * a clean shell exit. - * - `bytesDropped` — PTY ring-buffer bytes lost during the most recent disconnect, as - * reported by the server in the reconnect confirmation frame. - * Zero if no overflow occurred or on a fresh connection. * - `exitCode` — Shell process exit code. `null` until the shell exits; `0` for a clean * exit. Check this after the `for await` loop exits alongside `kicked`. */ export class ShellSession implements AsyncIterable { private _shellId: string private _sessionId: string - private _reconnected = false private _kicked = false - private _bytesDropped = 0 private _exitCode: number | null = null /** Server-confirmed shell identifier. */ @@ -157,10 +148,6 @@ export class ShellSession implements AsyncIterable { get sessionId(): string { return this._sessionId } - /** True when the most recent connect reattached an existing PTY. */ - get reconnected(): boolean { - return this._reconnected - } /** * True when another client connected with the same shellId (close 4000). * Check after the `for await` loop exits to distinguish a kick from a clean exit. @@ -168,13 +155,6 @@ export class ShellSession implements AsyncIterable { get kicked(): boolean { return this._kicked } - /** - * PTY ring-buffer bytes lost during the most recent disconnect. - * Zero when no overflow occurred or on a fresh connection. - */ - get bytesDropped(): number { - return this._bytesDropped - } /** * Shell process exit code. `null` until the shell exits; `0` for a clean exit. * Check after the `for await` loop exits alongside `kicked`. @@ -280,9 +260,10 @@ export class ShellSession implements AsyncIterable { return this._state.ws } - /** Send a CLOSE frame (0xFF) to permanently kill the shell, then close the WebSocket. - * The server kills the shell process (SIGHUP → SIGKILL) and responds with its own [0xFF]. - * Unlike dropping the WebSocket (which detaches and allows reconnection), this is permanent. */ + /** Disconnect from the shell session by closing the WebSocket. + * The shell process stays alive on the server for the reconnect window, allowing + * later reconnection with the same shellId. Unlike the old behavior (which sent 0xFF + * to permanently kill the shell), this just detaches the client. */ async close(): Promise { const prev = this._state if (prev.status === 'closed') { @@ -304,15 +285,6 @@ export class ShellSession implements AsyncIterable { } } else if (prev.status === 'open') { this._stopKeepalive(prev.keepaliveTimer) - // ws.send() throws synchronously if the socket is already closing/closed. - // Swallow it — the intent is best-effort notification, not guaranteed delivery. - try { - prev.ws.send(this.framer.encodeClose()) - } catch (err) { - this.log.debug( - `ShellSession: CLOSE frame not sent — socket already closing/closed (shellId=${this.shellId}): ${String(err)}` - ) - } try { prev.ws.close() } catch (err) { @@ -335,14 +307,13 @@ export class ShellSession implements AsyncIterable { * Async iterator — yields inbound ShellFrames, reconnecting on drop if configured. * * The loop exits silently (no throw) in three cases: shell exit, kicked by a new - * client, or reconnect budget exhausted. Check `exitCode`, `kicked`, and - * `bytesDropped` after the loop to distinguish them: + * client, or reconnect budget exhausted. Check `exitCode` and `kicked` + * after the loop to distinguish them: * * ```typescript * for await (const frame of shell) { ... } * if (shell.kicked) { ... } // another client took over * if (shell.exitCode !== null) { ... } // shell process exited - * if (shell.bytesDropped > 0) { ... } // ring-buffer overflow on reconnect * ``` */ [Symbol.asyncIterator](): AsyncIterator { @@ -407,15 +378,9 @@ export class ShellSession implements AsyncIterable { this._abortController?.abort() // A reconnect attempt enters here with status 'reconnecting'; a fresh connect with 'idle'. - const isReconnect = this._state.status === 'reconnecting' this._closeError = null - this._reconnected = false this._kicked = false this._exitCode = null - // bytesDropped reports loss for the most recent disconnect (per-disconnect, not - // cumulative). Reset it only on a FRESH connect — resetting on every reconnect attempt - // would clobber a value the iterator has not yet read when reconnects happen back-to-back. - if (!isReconnect) this._bytesDropped = 0 this._state = { status: 'connecting', ws: null } let connectResult: { url: string; headers: Record; protocols?: string[] } @@ -446,11 +411,7 @@ export class ShellSession implements AsyncIterable { // their own controller — not this._abortController, which is replaced on reconnect. this._abortController = new AbortController() const controller = this._abortController - // Single iterator used for both the STATUS handshake and subsequent frame reads. - // On the timeout path, one frame may be lost (the abandoned .next() from the - // Promise.race in _readMetadataFrame consumes it), but that is a degraded scenario. - // Using two independent iterators would cause every pre-STATUS frame to be yielded - // twice (once via pendingFrames, once via the independent listener queue). + // Single iterator for all frame reads after the WebSocket opens. const messageIterator = on(ws, 'message', { signal: controller.signal }) as AsyncIterableIterator ws.on('close', (code: number, reason: Buffer) => { @@ -514,25 +475,10 @@ export class ShellSession implements AsyncIterable { openRaceAc.abort() } - let pendingFrames: ShellFrame[] - try { - pendingFrames = await this._readMetadataFrame(messageIterator) - } catch (err) { - // Server closed before sending STATUS, or close() fired during handshake. - ws.terminate() - if (!this._isClosed()) this._state = { status: 'idle' } - throw err - } - - // close() may have fired during _readMetadataFrame — terminate and bail out. - if (this._isClosed()) { - ws.terminate() - return - } - const keepaliveTimer = this._startKeepalive(ws, controller.signal) - // Atomic promotion: all connection objects become available together. - this._state = { status: 'open', ws, messageIterator, keepaliveTimer, pendingFrames } + // Connection is ready immediately after WebSocket opens — shellId comes from 101 header. + // No longer blocking on the 0x03 confirmation frame. + this._state = { status: 'open', ws, messageIterator, keepaliveTimer } } /** Receive one raw binary message from the WebSocket. */ @@ -549,95 +495,11 @@ export class ShellSession implements AsyncIterable { } } - /** - * Consume frames until a STATUS confirmation is found, stashing others in pendingFrames. - * Returns the accumulated pending frames to be stored in the 'open' state. - */ - private async _readMetadataFrame(messageIterator: AsyncIterableIterator): Promise { - // Use an explicit AbortController so the timer can be cancelled as soon as - // STATUS arrives — AbortSignal.timeout() creates a timer that outlives the - // fast-path read and accumulates orphan wakeups on rapid reconnects. - const timeoutAc = new AbortController() - const timer = globalThis.setTimeout(() => timeoutAc.abort(new Error('timeout')), DEFAULT_METADATA_TIMEOUT) - const timeoutP = new Promise((_, rej) => - timeoutAc.signal.addEventListener('abort', () => rej(timeoutAc.signal.reason as Error), { once: true }) - ) - - const pendingFrames: ShellFrame[] = [] - - try { - while (true) { - let raw: Buffer - try { - raw = await Promise.race([this._recvRaw(messageIterator), timeoutP]) - } catch (err: unknown) { - if (timeoutAc.signal.aborted) { - // If the WebSocket also closed concurrently (race between the 10s timer - // and a server close event), prefer the real close error — promoting a - // dead messageIterator to 'open' would mislead callers and lose the cause. - if (this._closeError !== null) throw this._closeError - this.log.warn(`ShellSession: Timed out waiting for STATUS confirmation (shellId=${this.shellId})`) - return pendingFrames - } - throw err - } - - const frame = this.framer.decode(raw) - - if (frame.channel === ShellChannel.STATUS) { - try { - const meta = (frame.json()['metadata'] ?? {}) as Record - if (meta['shellId']) { - // Confirmation frame — update shellId and we're done. On a reconnect this is - // the single reconnection confirmation ("sent once, before replay begins"), - // so read bytesDropped here — it is delivered on THIS frame and nowhere else. - this._shellId = String(meta['shellId']) - this._reconnected = Boolean(meta['reconnected']) - this._recordBytesDropped(meta) - return pendingFrames - } - } catch (err) { - this.log.debug( - `ShellSession: malformed STATUS frame, proceeding with client-generated shellId=${this.shellId}: ${String(err)}` - ) - return pendingFrames - } - // No shellId → termination frame (shell died before confirmation). - // Stash it so _iterate can set exitCode and return cleanly. - this.log.debug(`ShellSession: termination STATUS received before confirmation (shellId=${this.shellId})`) - pendingFrames.push(frame) - return pendingFrames - } - - pendingFrames.push(frame) - } - } finally { - globalThis.clearTimeout(timer) - } - } - private _isConfirmationStatus(status: Record): boolean { const meta = status['metadata'] as Record | undefined return Boolean(meta?.['shellId']) } - /** - * Record `bytesDropped` from a reconnection confirmation frame's metadata, if present. - * `bytesDropped` reports PTY output lost from ring-buffer overflow during THIS disconnect - * (per-disconnect, not session-cumulative — assign, don't accumulate). Present - * only when greater than 0; absent on a clean reconnect. - */ - private _recordBytesDropped(meta: Record | undefined): void { - const dropped = meta?.['bytesDropped'] - if (typeof dropped === 'number' && dropped > 0) { - this._bytesDropped = dropped - this.log.warn( - `ShellSession: ${dropped} bytes of PTY output lost during disconnect ` + - `(ring buffer overflow, shellId=${this.shellId})` - ) - } - } - private _isTerminationStatus(status: Record): boolean { if (this._isConfirmationStatus(status)) return false return status['status'] === 'Success' || status['status'] === 'Failure' @@ -670,37 +532,6 @@ export class ShellSession implements AsyncIterable { // concurrently while the generator is suspended at a yield or await point. state = this._state - // Drain frames buffered during the metadata handshake first. - while (state.pendingFrames.length > 0) { - if (this._isClosed()) return - const frame = state.pendingFrames.shift()! - // HEARTBEAT — server echo of client keepalive, not application data. - if (frame.channel === ShellChannel.HEARTBEAT) continue - if (frame.channel === ShellChannel.CLOSE) { - this._state = { status: 'idle' } - this.log.debug(`ShellSession: CLOSE frame received in pending queue (shellId=${this.shellId})`) - return - } - if (frame.channel === ShellChannel.STATUS) { - try { - const s = frame.json() - if (this._isTerminationStatus(s)) { - this._exitCode = this._parseExitCode(s) - this._state = { status: 'idle' } - yield frame - return - } - } catch (err) { - this.log.debug( - `ShellSession: malformed STATUS frame in pending queue (shellId=${this.shellId}): ${String(err)}` - ) - } - } - yield frame - } - - if (this._state.status !== 'open') return - let raw: Buffer try { raw = await this._recvRaw(state.messageIterator) @@ -745,12 +576,7 @@ export class ShellSession implements AsyncIterable { try { const s = frame.json() if (this._isConfirmationStatus(s)) { - // A confirmation frame in the data stream. Per spec §6 the reconnection - // confirmation (carrying bytesDropped) is sent once *before* replay and is - // consumed by _readMetadataFrame, not here — so this path normally does not - // fire on reconnect. Record bytesDropped defensively in case a confirmation - // reaches the stream, then swallow it (not application output). - this._recordBytesDropped(s['metadata'] as Record | undefined) + // Confirmation frame — silently swallow (not application output). continue } if (this._isTerminationStatus(s)) { @@ -798,7 +624,7 @@ export class ShellSession implements AsyncIterable { private async _waitForClose(ws: WebSocket): Promise { if (this._closeError !== null || this._isClosed()) return const ac = new AbortController() - const timer = globalThis.setTimeout(() => ac.abort(), DEFAULT_METADATA_TIMEOUT) + const timer = globalThis.setTimeout(() => ac.abort(), 10_000) try { await once(ws, 'close', { signal: ac.signal }) } catch { @@ -943,10 +769,10 @@ export class ShellSession implements AsyncIterable { // close() may have fired while connectFn was awaiting — _connectWithUpgrade returns void // rather than throwing in that case, so guard here before claiming a successful reconnect. if (this._isClosed()) return false - this.log.info(`ShellSession: Reconnected (reconnected=${this.reconnected}, shellId=${this.shellId})`) + this.log.info(`ShellSession: Reconnected (shellId=${this.shellId})`) if (cfg.onReconnect) { try { - await cfg.onReconnect(this.reconnected) + await cfg.onReconnect() } catch (err) { this.log.warn(`ShellSession: onReconnect callback threw (shellId=${this.shellId}): ${String(err)}`) }