think.ts

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Container, Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";

const THINK_TOOL_NAME = "think";

const THINK_GUIDELINES = [
	"Use think as your scratchpad: it is where your reasoning happens, and its content is private rather than part of the answer.",
	"Call think before the first action of a turn, and again before any step that is expensive to undo: an edit, a destructive command, or a final answer.",
	"In think, restate what is actually being asked and the given constraints. Split the task into ordered sub-problems and solve each explicitly, writing intermediate results instead of jumping to the conclusion. Enumerate and resolve every case split.",
	"Before answering, use think to verify each claim against the constraints, test a boundary or degenerate case, and look for the error you most plausibly made. If a check fails, redo that step instead of patching the conclusion.",
	"Call think again only for materially new state: a tool result that changes the plan, a failed check, or a newly discovered sub-problem. Do not use it to narrate progress or repeat reasoning already recorded.",
];

const RESPONSES_APIS = new Set([
	"openai-responses",
	"azure-openai-responses",
	"openai-codex-responses",
]);

const THINKING_EFFORT_GUIDANCE: Record<ThinkingLevel, string> = {
	off: "Keep the think scratchpad to one very short sentence confirming the task and constraints; do not explore alternatives unless required for correctness.",
	minimal:
		"Use a minimal think scratchpad: identify the immediate next step and perform one quick correctness check.",
	low: "Use a brief think scratchpad: outline the approach, resolve the main uncertainty, and check the likely failure point.",
	medium:
		"Use a moderately detailed think scratchpad: decompose the task, work through the important intermediate steps, and verify the result against the constraints.",
	high: "Use a thorough think scratchpad: explicitly derive the solution, compare plausible alternatives, examine edge cases, and independently verify the important conclusions before acting or answering.",
	xhigh:
		"Use a very thorough think scratchpad: deeply decompose the problem, resolve every meaningful case and ambiguity, test boundary conditions, and perform multiple independent checks before acting or answering.",
	max: "Use the most exhaustive think scratchpad warranted by the task: explore all credible approaches and failure modes, derive each important step explicitly, test adversarial and degenerate cases, and repeatedly verify the final plan or answer. Prefer completeness over speed.",
};

type Payload = Record<string, unknown>;
type ThinkingLevel =
	| "off"
	| "minimal"
	| "low"
	| "medium"
	| "high"
	| "xhigh"
	| "max";
type ReasoningModel = {
	id: string;
	api?: string;
	thinkingLevelMap?: Partial<Record<ThinkingLevel, string | null>>;
};

function isPayload(value: unknown): value is Payload {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

function getOffEffort(model: ReasoningModel | undefined): string {
	// GPT-5.6 accepts `none`, including through the Codex Responses endpoint.
	// Pi's current Codex catalog omits an `off` mapping, so do not interpret an
	// absent entry as meaning that reasoning cannot be disabled.
	return model?.thinkingLevelMap?.off ?? "none";
}

function isResponsesModel(model: ReasoningModel | undefined): boolean {
	return model?.api !== undefined && RESPONSES_APIS.has(model.api);
}

function isGpt56PlusResponsesModel(model: ReasoningModel | undefined): boolean {
	if (
		model?.api !== "openai-responses" &&
		model?.api !== "azure-openai-responses"
	)
		return false;
	const match = /(?:^|\/)gpt-(\d+)\.(\d+)/i.exec(model.id);
	if (!match) return false;
	const major = Number(match[1]);
	const minor = Number(match[2]);
	return major > 5 || (major === 5 && minor >= 6);
}

function addThinkingEffortInstruction(
	payload: Payload,
	level: ThinkingLevel | undefined,
): Payload {
	if (!Array.isArray(payload.input)) return payload;
	const instruction = THINKING_EFFORT_GUIDANCE[level ?? "medium"];
	const alreadyPresent = payload.input.some(
		(item) =>
			isPayload(item) &&
			item.role === "developer" &&
			Array.isArray(item.content) &&
			item.content.some(
				(block) =>
					isPayload(block) &&
					block.type === "input_text" &&
					block.text === instruction,
			),
	);
	if (alreadyPresent) return payload;
	return {
		...payload,
		input: [
			...payload.input,
			{
				role: "developer",
				content: [{ type: "input_text", text: instruction }],
			},
		],
	};
}

function payloadAdvertisesThink(payload: Payload): boolean {
	if (!Array.isArray(payload.tools)) return false;
	return payload.tools.some((tool) => {
		if (!isPayload(tool) || tool.type !== "function") return false;
		if (tool.name === THINK_TOOL_NAME) return true;
		return isPayload(tool.function) && tool.function.name === THINK_TOOL_NAME;
	});
}

function forceThinkToolChoice(payload: Payload): Payload {
	// Do not replace an explicit directive from the user or another extension.
	// Codex emits `auto` by default, so that value is safe to replace.
	if (payload.tool_choice !== undefined && payload.tool_choice !== "auto")
		return payload;
	return {
		...payload,
		tool_choice: { type: "function", name: THINK_TOOL_NAME },
	};
}

function removeClearThinkingStrategy(value: unknown): unknown {
	if (!isPayload(value) || !Array.isArray(value.edits)) return value;

	const edits = value.edits.filter(
		(edit) => !isPayload(edit) || edit.type !== "clear_thinking_20251015",
	);
	return edits.length > 0 ? { ...value, edits } : undefined;
}

function disableProviderReasoning(
	payload: unknown,
	model: ReasoningModel | undefined,
): unknown {
	if (!isPayload(payload)) return payload;

	const result: Payload = { ...payload };
	const offEffort = getOffEffort(model);

	// pi-messages keeps stream options nested in the provider payload.
	if (isPayload(result.options) && "reasoning" in result.options) {
		result.options = { ...result.options, reasoning: undefined };
	}

	// OpenAI Responses/OpenRouter/Together use a nested reasoning object.
	if (isPayload(result.reasoning)) {
		result.reasoning =
			"enabled" in result.reasoning
				? { enabled: false }
				: { effort: offEffort };
	}
	if ("reasoning_effort" in result) result.reasoning_effort = offEffort;
	if ("reasoningEffort" in result) delete result.reasoningEffort;

	// Anthropic, DeepSeek, Z.AI, and string-thinking compatible endpoints.
	if (isPayload(result.thinking)) {
		if (model?.thinkingLevelMap?.off === null) {
			// Some Anthropic models (for example Fable) default to adaptive thinking
			// and reject an explicit disabled mode.
			delete result.thinking;
		} else {
			result.thinking = { type: "disabled" };
			// This context-management strategy is only valid while Anthropic
			// thinking is enabled or adaptive.
			result.context_management = removeClearThinkingStrategy(
				result.context_management,
			);
		}
		delete result.output_config;
	} else if (typeof result.thinking === "string") {
		result.thinking = offEffort;
	}
	if ("enable_thinking" in result) result.enable_thinking = false;
	if ("thinking_token_budget" in result) delete result.thinking_token_budget;

	for (const key of ["chat_template_kwargs", "chat_template_args"] as const) {
		if (isPayload(result[key]) && "enable_thinking" in result[key]) {
			result[key] = { ...result[key], enable_thinking: false };
		}
	}

	// Google puts thinking configuration under config.
	if (isPayload(result.config) && "thinkingConfig" in result.config) {
		const modelId = model?.id.toLowerCase() ?? "";
		const thinkingConfig = /gemini-3(?:\.\d+)?-pro/.test(modelId)
			? { thinkingLevel: "LOW" }
			: /gemini-3(?:\.\d+)?-flash/.test(modelId) || /gemma-?4/.test(modelId)
				? { thinkingLevel: "MINIMAL" }
				: { thinkingBudget: 0 };
		result.config = { ...result.config, thinkingConfig };
	}

	// Bedrock carries Anthropic thinking fields in this nested object.
	if (isPayload(result.additionalModelRequestFields)) {
		const additional = { ...result.additionalModelRequestFields };
		delete additional.thinking;
		delete additional.output_config;
		delete additional.anthropic_beta;
		result.additionalModelRequestFields =
			Object.keys(additional).length > 0 ? additional : undefined;
	}

	return result;
}

export default function thinkExtension(pi: ExtensionAPI) {
	let forceThinkOnNextResponsesRequest = false;

	pi.registerTool({
		name: THINK_TOOL_NAME,
		label: "Think",
		description:
			"Use this private scratchpad to plan, derive, or check work before answering. Record only materially new reasoning. The user does not see this tool activity.",
		promptSnippet:
			"Record private intermediate reasoning in think before answering",
		promptGuidelines: THINK_GUIDELINES,
		parameters: Type.Object(
			{
				thoughts: Type.String({
					description:
						"Current reasoning, intermediate results, checks, or unresolved obligations",
				}),
			},
			{ additionalProperties: false },
		),

		async execute() {
			return {
				content: [{ type: "text", text: "------" }],
				details: { recorded: true },
			};
		},

		renderCall(args, theme) {
			const title = theme.fg("toolTitle", theme.bold("think"));
			if (!args.thoughts) return new Text(title, 0, 0);
			const thoughts = theme.fg("thinkingText", theme.italic(args.thoughts));
			return new Text(`${title}\n${thoughts}`, 0, 0);
		},

		renderResult() {
			// The useful content is the scratchpad itself; hide the acknowledgement.
			return new Container();
		},
	});

	// Arm one eager scratchpad call for each user-initiated agent run. The payload
	// hook below applies it only to a Responses request that actually advertises
	// think, so nested requests such as compaction cannot accidentally consume it.
	pi.on("before_agent_start", (_event, ctx) => {
		forceThinkOnNextResponsesRequest =
			pi.getActiveTools().includes(THINK_TOOL_NAME) &&
			isResponsesModel(ctx.model);
	});

	// Keep the selected thinking level intact in session/UI state. Suppress built-in
	// reasoning only in the serialized provider request while think is active.
	pi.on("before_provider_request", (event, ctx) => {
		if (!pi.getActiveTools().includes(THINK_TOOL_NAME)) {
			forceThinkOnNextResponsesRequest = false;
			return;
		}

		const disabledPayload = disableProviderReasoning(event.payload, ctx.model);
		if (!isPayload(disabledPayload)) return disabledPayload;

		const payloadWithEffortGuidance =
			isGpt56PlusResponsesModel(ctx.model) &&
			payloadAdvertisesThink(disabledPayload)
				? addThinkingEffortInstruction(disabledPayload, ctx.thinkingLevel)
				: disabledPayload;
		if (
			forceThinkOnNextResponsesRequest &&
			isResponsesModel(ctx.model) &&
			payloadAdvertisesThink(payloadWithEffortGuidance)
		) {
			forceThinkOnNextResponsesRequest = false;
			return forceThinkToolChoice(payloadWithEffortGuidance);
		}
		return payloadWithEffortGuidance;
	});

	pi.on("agent_end", () => {
		forceThinkOnNextResponsesRequest = false;
	});
}