From c54b580902f5b33d689b7721e86375c36593f9c8 Mon Sep 17 00:00:00 2001 From: Josh Snyder Date: Thu, 21 May 2026 05:53:57 +0000 Subject: [PATCH] feat(agent): add-on system with ztk as the first built-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces a per-task add-on system that lets users enable opt-in agent capabilities by name, configured from a `task.options.add_ons` JSON blob. The first built-in is `ztk`, a Claude PreToolUse hook that wraps Bash commands through the ztk token-reduction CLI. Architecture - New `packages/agent/src/add-ons/` module with a typed `AddOnDefinition` contract, a registry that resolves names → contributions for the active adapter, and a process-wide default registry populated with the built-in add-ons. - Unified registry, per-adapter capability honoring. Claude consumes every contribution slot (env, systemPromptAppend, preToolUse, postToolUse) via `buildSessionOptions`. Codex honors only `systemPromptAppend`; `preToolUse` / `postToolUse` have no equivalent in the upstream `codex-acp` binary, so add-ons that need command interception declare `supportedAdapters: ["claude"]` and are skipped on Codex with a logged warning. - Add-on configuration travels through `_meta.addOns` on the ACP `newSession` request, parallel to the existing `claudeCode` and `jsonSchema` channels. Resolved by the adapter's `newSession`, `loadSession`, `unstable_resumeSession`, and `unstable_forkSession`. Cloud parity - `AgentServerConfig.addOns` plus a matching `--addOns ` CLI flag on the sandbox agent server forwards the config to the cloud `newSession`, so a task carries its add-ons whether it runs locally or on the cloud sandbox. Task model - New `Task.options` JSONField on the Django Task model is the client-side home for add-on config (`task.options.add_ons`). The Django migration is not in this PR — until it lands the field is absent on every API response and add-ons can only be exercised programmatically. ztk specifics - `ztk` add-on contributes a single Claude `PreToolUse` hook that rewrites `tool_input.command` for `Bash` tool calls to `ztk run [--skip-permissions] -- `. Options: `binaryPath`, `skipPermissions`. Binary resolution checks the configured path, then PATH, then known install locations, then `~/.local/bin`. Skipped silently on Codex. Tests - 19 new tests across `registry.test.ts` (merge semantics, adapter gating, unknown-name handling, options-validation failure, prepare ordering) and `ztk.test.ts` (Bash rewriting, double-wrap protection, shell-quote escaping, --skip-permissions, Codex skip). Generated-By: PostHog Code Task-Id: de63edf0-0b52-40b4-9ffc-4758850b10d3 --- apps/code/src/shared/types.ts | 15 ++ packages/agent/package.json | 16 ++ .../agent/src/adapters/claude/claude-agent.ts | 8 + .../src/adapters/claude/session/options.ts | 83 ++++++-- packages/agent/src/adapters/claude/types.ts | 8 + .../agent/src/adapters/codex/codex-agent.ts | 94 ++++++++- .../agent/src/add-ons/default-registry.ts | 9 + packages/agent/src/add-ons/registry.test.ts | 190 ++++++++++++++++++ packages/agent/src/add-ons/registry.ts | 106 ++++++++++ packages/agent/src/add-ons/types.ts | 66 ++++++ packages/agent/src/add-ons/ztk.test.ts | 189 +++++++++++++++++ packages/agent/src/add-ons/ztk.ts | 112 +++++++++++ packages/agent/src/server/agent-server.ts | 1 + packages/agent/src/server/bin.ts | 16 +- packages/agent/src/server/schemas.ts | 10 + packages/agent/src/server/types.ts | 8 + packages/agent/src/types.ts | 16 ++ packages/agent/tsup.config.ts | 4 + 18 files changed, 928 insertions(+), 23 deletions(-) create mode 100644 packages/agent/src/add-ons/default-registry.ts create mode 100644 packages/agent/src/add-ons/registry.test.ts create mode 100644 packages/agent/src/add-ons/registry.ts create mode 100644 packages/agent/src/add-ons/types.ts create mode 100644 packages/agent/src/add-ons/ztk.test.ts create mode 100644 packages/agent/src/add-ons/ztk.ts diff --git a/apps/code/src/shared/types.ts b/apps/code/src/shared/types.ts index f062f670c..89b4f79ab 100644 --- a/apps/code/src/shared/types.ts +++ b/apps/code/src/shared/types.ts @@ -34,6 +34,20 @@ interface UserBasic { is_email_verified?: boolean | null; } +/** + * Per-task configuration stored on the Django `Task.options` JSONField. + * Open-ended so we can add new keys without OpenAPI churn; today only + * `add_ons` is consumed by the agent runtime. + * + * Server-side requires `options = models.JSONField(default=dict, blank=True)` + * on the Task model plus a matching serializer entry. Until that migration + * lands, this field will be absent on every Task returned by the API. + */ +export interface TaskOptions { + /** Keys are add-on names registered with `@posthog/agent`'s AddOnRegistry. */ + add_ons?: Record>; +} + export interface Task { id: string; task_number: number | null; @@ -51,6 +65,7 @@ export interface Task { json_schema?: Record | null; signal_report?: string | null; internal?: boolean; + options?: TaskOptions; latest_run?: TaskRun; } diff --git a/packages/agent/package.json b/packages/agent/package.json index d836afe26..4a5a09ab6 100644 --- a/packages/agent/package.json +++ b/packages/agent/package.json @@ -28,6 +28,22 @@ "types": "./dist/types.d.ts", "import": "./dist/types.js" }, + "./add-ons/types": { + "types": "./dist/add-ons/types.d.ts", + "import": "./dist/add-ons/types.js" + }, + "./add-ons/registry": { + "types": "./dist/add-ons/registry.d.ts", + "import": "./dist/add-ons/registry.js" + }, + "./add-ons/default-registry": { + "types": "./dist/add-ons/default-registry.d.ts", + "import": "./dist/add-ons/default-registry.js" + }, + "./add-ons/ztk": { + "types": "./dist/add-ons/ztk.d.ts", + "import": "./dist/add-ons/ztk.js" + }, "./adapters/claude/questions/utils": { "types": "./dist/adapters/claude/questions/utils.d.ts", "import": "./dist/adapters/claude/questions/utils.js" diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 2b516b72d..228d84396 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -51,6 +51,7 @@ import { POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "../../acp-extensions"; +import { defaultAddOnRegistry } from "../../add-ons/default-registry"; import { createEnrichment, type Enrichment, @@ -1140,6 +1141,12 @@ export class ClaudeAcpAgent extends BaseAcpAgent { ? (meta.permissionMode as CodeExecutionMode) : "default"; + const addOnContribution = await defaultAddOnRegistry.collect(meta?.addOns, { + cwd, + adapter: "claude", + logger: this.logger, + }); + const options = buildSessionOptions({ cwd, mcpServers, @@ -1147,6 +1154,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { canUseTool: this.createCanUseTool(sessionId, meta?.allowedDomains), logger: this.logger, systemPrompt, + addOnContribution, userProvidedOptions: meta?.claudeCode?.options, sessionId, isResume, diff --git a/packages/agent/src/adapters/claude/session/options.ts b/packages/agent/src/adapters/claude/session/options.ts index 1ad1bf255..241703df1 100644 --- a/packages/agent/src/adapters/claude/session/options.ts +++ b/packages/agent/src/adapters/claude/session/options.ts @@ -4,12 +4,14 @@ import * as os from "node:os"; import * as path from "node:path"; import type { CanUseTool, + HookCallback, McpServerConfig, Options, OutputFormat, SpawnedProcess, SpawnOptions, } from "@anthropic-ai/claude-agent-sdk"; +import type { AddOnContribution } from "../../../add-ons/types"; import type { FileEnrichmentDeps } from "../../../enrichment/file-enricher"; import { IS_ROOT } from "../../../utils/common"; import type { Logger } from "../../../utils/logger"; @@ -55,6 +57,12 @@ export interface BuildOptionsParams { effort?: EffortLevel; enrichmentDeps?: FileEnrichmentDeps; enrichedReadCache?: EnrichedReadCache; + /** + * Pre-resolved contribution from the {@link AddOnRegistry}. Merged into + * `hooks`, `env`, and `systemPrompt` so add-ons can prepend Bash rewriters, + * inject env vars, and append system-prompt text. + */ + addOnContribution?: AddOnContribution; } export function buildSystemPrompt( @@ -89,6 +97,26 @@ export function buildSystemPrompt( return defaultPrompt; } +function appendToSystemPrompt( + systemPrompt: Options["systemPrompt"], + extra: string | undefined, +): Options["systemPrompt"] { + if (!extra) return systemPrompt; + if (typeof systemPrompt === "string") return systemPrompt + extra; + if ( + typeof systemPrompt === "object" && + systemPrompt !== null && + "type" in systemPrompt && + systemPrompt.type === "preset" + ) { + return { + ...systemPrompt, + append: (systemPrompt.append ?? "") + extra, + }; + } + return systemPrompt; +} + function buildMcpServers( userServers: Record | undefined, acpServers: Record, @@ -101,7 +129,9 @@ function buildMcpServers( }; } -function buildEnvironment(): Record { +function buildEnvironment( + addOnEnv?: Record, +): Record { const bedrockFallbackHeader = "x-posthog-use-bedrock-fallback: true"; const existingCustomHeaders = process.env.ANTHROPIC_CUSTOM_HEADERS; const customHeaders = existingCustomHeaders @@ -118,6 +148,9 @@ function buildEnvironment(): Record { CLAUDE_CODE_EMIT_SESSION_STATE_EVENTS: "1", // Route to AWS Bedrock as a fallback when Anthropic returns 5xx ANTHROPIC_CUSTOM_HEADERS: customHeaders, + // Add-on contributions win over our defaults so they can override e.g. + // ANTHROPIC_CUSTOM_HEADERS for proxy add-ons. + ...(addOnEnv ?? {}), }; } @@ -129,6 +162,8 @@ function buildHooks( enrichmentDeps: FileEnrichmentDeps | undefined, enrichedReadCache: EnrichedReadCache | undefined, registeredAgents: ReadonlySet, + addOnPreToolUse: HookCallback[] | undefined, + addOnPostToolUse: HookCallback[] | undefined, ): Options["hooks"] { const postToolUseHooks = [createPostToolUseHook({ onModeChange })]; if (enrichmentDeps && enrichedReadCache) { @@ -137,21 +172,30 @@ function buildHooks( ); } + // Add-on PreToolUse hooks run BEFORE the built-in permission gate so they + // can rewrite tool input (e.g. ztk wrapping a Bash command) and have those + // changes reflected in the permission-check pass. + const preToolUseGroups: NonNullable["PreToolUse"] = [ + ...(userHooks?.PreToolUse || []), + ...(addOnPreToolUse?.length ? [{ hooks: addOnPreToolUse }] : []), + { + hooks: [ + createPreToolUseHook(settingsManager, logger), + createSubagentRewriteHook(logger, registeredAgents), + ], + }, + ]; + + const postToolUseGroups: NonNullable["PostToolUse"] = [ + ...(userHooks?.PostToolUse || []), + { hooks: postToolUseHooks }, + ...(addOnPostToolUse?.length ? [{ hooks: addOnPostToolUse }] : []), + ]; + return { ...userHooks, - PostToolUse: [ - ...(userHooks?.PostToolUse || []), - { hooks: postToolUseHooks }, - ], - PreToolUse: [ - ...(userHooks?.PreToolUse || []), - { - hooks: [ - createPreToolUseHook(settingsManager, logger), - createSubagentRewriteHook(logger, registeredAgents), - ], - }, - ], + PostToolUse: postToolUseGroups, + PreToolUse: preToolUseGroups, }; } @@ -320,10 +364,15 @@ export function buildSessionOptions(params: BuildOptionsParams): Options { const agents = buildAgents(params.userProvidedOptions?.agents); const registeredAgentNames = new Set(Object.keys(agents)); + const resolvedSystemPrompt = appendToSystemPrompt( + params.systemPrompt ?? buildSystemPrompt(), + params.addOnContribution?.systemPromptAppend, + ); + const options: Options = { ...params.userProvidedOptions, betas: ["context-1m-2025-08-07"], - systemPrompt: params.systemPrompt ?? buildSystemPrompt(), + systemPrompt: resolvedSystemPrompt, settingSources: ["user", "project", "local"], stderr: (err) => params.logger.error(err), cwd: params.cwd, @@ -343,7 +392,7 @@ export function buildSessionOptions(params: BuildOptionsParams): Options { params.mcpServers, loadUserClaudeJsonMcpServers(params.cwd, params.logger), ), - env: buildEnvironment(), + env: buildEnvironment(params.addOnContribution?.env), hooks: buildHooks( params.userProvidedOptions?.hooks, params.onModeChange, @@ -352,6 +401,8 @@ export function buildSessionOptions(params: BuildOptionsParams): Options { params.enrichmentDeps, params.enrichedReadCache, registeredAgentNames, + params.addOnContribution?.preToolUse, + params.addOnContribution?.postToolUse, ), outputFormat: params.outputFormat, abortController: getAbortController( diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts index 1898a2360..9c962629a 100644 --- a/packages/agent/src/adapters/claude/types.ts +++ b/packages/agent/src/adapters/claude/types.ts @@ -8,6 +8,7 @@ import type { Query, SDKUserMessage, } from "@anthropic-ai/claude-agent-sdk"; +import type { AddOnConfig } from "../../add-ons/types"; import type { Pushable } from "../../utils/streams"; import type { BaseSession } from "../base-acp-agent"; import type { McpToolApprovals } from "./mcp/tool-metadata"; @@ -135,4 +136,11 @@ export type NewSessionMeta = { options?: Options; emitRawSDKMessages?: boolean | SDKMessageFilter[]; }; + /** + * Add-on configuration sourced from `task.options.add_ons`. Keys are + * add-on names registered in the default `AddOnRegistry`; values are + * opaque options blobs validated per-add-on at session start. Unknown + * or unsupported names are skipped with a warning. + */ + addOns?: AddOnConfig; }; diff --git a/packages/agent/src/adapters/codex/codex-agent.ts b/packages/agent/src/adapters/codex/codex-agent.ts index 4313d29c0..c6b7f0558 100644 --- a/packages/agent/src/adapters/codex/codex-agent.ts +++ b/packages/agent/src/adapters/codex/codex-agent.ts @@ -44,6 +44,8 @@ import { POSTHOG_METHODS, POSTHOG_NOTIFICATIONS, } from "../../acp-extensions"; +import { defaultAddOnRegistry } from "../../add-ons/default-registry"; +import type { AddOnConfig, AddOnContribution } from "../../add-ons/types"; import { createEnrichment, type Enrichment, @@ -100,6 +102,13 @@ interface NewSessionMeta { disableBuiltInTools?: boolean; allowedDomains?: string[]; jsonSchema?: Record | null; + /** + * Add-on configuration sourced from `task.options.add_ons`. Only the + * `systemPromptAppend` slot is honored on Codex today; add-ons that + * require `preToolUse`/`postToolUse` hooks declare + * `supportedAdapters: ["claude"]` and are skipped here. + */ + addOns?: AddOnConfig; } export interface CodexAcpAgentOptions { @@ -338,7 +347,12 @@ export class CodexAcpAgent extends BaseAcpAgent { const meta = params._meta as NewSessionMeta | undefined; const requestedPermissionMode = toCodexPermissionMode(meta?.permissionMode); - const injectedParams = this.applyStructuredOutput(params, meta); + const addOnContribution = await this.collectAddOnContribution( + meta?.addOns, + params.cwd, + ); + const withAddOns = this.applyAddOnContribution(params, addOnContribution); + const injectedParams = this.applyStructuredOutput(withAddOns, meta); const response = await this.codexConnection.newSession(injectedParams); response.configOptions = normalizeCodexConfigOptions( response.configOptions, @@ -380,7 +394,12 @@ export class CodexAcpAgent extends BaseAcpAgent { async loadSession(params: LoadSessionRequest): Promise { const meta = params._meta as NewSessionMeta | undefined; - const injectedParams = this.applyStructuredOutput(params, meta); + const addOnContribution = await this.collectAddOnContribution( + meta?.addOns, + params.cwd, + ); + const withAddOns = this.applyAddOnContribution(params, addOnContribution); + const injectedParams = this.applyStructuredOutput(withAddOns, meta); const response = await this.codexConnection.loadSession(injectedParams); response.configOptions = normalizeCodexConfigOptions( response.configOptions, @@ -418,15 +437,20 @@ export class CodexAcpAgent extends BaseAcpAgent { params: ResumeSessionRequest, ): Promise { const meta = params._meta as NewSessionMeta | undefined; - const injectedParams = this.applyStructuredOutput( + const addOnContribution = await this.collectAddOnContribution( + meta?.addOns, + params.cwd, + ); + const withAddOns = this.applyAddOnContribution( { sessionId: params.sessionId, cwd: params.cwd, mcpServers: params.mcpServers ?? [], _meta: params._meta, }, - meta, + addOnContribution, ); + const injectedParams = this.applyStructuredOutput(withAddOns, meta); // codex-acp doesn't support resume natively, use loadSession instead const loadResponse = await this.codexConnection.loadSession(injectedParams); @@ -465,14 +489,19 @@ export class CodexAcpAgent extends BaseAcpAgent { params: ForkSessionRequest, ): Promise { const meta = params._meta as NewSessionMeta | undefined; - const injectedParams = this.applyStructuredOutput( + const addOnContribution = await this.collectAddOnContribution( + meta?.addOns, + params.cwd, + ); + const withAddOns = this.applyAddOnContribution( { cwd: params.cwd, mcpServers: params.mcpServers ?? [], _meta: params._meta, }, - meta, + addOnContribution, ); + const injectedParams = this.applyStructuredOutput(withAddOns, meta); // Create a new session via codex-acp (fork isn't natively supported) const newResponse = await this.codexConnection.newSession(injectedParams); @@ -499,6 +528,59 @@ export class CodexAcpAgent extends BaseAcpAgent { return newResponse; } + /** + * Resolve the add-on contribution for this session. Add-ons that declare + * `supportedAdapters: ["claude"]` (like ztk) are silently skipped here — + * `preToolUse`/`postToolUse` slots they would emit have no Codex equivalent. + */ + private collectAddOnContribution( + addOns: AddOnConfig | undefined, + cwd: string, + ): Promise { + return defaultAddOnRegistry.collect(addOns, { + cwd, + adapter: "codex", + logger: this.logger, + }); + } + + /** + * Apply the supported slots of an add-on contribution to an outbound ACP + * session request: `systemPromptAppend` is appended to `_meta.systemPrompt`. + * `env` is ignored because the `codex-acp` subprocess has already been + * spawned with its environment fixed; add-ons that need env vars in Codex + * must declare themselves Claude-only. + */ + private applyAddOnContribution( + request: T, + contribution: AddOnContribution, + ): T { + if (!contribution.systemPromptAppend && !contribution.env) { + return request; + } + if (contribution.env) { + this.logger.warn( + "Add-on contributed env vars but Codex env is fixed at spawn — ignoring", + { keys: Object.keys(contribution.env) }, + ); + } + if (!contribution.systemPromptAppend) { + return request; + } + const existingMeta = (request._meta ?? {}) as Record; + const existingSystemPrompt = + typeof existingMeta.systemPrompt === "string" + ? existingMeta.systemPrompt + : ""; + return { + ...request, + _meta: { + ...existingMeta, + systemPrompt: existingSystemPrompt + contribution.systemPromptAppend, + }, + }; + } + /** * When the caller wires up `onStructuredOutput` and provides a JSON schema * via `_meta.jsonSchema`, inject the stdio MCP server that exposes diff --git a/packages/agent/src/add-ons/default-registry.ts b/packages/agent/src/add-ons/default-registry.ts new file mode 100644 index 000000000..403379c04 --- /dev/null +++ b/packages/agent/src/add-ons/default-registry.ts @@ -0,0 +1,9 @@ +import { AddOnRegistry } from "./registry"; +import { ztkAddOn } from "./ztk"; + +/** + * Process-wide default registry. Built-in add-ons are registered here so + * adapters can resolve them without ceremony. + */ +export const defaultAddOnRegistry = new AddOnRegistry(); +defaultAddOnRegistry.register(ztkAddOn); diff --git a/packages/agent/src/add-ons/registry.test.ts b/packages/agent/src/add-ons/registry.test.ts new file mode 100644 index 000000000..3bc51107b --- /dev/null +++ b/packages/agent/src/add-ons/registry.test.ts @@ -0,0 +1,190 @@ +import { describe, expect, it, vi } from "vitest"; +import { Logger } from "../utils/logger"; +import { AddOnRegistry } from "./registry"; +import type { AddOnContext, AddOnContribution, AddOnDefinition } from "./types"; + +function makeCtx(adapter: "claude" | "codex" = "claude"): AddOnContext { + return { + cwd: "/tmp/fake-cwd", + adapter, + logger: new Logger(), + }; +} + +function makeDefinition( + overrides: Partial> = {}, +): AddOnDefinition<{ value?: string }> { + return { + name: "test", + parseOptions: (raw) => raw as { value?: string }, + contribute: () => ({}), + ...overrides, + }; +} + +describe("AddOnRegistry", () => { + it("returns an empty contribution when no config is provided", async () => { + const registry = new AddOnRegistry(); + const result = await registry.collect(undefined, makeCtx()); + expect(result).toEqual({}); + }); + + it("merges env vars from multiple add-ons (later wins on conflict)", async () => { + const registry = new AddOnRegistry(); + registry.register( + makeDefinition({ + name: "a", + contribute: () => ({ env: { SHARED: "from-a", A_ONLY: "1" } }), + }), + ); + registry.register( + makeDefinition({ + name: "b", + contribute: () => ({ env: { SHARED: "from-b", B_ONLY: "1" } }), + }), + ); + + const result = await registry.collect({ a: {}, b: {} }, makeCtx()); + expect(result.env).toEqual({ + SHARED: "from-b", + A_ONLY: "1", + B_ONLY: "1", + }); + }); + + it("concatenates systemPromptAppend across add-ons in iteration order", async () => { + const registry = new AddOnRegistry(); + registry.register( + makeDefinition({ + name: "a", + contribute: () => ({ systemPromptAppend: "AA" }), + }), + ); + registry.register( + makeDefinition({ + name: "b", + contribute: () => ({ systemPromptAppend: "BB" }), + }), + ); + + const result = await registry.collect({ a: {}, b: {} }, makeCtx()); + expect(result.systemPromptAppend).toBe("AABB"); + }); + + it("aggregates preToolUse and postToolUse hooks", async () => { + const registry = new AddOnRegistry(); + const hookA = vi.fn(); + const hookB = vi.fn(); + registry.register( + makeDefinition({ + name: "a", + contribute: (): AddOnContribution => ({ + preToolUse: [hookA], + postToolUse: [hookB], + }), + }), + ); + + const result = await registry.collect({ a: {} }, makeCtx()); + expect(result.preToolUse).toEqual([hookA]); + expect(result.postToolUse).toEqual([hookB]); + }); + + it("skips add-ons not supported on the current adapter", async () => { + const registry = new AddOnRegistry(); + registry.register( + makeDefinition({ + name: "claude-only", + supportedAdapters: ["claude"], + contribute: () => ({ systemPromptAppend: "should-not-appear" }), + }), + ); + + const result = await registry.collect( + { "claude-only": {} }, + makeCtx("codex"), + ); + expect(result.systemPromptAppend).toBeUndefined(); + }); + + it("skips unknown add-on names with a warning instead of throwing", async () => { + const registry = new AddOnRegistry(); + const ctx = makeCtx(); + const warnSpy = vi.spyOn(ctx.logger, "warn").mockImplementation(() => {}); + + await expect( + registry.collect({ "does-not-exist": {} }, ctx), + ).resolves.toEqual({}); + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining("does-not-exist"), + expect.objectContaining({ addOn: "does-not-exist" }), + ); + }); + + it("skips add-ons whose options fail parsing instead of throwing", async () => { + const registry = new AddOnRegistry(); + registry.register( + makeDefinition({ + name: "strict", + parseOptions: () => { + throw new Error("bad options"); + }, + contribute: () => ({ env: { SHOULD_NOT: "1" } }), + }), + ); + const ctx = makeCtx(); + const warnSpy = vi.spyOn(ctx.logger, "warn").mockImplementation(() => {}); + + const result = await registry.collect({ strict: { x: 1 } }, ctx); + expect(result.env).toBeUndefined(); + expect(warnSpy).toHaveBeenCalledWith( + expect.stringContaining("strict"), + expect.objectContaining({ addOn: "strict" }), + ); + }); + + it("awaits prepare() before contribute()", async () => { + const registry = new AddOnRegistry(); + const order: string[] = []; + registry.register( + makeDefinition({ + name: "ordered", + prepare: async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + order.push("prepare"); + }, + contribute: () => { + order.push("contribute"); + return {}; + }, + }), + ); + + await registry.collect({ ordered: {} }, makeCtx()); + expect(order).toEqual(["prepare", "contribute"]); + }); + + it("propagates prepare() failures so missing prerequisites surface early", async () => { + const registry = new AddOnRegistry(); + registry.register( + makeDefinition({ + name: "needs-binary", + prepare: () => { + throw new Error("binary missing"); + }, + }), + ); + + await expect( + registry.collect({ "needs-binary": {} }, makeCtx()), + ).rejects.toThrow("binary missing"); + }); + + it("rejects duplicate registrations", () => { + const registry = new AddOnRegistry(); + registry.register(makeDefinition({ name: "dup" })); + expect(() => registry.register(makeDefinition({ name: "dup" }))).toThrow( + /already registered/, + ); + }); +}); diff --git a/packages/agent/src/add-ons/registry.ts b/packages/agent/src/add-ons/registry.ts new file mode 100644 index 000000000..308c659d7 --- /dev/null +++ b/packages/agent/src/add-ons/registry.ts @@ -0,0 +1,106 @@ +import type { + AddOnConfig, + AddOnContext, + AddOnContribution, + AddOnDefinition, +} from "./types"; + +// biome-ignore lint/suspicious/noExplicitAny: registry erases per-definition option types +type AnyAddOnDefinition = AddOnDefinition; + +export class AddOnRegistry { + private definitions = new Map(); + + register(definition: AddOnDefinition): void { + if (this.definitions.has(definition.name)) { + throw new Error( + `AddOn "${definition.name}" is already registered. Names must be unique.`, + ); + } + this.definitions.set(definition.name, definition); + } + + get(name: string): AnyAddOnDefinition | undefined { + return this.definitions.get(name); + } + + list(): AnyAddOnDefinition[] { + return [...this.definitions.values()]; + } + + /** + * Resolve every enabled add-on for the current adapter and merge their + * contributions into a single object. Unknown names, unsupported adapters, + * and option-parse failures are logged and skipped — never throw out of + * `collect()`, since one misconfigured add-on should not break the session. + * `prepare()` failures DO throw, because they signal a missing prerequisite + * the user must fix. + */ + async collect( + config: AddOnConfig | undefined, + ctx: AddOnContext, + ): Promise { + const merged: AddOnContribution = {}; + if (!config) return merged; + + for (const [name, rawOptions] of Object.entries(config)) { + const definition = this.definitions.get(name); + if (!definition) { + ctx.logger.warn(`Unknown add-on "${name}" — skipping`, { + addOn: name, + }); + continue; + } + + if ( + definition.supportedAdapters && + !definition.supportedAdapters.includes(ctx.adapter) + ) { + ctx.logger.info( + `Add-on "${name}" is not supported on adapter "${ctx.adapter}" — skipping`, + { addOn: name, adapter: ctx.adapter }, + ); + continue; + } + + let options: unknown; + try { + options = definition.parseOptions(rawOptions ?? {}); + } catch (error) { + ctx.logger.warn(`Add-on "${name}" rejected options — skipping`, { + addOn: name, + error, + }); + continue; + } + + if (definition.prepare) { + await definition.prepare(ctx, options); + } + + const contribution = await definition.contribute(ctx, options); + mergeContribution(merged, contribution); + } + + return merged; + } +} + +function mergeContribution( + target: AddOnContribution, + source: AddOnContribution, +): void { + if (source.env) { + target.env = { ...(target.env ?? {}), ...source.env }; + } + if (source.systemPromptAppend) { + target.systemPromptAppend = + (target.systemPromptAppend ?? "") + source.systemPromptAppend; + } + if (source.preToolUse?.length) { + target.preToolUse = [...(target.preToolUse ?? []), ...source.preToolUse]; + } + if (source.postToolUse?.length) { + target.postToolUse = [...(target.postToolUse ?? []), ...source.postToolUse]; + } +} diff --git a/packages/agent/src/add-ons/types.ts b/packages/agent/src/add-ons/types.ts new file mode 100644 index 000000000..d76da5459 --- /dev/null +++ b/packages/agent/src/add-ons/types.ts @@ -0,0 +1,66 @@ +import type { HookCallback } from "@anthropic-ai/claude-agent-sdk"; +import type { Logger } from "../utils/logger"; + +/** + * Shape of `task.options.add_ons` as it travels from the Django Task model + * through `_meta.addOns` on a `newSession` ACP request. Keys are add-on names + * registered with the {@link AddOnRegistry}; values are opaque option blobs + * that each add-on validates with its own `parseOptions` implementation. + */ +export type AddOnConfig = Record>; + +export type AddOnAdapter = "claude" | "codex"; + +/** + * The contribution slots an add-on may emit. Adapters honor only the slots + * they support — Claude honors everything, Codex skips `preToolUse` and + * `postToolUse` because the upstream `codex-acp` binary exposes no pre- or + * post-tool interception. Add-ons that require interception declare + * `supportedAdapters: ["claude"]` to opt out on Codex. + */ +export type AddOnCapability = + | "env" + | "systemPrompt" + | "preToolUse" + | "postToolUse"; + +export interface AddOnContext { + cwd: string; + adapter: AddOnAdapter; + logger: Logger; +} + +export interface AddOnContribution { + env?: Record; + systemPromptAppend?: string; + preToolUse?: HookCallback[]; + postToolUse?: HookCallback[]; +} + +export interface AddOnDefinition> { + /** Unique name. Matches the key under `task.options.add_ons`. */ + name: string; + /** Capability slots this add-on actually uses. Informational. */ + requires?: AddOnCapability[]; + /** + * Adapters this add-on supports. Omit to support every adapter; specify + * a subset (e.g. `["claude"]`) to be silently skipped on unsupported ones. + */ + supportedAdapters?: AddOnAdapter[]; + /** + * Validate and shape the raw options blob. Throw to signal invalid input — + * the registry will skip the add-on and log a warning rather than abort the + * whole session. + */ + parseOptions(rawOptions: unknown): TOptions; + /** + * Idempotent setup that must complete before the session starts. Use for + * resolving binaries on disk, downloading assets, etc. Throw to fail loudly. + */ + prepare?(ctx: AddOnContext, options: TOptions): Promise | void; + /** Produce the session-level contribution. */ + contribute( + ctx: AddOnContext, + options: TOptions, + ): Promise | AddOnContribution; +} diff --git a/packages/agent/src/add-ons/ztk.test.ts b/packages/agent/src/add-ons/ztk.test.ts new file mode 100644 index 000000000..cc97998bd --- /dev/null +++ b/packages/agent/src/add-ons/ztk.test.ts @@ -0,0 +1,189 @@ +import { mkdirSync, mkdtempSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { + HookCallback, + HookInput, + HookJSONOutput, +} from "@anthropic-ai/claude-agent-sdk"; +import { beforeEach, describe, expect, it } from "vitest"; +import { Logger } from "../utils/logger"; +import { AddOnRegistry } from "./registry"; +import type { AddOnContext, AddOnContribution } from "./types"; +import { ztkAddOn } from "./ztk"; + +function firstHook(contribution: AddOnContribution): HookCallback { + const hook = contribution.preToolUse?.[0]; + if (!hook) { + throw new Error("expected a PreToolUse hook on contribution"); + } + return hook; +} + +function makeFakeZtkBinary(): string { + const dir = mkdtempSync(join(tmpdir(), "ztk-test-")); + const binaryPath = join(dir, "ztk"); + writeFileSync(binaryPath, "#!/bin/sh\nexit 0\n", { mode: 0o755 }); + return binaryPath; +} + +function makeCtx(overrides: Partial = {}): AddOnContext { + return { + cwd: "/tmp/fake", + adapter: "claude", + logger: new Logger(), + ...overrides, + }; +} + +async function runHook( + hook: HookCallback, + partial: { tool_name: string; tool_input?: Record }, +): Promise { + const input = { + hook_event_name: "PreToolUse", + session_id: "s", + transcript_path: "/tmp/t", + cwd: "/tmp", + tool_name: partial.tool_name, + tool_input: partial.tool_input ?? {}, + } as HookInput; + return hook(input, "tool-use-id", { signal: new AbortController().signal }); +} + +describe("ztk add-on", () => { + let binaryPath: string; + + beforeEach(() => { + binaryPath = makeFakeZtkBinary(); + }); + + it("rejects unknown options keys", () => { + expect(() => ztkAddOn.parseOptions({ binaryPath: 42 })).toThrow(); + }); + + it("throws from prepare() when the configured binary is missing", () => { + expect(() => + ztkAddOn.prepare?.( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath: "/does/not/exist/ztk" }), + ), + ).toThrow(/binary not found/); + }); + + it("contributes a PreToolUse hook that rewrites Bash commands", async () => { + const contribution = await ztkAddOn.contribute( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath }), + ); + expect(contribution.preToolUse).toHaveLength(1); + + const hook = firstHook(contribution); + const result = (await runHook(hook, { + tool_name: "Bash", + tool_input: { command: "ls -la" }, + })) as { + hookSpecificOutput?: { + hookEventName: string; + updatedInput: { command: string }; + }; + }; + + expect(result.hookSpecificOutput?.hookEventName).toBe("PreToolUse"); + expect(result.hookSpecificOutput?.updatedInput.command).toBe( + `'${binaryPath}' run -- ls -la`, + ); + }); + + it("passes through non-Bash tool calls unchanged", async () => { + const contribution = await ztkAddOn.contribute( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath }), + ); + const hook = firstHook(contribution); + const result = (await runHook(hook, { + tool_name: "Read", + tool_input: { file_path: "/etc/hosts" }, + })) as { continue: boolean; hookSpecificOutput?: unknown }; + + expect(result.continue).toBe(true); + expect(result.hookSpecificOutput).toBeUndefined(); + }); + + it("appends --skip-permissions when configured", async () => { + const contribution = await ztkAddOn.contribute( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath, skipPermissions: true }), + ); + const hook = firstHook(contribution); + const result = (await runHook(hook, { + tool_name: "Bash", + tool_input: { command: "echo hi" }, + })) as { + hookSpecificOutput?: { updatedInput: { command: string } }; + }; + expect(result.hookSpecificOutput?.updatedInput.command).toBe( + `'${binaryPath}' run --skip-permissions -- echo hi`, + ); + }); + + it("declares itself unsupported on Codex so the registry skips it", async () => { + const registry = new AddOnRegistry(); + registry.register(ztkAddOn); + + const result = await registry.collect( + { ztk: { binaryPath } }, + makeCtx({ adapter: "codex" }), + ); + expect(result.preToolUse).toBeUndefined(); + }); + + it("end-to-end: resolves through the registry on Claude", async () => { + const registry = new AddOnRegistry(); + registry.register(ztkAddOn); + + const result = await registry.collect( + { ztk: { binaryPath } }, + makeCtx({ adapter: "claude" }), + ); + expect(result.preToolUse).toHaveLength(1); + }); + + it("does not double-wrap an already-wrapped command", async () => { + const contribution = await ztkAddOn.contribute( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath }), + ); + const hook = firstHook(contribution); + const wrapped = `'${binaryPath}' run -- echo hi`; + const result = (await runHook(hook, { + tool_name: "Bash", + tool_input: { command: wrapped }, + })) as { continue: boolean; hookSpecificOutput?: unknown }; + + expect(result.hookSpecificOutput).toBeUndefined(); + expect(result.continue).toBe(true); + }); + + it("escapes single quotes inside the binary path", async () => { + // Make a binary at a path that already contains a quote — verifies escaping. + const dir = mkdtempSync(join(tmpdir(), "ztk-quote-")); + const quoted = join(dir, "weird'name"); + mkdirSync(quoted, { recursive: true }); + const bp = join(quoted, "ztk"); + writeFileSync(bp, "#!/bin/sh\nexit 0\n", { mode: 0o755 }); + + const contribution = await ztkAddOn.contribute( + makeCtx(), + ztkAddOn.parseOptions({ binaryPath: bp }), + ); + const hook = firstHook(contribution); + const result = (await runHook(hook, { + tool_name: "Bash", + tool_input: { command: "ls" }, + })) as { + hookSpecificOutput?: { updatedInput: { command: string } }; + }; + expect(result.hookSpecificOutput?.updatedInput.command).toContain("'\\''"); + }); +}); diff --git a/packages/agent/src/add-ons/ztk.ts b/packages/agent/src/add-ons/ztk.ts new file mode 100644 index 000000000..76cddb3cf --- /dev/null +++ b/packages/agent/src/add-ons/ztk.ts @@ -0,0 +1,112 @@ +import { existsSync } from "node:fs"; +import { join } from "node:path"; +import type { HookCallback } from "@anthropic-ai/claude-agent-sdk"; +import { z } from "zod"; +import type { AddOnContext, AddOnDefinition } from "./types"; + +const ztkOptionsSchema = z.object({ + /** + * Absolute path to the ztk binary. When omitted, the add-on looks on + * `$PATH` and a small set of known install locations. + */ + binaryPath: z.string().optional(), + /** + * Forwarded to `ztk run --skip-permissions`. Tells ztk to bypass its own + * approval prompts because Claude's permission model already gates Bash. + */ + skipPermissions: z.boolean().optional(), +}); + +export type ZtkOptions = z.infer; + +const KNOWN_INSTALL_DIRS = ["/opt/homebrew/bin", "/usr/local/bin", "/usr/bin"]; + +function resolveZtkBinary(options: ZtkOptions): string { + if (options.binaryPath) { + if (!existsSync(options.binaryPath)) { + throw new Error( + `ztk add-on: binary not found at configured path "${options.binaryPath}"`, + ); + } + return options.binaryPath; + } + + const pathDirs = (process.env.PATH ?? "").split(":").filter(Boolean); + const homeBin = process.env.HOME + ? join(process.env.HOME, ".local/bin") + : null; + const candidates = [...pathDirs, ...KNOWN_INSTALL_DIRS]; + if (homeBin) candidates.push(homeBin); + + for (const dir of candidates) { + const candidate = join(dir, "ztk"); + if (existsSync(candidate)) return candidate; + } + + throw new Error( + 'ztk add-on: binary "ztk" not found on PATH. Install it from ' + + "https://github.com/codejunkie99/ztk or set add-on option `binaryPath` to an absolute path.", + ); +} + +function shellEscape(value: string): string { + return `'${value.replace(/'/g, "'\\''")}'`; +} + +function makeZtkBashHook( + binaryPath: string, + options: ZtkOptions, +): HookCallback { + const flag = options.skipPermissions ? " --skip-permissions" : ""; + const escapedBinary = shellEscape(binaryPath); + return async (input) => { + if (input.hook_event_name !== "PreToolUse") return { continue: true }; + if (input.tool_name !== "Bash") return { continue: true }; + + const toolInput = (input.tool_input ?? {}) as { command?: unknown }; + const command = toolInput.command; + if (typeof command !== "string" || command.length === 0) { + return { continue: true }; + } + + // Already wrapped — don't double-wrap if the hook runs twice for some reason. + if (command.startsWith(`${escapedBinary} run`)) { + return { continue: true }; + } + + const rewritten = `${escapedBinary} run${flag} -- ${command}`; + return { + continue: true, + hookSpecificOutput: { + hookEventName: "PreToolUse", + updatedInput: { + ...(input.tool_input as Record), + command: rewritten, + }, + }, + }; + }; +} + +export const ztkAddOn: AddOnDefinition = { + name: "ztk", + requires: ["preToolUse"], + // ztk needs to mutate Bash tool input before execution. Only Claude exposes + // a PreToolUse hook with `updatedInput`; the codex-acp binary has no + // equivalent interception point, so the add-on is silently skipped there. + supportedAdapters: ["claude"], + parseOptions(rawOptions) { + return ztkOptionsSchema.parse(rawOptions ?? {}); + }, + prepare(_ctx: AddOnContext, options: ZtkOptions) { + // Resolve eagerly so the user sees a missing-binary error at session + // start rather than on the first Bash call. + resolveZtkBinary(options); + }, + contribute(_ctx: AddOnContext, options: ZtkOptions) { + const binaryPath = resolveZtkBinary(options); + return { + preToolUse: [makeZtkBashHook(binaryPath, options)], + }; + }, +}; diff --git a/packages/agent/src/server/agent-server.ts b/packages/agent/src/server/agent-server.ts index af4998a4a..d5df8346b 100644 --- a/packages/agent/src/server/agent-server.ts +++ b/packages/agent/src/server/agent-server.ts @@ -954,6 +954,7 @@ export class AgentServer { allowedDomains: this.config.allowedDomains, jsonSchema: preTask?.json_schema ?? null, permissionMode: initialPermissionMode, + ...(this.config.addOns && { addOns: this.config.addOns }), ...(this.config.claudeCode?.plugins?.length && { claudeCode: { options: { diff --git a/packages/agent/src/server/bin.ts b/packages/agent/src/server/bin.ts index d72a91ba3..ba371b613 100644 --- a/packages/agent/src/server/bin.ts +++ b/packages/agent/src/server/bin.ts @@ -3,7 +3,11 @@ import { Command } from "commander"; import { z } from "zod/v4"; import { isSupportedReasoningEffort } from "../adapters/reasoning-effort"; import { AgentServer } from "./agent-server"; -import { claudeCodeConfigSchema, mcpServersSchema } from "./schemas"; +import { + addOnsConfigSchema, + claudeCodeConfigSchema, + mcpServersSchema, +} from "./schemas"; const envSchema = z.object({ JWT_PUBLIC_KEY: z @@ -96,6 +100,10 @@ program "--allowedDomains ", "Comma-separated list of domains allowed for web tools (WebFetch, WebSearch)", ) + .option( + "--addOns ", + "Add-on config as JSON object mapping add-on name → options (sourced from task.options.add_ons)", + ) .action(async (options) => { const envResult = envSchema.safeParse(process.env); @@ -122,6 +130,11 @@ program claudeCodeConfigSchema, "--claudeCodeConfig", ); + const addOns = parseJsonOption( + options.addOns, + addOnsConfigSchema, + "--addOns", + ); const allowedDomains = options.allowedDomains ? options.allowedDomains @@ -163,6 +176,7 @@ program runtimeAdapter: env.POSTHOG_CODE_RUNTIME_ADAPTER, model: env.POSTHOG_CODE_MODEL, reasoningEffort: env.POSTHOG_CODE_REASONING_EFFORT, + addOns, }); process.on("SIGINT", async () => { diff --git a/packages/agent/src/server/schemas.ts b/packages/agent/src/server/schemas.ts index 2dfa791a9..cbecd2620 100644 --- a/packages/agent/src/server/schemas.ts +++ b/packages/agent/src/server/schemas.ts @@ -29,6 +29,16 @@ export const mcpServersSchema = z.array(remoteMcpServerSchema); export type RemoteMcpServer = z.infer; +/** + * Per-add-on options blob; opaque at the transport layer. Each add-on + * validates its own options via its `parseOptions` implementation at + * session start. + */ +export const addOnsConfigSchema = z.record( + z.string(), + z.record(z.string(), z.unknown()), +); + export const claudeCodeConfigSchema = z.object({ systemPrompt: z .union([ diff --git a/packages/agent/src/server/types.ts b/packages/agent/src/server/types.ts index 10cf96fc7..57483ac65 100644 --- a/packages/agent/src/server/types.ts +++ b/packages/agent/src/server/types.ts @@ -1,3 +1,4 @@ +import type { AddOnConfig } from "../add-ons/types"; import type { AgentMode } from "../types"; import type { RemoteMcpServer } from "./schemas"; @@ -27,4 +28,11 @@ export interface AgentServerConfig { runtimeAdapter?: "claude" | "codex"; model?: string; reasoningEffort?: "low" | "medium" | "high" | "xhigh" | "max"; + /** + * Add-on configuration sourced from `task.options.add_ons`. Forwarded + * verbatim onto `_meta.addOns` of the cloud `newSession` call where the + * adapter's add-on registry resolves it. Names not registered on the + * sandbox-side `defaultAddOnRegistry` are skipped with a warning. + */ + addOns?: AddOnConfig; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 18e5572c0..6ec3583ee 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -2,6 +2,7 @@ import type { GitHandoffCheckpoint, HandoffLocalGitState as GitHandoffLocalGitState, } from "@posthog/git/handoff"; +import type { AddOnConfig } from "./add-ons/types"; /** * Stored custom notification following ACP extensibility model. @@ -25,6 +26,20 @@ export interface StoredNotification { */ export type StoredEntry = StoredNotification; +/** + * Per-task configuration blob stored on the Django `Task.options` JSONField. + * Open-ended so we can add new keys without OpenAPI churn. Today only + * `add_ons` is consumed by the agent runtime. + * + * Server-side requires a `options = models.JSONField(default=dict, blank=True)` + * field on the Task model and a matching serializer entry. Until that migration + * lands, this field will be absent on every Task returned by the API. + */ +export interface TaskOptions { + /** Keys are add-on names; values are opaque option blobs validated per-add-on. */ + add_ons?: AddOnConfig; +} + // PostHog Task model (matches PostHog Code's OpenAPI schema) export interface Task { id: string; @@ -42,6 +57,7 @@ export interface Task { repository: string; // Format: "organization/repository" (e.g., "posthog/posthog-js") json_schema?: Record | null; // JSON schema for task output validation internal?: boolean; + options?: TaskOptions; created_at: string; updated_at: string; created_by?: { diff --git a/packages/agent/tsup.config.ts b/packages/agent/tsup.config.ts index ef30a1681..0b6f354ba 100644 --- a/packages/agent/tsup.config.ts +++ b/packages/agent/tsup.config.ts @@ -79,6 +79,10 @@ export default defineConfig([ "src/pr-url-detector.ts", "src/resume.ts", "src/types.ts", + "src/add-ons/types.ts", + "src/add-ons/registry.ts", + "src/add-ons/default-registry.ts", + "src/add-ons/ztk.ts", "src/adapters/claude/questions/utils.ts", "src/adapters/claude/permissions/permission-options.ts", "src/adapters/claude/tools.ts",