From 23420012abd670e7ff74dd47c8a81103ca97c0a9 Mon Sep 17 00:00:00 2001
From: Mario Zechner <badlogicgames@gmail.com>
Date: Mon, 4 May 2026 00:45:56 +0200
Subject: [PATCH] fix(coding-agent): decouple codex session cleanup

---
 packages/ai/src/index.ts                      |  6 +----
 .../src/providers/openai-codex-responses.ts   |  3 +++
 packages/ai/src/session-resources.ts          | 24 +++++++++++++++++++
 .../coding-agent/src/core/agent-session.ts    |  4 ++--
 4 files changed, 30 insertions(+), 7 deletions(-)
 create mode 100644 packages/ai/src/session-resources.ts

diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts
index f78cd7c6..9f996059 100644
--- a/packages/ai/src/index.ts
+++ b/packages/ai/src/index.ts
@@ -16,14 +16,10 @@ export type {
 	OpenAICodexResponsesOptions,
 	OpenAICodexWebSocketDebugStats,
 } from "./providers/openai-codex-responses.js";
-export {
-	closeOpenAICodexWebSocketSessions,
-	getOpenAICodexWebSocketDebugStats,
-	resetOpenAICodexWebSocketDebugStats,
-} from "./providers/openai-codex-responses.js";
 export type { OpenAICompletionsOptions } from "./providers/openai-completions.js";
 export type { OpenAIResponsesOptions } from "./providers/openai-responses.js";
 export * from "./providers/register-builtins.js";
+export * from "./session-resources.js";
 export * from "./stream.js";
 export * from "./types.js";
 export * from "./utils/diagnostics.js";
diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts
index a6a60a45..80eaf496 100644
--- a/packages/ai/src/providers/openai-codex-responses.ts
+++ b/packages/ai/src/providers/openai-codex-responses.ts
@@ -22,6 +22,7 @@ if (typeof process !== "undefined" && (process.versions?.node || process.version
 
 import { getEnvApiKey } from "../env-api-keys.js";
 import { clampThinkingLevel } from "../models.js";
+import { registerSessionResourceCleanup } from "../session-resources.js";
 import type {
 	Api,
 	AssistantMessage,
@@ -682,6 +683,8 @@ export function closeOpenAICodexWebSocketSessions(sessionId?: string): void {
 	websocketSessionCache.clear();
 }
 
+registerSessionResourceCleanup(closeOpenAICodexWebSocketSessions);
+
 function isWebSocketSseFallbackActive(sessionId: string | undefined): boolean {
 	return sessionId ? websocketSseFallbackSessions.has(sessionId) : false;
 }
diff --git a/packages/ai/src/session-resources.ts b/packages/ai/src/session-resources.ts
new file mode 100644
index 00000000..d6300926
--- /dev/null
+++ b/packages/ai/src/session-resources.ts
@@ -0,0 +1,24 @@
+export type SessionResourceCleanup = (sessionId?: string) => void;
+
+const sessionResourceCleanups = new Set<SessionResourceCleanup>();
+
+export function registerSessionResourceCleanup(cleanup: SessionResourceCleanup): () => void {
+	sessionResourceCleanups.add(cleanup);
+	return () => {
+		sessionResourceCleanups.delete(cleanup);
+	};
+}
+
+export function cleanupSessionResources(sessionId?: string): void {
+	const errors: unknown[] = [];
+	for (const cleanup of sessionResourceCleanups) {
+		try {
+			cleanup(sessionId);
+		} catch (error) {
+			errors.push(error);
+		}
+	}
+	if (errors.length > 0) {
+		throw new AggregateError(errors, "Failed to cleanup session resources");
+	}
+}
diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts
index 5e93e5f1..db2d6f0d 100644
--- a/packages/coding-agent/src/core/agent-session.ts
+++ b/packages/coding-agent/src/core/agent-session.ts
@@ -26,12 +26,12 @@ import type {
 import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@mariozechner/pi-ai";
 import {
 	clampThinkingLevel,
+	cleanupSessionResources,
 	getSupportedThinkingLevels,
 	isContextOverflow,
 	modelsAreEqual,
 	resetApiProviders,
 } from "@mariozechner/pi-ai";
-import { closeOpenAICodexWebSocketSessions } from "@mariozechner/pi-ai/openai-codex-responses";
 import { theme } from "../modes/interactive/theme/theme.js";
 import { stripFrontmatter } from "../utils/frontmatter.js";
 import { sleep } from "../utils/sleep.js";
@@ -753,7 +753,7 @@ export class AgentSession {
 		);
 		this._disconnectFromAgent();
 		this._eventListeners = [];
-		closeOpenAICodexWebSocketSessions(this.sessionId);
+		cleanupSessionResources(this.sessionId);
 	}
 
 	// =========================================================================