import { createHash } from "node:crypto";
import { homedir } from "node:os";
import { join } from "node:path";

import type { AgentToolResult, ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";

import { parseGoalCommand } from "./goal/command.js";
import { shouldQueueGoalContinuationAfterAgentEnd, shouldQueueGoalContinuationWhenIdle } from "./goal/continuation.js";
import { formatGoalForTool, formatGoalToolResponse, goalStatusLabel } from "./goal/format.js";
import { buildBudgetLimitedPrompt, buildContinuationPrompt } from "./goal/prompt.js";
import { GoalAlreadyExistsError } from "./goal/errors.js";
import { accountGoalUsage, clearGoal, createGoal, readGoal, updateGoal } from "./goal/store.js";
import type { Goal, GoalAccountingMode, GoalStoreRef, TokenUsageSnapshot } from "./goal/types.js";
import { COMPLETABLE_GOAL_STATUS_VALUES, isRecord } from "./goal/types.js";
import { updateGoalUi } from "./goal/ui.js";

const GOAL_USAGE = "Usage: /goal <objective>";
const GOAL_EMPTY_HINT = "No goal is currently set.";
const GOAL_CONTINUATION_MESSAGE_TYPE = "pi-goal-continuation";
const GOAL_BUDGET_LIMIT_MESSAGE_TYPE = "pi-goal-budget-limit";
const REPLACE_GOAL_CHOICE = "Replace current goal";
const CANCEL_REPLACE_GOAL_CHOICE = "Cancel";
const RESUME_GOAL_CHOICE = "Resume goal";
const LEAVE_GOAL_PAUSED_CHOICE = "Leave paused";
const EMPTY_USAGE: TokenUsageSnapshot = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 };
const STALE_EXTENSION_CONTEXT_ERROR_PREFIX = "This extension ctx is stale after session replacement or reload.";

type GoalToolResult = AgentToolResult<Record<string, never>> & { isError?: boolean };
type AssistantUsageMessage = {
	role: "assistant";
	usage: Record<string, unknown>;
};
type AgentGoalAccounting = {
	goalId: string;
	measuredFromMilliseconds: number;
};

export default function (pi: ExtensionAPI): void {
	let agentTurnInProgress = false;
	let agentGoalAccounting: AgentGoalAccounting | null = null;
	let completedThisTurnGoalId: string | null = null;
	let budgetLimitReportedGoalId: string | null = null;

	pi.registerTool({
		name: "create_goal",
		label: "Create Goal",
		description:
			"Create a goal only when explicitly requested by the user or system/developer instructions; do not infer goals from ordinary tasks.\nSet token_budget only when an explicit token budget is requested. Fails if an unfinished goal exists; use update_goal only for status.",
		parameters: Type.Object(
			{
				objective: Type.String({
					description:
						"Required. The concrete objective to start pursuing. This starts a new active goal when no goal exists or replaces the current goal when it is complete.",
				}),
				token_budget: Type.Optional(
					Type.Integer({ description: "Positive token budget for the new goal. Omit unless explicitly requested." }),
				),
			},
			{ additionalProperties: false },
		),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			const ref = goalStoreRef(ctx);
			let goal: Goal;
			try {
				goal = await createGoal(ref, params.objective, params.token_budget);
			} catch (error) {
				if (error instanceof GoalAlreadyExistsError) return toolText(error.message, true);
				throw error;
			}
			beginAgentGoalAccounting(goal);
			updateGoalUi(ctx, goal);
			return toolText(formatGoalToolResponse(goal, false));
		},
	});

	pi.registerTool({
		name: "update_goal",
		label: "Update Goal",
		description:
			"Update the existing goal.\nUse this tool only to mark the goal achieved or genuinely blocked.\nSet status to `complete` only when the objective has actually been achieved and no required work remains.\nSet status to `blocked` only when the same blocking condition has repeated for at least three consecutive goal turns, counting the original/user-triggered turn and any automatic continuations, and the agent cannot make meaningful progress without user input or an external-state change.\nIf the user resumes a goal that was previously marked `blocked`, treat the resumed run as a fresh blocked audit. If the same blocking condition then repeats for at least three consecutive resumed goal turns, set status to `blocked` again.\nOnce the blocked threshold is satisfied, do not keep reporting that you are still blocked while leaving the goal active; set status to `blocked`.\nDo not use `blocked` merely because the work is hard, slow, uncertain, incomplete, or would benefit from clarification.\nDo not mark a goal complete merely because its budget is nearly exhausted or because you are stopping work.\nYou cannot use this tool to pause, resume, budget-limit, or usage-limit a goal; those status changes are controlled by the user or system.\nWhen marking a budgeted goal achieved with status `complete`, report the final token usage from the tool result to the user.",
		parameters: Type.Object(
			{
				status: Type.Union(
					COMPLETABLE_GOAL_STATUS_VALUES.map((status) => Type.Literal(status)),
					{
						description:
							"Required. Set to `complete` only when the objective is achieved and no required work remains. Set to `blocked` only after the same blocking condition has recurred for at least three consecutive goal turns and the agent is at an impasse. After a previously blocked goal is resumed, the resumed run starts a fresh blocked audit.",
					},
				),
			},
			{ additionalProperties: false },
		),
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			if (params.status !== "complete" && params.status !== "blocked") {
				return toolText(
					"update_goal can only mark the existing goal complete or blocked; pause, resume, budget-limited, and usage-limited status changes are controlled by the user or system",
					true,
				);
			}
			await accountCurrentAgentTurn(ctx, EMPTY_USAGE, "active");
			const goal = await updateGoal(goalStoreRef(ctx), { status: params.status });
			if (params.status === "complete") {
				markGoalCompletedThisTurn(goal);
			} else {
				stopAgentGoalAccounting(goal.id);
			}
			updateGoalUi(ctx, goal);
			return toolText(formatGoalToolResponse(goal, params.status === "complete"));
		},
	});

	pi.registerTool({
		name: "get_goal",
		label: "Get Goal",
		description:
			"Get the current goal for this thread, including status, budgets, token and elapsed-time usage, and remaining token budget.",
		parameters: Type.Object({}, { additionalProperties: false }),
		async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
			const goal = await readGoal(goalStoreRef(ctx));
			updateGoalUi(ctx, goal);
			return toolText(formatGoalToolResponse(goal, false));
		},
	});

	pi.registerCommand("goal", {
		description: "Set, inspect, pause, resume, or clear the persistent goal",
		handler: async (rawArgs, ctx) => {
			const command = parseGoalCommand(rawArgs);
			try {
				switch (command.kind) {
					case "show": {
						const goal = await readGoal(goalStoreRef(ctx));
						updateGoalUi(ctx, goal);
						ctx.ui.notify(
							goal === null ? `${GOAL_USAGE}\n${GOAL_EMPTY_HINT}` : formatGoalForTool(goal),
							goal ? "info" : "warning",
						);
						return;
					}
					case "setObjective": {
						await setGoalObjective(pi, ctx, command.objective);
						return;
					}
					case "setStatus": {
						if (command.status === "paused") {
							await accountCurrentAgentTurn(ctx, EMPTY_USAGE, "active");
						}
						const goal = await updateGoal(goalStoreRef(ctx), { status: command.status });
						if (goal.status === "active") {
							beginAgentGoalAccounting(goal);
						} else {
							stopAgentGoalAccounting(goal.id);
						}
						updateGoalUi(ctx, goal);
						ctx.ui.notify(`Goal ${goalStatusLabel(goal.status)}\n${formatGoalForTool(goal)}`, "info");
						queueGoalContinuation(pi, ctx, goal);
						return;
					}
					case "clear": {
						await accountCurrentAgentTurn(ctx, EMPTY_USAGE, "active");
						const cleared = await clearGoal(goalStoreRef(ctx));
						clearAgentGoalAccounting();
						updateGoalUi(ctx, null);
						ctx.ui.notify(
							cleared ? "Goal cleared" : "No goal to clear\nThis thread does not currently have a goal.",
							cleared ? "info" : "warning",
						);
						return;
					}
				}
			} catch (error) {
				ctx.ui.notify(errorMessage(error), "error");
			}
		},
	});

	pi.on("session_start", async (event, ctx) => {
		const goal = await readGoal(goalStoreRef(ctx));
		if (goal?.status === "active") {
			beginAgentGoalAccounting(goal);
		} else {
			clearAgentGoalAccounting();
		}
		updateGoalUi(ctx, goal);
		if (await maybePromptResumePausedGoal(pi, ctx, event.reason, goal)) {
			return;
		}
		if (shouldQueueGoalContinuationWhenIdle(goal, ctx.isIdle(), ctx.hasPendingMessages())) {
			queueHiddenGoalPrompt(pi, GOAL_CONTINUATION_MESSAGE_TYPE, buildContinuationPrompt(goal));
		}
	});

	pi.on("agent_start", async (_event, ctx) => {
		agentTurnInProgress = true;
		completedThisTurnGoalId = null;
		const goal = await readGoal(goalStoreRef(ctx));
		if (goal?.status === "active") {
			beginAgentGoalAccounting(goal);
		} else {
			agentGoalAccounting = null;
		}
	});

	pi.on("agent_end", async (event, ctx) => {
		const mode: GoalAccountingMode = completedThisTurnGoalId === null ? "active" : "activeOrComplete";
		const goal = await accountCurrentAgentTurn(ctx, collectAssistantUsage(event.messages), mode);
		agentTurnInProgress = false;
		completedThisTurnGoalId = null;
		if (goal?.status === "active") {
			beginAgentGoalAccounting(goal);
		} else {
			clearAgentGoalAccounting();
		}
		updateGoalUiBestEffort(ctx, goal);
		if (goal?.status === "budgetLimited") {
			if (!ctx.hasPendingMessages() && budgetLimitReportedGoalId !== goal.id) {
				budgetLimitReportedGoalId = goal.id;
				queueHiddenGoalPrompt(pi, GOAL_BUDGET_LIMIT_MESSAGE_TYPE, buildBudgetLimitedPrompt(goal));
			}
			return;
		}
		if (
			goal?.status === "active" &&
			shouldQueueGoalContinuationAfterAgentEnd(goal, ctx.hasPendingMessages(), event.messages)
		) {
			queueHiddenGoalPrompt(pi, GOAL_CONTINUATION_MESSAGE_TYPE, buildContinuationPrompt(goal));
		}
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		if (agentGoalAccounting !== null) {
			await accountCurrentAgentTurn(ctx, EMPTY_USAGE, "active");
		}
		clearAgentGoalAccounting();
	});

	async function setGoalObjective(pi: ExtensionAPI, ctx: ExtensionContext, objective: string): Promise<void> {
		const ref = goalStoreRef(ctx);
		const current = await readGoal(ref);
		if (current !== null) {
			const shouldReplace = await confirmReplaceGoal(ctx, objective);
			if (!shouldReplace) return;
		}

		if (current?.status === "active") {
			await accountCurrentAgentTurn(ctx, EMPTY_USAGE, "active");
		}
		const goal = current === null ? await createGoal(ref, objective) : await updateGoal(ref, { objective });
		if (goal.status === "active") beginAgentGoalAccounting(goal);
		updateGoalUi(ctx, goal);
		ctx.ui.notify(`Goal ${goalStatusLabel(goal.status)}\n${formatGoalForTool(goal)}`, "info");
		queueGoalContinuation(pi, ctx, goal);
	}

	async function confirmReplaceGoal(ctx: ExtensionContext, objective: string): Promise<boolean> {
		if (!ctx.hasUI) return true;
		const choice = await ctx.ui.select(`Replace goal?\nNew objective: ${objective}`, [
			REPLACE_GOAL_CHOICE,
			CANCEL_REPLACE_GOAL_CHOICE,
		]);
		return choice === REPLACE_GOAL_CHOICE;
	}

	async function maybePromptResumePausedGoal(
		pi: ExtensionAPI,
		ctx: ExtensionContext,
		sessionStartReason: string,
		goal: Goal | null,
	): Promise<boolean> {
		if (!isResumeOfPausedGoal(ctx, sessionStartReason, goal)) {
			return false;
		}

		const choice = await ctx.ui.select(`Resume paused goal?\nGoal: ${goal.objective}`, [
			RESUME_GOAL_CHOICE,
			LEAVE_GOAL_PAUSED_CHOICE,
		]);
		if (choice !== RESUME_GOAL_CHOICE) return true;

		const resumed = await updateGoal(goalStoreRef(ctx), { status: "active" });
		beginAgentGoalAccounting(resumed);
		updateGoalUi(ctx, resumed);
		ctx.ui.notify(`Goal ${goalStatusLabel(resumed.status)}\n${formatGoalForTool(resumed)}`, "info");
		queueGoalContinuation(pi, ctx, resumed);
		return true;
	}

	function beginAgentGoalAccounting(goal: Goal): void {
		if (goal.status !== "active") return;
		if (agentGoalAccounting?.goalId === goal.id) return;
		agentGoalAccounting = { goalId: goal.id, measuredFromMilliseconds: Date.now() };
	}

	function markGoalCompletedThisTurn(goal: Goal): void {
		if (!agentTurnInProgress) return;
		completedThisTurnGoalId = goal.id;
		agentGoalAccounting = { goalId: goal.id, measuredFromMilliseconds: Date.now() };
	}

	function stopAgentGoalAccounting(goalId: string): void {
		if (agentGoalAccounting?.goalId === goalId) {
			agentGoalAccounting = null;
		}
		if (completedThisTurnGoalId === goalId) {
			completedThisTurnGoalId = null;
		}
	}

	function clearAgentGoalAccounting(): void {
		agentGoalAccounting = null;
		completedThisTurnGoalId = null;
	}

	async function accountCurrentAgentTurn(
		ctx: ExtensionContext,
		usage: TokenUsageSnapshot,
		mode: GoalAccountingMode,
	): Promise<Goal | null> {
		const accounting = agentGoalAccounting;
		const ref = goalStoreRef(ctx);
		if (accounting === null) return readGoal(ref);

		const now = Date.now();
		const elapsedSeconds = Math.max(0, Math.round((now - accounting.measuredFromMilliseconds) / 1000));
		const goal = await accountGoalUsage(ref, usage, elapsedSeconds, mode, accounting.goalId);
		if (goal?.id === accounting.goalId) {
			agentGoalAccounting = { goalId: accounting.goalId, measuredFromMilliseconds: now };
		} else {
			clearAgentGoalAccounting();
		}
		return goal;
	}
}

function updateGoalUiBestEffort(ctx: ExtensionContext, goal: Goal | null): void {
	try {
		updateGoalUi(ctx, goal);
	} catch (error) {
		if (error instanceof Error && error.message.startsWith(STALE_EXTENSION_CONTEXT_ERROR_PREFIX)) {
			return;
		}
		throw error;
	}
}

function isResumeOfPausedGoal(ctx: ExtensionContext, sessionStartReason: string, goal: Goal | null): goal is Goal {
	return (
		sessionStartReason === "resume" &&
		goal?.status === "paused" &&
		ctx.hasUI &&
		ctx.isIdle() &&
		!ctx.hasPendingMessages()
	);
}

function queueGoalContinuation(pi: ExtensionAPI, ctx: ExtensionContext, goal: Goal): void {
	if (shouldQueueGoalContinuationWhenIdle(goal, ctx.isIdle(), ctx.hasPendingMessages())) {
		queueHiddenGoalPrompt(pi, GOAL_CONTINUATION_MESSAGE_TYPE, buildContinuationPrompt(goal));
	}
}

function queueHiddenGoalPrompt(pi: ExtensionAPI, customType: string, content: string): void {
	pi.sendMessage({ customType, content, display: false }, { triggerTurn: true, deliverAs: "followUp" });
}

function goalStoreRef(ctx: ExtensionContext): GoalStoreRef {
	const sessionFile = ctx.sessionManager.getSessionFile();
	const baseDir =
		sessionFile === undefined
			? join(agentDir(), "extensions", "pi-goal", "no-session", cwdStoreKey(ctx.cwd))
			: join(ctx.sessionManager.getSessionDir(), "extensions", "pi-goal");

	return {
		baseDir,
		threadId: ctx.sessionManager.getSessionId(),
	};
}

function agentDir(): string {
	return process.env["PI_CODING_AGENT_DIR"] ?? join(homedir(), ".pi", "agent");
}

function cwdStoreKey(cwd: string): string {
	return createHash("sha256").update(cwd).digest("hex").slice(0, 24);
}

function toolText(text: string, isError = false): GoalToolResult {
	return { content: [{ type: "text" as const, text }], details: {}, isError };
}

function collectAssistantUsage(messages: unknown[]): TokenUsageSnapshot {
	const usage: TokenUsageSnapshot = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0 };
	for (const message of messages) {
		if (!isAssistantUsageMessage(message)) continue;
		usage.input += numericUsageField(message.usage, "input");
		usage.output += numericUsageField(message.usage, "output");
		usage.cacheRead += numericUsageField(message.usage, "cacheRead");
		usage.cacheWrite += numericUsageField(message.usage, "cacheWrite");
		usage.totalTokens += numericUsageField(message.usage, "totalTokens");
	}
	return usage;
}

function isAssistantUsageMessage(message: unknown): message is AssistantUsageMessage {
	if (!isRecord(message)) return false;
	return message["role"] === "assistant" && isRecord(message["usage"]);
}

function numericUsageField(usage: Record<string, unknown>, key: string): number {
	const value = usage[key];
	return typeof value === "number" && Number.isFinite(value) ? value : 0;
}

function errorMessage(error: unknown): string {
	return error instanceof Error ? error.message : String(error);
}

import { mkdtemp, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { afterEach, describe, expect, it } from "bun:test";

import { accountGoalUsage, createGoal, readGoal, updateGoal } from "../src/goal/store.js";
import type { TokenUsageSnapshot } from "../src/goal/types.js";
import type { GoalStoreRef } from "../src/goal/types.js";
import { GOAL_STATUS_VALUES } from "../src/goal/types.js";

const tempDirs: string[] = [];

async function tempStore(threadId: string): Promise<GoalStoreRef> {
	const dir = await mkdtemp(join(tmpdir(), "pi-goal-codex-"));
	tempDirs.push(dir);
	return { baseDir: dir, threadId };
}

describe("codex alignment: blocked status", () => {
	afterEach(async () => {
		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
	});

	it("includes blocked in the goal status vocabulary", () => {
		expect(GOAL_STATUS_VALUES).toContain("blocked");
	});

	it("marks an active goal blocked and clears its lastStartedAt", async () => {
		const ref = await tempStore("thread-block");
		await createGoal(ref, "Pursue the objective");

		const blocked = await updateGoal(ref, { status: "blocked" });

		expect(blocked.status).toBe("blocked");
		expect(blocked.lastStartedAt).toBeUndefined();
		expect((await readGoal(ref))?.status).toBe("blocked");
	});

	it("preserves budgetLimited when a blocked update is requested", async () => {
		const ref = await tempStore("thread-block-budget");
		await createGoal(ref, "Budget goal", 10);
		const overBudget: TokenUsageSnapshot = { input: 20, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 20 };
		const limited = await accountGoalUsage(ref, overBudget, 0, "active");
		expect(limited?.status).toBe("budgetLimited");

		const afterBlock = await updateGoal(ref, { status: "blocked" });

		expect(afterBlock.status).toBe("budgetLimited");
	});
});

describe("codex alignment: create replaces only a complete goal", () => {
	afterEach(async () => {
		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
	});

	it("replaces a completed goal with a fresh active goal", async () => {
		const ref = await tempStore("thread-replace-complete");
		const first = await createGoal(ref, "First objective");
		await updateGoal(ref, { status: "complete" });

		const replacement = await createGoal(ref, "Second objective");

		expect(replacement.status).toBe("active");
		expect(replacement.objective).toBe("Second objective");
		expect(replacement.id).not.toBe(first.id);
		expect(replacement.tokensUsed).toBe(0);
	});

	it("rejects creating a goal while an unfinished goal exists with the codex message", async () => {
		const ref = await tempStore("thread-unfinished");
		await createGoal(ref, "Active objective");

		await expect(createGoal(ref, "Another objective")).rejects.toThrow(
			"cannot create a new goal because this thread has an unfinished goal; complete the existing goal first",
		);
	});

	it("rejects creating a goal while a blocked goal exists", async () => {
		const ref = await tempStore("thread-blocked-exists");
		await createGoal(ref, "Blocked objective");
		await updateGoal(ref, { status: "blocked" });

		await expect(createGoal(ref, "New objective")).rejects.toThrow("has an unfinished goal");
	});
});

import { describe, expect, it } from "bun:test";

import { parseGoalCommand } from "../src/goal/command.js";

describe("goal command parsing", () => {
	it("treats bare /goal as a summary request", () => {
		expect(parseGoalCommand("")).toEqual({ kind: "show" });
	});

	it("treats arbitrary text after /goal as the objective", () => {
		expect(parseGoalCommand("ship the Codex style flow --token-budget 88")).toEqual({
			kind: "setObjective",
			objective: "ship the Codex style flow --token-budget 88",
		});
	});

	it("does not require or special-case a set subcommand", () => {
		expect(parseGoalCommand("set up the release")).toEqual({
			kind: "setObjective",
			objective: "set up the release",
		});
	});

	it("keeps Codex-style control commands reserved", () => {
		expect(parseGoalCommand("pause")).toEqual({ kind: "setStatus", status: "paused" });
		expect(parseGoalCommand("resume")).toEqual({ kind: "setStatus", status: "active" });
		expect(parseGoalCommand("clear")).toEqual({ kind: "clear" });
	});

	it("treats non-Codex control words as objectives", () => {
		expect(parseGoalCommand("status")).toEqual({ kind: "setObjective", objective: "status" });
		expect(parseGoalCommand("complete")).toEqual({ kind: "setObjective", objective: "complete" });
		expect(parseGoalCommand("help")).toEqual({ kind: "setObjective", objective: "help" });
	});
});

import { describe, expect, it } from "bun:test";

import {
	shouldQueueGoalContinuationAfterAgentEnd,
	shouldQueueGoalContinuationWhenIdle,
} from "../src/goal/continuation.js";
import type { Goal } from "../src/goal/types.js";

const cleanTurn = [{ role: "assistant", stopReason: "stop", content: [{ type: "text", text: "done" }] }];

describe("goal continuation policy", () => {
	it("continues an active goal after a clean agent turn when no user work is pending", () => {
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "active" }), false, cleanTurn)).toBe(true);
	});

	it("does not continue after an agent turn when another message is already pending", () => {
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "active" }), true, cleanTurn)).toBe(false);
	});

	it("only auto-continues active goals after an agent turn", () => {
		expect(shouldQueueGoalContinuationAfterAgentEnd(null, false, cleanTurn)).toBe(false);
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "paused" }), false, cleanTurn)).toBe(false);
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "budgetLimited" }), false, cleanTurn)).toBe(
			false,
		);
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "complete" }), false, cleanTurn)).toBe(false);
	});

	it("does not continue after a turn that ended with a provider error", () => {
		const erroredTurn = [
			{ role: "assistant", stopReason: "error", errorMessage: "boom", content: [] },
		];
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "active" }), false, erroredTurn)).toBe(false);
	});

	it("does not continue after a turn whose last tool result was aborted", () => {
		const abortedTurn = [
			{ role: "assistant", stopReason: "toolUse", content: [{ type: "toolCall", id: "1", name: "bash" }] },
			{ role: "toolResult", isError: true, content: [{ type: "text", text: "Aborted by user" }] },
		];
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "active" }), false, abortedTurn)).toBe(false);
	});

	it("continues after a turn whose tool failed for a non-abort reason", () => {
		const toolFailureTurn = [
			{ role: "assistant", stopReason: "toolUse", content: [{ type: "toolCall", id: "1", name: "bash" }] },
			{ role: "toolResult", isError: true, content: [{ type: "text", text: "command not found" }] },
			{ role: "assistant", stopReason: "stop", content: [{ type: "text", text: "recovered" }] },
		];
		expect(shouldQueueGoalContinuationAfterAgentEnd(testGoal({ status: "active" }), false, toolFailureTurn)).toBe(
			true,
		);
	});

	it("requires idle state for command and session-start continuation", () => {
		expect(shouldQueueGoalContinuationWhenIdle(testGoal({ status: "active" }), true, false)).toBe(true);
		expect(shouldQueueGoalContinuationWhenIdle(testGoal({ status: "active" }), false, false)).toBe(false);
		expect(shouldQueueGoalContinuationWhenIdle(testGoal({ status: "active" }), true, true)).toBe(false);
	});
});

function testGoal(overrides: Partial<Goal> = {}): Goal {
	return {
		id: "goal-1",
		threadId: "thread-1",
		objective: "Keep going until complete",
		status: "active",
		tokensUsed: 0,
		timeUsedSeconds: 0,
		createdAt: 1_777_766_400,
		updatedAt: 1_777_766_400,
		...overrides,
	};
}

import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { dirname, join } from "node:path";

import type { AgentToolResult, ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { afterEach, describe, expect, it, setSystemTime } from "bun:test";
import { goalFilePath, readGoal } from "../src/goal/store.js";
import type { GoalStoreRef } from "../src/goal/types.js";
import { goalFooterIndicator } from "../src/goal/ui.js";
import piGoalExtension from "../src/index.js";

type ToolResult = AgentToolResult<unknown>;

type GoalContext = {
	hasUI: boolean;
	ui: MockUi;
	cwd: string;
	sessionManager: {
		getSessionFile(): string;
		getSessionDir(): string;
		getSessionId(): string;
	};
	isIdle(): boolean;
	hasPendingMessages(): boolean;
};

type RegisteredTool = {
	name: string;
	description: string;
	parameters: unknown;
	execute(
		toolCallId: string,
		params: Record<string, unknown>,
		signal: AbortSignal | undefined,
		onUpdate: undefined,
		ctx: GoalContext,
	): Promise<ToolResult>;
};

type RegisteredCommand = {
	handler(args: string, ctx: GoalContext): Promise<void>;
};

type EventPayload = {
	type: string;
	reason?: string;
	messages?: unknown[];
};

type EventHandler = (event: EventPayload, ctx: GoalContext) => unknown | Promise<unknown>;

type NotifyType = "info" | "warning" | "error";
type SelectCall = { title: string; options: string[] };
type ConfirmCall = { title: string; message: string };
type NotifyCall = { message: string; type: NotifyType | undefined };
type MockUi = {
	selectCalls: SelectCall[];
	confirmCalls: ConfirmCall[];
	notifyCalls: NotifyCall[];
	select(title: string, options: string[]): Promise<string | undefined>;
	confirm(title: string, message: string): Promise<boolean>;
	notify(message: string, type?: NotifyType): void;
	setWidget(key: string, content: string[] | undefined): void;
	setStatus(key: string, text: string | undefined): void;
	setFooter(factory: unknown): void;
};
type SentMessage = {
	message: { customType: string; content: string; display: boolean };
	options: Record<string, unknown>;
};

const tempDirs: string[] = [];

let mockedClockMilliseconds: number | null = null;

// Bun's setSystemTime treats an epoch-0 Date as "reset to real time", so the
// fake clock is anchored at 2020-01-01T00:00:00Z and advanced from there.
const CLOCK_EPOCH_MILLISECONDS = 1_577_836_800_000;

function setClock(milliseconds: number): void {
	mockedClockMilliseconds = milliseconds;
	setSystemTime(new Date(CLOCK_EPOCH_MILLISECONDS + milliseconds));
}

function advanceClock(milliseconds: number): void {
	setClock((mockedClockMilliseconds ?? 0) + milliseconds);
}

function resetClock(): void {
	mockedClockMilliseconds = null;
	setSystemTime();
}

describe("pi-goal extension tool contract", () => {
	it("exposes the Codex goal tools with matching descriptions and schemas", () => {
		const harness = createHarness();

		expect(toolContract(harness.tool("get_goal"))).toEqual({
			name: "get_goal",
			description:
				"Get the current goal for this thread, including status, budgets, token and elapsed-time usage, and remaining token budget.",
			parameters: {
				type: "object",
				properties: {},
				additionalProperties: false,
			},
		});
		expect(toolContract(harness.tool("create_goal"))).toEqual({
			name: "create_goal",
			description:
				"Create a goal only when explicitly requested by the user or system/developer instructions; do not infer goals from ordinary tasks.\nSet token_budget only when an explicit token budget is requested. Fails if an unfinished goal exists; use update_goal only for status.",
			parameters: {
				type: "object",
				required: ["objective"],
				properties: {
					objective: {
						type: "string",
						description:
							"Required. The concrete objective to start pursuing. This starts a new active goal when no goal exists or replaces the current goal when it is complete.",
					},
					token_budget: {
						type: "integer",
						description: "Positive token budget for the new goal. Omit unless explicitly requested.",
					},
				},
				additionalProperties: false,
			},
		});
		expect(toolContract(harness.tool("update_goal"))).toEqual({
			name: "update_goal",
			description:
				"Update the existing goal.\nUse this tool only to mark the goal achieved or genuinely blocked.\nSet status to `complete` only when the objective has actually been achieved and no required work remains.\nSet status to `blocked` only when the same blocking condition has repeated for at least three consecutive goal turns, counting the original/user-triggered turn and any automatic continuations, and the agent cannot make meaningful progress without user input or an external-state change.\nIf the user resumes a goal that was previously marked `blocked`, treat the resumed run as a fresh blocked audit. If the same blocking condition then repeats for at least three consecutive resumed goal turns, set status to `blocked` again.\nOnce the blocked threshold is satisfied, do not keep reporting that you are still blocked while leaving the goal active; set status to `blocked`.\nDo not use `blocked` merely because the work is hard, slow, uncertain, incomplete, or would benefit from clarification.\nDo not mark a goal complete merely because its budget is nearly exhausted or because you are stopping work.\nYou cannot use this tool to pause, resume, budget-limit, or usage-limit a goal; those status changes are controlled by the user or system.\nWhen marking a budgeted goal achieved with status `complete`, report the final token usage from the tool result to the user.",
			parameters: {
				type: "object",
				required: ["status"],
				properties: {
					status: {
						anyOf: [
							{ type: "string", const: "complete" },
							{ type: "string", const: "blocked" },
						],
						description:
							"Required. Set to `complete` only when the objective is achieved and no required work remains. Set to `blocked` only after the same blocking condition has recurred for at least three consecutive goal turns and the agent is at an impasse. After a previously blocked goal is resumed, the resumed run starts a fresh blocked audit.",
					},
				},
				additionalProperties: false,
			},
		});
	});
});

describe("pi-goal extension accounting", () => {
	afterEach(async () => {
		resetClock();
		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
	});

	it("starts elapsed-time accounting when a goal is created during an active agent turn", async () => {
		setClock(0);
		const harness = createHarness();
		const ctx = await createContext("thread-create-during-turn");

		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		advanceClock(30_000);
		await harness
			.tool("create_goal")
			.execute("create-goal", { objective: "created after the turn started" }, undefined, undefined, ctx);
		advanceClock(10_000);

		await harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx);

		const goal = await readGoal(refForContext(ctx));
		expect(goal?.timeUsedSeconds).toBe(10);
	});

	it("accounts resumed active goal time from session start without counting offline time", async () => {
		setClock(0);
		const harness = createHarness();
		const ctx = await createContext("thread-resume-active-accounting");

		// given
		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		await harness.tool("create_goal").execute("create-goal", { objective: "Resume work" }, undefined, undefined, ctx);
		advanceClock(20_000);
		await harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx);
		await harness.emit("session_shutdown", { type: "session_shutdown" }, ctx);
		advanceClock(80_000);

		// when
		await harness.emit("session_start", { type: "session_start", reason: "resume" }, ctx);
		advanceClock(7_000);
		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		advanceClock(11_000);
		await harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx);

		// then
		const goal = await readGoal(refForContext(ctx));
		expect(goal?.timeUsedSeconds).toBe(38);
	});

	it("finalizes elapsed time and usage when update_goal completes an active turn", async () => {
		setClock(0);
		const harness = createHarness();
		const ctx = await createContext("thread-complete-during-turn");

		await harness
			.tool("create_goal")
			.execute("create-goal", { objective: "finish in this turn" }, undefined, undefined, ctx);
		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		advanceClock(65_000);

		const completion = await harness
			.tool("update_goal")
			.execute("complete-goal", { status: "complete" }, undefined, undefined, ctx);

		const completedGoal = await readGoal(refForContext(ctx));
		expect(completedGoal?.status).toBe("complete");
		expect(completedGoal?.timeUsedSeconds).toBe(65);
		expect(completedGoal === null ? "" : goalFooterIndicator(completedGoal).text).toBe("Goal achieved (1m)");
		expect(toolResultText(completion)).toContain('"timeUsedSeconds": 65');

		advanceClock(5_000);
		await harness.emit(
			"agent_end",
			{
				type: "agent_end",
				messages: [
					{
						role: "assistant",
						usage: { input: 100, output: 20, cacheRead: 60, cacheWrite: 0, totalTokens: 120 },
					},
				],
			},
			ctx,
		);

		const finalizedGoal = await readGoal(refForContext(ctx));
		expect(finalizedGoal?.tokensUsed).toBe(120);
		expect(finalizedGoal?.timeUsedSeconds).toBe(70);
	});

	it("does not check pending messages after a goal completes", async () => {
		const harness = createHarness();
		const ctx = await createContext("thread-complete-with-stale-pending");

		await harness
			.tool("create_goal")
			.execute("create-goal", { objective: "finish without continuation checks" }, undefined, undefined, ctx);
		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		await harness.tool("update_goal").execute("complete-goal", { status: "complete" }, undefined, undefined, ctx);
		ctx.hasPendingMessages = () => {
			throw new Error("stale pending messages");
		};

		await expect(harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx)).resolves.toBeUndefined();
	});

	it("does not fail completed accounting when the UI ctx is stale", async () => {
		const harness = createHarness();
		const ctx = await createContext("thread-complete-with-stale-ui");

		await harness
			.tool("create_goal")
			.execute("create-goal", { objective: "finish with stale ui" }, undefined, undefined, ctx);
		await harness.emit("agent_start", { type: "agent_start" }, ctx);
		await harness.tool("update_goal").execute("complete-goal", { status: "complete" }, undefined, undefined, ctx);
		Object.defineProperty(ctx, "hasUI", {
			get() {
				throw new Error("This extension ctx is stale after session replacement or reload.");
			},
		});

		await expect(
			harness.emit(
				"agent_end",
				{
					type: "agent_end",
					messages: [
						{
							role: "assistant",
							usage: { input: 100, output: 20, cacheRead: 60, cacheWrite: 0, totalTokens: 120 },
						},
					],
				},
				ctx,
			),
		).resolves.toBeUndefined();

		const goal = await readGoal(refForContext(ctx));
		expect(goal?.tokensUsed).toBe(120);
	});

	it("does not reread goal state during shutdown when no accounting is active", async () => {
		const harness = createHarness();
		const ctx = await createContext("thread-shutdown-no-accounting");
		const filePath = goalFilePath(refForContext(ctx));
		await mkdir(dirname(filePath), { recursive: true });
		await writeFile(filePath, "", "utf8");

		await expect(harness.emit("session_shutdown", { type: "session_shutdown" }, ctx)).resolves.toBeUndefined();
	});

	it("does not touch stale shutdown ctx when no accounting is active", async () => {
		const harness = createHarness();
		const ctx = await createContext("thread-shutdown-stale-ctx");
		Object.defineProperty(ctx, "hasUI", {
			get() {
				throw new Error("stale ctx");
			},
		});

		await expect(harness.emit("session_shutdown", { type: "session_shutdown" }, ctx)).resolves.toBeUndefined();
	});
});

describe("pi-goal extension command UI parity", () => {
	afterEach(async () => {
		resetClock();
		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
	});

	it("shows Codex-style usage text for a bare /goal without a goal", async () => {
		const harness = createHarness();
		const ui = createMockUi();
		const ctx = await createContext("thread-show-no-goal", { hasUI: true, ui });

		await harness.command("goal").handler("", ctx);

		expect(ui.notifyCalls).toContainEqual({
			message: "Usage: /goal <objective>\nNo goal is currently set.",
			type: "warning",
		});
	});

	it("shows Codex-style clear feedback when no goal exists", async () => {
		const harness = createHarness();
		const ui = createMockUi();
		const ctx = await createContext("thread-clear-no-goal", { hasUI: true, ui });

		await harness.command("goal").handler("clear", ctx);

		expect(ui.notifyCalls).toContainEqual({
			message: "No goal to clear\nThis thread does not currently have a goal.",
			type: "warning",
		});
	});

	it("asks with Codex-style choices before replacing an existing goal", async () => {
		const harness = createHarness();
		const ui = createMockUi({ selectResponses: ["Cancel"] });
		const ctx = await createContext("thread-replace-cancel", { hasUI: true, ui });
		await harness.tool("create_goal").execute("create-goal", { objective: "Original" }, undefined, undefined, ctx);

		await harness.command("goal").handler("Replacement", ctx);

		expect(ui.selectCalls).toContainEqual({
			title: "Replace goal?\nNew objective: Replacement",
			options: ["Replace current goal", "Cancel"],
		});
		expect(ui.confirmCalls).toHaveLength(0);
		expect(await readGoal(refForContext(ctx))).toMatchObject({ objective: "Original" });
	});

	it("replaces an existing goal only after the replace choice is selected", async () => {
		const harness = createHarness();
		const ui = createMockUi({ selectResponses: ["Replace current goal"] });
		const ctx = await createContext("thread-replace-confirm", { hasUI: true, ui });
		await harness.tool("create_goal").execute("create-goal", { objective: "Original" }, undefined, undefined, ctx);

		await harness.command("goal").handler("Replacement", ctx);

		expect(await readGoal(refForContext(ctx))).toMatchObject({
			objective: "Replacement",
			status: "active",
			tokensUsed: 0,
			timeUsedSeconds: 0,
		});
		expect(ui.notifyCalls.at(-1)).toMatchObject({
			message: expect.stringContaining("Goal active\nObjective: Replacement"),
			type: "info",
		});
	});

	it("prompts to resume a paused goal when a session is resumed", async () => {
		const harness = createHarness();
		const ui = createMockUi({ selectResponses: ["Resume goal"] });
		const ctx = await createContext("thread-resume-paused", { hasUI: true, ui });
		await harness.tool("create_goal").execute("create-goal", { objective: "Paused work" }, undefined, undefined, ctx);
		await harness.command("goal").handler("pause", ctx);

		await harness.emit("session_start", { type: "session_start", reason: "resume" }, ctx);

		expect(ui.selectCalls).toContainEqual({
			title: "Resume paused goal?\nGoal: Paused work",
			options: ["Resume goal", "Leave paused"],
		});
		expect(await readGoal(refForContext(ctx))).toMatchObject({ objective: "Paused work", status: "active" });
		expect(harness.sentMessages).toHaveLength(1);
		expect(harness.sentMessages[0]?.message.customType).toBe("pi-goal-continuation");
	});

	it("does not prompt to resume a paused goal on non-resume session starts", async () => {
		const harness = createHarness();
		const ui = createMockUi({ selectResponses: ["Resume goal"] });
		const ctx = await createContext("thread-startup-paused", { hasUI: true, ui });
		await harness.tool("create_goal").execute("create-goal", { objective: "Paused work" }, undefined, undefined, ctx);
		await harness.command("goal").handler("pause", ctx);

		await harness.emit("session_start", { type: "session_start", reason: "startup" }, ctx);

		expect(ui.selectCalls).toHaveLength(0);
		expect(await readGoal(refForContext(ctx))).toMatchObject({ objective: "Paused work", status: "paused" });
		expect(harness.sentMessages).toHaveLength(0);
	});

	it("leaves a paused resumed-session goal paused when that choice is selected", async () => {
		const harness = createHarness();
		const ui = createMockUi({ selectResponses: ["Leave paused"] });
		const ctx = await createContext("thread-leave-paused", { hasUI: true, ui });
		await harness.tool("create_goal").execute("create-goal", { objective: "Paused work" }, undefined, undefined, ctx);
		await harness.command("goal").handler("pause", ctx);

		await harness.emit("session_start", { type: "session_start", reason: "resume" }, ctx);

		expect(await readGoal(refForContext(ctx))).toMatchObject({ objective: "Paused work", status: "paused" });
		expect(harness.sentMessages).toHaveLength(0);
	});
});

function createHarness(): {
	tool(name: string): RegisteredTool;
	command(name: string): RegisteredCommand;
	emit(event: string, payload: EventPayload, ctx: GoalContext): Promise<void>;
	sentMessages: SentMessage[];
} {
	const tools = new Map<string, RegisteredTool>();
	const commands = new Map<string, RegisteredCommand>();
	const handlers = new Map<string, EventHandler[]>();
	const sentMessages: SentMessage[] = [];

	piGoalExtension(createExtensionApi(tools, commands, handlers, sentMessages));

	return {
		tool(name) {
			const tool = tools.get(name);
			if (tool === undefined) throw new Error(`tool not registered: ${name}`);
			return tool;
		},
		command(name) {
			const command = commands.get(name);
			if (command === undefined) throw new Error(`command not registered: ${name}`);
			return command;
		},
		async emit(event, payload, ctx) {
			for (const handler of handlers.get(event) ?? []) {
				await handler(payload, ctx);
			}
		},
		sentMessages,
	};
}

function createExtensionApi(
	tools: Map<string, RegisteredTool>,
	commands: Map<string, RegisteredCommand>,
	handlers: Map<string, EventHandler[]>,
	sentMessages: SentMessage[],
): ExtensionAPI {
	return {
		on(event, handler) {
			const eventHandlers = handlers.get(event) ?? [];
			eventHandlers.push((payload, ctx) => handler(payload as never, ctx as never));
			handlers.set(event, eventHandlers);
		},
		registerTool(tool) {
			tools.set(tool.name, {
				name: tool.name,
				description: tool.description,
				parameters: tool.parameters,
				execute(toolCallId, params, signal, onUpdate, ctx) {
					return tool.execute(toolCallId, params as never, signal, onUpdate, ctx as never);
				},
			});
		},
		registerCommand(name, options) {
			commands.set(name, {
				handler(args, ctx) {
					return options.handler(args, ctx as never);
				},
			});
		},
		registerShortcut() {},
		registerFlag() {},
		getFlag() {
			return undefined;
		},
		registerMessageRenderer() {},
		sendMessage(message, options) {
			sentMessages.push({
				message: {
					customType: message.customType,
					content: String(message.content),
					display: message.display,
				},
				options: options ?? {},
			});
		},
		sendUserMessage() {},
		appendEntry() {},
		setSessionName() {},
		getSessionName() {
			return undefined;
		},
		setLabel() {},
		async exec() {
			return { stdout: "", stderr: "", code: 0, killed: false };
		},
		getActiveTools() {
			return [];
		},
		getAllTools() {
			return [];
		},
		setActiveTools() {},
		getCommands() {
			return [];
		},
		async setModel() {
			return false;
		},
		getThinkingLevel() {
			return "medium";
		},
		setThinkingLevel() {},
		registerProvider() {},
		unregisterProvider() {},
		events: {
			emit() {},
			on() {
				return () => {};
			},
		},
	};
}

type ContextOptions = {
	hasUI?: boolean;
	ui?: MockUi;
	isIdle?: boolean;
	hasPendingMessages?: boolean;
};

async function createContext(threadId: string, options: ContextOptions = {}): Promise<GoalContext> {
	const sessionDir = await mkdtemp(join(tmpdir(), "pi-goal-extension-"));
	tempDirs.push(sessionDir);
	return {
		hasUI: options.hasUI ?? false,
		ui: options.ui ?? createMockUi(),
		cwd: sessionDir,
		sessionManager: {
			getSessionFile: () => join(sessionDir, "session.json"),
			getSessionDir: () => sessionDir,
			getSessionId: () => threadId,
		},
		isIdle: () => options.isIdle ?? true,
		hasPendingMessages: () => options.hasPendingMessages ?? false,
	};
}

function createMockUi(
	options: { selectResponses?: (string | undefined)[]; confirmResponses?: boolean[] } = {},
): MockUi {
	const selectResponses = [...(options.selectResponses ?? [])];
	const confirmResponses = [...(options.confirmResponses ?? [])];
	return {
		selectCalls: [],
		confirmCalls: [],
		notifyCalls: [],
		async select(title, choices) {
			this.selectCalls.push({ title, options: choices });
			return selectResponses.shift();
		},
		async confirm(title, message) {
			this.confirmCalls.push({ title, message });
			return confirmResponses.shift() ?? false;
		},
		notify(message, type) {
			this.notifyCalls.push({ message, type });
		},
		setWidget() {},
		setStatus() {},
		setFooter() {},
	};
}

function refForContext(ctx: GoalContext): GoalStoreRef {
	return {
		baseDir: join(ctx.sessionManager.getSessionDir(), "extensions", "pi-goal"),
		threadId: ctx.sessionManager.getSessionId(),
	};
}

function toolResultText(result: ToolResult): string {
	const firstContent = result.content[0];
	if (firstContent?.type !== "text") throw new Error("tool result had no text content");
	return firstContent.text;
}

function toolContract(tool: RegisteredTool): Pick<RegisteredTool, "name" | "description" | "parameters"> {
	return {
		name: tool.name,
		description: tool.description,
		parameters: tool.parameters,
	};
}

import { describe, expect, it } from "bun:test";

import { formatGoalElapsedSeconds, goalToolResponse, goalUsageSummary } from "../src/goal/format.js";
import type { Goal } from "../src/goal/types.js";

describe("goal display formatting", () => {
	it("formats elapsed seconds like Codex TUI", () => {
		expect(formatGoalElapsedSeconds(0)).toBe("0s");
		expect(formatGoalElapsedSeconds(59)).toBe("59s");
		expect(formatGoalElapsedSeconds(60)).toBe("1m");
		expect(formatGoalElapsedSeconds(30 * 60)).toBe("30m");
		expect(formatGoalElapsedSeconds(90 * 60)).toBe("1h 30m");
		expect(formatGoalElapsedSeconds(2 * 60 * 60)).toBe("2h");
		expect(formatGoalElapsedSeconds(24 * 60 * 60 - 1)).toBe("23h 59m");
		expect(formatGoalElapsedSeconds(24 * 60 * 60)).toBe("1d 0h 0m");
		expect(formatGoalElapsedSeconds(2 * 24 * 60 * 60 + 23 * 60 * 60 + 42 * 60)).toBe("2d 23h 42m");
	});

	it("summarizes goal time and budgeted tokens", () => {
		expect(goalUsageSummary(testGoal({ tokenBudget: 50_000, tokensUsed: 63_876 }))).toBe(
			"Objective: Port /goal as a pi extension Time: 2m. Tokens: 63.9K/50K.",
		);
	});

	it("returns Codex-style tool response budget report for completed budgeted goals", () => {
		expect(
			goalToolResponse(
				testGoal({
					status: "complete",
					tokenBudget: 10_000,
					tokensUsed: 3_250,
					timeUsedSeconds: 75,
				}),
				true,
			),
		).toMatchObject({
			goal: {
				threadId: "thread-1",
				status: "complete",
				tokenBudget: 10_000,
				createdAt: 1_777_766_400,
			},
			remainingTokens: 6_750,
			completionBudgetReport:
				"Goal achieved. Report final usage from this tool result's structured goal fields. If `goal.tokenBudget` is present, include token usage from `goal.tokensUsed` and `goal.tokenBudget`. If `goal.timeUsedSeconds` is greater than 0, summarize elapsed time in a concise, human-friendly form appropriate to the response language.",
		});
	});
});

function testGoal(overrides: Partial<Goal> = {}): Goal {
	return {
		id: "goal-1",
		threadId: "thread-1",
		objective: "Port /goal as a pi extension",
		status: "active",
		tokensUsed: 0,
		timeUsedSeconds: 120,
		createdAt: 1_777_766_400,
		updatedAt: 1_777_766_400,
		...overrides,
	};
}

import { describe, expect, it } from "bun:test";

import { buildBudgetLimitedPrompt, buildContinuationPrompt } from "../src/goal/prompt.js";
import type { Goal } from "../src/goal/types.js";

describe("goal prompts", () => {
	it("escapes the continuation objective at its data boundary", () => {
		const prompt = buildContinuationPrompt(testGoal("A & B < C > D", { tokenBudget: 100 }));

		expect(prompt).toContain("<objective>\nA &amp; B &lt; C &gt; D\n</objective>");
		expect(prompt).not.toContain("<untrusted_objective>");
	});

	it("reflects token accounting inputs without pinning their presentation", () => {
		const baseline = buildContinuationPrompt(testGoal("Objective", { tokensUsed: 7 }));
		const changedUsage = buildContinuationPrompt(testGoal("Objective", { tokensUsed: 8 }));
		const bounded = buildContinuationPrompt(testGoal("Objective", { tokensUsed: 7, tokenBudget: 100 }));

		expect(changedUsage).not.toBe(baseline);
		expect(bounded).not.toBe(baseline);
	});

	it("escapes budget-limit objectives and reflects accounting inputs", () => {
		const prompt = buildBudgetLimitedPrompt(
			testGoal("A & B < C > D", { status: "budgetLimited", tokenBudget: 10, tokensUsed: 12 }),
		);
		const changedAccounting = buildBudgetLimitedPrompt(
			testGoal("A & B < C > D", {
				status: "budgetLimited",
				tokenBudget: 11,
				tokensUsed: 13,
				timeUsedSeconds: 21,
			}),
		);

		expect(prompt).toContain("<objective>\nA &amp; B &lt; C &gt; D\n</objective>");
		expect(prompt).not.toContain("<untrusted_objective>");
		expect(changedAccounting).not.toBe(prompt);
	});
});

function testGoal(objective: string, overrides: Partial<Goal> = {}): Goal {
	return {
		id: "goal-1",
		threadId: "thread-1",
		objective,
		status: "active",
		tokensUsed: 10,
		timeUsedSeconds: 20,
		createdAt: 1_777_766_400,
		updatedAt: 1_777_766_400,
		...overrides,
	};
}

import { mkdtemp, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";

import { afterEach, describe, expect, it } from "bun:test";

import { accountGoalUsage, clearGoal, createGoal, goalFilePath, readGoal, updateGoal } from "../src/goal/store.js";
import type { GoalStoreRef } from "../src/goal/types.js";

const tempDirs: string[] = [];

describe("goal store", () => {
	afterEach(async () => {
		await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true })));
	});

	it("creates a persisted active goal", async () => {
		const ref = await tempStore("thread-create");
		const goal = await createGoal(ref, "  Ship the extension  ", 10_000);

		expect(goal.threadId).toBe("thread-create");
		expect(goal.objective).toBe("Ship the extension");
		expect(goal.status).toBe("active");
		expect(goal.tokenBudget).toBe(10_000);
		expect(await readGoal(ref)).toMatchObject({ id: goal.id, objective: "Ship the extension" });
		// Normalize separators so the path assertion holds on Windows (backslashes) too.
		expect(goalFilePath(ref).replaceAll("\\", "/")).toContain("extensions/pi-goal/thread-create.json");
		expect(goalFilePath(ref)).not.toContain(".pi");
		expect(await readFile(goalFilePath(ref), "utf8")).toContain('"version": 1');
	});

	it("does not replace an existing goal when createGoal is called again", async () => {
		const ref = await tempStore("thread-duplicate-create");
		const original = await createGoal(ref, "Original", 10_000);

		await expect(createGoal(ref, "Replacement", 20_000)).rejects.toThrow(
			"cannot create a new goal because this thread has an unfinished goal; complete the existing goal first",
		);

		expect(await readGoal(ref)).toMatchObject({
			id: original.id,
			objective: "Original",
			tokenBudget: 10_000,
		});
	});

	it("replaces changed objectives and preserves usage for status updates", async () => {
		const ref = await tempStore();
		const first = await createGoal(ref, "Original");
		await accountGoalUsage(ref, { input: 23, output: 2, cacheRead: 0, cacheWrite: 4, totalTokens: 25 }, 70);

		const paused = await updateGoal(ref, { status: "paused" });
		expect(paused.id).toBe(first.id);
		expect(paused.tokensUsed).toBe(25);
		expect(paused.timeUsedSeconds).toBe(70);

		const replaced = await updateGoal(ref, { objective: "Replacement" });
		expect(replaced.id).not.toBe(first.id);
		expect(replaced.tokensUsed).toBe(0);
		expect(replaced.timeUsedSeconds).toBe(0);
		expect(replaced.status).toBe("active");
	});

	it("resumes a matching nonterminal goal when the objective is set again", async () => {
		const ref = await tempStore();
		const first = await createGoal(ref, "Same");
		const paused = await updateGoal(ref, { status: "paused" });

		const resumed = await updateGoal(ref, { objective: "Same" });

		expect(paused.id).toBe(first.id);
		expect(resumed.id).toBe(first.id);
		expect(resumed.status).toBe("active");
	});

	it("counts Pi non-cached input plus output tokens like Codex", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Budgeted");

		const goal = await accountGoalUsage(
			ref,
			{ input: 100, output: 20, cacheRead: 70, cacheWrite: 0, totalTokens: 999 },
			0,
		);

		expect(goal).toMatchObject({ tokensUsed: 120 });
	});

	it("marks active goals budgetLimited when accounting reaches budget", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Budgeted", 50);

		const goal = await accountGoalUsage(
			ref,
			{ input: 31, output: 20, cacheRead: 0, cacheWrite: 0, totalTokens: 51 },
			4,
		);

		expect(goal).toMatchObject({ status: "budgetLimited", tokensUsed: 51, timeUsedSeconds: 4 });
	});

	it("continues accounting budget-limited goals for in-flight active usage", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Budgeted", 20);
		await accountGoalUsage(ref, { input: 5, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 5 }, 7);
		await accountGoalUsage(ref, { input: 15, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 15 }, 3);

		const goal = await accountGoalUsage(
			ref,
			{ input: 5, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 5 },
			5,
			"active",
		);

		expect(goal).toMatchObject({ status: "budgetLimited", tokensUsed: 25, timeUsedSeconds: 15 });
	});

	it("keeps budget-limited goals terminal when paused or reactivated over budget", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Budgeted", 20);
		await accountGoalUsage(ref, { input: 25, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 25 }, 1);

		const paused = await updateGoal(ref, { status: "paused" });
		expect(paused).toMatchObject({ status: "budgetLimited", tokensUsed: 25, tokenBudget: 20 });

		const reactivated = await updateGoal(ref, { status: "active" });
		expect(reactivated).toMatchObject({ status: "budgetLimited", tokensUsed: 25, tokenBudget: 20 });
	});

	it("immediately budget-limits active goals when a lowered budget is already exceeded", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Budgeted", 100);
		await accountGoalUsage(ref, { input: 50, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 50 }, 1);

		const lowered = await updateGoal(ref, { tokenBudget: 40 });

		expect(lowered).toMatchObject({ status: "budgetLimited", tokensUsed: 50, tokenBudget: 40 });
	});

	it("can finalize paused in-flight usage and promote stopped goals over budget", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Stopped", 20);
		await updateGoal(ref, { status: "paused" });

		const activeOnly = await accountGoalUsage(
			ref,
			{ input: 25, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 25 },
			3,
			"active",
		);
		expect(activeOnly).toMatchObject({ status: "paused", tokensUsed: 0, timeUsedSeconds: 0 });

		const stopped = await accountGoalUsage(
			ref,
			{ input: 25, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 25 },
			3,
			"activeOrStopped",
		);
		expect(stopped).toMatchObject({ status: "budgetLimited", tokensUsed: 25, timeUsedSeconds: 3 });
	});

	it("clears the store while preserving the versioned file", async () => {
		const ref = await tempStore();
		await createGoal(ref, "Temporary");

		expect(await clearGoal(ref)).toBe(true);
		expect(await readGoal(ref)).toBeNull();
	});
});

async function tempStore(threadId = "thread-test"): Promise<GoalStoreRef> {
	const dir = await mkdtemp(join(tmpdir(), "pi-goal-"));
	tempDirs.push(dir);
	return { baseDir: join(dir, "extensions", "pi-goal"), threadId };
}

import { describe, expect, it } from "bun:test";
import type { Goal } from "../src/goal/types.js";
import { composeFooterStatusLine, goalFooterIndicator } from "../src/goal/ui.js";

describe("goal footer UI", () => {
	it("formats Codex-style goal indicator labels", () => {
		expect(goalFooterIndicator(testGoal()).text).toBe("Pursuing goal (2m)");
		expect(goalFooterIndicator(testGoal({ tokenBudget: 50_000, tokensUsed: 12_500 })).text).toBe(
			"Pursuing goal (12.5K / 50K)",
		);
		expect(goalFooterIndicator(testGoal({ status: "paused" })).text).toBe("Goal paused (/goal resume)");
		expect(
			goalFooterIndicator(testGoal({ status: "budgetLimited", tokenBudget: 50_000, tokensUsed: 63_876 })).text,
		).toBe("Goal unmet (63.9K / 50K tokens)");
		expect(goalFooterIndicator(testGoal({ status: "complete", tokenBudget: 10_000, tokensUsed: 3_250 })).text).toBe(
			"Goal achieved (3.3K tokens)",
		);
	});

	it("right-aligns the goal indicator on the bottom footer line", () => {
		const line = composeFooterStatusLine("", "Pursuing goal (2m)", 32);

		expect(line).toHaveLength(32);
		expect(line.endsWith("Pursuing goal (2m)")).toBe(true);
		expect(line.trimStart()).toBe("Pursuing goal (2m)");
	});

	it("keeps other extension statuses on the left when the goal indicator fits", () => {
		const line = composeFooterStatusLine("review ready", "Goal paused (/goal resume)", 52);

		expect(line).toHaveLength(52);
		expect(line.startsWith("review ready")).toBe(true);
		expect(line.endsWith("Goal paused (/goal resume)")).toBe(true);
	});
});

function testGoal(overrides: Partial<Goal> = {}): Goal {
	return {
		id: "goal-1",
		threadId: "thread-1",
		objective: "Port /goal as a pi extension",
		status: "active",
		tokensUsed: 0,
		timeUsedSeconds: 120,
		createdAt: 1_777_766_400,
		updatedAt: 1_777_766_400,
		...overrides,
	};
}

import { describe, expect, it } from "bun:test";

import { MAX_OBJECTIVE_LENGTH, validateObjective } from "../src/goal/validation.js";

describe("validateObjective", () => {
	it("accepts objective when at Codex character limit", () => {
		const objective = "a".repeat(MAX_OBJECTIVE_LENGTH);

		expect(validateObjective(objective)).toBe(objective);
	});

	it("throws Codex-style file hint when objective exceeds limit", () => {
		const objective = "a".repeat(MAX_OBJECTIVE_LENGTH + 1);

		expect(() => validateObjective(objective)).toThrow("Put longer instructions in a file");
	});
});