1  import assert from "node:assert/strict";
  2  import { afterEach, test } from "node:test";
  3  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
  4  
  5  type ExecResult = {
  6  	code: number;
  7  	stdout: string;
  8  	stderr?: string;
  9  };
 10  
 11  type EventHandler = (event: unknown, ctx: TestContext) => Promise<void> | void;
 12  
 13  type TestContext = {
 14  	hasUI: boolean;
 15  	model?: { id: string };
 16  	sessionManager: {
 17  		getBranch(): unknown[];
 18  	};
 19  	getContextUsage(): { tokens?: number } | undefined;
 20  };
 21  
 22  type Harness = {
 23  	handlers: Map<string, EventHandler>;
 24  	execCalls: Array<{ command: string; args: string[] }>;
 25  	pi: ExtensionAPI;
 26  };
 27  
 28  const ORIGINAL_CMUX_SOCKET_PATH = process.env.CMUX_SOCKET_PATH;
 29  const ORIGINAL_CMUX_WORKSPACE_ID = process.env.CMUX_WORKSPACE_ID;
 30  
 31  function createContext(): TestContext {
 32  	return {
 33  		hasUI: true,
 34  		sessionManager: {
 35  			getBranch: () => [],
 36  		},
 37  		getContextUsage: () => undefined,
 38  	};
 39  }
 40  
 41  function createHarness(treeOutputs: string[]): Harness {
 42  	const handlers = new Map<string, EventHandler>();
 43  	const execCalls: Array<{ command: string; args: string[] }> = [];
 44  	let treeIndex = 0;
 45  
 46  	const pi = {
 47  		on(event: string, handler: EventHandler): void {
 48  			handlers.set(event, handler);
 49  		},
 50  		exec: async (command: string, args: string[]): Promise<ExecResult> => {
 51  			execCalls.push({ command, args });
 52  			assert.equal(command, "cmux");
 53  
 54  			if (args[0] === "tree") {
 55  				const stdout = treeOutputs[treeIndex] ?? treeOutputs.at(-1) ?? "";
 56  				treeIndex += 1;
 57  				return { code: 0, stdout };
 58  			}
 59  
 60  			return { code: 0, stdout: "" };
 61  		},
 62  		getSessionName: (): string => "alpha-session",
 63  		getThinkingLevel: (): "off" => "off",
 64  	} as unknown as ExtensionAPI;
 65  
 66  	return { handlers, execCalls, pi };
 67  }
 68  
 69  async function loadExtension() {
 70  	process.env.CMUX_SOCKET_PATH = "/tmp/cmux.sock";
 71  	process.env.CMUX_WORKSPACE_ID = "workspace-shell-123";
 72  	const moduleUrl = new URL(`../index.ts?test=${Date.now()}`, import.meta.url);
 73  	return (await import(moduleUrl.href)).default as (pi: ExtensionAPI) => void;
 74  }
 75  
 76  afterEach(() => {
 77  	if (ORIGINAL_CMUX_SOCKET_PATH === undefined) {
 78  		delete process.env.CMUX_SOCKET_PATH;
 79  	} else {
 80  		process.env.CMUX_SOCKET_PATH = ORIGINAL_CMUX_SOCKET_PATH;
 81  	}
 82  
 83  	if (ORIGINAL_CMUX_WORKSPACE_ID === undefined) {
 84  		delete process.env.CMUX_WORKSPACE_ID;
 85  	} else {
 86  		process.env.CMUX_WORKSPACE_ID = ORIGINAL_CMUX_WORKSPACE_ID;
 87  	}
 88  });
 89  
 90  test("agent_end sync uses the shell workspace id instead of cmux current-workspace", async () => {
 91  	const cmuxExtension = await loadExtension();
 92  	const harness = createHarness([
 93  		'workspace workspace:1 "alpha-session"\n  pane pane:1\n    surface surface:1\n',
 94  		'workspace workspace:1 "original-title"\n  pane pane:1\n    surface surface:1\n',
 95  	]);
 96  	cmuxExtension(harness.pi);
 97  
 98  	const sessionStartHandler = harness.handlers.get("session_start");
 99  	const agentEndHandler = harness.handlers.get("agent_end");
100  	assert.ok(sessionStartHandler, "expected session_start handler");
101  	assert.ok(agentEndHandler, "expected agent_end handler");
102  
103  	await sessionStartHandler({}, createContext());
104  	harness.execCalls.length = 0;
105  
106  	await agentEndHandler({}, createContext());
107  
108  	assert.equal(
109  		harness.execCalls.some((call) => call.command === "cmux" && call.args[0] === "current-workspace"),
110  		false,
111  	);
112  
113  	assert.deepEqual(
114  		harness.execCalls.find((call) => call.command === "cmux" && call.args[0] === "tree")?.args,
115  		["tree", "--workspace", "workspace-shell-123"],
116  	);
117  	assert.deepEqual(
118  		harness.execCalls.find((call) => call.command === "cmux" && call.args[0] === "rename-workspace")?.args,
119  		["rename-workspace", "--workspace", "workspace-shell-123", "alpha-session"],
120  	);
121  });