Skip to content

Commit de303fb

Browse files
kilingzhangkilingzhang
committed
feat: recover ACP runtime service sessions
1 parent 35146cf commit de303fb

11 files changed

Lines changed: 2473 additions & 159 deletions

File tree

‎apps/workbench/src/features/sessions/SessionsScreen.tsx‎

Lines changed: 117 additions & 98 deletions
Large diffs are not rendered by default.

‎src/bin.ts‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,11 @@ import { spawn } from "node:child_process";
55
import { runFreeAuthCommand } from "./auth-bin.js";
66
import { runFreeBridgeCommand } from "./client/relay-bridge.js";
77
import { runFreeHostCommand } from "./host/bin.js";
8+
import {
9+
connectAcpRuntimeServiceClient,
10+
runAcpRuntimeService,
11+
type AcpRuntimeServiceManagedSession,
12+
} from "./host/runtime-service.js";
813
import { resolveCurrentFreeExecutablePath } from "./launcher.js";
914

1015
async function main(argv: readonly string[]): Promise<void> {
@@ -38,9 +43,139 @@ async function main(argv: readonly string[]): Promise<void> {
3843
await runFreeHostCommand(rest);
3944
return;
4045
}
46+
if (command === "runtime") {
47+
await runFreeRuntimeCommand(rest);
48+
return;
49+
}
50+
if (command === "runtime-service") {
51+
if (rest[0] && rest[0] !== "run") {
52+
throw new Error(`Unknown runtime-service command: ${rest[0]}`);
53+
}
54+
await runAcpRuntimeService();
55+
return;
56+
}
4157
throw new Error(`Unknown free command: ${command}`);
4258
}
4359

60+
async function runFreeRuntimeCommand(argv: readonly string[]): Promise<void> {
61+
const [command, ...rest] = argv;
62+
if (!command || command === "--help" || command === "-h") {
63+
printRuntimeHelp();
64+
return;
65+
}
66+
if (command === "status") {
67+
const client = await connectAcpRuntimeServiceClient({});
68+
try {
69+
const status = await client.management.status();
70+
process.stdout.write(
71+
[
72+
`runtime: ${status.instanceId}`,
73+
`sessions: ${status.sessionCount}`,
74+
`active turns: ${status.activeTurns}`,
75+
`attached clients: ${status.peerCount}`,
76+
].join("\n") + "\n",
77+
);
78+
} finally {
79+
client.close();
80+
}
81+
return;
82+
}
83+
if (command === "sessions") {
84+
await runFreeRuntimeSessionsCommand(rest);
85+
return;
86+
}
87+
throw new Error(`Unknown runtime command: ${command}`);
88+
}
89+
90+
async function runFreeRuntimeSessionsCommand(argv: readonly string[]): Promise<void> {
91+
const [command, ...rest] = argv;
92+
if (!command || command === "--help" || command === "-h") {
93+
printRuntimeSessionsHelp();
94+
return;
95+
}
96+
const client = await connectAcpRuntimeServiceClient({});
97+
try {
98+
if (command === "list") {
99+
printManagedSessions(await client.management.listSessions());
100+
return;
101+
}
102+
if (command === "close") {
103+
if (rest[0] === "--all") {
104+
const sessions = await client.management.listSessions();
105+
for (const session of sessions) {
106+
await client.management.closeSession(session.id);
107+
}
108+
process.stdout.write(`Closed ${sessions.length} runtime session(s).\n`);
109+
return;
110+
}
111+
const sessionId = rest[0];
112+
if (!sessionId) {
113+
throw new Error("Missing runtime session id.");
114+
}
115+
await client.management.closeSession(sessionId);
116+
process.stdout.write(`Closed runtime session ${sessionId}.\n`);
117+
return;
118+
}
119+
} finally {
120+
client.close();
121+
}
122+
throw new Error(`Unknown runtime sessions command: ${command}`);
123+
}
124+
125+
function printManagedSessions(sessions: readonly AcpRuntimeServiceManagedSession[]): void {
126+
if (sessions.length === 0) {
127+
process.stdout.write("No managed ACP runtime sessions.\n");
128+
return;
129+
}
130+
const rows = sessions.map((session) => [
131+
session.id,
132+
session.status,
133+
String(session.activeTurns),
134+
session.updatedAt ?? "-",
135+
session.title ?? "-",
136+
]);
137+
const headers = ["SESSION", "STATUS", "ACTIVE_TURNS", "UPDATED", "TITLE"];
138+
const widths = headers.map((header, index) =>
139+
Math.max(header.length, ...rows.map((row) => row[index]?.length ?? 0)),
140+
);
141+
process.stdout.write(`${formatRuntimeRow(headers, widths)}\n`);
142+
for (const row of rows) {
143+
process.stdout.write(`${formatRuntimeRow(row, widths)}\n`);
144+
}
145+
}
146+
147+
function formatRuntimeRow(row: readonly string[], widths: readonly number[]): string {
148+
return row.map((value, index) => value.padEnd(widths[index] ?? 0)).join(" ");
149+
}
150+
151+
function printRuntimeHelp(): void {
152+
process.stdout.write(
153+
[
154+
"Usage:",
155+
" free runtime status",
156+
" free runtime sessions list",
157+
" free runtime sessions close <session-id>",
158+
" free runtime sessions close --all",
159+
"",
160+
"Lifecycle:",
161+
" Runtime sessions are owned by the local ACP runtime service.",
162+
" Host restarts do not close managed ACP sessions.",
163+
" ACP session/close and runtime sessions close release the underlying ACP agent process.",
164+
].join("\n") + "\n",
165+
);
166+
}
167+
168+
function printRuntimeSessionsHelp(): void {
169+
process.stdout.write(
170+
[
171+
"Usage:",
172+
" free runtime sessions list",
173+
" free runtime sessions close <session-id>",
174+
" free runtime sessions close --all",
175+
].join("\n") + "\n",
176+
);
177+
}
178+
44179
function runSupervisedBridge(args: readonly string[]): Promise<void> {
45180
return new Promise((resolve, reject) => {
46181
const launcher = resolveCurrentFreeLauncher();

‎src/client/stdio-bridge.test.ts‎

Lines changed: 92 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -157,14 +157,13 @@ describe("createAcpRemoteStdioBridge", () => {
157157
}
158158
});
159159

160-
it("auto-authorizes the matching session selection before forwarding session/new", async () => {
160+
it("opens authorization for session/new so workspace selection remains explicit", async () => {
161161
const input = new PassThrough();
162162
const output = new PassThrough();
163-
const fetchBodies: unknown[] = [];
164-
const fetchUrls: string[] = [];
163+
const openedUrls: string[] = [];
165164
const fetchMock = vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
166-
fetchUrls.push(String(url));
167-
fetchBodies.push(JSON.parse(String(init?.body ?? "{}")) as unknown);
165+
void url;
166+
void init;
168167
return new Response(JSON.stringify({ ok: true }), {
169168
headers: { "content-type": "application/json" },
170169
status: 200,
@@ -180,6 +179,9 @@ describe("createAcpRemoteStdioBridge", () => {
180179
clientId: "client-1",
181180
connectionId: "connection-1",
182181
input,
182+
openAuthUrl(url) {
183+
openedUrls.push(url);
184+
},
183185
output,
184186
relayUrl: "ws://127.0.0.1:8791",
185187
socketFactory() {
@@ -217,7 +219,6 @@ describe("createAcpRemoteStdioBridge", () => {
217219
})}\n`,
218220
);
219221

220-
await waitFor(() => fetchMock.mock.calls.length === 2);
221222
await waitFor(() =>
222223
sockets[0]?.sent.some((message) => {
223224
try {
@@ -227,11 +228,6 @@ describe("createAcpRemoteStdioBridge", () => {
227228
}
228229
}) ?? false,
229230
);
230-
const sessionAuthorizeUrl = new URL(fetchUrls[1]!);
231-
const sessionAuthorizeBody = fetchBodies[1] as {
232-
hostId?: string;
233-
sessionSelectionId?: string;
234-
};
235231
const outboundMessage = sockets[0]!.sent.find((message) => {
236232
try {
237233
return JSON.parse(message).method === "session/new";
@@ -245,13 +241,95 @@ describe("createAcpRemoteStdioBridge", () => {
245241
const outboundSelectionId =
246242
outbound.params?._meta?.["acp-runtime/remote/sessionSelectionId"];
247243

244+
expect(fetchMock).toHaveBeenCalledTimes(1);
245+
expect(openedUrls).toHaveLength(1);
246+
const sessionAuthorizeUrl = new URL(openedUrls[0]!);
248247
expect(sessionAuthorizeUrl.searchParams.get("sessionSelectionId")).toBe(
249248
outboundSelectionId,
250249
);
251-
expect(sessionAuthorizeBody).toMatchObject({
252-
hostId: "host-1",
253-
sessionSelectionId: outboundSelectionId,
250+
} finally {
251+
bridge.close();
252+
vi.unstubAllGlobals();
253+
}
254+
});
255+
256+
it("opens authorization for session/new after bridge restart before initialize provides auth url", async () => {
257+
const input = new PassThrough();
258+
const output = new PassThrough();
259+
const openedUrls: string[] = [];
260+
const fetchMock = vi.fn(async (url: string | URL | Request, init?: RequestInit) => {
261+
void url;
262+
void init;
263+
return new Response(JSON.stringify({ ok: true }), {
264+
headers: { "content-type": "application/json" },
265+
status: 200,
254266
});
267+
});
268+
vi.stubGlobal("fetch", fetchMock);
269+
const sockets: TestSocket[] = [];
270+
const bridge = createAcpRemoteStdioBridge({
271+
autoAuthorize: {
272+
accountSession: "account-session-1",
273+
hostId: "host-1",
274+
},
275+
clientId: "client-1",
276+
connectionId: "connection-1",
277+
input,
278+
openAuthUrl(url) {
279+
openedUrls.push(url);
280+
},
281+
output,
282+
relayUrl: "ws://127.0.0.1:8791",
283+
socketFactory() {
284+
const socket = new TestSocket();
285+
sockets.push(socket);
286+
return socket;
287+
},
288+
});
289+
290+
try {
291+
input.write(
292+
`${JSON.stringify({
293+
id: 3,
294+
jsonrpc: "2.0",
295+
method: "session/new",
296+
params: { cwd: "/tmp/project", mcpServers: [] },
297+
})}\n`,
298+
);
299+
300+
await waitFor(() =>
301+
sockets[0]?.sent.some((message) => {
302+
try {
303+
return JSON.parse(message).method === "session/new";
304+
} catch {
305+
return false;
306+
}
307+
}) ?? false,
308+
);
309+
const outboundMessage = sockets[0]!.sent.find((message) => {
310+
try {
311+
return JSON.parse(message).method === "session/new";
312+
} catch {
313+
return false;
314+
}
315+
})!;
316+
const outbound = JSON.parse(outboundMessage) as {
317+
params?: { _meta?: Record<string, string> };
318+
};
319+
const outboundSelectionId =
320+
outbound.params?._meta?.["acp-runtime/remote/sessionSelectionId"];
321+
322+
expect(fetchMock).not.toHaveBeenCalled();
323+
expect(openedUrls).toHaveLength(1);
324+
const sessionAuthorizeUrl = new URL(openedUrls[0]!);
325+
expect(sessionAuthorizeUrl.origin).toBe("http://127.0.0.1:8790");
326+
expect(sessionAuthorizeUrl.pathname).toBe("/authorize");
327+
expect(sessionAuthorizeUrl.searchParams.get("connectionId")).toBe(
328+
"connection-1",
329+
);
330+
expect(sessionAuthorizeUrl.searchParams.get("sessionSelectionId")).toBe(
331+
outboundSelectionId,
332+
);
255333
} finally {
256334
bridge.close();
257335
vi.unstubAllGlobals();

‎src/client/stdio-bridge.ts‎

Lines changed: 46 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -947,45 +947,49 @@ export function createAcpRemoteStdioBridge(
947947
);
948948
const isAuthenticate = isRelayAuthenticateRequest(outbound);
949949
const isSessionNew = isRelaySessionNewRequest(outbound);
950+
const authorizationUrl =
951+
authUrl ??
952+
(isSessionNew && options.autoAuthorize
953+
? createConnectionAuthorizationUrl({
954+
connectionId,
955+
relayUrl: String(options.relayUrl),
956+
})
957+
: undefined);
950958
if (isAuthenticate || isSessionNew) {
951-
if (authUrl && options.autoAuthorize) {
959+
if (authorizationUrl && options.autoAuthorize) {
952960
const selectionId = sessionSelection?.selectionId;
953961
if (selectionId) {
954962
const urlWithSelection = addSessionSelectionIdToAuthUrl(
955-
authUrl,
963+
authorizationUrl,
956964
selectionId,
957965
);
958-
debugLog("auto-authorize relay session selection");
959-
await authorizeRelay({
960-
authUrl: urlWithSelection,
961-
sessionSelectionId: selectionId,
962-
...options.autoAuthorize,
963-
});
964-
debugLog("auto-authorize relay session selection completed");
966+
debugLog("open authorization url trigger=session/new");
967+
openAuthUrl(toWorkbenchAuthorizationUrl(urlWithSelection));
968+
armAuthRequestTimeout(outbound, urlWithSelection);
965969
} else {
966970
if (!authorizePromise) {
967971
debugLog("auto-authorize relay browser authentication");
968972
authorizePromise = authorizeRelay({
969-
authUrl,
973+
authUrl: authorizationUrl,
970974
...options.autoAuthorize,
971975
});
972976
}
973977
await authorizePromise;
974978
}
975979
} else if (authorizePromise) {
976980
await authorizePromise;
977-
} else if (authUrl) {
981+
} else if (authorizationUrl) {
978982
debugLog(
979983
`open authorization url trigger=${
980984
isSessionNew ? "session/new" : "authenticate"
981985
}`,
982986
);
983987
const urlWithSelection = sessionSelection?.selectionId
984988
? addSessionSelectionIdToAuthUrl(
985-
authUrl,
989+
authorizationUrl,
986990
sessionSelection.selectionId,
987991
)
988-
: authUrl;
992+
: authorizationUrl;
989993
openAuthUrl(toWorkbenchAuthorizationUrl(urlWithSelection));
990994
armAuthRequestTimeout(outbound, urlWithSelection);
991995
}
@@ -1344,6 +1348,35 @@ function addSessionSelectionIdToAuthUrl(
13441348
}
13451349
}
13461350

1351+
function createConnectionAuthorizationUrl(input: {
1352+
connectionId: string;
1353+
relayUrl: string;
1354+
}): string | undefined {
1355+
const workbenchOrigin = resolveFreeWorkbenchOriginForRelayUrl({
1356+
relayUrl: input.relayUrl,
1357+
});
1358+
const origin = workbenchOrigin ?? relayUrlToHttpOrigin(input.relayUrl);
1359+
if (!origin) {
1360+
return undefined;
1361+
}
1362+
try {
1363+
const url = new URL("/authorize", origin);
1364+
url.searchParams.set("connectionId", input.connectionId);
1365+
return url.toString();
1366+
} catch {
1367+
return undefined;
1368+
}
1369+
}
1370+
1371+
function relayUrlToHttpOrigin(relayUrl: string): string | undefined {
1372+
try {
1373+
const url = new URL(relayUrl.replace(/^ws(s?):\/\//, "http$1://"));
1374+
return url.origin;
1375+
} catch {
1376+
return undefined;
1377+
}
1378+
}
1379+
13471380
function toWorkbenchAuthorizationUrl(authUrl: string): string {
13481381
try {
13491382
const relayUrl = new URL(authUrl);

0 commit comments

Comments
 (0)