feat(omo-codex): block budgeted create_goal calls

This commit is contained in:
YeonGyu-Kim
2026-05-28 14:19:14 +09:00
parent 68848233c4
commit 1d3847a6d3
7 changed files with 219 additions and 1 deletions
@@ -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"
}
]
}
]
}
}
@@ -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 <subcommand> [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<number> {
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;
}
@@ -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<string> {
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<void> {
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<void> {
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<string, unknown> {
return typeof value === "object" && value !== null && !Array.isArray(value);
}
@@ -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");
});
});
@@ -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", () => {
@@ -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)$",
@@ -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");