Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 75 additions & 0 deletions src/agent/background-shell-tool.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,81 @@ describe("background shell through the agent toolset", () => {
}
});

test("shell_collect with an aborted signal releases as running without killing the child", async () => {
const toolset = await createAgentToolset({
cwd: process.cwd(),
permissionGate: gate(process.cwd()),
});
try {
const started = await toolset.dynamicRunner.run(
{
id: "a-start",
name: "run_shell",
arguments: {
command: "sleep 60",
background: true,
},
},
new AbortController().signal,
);
const { shell_id } = JSON.parse(String(started.content)) as {
shell_id: string;
};
const aborted = new AbortController();
aborted.abort(new Error("interrupted by interrupt_agent"));
const out = await toolset.dynamicRunner.run(
{
id: "a-collect",
name: "shell_collect",
arguments: { shell_id, action: "collect", wait_ms: 60_000 },
},
aborted.signal,
);
expect(JSON.parse(String(out.content))).toMatchObject({
status: "running",
});
// A live-signal collect still sees the child running: the abort above
// released the waiter instead of killing the process. Had the abort
// killed it, this would already report completed.
const stillThere = await toolset.dynamicRunner.run(
{
id: "a-collect2",
name: "shell_collect",
arguments: { shell_id, action: "collect" },
},
new AbortController().signal,
);
expect(JSON.parse(String(stillThere.content))).toMatchObject({
status: "running",
});
// And the child is still killable: cancel settles it as completed.
const cancelled = await toolset.dynamicRunner.run(
{
id: "a-cancel",
name: "shell_collect",
arguments: { shell_id, action: "cancel" },
},
new AbortController().signal,
);
expect(JSON.parse(String(cancelled.content))).toMatchObject({
status: "cancelling",
});
const final = await toolset.dynamicRunner.run(
{
id: "a-collect3",
name: "shell_collect",
arguments: { shell_id, action: "collect", wait_ms: 5_000 },
},
new AbortController().signal,
);
expect(JSON.parse(String(final.content))).toMatchObject({
status: "completed",
});
} finally {
await toolset.dispose();
}
});

test("toolset dispose kills every live background process group", async () => {
const token = `ic_toolset_dispose_${randomUUID()}`;
const toolset = await createAgentToolset({
Expand Down
11 changes: 9 additions & 2 deletions src/agent/background-shell-tool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,10 @@ export function createShellCollectTool(
) {
return {
definition: shellCollectDefinition,
handler: async (rawArgs: Record<string, unknown>): Promise<string> => {
handler: async (
rawArgs: Record<string, unknown>,
signal?: AbortSignal,
): Promise<string> => {
const parsed = ShellCollectArgs(rawArgs);
if (parsed instanceof type.errors) {
return "Error: shell_collect requires shell_id (string) and action ('collect' | 'cancel').";
Expand All @@ -91,7 +94,11 @@ export function createShellCollectTool(
}
return JSON.stringify({ shell_id, status: "cancelling" });
}
const snapshot = await registry.collect(shell_id, parsed.wait_ms ?? 0);
const snapshot = await registry.collect(
shell_id,
parsed.wait_ms ?? 0,
signal,
);
if (snapshot.state === "running") {
return JSON.stringify({ shell_id, status: "running" });
}
Expand Down
49 changes: 49 additions & 0 deletions src/shell/background-shell.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -122,4 +122,53 @@ describe("background shell registry", () => {
const after = await registry.collect(started.id, 0);
expect(after.state).toBe("not-found");
});

test("collect with an aborted signal releases as running without killing the child", async () => {
const token = `ic_bg_abort_${randomUUID()}`;
const registry = createBackgroundShellRegistry();
const started = registry.start({
command: `bash -c 'exec -a ${token} sleep 600'`,
cwd: tmpCwd,
});
if ("error" in started) throw new Error(started.error);
try {
const aborted = new AbortController();
aborted.abort(new Error("interrupted by interrupt_agent"));
const snapshot = await registry.collect(
started.id,
60_000,
aborted.signal,
);
expect(snapshot.state).toBe("running");
const stillThere = await registry.collect(started.id, 0);
expect(stillThere.state).toBe("running");
const probe = spawnSync("pgrep", ["-f", token], { encoding: "utf8" });
expect(probe.status).toBe(0);
} finally {
registry.disposeAll("test done");
}
});

test("releaseWaiters wakes a parked collect without killing the child", async () => {
const token = `ic_bg_release_${randomUUID()}`;
const registry = createBackgroundShellRegistry();
const started = registry.start({
command: `bash -c 'exec -a ${token} sleep 600'`,
cwd: tmpCwd,
});
if ("error" in started) throw new Error(started.error);
try {
const pending = registry.collect(started.id, 60_000);
await new Promise((r) => setTimeout(r, 100));
registry.releaseWaiters();
const snapshot = await pending;
expect(snapshot.state).toBe("running");
const stillThere = await registry.collect(started.id, 0);
expect(stillThere.state).toBe("running");
const probe = spawnSync("pgrep", ["-f", token], { encoding: "utf8" });
expect(probe.status).toBe(0);
} finally {
registry.disposeAll("test done");
}
});
});
45 changes: 40 additions & 5 deletions src/shell/background-shell.ts
Original file line number Diff line number Diff line change
Expand Up @@ -64,9 +64,14 @@ export type BackgroundShellSnapshot =

export interface BackgroundShellRegistry {
start: (args: StartBackgroundShellArgs) => { id: string } | { error: string };
collect: (id: string, waitMs?: number) => Promise<BackgroundShellSnapshot>;
collect: (
id: string,
waitMs?: number,
signal?: AbortSignal,
) => Promise<BackgroundShellSnapshot>;
cancel: (id: string) => boolean;
disposeAll: (reason: string) => void;
releaseWaiters: () => void;
runningCount: () => number;
}

Expand Down Expand Up @@ -151,14 +156,33 @@ export function createBackgroundShellRegistry(
const collect = async (
id: string,
waitMs = 0,
signal?: AbortSignal,
): Promise<BackgroundShellSnapshot> => {
const done = completed.get(id);
if (done !== undefined) return { state: "completed", exit: done };
if (!running.has(id)) return { state: "not-found" };
if (waitMs > 0) {
// An already-aborted collect releases immediately as still-running:
// interrupt must not park the session on a live descendant, and must
// not kill it either — the child belongs to the still-alive session.
if (signal?.aborted === true) return { state: "running" };
await new Promise<void>((resolve) => {
exitWaiters.set(id, resolve);
setTimeout(resolve, waitMs);
const timer = setTimeout(() => {
exitWaiters.delete(id);
signal?.removeEventListener("abort", onAbort);
resolve();
}, waitMs);
const onAbort = (): void => {
clearTimeout(timer);
exitWaiters.delete(id);
resolve();
};
signal?.addEventListener("abort", onAbort, { once: true });
exitWaiters.set(id, () => {
clearTimeout(timer);
signal?.removeEventListener("abort", onAbort);
resolve();
});
}).finally(() => exitWaiters.delete(id));
const finished = completed.get(id);
if (finished !== undefined) return { state: "completed", exit: finished };
Expand All @@ -173,12 +197,22 @@ export function createBackgroundShellRegistry(
return true;
};

/**
* Wake every parked `collect` waiter as still-running without touching the
* children. Interrupt paths call this so a live descendant releases the
* session instead of wedging teardown; close paths use `disposeAll`, which
* wakes waiters and then kills the trees.
*/
const releaseWaiters = (): void => {
for (const wake of exitWaiters.values()) wake();
exitWaiters.clear();
};

const disposeAll = (reason: string): void => {
for (const child of running.values()) killProcessTree(child);
running.clear();
completed.clear();
for (const wake of exitWaiters.values()) wake();
exitWaiters.clear();
releaseWaiters();
// onExit is intentionally not fired for disposed shells: the session is
// gone, so there is no later turn to deliver to (`reason` is for callers
// that log it).
Expand All @@ -190,6 +224,7 @@ export function createBackgroundShellRegistry(
collect,
cancel,
disposeAll,
releaseWaiters,
runningCount: () => running.size,
};
}
Loading
Loading