test: add test for server notifications

AlexandrSuhinin · AlexandrSuhinin · commit 36b124b887cd · 2025-12-12T15:50:43.000+02:00
diff --git a/src/CodexACPAgent.ts b/src/CodexACPAgent.ts
@@ -80,11 +80,16 @@ export class CodexACPAgent implements acp.Agent {
         return {};
     }
 
-    async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
-        const sessionState = this.sessions.get(params.sessionId);
+    getSessionState(sessionId: string): SessionState {
+        const sessionState = this.sessions.get(sessionId);
         if (!sessionState) {
-            throw new Error(`Session ${params.sessionId} not found`);
+            throw new Error(`Session ${sessionId} not found`);
         }
+        return sessionState;
+    }
+
+    async prompt(params: acp.PromptRequest): Promise<acp.PromptResponse> {
+        const sessionState = this.getSessionState(params.sessionId);
 
         sessionState.pendingPrompt?.abort();
         sessionState.pendingPrompt = new AbortController();
diff --git a/src/CodexAppServerClient.ts b/src/CodexAppServerClient.ts
@@ -26,9 +26,13 @@ export class CodexAppServerClient {
 
     constructor(connection: MessageConnection) {
         this.connection = connection;
-        this.onServerNotification((notification) => {
-            for (const callback of this.transportEventHandlers) {
-                callback({ eventType: "notification", ...notification});
+        this.connection.onUnhandledNotification((data) => {
+            const serverNotification = data as ServerNotification ?? null;
+            if (serverNotification) {
+                this.notify(serverNotification);
+            }
+            for (const callback of this.codexEventHandlers) {
+                callback({ eventType: "notification", ...serverNotification});
             }
         });
     }
@@ -82,35 +86,23 @@ export class CodexAppServerClient {
     }
 
     onServerNotification(callback: (event: ServerNotification) => void){
-        this.connection.onUnhandledNotification((data) => {
-            const serverNotification = data as ServerNotification ?? null;
-            if (serverNotification) {
-                callback(serverNotification)
-            }
-        });
+        this.notificationHandlers.push(callback);
     }
 
-    onUnhandledNotification(callback: (data: NotificationMessage) => void){
-        this.connection.onUnhandledNotification((data) => {
-            const event = this.getEventMessage(data);
-            if (!event) {
-                callback(data)
-            }
-        });
+    private codexEventHandlers: Array<(event: CodexConnectionEvent) => void> = [];
+    onClientTransportEvent(callback: (event: CodexConnectionEvent) => void){
+        this.codexEventHandlers.push(callback);
     }
 
-    private getEventMessage(data: NotificationMessage): EventMsg | null {
-        const params = data.params;
-        return (params as { msg?: EventMsg })?.msg ?? null;
-    }
-
-    private transportEventHandlers: Array<(event: ClientTransportEvent) => void> = [];
-    onClientTransportEvent(callback: (event: ClientTransportEvent) => void){
-        this.transportEventHandlers.push(callback);
+    private notificationHandlers: Array<(event: ServerNotification) => void> = [];
+    private notify(notification: ServerNotification) {
+        for (const notificationHandler of this.notificationHandlers) {
+            notificationHandler(notification);
+        }
     }
 
     private async sendRequest<R>(request: CodexRequest): Promise<R> {
-        for (const callback of this.transportEventHandlers) {
+        for (const callback of this.codexEventHandlers) {
             callback({ eventType: "request", ...request});
         }
         let result: any;
@@ -120,14 +112,14 @@ export class CodexAppServerClient {
         else {
             await this.connection.sendRequest<R>(request.method);
         }
-        for (const callback of this.transportEventHandlers) {
+        for (const callback of this.codexEventHandlers) {
             callback({ eventType: "response", ...result});
         }
         return result;
     }
 }
 
-export type ClientTransportEvent = { eventType: "request" } & CodexRequest | { eventType: "response" } & unknown | { eventType: "notification" } & ServerNotification;
+export type CodexConnectionEvent = { eventType: "request" } & CodexRequest | { eventType: "response" } & unknown | { eventType: "notification" } & ServerNotification;
 
 type CodexRequest = DistributiveOmit<ClientRequest, "id">
 
diff --git a/src/CodexEventHandler.ts b/src/CodexEventHandler.ts
@@ -13,6 +13,7 @@ import type {
     ThreadItem,
     TurnPlanUpdatedNotification
 } from "./app-server/v2";
+import {readFile} from "node:fs/promises";
 
 export class CodexEventHandler {
 
@@ -33,7 +34,6 @@ export class CodexEventHandler {
     }
 
     private async createUpdateEvent(notification: ServerNotification): Promise<UpdateSessionEvent | null> {
-        //TODO should take flow and return flow
         switch (notification.method) {
             case "item/agentMessage/delta":
                 return await this.createTextEvent(notification.params);
@@ -144,12 +144,7 @@ export class CodexEventHandler {
     }
 
     private async createPatchContent(change: FileUpdateChange): Promise<ToolCallContent | null> {
-        const textResponse = await this.connection.readTextFile({
-            sessionId: this.sessionState.sessionId,
-            path: change.path
-        });
-
-        const oldContent = textResponse.content;
+        const oldContent = await readFile(change.path, { encoding: "utf8" });
         const newContent = applyPatch(oldContent, change.diff);
         if (!newContent) {
             return null
diff --git a/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts b/src/__tests__/CodexACPAgent/CodexAcpClient.test.ts
@@ -1,6 +1,8 @@
-import {describe, expect, it, vi, beforeEach} from 'vitest';
+import {describe, expect, it, vi, beforeEach, afterEach} from 'vitest';
 import type {CodexAuthRequest} from "../../CodexAuthMethod";
 import {createTestFixture, type TestFixture} from "../acp-test-utils";
+import type {ServerNotification} from "../../app-server";
+import type {SessionState} from "../../CodexACPAgent";
 
 describe('ACP server test', () => {
 
@@ -20,7 +22,7 @@ describe('ACP server test', () => {
         const newSessionResponse = await codexAcpAgent.newSession({cwd: "", mcpServers: []});
         codexAcpAgent.prompt({ sessionId: newSessionResponse.sessionId, prompt: [{type: "text", text: "Hi!"}] });
 
-        const transportDump = fixture.getTransportDump(ignoredFields);
+        const transportDump = fixture.getCodexConnectionDump(ignoredFields);
         await expect(transportDump).toMatchFileSnapshot("data/start-conversation.json");
     });
 
@@ -29,13 +31,13 @@ describe('ACP server test', () => {
 
         await codexAcpAgent.initialize({protocolVersion: 1});
         await fixture.getCodexAcpClient().logout();
-        fixture.clearTransportDump();
+        fixture.clearCodexConnectionDump();
 
         await expect(
             codexAcpAgent.newSession({cwd: "", mcpServers: []})
         ).rejects.toThrow("Authentication required");
 
-        const transportDump = fixture.getTransportDump(ignoredFields);
+        const transportDump = fixture.getCodexConnectionDump(ignoredFields);
         await expect(transportDump).toMatchFileSnapshot("data/auth-failed.json");
     });
 
@@ -44,14 +46,56 @@ describe('ACP server test', () => {
 
         await codexAcpAgent.initialize({protocolVersion: 1});
         await fixture.getCodexAcpClient().logout();
-        fixture.clearTransportDump();
+        fixture.clearCodexConnectionDump();
 
         const authRequest: CodexAuthRequest = { methodId: "api-key", _meta: {apiKey: "TOKEN"} }
         await codexAcpAgent.authenticate(authRequest);
         const newSessionResponse = await codexAcpAgent.newSession({cwd: "", mcpServers: []});
         expect(newSessionResponse.sessionId).toBeDefined()
 
-        const transportDump = fixture.getTransportDump(ignoredFields);
+        const transportDump = fixture.getCodexConnectionDump(ignoredFields);
         await expect(transportDump).toMatchFileSnapshot("data/auth-with-key.json");
     });
+
+    function loadNotifications(){
+        //TODO load from file
+        const serverNotifications: ServerNotification[] = [
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "He", }},
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "ll", }},
+            { method: "item/agentMessage/delta", params: { threadId: "string", turnId: "string", itemId: "string", delta: "o!", }},
+        ];
+        function onServerNotification(callback: (event: ServerNotification) => void){
+            for (const notification of serverNotifications) {
+                callback(notification);
+            }
+        }
+        return onServerNotification;
+    }
+
+    it('should map events from dump', async () => {
+        fixture.getCodexAppServerClient().onServerNotification = loadNotifications();
+
+        const codexAcpAgent = fixture.getCodexAcpAgent();
+
+        fixture.getCodexAppServerClient().turnStart = vi.fn().mockResolvedValue(undefined);
+        fixture.getCodexAppServerClient().awaitTurnCompleted = vi.fn().mockResolvedValue(undefined);
+        fixture.getCodexAcpAgent().getSessionState = vi.fn().mockResolvedValue({ pendingPrompt: null, sessionId: "id" });
+
+        await codexAcpAgent.prompt({ sessionId: "id", prompt: [{type: "text", text: ""}] });
+
+        expect(fixture.getAcpConnectionDump([])).toMatchFileSnapshot("data/output-acp-events.json");
+
+    }, 90_000);
+
+    //dev-time test
+    it.skip('should convert session notification to acp events', async () => {
+        fixture.onCodexConnectionEvent((event) => {
+            console.log(JSON.stringify(event, null, 2));
+        });
+        const codexAcpAgent = fixture.getCodexAcpAgent();
+        await codexAcpAgent.initialize({protocolVersion: 1});
+        const newSessionResponse = await codexAcpAgent.newSession({cwd: "/home/alex/work/spring-petclinic/", mcpServers: []});
+        fixture.clearCodexConnectionDump();
+        await codexAcpAgent.prompt({ sessionId: newSessionResponse.sessionId, prompt: [{type: "text", text: "Add method `minus` to Math Utils."}] });
+    }, 90_000);
 });
diff --git a/src/__tests__/CodexACPAgent/data/input-server-events.json b/src/__tests__/CodexACPAgent/data/input-server-events.json
diff --git a/src/__tests__/CodexACPAgent/data/output-acp-events.json b/src/__tests__/CodexACPAgent/data/output-acp-events.json
@@ -0,0 +1,42 @@
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "He"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "ll"
+        }
+      }
+    }
+  ]
+}
+{
+  "method": "sessionUpdate",
+  "args": [
+    {
+      "update": {
+        "sessionUpdate": "agent_message_chunk",
+        "content": {
+          "type": "text",
+          "text": "o!"
+        }
+      }
+    }
+  ]
+}
diff --git a/src/__tests__/acp-test-utils.ts b/src/__tests__/acp-test-utils.ts
@@ -1,28 +1,54 @@
 import {CodexAcpClient} from '../CodexAcpClient';
-import {type ClientTransportEvent, CodexAppServerClient} from '../CodexAppServerClient';
+import {type CodexConnectionEvent, CodexAppServerClient} from '../CodexAppServerClient';
 import {startCodexConnection} from "../CodexJsonRpcConnection";
 import {CodexACPAgent} from "../CodexACPAgent";
+import type {AgentSideConnection} from "@agentclientprotocol/sdk";
+
+export type MethodCallEvent = { method: string; args: any[] };
+
+function createSmartMock<T extends object>(onCall: (event: MethodCallEvent) => void) {
+    return new Proxy({} as T, {
+        get(_, prop) {
+            return (...args: any[]) => {
+                onCall({ method: String(prop), args });
+                return { mock: "Mocked return" };
+            };
+        }
+    });
+}
 
 export interface TestFixture {
     getCodexAppServerClient(): CodexAppServerClient,
     getCodexAcpClient(): CodexAcpClient,
     getCodexAcpAgent(): CodexACPAgent,
-    getTransportEvents(): ClientTransportEvent[],
-    getTransportDump(ignoredFields: string[]): string,
-    clearTransportDump(): void
+
+    onCodexConnectionEvent(handler: (event: CodexConnectionEvent) => void): void,
+    getCodexConnectionDump(ignoredFields: string[]): string,
+    clearCodexConnectionDump(): void,
+
+    onAcpConnectionEvent(handler: (event: MethodCallEvent) => void): void,
+    getAcpConnectionDump(ignoredFields: string[]): string,
+    clearAcpConnectionDump(): void,
 }
 
 export function createTestFixture(): TestFixture {
     const pathToCodex = "././node_modules/.bin/codex"
-    const mockedAcpConnection = { } as any;
+    const acpConnectionEvents: MethodCallEvent[] = []
+    const acpEventHandlers: ((event: MethodCallEvent) => void)[] = [];
+    const acpConnection = createSmartMock<AgentSideConnection>((event) => {
+        acpConnectionEvents.push(event);
+        acpEventHandlers.forEach(handler => handler(event));
+    });
     const codexAppServerClient = new CodexAppServerClient(startCodexConnection(pathToCodex));
 
     const codexAcpClient = new CodexAcpClient(codexAppServerClient);
-    const codexAcpAgent = new CodexACPAgent(mockedAcpConnection, codexAcpClient);
+    const codexAcpAgent = new CodexACPAgent(acpConnection, codexAcpClient);
 
-    const transportEvents: ClientTransportEvent[] = []
+    const transportEvents: CodexConnectionEvent[] = []
+    const codexEventHandlers: ((event: CodexConnectionEvent) => void)[] = [];
     codexAppServerClient.onClientTransportEvent((event) => {
         transportEvents.push(event);
+        codexEventHandlers.forEach(handler => handler(event));
     });
 
     return {
@@ -32,23 +58,38 @@ export function createTestFixture(): TestFixture {
         getCodexAcpClient(): CodexAcpClient {
             return codexAcpClient;
         },
-        getTransportDump(ignoredFields: string[]): string {
-            function stringify(obj: any, anonymizedFields: string[] = []) {
-                function fieldAnonymizer(key: string, value: any): any {
-                    return anonymizedFields.includes(key) ? key : value;
-                }
-                return JSON.stringify(obj, fieldAnonymizer, 2);
-            }
-            return this.getTransportEvents().map(event => stringify(event, ignoredFields)).join("\n");
+        getCodexConnectionDump(ignoredFields: string[]): string {
+            return createArrayDump(transportEvents, ignoredFields);
         },
-        getTransportEvents(): ClientTransportEvent[] {
-            return transportEvents;
+        onCodexConnectionEvent(handler: (event: CodexConnectionEvent) => void): void {
+            codexEventHandlers.push(handler);
         },
         getCodexAppServerClient(): CodexAppServerClient {
             return codexAppServerClient;
         },
-        clearTransportDump(): void {
+        clearCodexConnectionDump(): void {
             transportEvents.splice(0, transportEvents.length);
+        },
+        onAcpConnectionEvent(handler: (event: MethodCallEvent) => void): void {
+            acpEventHandlers.push(handler);
+        },
+        getAcpConnectionDump(ignoredFields: string[]): string {
+            return createArrayDump(acpConnectionEvents, ignoredFields);
+        },
+        clearAcpConnectionDump(){
+            acpConnectionEvents.splice(0, acpConnectionEvents.length);
         }
     };
+}
+
+
+function createObjectDump(obj: any, anonymizedFields: string[] = []) {
+    function fieldAnonymizer(key: string, value: any): any {
+        return anonymizedFields.includes(key) ? key : value;
+    }
+    return JSON.stringify(obj, fieldAnonymizer, 2);
+}
+
+function createArrayDump(objects: any[], anonymizedFields: string[]): string {
+    return objects.map(event => createObjectDump(event, anonymizedFields)).join("\n");
 }
