protocol parity

Goes with 009e4e93c3
This commit is contained in:
Connor Peet
2026-03-13 14:50:22 -07:00
parent 706e49e462
commit f8311c303d
15 changed files with 271 additions and 230 deletions

View File

@@ -10,7 +10,7 @@ import { URI } from '../../../../base/common/uri.js';
import { ensureNoDisposablesAreLeakedInTestSuite } from '../../../../base/test/common/utils.js';
import { NullLogService } from '../../../log/common/log.js';
import type { ISessionAction } from '../../common/state/sessionActions.js';
import { isJsonRpcNotification, isJsonRpcResponse, type ICreateSessionParams, type IProtocolMessage, type IProtocolNotification, type IServerHelloParams, type IStateSnapshot } from '../../common/state/sessionProtocol.js';
import { isJsonRpcNotification, isJsonRpcResponse, type ICreateSessionParams, type IInitializeResult, type IProtocolMessage, type IProtocolNotification, type IReconnectResult, type IStateSnapshot } from '../../common/state/sessionProtocol.js';
import { SessionStatus, type ISessionSummary } from '../../common/state/sessionState.js';
import { PROTOCOL_VERSION } from '../../common/state/sessionCapabilities.js';
import type { IProtocolServer, IProtocolTransport } from '../../common/state/sessionTransport.js';
@@ -117,7 +117,7 @@ suite('ProtocolServerHandler', () => {
function connectClient(clientId: string, initialSubscriptions?: readonly URI[]): MockProtocolTransport {
const transport = new MockProtocolTransport();
server.simulateConnection(transport);
transport.simulateMessage(notification('initialize', {
transport.simulateMessage(request(1, 'initialize', {
protocolVersion: PROTOCOL_VERSION,
clientId,
initialSubscriptions,
@@ -144,14 +144,14 @@ suite('ProtocolServerHandler', () => {
ensureNoDisposablesAreLeakedInTestSuite();
test('handshake sends serverHello notification', () => {
test('handshake returns initialize response', () => {
const transport = connectClient('client-1');
const hello = findNotification(transport.sent, 'serverHello');
assert.ok(hello, 'should have sent serverHello');
const params = hello.params as IServerHelloParams;
assert.strictEqual(params.protocolVersion, PROTOCOL_VERSION);
assert.strictEqual(params.serverSeq, stateManager.serverSeq);
const resp = findResponse(transport.sent, 1);
assert.ok(resp, 'should have sent initialize response');
const result = (resp as { result: IInitializeResult }).result;
assert.strictEqual(result.protocolVersion, PROTOCOL_VERSION);
assert.strictEqual(result.serverSeq, stateManager.serverSeq);
});
test('handshake with initialSubscriptions returns snapshots', () => {
@@ -159,11 +159,11 @@ suite('ProtocolServerHandler', () => {
const transport = connectClient('client-1', [sessionUri]);
const hello = findNotification(transport.sent, 'serverHello');
assert.ok(hello);
const params = hello.params as IServerHelloParams;
assert.strictEqual(params.snapshots.length, 1);
assert.strictEqual(params.snapshots[0].resource.toString(), sessionUri.toString());
const resp = findResponse(transport.sent, 1);
assert.ok(resp);
const result = (resp as { result: IInitializeResult }).result;
assert.strictEqual(result.snapshots.length, 1);
assert.strictEqual(result.snapshots[0].resource.toString(), sessionUri.toString());
});
test('subscribe request returns snapshot', async () => {
@@ -249,8 +249,8 @@ suite('ProtocolServerHandler', () => {
stateManager.dispatchServerAction({ type: 'session/ready', session: sessionUri });
const transport1 = connectClient('client-r', [sessionUri]);
const hello = findNotification(transport1.sent, 'serverHello');
const helloSeq = (hello!.params as IServerHelloParams).serverSeq;
const resp = findResponse(transport1.sent, 1);
const initSeq = (resp as { result: IInitializeResult }).result.serverSeq;
transport1.simulateClose();
stateManager.dispatchServerAction({ type: 'session/titleChanged', session: sessionUri, title: 'Title A' });
@@ -258,14 +258,19 @@ suite('ProtocolServerHandler', () => {
const transport2 = new MockProtocolTransport();
server.simulateConnection(transport2);
transport2.simulateMessage(notification('reconnect', {
transport2.simulateMessage(request(1, 'reconnect', {
clientId: 'client-r',
lastSeenServerSeq: helloSeq,
lastSeenServerSeq: initSeq,
subscriptions: [sessionUri],
}));
const replayed = findNotifications(transport2.sent, 'action');
assert.strictEqual(replayed.length, 2);
const reconnectResp = findResponse(transport2.sent, 1);
assert.ok(reconnectResp, 'should have sent reconnect response');
const result = (reconnectResp as { result: IReconnectResult }).result;
assert.strictEqual(result.type, 'replay');
if (result.type === 'replay') {
assert.strictEqual(result.actions.length, 2);
}
});
test('reconnect sends fresh snapshots when gap too large', () => {
@@ -281,16 +286,19 @@ suite('ProtocolServerHandler', () => {
const transport2 = new MockProtocolTransport();
server.simulateConnection(transport2);
transport2.simulateMessage(notification('reconnect', {
transport2.simulateMessage(request(1, 'reconnect', {
clientId: 'client-g',
lastSeenServerSeq: 0,
subscriptions: [sessionUri],
}));
const reconnectResp = findNotification(transport2.sent, 'reconnectResponse');
assert.ok(reconnectResp, 'should receive a reconnectResponse');
const params = reconnectResp!.params as { snapshots: IStateSnapshot[] };
assert.ok(params.snapshots.length > 0, 'should contain snapshots');
const reconnectResp = findResponse(transport2.sent, 1);
assert.ok(reconnectResp, 'should have sent reconnect response');
const result = (reconnectResp as { result: IReconnectResult }).result;
assert.strictEqual(result.type, 'snapshot');
if (result.type === 'snapshot') {
assert.ok(result.snapshots.length > 0, 'should contain snapshots');
}
});
test('client disconnect cleans up', () => {

View File

@@ -16,13 +16,14 @@ import {
JSON_RPC_PARSE_ERROR,
type IActionBroadcastParams,
type IFetchTurnsResult,
type IInitializeResult,
type IJsonRpcErrorResponse,
type IJsonRpcSuccessResponse,
type IListSessionsResult,
type INotificationBroadcastParams,
type IProtocolMessage,
type IProtocolNotification,
type IServerHelloParams,
type IReconnectResult,
type IStateSnapshot,
} from '../../common/state/sessionProtocol.js';
import type { IDeltaAction, ISessionAddedNotification, ISessionRemovedNotification, IUsageAction } from '../../common/state/sessionActions.js';
@@ -244,8 +245,7 @@ function getActionParams(n: IProtocolNotification): IActionBroadcastParams {
/** Perform handshake, create a session, subscribe, and return its URI. */
async function createAndSubscribeSession(c: TestProtocolClient, clientId: string): Promise<URI> {
c.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId });
await c.waitForNotification(n => n.method === 'serverHello');
await c.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId });
await c.call('createSession', { session: nextSessionUri(), provider: 'mock' });
@@ -299,28 +299,25 @@ suite('Protocol WebSocket E2E', function () {
});
// 1. Handshake
test('handshake returns serverHello with protocol version', async function () {
test('handshake returns initialize response with protocol version', async function () {
this.timeout(5_000);
client.notify('initialize', {
const result = await client.call<IInitializeResult>('initialize', {
protocolVersion: PROTOCOL_VERSION,
clientId: 'test-handshake',
initialSubscriptions: [URI.from({ scheme: 'agenthost', path: '/root' })],
});
const hello = await client.waitForNotification(n => n.method === 'serverHello');
const params = hello.params as IServerHelloParams;
assert.strictEqual(params.protocolVersion, PROTOCOL_VERSION);
assert.ok(params.serverSeq >= 0);
assert.ok(params.snapshots.length >= 1, 'should have root state snapshot');
assert.strictEqual(result.protocolVersion, PROTOCOL_VERSION);
assert.ok(result.serverSeq >= 0);
assert.ok(result.snapshots.length >= 1, 'should have root state snapshot');
});
// 2. Create session
test('create session triggers sessionAdded notification', async function () {
this.timeout(10_000);
client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-create-session' });
await client.waitForNotification(n => n.method === 'serverHello');
await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-create-session' });
await client.call('createSession', { session: nextSessionUri(), provider: 'mock' });
@@ -411,8 +408,7 @@ suite('Protocol WebSocket E2E', function () {
test('listSessions returns sessions', async function () {
this.timeout(10_000);
client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-list-sessions' });
await client.waitForNotification(n => n.method === 'serverHello');
await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-list-sessions' });
await client.call('createSession', { session: nextSessionUri(), provider: 'mock' });
await client.waitForNotification(n =>
@@ -440,19 +436,16 @@ suite('Protocol WebSocket E2E', function () {
const client2 = new TestProtocolClient(server.port);
await client2.connect();
client2.notify('reconnect', {
const result = await client2.call<IReconnectResult>('reconnect', {
clientId: 'test-reconnect',
lastSeenServerSeq: missedFromSeq,
subscriptions: [sessionUri],
});
await new Promise(resolve => setTimeout(resolve, 500));
const replayed = client2.receivedNotifications();
assert.ok(replayed.length > 0, 'should receive replayed actions or reconnect response');
const hasActions = replayed.some(n => n.method === 'action');
const hasReconnect = replayed.some(n => n.method === 'reconnectResponse');
assert.ok(hasActions || hasReconnect);
assert.ok(result.type === 'replay' || result.type === 'snapshot', 'should receive replay or snapshot');
if (result.type === 'replay') {
assert.ok(result.actions.length > 0, 'should have replayed actions');
}
client2.close();
});
@@ -502,8 +495,7 @@ suite('Protocol WebSocket E2E', function () {
test('createSession with invalid provider does not crash server', async function () {
this.timeout(10_000);
client.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-invalid-create' });
await client.waitForNotification(n => n.method === 'serverHello');
await client.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-invalid-create' });
// This should return a JSON-RPC error
let gotError = false;
@@ -534,9 +526,9 @@ suite('Protocol WebSocket E2E', function () {
await new Promise(resolve => setTimeout(resolve, 200));
await client.waitForNotification(n => isActionNotification(n, 'session/turnComplete'));
const result = await client.call<IFetchTurnsResult>('fetchTurns', { session: sessionUri, startTurn: 0, count: 10 });
const result = await client.call<IFetchTurnsResult>('fetchTurns', { session: sessionUri, limit: 10 });
assert.ok(result.turns.length >= 2);
assert.ok(result.totalTurns >= 2);
assert.strictEqual(typeof result.hasMore, 'boolean');
});
// ---- Gap tests: coverage ---------------------------------------------------
@@ -599,8 +591,7 @@ suite('Protocol WebSocket E2E', function () {
const client2 = new TestProtocolClient(server.port);
await client2.connect();
client2.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-multi-client-2' });
await client2.waitForNotification(n => n.method === 'serverHello');
await client2.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-multi-client-2' });
await client2.call('subscribe', { resource: sessionUri });
client2.clearReceived();
@@ -627,8 +618,7 @@ suite('Protocol WebSocket E2E', function () {
const client2 = new TestProtocolClient(server.port);
await client2.connect();
client2.notify('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-unsub-helper' });
await client2.waitForNotification(n => n.method === 'serverHello');
await client2.call('initialize', { protocolVersion: PROTOCOL_VERSION, clientId: 'test-unsub-helper' });
await client2.call('subscribe', { resource: sessionUri });
dispatchTurnStarted(client2, sessionUri, 'turn-unsub', 'hello', 1);