Skip to content

Commit a17e709

Browse files
committed
add finish tool for model-driven run shutdown
1 parent 1ee1eb5 commit a17e709

7 files changed

Lines changed: 214 additions & 2 deletions

File tree

packages/agent/src/adapters/claude/claude-agent.ts

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ import {
5757
type Enrichment,
5858
type FileEnrichmentDeps,
5959
} from "../../enrichment/file-enricher";
60+
import { PostHogAPIClient } from "../../posthog-api";
6061
import {
6162
classifyPostHogExecCall,
6263
isUnclassifiedPostHogSubTool,
@@ -74,7 +75,7 @@ import { resolveGithubToken } from "../../utils/github-token";
7475
import { Logger } from "../../utils/logger";
7576
import { Pushable } from "../../utils/streams";
7677
import { BaseAcpAgent } from "../base-acp-agent";
77-
import { LOCAL_TOOLS_MCP_NAME } from "../local-tools";
78+
import { LOCAL_TOOLS_MCP_NAME, type LocalToolCtx } from "../local-tools";
7879
import { resolveSpokenNarration, resolveTaskId } from "../session-meta";
7980
import {
8081
buildBreakdown,
@@ -1834,6 +1835,30 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
18341835
}
18351836
}
18361837

1838+
// Backs the `finish` local tool: marks the task run terminal so the Temporal
1839+
// workflow tears the sandbox down. Only wired when we have both the run
1840+
// identifiers and a PostHog API config, i.e. a real cloud run.
1841+
private buildRequestFinish(
1842+
taskId: string | undefined,
1843+
taskRunId: string | undefined,
1844+
): LocalToolCtx["requestFinish"] {
1845+
const config = this.options?.posthogApiConfig;
1846+
if (!config || !taskId || !taskRunId) {
1847+
return undefined;
1848+
}
1849+
return async (status, message) => {
1850+
try {
1851+
await new PostHogAPIClient(config).updateTaskRun(taskId, taskRunId, {
1852+
status,
1853+
...(status === "failed" && message ? { error_message: message } : {}),
1854+
});
1855+
} catch (error) {
1856+
this.logger.error("finish tool failed to mark run terminal", error);
1857+
throw error;
1858+
}
1859+
};
1860+
}
1861+
18371862
private async createSession(
18381863
params: {
18391864
cwd: string;
@@ -1892,6 +1917,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
18921917
const baseBranch = meta?.baseBranch;
18931918
const environment = meta?.environment;
18941919
const spokenNarration = resolveSpokenNarration(meta);
1920+
const requestFinish = this.buildRequestFinish(taskId, meta?.taskRunId);
18951921
const buildInProcessMcpServers = (): Record<
18961922
string,
18971923
McpSdkServerConfigWithInstance
@@ -1903,8 +1929,13 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
19031929
taskId,
19041930
taskRunId: meta?.taskRunId,
19051931
baseBranch,
1932+
requestFinish,
1933+
},
1934+
{
1935+
environment,
1936+
spokenNarration,
1937+
background: meta?.mode === "background",
19061938
},
1907-
{ environment, spokenNarration },
19081939
);
19091940
return server ? { [LOCAL_TOOLS_MCP_NAME]: server } : {};
19101941
};

packages/agent/src/adapters/claude/types.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import type {
1111
SDKUserMessage,
1212
} from "@anthropic-ai/claude-agent-sdk";
1313
import type { PostHogProductId } from "../../posthog-products";
14+
import type { AgentMode } from "../../types";
1415
import type { Pushable } from "../../utils/streams";
1516
import type { BaseSession } from "../base-acp-agent";
1617
import type { ContextBreakdownBaseline } from "./context-breakdown";
@@ -177,6 +178,12 @@ export type NewSessionMeta = {
177178
taskRunId?: string;
178179
taskId?: string;
179180
environment?: "local" | "cloud";
181+
/**
182+
* Run mode. "background" means unattended (loops, durable ingest) — no human
183+
* drives the turns, so the agent may end its own run via the `finish` tool.
184+
* "interactive" runs are driven turn-by-turn and are ended by the human.
185+
*/
186+
mode?: AgentMode;
180187
disableBuiltInTools?: boolean;
181188
systemPrompt?: unknown;
182189
sessionId?: string;

packages/agent/src/adapters/local-tools/index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import type { LocalTool, LocalToolCtx, LocalToolGateMeta } from "./registry";
22
import { cloneRepoTool } from "./tools/clone-repo";
3+
import { finishTool } from "./tools/finish";
34
import { listReposTool } from "./tools/list-repos";
45
import { signedCommitTool } from "./tools/signed-commit";
56
import { signedMergeTool } from "./tools/signed-merge";
@@ -23,6 +24,7 @@ export const LOCAL_TOOLS: LocalTool[] = [
2324
listReposTool,
2425
cloneRepoTool,
2526
speakTool,
27+
finishTool,
2628
];
2729

2830
/** Tools whose gate passes for the given context — the set to actually expose. */

packages/agent/src/adapters/local-tools/registry.ts

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,16 @@ export interface LocalToolCtx {
2121
* back to origin/HEAD detection when unset.
2222
*/
2323
baseBranch?: string;
24+
/**
25+
* Marks this run terminal and lets the workflow tear the sandbox down. Set by
26+
* the adapter for cloud runs that own a sandbox; absent for local sessions.
27+
* Powers the `finish` tool so an unattended run ends the moment the agent
28+
* decides it's done, instead of idling until a timeout.
29+
*/
30+
requestFinish?: (
31+
status: "completed" | "failed",
32+
message?: string,
33+
) => Promise<void>;
2434
}
2535

2636
/** Minimal session-meta shape needed to gate tools (e.g. cloud-only). */
@@ -30,6 +40,12 @@ export interface LocalToolGateMeta {
3040
channelMode?: boolean;
3141
/** Spoken narration is on for this session: enables the speak tool. */
3242
spokenNarration?: boolean;
43+
/**
44+
* Unattended run (no human driving turns): enables the `finish` tool so the
45+
* agent can end its own run. False/undefined for interactive runs, which a
46+
* human ends.
47+
*/
48+
background?: boolean;
3349
}
3450

3551
/**
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import { describe, expect, it, vi } from "vitest";
2+
import { enabledLocalTools } from "../index";
3+
import type { LocalToolCtx, LocalToolGateMeta } from "../registry";
4+
import { FINISH_TOOL_NAME, finishSchema, finishTool } from "./finish";
5+
6+
describe("finish tool", () => {
7+
const requestFinish = async (): Promise<void> => {};
8+
9+
it.each([
10+
{
11+
name: "background cloud run with requestFinish",
12+
ctx: { cwd: "/repo", requestFinish },
13+
meta: { environment: "cloud", background: true },
14+
expected: true,
15+
},
16+
{
17+
name: "background cloud run without requestFinish",
18+
ctx: { cwd: "/repo" },
19+
meta: { environment: "cloud", background: true },
20+
expected: false,
21+
},
22+
{
23+
name: "interactive cloud run (background unset)",
24+
ctx: { cwd: "/repo", requestFinish },
25+
meta: { environment: "cloud" },
26+
expected: false,
27+
},
28+
{
29+
name: "interactive cloud run (background false)",
30+
ctx: { cwd: "/repo", requestFinish },
31+
meta: { environment: "cloud", background: false },
32+
expected: false,
33+
},
34+
{
35+
name: "background local run",
36+
ctx: { cwd: "/repo", requestFinish },
37+
meta: { environment: "local", background: true },
38+
expected: false,
39+
},
40+
{
41+
name: "no gate meta",
42+
ctx: { cwd: "/repo", requestFinish },
43+
meta: undefined,
44+
expected: false,
45+
},
46+
] as {
47+
name: string;
48+
ctx: LocalToolCtx;
49+
meta: LocalToolGateMeta | undefined;
50+
expected: boolean;
51+
}[])("is exposed only for $name → $expected", ({ ctx, meta, expected }) => {
52+
const tools = enabledLocalTools(ctx, meta);
53+
expect(tools.some((t) => t.name === FINISH_TOOL_NAME)).toBe(expected);
54+
});
55+
56+
it("stays visible without ToolSearch (alwaysLoad)", () => {
57+
expect(finishTool.alwaysLoad).toBe(true);
58+
});
59+
60+
it("defaults status to completed", () => {
61+
expect(finishSchema.status.parse(undefined)).toBe("completed");
62+
});
63+
64+
it("rejects an unknown status", () => {
65+
expect(finishSchema.status.safeParse("aborted").success).toBe(false);
66+
});
67+
68+
it("marks the run terminal via requestFinish", async () => {
69+
const spy = vi.fn(async () => {});
70+
const result = await finishTool.handler(
71+
{ cwd: "/repo", requestFinish: spy },
72+
{ status: "failed", reason: "ran out of quota" },
73+
);
74+
expect(spy).toHaveBeenCalledWith("failed", "ran out of quota");
75+
expect(result.isError).toBeUndefined();
76+
});
77+
78+
it("errors when requestFinish is unavailable", async () => {
79+
const result = await finishTool.handler(
80+
{ cwd: "/repo" },
81+
{ status: "completed" },
82+
);
83+
expect(result.isError).toBe(true);
84+
});
85+
});
Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
import { z } from "zod";
2+
import { defineLocalTool, type LocalToolResult } from "../registry";
3+
4+
export const FINISH_TOOL_NAME = "finish";
5+
6+
export const finishSchema = {
7+
status: z
8+
.enum(["completed", "failed"])
9+
.default("completed")
10+
.describe(
11+
"How the run ended. 'completed' (default) for a normal, successful " +
12+
"finish; 'failed' only if you hit something you could not get past and " +
13+
"are stopping short of the goal.",
14+
),
15+
reason: z
16+
.string()
17+
.max(500)
18+
.optional()
19+
.describe(
20+
"Short note on why you're stopping — recorded on the run. Required-in- " +
21+
"spirit for 'failed': say what blocked you so a human can pick it up.",
22+
),
23+
};
24+
25+
export const FINISH_TOOL_DESCRIPTION =
26+
"End this run and release the sandbox. This is an unattended background run: " +
27+
"nothing else will stop it promptly, so calling `finish` is how the machine " +
28+
"is reclaimed instead of sitting idle until a timeout fires. Call it once — " +
29+
"and only once — you are genuinely done: every sub-agent has returned, any CI " +
30+
"or checks you were waiting on have settled, and you've delivered whatever " +
31+
"your instructions asked for (or deliberately skipped delivery per those " +
32+
"instructions). Do NOT call it while you're still working or still waiting on " +
33+
"something to finish. After it returns, stop — the run is over.";
34+
35+
/**
36+
* Lets the model end its own background run. The handler calls back into the
37+
* adapter's `requestFinish`, which marks the task run terminal via the PostHog
38+
* API; the Temporal workflow observes the terminal status and tears the sandbox
39+
* down. Gated to cloud runs that actually own a sandbox — local sessions have
40+
* no `requestFinish` and no sandbox to reclaim, so the tool stays hidden there.
41+
*/
42+
export const finishTool = defineLocalTool({
43+
name: FINISH_TOOL_NAME,
44+
description: FINISH_TOOL_DESCRIPTION,
45+
schema: finishSchema,
46+
alwaysLoad: true,
47+
isEnabled: (ctx, meta) =>
48+
meta?.environment === "cloud" &&
49+
meta?.background === true &&
50+
ctx.requestFinish !== undefined,
51+
handler: async (ctx, args): Promise<LocalToolResult> => {
52+
if (!ctx.requestFinish) {
53+
return {
54+
content: [
55+
{ type: "text", text: "finish is not available in this session." },
56+
],
57+
isError: true,
58+
};
59+
}
60+
await ctx.requestFinish(args.status, args.reason);
61+
return {
62+
content: [
63+
{
64+
type: "text",
65+
text: `Run marked ${args.status}; shutting the sandbox down. Stop here.`,
66+
},
67+
],
68+
};
69+
},
70+
});

packages/agent/src/server/agent-server.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1416,6 +1416,7 @@ export class AgentServer {
14161416
taskRunId: payload.run_id,
14171417
taskId: payload.task_id,
14181418
environment: "cloud",
1419+
mode: this.getEffectiveMode(payload),
14191420
systemPrompt: sessionSystemPrompt,
14201421
...(this.config.model && { model: this.config.model }),
14211422
allowedDomains: this.config.allowedDomains,

0 commit comments

Comments
 (0)