@aliaksei-raketski/pi-fast-mode 0.1.0 → 0.2.0

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Aliaksei Raketski
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -22,6 +22,12 @@ pi install npm:@aliaksei-raketski/pi-fast-mode
 pi install -l npm:@aliaksei-raketski/pi-fast-mode
 ```
 
+Try locally from this repository:
+
+```bash
+pi -e ./packages/fast-mode
+```
+
 ## Behavior
 
 - Start Pi with fast mode enabled:
@@ -39,4 +45,6 @@ When fast mode is enabled but the current model is not supported, status stays a
 
 - `fast on` in gray (so you can see it is enabled, but inactive for current model)
 
-When you switch to a supported model, fast mode is applied automatically if it is enabled.
+When you switch to a supported model, fast mode is applied automatically if it is enabled.
+
+The fast mode toggle is stored in the current session, so it survives `/reload`, resume, and branch navigation.

package/package.json

@@ -1,40 +1,41 @@
 {
-  "name": "@aliaksei-raketski/pi-fast-mode",
-  "version": "0.1.0",
-  "description": "Pi extension that enables fast-mode payload tuning for supported Claude/OpenAI models.",
-  "keywords": [
-    "pi-package",
-    "pi",
-    "pi-extension",
-    "fast-mode"
-  ],
-  "license": "MIT",
-  "homepage": "https://github.com/aliaksei-raketski/pi-packages/tree/main/fast-mode",
-  "bugs": {
-    "url": "https://github.com/aliaksei-raketski/pi-packages/issues"
-  },
-  "repository": {
-    "type": "git",
-    "url": "git+https://github.com/aliaksei-raketski/pi-packages.git",
-    "directory": "fast-mode"
-  },
-  "publishConfig": {
-    "access": "public",
-    "registry": "https://registry.npmjs.org/"
-  },
-  "engines": {
-    "node": ">=18"
-  },
-  "peerDependencies": {
-    "@earendil-works/pi-coding-agent": "*"
-  },
-  "pi": {
-    "extensions": [
-      "./index.ts"
-    ]
-  },
-  "files": [
-    "index.ts",
-    "README.md"
-  ]
+	"name": "@aliaksei-raketski/pi-fast-mode",
+	"version": "0.2.0",
+	"description": "Pi extension that enables fast-mode payload tuning for supported Claude/OpenAI models.",
+	"keywords": [
+		"pi-package",
+		"pi",
+		"pi-extension",
+		"fast-mode"
+	],
+	"license": "MIT",
+	"homepage": "https://github.com/aliaksei-raketski/pi-packages/tree/main/packages/fast-mode",
+	"bugs": {
+		"url": "https://github.com/aliaksei-raketski/pi-packages/issues"
+	},
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/aliaksei-raketski/pi-packages.git",
+		"directory": "packages/fast-mode"
+	},
+	"publishConfig": {
+		"access": "public",
+		"registry": "https://registry.npmjs.org/"
+	},
+	"engines": {
+		"node": ">=18"
+	},
+	"peerDependencies": {
+		"@earendil-works/pi-coding-agent": "*"
+	},
+	"pi": {
+		"extensions": [
+			"./src/index.ts"
+		]
+	},
+	"files": [
+		"src",
+		"README.md",
+		"LICENSE"
+	]
 }

package/{index.ts → src/core.ts}

@@ -1,11 +1,10 @@
-import { type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
+export const FAST_COMMAND = "fast";
+export const FAST_FLAG = "fast";
+export const FAST_STATUS_KEY = "fast";
+export const FAST_STATE_CUSTOM_TYPE = "fast";
 
-const FAST_COMMAND = "fast";
-const FAST_FLAG = "fast";
-const FAST_STATUS_KEY = "fast";
-
-const FAST_ON_TEXT = "fast on";
-const FAST_OFF_TEXT = "fast off";
+export const FAST_ON_TEXT = "fast on";
+export const FAST_OFF_TEXT = "fast off";
 
 const FAST_SPEED = "fast";
 const FAST_BETA = "fast-mode-2026-02-01";
@@ -17,41 +16,65 @@ const CLAUDE_API = "anthropic-messages";
 const OPENAI_PROVIDER = "openai-codex";
 const OPENAI_API = "openai-codex-responses";
 
-type FastFeature = {
+export type FastModel = {
+	provider: string;
+	api?: string;
+	id: string;
+	headers?: Record<string, string>;
+};
+
+export type FastContext = {
+	model?: FastModel;
+	modelRegistry: {
+		isUsingOAuth(model: FastModel): boolean;
+	};
+};
+
+export type FastFeature = {
 	provider: string;
 	api: string;
 	supportedModels: Set<string>;
 	injectionKey: string;
 	injectionValue: string;
 	unsupportedModelMessage: string;
-	isEligible?: (ctx: ExtensionContext) => string | undefined;
+	isEligible?: (ctx: FastContext) => string | undefined;
+};
+
+export type FastModeState = {
+	enabled: boolean;
 };
 
-type FastModeState = {
+export type FastStateEntryData = {
 	enabled: boolean;
 };
 
-type CurrentModelStatus = {
+export type FastSessionEntry = {
+	type: string;
+	customType?: string;
+	data?: unknown;
+};
+
+export type CurrentModelStatus = {
 	feature?: FastFeature;
 	isSupported: boolean;
 	reason?: string;
 };
 
-type PayloadRecord = Record<string, unknown>;
-
-type HeaderModel = {
-	headers?: Record<string, string>;
+export type FastStatusView = {
+	text: string;
+	color: "accent" | "muted";
 };
 
-const FEATURES: FastFeature[] = [
+type PayloadRecord = Record<string, unknown>;
+
+export const FEATURES: FastFeature[] = [
 	{
 		provider: CLAUDE_PROVIDER,
 		api: CLAUDE_API,
 		supportedModels: new Set(["claude-opus-4-6", "claude-opus-4-7", "claude-opus-4-8"]),
 		injectionKey: "speed",
 		injectionValue: FAST_SPEED,
-		unsupportedModelMessage:
-			"Fast mode is only available for Claude Opus 4.6, 4.7, and 4.8",
+		unsupportedModelMessage: "Fast mode is only available for Claude Opus 4.6, 4.7, and 4.8",
 	},
 	{
 		provider: OPENAI_PROVIDER,
@@ -67,29 +90,30 @@ const FEATURES: FastFeature[] = [
 	},
 ];
 
-const sessionStates = new WeakMap<object, FastModeState>();
-
-function getSessionState(ctx: ExtensionContext): FastModeState {
-	let state = sessionStates.get(ctx.sessionManager);
-	if (!state) {
-		state = { enabled: false };
-		sessionStates.set(ctx.sessionManager, state);
-	}
-	return state;
+export function createFastModeState(enabled = false): FastModeState {
+	return { enabled };
 }
 
-function isPayloadRecord(payload: unknown): payload is PayloadRecord {
-	return typeof payload === "object" && payload !== null && !Array.isArray(payload);
+export function createFastStateEntryData(state: FastModeState): FastStateEntryData {
+	return { enabled: state.enabled };
 }
 
-function splitBetaHeader(value: string | undefined): string[] {
-	return (value ?? "")
-		.split(",")
-		.map((value) => value.trim())
-		.filter(Boolean);
+export function restoreFastModeState(
+	entries: Iterable<FastSessionEntry>,
+	defaultEnabled = false,
+): FastModeState {
+	let enabled = defaultEnabled;
+
+	for (const entry of entries) {
+		if (entry.type !== "custom" || entry.customType !== FAST_STATE_CUSTOM_TYPE) continue;
+		if (!isPayloadRecord(entry.data) || typeof entry.data.enabled !== "boolean") continue;
+		enabled = entry.data.enabled;
+	}
+
+	return createFastModeState(enabled);
 }
 
-function getCurrentModelStatus(ctx: ExtensionContext): CurrentModelStatus {
+export function getCurrentModelStatus(ctx: FastContext): CurrentModelStatus {
 	const model = ctx.model;
 	if (!model) {
 		return {
@@ -110,11 +134,15 @@ function getCurrentModelStatus(ctx: ExtensionContext): CurrentModelStatus {
 		};
 	}
 
-	const matchingFeature = featuresForBackend.find((feature) => feature.supportedModels.has(model.id));
+	const matchingFeature = featuresForBackend.find((feature) =>
+		feature.supportedModels.has(model.id),
+	);
 	if (!matchingFeature) {
 		return {
 			isSupported: false,
-			reason: featuresForBackend[0]?.unsupportedModelMessage ?? "Current model does not support fast mode",
+			reason:
+				featuresForBackend[0]?.unsupportedModelMessage ??
+				"Current model does not support fast mode",
 		};
 	}
 
@@ -135,15 +163,50 @@ function getCurrentModelStatus(ctx: ExtensionContext): CurrentModelStatus {
 	};
 }
 
+export function syncFeatureState(ctx: FastContext, state: FastModeState): CurrentModelStatus {
+	const modelStatus = getCurrentModelStatus(ctx);
+	syncClaudeBetaHeader(ctx, state, modelStatus);
+	return modelStatus;
+}
+
+export function getStatusView(
+	state: FastModeState,
+	modelStatus: CurrentModelStatus,
+): FastStatusView {
+	return {
+		text: state.enabled ? FAST_ON_TEXT : FAST_OFF_TEXT,
+		color: state.enabled && modelStatus.isSupported ? "accent" : "muted",
+	};
+}
+
+export function getFastPayload(
+	payload: unknown,
+	ctx: FastContext,
+	state: FastModeState,
+	modelStatus: CurrentModelStatus,
+): PayloadRecord | undefined {
+	if (!state.enabled) return undefined;
+	if (!modelStatus.isSupported || !modelStatus.feature) return undefined;
+	if (!isPayloadRecord(payload)) return undefined;
+	if (payload.model !== ctx.model?.id) return undefined;
+	if (modelStatus.feature.injectionKey in payload) return undefined;
+
+	return {
+		...payload,
+		[modelStatus.feature.injectionKey]: modelStatus.feature.injectionValue,
+	};
+}
+
 function syncClaudeBetaHeader(
-	ctx: ExtensionContext,
+	ctx: FastContext,
 	state: FastModeState,
 	modelStatus: CurrentModelStatus,
 ): void {
-	const model = ctx.model as (typeof ctx.model & HeaderModel) | undefined;
+	const model = ctx.model;
 	if (!model || model.provider !== CLAUDE_PROVIDER || model.api !== CLAUDE_API) return;
 
-	const shouldEnable = state.enabled && modelStatus.isSupported && modelStatus.feature?.provider === CLAUDE_PROVIDER;
+	const shouldEnable =
+		state.enabled && modelStatus.isSupported && modelStatus.feature?.provider === CLAUDE_PROVIDER;
 	const headers = { ...(model.headers ?? {}) };
 	const existing = splitBetaHeader(headers["anthropic-beta"] ?? headers["Anthropic-Beta"]);
 	const requiredBase = ctx.modelRegistry.isUsingOAuth(model) ? CLAUDE_CODE_OAUTH_BETAS : [];
@@ -160,97 +223,13 @@ function syncClaudeBetaHeader(
 	model.headers = headers;
 }
 
-function syncFeatureState(ctx: ExtensionContext, state: FastModeState): CurrentModelStatus {
-	const modelStatus = getCurrentModelStatus(ctx);
-	syncClaudeBetaHeader(ctx, state, modelStatus);
-	return modelStatus;
-}
-
-function updateStatus(ctx: ExtensionContext, state: FastModeState, modelStatus: CurrentModelStatus): void {
-	if (!ctx.hasUI) return;
-
-	const statusText = state.enabled ? FAST_ON_TEXT : FAST_OFF_TEXT;
-	const isActiveForCurrentModel = state.enabled && modelStatus.isSupported;
-	ctx.ui.setStatus(
-		FAST_STATUS_KEY,
-		ctx.ui.theme.fg(isActiveForCurrentModel ? "accent" : "muted", statusText),
-	);
-}
-
-function getFastPayload(
-	payload: unknown,
-	ctx: ExtensionContext,
-	state: FastModeState,
-	modelStatus: CurrentModelStatus,
-): PayloadRecord | undefined {
-	if (!state.enabled) return undefined;
-	if (!modelStatus.isSupported || !modelStatus.feature) return undefined;
-	if (!isPayloadRecord(payload)) return undefined;
-	if (payload.model !== ctx.model?.id) return undefined;
-	if (modelStatus.feature.injectionKey in payload) return undefined;
-
-	return {
-		...payload,
-		[modelStatus.feature.injectionKey]: modelStatus.feature.injectionValue,
-	};
+function splitBetaHeader(value: string | undefined): string[] {
+	return (value ?? "")
+		.split(",")
+		.map((value) => value.trim())
+		.filter(Boolean);
 }
 
-export default function fastMode(pi: ExtensionAPI) {
-	pi.registerFlag(FAST_FLAG, {
-		description: "Start with fast mode enabled",
-		type: "boolean",
-		default: false,
-	});
-
-	pi.on("session_start", (_event, ctx) => {
-		const state = getSessionState(ctx);
-		state.enabled = pi.getFlag(FAST_FLAG) === true;
-		const modelStatus = syncFeatureState(ctx, state);
-		updateStatus(ctx, state, modelStatus);
-	});
-
-	pi.on("model_select", (_event, ctx) => {
-		const state = getSessionState(ctx);
-		const modelStatus = syncFeatureState(ctx, state);
-		updateStatus(ctx, state, modelStatus);
-	});
-
-	pi.on("before_provider_request", (event, ctx) => {
-		const state = getSessionState(ctx);
-		const modelStatus = syncFeatureState(ctx, state);
-		updateStatus(ctx, state, modelStatus);
-		return getFastPayload(event.payload, ctx, state, modelStatus);
-	});
-
-	pi.on("session_shutdown", (_event, ctx) => {
-		if (!ctx.hasUI) return;
-		ctx.ui.setStatus(FAST_STATUS_KEY, undefined);
-	});
-
-	pi.registerCommand(FAST_COMMAND, {
-		description: "Toggle fast mode",
-		getArgumentCompletions: () => null,
-		handler: async (args, ctx) => {
-			if (args.trim()) {
-				ctx.ui.notify("Usage: /fast", "warning");
-				return;
-			}
-
-			const state = getSessionState(ctx);
-			state.enabled = !state.enabled;
-
-			const modelStatus = syncFeatureState(ctx, state);
-			updateStatus(ctx, state, modelStatus);
-
-			ctx.ui.notify(`Fast mode is now ${state.enabled ? "on" : "off"}.`, "info");
-
-			if (state.enabled && !modelStatus.isSupported) {
-				const detail = modelStatus.reason ? ` (${modelStatus.reason})` : "";
-				ctx.ui.notify(
-					`Current model is not supported for fast mode${detail}. Fast mode will apply automatically once you switch to a supported model.`,
-					"warning",
-				);
-			}
-		},
-	});
+function isPayloadRecord(payload: unknown): payload is PayloadRecord {
+	return typeof payload === "object" && payload !== null && !Array.isArray(payload);
 }

package/src/index.ts

@@ -0,0 +1,122 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import {
+	FAST_COMMAND,
+	FAST_FLAG,
+	FAST_STATE_CUSTOM_TYPE,
+	FAST_STATUS_KEY,
+	createFastModeState,
+	createFastStateEntryData,
+	getFastPayload,
+	getStatusView,
+	restoreFastModeState,
+	syncFeatureState,
+	type FastContext,
+	type FastModeState,
+	type FastModel,
+	type CurrentModelStatus,
+} from "./core.ts";
+
+const sessionStates = new WeakMap<object, FastModeState>();
+
+function getSessionState(ctx: ExtensionContext): FastModeState {
+	let state = sessionStates.get(ctx.sessionManager);
+	if (!state) {
+		state = createFastModeState();
+		sessionStates.set(ctx.sessionManager, state);
+	}
+	return state;
+}
+
+function restoreSessionState(ctx: ExtensionContext, defaultEnabled: boolean): FastModeState {
+	const state = restoreFastModeState(ctx.sessionManager.getBranch(), defaultEnabled);
+	sessionStates.set(ctx.sessionManager, state);
+	return state;
+}
+
+function toFastContext(ctx: ExtensionContext): FastContext {
+	return {
+		model: ctx.model as FastModel | undefined,
+		modelRegistry: {
+			isUsingOAuth: (model) =>
+				ctx.modelRegistry.isUsingOAuth(model as NonNullable<typeof ctx.model>),
+		},
+	};
+}
+
+function updateStatus(
+	ctx: ExtensionContext,
+	state: FastModeState,
+	modelStatus: CurrentModelStatus,
+): void {
+	if (!ctx.hasUI) return;
+
+	const status = getStatusView(state, modelStatus);
+	ctx.ui.setStatus(FAST_STATUS_KEY, ctx.ui.theme.fg(status.color, status.text));
+}
+
+export default function fastMode(pi: ExtensionAPI) {
+	pi.registerFlag(FAST_FLAG, {
+		description: "Start with fast mode enabled",
+		type: "boolean",
+		default: false,
+	});
+
+	pi.on("session_start", (_event, ctx) => {
+		const state = restoreSessionState(ctx, pi.getFlag(FAST_FLAG) === true);
+		const modelStatus = syncFeatureState(toFastContext(ctx), state);
+		updateStatus(ctx, state, modelStatus);
+	});
+
+	pi.on("model_select", (_event, ctx) => {
+		const state = getSessionState(ctx);
+		const modelStatus = syncFeatureState(toFastContext(ctx), state);
+		updateStatus(ctx, state, modelStatus);
+	});
+
+	pi.on("session_tree", (_event, ctx) => {
+		const state = restoreSessionState(ctx, pi.getFlag(FAST_FLAG) === true);
+		const modelStatus = syncFeatureState(toFastContext(ctx), state);
+		updateStatus(ctx, state, modelStatus);
+	});
+
+	pi.on("before_provider_request", (event, ctx) => {
+		const state = getSessionState(ctx);
+		const fastContext = toFastContext(ctx);
+		const modelStatus = syncFeatureState(fastContext, state);
+		updateStatus(ctx, state, modelStatus);
+		return getFastPayload(event.payload, fastContext, state, modelStatus);
+	});
+
+	pi.on("session_shutdown", (_event, ctx) => {
+		if (!ctx.hasUI) return;
+		ctx.ui.setStatus(FAST_STATUS_KEY, undefined);
+	});
+
+	pi.registerCommand(FAST_COMMAND, {
+		description: "Toggle fast mode",
+		getArgumentCompletions: () => null,
+		handler: async (args, ctx) => {
+			if (args.trim()) {
+				ctx.ui.notify("Usage: /fast", "warning");
+				return;
+			}
+
+			const state = getSessionState(ctx);
+			state.enabled = !state.enabled;
+			pi.appendEntry(FAST_STATE_CUSTOM_TYPE, createFastStateEntryData(state));
+
+			const modelStatus = syncFeatureState(toFastContext(ctx), state);
+			updateStatus(ctx, state, modelStatus);
+
+			ctx.ui.notify(`Fast mode is now ${state.enabled ? "on" : "off"}.`, "info");
+
+			if (state.enabled && !modelStatus.isSupported) {
+				const detail = modelStatus.reason ? ` (${modelStatus.reason})` : "";
+				ctx.ui.notify(
+					`Current model is not supported for fast mode${detail}. Fast mode will apply automatically once you switch to a supported model.`,
+					"warning",
+				);
+			}
+		},
+	});
+}