Add extension mode context

This commit is contained in:
Mario Zechner
2026-06-01 18:31:44 +02:00
parent 335e09ba0d
commit e56521e323
25 changed files with 98 additions and 25 deletions

View File

@@ -2,6 +2,10 @@
## [Unreleased] ## [Unreleased]
### Added
- Added `ctx.mode` to extension contexts so extensions can distinguish TUI, RPC, JSON, and print mode.
### Fixed ### Fixed
- Fixed opening and listing very large JSONL session files by reading session entries line-by-line instead of materializing the full file as one string ([#5231](https://github.com/earendil-works/pi/issues/5231)). - Fixed opening and listing very large JSONL session files by reading session entries line-by-line instead of materializing the full file as one string ([#5231](https://github.com/earendil-works/pi/issues/5231)).

View File

@@ -860,9 +860,13 @@ All handlers receive `ctx: ExtensionContext`.
UI methods for user interaction. See [Custom UI](#custom-ui) for full details. UI methods for user interaction. See [Custom UI](#custom-ui) for full details.
### ctx.mode
Current run mode: `"tui"`, `"rpc"`, `"json"`, or `"print"`. Use `ctx.mode === "tui"` to guard terminal-only features such as `custom()`, component factories, terminal input, and direct TUI rendering.
### ctx.hasUI ### ctx.hasUI
`false` in print mode (`-p`) and JSON mode. `true` in interactive and RPC mode. In RPC mode, dialog methods (`select`, `confirm`, `input`, `editor`) work via the extension UI sub-protocol, and fire-and-forget methods (`notify`, `setStatus`, `setWidget`, `setTitle`, `setEditorText`) emit requests to the client. Some TUI-specific methods are no-ops or return defaults (see [rpc.md](rpc.md#extension-ui-protocol)). `true` in TUI and RPC modes. `false` in print mode (`-p`) and JSON mode. Use this to guard dialog methods (`select`, `confirm`, `input`, `editor`) and fire-and-forget methods (`notify`, `setStatus`, `setWidget`, `setTitle`, `setEditorText`) that work in both TUI and RPC modes. In RPC mode, some TUI-specific methods are no-ops or return defaults (see [rpc.md](rpc.md#extension-ui-protocol)).
### ctx.cwd ### ctx.cwd
@@ -2516,14 +2520,14 @@ const highlighted = highlightCode(code, lang, theme);
## Mode Behavior ## Mode Behavior
| Mode | UI Methods | Notes | | Mode | `ctx.mode` | `ctx.hasUI` | Notes |
|------|-----------|-------| |------|------------|-------------|-------|
| Interactive | Full TUI | Normal operation | | Interactive | `"tui"` | `true` | Full TUI with terminal rendering |
| RPC (`--mode rpc`) | JSON protocol | Host handles UI, see [rpc.md](rpc.md) | | RPC (`--mode rpc`) | `"rpc"` | `true` | Dialogs and notifications via JSON protocol; `custom()` returns `undefined`. See [rpc.md](rpc.md) |
| JSON (`--mode json`) | No-op | Event stream to stdout, see [json.md](json.md) | | JSON (`--mode json`) | `"json"` | `false` | Event stream to stdout; UI methods are no-ops |
| Print (`-p`) | No-op | Extensions run but can't prompt | | Print (`-p`) | `"print"` | `false` | Extensions run but can't prompt |
In non-interactive modes, check `ctx.hasUI` before using UI methods. Use `ctx.mode === "tui"` before TUI-specific features (`custom()`, component factories, terminal input). Use `ctx.hasUI` before dialog and notification methods that work in both TUI and RPC modes.
## Examples Reference ## Examples Reference

View File

@@ -1003,7 +1003,7 @@ Some `ExtensionUIContext` methods are not supported or degraded in RPC mode beca
- `getTheme()` returns `undefined` - `getTheme()` returns `undefined`
- `setTheme()` returns `{ success: false, error: "..." }` - `setTheme()` returns `{ success: false, error: "..." }`
Note: `ctx.hasUI` is `true` in RPC mode because the dialog and fire-and-forget methods are functional via the extension UI sub-protocol. Note: `ctx.mode` is `"rpc"` and `ctx.hasUI` is `true` in RPC mode because the dialog and fire-and-forget methods are functional via the extension UI sub-protocol. Use `ctx.mode === "tui"` to guard TUI-specific features like `custom()` that require a real terminal.
### Extension UI Requests (stdout) ### Extension UI Requests (stdout)

View File

@@ -47,7 +47,7 @@ function getPiMascot(theme: Theme): string[] {
export default function (pi: ExtensionAPI) { export default function (pi: ExtensionAPI) {
// Set custom header immediately on load (if UI is available) // Set custom header immediately on load (if UI is available)
pi.on("session_start", async (_event, ctx) => { pi.on("session_start", async (_event, ctx) => {
if (ctx.hasUI) { if (ctx.mode === "tui") {
ctx.ui.setHeader((_tui, theme) => { ctx.ui.setHeader((_tui, theme) => {
return { return {
render(_width: number): string[] { render(_width: number): string[] {

View File

@@ -23,7 +23,7 @@ export default function (pi: ExtensionAPI) {
description: "Play DOOM as an overlay. Q to pause and exit.", description: "Play DOOM as an overlay. Q to pause and exit.",
handler: async (args, ctx) => { handler: async (args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("DOOM requires interactive mode", "error"); ctx.ui.notify("DOOM requires interactive mode", "error");
return; return;
} }

View File

@@ -81,7 +81,7 @@ export default function (pi: ExtensionAPI) {
pi.registerCommand("handoff", { pi.registerCommand("handoff", {
description: "Transfer context to a new focused session", description: "Transfer context to a new focused session",
handler: async (args, ctx) => { handler: async (args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("handoff requires interactive mode", "error"); ctx.ui.notify("handoff requires interactive mode", "error");
return; return;
} }

View File

@@ -146,7 +146,7 @@ export default function (pi: ExtensionAPI) {
} }
// No UI available (print mode, RPC, etc.) // No UI available (print mode, RPC, etc.)
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
return { return {
result: { output: "(interactive commands require TUI)", exitCode: 1, cancelled: false, truncated: false }, result: { output: "(interactive commands require TUI)", exitCode: 1, cancelled: false, truncated: false },
}; };

View File

@@ -31,7 +31,7 @@ export default function (pi: ExtensionAPI) {
pi.registerCommand("qna", { pi.registerCommand("qna", {
description: "Extract questions from last assistant message into editor", description: "Extract questions from last assistant message into editor",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("qna requires interactive mode", "error"); ctx.ui.notify("qna requires interactive mode", "error");
return; return;
} }

View File

@@ -41,7 +41,7 @@ export default function question(pi: ExtensionAPI) {
parameters: QuestionParams, parameters: QuestionParams,
async execute(_toolCallId, params, _signal, _onUpdate, ctx) { async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
return { return {
content: [{ type: "text", text: "Error: UI not available (running in non-interactive mode)" }], content: [{ type: "text", text: "Error: UI not available (running in non-interactive mode)" }],
details: { details: {

View File

@@ -82,7 +82,7 @@ export default function questionnaire(pi: ExtensionAPI) {
parameters: QuestionnaireParams, parameters: QuestionnaireParams,
async execute(_toolCallId, params, _signal, _onUpdate, ctx) { async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
return errorResult("Error: UI not available (running in non-interactive mode)"); return errorResult("Error: UI not available (running in non-interactive mode)");
} }
if (params.questions.length === 0) { if (params.questions.length === 0) {

View File

@@ -311,7 +311,7 @@ export default function (pi: ExtensionAPI) {
description: "Play Snake!", description: "Play Snake!",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("Snake requires interactive mode", "error"); ctx.ui.notify("Snake requires interactive mode", "error");
return; return;
} }

View File

@@ -529,7 +529,7 @@ export default function (pi: ExtensionAPI) {
description: "Play Space Invaders!", description: "Play Space Invaders!",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("Space Invaders requires interactive mode", "error"); ctx.ui.notify("Space Invaders requires interactive mode", "error");
return; return;
} }

View File

@@ -115,7 +115,7 @@ const buildSummaryPrompt = (conversationText: string): string =>
].join("\n"); ].join("\n");
const showSummaryUi = async (summary: string, ctx: ExtensionCommandContext) => { const showSummaryUi = async (summary: string, ctx: ExtensionCommandContext) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
return; return;
} }

View File

@@ -779,7 +779,7 @@ Decide the target cell first, then dump every action for the turn in one go.
description: "Play tic-tac-toe against the agent", description: "Play tic-tac-toe against the agent",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("Tic-tac-toe requires interactive mode", "error"); ctx.ui.notify("Tic-tac-toe requires interactive mode", "error");
return; return;
} }

View File

@@ -284,7 +284,7 @@ export default function (pi: ExtensionAPI) {
pi.registerCommand("todos", { pi.registerCommand("todos", {
description: "Show all todos on the current branch", description: "Show all todos on the current branch",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (!ctx.hasUI) { if (ctx.mode !== "tui") {
ctx.ui.notify("/todos requires interactive mode", "error"); ctx.ui.notify("/todos requires interactive mode", "error");
return; return;
} }

View File

@@ -67,6 +67,11 @@ export default function toolsExtension(pi: ExtensionAPI) {
pi.registerCommand("tools", { pi.registerCommand("tools", {
description: "Enable/disable tools", description: "Enable/disable tools",
handler: async (_args, ctx) => { handler: async (_args, ctx) => {
if (ctx.mode !== "tui") {
ctx.ui.notify("/tools requires TUI mode", "error");
return;
}
// Refresh tool list // Refresh tool list
allTools = pi.getAllTools(); allTools = pi.getAllTools();

View File

@@ -56,6 +56,7 @@ import {
type ContextUsage, type ContextUsage,
type ExtensionCommandContextActions, type ExtensionCommandContextActions,
type ExtensionErrorListener, type ExtensionErrorListener,
type ExtensionMode,
ExtensionRunner, ExtensionRunner,
type ExtensionUIContext, type ExtensionUIContext,
type InputSource, type InputSource,
@@ -187,6 +188,7 @@ export interface AgentSessionConfig {
export interface ExtensionBindings { export interface ExtensionBindings {
uiContext?: ExtensionUIContext; uiContext?: ExtensionUIContext;
mode?: ExtensionMode;
commandContextActions?: ExtensionCommandContextActions; commandContextActions?: ExtensionCommandContextActions;
abortHandler?: () => void; abortHandler?: () => void;
shutdownHandler?: ShutdownHandler; shutdownHandler?: ShutdownHandler;
@@ -300,6 +302,7 @@ export class AgentSession {
private _baseToolsOverride?: Record<string, AgentTool>; private _baseToolsOverride?: Record<string, AgentTool>;
private _sessionStartEvent: SessionStartEvent; private _sessionStartEvent: SessionStartEvent;
private _extensionUIContext?: ExtensionUIContext; private _extensionUIContext?: ExtensionUIContext;
private _extensionMode: ExtensionMode = "print";
private _extensionCommandContextActions?: ExtensionCommandContextActions; private _extensionCommandContextActions?: ExtensionCommandContextActions;
private _extensionAbortHandler?: () => void; private _extensionAbortHandler?: () => void;
private _extensionShutdownHandler?: ShutdownHandler; private _extensionShutdownHandler?: ShutdownHandler;
@@ -2064,6 +2067,9 @@ export class AgentSession {
if (bindings.uiContext !== undefined) { if (bindings.uiContext !== undefined) {
this._extensionUIContext = bindings.uiContext; this._extensionUIContext = bindings.uiContext;
} }
if (bindings.mode !== undefined) {
this._extensionMode = bindings.mode;
}
if (bindings.commandContextActions !== undefined) { if (bindings.commandContextActions !== undefined) {
this._extensionCommandContextActions = bindings.commandContextActions; this._extensionCommandContextActions = bindings.commandContextActions;
} }
@@ -2136,7 +2142,7 @@ export class AgentSession {
} }
private _applyExtensionBindings(runner: ExtensionRunner): void { private _applyExtensionBindings(runner: ExtensionRunner): void {
runner.setUIContext(this._extensionUIContext); runner.setUIContext(this._extensionUIContext, this._extensionMode);
runner.bindCommandContext(this._extensionCommandContextActions); runner.bindCommandContext(this._extensionCommandContextActions);
this._extensionErrorUnsubscriber?.(); this._extensionErrorUnsubscriber?.();

View File

@@ -66,6 +66,7 @@ export type {
ExtensionFactory, ExtensionFactory,
ExtensionFlag, ExtensionFlag,
ExtensionHandler, ExtensionHandler,
ExtensionMode,
// Runtime // Runtime
ExtensionRuntime, ExtensionRuntime,
ExtensionShortcut, ExtensionShortcut,

View File

@@ -28,6 +28,7 @@ import type {
ExtensionError, ExtensionError,
ExtensionEvent, ExtensionEvent,
ExtensionFlag, ExtensionFlag,
ExtensionMode,
ExtensionRuntime, ExtensionRuntime,
ExtensionShortcut, ExtensionShortcut,
ExtensionUIContext, ExtensionUIContext,
@@ -225,6 +226,7 @@ export class ExtensionRunner {
private extensions: Extension[]; private extensions: Extension[];
private runtime: ExtensionRuntime; private runtime: ExtensionRuntime;
private uiContext: ExtensionUIContext; private uiContext: ExtensionUIContext;
private mode: ExtensionMode = "print";
private cwd: string; private cwd: string;
private sessionManager: SessionManager; private sessionManager: SessionManager;
private modelRegistry: ModelRegistry; private modelRegistry: ModelRegistry;
@@ -354,8 +356,9 @@ export class ExtensionRunner {
this.reloadHandler = async () => {}; this.reloadHandler = async () => {};
} }
setUIContext(uiContext?: ExtensionUIContext): void { setUIContext(uiContext?: ExtensionUIContext, mode: ExtensionMode = "print"): void {
this.uiContext = uiContext ?? noOpUIContext; this.uiContext = uiContext ?? noOpUIContext;
this.mode = mode;
} }
getUIContext(): ExtensionUIContext { getUIContext(): ExtensionUIContext {
@@ -578,6 +581,10 @@ export class ExtensionRunner {
runner.assertActive(); runner.assertActive();
return runner.uiContext; return runner.uiContext;
}, },
get mode() {
runner.assertActive();
return runner.mode;
},
get hasUI() { get hasUI() {
runner.assertActive(); runner.assertActive();
return runner.hasUI(); return runner.hasUI();

View File

@@ -295,10 +295,14 @@ export interface CompactOptions {
/** /**
* Context passed to extension event handlers. * Context passed to extension event handlers.
*/ */
export type ExtensionMode = "tui" | "rpc" | "json" | "print";
export interface ExtensionContext { export interface ExtensionContext {
/** UI methods for user interaction */ /** UI methods for user interaction */
ui: ExtensionUIContext; ui: ExtensionUIContext;
/** Whether UI is available (false in print/RPC mode) */ /** Current run mode. Use "tui" to guard terminal-only UI such as custom components. */
mode: ExtensionMode;
/** Whether dialog-capable UI is available (true in TUI and RPC modes) */
hasUI: boolean; hasUI: boolean;
/** Current working directory */ /** Current working directory */
cwd: string; cwd: string;

View File

@@ -1508,6 +1508,7 @@ export class InteractiveMode {
const uiContext = this.createExtensionUIContext(); const uiContext = this.createExtensionUIContext();
await this.session.bindExtensions({ await this.session.bindExtensions({
uiContext, uiContext,
mode: "tui",
abortHandler: () => { abortHandler: () => {
this.restoreQueuedMessagesToEditor({ abort: true }); this.restoreQueuedMessagesToEditor({ abort: true });
}, },
@@ -1654,6 +1655,7 @@ export class InteractiveMode {
// Create a context for shortcut handlers // Create a context for shortcut handlers
const createContext = (): ExtensionContext => ({ const createContext = (): ExtensionContext => ({
ui: this.createExtensionUIContext(), ui: this.createExtensionUIContext(),
mode: "tui",
hasUI: true, hasUI: true,
cwd: this.sessionManager.getCwd(), cwd: this.sessionManager.getCwd(),
sessionManager: this.sessionManager, sessionManager: this.sessionManager,

View File

@@ -71,6 +71,7 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr
const rebindSession = async (): Promise<void> => { const rebindSession = async (): Promise<void> => {
session = runtimeHost.session; session = runtimeHost.session;
await session.bindExtensions({ await session.bindExtensions({
mode: mode === "json" ? "json" : "print",
commandContextActions: { commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(), waitForIdle: () => session.agent.waitForIdle(),
newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions), newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions),

View File

@@ -317,6 +317,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise<neve
session = runtimeHost.session; session = runtimeHost.session;
await session.bindExtensions({ await session.bindExtensions({
uiContext: createExtensionUIContext(), uiContext: createExtensionUIContext(),
mode: "rpc",
commandContextActions: { commandContextActions: {
waitForIdle: () => session.agent.waitForIdle(), waitForIdle: () => session.agent.waitForIdle(),
newSession: async (options) => runtimeHost.newSession(options), newSession: async (options) => runtimeHost.newSession(options),

View File

@@ -9,7 +9,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { AuthStorage } from "../src/core/auth-storage.ts"; import { AuthStorage } from "../src/core/auth-storage.ts";
import { createExtensionRuntime, discoverAndLoadExtensions } from "../src/core/extensions/loader.ts"; import { createExtensionRuntime, discoverAndLoadExtensions } from "../src/core/extensions/loader.ts";
import { ExtensionRunner } from "../src/core/extensions/runner.ts"; import { ExtensionRunner } from "../src/core/extensions/runner.ts";
import type { ExtensionActions, ExtensionContextActions, ProviderConfig } from "../src/core/extensions/types.ts"; import type {
ExtensionActions,
ExtensionContextActions,
ExtensionUIContext,
ProviderConfig,
} from "../src/core/extensions/types.ts";
import { KeybindingsManager, type KeyId } from "../src/core/keybindings.ts"; import { KeybindingsManager, type KeyId } from "../src/core/keybindings.ts";
import { ModelRegistry } from "../src/core/model-registry.ts"; import { ModelRegistry } from "../src/core/model-registry.ts";
import { SessionManager } from "../src/core/session-manager.ts"; import { SessionManager } from "../src/core/session-manager.ts";
@@ -441,6 +446,38 @@ describe("ExtensionRunner", () => {
controller.abort(); controller.abort();
expect(ctx.signal?.aborted).toBe(true); expect(ctx.signal?.aborted).toBe(true);
}); });
it("exposes print mode and hasUI false by default", async () => {
const result = await discoverAndLoadExtensions([], tempDir, tempDir);
const runner = new ExtensionRunner(result.extensions, result.runtime, tempDir, sessionManager, modelRegistry);
runner.bindCore(extensionActions, extensionContextActions);
const ctx = runner.createContext();
expect(ctx.mode).toBe("print");
expect(ctx.hasUI).toBe(false);
});
it("exposes rpc mode with hasUI true when an RPC UI context is provided", async () => {
const result = await discoverAndLoadExtensions([], tempDir, tempDir);
const runner = new ExtensionRunner(result.extensions, result.runtime, tempDir, sessionManager, modelRegistry);
runner.bindCore(extensionActions, extensionContextActions);
runner.setUIContext({} as ExtensionUIContext, "rpc");
const ctx = runner.createContext();
expect(ctx.mode).toBe("rpc");
expect(ctx.hasUI).toBe(true);
});
it("exposes tui mode with hasUI true when a TUI UI context is provided", async () => {
const result = await discoverAndLoadExtensions([], tempDir, tempDir);
const runner = new ExtensionRunner(result.extensions, result.runtime, tempDir, sessionManager, modelRegistry);
runner.bindCore(extensionActions, extensionContextActions);
runner.setUIContext({} as ExtensionUIContext, "tui");
const ctx = runner.createContext();
expect(ctx.mode).toBe("tui");
expect(ctx.hasUI).toBe(true);
});
}); });
describe("error handling", () => { describe("error handling", () => {

View File

@@ -4,6 +4,7 @@ import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "..
function createContext(tokens: number | null, compact = vi.fn()): ExtensionContext { function createContext(tokens: number | null, compact = vi.fn()): ExtensionContext {
return { return {
mode: "print",
hasUI: false, hasUI: false,
ui: {} as ExtensionContext["ui"], ui: {} as ExtensionContext["ui"],
cwd: process.cwd(), cwd: process.cwd(),