From 4a05d02f012c229d6c828ff0ec40e8511eb24f58 Mon Sep 17 00:00:00 2001
From: can1357 <me@can.ac>
Date: Thu, 30 Jul 2026 01:18:45 +0200
Subject: [PATCH] test(extensions): cover session async job wiring

(cherry picked from commit e79cc9e79a68b53332eb6ca04f2e555e25d1b7e0)
---
 .../controllers/extension-ui-controller.ts    |  1 -
 .../sdk-async-job-manager-singleton.test.ts   | 46 ++++++++++++++++++-
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
index 1c49168ae..750bdd193 100644
--- a/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
+++ b/packages/coding-agent/src/modes/controllers/extension-ui-controller.ts
@@ -532,7 +532,6 @@ export class ExtensionUiController {
 					await registeredTool.definition.onSession(event, {
 						...runner!.createContext(),
 						ui: uiContext,
-
 						hasUI: true,
 						compact: instructionsOrOptions => this.#compactSession(instructionsOrOptions),
 					});
diff --git a/packages/coding-agent/test/sdk-async-job-manager-singleton.test.ts b/packages/coding-agent/test/sdk-async-job-manager-singleton.test.ts
index ca66ba796..a6c8bb479 100644
--- a/packages/coding-agent/test/sdk-async-job-manager-singleton.test.ts
+++ b/packages/coding-agent/test/sdk-async-job-manager-singleton.test.ts
@@ -5,9 +5,11 @@ import * as path from "node:path";
 import { AsyncJobManager } from "@oh-my-pi/pi-coding-agent/async/job-manager";
 import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
 import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
-import { createAgentSession } from "@oh-my-pi/pi-coding-agent/sdk";
+import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk";
 import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
+import type { AsyncJobSnapshot } from "@oh-my-pi/pi-coding-agent/session/agent-session";
 import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils";
+import { type } from "arktype";
 
 describe("AsyncJobManager singleton across concurrent top-level sessions", () => {
 	const tempDirs: string[] = [];
@@ -38,7 +40,10 @@ describe("AsyncJobManager singleton across concurrent top-level sessions", () =>
 		AsyncJobManager.resetForTests();
 	});
 
-	async function spawnTopLevelSession(extraSettings?: Record<string, unknown>) {
+	async function spawnTopLevelSession(
+		extraSettings?: Record<string, unknown>,
+		extensions: ExtensionFactory[] = [],
+	) {
 		const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), `pi-sdk-async-singleton-${Snowflake.next()}-`));
 		tempDirs.push(tempDir);
 		const cwd = path.join(tempDir, `project-${Snowflake.next()}`);
@@ -49,6 +54,7 @@ describe("AsyncJobManager singleton across concurrent top-level sessions", () =>
 			agentDir,
 			settings: Settings.isolated({ "bash.autoBackground.enabled": true, ...(extraSettings ?? {}) }),
 			disableExtensionDiscovery: true,
+			extensions,
 			skills: [],
 			contextFiles: [],
 			promptTemplates: [],
@@ -130,6 +136,42 @@ describe("AsyncJobManager singleton across concurrent top-level sessions", () =>
 		}
 	}, 60000);
 
+	it("exposes the owning session's jobs through a production extension context", async () => {
+		let observedSnapshot: AsyncJobSnapshot | null | undefined;
+		const snapshotExtension: ExtensionFactory = pi => {
+			pi.registerTool({
+				name: "capture_async_job_snapshot",
+				label: "Capture async job snapshot",
+				description: "Capture the session-owned async job snapshot for this test.",
+				parameters: type({}),
+				approval: "read",
+				async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
+					observedSnapshot = ctx.getAsyncJobSnapshot();
+					return { content: [{ type: "text", text: "captured" }] };
+				},
+			});
+		};
+		const session = await spawnTopLevelSession(undefined, [snapshotExtension]);
+		const manager = AsyncJobManager.instance();
+		expect(manager).toBeDefined();
+		const release = Promise.withResolvers<string>();
+		const jobId = manager!.register("bash", "extension snapshot test", async () => release.promise, {
+			ownerId: "Main",
+		});
+
+		try {
+			const snapshotTool = session.getToolByName("capture_async_job_snapshot");
+			expect(snapshotTool).toBeDefined();
+			await snapshotTool!.execute("call-snapshot", {});
+
+			expect(observedSnapshot?.running.some(job => job.id === jobId)).toBe(true);
+		} finally {
+			release.resolve("done");
+			await manager!.waitForAll();
+			await session.dispose();
+		}
+	}, 60000);
+
 	it("refuses async bash from a secondary session instead of routing it to the primary's manager", async () => {
 		const primary = await spawnTopLevelSession({ "async.enabled": true });
 		try {