feat(goals): added per-session goal mode with token budget tracking
- Added GoalRuntime with wall-clock and token accounting, budget steering, and lifecycle operations (create, pause, resume, drop, complete). - Exposed goal tool as a hidden agent tool, activated only when goal mode is enabled. - Integrated goal continuation loop in InteractiveMode with auto-submit between turns. - Added status line segment and theme icons for goal mode state.
This commit is contained in:
can1357
@@ -0,0 +1,271 @@
|
||||
import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test";
|
||||
import * as path from "node:path";
|
||||
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
|
||||
import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry";
|
||||
import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
|
||||
import { GoalTool } from "@oh-my-pi/pi-coding-agent/goals/tools/goal-tool";
|
||||
import { InteractiveMode } from "@oh-my-pi/pi-coding-agent/modes/interactive-mode";
|
||||
import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
|
||||
import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session";
|
||||
import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage";
|
||||
import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
|
||||
import { createTools, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
|
||||
import { TempDir } from "@oh-my-pi/pi-utils";
|
||||
|
||||
function createToolSession(cwd: string, settings: Settings, overrides: Partial<ToolSession> = {}): ToolSession {
|
||||
return {
|
||||
cwd,
|
||||
hasUI: false,
|
||||
getSessionFile: () => null,
|
||||
getSessionSpawns: () => "*",
|
||||
settings,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
type GoalHarness = {
|
||||
tempDir: TempDir;
|
||||
authStorage: AuthStorage;
|
||||
settings: Settings;
|
||||
session: AgentSession;
|
||||
mode: InteractiveMode;
|
||||
toolSession: ToolSession;
|
||||
cleanup: () => Promise<void>;
|
||||
};
|
||||
|
||||
async function createGoalHarness(): Promise<GoalHarness> {
|
||||
resetSettingsForTest();
|
||||
const tempDir = TempDir.createSync("@pi-goal-mode-");
|
||||
await Settings.init({ inMemory: true, cwd: tempDir.path() });
|
||||
const authStorage = await AuthStorage.create(path.join(tempDir.path(), "testauth.db"));
|
||||
const modelRegistry = new ModelRegistry(authStorage);
|
||||
const model = modelRegistry.find("anthropic", "claude-sonnet-4-5");
|
||||
if (!model) {
|
||||
throw new Error("Expected claude-sonnet-4-5 to exist in registry");
|
||||
}
|
||||
|
||||
const settings = Settings.isolated({
|
||||
"compaction.enabled": false,
|
||||
"goal.enabled": true,
|
||||
"plan.enabled": true,
|
||||
});
|
||||
const bootstrapToolSession = createToolSession(tempDir.path(), settings);
|
||||
const initialTools = await createTools(bootstrapToolSession, ["read"]);
|
||||
const toolRegistry = new Map<string, AgentTool>(initialTools.map(tool => [tool.name, tool] as const));
|
||||
|
||||
const session = new AgentSession({
|
||||
agent: new Agent({
|
||||
initialState: {
|
||||
model,
|
||||
systemPrompt: ["Test"],
|
||||
tools: initialTools,
|
||||
messages: [],
|
||||
},
|
||||
}),
|
||||
sessionManager: SessionManager.create(tempDir.path(), tempDir.path()),
|
||||
settings,
|
||||
modelRegistry,
|
||||
toolRegistry,
|
||||
rebuildSystemPrompt: async () => ({ systemPrompt: ["Test"] }),
|
||||
});
|
||||
const mode = new InteractiveMode(session, "test");
|
||||
const toolSession = createToolSession(tempDir.path(), settings, {
|
||||
getGoalModeState: () => session.getGoalModeState(),
|
||||
getGoalRuntime: () => session.goalRuntime,
|
||||
});
|
||||
toolRegistry.set("goal", new GoalTool(toolSession));
|
||||
|
||||
return {
|
||||
tempDir,
|
||||
authStorage,
|
||||
settings,
|
||||
session,
|
||||
mode,
|
||||
toolSession,
|
||||
cleanup: async () => {
|
||||
mode.stop();
|
||||
await session.dispose();
|
||||
authStorage.close();
|
||||
tempDir.removeSync();
|
||||
resetSettingsForTest();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function toolNamesFor(harness: GoalHarness): Promise<string[]> {
|
||||
return (await createTools(harness.toolSession, harness.session.getActiveToolNames())).map(tool => tool.name);
|
||||
}
|
||||
|
||||
describe("InteractiveMode goal mode integration", () => {
|
||||
let harness: GoalHarness;
|
||||
|
||||
beforeAll(() => {
|
||||
initTheme();
|
||||
});
|
||||
|
||||
beforeEach(async () => {
|
||||
harness = await createGoalHarness();
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
vi.restoreAllMocks();
|
||||
await harness.cleanup();
|
||||
});
|
||||
|
||||
it("toggles goal tool exposure when goal mode enters and pauses", async () => {
|
||||
expect(await toolNamesFor(harness)).not.toContain("goal");
|
||||
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
|
||||
expect(harness.mode.goalModeEnabled).toBe(true);
|
||||
expect(harness.session.getGoalModeState()?.enabled).toBe(true);
|
||||
expect(await toolNamesFor(harness)).toContain("goal");
|
||||
|
||||
vi.spyOn(harness.mode, "showHookSelector").mockResolvedValue("Pause");
|
||||
await harness.mode.handleGoalModeCommand();
|
||||
|
||||
expect(harness.mode.goalModeEnabled).toBe(false);
|
||||
expect(harness.mode.goalModePaused).toBe(true);
|
||||
expect(harness.session.getGoalModeState()?.goal.status).toBe("paused");
|
||||
expect(await toolNamesFor(harness)).not.toContain("goal");
|
||||
});
|
||||
|
||||
it("refuses /goal while plan mode is active", async () => {
|
||||
const showWarning = vi.spyOn(harness.mode, "showWarning");
|
||||
harness.mode.planModeEnabled = true;
|
||||
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
|
||||
expect(showWarning).toHaveBeenCalledWith("Exit plan mode first.");
|
||||
expect(harness.session.getGoalModeState()).toBeUndefined();
|
||||
});
|
||||
|
||||
it("refuses /plan while goal mode is active", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
const showWarning = vi.spyOn(harness.mode, "showWarning");
|
||||
|
||||
await harness.mode.handlePlanModeCommand();
|
||||
|
||||
expect(showWarning).toHaveBeenCalledWith("Exit goal mode first.");
|
||||
expect(harness.mode.planModeEnabled).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects a new /goal objective while paused", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
vi.spyOn(harness.mode, "showHookSelector").mockResolvedValue("Pause");
|
||||
await harness.mode.handleGoalModeCommand();
|
||||
const showWarning = vi.spyOn(harness.mode, "showWarning");
|
||||
|
||||
await harness.mode.handleGoalModeCommand("Replace the objective");
|
||||
|
||||
expect(showWarning).toHaveBeenCalledWith(
|
||||
"Resume the current goal first, or drop it before setting a new objective.",
|
||||
);
|
||||
expect(harness.session.getGoalModeState()?.enabled).toBe(false);
|
||||
expect(harness.session.getGoalModeState()?.goal.objective).toBe("Ship the release");
|
||||
expect(harness.session.getGoalModeState()?.goal.status).toBe("paused");
|
||||
});
|
||||
|
||||
it("resumes the paused goal via the bare /goal menu", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
const selector = vi.spyOn(harness.mode, "showHookSelector").mockResolvedValueOnce("Pause");
|
||||
await harness.mode.handleGoalModeCommand();
|
||||
expect(harness.mode.goalModePaused).toBe(true);
|
||||
selector.mockResolvedValueOnce("Resume");
|
||||
const showStatus = vi.spyOn(harness.mode, "showStatus");
|
||||
|
||||
await harness.mode.handleGoalModeCommand();
|
||||
|
||||
expect(showStatus).toHaveBeenCalledWith("Goal mode resumed.");
|
||||
expect(harness.mode.goalModeEnabled).toBe(true);
|
||||
expect(harness.mode.goalModePaused).toBe(false);
|
||||
expect(harness.session.getGoalModeState()?.enabled).toBe(true);
|
||||
expect(harness.session.getGoalModeState()?.goal.objective).toBe("Ship the release");
|
||||
expect(harness.session.getGoalModeState()?.goal.status).toBe("active");
|
||||
expect(await toolNamesFor(harness)).toContain("goal");
|
||||
});
|
||||
|
||||
it("mutates the goal token budget via /goal budget without resetting accumulated usage", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
// Seed accumulated usage by driving the runtime directly — equivalent to a turn's flush.
|
||||
const goal = harness.session.getGoalModeState()?.goal;
|
||||
if (!goal) throw new Error("expected active goal");
|
||||
goal.tokensUsed = 42;
|
||||
goal.timeUsedSeconds = 5;
|
||||
|
||||
await harness.mode.handleGoalModeCommand("budget 123");
|
||||
|
||||
const after = harness.session.getGoalModeState();
|
||||
expect(after?.goal.tokenBudget).toBe(123);
|
||||
// Accumulated counters are preserved across the mutation.
|
||||
expect(after?.goal.tokensUsed).toBe(42);
|
||||
expect(after?.goal.timeUsedSeconds).toBe(5);
|
||||
|
||||
await harness.mode.handleGoalModeCommand("budget off");
|
||||
expect(harness.session.getGoalModeState()?.goal.tokenBudget).toBeUndefined();
|
||||
expect(harness.session.getGoalModeState()?.goal.tokensUsed).toBe(42);
|
||||
});
|
||||
|
||||
it("refuses /goal budget while only a paused goal exists (fix #5)", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
vi.spyOn(harness.mode, "showHookSelector").mockResolvedValue("Pause");
|
||||
await harness.mode.handleGoalModeCommand();
|
||||
expect(harness.mode.goalModePaused).toBe(true);
|
||||
const showWarning = vi.spyOn(harness.mode, "showWarning");
|
||||
|
||||
await harness.mode.handleGoalModeCommand("budget 99");
|
||||
|
||||
expect(showWarning).toHaveBeenCalledWith("Resume the goal before adjusting the budget.");
|
||||
// Mutation must not have run while the goal is paused.
|
||||
expect(harness.session.getGoalModeState()?.goal.tokenBudget).toBeUndefined();
|
||||
});
|
||||
|
||||
it("returns the completion report from the goal tool and exits goal mode before the next turn rebuild", async () => {
|
||||
await harness.mode.handleGoalModeCommand("Ship the release");
|
||||
await harness.mode.handleGoalModeCommand("budget 50");
|
||||
const appendCustomEntry = vi.spyOn(harness.session.sessionManager, "appendCustomEntry");
|
||||
const goalTool = (await createTools(harness.toolSession, harness.session.getActiveToolNames())).find(
|
||||
tool => tool.name === "goal",
|
||||
);
|
||||
if (!goalTool) {
|
||||
throw new Error("Expected goal tool to be active");
|
||||
}
|
||||
|
||||
const result = await goalTool.execute("call-1", { op: "complete" });
|
||||
const completionText = JSON.stringify(result.content);
|
||||
|
||||
expect(result.details?.completionBudgetReport).toBe(
|
||||
"Goal achieved. Report final budget usage to the user: tokens used: 0 of 50.",
|
||||
);
|
||||
expect(completionText).toContain("Goal achieved. Report final budget usage to the user: tokens used: 0 of 50.");
|
||||
expect(harness.session.getGoalModeState()?.mode).toBe("exiting");
|
||||
// Per fix #1: completeGoalFromTool clears state.enabled so subsequent createTools
|
||||
// calls (e.g. mid-turn refreshes) no longer advertise the goal tool. The model's
|
||||
// existing toolset for the in-flight turn is unaffected — what we care about here
|
||||
// is that the next createTools observation reflects the deactivation.
|
||||
expect(harness.session.getGoalModeState()?.enabled).toBe(false);
|
||||
expect(await toolNamesFor(harness)).not.toContain("goal");
|
||||
|
||||
const nextTurn = harness.mode.getUserInput();
|
||||
// getUserInput observes mode === "exiting" and awaits #exitGoalMode before
|
||||
// arming onInputCallback. Drain microtasks until that side-effect lands.
|
||||
for (let i = 0; i < 100 && harness.session.getGoalModeState() !== undefined; i++) {
|
||||
await Bun.sleep(0);
|
||||
}
|
||||
expect(harness.mode.goalModeEnabled).toBe(false);
|
||||
expect(harness.mode.goalModePaused).toBe(false);
|
||||
expect(harness.session.getGoalModeState()).toBeUndefined();
|
||||
expect(await toolNamesFor(harness)).not.toContain("goal");
|
||||
expect(appendCustomEntry).toHaveBeenCalledWith(
|
||||
"goal-completed",
|
||||
expect.objectContaining({
|
||||
objective: "Ship the release",
|
||||
tokenBudget: 50,
|
||||
tokensUsed: 0,
|
||||
}),
|
||||
);
|
||||
|
||||
harness.mode.onInputCallback?.(harness.mode.startPendingSubmission({ text: "next turn" }));
|
||||
await nextTurn;
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,299 @@
|
||||
import { describe, expect, it } from "bun:test";
|
||||
import {
|
||||
escapeXmlText,
|
||||
GoalRuntime,
|
||||
type GoalRuntimeHost,
|
||||
goalTokenDelta,
|
||||
renderGoalPrompt,
|
||||
renderTrustedObjective,
|
||||
} from "@oh-my-pi/pi-coding-agent/goals/runtime";
|
||||
import type { Goal, GoalModeState, GoalRuntimeEvent, GoalTokenUsage } from "@oh-my-pi/pi-coding-agent/goals/state";
|
||||
|
||||
function createUsage(overrides: Partial<GoalTokenUsage> = {}): GoalTokenUsage {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function createGoal(overrides: Partial<Goal> = {}): Goal {
|
||||
return {
|
||||
id: "goal-1",
|
||||
objective: "Ship <fast> & safely",
|
||||
status: "active",
|
||||
tokenBudget: undefined,
|
||||
tokensUsed: 0,
|
||||
timeUsedSeconds: 0,
|
||||
createdAt: 0,
|
||||
updatedAt: 0,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function cloneGoal(goal: Goal): Goal {
|
||||
return { ...goal };
|
||||
}
|
||||
|
||||
function cloneState(state: GoalModeState | undefined): GoalModeState | undefined {
|
||||
return state ? { ...state, goal: cloneGoal(state.goal) } : undefined;
|
||||
}
|
||||
|
||||
function cloneEvent(event: GoalRuntimeEvent): GoalRuntimeEvent {
|
||||
if (event.type === "goal_updated") {
|
||||
return {
|
||||
...event,
|
||||
goal: event.goal ? cloneGoal(event.goal) : null,
|
||||
state: cloneState(event.state),
|
||||
};
|
||||
}
|
||||
return { ...event };
|
||||
}
|
||||
|
||||
function createHarness(initial: { state?: GoalModeState; usage?: GoalTokenUsage; now?: number } = {}) {
|
||||
let state = cloneState(initial.state);
|
||||
let usage = createUsage(initial.usage);
|
||||
let now = initial.now ?? 0;
|
||||
const events: GoalRuntimeEvent[] = [];
|
||||
const persists: Array<{ mode: "goal" | "goal_paused" | "none"; state?: GoalModeState }> = [];
|
||||
const hiddenMessages: Array<{ customType: string; content: string; deliverAs?: "steer" | "followUp" | "nextTurn" }> =
|
||||
[];
|
||||
const host: GoalRuntimeHost = {
|
||||
getState: () => cloneState(state),
|
||||
setState: next => {
|
||||
state = cloneState(next);
|
||||
},
|
||||
getCurrentUsage: () => createUsage(usage),
|
||||
emit: async event => {
|
||||
events.push(cloneEvent(event));
|
||||
},
|
||||
persist: (mode, persistedState) => {
|
||||
persists.push({ mode, state: cloneState(persistedState) });
|
||||
},
|
||||
sendHiddenMessage: async message => {
|
||||
hiddenMessages.push({ ...message });
|
||||
},
|
||||
now: () => now,
|
||||
};
|
||||
return {
|
||||
runtime: new GoalRuntime(host),
|
||||
getState: () => cloneState(state),
|
||||
setUsage: (next: Partial<GoalTokenUsage>) => {
|
||||
usage = createUsage(next);
|
||||
},
|
||||
advance: (ms: number) => {
|
||||
now += ms;
|
||||
},
|
||||
events,
|
||||
persists,
|
||||
hiddenMessages,
|
||||
};
|
||||
}
|
||||
|
||||
describe("goal runtime", () => {
|
||||
it("counts cache writes but ignores cache reads in token deltas", () => {
|
||||
expect(
|
||||
goalTokenDelta(
|
||||
createUsage({ input: 13, output: 6, cacheRead: 999, cacheWrite: 8 }),
|
||||
createUsage({ input: 10, output: 4, cacheRead: 1, cacheWrite: 5 }),
|
||||
),
|
||||
).toBe(8);
|
||||
});
|
||||
|
||||
it("clamps token deltas at zero across usage resets", () => {
|
||||
expect(
|
||||
goalTokenDelta(
|
||||
createUsage({ input: 10, output: 5, cacheRead: 0, cacheWrite: 2 }),
|
||||
createUsage({ input: 100, output: 50, cacheRead: 500, cacheWrite: 20 }),
|
||||
),
|
||||
).toBe(0);
|
||||
});
|
||||
|
||||
it("advances wall-clock accounting only by persisted whole seconds", async () => {
|
||||
const harness = createHarness({
|
||||
state: { enabled: true, mode: "active", goal: createGoal() },
|
||||
});
|
||||
|
||||
harness.runtime.onTurnStart("turn-1", createUsage());
|
||||
harness.advance(2_500);
|
||||
await harness.runtime.flushUsage("suppressed");
|
||||
expect(harness.getState()?.goal.timeUsedSeconds).toBe(2);
|
||||
expect(harness.runtime.snapshot.wallClock.lastAccountedAt).toBe(2_000);
|
||||
expect(harness.persists).toHaveLength(1);
|
||||
|
||||
harness.advance(400);
|
||||
await harness.runtime.flushUsage("suppressed");
|
||||
expect(harness.getState()?.goal.timeUsedSeconds).toBe(2);
|
||||
expect(harness.runtime.snapshot.wallClock.lastAccountedAt).toBe(2_000);
|
||||
expect(harness.persists).toHaveLength(1);
|
||||
|
||||
harness.advance(700);
|
||||
await harness.runtime.flushUsage("suppressed");
|
||||
expect(harness.getState()?.goal.timeUsedSeconds).toBe(3);
|
||||
expect(harness.runtime.snapshot.wallClock.lastAccountedAt).toBe(3_000);
|
||||
expect(harness.persists).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("steers only once until a budget mutation resets the cycle", async () => {
|
||||
const harness = createHarness({
|
||||
state: {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ tokenBudget: 10, tokensUsed: 8 }),
|
||||
},
|
||||
});
|
||||
|
||||
harness.runtime.onTurnStart("turn-1", createUsage());
|
||||
harness.setUsage({ input: 2 });
|
||||
await harness.runtime.flushUsage("allowed");
|
||||
expect(harness.getState()?.goal.status).toBe("budget-limited");
|
||||
expect(harness.hiddenMessages).toHaveLength(1);
|
||||
expect(harness.hiddenMessages[0]).toMatchObject({
|
||||
customType: "goal-budget-limit",
|
||||
deliverAs: "steer",
|
||||
});
|
||||
|
||||
harness.setUsage({ input: 5 });
|
||||
await harness.runtime.flushUsage("allowed");
|
||||
expect(harness.hiddenMessages).toHaveLength(1);
|
||||
|
||||
await harness.runtime.onBudgetMutated(20);
|
||||
expect(harness.getState()?.enabled).toBe(true);
|
||||
expect(harness.getState()?.goal.status).toBe("active");
|
||||
expect(harness.getState()?.goal.tokenBudget).toBe(20);
|
||||
expect(harness.hiddenMessages).toHaveLength(1);
|
||||
|
||||
harness.setUsage({ input: 15 });
|
||||
await harness.runtime.flushUsage("allowed");
|
||||
expect(harness.getState()?.goal.status).toBe("budget-limited");
|
||||
expect(harness.hiddenMessages).toHaveLength(2);
|
||||
});
|
||||
|
||||
it("pauses an active goal when an interruption aborts the task", async () => {
|
||||
const harness = createHarness({
|
||||
state: { enabled: true, mode: "active", goal: createGoal() },
|
||||
});
|
||||
|
||||
harness.runtime.onTurnStart("turn-1", createUsage());
|
||||
harness.advance(1_000);
|
||||
harness.setUsage({ output: 4 });
|
||||
await harness.runtime.onTaskAborted({ reason: "interrupted" });
|
||||
|
||||
const state = harness.getState();
|
||||
expect(state?.enabled).toBe(false);
|
||||
expect(state?.goal.status).toBe("paused");
|
||||
expect(state?.goal.tokensUsed).toBe(4);
|
||||
expect(state?.goal.timeUsedSeconds).toBe(1);
|
||||
expect(harness.persists.at(-1)?.mode).toBe("goal_paused");
|
||||
});
|
||||
|
||||
it("auto-pauses active goals when a thread resumes", async () => {
|
||||
const harness = createHarness({
|
||||
state: { enabled: true, mode: "active", goal: createGoal() },
|
||||
});
|
||||
|
||||
const resumed = await harness.runtime.onThreadResumed();
|
||||
expect(resumed?.enabled).toBe(false);
|
||||
expect(resumed?.goal.status).toBe("paused");
|
||||
expect(harness.getState()?.enabled).toBe(false);
|
||||
expect(harness.getState()?.goal.status).toBe("paused");
|
||||
expect(harness.persists.at(-1)?.mode).toBe("goal_paused");
|
||||
});
|
||||
|
||||
it("escapes XML in goal helpers and rendered prompts", () => {
|
||||
const objective = "Fix <root>&keep>safe";
|
||||
const goal = createGoal({ objective });
|
||||
const prompt = renderGoalPrompt("active", goal);
|
||||
|
||||
expect(renderTrustedObjective(objective)).toBe("<objective>\nFix <root>&keep>safe\n</objective>");
|
||||
expect(prompt).toContain("Fix <root>&keep>safe");
|
||||
expect(prompt).not.toContain(objective);
|
||||
});
|
||||
|
||||
it("returns the input verbatim when escapeXmlText has nothing to escape", () => {
|
||||
const input = "plain text — with 'quotes' and \"double\" plus unicode ✓";
|
||||
expect(escapeXmlText(input)).toBe(input);
|
||||
// fast-path identity: the helper should not allocate a new string when nothing changed
|
||||
expect(escapeXmlText(input)).toBe(escapeXmlText(input));
|
||||
});
|
||||
|
||||
it("escapeXmlText escapes only the XML-significant trio and leaves other characters untouched", () => {
|
||||
expect(escapeXmlText("a & b < c > d")).toBe("a & b < c > d");
|
||||
expect(escapeXmlText("'\"`")).toBe("'\"`");
|
||||
});
|
||||
|
||||
it("onBudgetMutated downward to below current usage flips active to budget-limited and steers", async () => {
|
||||
const harness = createHarness({
|
||||
state: {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ tokenBudget: 100, tokensUsed: 30, status: "active" }),
|
||||
},
|
||||
});
|
||||
|
||||
const next = await harness.runtime.onBudgetMutated(20);
|
||||
|
||||
expect(next?.goal.status).toBe("budget-limited");
|
||||
expect(next?.goal.tokenBudget).toBe(20);
|
||||
expect(next?.goal.tokensUsed).toBe(30);
|
||||
expect(harness.hiddenMessages).toHaveLength(1);
|
||||
expect(harness.hiddenMessages[0]?.customType).toBe("goal-budget-limit");
|
||||
});
|
||||
|
||||
it("completeGoalFromTool clears enabled and flips status to complete with mode exiting (fix #1)", async () => {
|
||||
const harness = createHarness({
|
||||
state: {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ tokenBudget: 100, tokensUsed: 42, timeUsedSeconds: 7 }),
|
||||
},
|
||||
});
|
||||
|
||||
const completed = await harness.runtime.completeGoalFromTool();
|
||||
|
||||
expect(completed.status).toBe("complete");
|
||||
const state = harness.getState();
|
||||
expect(state?.enabled).toBe(false);
|
||||
expect(state?.mode).toBe("exiting");
|
||||
expect(state?.reason).toBe("completed");
|
||||
expect(state?.goal.status).toBe("complete");
|
||||
});
|
||||
|
||||
it("dropGoal emits goal_updated with the dropped goal and clears persisted state", async () => {
|
||||
const harness = createHarness({
|
||||
state: {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ id: "g-99", objective: "Ship soon" }),
|
||||
},
|
||||
});
|
||||
|
||||
const dropped = await harness.runtime.dropGoal();
|
||||
|
||||
expect(dropped?.status).toBe("dropped");
|
||||
expect(dropped?.id).toBe("g-99");
|
||||
expect(harness.getState()).toBeUndefined();
|
||||
const lastEvent = harness.events.at(-1);
|
||||
if (lastEvent?.type !== "goal_updated") {
|
||||
throw new Error("expected goal_updated event after dropGoal");
|
||||
}
|
||||
expect(lastEvent.goal?.status).toBe("dropped");
|
||||
expect(lastEvent.state?.enabled).toBe(false);
|
||||
});
|
||||
|
||||
it("rejects op=create on the runtime when a non-dropped goal already exists", async () => {
|
||||
const harness = createHarness({
|
||||
state: {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ objective: "Existing" }),
|
||||
},
|
||||
});
|
||||
|
||||
await expect(harness.runtime.createGoal({ objective: "Second" })).rejects.toThrow(
|
||||
"cannot create a new goal because this session already has a goal",
|
||||
);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,210 @@
|
||||
import { describe, expect, it, vi } from "bun:test";
|
||||
import { completionBudgetReport, GoalRuntime } from "@oh-my-pi/pi-coding-agent/goals/runtime";
|
||||
import type { Goal, GoalModeState, GoalTokenUsage } from "@oh-my-pi/pi-coding-agent/goals/state";
|
||||
import { GoalTool } from "@oh-my-pi/pi-coding-agent/goals/tools/goal-tool";
|
||||
import type { ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
|
||||
|
||||
function createUsage(overrides: Partial<GoalTokenUsage> = {}): GoalTokenUsage {
|
||||
return {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function createGoal(overrides: Partial<Goal> = {}): Goal {
|
||||
return {
|
||||
id: "goal-1",
|
||||
objective: "Ship it",
|
||||
status: "active",
|
||||
tokenBudget: undefined,
|
||||
tokensUsed: 0,
|
||||
timeUsedSeconds: 0,
|
||||
createdAt: 0,
|
||||
updatedAt: 0,
|
||||
...overrides,
|
||||
};
|
||||
}
|
||||
|
||||
function cloneState(state: GoalModeState | undefined): GoalModeState | undefined {
|
||||
return state ? { ...state, goal: { ...state.goal } } : undefined;
|
||||
}
|
||||
|
||||
function createToolSession(overrides: Partial<ToolSession>): ToolSession {
|
||||
return overrides as ToolSession;
|
||||
}
|
||||
|
||||
function createRuntimeHarness(initialState?: GoalModeState) {
|
||||
let state = cloneState(initialState);
|
||||
const runtime = new GoalRuntime({
|
||||
getState: () => cloneState(state),
|
||||
setState: next => {
|
||||
state = cloneState(next);
|
||||
},
|
||||
getCurrentUsage: () => createUsage(),
|
||||
emit: async () => {},
|
||||
persist: (_mode, _state) => {},
|
||||
sendHiddenMessage: async _message => {},
|
||||
now: () => 0,
|
||||
});
|
||||
return {
|
||||
runtime,
|
||||
getState: () => cloneState(state),
|
||||
};
|
||||
}
|
||||
|
||||
describe("GoalTool", () => {
|
||||
it("routes create/get/complete operations and returns completion budget details", async () => {
|
||||
const createGoalState: GoalModeState = {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ objective: "Create route", tokenBudget: 10 }),
|
||||
};
|
||||
const getGoalState: GoalModeState = {
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ objective: "Get route", tokensUsed: 4, tokenBudget: 10 }),
|
||||
};
|
||||
const completedGoal = createGoal({
|
||||
objective: "Complete route",
|
||||
status: "complete",
|
||||
tokensUsed: 7,
|
||||
timeUsedSeconds: 3,
|
||||
tokenBudget: 10,
|
||||
});
|
||||
const runtime = {
|
||||
createGoal: vi.fn(async () => createGoalState),
|
||||
completeGoalFromTool: vi.fn(async () => completedGoal),
|
||||
};
|
||||
const getGoalModeState = vi.fn(() => getGoalState);
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => runtime as unknown as GoalRuntime,
|
||||
getGoalModeState,
|
||||
}),
|
||||
);
|
||||
|
||||
const created = await tool.execute("call-create", {
|
||||
op: "create",
|
||||
objective: " Create route ",
|
||||
token_budget: 10,
|
||||
});
|
||||
expect(runtime.createGoal).toHaveBeenCalledWith({ objective: "Create route", tokenBudget: 10 });
|
||||
expect(created.details).toMatchObject({
|
||||
op: "create",
|
||||
goal: createGoalState.goal,
|
||||
remainingTokens: 10,
|
||||
completionBudgetReport: null,
|
||||
});
|
||||
|
||||
const fetched = await tool.execute("call-get", { op: "get" });
|
||||
expect(getGoalModeState).toHaveBeenCalledTimes(1);
|
||||
expect(fetched.details).toMatchObject({
|
||||
op: "get",
|
||||
goal: getGoalState.goal,
|
||||
remainingTokens: 6,
|
||||
completionBudgetReport: null,
|
||||
});
|
||||
expect(runtime.completeGoalFromTool).not.toHaveBeenCalled();
|
||||
|
||||
const completed = await tool.execute("call-complete", { op: "complete" });
|
||||
expect(runtime.completeGoalFromTool).toHaveBeenCalledTimes(1);
|
||||
expect(completed.details).toMatchObject({
|
||||
op: "complete",
|
||||
goal: completedGoal,
|
||||
remainingTokens: 3,
|
||||
completionBudgetReport: completionBudgetReport(completedGoal),
|
||||
});
|
||||
expect(completed.content[0]).toEqual({
|
||||
type: "text",
|
||||
text: "Goal: Complete route\nStatus: complete\nTokens: 7 used / 10 budget\nRemaining tokens: 3\n\nGoal achieved. Report final budget usage to the user: tokens used: 7 of 10; time used: 3 seconds.",
|
||||
});
|
||||
});
|
||||
|
||||
it("rejects create when a goal already exists", async () => {
|
||||
const harness = createRuntimeHarness({
|
||||
enabled: true,
|
||||
mode: "active",
|
||||
goal: createGoal({ objective: "Existing" }),
|
||||
});
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => harness.runtime,
|
||||
getGoalModeState: () => harness.getState(),
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(
|
||||
tool.execute("call-create", { op: "create", objective: "New goal", token_budget: 10 }),
|
||||
).rejects.toThrow("cannot create a new goal because this session already has a goal");
|
||||
});
|
||||
|
||||
it("rejects complete when no goal is active", async () => {
|
||||
const harness = createRuntimeHarness();
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => harness.runtime,
|
||||
getGoalModeState: () => harness.getState(),
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(tool.execute("call-complete", { op: "complete" })).rejects.toThrow(
|
||||
"cannot complete goal because goal mode is not active",
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects op=create when the objective is missing or only whitespace", async () => {
|
||||
const harness = createRuntimeHarness();
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => harness.runtime,
|
||||
getGoalModeState: () => harness.getState(),
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(tool.execute("call-empty", { op: "create", objective: " \t\n" })).rejects.toThrow(
|
||||
"objective is required when op=create",
|
||||
);
|
||||
expect(harness.getState()).toBeUndefined();
|
||||
});
|
||||
|
||||
it("rejects op=create when the token_budget is zero or negative", async () => {
|
||||
const harness = createRuntimeHarness();
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => harness.runtime,
|
||||
getGoalModeState: () => harness.getState(),
|
||||
}),
|
||||
);
|
||||
|
||||
await expect(tool.execute("call-zero", { op: "create", objective: "Ship it", token_budget: 0 })).rejects.toThrow(
|
||||
"token_budget must be a positive integer when provided",
|
||||
);
|
||||
await expect(tool.execute("call-neg", { op: "create", objective: "Ship it", token_budget: -5 })).rejects.toThrow(
|
||||
"token_budget must be a positive integer when provided",
|
||||
);
|
||||
expect(harness.getState()).toBeUndefined();
|
||||
});
|
||||
|
||||
it("flips state to exiting and clears enabled when op=complete succeeds (fix #1)", async () => {
|
||||
const harness = createRuntimeHarness();
|
||||
await harness.runtime.createGoal({ objective: "Ship the release", tokenBudget: 100 });
|
||||
const tool = new GoalTool(
|
||||
createToolSession({
|
||||
getGoalRuntime: () => harness.runtime,
|
||||
getGoalModeState: () => harness.getState(),
|
||||
}),
|
||||
);
|
||||
|
||||
const result = await tool.execute("call-complete", { op: "complete" });
|
||||
|
||||
expect(result.details).toMatchObject({ op: "complete" });
|
||||
const after = harness.getState();
|
||||
expect(after?.enabled).toBe(false);
|
||||
expect(after?.mode).toBe("exiting");
|
||||
expect(after?.reason).toBe("completed");
|
||||
expect(after?.goal.status).toBe("complete");
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user