diff --git a/src/vs/platform/agentHost/test/node/e2e/harness/agentHostE2ETestHarness.ts b/src/vs/platform/agentHost/test/node/e2e/harness/agentHostE2ETestHarness.ts index 7ed1d5cff4f7d..b43b328550dcf 100644 --- a/src/vs/platform/agentHost/test/node/e2e/harness/agentHostE2ETestHarness.ts +++ b/src/vs/platform/agentHost/test/node/e2e/harness/agentHostE2ETestHarness.ts @@ -12,7 +12,7 @@ import { execSync } from 'child_process'; import { mkdtempSync, rmSync } from 'fs'; import { homedir, tmpdir, userInfo } from 'os'; import { fileURLToPath } from 'url'; -import { raceTimeout, timeout } from '../../../../../../base/common/async.js'; +import { timeout } from '../../../../../../base/common/async.js'; import { join } from '../../../../../../base/common/path.js'; import { removeAnsiEscapeCodes } from '../../../../../../base/common/strings.js'; import { URI } from '../../../../../../base/common/uri.js'; @@ -30,7 +30,7 @@ import { import { CopilotCliConfigKey } from '../../../../common/copilotCliConfig.js'; import { CapiReplayMode } from './capiReplayProxy.js'; import { - getActionEnvelope, isActionNotification, IServerHandle, startRealServer, TestProtocolClient, + getActionEnvelope, isActionNotification, IServerHandle, startRealServer, stopServer, TestProtocolClient, } from '../../serverIntegrationTestHelpers.js'; import { createProviderSession, dispatchTurn, dispatchTurnWithAttachments } from '../../providerIntegrationTestHelpers.js'; import { AgentHostUpdateSnapshotsEnvVar, AhpSnapshotScenario } from './ahpSnapshot.js'; @@ -45,7 +45,6 @@ import { AgentHostUpdateSnapshotsEnvVar, AhpSnapshotScenario } from './ahpSnapsh const UPDATE_SNAPSHOTS = process.env[AgentHostUpdateSnapshotsEnvVar] === '1'; const RECORD = process.env['AGENT_HOST_REPLAY_RECORD'] === '1' || UPDATE_SNAPSHOTS; const REPLAY_MODE: CapiReplayMode = RECORD ? 'record' : 'replay'; -const SERVER_SHUTDOWN_TIMEOUT_MS = 30_000; /** * Upper bound on tests served by a single shared replay server before it is @@ -58,20 +57,6 @@ const TEMP_DIR_CLEANUP_TIMEOUT_MS = 30_000; export const REPLAY_PLACEHOLDER_TOKEN = 'replay-no-token'; export type AgentHostE2EModelTraffic = 'recorded' | 'none'; -async function stopServer(server: IServerHandle | undefined): Promise { - const serverProcess = server?.process; - if (!serverProcess || serverProcess.exitCode !== null || serverProcess.signalCode !== null) { - return; - } - - const serverExit = new Promise(resolve => serverProcess.once('exit', () => resolve())); - serverProcess.stdin?.end(); - if (!await raceTimeout(serverExit.then(() => true), SERVER_SHUTDOWN_TIMEOUT_MS)) { - serverProcess.kill(); - await serverExit; - } -} - export async function removeTempDirs(tempDirs: string[]): Promise { const pendingDirs = tempDirs.splice(0); const errors = new Map(); diff --git a/src/vs/platform/agentHost/test/node/protocol/agentHostServer.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/agentHostServer.integrationTest.ts index ecdf954e7d595..dd4b4681c6c1e 100644 --- a/src/vs/platform/agentHost/test/node/protocol/agentHostServer.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/agentHostServer.integrationTest.ts @@ -4,7 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; -import { getAgentHostE2ETestTimeout, IServerHandle, startServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; +import { getAgentHostE2ETestTimeout, IServerHandle, startServer, stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; suite('Agent Host Server', function () { @@ -15,8 +15,9 @@ suite('Agent Host Server', function () { server = await startServer({ quiet: false }); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); test('starts with production agent services registered', async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/clientTools.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/clientTools.integrationTest.ts index cceb45b660653..9dbbcb38a88ba 100644 --- a/src/vs/platform/agentHost/test/node/protocol/clientTools.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/clientTools.integrationTest.ts @@ -27,6 +27,7 @@ import { IServerHandle, isActionNotification, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -40,8 +41,9 @@ suite('Protocol WebSocket — Client Tools', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/handshake.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/handshake.integrationTest.ts index 1356ad5ec94a1..c9712dd3b486d 100644 --- a/src/vs/platform/agentHost/test/node/protocol/handshake.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/handshake.integrationTest.ts @@ -11,7 +11,7 @@ import { type JsonRpcErrorResponse, } from '../../../common/state/sessionProtocol.js'; import { ROOT_STATE_URI } from '../../../common/state/sessionState.js'; -import { getAgentHostE2ETestTimeout, IServerHandle, nextSessionUri, startServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; +import { getAgentHostE2ETestTimeout, IServerHandle, nextSessionUri, startServer, stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; suite('Protocol WebSocket — Handshake & Errors', function () { @@ -23,8 +23,9 @@ suite('Protocol WebSocket — Handshake & Errors', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { @@ -38,7 +39,7 @@ suite('Protocol WebSocket — Handshake & Errors', function () { }); test('handshake returns initialize response with protocol version', async function () { - this.timeout(5_000); + this.timeout(getAgentHostE2ETestTimeout(5_000, 20_000)); const result = await client.call('initialize', { protocolVersions: [PROTOCOL_VERSION], @@ -55,17 +56,19 @@ suite('Protocol WebSocket — Handshake & Errors', function () { this.timeout(10_000); const raw = new TestProtocolClient(server.port); - await raw.connect(); - - const responsePromise = raw.waitForRawMessage(); - raw.sendRaw('this is not valid json{{{'); + try { + await raw.connect(); - const response = await responsePromise as JsonRpcErrorResponse; - assert.strictEqual(response.jsonrpc, '2.0'); - assert.strictEqual(response.id, null); - assert.strictEqual(response.error.code, JSON_RPC_PARSE_ERROR); + const responsePromise = raw.waitForRawMessage(); + raw.sendRaw('this is not valid json{{{'); - raw.close(); + const response = await responsePromise as JsonRpcErrorResponse; + assert.strictEqual(response.jsonrpc, '2.0'); + assert.strictEqual(response.id, null); + assert.strictEqual(response.error.code, JSON_RPC_PARSE_ERROR); + } finally { + raw.close(); + } }); test('createSession with invalid provider does not crash server', async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/multiClient.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/multiClient.integrationTest.ts index 98d937e5a944d..faa10caf012d3 100644 --- a/src/vs/platform/agentHost/test/node/protocol/multiClient.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/multiClient.integrationTest.ts @@ -5,6 +5,7 @@ import assert from 'assert'; import type { SessionAddedParams, SessionRemovedParams } from '../../../common/state/protocol/notifications.js'; +import { ActionType } from '../../../common/state/sessionActions.js'; import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; import type { ReconnectResult } from '../../../common/state/sessionProtocol.js'; import { ROOT_STATE_URI, buildDefaultChatUri } from '../../../common/state/sessionState.js'; @@ -18,6 +19,7 @@ import { isActionNotification, nextSessionUri, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -25,24 +27,35 @@ suite('Protocol WebSocket — Multi-Client', function () { let server: IServerHandle; let client: TestProtocolClient; + const secondaryClients: TestProtocolClient[] = []; + + function createSecondaryClient(): TestProtocolClient { + const secondaryClient = new TestProtocolClient(server.port); + secondaryClients.push(secondaryClient); + return secondaryClient; + } suiteSetup(async function () { this.timeout(getAgentHostE2ETestTimeout(15_000, 60_000)); server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); client = new TestProtocolClient(server.port); await client.connect(); }); teardown(function () { client.close(); + for (const secondaryClient of secondaryClients.splice(0)) { + secondaryClient.close(); + } }); test('sessionAdded notification is broadcast to all connected clients', async function () { @@ -50,7 +63,7 @@ suite('Protocol WebSocket — Multi-Client', function () { await client.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-broadcast-add-1' }); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-broadcast-add-2' }); @@ -72,15 +85,14 @@ suite('Protocol WebSocket — Multi-Client', function () { const uri2 = (n2.params as SessionAddedParams).summary.resource; assert.strictEqual(uri1, uri2, 'both clients should see the same session URI'); - client2.close(); }); test('sessionRemoved notification is broadcast to all connected clients', async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); const sessionUri = await createAndSubscribeSession(client, 'test-broadcast-remove-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-broadcast-remove-2' }); client2.clearReceived(); @@ -101,7 +113,6 @@ suite('Protocol WebSocket — Multi-Client', function () { assert.strictEqual(removed1.session.toString(), sessionUri.toString()); assert.strictEqual(removed2.session.toString(), sessionUri.toString()); - client2.close(); }); test('two clients on same session both see actions', async function () { @@ -109,7 +120,7 @@ suite('Protocol WebSocket — Multi-Client', function () { const sessionUri = await createAndSubscribeSession(client, 'test-multi-client-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-multi-client-2' }); await client2.call('subscribe', { channel: sessionUri }); @@ -126,7 +137,6 @@ suite('Protocol WebSocket — Multi-Client', function () { await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); await client2.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); - client2.close(); }); test('client B sends message on session created by client A', async function () { @@ -134,7 +144,7 @@ suite('Protocol WebSocket — Multi-Client', function () { const sessionUri = await createAndSubscribeSession(client, 'test-cross-msg-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-cross-msg-2' }); await client2.call('subscribe', { channel: sessionUri }); @@ -153,7 +163,6 @@ suite('Protocol WebSocket — Multi-Client', function () { await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); await client2.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); - client2.close(); }); test('both clients receive full tool progress updates', async function () { @@ -161,7 +170,7 @@ suite('Protocol WebSocket — Multi-Client', function () { const sessionUri = await createAndSubscribeSession(client, 'test-tool-progress-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-tool-progress-2' }); await client2.call('subscribe', { channel: sessionUri }); @@ -180,39 +189,41 @@ suite('Protocol WebSocket — Multi-Client', function () { await c.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); } - client2.close(); }); test('unsubscribe stops receiving session actions', async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); const sessionUri = await createAndSubscribeSession(client, 'test-unsubscribe'); client.notify('unsubscribe', { channel: sessionUri }); - await new Promise(resolve => setTimeout(resolve, 100)); + await client.call('ping'); client.clearReceived(); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-unsub-helper' }); await client2.call('subscribe', { channel: sessionUri }); - await client2.call('subscribe', { channel: buildDefaultChatUri(sessionUri) }); - - dispatchTurnStarted(client2, sessionUri, 'turn-unsub', 'hello', 1); - await client2.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); + client2.dispatch({ + channel: sessionUri, + clientSeq: 1, + action: { type: ActionType.SessionTitleChanged, title: 'Updated by subscribed client' }, + }); + await client2.waitForNotification(n => + isActionNotification(n, ActionType.SessionTitleChanged) + && getActionEnvelope(n).channel === sessionUri + ); - await new Promise(resolve => setTimeout(resolve, 300)); - const sessionActions = client.receivedNotifications(n => isActionNotification(n, 'session/')); + await client.call('ping'); + const sessionActions = client.receivedNotifications(n => n.method === 'action' && getActionEnvelope(n).channel === sessionUri); assert.strictEqual(sessionActions.length, 0, 'unsubscribed client should not receive session actions'); - - client2.close(); }); test('unsubscribed client receives no actions but still gets notifications', async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); const sessionUri = await createAndSubscribeSession(client, 'test-scoping-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-scoping-2' }); // Client 2 does NOT subscribe to the session @@ -221,8 +232,7 @@ suite('Protocol WebSocket — Multi-Client', function () { dispatchTurnStarted(client, sessionUri, 'turn-scoped', 'hello', 1); await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); - // Give some time for any stray actions to arrive - await new Promise(resolve => setTimeout(resolve, 300)); + await client2.call('ping'); const sessionActions = client2.receivedNotifications(n => n.method === 'action'); assert.strictEqual(sessionActions.length, 0, 'unsubscribed client should receive no session actions'); @@ -235,7 +245,6 @@ suite('Protocol WebSocket — Multi-Client', function () { ); assert.ok(removed, 'unsubscribed client should still receive sessionRemoved notification'); - client2.close(); }); test('late subscriber gets current state via snapshot', async function () { @@ -246,7 +255,7 @@ suite('Protocol WebSocket — Multi-Client', function () { await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); // Client 2 joins after the turn has completed - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-late-sub-2' }); @@ -255,7 +264,6 @@ suite('Protocol WebSocket — Multi-Client', function () { assert.strictEqual(state.turns[0].id, 'turn-late'); assert.strictEqual(state.turns[0].state, 'complete'); - client2.close(); }); test('permission flow: client B confirms tool started by client A', async function () { @@ -263,7 +271,7 @@ suite('Protocol WebSocket — Multi-Client', function () { const sessionUri = await createAndSubscribeSession(client, 'test-cross-perm-1'); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-cross-perm-2' }); await client2.call('subscribe', { channel: sessionUri }); @@ -298,7 +306,6 @@ suite('Protocol WebSocket — Multi-Client', function () { await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); await client2.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); - client2.close(); }); test('reconnect replays missed actions', async function () { @@ -314,7 +321,7 @@ suite('Protocol WebSocket — Multi-Client', function () { client.close(); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); const result = await client2.call('reconnect', { clientId: 'test-reconnect', @@ -327,6 +334,5 @@ suite('Protocol WebSocket — Multi-Client', function () { assert.ok(result.actions.length > 0, 'should have replayed actions'); } - client2.close(); }); }); diff --git a/src/vs/platform/agentHost/test/node/protocol/networkDiagnostics.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/networkDiagnostics.integrationTest.ts index ece890248b35a..15bc973d21f56 100644 --- a/src/vs/platform/agentHost/test/node/protocol/networkDiagnostics.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/networkDiagnostics.integrationTest.ts @@ -9,7 +9,7 @@ import type { AddressInfo } from 'net'; import type { IAgentHostNetworkDiagnosticsInfo, IAgentHostNetworkFetchResult } from '../../../common/agentService.js'; import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; import { ROOT_STATE_URI } from '../../../common/state/sessionState.js'; -import { getAgentHostE2ETestTimeout, IServerHandle, startServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; +import { getAgentHostE2ETestTimeout, IServerHandle, startServer, stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; suite('Protocol WebSocket - quiet network diagnostics', function () { @@ -34,8 +34,9 @@ suite('Protocol WebSocket - quiet network diagnostics', function () { }); suiteTeardown(async function () { - server.process.kill(); + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); await new Promise(resolve => target.close(() => resolve())); + await stopServer(server); }); setup(async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/otlpLogs.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/otlpLogs.integrationTest.ts index c1296b5b0fafb..edaf999f4dbd5 100644 --- a/src/vs/platform/agentHost/test/node/protocol/otlpLogs.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/otlpLogs.integrationTest.ts @@ -10,7 +10,7 @@ import type { InitializeResult } from '../../../common/state/sessionProtocol.js' import type { TelemetryCapabilities } from '../../../common/state/protocol/channels-otlp/state.js'; import type { OtlpExportLogsParams } from '../../../common/state/protocol/channels-otlp/notifications.js'; import { OTLP_LOGS_CHANNEL_TEMPLATE, iterateOtlpLogRecords } from '../../../common/otlp/otlpLogEmitter.js'; -import { getAgentHostE2ETestTimeout, IServerHandle, startServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; +import { getAgentHostE2ETestTimeout, IServerHandle, startServer, stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; /** * End-to-end checks that the agent host server actually advertises and @@ -34,8 +34,9 @@ suite('Protocol WebSocket — OTLP logs channel', function () { server = await startServer({ quiet: true }); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { @@ -49,7 +50,7 @@ suite('Protocol WebSocket — OTLP logs channel', function () { }); test('initialize advertises the logs channel template', async function () { - this.timeout(5_000); + this.timeout(getAgentHostE2ETestTimeout(5_000, 20_000)); const result = await client.call('initialize', { protocolVersions: [PROTOCOL_VERSION], @@ -61,7 +62,7 @@ suite('Protocol WebSocket — OTLP logs channel', function () { }); test('subscribe on the logs channel returns a stateless empty result', async function () { - this.timeout(5_000); + this.timeout(getAgentHostE2ETestTimeout(5_000, 20_000)); await client.call('initialize', { protocolVersions: [PROTOCOL_VERSION], diff --git a/src/vs/platform/agentHost/test/node/protocol/resourceOperations.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/resourceOperations.integrationTest.ts index 390139917a850..d7b220f76746a 100644 --- a/src/vs/platform/agentHost/test/node/protocol/resourceOperations.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/resourceOperations.integrationTest.ts @@ -12,12 +12,13 @@ import { URI } from '../../../../../base/common/uri.js'; import { ContentEncoding, ResourceType, ResourceWriteMode, type ResourceListResult, type ResourceReadResult, type ResourceResolveResult } from '../../../common/state/protocol/common/commands.js'; import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; import { ROOT_STATE_URI } from '../../../common/state/sessionState.js'; -import { getActionEnvelope, getAgentHostE2ETestTimeout, isActionNotification, type IServerHandle, startServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; +import { getActionEnvelope, getAgentHostE2ETestTimeout, isActionNotification, type IServerHandle, startServer, stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; suite('Protocol WebSocket - Resource Operations', function () { let server: IServerHandle; let client: TestProtocolClient; + const secondaryClients: TestProtocolClient[] = []; let testDirectory: string; let clientCounter = 0; @@ -26,12 +27,13 @@ suite('Protocol WebSocket - Resource Operations', function () { server = await startServer({ startupTimeoutMs: getAgentHostE2ETestTimeout(30_000, 50_000) }); }); - suiteTeardown(function () { - server?.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); testDirectory = mkdtempSync(join(tmpdir(), 'agent-host-resource-')); client = new TestProtocolClient(server.port); await client.connect(); @@ -40,7 +42,10 @@ suite('Protocol WebSocket - Resource Operations', function () { teardown(function () { client.close(); - rmSync(testDirectory, { recursive: true, force: true }); + for (const secondaryClient of secondaryClients.splice(0)) { + secondaryClient.close(); + } + rmSync(testDirectory, { recursive: true, force: true, maxRetries: 5, retryDelay: 100 }); }); function resource(name: string): string { @@ -405,6 +410,7 @@ suite('Protocol WebSocket - Resource Operations', function () { const first = await client.call<{ snapshot: { state: { root: string } } }>('subscribe', { channel: watch.channel }); const secondClient = new TestProtocolClient(server.port); + secondaryClients.push(secondClient); await secondClient.connect(); await secondClient.call('initialize', { protocolVersions: [PROTOCOL_VERSION], clientId: `resource-client-${++clientCounter}` }); const second = await secondClient.call<{ snapshot: { state: { root: string } } }>('subscribe', { channel: watch.channel }); @@ -415,6 +421,5 @@ suite('Protocol WebSocket - Resource Operations', function () { ]); client.notify('unsubscribe', { channel: watch.channel }); secondClient.notify('unsubscribe', { channel: watch.channel }); - secondClient.close(); }); }); diff --git a/src/vs/platform/agentHost/test/node/protocol/sessionConfig.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/sessionConfig.integrationTest.ts index 93ac384079693..16094c86362c2 100644 --- a/src/vs/platform/agentHost/test/node/protocol/sessionConfig.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/sessionConfig.integrationTest.ts @@ -18,6 +18,7 @@ import { IServerHandle, nextSessionUri, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -31,8 +32,9 @@ suite('Protocol WebSocket - Session Config', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { @@ -206,14 +208,7 @@ suite('Protocol WebSocket - Session Config persistence across restarts', functio client1.close(); } finally { - // Trigger graceful shutdown by closing stdin rather than sending - // SIGTERM — on Windows, `child.kill()` (SIGTERM) unconditionally - // terminates the process without invoking the shutdown handler, - // so in-flight `setMetadata` writes never reach SQLite. Closing - // stdin fires `process.stdin.on('end', shutdown)` in the server - // on every platform. - server1.process.stdin!.end(); - await new Promise(resolve => server1.process.once('exit', () => resolve())); + await stopServer(server1); } // ---- Phase 2: restart server, subscribe, verify restored config ---- @@ -243,8 +238,7 @@ suite('Protocol WebSocket - Session Config persistence across restarts', functio client2.close(); } finally { - server2.process.stdin!.end(); - await new Promise(resolve => server2.process.once('exit', () => resolve())); + await stopServer(server2); } }); }); diff --git a/src/vs/platform/agentHost/test/node/protocol/sessionDiffs.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/sessionDiffs.integrationTest.ts index 3d0ad44bc062c..b49717abf78e7 100644 --- a/src/vs/platform/agentHost/test/node/protocol/sessionDiffs.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/sessionDiffs.integrationTest.ts @@ -20,6 +20,7 @@ import { isActionNotification, nextSessionUri, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -38,8 +39,9 @@ const hasGit = (() => { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/sessionFeatures.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/sessionFeatures.integrationTest.ts index d69b4d115d481..d9eaef8de1459 100644 --- a/src/vs/platform/agentHost/test/node/protocol/sessionFeatures.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/sessionFeatures.integrationTest.ts @@ -22,6 +22,7 @@ import { IServerHandle, nextSessionUri, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -35,8 +36,9 @@ suite('Protocol WebSocket — Session Features', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/sessionLifecycle.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/sessionLifecycle.integrationTest.ts index a2681d77213b7..f7cdbbaa3ad86 100644 --- a/src/vs/platform/agentHost/test/node/protocol/sessionLifecycle.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/sessionLifecycle.integrationTest.ts @@ -21,6 +21,7 @@ import { IServerHandle, nextSessionUri, startServer, + stopServer, TestProtocolClient } from '../serverIntegrationTestHelpers.js'; @@ -28,6 +29,13 @@ suite('Protocol WebSocket — Session Lifecycle', function () { let server: IServerHandle; let client: TestProtocolClient; + const secondaryClients: TestProtocolClient[] = []; + + function createSecondaryClient(): TestProtocolClient { + const secondaryClient = new TestProtocolClient(server.port); + secondaryClients.push(secondaryClient); + return secondaryClient; + } // Short idle-release grace so the release/restore test exercises a real // release promptly. Safe on the shared server because the mock agent's @@ -39,8 +47,9 @@ suite('Protocol WebSocket — Session Lifecycle', function () { server = await startServer({ env: { [AgentHostSessionReleaseGraceMsEnvVar]: String(RELEASE_GRACE_MS) } }); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { @@ -51,6 +60,9 @@ suite('Protocol WebSocket — Session Lifecycle', function () { teardown(function () { client.close(); + for (const secondaryClient of secondaryClients.splice(0)) { + secondaryClient.close(); + } }); test('create session triggers sessionAdded notification', async function () { @@ -97,7 +109,7 @@ suite('Protocol WebSocket — Session Lifecycle', function () { }); test('subscribe to a pre-existing session restores turns from agent history', async function () { - this.timeout(10_000); + this.timeout(getAgentHostE2ETestTimeout(10_000, 30_000)); await client.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-restore' }); @@ -130,7 +142,7 @@ suite('Protocol WebSocket — Session Lifecycle', function () { // Restoring should NOT emit a duplicate sessionAdded notification // (the session is already known to clients via listSessions). - await new Promise(resolve => setTimeout(resolve, 200)); + await client.call('ping'); const sessionAddedNotifs = client.receivedNotifications(n => n.method === 'root/sessionAdded' ); @@ -204,7 +216,7 @@ suite('Protocol WebSocket — Session Lifecycle', function () { // Poll listSessions until the persisted flags appear (async DB write) client.close(); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-read-archived-flags-2' }); @@ -221,7 +233,6 @@ suite('Protocol WebSocket — Session Lifecycle', function () { assert.ok(session.status & SessionStatus.IsArchived, 'IsArchived should be persisted in listSessions'); assert.ok(session.status & SessionStatus.IsRead, 'IsRead should be persisted in listSessions'); - client2.close(); }); test('dispatching isRead=false explicitly persists as false', async function () { @@ -244,7 +255,7 @@ suite('Protocol WebSocket — Session Lifecycle', function () { await client.waitForNotification(n => isActionNotification(n, 'session/isReadChanged')); client.close(); - const client2 = new TestProtocolClient(server.port); + const client2 = createSecondaryClient(); await client2.connect(); await client2.call('initialize', { channel: ROOT_STATE_URI, protocolVersions: [PROTOCOL_VERSION], clientId: 'test-isread-false-2' }); @@ -260,6 +271,5 @@ suite('Protocol WebSocket — Session Lifecycle', function () { assert.ok(session, 'session should appear in listSessions'); assert.strictEqual(session.status & SessionStatus.IsRead, 0, 'IsRead flag should not be set'); - client2.close(); }); }); diff --git a/src/vs/platform/agentHost/test/node/protocol/toolApproval.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/toolApproval.integrationTest.ts index bf39759dec2a7..2f4474bceb9f4 100644 --- a/src/vs/platform/agentHost/test/node/protocol/toolApproval.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/toolApproval.integrationTest.ts @@ -15,6 +15,7 @@ import { IServerHandle, isActionNotification, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; @@ -28,8 +29,9 @@ suite('Protocol WebSocket — Permissions & Auto-Approve', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { diff --git a/src/vs/platform/agentHost/test/node/protocol/turnExecution.integrationTest.ts b/src/vs/platform/agentHost/test/node/protocol/turnExecution.integrationTest.ts index 3971009187e59..d5a8c0f4959d3 100644 --- a/src/vs/platform/agentHost/test/node/protocol/turnExecution.integrationTest.ts +++ b/src/vs/platform/agentHost/test/node/protocol/turnExecution.integrationTest.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { timeout } from '../../../../../base/common/async.js'; import { SubscribeResult } from '../../../common/state/protocol/commands.js'; import { PROTOCOL_VERSION } from '../../../common/state/protocol/version/registry.js'; import type { IResponsePartAction } from '../../../common/state/sessionActions.js'; @@ -19,9 +20,15 @@ import { IServerHandle, isActionNotification, startServer, + stopServer, TestProtocolClient, } from '../serverIntegrationTestHelpers.js'; +function isTurnComplete(notification: Parameters[0], turnId: string): boolean { + return isActionNotification(notification, 'chat/turnComplete') + && (getActionEnvelope(notification).action as { turnId?: string }).turnId === turnId; +} + suite('Protocol WebSocket — Turn Execution', function () { let server: IServerHandle; @@ -32,8 +39,9 @@ suite('Protocol WebSocket — Turn Execution', function () { server = await startServer(); }); - suiteTeardown(function () { - server.process.kill(); + suiteTeardown(async function () { + this.timeout(getAgentHostE2ETestTimeout(20_000, 50_000)); + await stopServer(server); }); setup(async function () { @@ -115,11 +123,10 @@ suite('Protocol WebSocket — Turn Execution', function () { const sessionUri = await createAndSubscribeSession(client, 'test-multi-turns'); dispatchTurnStarted(client, sessionUri, 'turn-m1', 'hello', 1); - await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); + await client.waitForNotification(n => isTurnComplete(n, 'turn-m1')); dispatchTurnStarted(client, sessionUri, 'turn-m2', 'hello', 2); - await new Promise(resolve => setTimeout(resolve, 200)); - await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); + await client.waitForNotification(n => isTurnComplete(n, 'turn-m2')); const state = await fetchSessionWithChat(client, sessionUri); assert.ok(state.turns.length >= 2, `expected >= 2 turns but got ${state.turns.length}`); @@ -133,11 +140,10 @@ suite('Protocol WebSocket — Turn Execution', function () { const sessionUri = await createAndSubscribeSession(client, 'test-fetchTurns'); dispatchTurnStarted(client, sessionUri, 'turn-ft-1', 'hello', 1); - await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); + await client.waitForNotification(n => isTurnComplete(n, 'turn-ft-1')); dispatchTurnStarted(client, sessionUri, 'turn-ft-2', 'hello', 2); - await new Promise(resolve => setTimeout(resolve, 200)); - await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); + await client.waitForNotification(n => isTurnComplete(n, 'turn-ft-2')); const loadedPromise = client.waitForNotification(n => isActionNotification(n, 'chat/turnsLoaded')); const result = await client.call('fetchTurns', { channel: defaultChatChannel(sessionUri) }); @@ -192,7 +198,7 @@ suite('Protocol WebSocket — Turn Execution', function () { const initialSnapshot = await client.call('subscribe', { channel: sessionUri }); const initialModifiedAt = Date.parse((initialSnapshot.snapshot!.state as ISessionWithDefaultChat).chats[0].modifiedAt); - await new Promise(resolve => setTimeout(resolve, 50)); + await timeout(50); dispatchTurnStarted(client, sessionUri, 'turn-mod', 'hello', 1); await client.waitForNotification(n => isActionNotification(n, 'chat/turnComplete')); diff --git a/src/vs/platform/agentHost/test/node/serverIntegrationTestHelpers.ts b/src/vs/platform/agentHost/test/node/serverIntegrationTestHelpers.ts index 2a4a77743e1d8..f00e79fcc05bf 100644 --- a/src/vs/platform/agentHost/test/node/serverIntegrationTestHelpers.ts +++ b/src/vs/platform/agentHost/test/node/serverIntegrationTestHelpers.ts @@ -5,6 +5,7 @@ import { ChildProcess, fork } from 'child_process'; import { cp, lstat, mkdir, readFile, readdir, realpath, rename, rm, stat, writeFile } from 'fs/promises'; +import { raceTimeout } from '../../../../base/common/async.js'; import { Schemas } from '../../../../base/common/network.js'; import { createRequire } from 'module'; import { mkdirSync } from 'fs'; @@ -67,6 +68,13 @@ interface IPendingCall { reject: (err: Error) => void; } +function getProtocolOperationTimeout(): number { + if (process.env['AGENT_HOST_E2E_COVERAGE'] === '1') { + return 30_000; + } + return isWindows ? 8_000 : 5_000; +} + type ReverseRequestMethod = | 'createResourceWatch' | 'resourceRequest' @@ -113,6 +121,7 @@ export class TestProtocolClient { private readonly _notifications: AhpNotification[] = []; private readonly _notifWaiters: { predicate: (n: AhpNotification) => boolean; resolve: (n: AhpNotification) => void; reject: (err: Error) => void; dispose: () => void }[] = []; private _nextWatchId = 1; + private _closed = false; constructor( port: number, @@ -411,11 +420,10 @@ export class TestProtocolClient { } /** Send a JSON-RPC request and await the response. */ - call(method: string, params?: unknown, timeoutMs = 5000): Promise { + call(method: string, params?: unknown, timeoutMs = getProtocolOperationTimeout()): Promise { const id = this._nextId++; const message: JsonRpcRequest = { jsonrpc: '2.0', id, method, params }; this._ahpSnapshot.record('c2s', message); - this._ws.send(JSON.stringify(message)); return new Promise((resolve, reject) => { const timer = setTimeout(() => { this._pendingCalls.delete(id); @@ -426,11 +434,18 @@ export class TestProtocolClient { resolve: result => { clearTimeout(timer); resolve(result as T); }, reject: err => { clearTimeout(timer); reject(err); }, }); + try { + this._ws.send(JSON.stringify(message)); + } catch (error) { + this._pendingCalls.delete(id); + clearTimeout(timer); + reject(error); + } }); } /** Wait for a server notification matching a predicate. */ - waitForNotification(predicate: (n: AhpNotification) => boolean, timeoutMs = 5000): Promise { + waitForNotification(predicate: (n: AhpNotification) => boolean, timeoutMs = getProtocolOperationTimeout()): Promise { const existing = this._notifications.find(predicate); if (existing) { return Promise.resolve(existing); @@ -486,7 +501,7 @@ export class TestProtocolClient { } /** Wait for the next raw message from the server. */ - waitForRawMessage(timeoutMs = 5000): Promise { + waitForRawMessage(timeoutMs = getProtocolOperationTimeout()): Promise { return new Promise((resolve, reject) => { const timer = setTimeout(() => { cleanup(); @@ -506,6 +521,10 @@ export class TestProtocolClient { } close(): void { + if (this._closed) { + return; + } + this._closed = true; for (const w of this._notifWaiters) { w.dispose(); w.reject(new Error('Client closed')); @@ -562,6 +581,41 @@ export interface IServerHandle { capiReplay?: CapiReplayProxy; } +const SERVER_SHUTDOWN_TIMEOUT_MS = isWindows || process.env['AGENT_HOST_E2E_COVERAGE'] === '1' ? 30_000 : 5_000; + +/** Gracefully stop an Agent Host test server, killing it if shutdown stalls. */ +export async function stopServer(server: IServerHandle | undefined): Promise { + const serverProcess = server?.process; + if (!serverProcess || serverProcess.exitCode !== null || serverProcess.signalCode !== null) { + return; + } + + const serverExit = new Promise(resolve => { + const onExit = () => resolve(); + serverProcess.once('exit', onExit); + if (serverProcess.exitCode !== null || serverProcess.signalCode !== null) { + serverProcess.removeListener('exit', onExit); + resolve(); + } + }); + serverProcess.stdin?.end(); + if (!await raceTimeout(serverExit.then(() => true), SERVER_SHUTDOWN_TIMEOUT_MS)) { + try { + if (serverProcess.exitCode === null && serverProcess.signalCode === null) { + const killed = serverProcess.kill('SIGKILL'); + if (!killed && serverProcess.exitCode === null && serverProcess.signalCode === null) { + throw new Error('Failed to terminate Agent Host test server'); + } + } + } catch (error) { + if (serverProcess.exitCode === null && serverProcess.signalCode === null) { + throw error; + } + } + await serverExit; + } +} + interface IMockLlmServerHandle { readonly url: string; requestCount(): number; @@ -586,8 +640,8 @@ export interface IMockScenario { const AGENT_HOST_E2E_COVERAGE = process.env['AGENT_HOST_E2E_COVERAGE'] === '1'; -export function getAgentHostE2ETestTimeout(normalTimeoutMs: number, coverageTimeoutMs: number): number { - return AGENT_HOST_E2E_COVERAGE ? coverageTimeoutMs : normalTimeoutMs; +export function getAgentHostE2ETestTimeout(normalTimeoutMs: number, extendedTimeoutMs: number): number { + return AGENT_HOST_E2E_COVERAGE || isWindows ? extendedTimeoutMs : normalTimeoutMs; } function withAgentHostCoverage(environment: NodeJS.ProcessEnv): NodeJS.ProcessEnv {