diff --git a/packages/omo-codex/plugin/components/ultragoal/hooks/hooks.json b/packages/omo-codex/plugin/components/ultragoal/hooks/hooks.json index f3e52faca..f674eabff 100644 --- a/packages/omo-codex/plugin/components/ultragoal/hooks/hooks.json +++ b/packages/omo-codex/plugin/components/ultragoal/hooks/hooks.json @@ -11,6 +11,19 @@ } ] } + ], + "PreToolUse": [ + { + "matcher": "^create_goal$", + "hooks": [ + { + "type": "command", + "command": "node \"${PLUGIN_ROOT}/dist/cli.js\" hook pre-tool-use", + "timeout": 5, + "statusMessage": "enforcing unlimited ultragoal budget" + } + ] + } ] } } diff --git a/packages/omo-codex/plugin/components/ultragoal/src/cli.ts b/packages/omo-codex/plugin/components/ultragoal/src/cli.ts index 4570f6921..cc1ec57d4 100644 --- a/packages/omo-codex/plugin/components/ultragoal/src/cli.ts +++ b/packages/omo-codex/plugin/components/ultragoal/src/cli.ts @@ -1,6 +1,6 @@ #!/usr/bin/env node import { ultragoalCommand } from "./cli-commands.js"; -import { runUltragoalHookCli } from "./codex-hook.js"; +import { runPreToolUseGoalBudgetGuardCli, runUltragoalHookCli } from "./codex-hook.js"; const TOP_LEVEL_HELP = "Usage:\n omo ultragoal [args]\n omo hook user-prompt-submit (Codex UserPromptSubmit hook)\n omo help | --help | -h (this message)\n\nRun `omo ultragoal help` for ultragoal subcommands.\n"; @@ -19,6 +19,10 @@ async function main(): Promise { await runUltragoalHookCli(process.stdin, process.stdout); return 0; } + if (sub === "pre-tool-use") { + await runPreToolUseGoalBudgetGuardCli(process.stdin, process.stdout); + return 0; + } process.stderr.write(`[omo] unknown hook subcommand: ${sub ?? "(none)"}\n`); return 1; } diff --git a/packages/omo-codex/plugin/components/ultragoal/src/codex-hook.ts b/packages/omo-codex/plugin/components/ultragoal/src/codex-hook.ts index 3b6a7687f..20bab204b 100644 --- a/packages/omo-codex/plugin/components/ultragoal/src/codex-hook.ts +++ b/packages/omo-codex/plugin/components/ultragoal/src/codex-hook.ts @@ -11,6 +11,32 @@ export interface UserPromptSubmitPayload { readonly turn_id?: string; } +export interface PreToolUsePayload { + readonly cwd: string; + readonly hook_event_name: "PreToolUse"; + readonly model: string; + readonly permission_mode: string; + readonly session_id: string; + readonly tool_input: unknown; + readonly tool_name: string; + readonly tool_use_id: string; + readonly transcript_path: string | null; + readonly turn_id: string; +} + +interface PreToolUseHookOutput { + readonly hookSpecificOutput: { + readonly hookEventName: "PreToolUse"; + readonly permissionDecision: "deny"; + readonly permissionDecisionReason: string; + readonly additionalContext: string; + }; +} + +const CREATE_GOAL_TOOL_NAME = "create_goal"; +const GOAL_BUDGET_WARNING = + "Do not set token_budget on create_goal. Omit the budget field so the goal stays unlimited; ultrawork and ultragoal runs must always use unlimited goals."; + export function parseUserPromptSubmitPayload(raw: string): UserPromptSubmitPayload | null { if (raw.trim().length === 0) return null; try { @@ -22,6 +48,17 @@ export function parseUserPromptSubmitPayload(raw: string): UserPromptSubmitPaylo } } +export function parsePreToolUsePayload(raw: string): PreToolUsePayload | null { + if (raw.trim().length === 0) return null; + try { + const parsed: unknown = JSON.parse(raw); + return isPreToolUsePayload(parsed) ? parsed : null; + } catch (error) { + if (error instanceof SyntaxError) return null; + return null; + } +} + export async function applyUserPromptUltragoalSteering(payload: UserPromptSubmitPayload): Promise { try { if (payload.hook_event_name !== "UserPromptSubmit") return ""; @@ -41,6 +78,21 @@ export async function applyUserPromptUltragoalSteering(payload: UserPromptSubmit } } +export function applyPreToolUseGoalBudgetGuard(payload: PreToolUsePayload): string { + if (payload.hook_event_name !== "PreToolUse") return ""; + if (payload.tool_name !== CREATE_GOAL_TOOL_NAME) return ""; + if (!hasGoalBudgetInput(payload.tool_input)) return ""; + const output: PreToolUseHookOutput = { + hookSpecificOutput: { + hookEventName: "PreToolUse", + permissionDecision: "deny", + permissionDecisionReason: GOAL_BUDGET_WARNING, + additionalContext: GOAL_BUDGET_WARNING, + }, + }; + return `${JSON.stringify(output)}\n`; +} + export async function runUltragoalHookCli(stdin: NodeJS.ReadableStream, stdout: NodeJS.WritableStream): Promise { try { const payload = parseUserPromptSubmitPayload(await readAll(stdin)); @@ -53,6 +105,21 @@ export async function runUltragoalHookCli(stdin: NodeJS.ReadableStream, stdout: } } +export async function runPreToolUseGoalBudgetGuardCli( + stdin: NodeJS.ReadableStream, + stdout: NodeJS.WritableStream, +): Promise { + try { + const payload = parsePreToolUsePayload(await readAll(stdin)); + if (payload === null) return; + const output = applyPreToolUseGoalBudgetGuard(payload); + if (output.length > 0) stdout.write(output); + } catch (error) { + if (error instanceof Error) return; + return; + } +} + function isUserPromptSubmitPayload(value: unknown): value is UserPromptSubmitPayload { if (!isRecord(value)) return false; return ( @@ -64,6 +131,26 @@ function isUserPromptSubmitPayload(value: unknown): value is UserPromptSubmitPay ); } +function isPreToolUsePayload(value: unknown): value is PreToolUsePayload { + if (!isRecord(value)) return false; + return ( + value["hook_event_name"] === "PreToolUse" && + typeof value["cwd"] === "string" && + typeof value["model"] === "string" && + typeof value["permission_mode"] === "string" && + typeof value["session_id"] === "string" && + typeof value["tool_name"] === "string" && + typeof value["tool_use_id"] === "string" && + (value["transcript_path"] === null || typeof value["transcript_path"] === "string") && + typeof value["turn_id"] === "string" && + Object.hasOwn(value, "tool_input") + ); +} + +function hasGoalBudgetInput(value: unknown): boolean { + return isRecord(value) && (Object.hasOwn(value, "token_budget") || Object.hasOwn(value, "tokenBudget")); +} + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } diff --git a/packages/omo-codex/plugin/components/ultragoal/test/codex-hook.test.ts b/packages/omo-codex/plugin/components/ultragoal/test/codex-hook.test.ts index c6be05e81..fa524dfdd 100644 --- a/packages/omo-codex/plugin/components/ultragoal/test/codex-hook.test.ts +++ b/packages/omo-codex/plugin/components/ultragoal/test/codex-hook.test.ts @@ -5,8 +5,11 @@ import { Readable, Writable } from "node:stream"; import { describe, expect, it } from "vitest"; import { + applyPreToolUseGoalBudgetGuard, applyUserPromptUltragoalSteering, + type PreToolUsePayload, parseUserPromptSubmitPayload, + runPreToolUseGoalBudgetGuardCli, runUltragoalHookCli, type UserPromptSubmitPayload, } from "../src/codex-hook.js"; @@ -50,6 +53,21 @@ function payload(prompt: string, cwd: string): UserPromptSubmitPayload { return { cwd, hook_event_name: "UserPromptSubmit", prompt, session_id: "s1" }; } +function preToolPayload(toolName: string, toolInput: unknown): PreToolUsePayload { + return { + cwd: "/repo", + hook_event_name: "PreToolUse", + model: "gpt-5.5", + permission_mode: "default", + session_id: "s1", + tool_input: toolInput, + tool_name: toolName, + tool_use_id: "call-1", + transcript_path: null, + turn_id: "turn-1", + }; +} + function payloadWithRuntimeEvent(hookEventName: string): UserPromptSubmitPayload { const input = payload( 'OMO_ULTRAGOAL_STEER: {"kind":"annotate_ledger","source":"user_prompt_submit","evidence":"x","rationale":"y"}', @@ -185,3 +203,64 @@ describe("runUltragoalHookCli (stdin/stdout integration)", () => { expect(capture.read()).toBe(""); }); }); + +describe("applyPreToolUseGoalBudgetGuard", () => { + it("#given create_goal sets token_budget #when PreToolUse runs #then it blocks with unlimited-goal warning", () => { + // given + const input = preToolPayload("create_goal", { objective: "Ship the feature", token_budget: 5000 }); + + // when + const output = applyPreToolUseGoalBudgetGuard(input); + + // then + const parsed = JSON.parse(output); + expect(parsed).toMatchObject({ + hookSpecificOutput: { + hookEventName: "PreToolUse", + permissionDecision: "deny", + }, + }); + expect(parsed.hookSpecificOutput.permissionDecisionReason).toContain("Do not set token_budget on create_goal"); + expect(parsed.hookSpecificOutput.permissionDecisionReason).toContain("unlimited"); + }); + + it("#given create_goal omits token_budget #when PreToolUse runs #then it stays silent", () => { + // given + const input = preToolPayload("create_goal", { objective: "Ship the feature" }); + + // when + const output = applyPreToolUseGoalBudgetGuard(input); + + // then + expect(output).toBe(""); + }); + + it("#given a neighboring tool includes token_budget text #when PreToolUse runs #then it stays silent", () => { + // given + const input = preToolPayload("update_goal", { status: "complete", token_budget: 5000 }); + + // when + const output = applyPreToolUseGoalBudgetGuard(input); + + // then + expect(output).toBe(""); + }); +}); + +describe("runPreToolUseGoalBudgetGuardCli", () => { + it("#given Codex PreToolUse stdin with budgeted create_goal #when CLI hook runs #then it writes blocking JSON", async () => { + // given + const stdin = Readable.from([ + JSON.stringify(preToolPayload("create_goal", { objective: "Ship", token_budget: 1 })), + ]); + const capture = captureStdout(); + + // when + await runPreToolUseGoalBudgetGuardCli(stdin, capture.stdout); + + // then + const parsed = JSON.parse(capture.read()); + expect(parsed.hookSpecificOutput.permissionDecision).toBe("deny"); + expect(parsed.hookSpecificOutput.permissionDecisionReason).toContain("unlimited"); + }); +}); diff --git a/packages/omo-codex/plugin/components/ultragoal/test/package-smoke.test.ts b/packages/omo-codex/plugin/components/ultragoal/test/package-smoke.test.ts index a1181e524..287e76b8f 100644 --- a/packages/omo-codex/plugin/components/ultragoal/test/package-smoke.test.ts +++ b/packages/omo-codex/plugin/components/ultragoal/test/package-smoke.test.ts @@ -54,6 +54,14 @@ describe("hooks/hooks.json", () => { expect(command).toContain("dist/cli.js"); expect(command).toContain("hook user-prompt-submit"); }); + + it("#given ultragoal component is enabled #when hooks are inspected #then create_goal PreToolUse guard is registered", async () => { + const text = await readText("hooks/hooks.json"); + + expect(text).toContain('"PreToolUse"'); + expect(text).toContain('"matcher": "^create_goal$"'); + expect(text).toContain("hook pre-tool-use"); + }); }); describe("src/cli.ts", () => { diff --git a/packages/omo-codex/plugin/hooks/hooks.json b/packages/omo-codex/plugin/hooks/hooks.json index d93a5d634..cbef5941b 100644 --- a/packages/omo-codex/plugin/hooks/hooks.json +++ b/packages/omo-codex/plugin/hooks/hooks.json @@ -52,6 +52,19 @@ ] } ], + "PreToolUse": [ + { + "matcher": "^create_goal$", + "hooks": [ + { + "type": "command", + "command": "node \"${PLUGIN_ROOT}/components/ultragoal/dist/cli.js\" hook pre-tool-use", + "timeout": 5, + "statusMessage": "enforcing OMO unlimited goal budget" + } + ] + } + ], "PostToolUse": [ { "matcher": "^(apply_patch|write|Write|edit|Edit|multi_edit|multiedit|MultiEdit)$", diff --git a/packages/omo-codex/plugin/test/aggregate.test.mjs b/packages/omo-codex/plugin/test/aggregate.test.mjs index cf7630d52..504b78a4d 100644 --- a/packages/omo-codex/plugin/test/aggregate.test.mjs +++ b/packages/omo-codex/plugin/test/aggregate.test.mjs @@ -57,6 +57,20 @@ test("#given isolated components #when hooks are inspected #then commands stay i assert.doesNotMatch(text, /codex-(comment-checker|lsp|rules|telemetry|ultragoal|ultrawork)@/); }); +test("#given aggregate OMO plugin is enabled #when hooks are inspected #then ultragoal guards budgeted create_goal calls", async () => { + // given + const hooks = await readJson("hooks/hooks.json"); + const text = JSON.stringify(hooks); + + // when + const preToolUseGroups = hooks.hooks.PreToolUse; + + // then + assert.match(text, /components\/ultragoal\/dist\/cli\.js/); + assert.match(text, /hook pre-tool-use/); + assert.deepEqual(preToolUseGroups.map((group) => group.matcher), ["^create_goal$"]); +}); + test("#given aggregate MCP config #when inspected #then lsp server stays component isolated", async () => { // given const mcp = await readJson(".mcp.json");