diff --git a/src/cli/aws/__tests__/connect-shell.test.ts b/src/cli/aws/__tests__/connect-shell.test.ts index 3f20118e0..2b90bddc9 100644 --- a/src/cli/aws/__tests__/connect-shell.test.ts +++ b/src/cli/aws/__tests__/connect-shell.test.ts @@ -1,6 +1,5 @@ import { ShellKickedError } from '../../../lib/errors/types.js'; import { buildShellUrl, connectShell, startKeepalive } from '../connect-shell.js'; -import { ShellChannel } from '../shell-framer.js'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; // --------------------------------------------------------------------------- @@ -19,14 +18,14 @@ vi.mock('../account', () => ({ const wsState = vi.hoisted(() => { return { calls: [] as string[], - messageHandler: undefined as ((data: Buffer) => void) | undefined, + openHandler: undefined as (() => void) | undefined, closeHandler: undefined as ((code: number) => void) | undefined, errorHandler: undefined as ((err: Error) => void) | undefined, upgradeHandler: undefined as ((response: { headers: Record }) => void) | undefined, terminateCalled: false, reset() { this.calls = []; - this.messageHandler = undefined; + this.openHandler = undefined; this.closeHandler = undefined; this.errorHandler = undefined; this.upgradeHandler = undefined; @@ -41,7 +40,7 @@ vi.mock('ws', () => ({ wsState.calls.push(url); } on(event: string, handler: (...args: unknown[]) => void) { - if (event === 'message') wsState.messageHandler = handler as (data: Buffer) => void; + if (event === 'open') wsState.openHandler = handler as () => void; if (event === 'close') wsState.closeHandler = handler as (code: number) => void; if (event === 'error') wsState.errorHandler = handler as (err: Error) => void; if (event === 'upgrade') @@ -107,133 +106,122 @@ describe('buildShellUrl', () => { }); // --------------------------------------------------------------------------- -// connectShell +// connectShell — immediate connect (no confirmation frame wait) // --------------------------------------------------------------------------- describe('connectShell', () => { - function makeConfirmationFrame(shellId: string, reconnected = false): Buffer { - const payload = JSON.stringify({ - kind: 'Status', - apiVersion: 'v1', - metadata: { shellId, reconnected }, - status: 'Success', - }); - return Buffer.concat([Buffer.from([ShellChannel.STATUS]), Buffer.from(payload)]); - } - beforeEach(() => { wsState.reset(); }); - it('resolves with shellId from X-Amzn-Bedrock-AgentCore-Shell-Id 101 header (primary)', async () => { + it('resolves with shellId from X-Amzn-Bedrock-AgentCore-Shell-Id 101 header', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - // Header fires first (101 upgrade), then STATUS frame arrives + // Header fires first (101 upgrade), then open event wsState.upgradeHandler?.({ headers: { 'x-amzn-bedrock-agentcore-shell-id': 'header-shell-id' } }); - wsState.messageHandler?.(makeConfirmationFrame('frame-shell-id')); + wsState.openHandler?.(); const conn = await connectPromise; - // Header takes precedence over STATUS frame expect(conn.shellId).toBe('header-shell-id'); }); - it('falls back to shellId from STATUS frame when header is absent', async () => { + it('uses provided shellId as fallback when header is absent', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', + shellId: 'my-fallback-shell', }); await new Promise(r => setTimeout(r, 0)); - // No upgrade event fired — STATUS frame is the only source - wsState.messageHandler?.(makeConfirmationFrame('frame-shell-id')); + // No upgrade event — open fires without header + wsState.openHandler?.(); const conn = await connectPromise; - expect(conn.shellId).toBe('frame-shell-id'); + expect(conn.shellId).toBe('my-fallback-shell'); }); - it('resolves with shellId from STATUS confirmation frame', async () => { + it('fails when a new connection has no shell ID (no header and no provided shellId)', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', + // No shellId provided (fresh connect) and the upgrade header never arrives, so there is no + // usable shell ID. Resolving with '' would silently break reconnect — connect must fail. }); - // Flush microtask queue (SigV4 signing chain needs >1 tick before WS is constructed) await new Promise(r => setTimeout(r, 0)); - wsState.messageHandler?.(makeConfirmationFrame('server-assigned-id')); + // No upgrade event fires — open arrives without a shell-id header. + wsState.openHandler?.(); - const conn = await connectPromise; - expect(conn.shellId).toBe('server-assigned-id'); - expect(conn.reconnected).toBe(false); + await expect(connectPromise).rejects.toThrow(/did not return a shell ID/); }); - it('sets reconnected=true from STATUS frame metadata', async () => { + it('resolves immediately on open event (no confirmation frame needed)', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - wsState.messageHandler?.(makeConfirmationFrame('existing-shell', true)); + wsState.upgradeHandler?.({ headers: { 'x-amzn-bedrock-agentcore-shell-id': 'fast-shell' } }); + wsState.openHandler?.(); const conn = await connectPromise; - expect(conn.reconnected).toBe(true); + expect(conn.shellId).toBe('fast-shell'); }); - it('throws ShellKickedError when WS closes with code 4000', async () => { + it('does not have reconnected or bytesDropped on ShellConnection', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - wsState.closeHandler?.(4000); + wsState.upgradeHandler?.({ headers: { 'x-amzn-bedrock-agentcore-shell-id': 'shell-1' } }); + wsState.openHandler?.(); - await expect(connectPromise).rejects.toThrow(ShellKickedError); + const conn = await connectPromise; + expect(conn).not.toHaveProperty('reconnected'); + expect(conn).not.toHaveProperty('bytesDropped'); }); - it('throws generic error when WS closes with non-4000 code before confirmation', async () => { + it('throws ShellKickedError when WS closes with code 4000', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - wsState.closeHandler?.(1006); + wsState.closeHandler?.(4000); - await expect(connectPromise).rejects.toThrow(/closed before confirmation/); + await expect(connectPromise).rejects.toThrow(ShellKickedError); }); - it('throws on WS error before confirmation', async () => { + it('throws generic error when WS closes with non-4000 code before open', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - wsState.errorHandler?.(new Error('ECONNREFUSED')); + wsState.closeHandler?.(1006); - await expect(connectPromise).rejects.toThrow('ECONNREFUSED'); + await expect(connectPromise).rejects.toThrow(/closed before open/); }); - it('ignores non-STATUS frames before confirmation', async () => { + it('throws on WS error before open', async () => { const connectPromise = connectShell({ region: 'us-east-1', runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', }); await new Promise(r => setTimeout(r, 0)); - // Send a STDOUT frame first — should be ignored - const stdout = Buffer.concat([Buffer.from([ShellChannel.STDOUT]), Buffer.from('noise')]); - wsState.messageHandler?.(stdout); - // Then send confirmation - wsState.messageHandler?.(makeConfirmationFrame('abc')); + wsState.errorHandler?.(new Error('ECONNREFUSED')); - const conn = await connectPromise; - expect(conn.shellId).toBe('abc'); + await expect(connectPromise).rejects.toThrow('ECONNREFUSED'); }); it('does not retry after ShellKickedError (close code 4000)', async () => { @@ -259,72 +247,15 @@ describe('connectShell', () => { }); await new Promise(r => setTimeout(r, 0)); - wsState.messageHandler?.(makeConfirmationFrame('reconnect-id', true)); + wsState.upgradeHandler?.({ headers: { 'x-amzn-bedrock-agentcore-shell-id': 'reconnect-id' } }); + wsState.openHandler?.(); const conn = await connectPromise; expect(conn.shellId).toBe('reconnect-id'); - expect(conn.reconnected).toBe(true); - expect(wsState.calls[0]).toContain('shellId=reconnect-id'); }); }); -// --------------------------------------------------------------------------- -// confirmationTimeoutMs — rejects if STATUS frame never arrives -// --------------------------------------------------------------------------- - -describe('connectShell confirmationTimeoutMs', () => { - beforeEach(() => { - wsState.reset(); - vi.useFakeTimers(); - }); - - afterEach(() => { - vi.useRealTimers(); - }); - - it('rejects with timeout message when STATUS frame never arrives', async () => { - // Use bearerToken path to bypass async SigV4 signing — WS is created synchronously - // so the confirmation timer is registered before we advance fake timers. - const connectPromise = connectShell({ - region: 'us-east-1', - runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', - bearerToken: 'test-token', - confirmationTimeoutMs: 5_000, - }); - - // One tick for the Promise constructor inside openWebSocket to run - await Promise.resolve(); - vi.advanceTimersByTime(5_001); - - await expect(connectPromise).rejects.toThrow(/Timed out waiting for shell confirmation \(5s\)/); - }); - - it('does not reject when STATUS frame arrives before the timeout', async () => { - // Use real timers for this test — fake timers interfere with the async signing chain. - vi.useRealTimers(); - - const connectPromise = connectShell({ - region: 'us-east-1', - runtimeArn: 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', - confirmationTimeoutMs: 5_000, - }); - - // Wait for the SigV4 signing microtasks + WS construction to complete - await new Promise(r => setTimeout(r, 0)); - const payload = JSON.stringify({ - kind: 'Status', - apiVersion: 'v1', - metadata: { shellId: 'fast-shell', reconnected: false }, - status: 'Success', - }); - wsState.messageHandler?.(Buffer.concat([Buffer.from([ShellChannel.STATUS]), Buffer.from(payload)])); - - const conn = await connectPromise; - expect(conn.shellId).toBe('fast-shell'); - }); -}); - // --------------------------------------------------------------------------- // AGENTCORE_STAGE case-insensitivity // --------------------------------------------------------------------------- @@ -348,7 +279,7 @@ describe('buildShellUrl AGENTCORE_STAGE case-insensitivity', () => { }); // --------------------------------------------------------------------------- -// Gap 1 — serviceEndpoint() in buildShellUrl (partition-aware prod URL) +// Partition-aware prod URL // --------------------------------------------------------------------------- describe('buildShellUrl partition-aware hostname', () => { @@ -363,28 +294,16 @@ describe('buildShellUrl partition-aware hostname', () => { it('uses the region-specific DNS suffix for GovCloud (us-gov-west-1)', () => { const url = buildShellUrl('us-gov-west-1', 'arn:aws-us-gov:bedrock-agentcore:us-gov-west-1:123:runtime/r'); - // GovCloud partition dnsSuffix is 'amazonaws.com' per @aws-sdk/util-endpoints expect(url.hostname).toBe('bedrock-agentcore.us-gov-west-1.amazonaws.com'); - // Confirm partition name is aws-us-gov (i.e. serviceEndpoint was used, not a hardcoded domain) expect(url.hostname).toMatch(/^bedrock-agentcore\.us-gov-west-1\./); }); }); // --------------------------------------------------------------------------- -// Gap 2 — HTTP upgrade error translation (indirect via WS mock) +// HTTP upgrade error translation // --------------------------------------------------------------------------- describe('connectShell error translation', () => { - function _makeConfirmationFrame(shellId: string, reconnected = false): Buffer { - const payload = JSON.stringify({ - kind: 'Status', - apiVersion: 'v1', - metadata: { shellId, reconnected }, - status: 'Success', - }); - return Buffer.concat([Buffer.from([ShellChannel.STATUS]), Buffer.from(payload)]); - } - beforeEach(() => { wsState.reset(); }); @@ -439,20 +358,10 @@ describe('connectShell error translation', () => { }); // --------------------------------------------------------------------------- -// Gap 3 — Reconnect UX callbacks +// Reconnect UX callbacks // --------------------------------------------------------------------------- describe('connectShell reconnect callbacks', () => { - function makeConfirmationFrame(shellId: string, reconnected = false): Buffer { - const payload = JSON.stringify({ - kind: 'Status', - apiVersion: 'v1', - metadata: { shellId, reconnected }, - status: 'Success', - }); - return Buffer.concat([Buffer.from([ShellChannel.STATUS]), Buffer.from(payload)]); - } - beforeEach(() => { wsState.reset(); }); @@ -472,7 +381,7 @@ describe('connectShell reconnect callbacks', () => { expect(onKicked).toHaveBeenCalledTimes(1); }); - it('calls onAttempt(1, reason) on first retry when WS fails before confirmation', async () => { + it('calls onAttempt(1, reason) on first retry when WS fails before open', async () => { const onAttempt = vi.fn(); // Use a very short base delay so the test doesn't wait @@ -490,8 +399,9 @@ describe('connectShell reconnect callbacks', () => { // Wait for backoff + second WS to be constructed await new Promise(r => setTimeout(r, 50)); - // Second attempt: send confirmation - wsState.messageHandler?.(makeConfirmationFrame('new-shell-id')); + // Second attempt: fire open + wsState.upgradeHandler?.({ headers: { 'x-amzn-bedrock-agentcore-shell-id': 'new-shell-id' } }); + wsState.openHandler?.(); await connectPromise; @@ -500,7 +410,7 @@ describe('connectShell reconnect callbacks', () => { }); // --------------------------------------------------------------------------- -// Gap 5 — startKeepalive +// startKeepalive // --------------------------------------------------------------------------- function makeMockWs() { diff --git a/src/cli/aws/__tests__/shell-framer.test.ts b/src/cli/aws/__tests__/shell-framer.test.ts index 5a5bf0abe..5c877aa28 100644 --- a/src/cli/aws/__tests__/shell-framer.test.ts +++ b/src/cli/aws/__tests__/shell-framer.test.ts @@ -131,13 +131,6 @@ describe('ShellFramer.encodeHeartbeat', () => { }); }); -describe('ShellFramer.encodeClose', () => { - it('returns single-byte CLOSE frame', () => { - const buf = framer.encodeClose(); - expect(buf).toEqual(Buffer.from([ShellChannel.CLOSE])); - }); -}); - describe('parseStatusFrame', () => { it('identifies a confirmation frame with shellId', () => { const payload = JSON.stringify({ diff --git a/src/cli/aws/connect-shell.ts b/src/cli/aws/connect-shell.ts index 596db5300..d46841bf2 100644 --- a/src/cli/aws/connect-shell.ts +++ b/src/cli/aws/connect-shell.ts @@ -1,6 +1,5 @@ import { ShellKickedError } from '../../lib/errors/types'; import { getCredentialProvider } from './account'; -import { ShellChannel, ShellFramer, parseStatusFrame } from './shell-framer'; import { dataPlaneEndpoint } from './stage-endpoint'; import { Sha256 } from '@aws-crypto/sha256-js'; import { HttpRequest } from '@smithy/protocol-http'; @@ -22,10 +21,6 @@ export interface ShellReconnectOptions { onAttempt?: (attempt: number, reason: string) => void; /** Called when close code 4000 is received — another client took the session. */ onKicked?: () => void; - /** Called when reconnect yields a fresh shell (previous session expired). */ - onNewSession?: (shellId: string) => void; - /** Called when the confirmation frame reports bytes lost during disconnect. */ - onBytesDropped?: (n: number) => void; } export interface ConnectShellOptions { @@ -41,18 +36,13 @@ export interface ConnectShellOptions { reconnect?: ShellReconnectOptions; /** Bearer token for CUSTOM_JWT auth. When set, authenticates via WebSocket subprotocol instead of SigV4. */ bearerToken?: string; - /** Milliseconds to wait for the STATUS confirmation frame before failing. Default: 10_000 */ - confirmationTimeoutMs?: number; } export interface ShellConnection { ws: WebSocket; - /** The server-assigned shell identifier (from wire `shellId`). */ + /** The server-assigned shell identifier (from the 101 upgrade response header). */ shellId: string; sessionId?: string; - reconnected: boolean; - /** Bytes of output lost during a disconnect, reported in the confirmation frame. */ - bytesDropped?: number; } // --------------------------------------------------------------------------- @@ -138,7 +128,7 @@ function translateUpgradeError(err: Error): Error { // --------------------------------------------------------------------------- async function openWebSocket(options: ConnectShellOptions): Promise { - const { region, runtimeArn, shellId, sessionId, bearerToken, confirmationTimeoutMs = 10_000 } = options; + const { region, runtimeArn, shellId, sessionId, bearerToken } = options; const url = buildShellUrl(region, runtimeArn, shellId); let ws: WebSocket; @@ -166,28 +156,19 @@ async function openWebSocket(options: ConnectShellOptions): Promise((resolve, reject) => { - const framer = new ShellFramer(); let settled = false; - // Shell ID from the 101 response header — preferred over the STATUS frame per spec. + // Shell ID from the 101 response header — the primary (and now only) source. let shellIdFromHeader: string | undefined; const fail = (err: Error) => { if (!settled) { settled = true; - clearTimeout(confirmationTimer); ws.terminate(); reject(translateUpgradeError(err)); } }; - // Fail fast if the server never sends the STATUS confirmation frame - const confirmationTimer = setTimeout( - () => fail(new Error(`Timed out waiting for shell confirmation (${confirmationTimeoutMs / 1000}s)`)), - confirmationTimeoutMs - ); - - // Read shellId from the 101 Switching Protocols response headers (primary source). - // The STATUS frame (0x03) is the fallback for browser clients that cannot read headers. + // Read shellId from the 101 Switching Protocols response headers. ws.on('upgrade', (response: { headers: Record }) => { const raw = response.headers['x-amzn-bedrock-agentcore-shell-id']; if (raw) { @@ -202,42 +183,36 @@ async function openWebSocket(options: ConnectShellOptions): Promise { + // Connection is ready immediately after WebSocket opens — no confirmation frame wait. + ws.on('open', () => { if (settled) return; - let frame; - try { - frame = framer.decode(Buffer.isBuffer(data) ? data : Buffer.from(data as ArrayBuffer)); - } catch { - return; // malformed — wait for status frame - } - - if (frame.channel !== ShellChannel.STATUS) return; - - const parsed = parseStatusFrame(frame); - if (parsed.type === 'confirmation') { - settled = true; - clearTimeout(confirmationTimer); - const conn: ShellConnection = { - ws, - // Header is primary; STATUS frame is fallback for browser clients. - shellId: shellIdFromHeader ?? parsed.shellId, - sessionId: options.sessionId, - reconnected: parsed.reconnected, - }; - if (parsed.bytesDropped !== undefined) { - conn.bytesDropped = parsed.bytesDropped; - } - resolve(conn); - } - // termination before confirmation — treat as error - if (parsed.type === 'termination') { - fail(new Error('Shell terminated before confirmation frame')); + // The 101 upgrade header is the only server-provided source of the shell ID now that the + // 0x03 confirmation frame is gone; on a reconnect the caller-supplied shellId is a valid + // fallback. If neither is present we have no usable ID — fail instead of resolving with '', + // because an empty shell ID silently breaks reconnect (the session cannot be reattached and + // the reconnect hint is suppressed). The server returns this header on every successful + // upgrade, so its absence is a real anomaly, not a normal state. + const resolvedShellId = shellIdFromHeader ?? shellId; + if (!resolvedShellId) { + fail( + new Error( + 'Shell session could not be established: the server did not return a shell ID. ' + + 'Retry, or run `agentcore status` to check the runtime.' + ) + ); + return; } + settled = true; + resolve({ + ws, + shellId: resolvedShellId, + sessionId: options.sessionId, + }); }); }); } @@ -306,15 +281,7 @@ export async function connectShell(options: ConnectShellOptions): Promise 0 && !conn.reconnected && onNewSession) { - onNewSession(conn.shellId); - } - // Carry shellId forward so subsequent reconnects reattach to the same PTY currentShellId = conn.shellId; return conn; diff --git a/src/cli/aws/shell-framer.ts b/src/cli/aws/shell-framer.ts index 5e7c0d14f..0ab634c4c 100644 --- a/src/cli/aws/shell-framer.ts +++ b/src/cli/aws/shell-framer.ts @@ -76,10 +76,6 @@ export class ShellFramer { encodeHeartbeat(): Buffer { return Buffer.from([ShellChannel.HEARTBEAT]); } - - encodeClose(): Buffer { - return Buffer.from([ShellChannel.CLOSE]); - } } export class ValueError extends Error { diff --git a/src/cli/commands/exec/__tests__/action.test.ts b/src/cli/commands/exec/__tests__/action.test.ts index e8d4c57d8..74465e421 100644 --- a/src/cli/commands/exec/__tests__/action.test.ts +++ b/src/cli/commands/exec/__tests__/action.test.ts @@ -235,7 +235,6 @@ describe('handleShellSession banner messages', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs as unknown as import('ws').default, shellId: 'test-shell-id', - reconnected: false, }); // eslint-disable-next-line @typescript-eslint/no-empty-function @@ -260,7 +259,6 @@ describe('handleShellSession banner messages', () => { return Promise.resolve({ ws: mockWs as unknown as import('ws').default, shellId: 'test-shell-id', - reconnected: false, }); }); @@ -321,31 +319,6 @@ describe('handleShellSession banner messages', () => { expect(stderrCalls.some(msg => msg.includes('[session closed · exit 0]'))).toBe(true); expect(result.success).toBe(true); }); - - it('writes "[info] Previous shell session has ended..." when shellId passed but reconnected=false', async () => { - vi.mocked(connectShell).mockResolvedValue({ - ws: mockWs as unknown as import('ws').default, - shellId: 'new-shell-id', - reconnected: false, - }); - - const options: ExecOptions = { - runtimeArn: CTX.runtimeArn, - region: CTX.region, - shellId: 'old-shell-id', - }; - - const sessionPromise = handleShellSession(CTX, options); - await new Promise(r => setTimeout(r, 0)); - - const stderrCalls = (stderrSpy.mock.calls as [string][]).map(c => c[0]); - expect(stderrCalls.some(msg => msg.includes('[info]') && msg.includes('Previous shell session has ended'))).toBe( - true - ); - - (mockWs as unknown as { _fire: (e: string, ...a: unknown[]) => void })._fire('close', 0); - await sessionPromise; - }); }); // --------------------------------------------------------------------------- @@ -381,7 +354,6 @@ describe('handleShellSession CLOSE frame (0xFF)', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs as unknown as import('ws').default, shellId: 'shell-abc', - reconnected: false, }); // eslint-disable-next-line @typescript-eslint/no-empty-function vi.mocked(startKeepalive).mockReturnValue(() => {}); @@ -440,7 +412,6 @@ describe('handleShellSession unknown channel byte', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs as unknown as import('ws').default, shellId: 'shell-xyz', - reconnected: false, }); // eslint-disable-next-line @typescript-eslint/no-empty-function vi.mocked(startKeepalive).mockReturnValue(() => {}); @@ -508,7 +479,6 @@ describe('handleShellSession startKeepalive integration', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs as unknown as import('ws').default, shellId: 'shell-keep', - reconnected: false, }); stopKeepalive = vi.fn(); @@ -584,7 +554,7 @@ describe('handleShellSession reconnect callbacks wired to connectShell', () => { vi.mocked(connectShell).mockImplementation(opts => { capturedOnAttempt = opts.reconnect?.onAttempt; capturedWs = makeMockWsForCallbacks(); - return Promise.resolve({ ws: capturedWs, shellId: 'shell-ra', reconnected: false }); + return Promise.resolve({ ws: capturedWs, shellId: 'shell-ra' }); }); const options: ExecOptions = { runtimeArn: CTX.runtimeArn, region: CTX.region }; @@ -607,7 +577,7 @@ describe('handleShellSession reconnect callbacks wired to connectShell', () => { vi.mocked(connectShell).mockImplementation(opts => { capturedOnKicked = opts.reconnect?.onKicked; capturedWs = makeMockWsForCallbacks(); - return Promise.resolve({ ws: capturedWs, shellId: 'shell-kick', reconnected: false }); + return Promise.resolve({ ws: capturedWs, shellId: 'shell-kick' }); }); const options: ExecOptions = { runtimeArn: CTX.runtimeArn, region: CTX.region }; @@ -629,7 +599,7 @@ describe('handleShellSession reconnect callbacks wired to connectShell', () => { vi.mocked(connectShell).mockImplementation(opts => { capturedOnKicked = opts.reconnect?.onKicked; capturedWs = makeMockWsForCallbacks(); - return Promise.resolve({ ws: capturedWs, shellId: 'shell-kick2', reconnected: false }); + return Promise.resolve({ ws: capturedWs, shellId: 'shell-kick2' }); }); const options: ExecOptions = { runtimeArn: CTX.runtimeArn, region: CTX.region }; @@ -650,7 +620,7 @@ describe('handleShellSession reconnect callbacks wired to connectShell', () => { vi.mocked(connectShell).mockImplementation(opts => { capturedOnAttempt = opts.reconnect?.onAttempt; capturedWs = makeMockWsForCallbacks(); - return Promise.resolve({ ws: capturedWs, shellId: 'shell-ra2', reconnected: false }); + return Promise.resolve({ ws: capturedWs, shellId: 'shell-ra2' }); }); const options: ExecOptions = { runtimeArn: CTX.runtimeArn, region: CTX.region }; @@ -1202,7 +1172,6 @@ describe('handleShellSession WS close code 1000 → clean exit', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs as unknown as import('ws').default, shellId: 'shell-1000', - reconnected: false, }); // eslint-disable-next-line @typescript-eslint/no-empty-function vi.mocked(startKeepalive).mockReturnValue(() => {}); @@ -1238,10 +1207,11 @@ describe('handleShellSession WS close code 1000 → clean exit', () => { expect(stderrData).not.toMatch(/disconnected/); }); - it('resolves success:true with exitCode 0 when WS closes with code 1006 and no STATUS frame', async () => { - // connectShell only resolves after STATUS confirmation, so any WS close without a STATUS - // termination frame (any close code) is treated as exit 0 — the shell ran to completion; - // the server just didn't send a termination frame. + it('resolves success:false with exitCode 1 when WS closes with abnormal code 1006 and no STATUS frame', async () => { + // connectShell now resolves as soon as the socket opens (the 0x03 confirmation-frame wait was + // removed), so an abnormal close such as 1006 can happen before the shell is usable. Without a + // STATUS termination frame, only code 1000 counts as a clean exit — any other code is a real + // failure and must NOT be reported as exit 0. const handlers2: Record void)[]> = {}; const fire2 = (event: string, ...args: unknown[]) => handlers2[event]?.forEach(fn => fn(...args)); const mockWs2 = { @@ -1261,7 +1231,6 @@ describe('handleShellSession WS close code 1000 → clean exit', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs2 as unknown as import('ws').default, shellId: 'shell-1006', - reconnected: false, }); const options: ExecOptions = { runtimeArn: CTX.runtimeArn, region: CTX.region }; @@ -1271,9 +1240,9 @@ describe('handleShellSession WS close code 1000 → clean exit', () => { fire2('close', 1006); const result = await sessionPromise; - expect(result.success).toBe(true); - // exitCode is 0: session confirmed, no STATUS termination frame → treated as clean exit - expect(result.exitCode).toBe(0); + expect(result.success).toBe(false); + // exitCode is 1: abnormal close with no STATUS termination frame → treated as a failure + expect(result.exitCode).toBe(1); }); }); @@ -1320,7 +1289,6 @@ describe('handleShellSession reconnect hint format', () => { vi.mocked(connectShell).mockResolvedValue({ ws: mockWs2 as unknown as import('ws').default, shellId: 'shell-hint', - reconnected: false, }); const stdinHandlers: Record void)[]> = {}; diff --git a/src/cli/commands/exec/__tests__/command.test.ts b/src/cli/commands/exec/__tests__/command.test.ts index 97932831f..cf370d538 100644 --- a/src/cli/commands/exec/__tests__/command.test.ts +++ b/src/cli/commands/exec/__tests__/command.test.ts @@ -54,7 +54,7 @@ vi.mock('../action.js', () => ({ async (recorder: { set: (attrs: Record) => void }) => { const sessionResult = await mockHandleShellSession(opts); recorder.set({ - is_reconnect: (sessionResult as Record).isReconnect ?? Boolean(opts.shellId), + is_reconnect: Boolean(opts.shellId), exit_code: (sessionResult as Record).exitCode ?? ((sessionResult as Record).success ? 0 : 1), @@ -343,7 +343,6 @@ describe('exec telemetry attributes', () => { exitCode: 0, reconnectAttempts: 0, wasKicked: false, - isReconnect: false, }); const program = new Command(); @@ -372,7 +371,6 @@ describe('exec telemetry attributes', () => { exitCode: 2, reconnectAttempts: 0, wasKicked: false, - isReconnect: false, }); const program = new Command(); @@ -399,7 +397,6 @@ describe('exec telemetry attributes', () => { exitCode: 0, reconnectAttempts: 3, wasKicked: true, - isReconnect: true, }); const program = new Command(); @@ -407,9 +404,12 @@ describe('exec telemetry attributes', () => { registerExec(program); await expect( - program.parseAsync(['exec', '--it', '--runtime', 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r'], { - from: 'user', - }) + program.parseAsync( + ['exec', '--it', '--runtime', 'arn:aws:bedrock-agentcore:us-east-1:123:runtime/r', '--shell-id', 'old-shell'], + { + from: 'user', + } + ) ).rejects.toThrow(); const telemetryCalls = vi.mocked(withCommandRunTelemetry).mock.calls; diff --git a/src/cli/commands/exec/action.ts b/src/cli/commands/exec/action.ts index 4f20d5b32..e2a2ba32a 100644 --- a/src/cli/commands/exec/action.ts +++ b/src/cli/commands/exec/action.ts @@ -317,12 +317,6 @@ export async function handleShellSession(ctx: ExecContext, options: ExecOptions) wasKicked = true; process.stderr.write('\r\n[session attached from another client · not reconnecting]\r\n'); }, - onNewSession: () => { - process.stderr.write('\r\n[new shell session (previous session expired)]\r\n'); - }, - onBytesDropped: n => { - process.stderr.write(`\r\n[${n} bytes of output lost during disconnect]\r\n`); - }, }, }); } catch (err) { @@ -330,16 +324,9 @@ export async function handleShellSession(ctx: ExecContext, options: ExecOptions) } const framer = new ShellFramer(); - const { ws, shellId, reconnected } = conn; + const { ws, shellId } = conn; let exitCode: number | null = null; - // Warn when the user requested a reconnect but the previous shell had already exited - if (options.shellId && !reconnected) { - process.stderr.write( - '[info] Previous shell session has ended. Starting a new shell (environment variables and history are not restored).\n' - ); - } - process.stderr.write(`[connected · session ${sessionId} · Ctrl+D or 'exit' to quit · Ctrl+] to detach]\n`); return new Promise(resolve => { @@ -402,7 +389,6 @@ export async function handleShellSession(ctx: ExecContext, options: ExecOptions) exitCode: code, reconnectAttempts, wasKicked, - isReconnect: reconnected, detached, }; @@ -461,6 +447,7 @@ export async function handleShellSession(ctx: ExecContext, options: ExecOptions) exitCode = parsed.exitCode; ws.close(); } + // Confirmation frames silently swallowed — server may still send them during transition break; } case ShellChannel.CLOSE: @@ -472,12 +459,25 @@ export async function handleShellSession(ctx: ExecContext, options: ExecOptions) }); ws.on('close', (code: number) => { - // If the STATUS termination frame arrived, use its exit code. - // Otherwise, treat non-kick closes as exit 0: the shell ran to completion but the server - // didn't send a STATUS termination frame (observed behavior on the beta runtime). - // connectShell only resolves after the STATUS confirmation frame, so the session is always - // active by the time we reach here — there are no unconfirmed closes. - const resolvedExitCode = exitCode ?? (code !== 4000 ? 0 : null); + // The STATUS termination frame is the authoritative exit signal — when it arrived, use its + // exit code regardless of the WebSocket close code. + // + // Without a STATUS frame, fall back to the WebSocket close code. connectShell now resolves as + // soon as the socket opens (the 0x03 confirmation-frame wait was removed), so an abnormal + // close such as 1006 can occur before the shell is usable. Only code 1000 (normal closure — + // the server's deliberate close after a clean shell exit) counts as success; any other code + // is a real failure and must not be reported as exit 0. Kick (4000) stays null so cleanup + // prints the reconnect hint instead of a spurious exit line. + let resolvedExitCode: number | null; + if (exitCode !== null) { + resolvedExitCode = exitCode; + } else if (code === 1000) { + resolvedExitCode = 0; + } else if (code === 4000) { + resolvedExitCode = null; + } else { + resolvedExitCode = 1; + } cleanup(resolvedExitCode); }); @@ -512,7 +512,7 @@ export async function runInteractiveShell(options: ExecOptions): Promise { const ctx = await loadExecContext(options); const r = await handleShellSession(ctx, options); recorder.set({ - is_reconnect: r.isReconnect ?? Boolean(options.shellId), + is_reconnect: Boolean(options.shellId), exit_code: r.exitCode ?? (r.success ? 0 : 1), reconnect_attempts: r.reconnectAttempts ?? 0, was_kicked: r.wasKicked ?? false, diff --git a/src/cli/commands/exec/command.tsx b/src/cli/commands/exec/command.tsx index a474f72b4..0653fb8fb 100644 --- a/src/cli/commands/exec/command.tsx +++ b/src/cli/commands/exec/command.tsx @@ -219,7 +219,7 @@ export async function runExecLoop(options: ExecOptions = {}): Promise { const ctx = await loadExecContext(shellOptions); const r = await handleShellSession(ctx, shellOptions); recorder.set({ - is_reconnect: r.isReconnect ?? Boolean(shellOptions.shellId), + is_reconnect: Boolean(shellOptions.shellId), exit_code: r.exitCode ?? (r.success ? 0 : 1), reconnect_attempts: r.reconnectAttempts ?? 0, was_kicked: r.wasKicked ?? false, diff --git a/src/cli/commands/exec/types.ts b/src/cli/commands/exec/types.ts index e96da106d..b235d3391 100644 --- a/src/cli/commands/exec/types.ts +++ b/src/cli/commands/exec/types.ts @@ -38,8 +38,6 @@ export type ExecResult = Result & { reconnectAttempts?: number; /** True if the session was kicked by another client (close code 4000). */ wasKicked?: boolean; - /** True if the initial connection reattached an existing shell. */ - isReconnect?: boolean; /** True if the user explicitly detached with Ctrl+] (shell is still alive on the VM). */ detached?: boolean; /** Buffered stdout from a one-shot command (populated when --json is set). */