From e56521e3234131a2c1639a74e2f15fff643acf30 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Mon, 1 Jun 2026 18:31:44 +0200 Subject: [PATCH] Add extension mode context --- packages/coding-agent/CHANGELOG.md | 4 ++ packages/coding-agent/docs/extensions.md | 20 ++++++---- packages/coding-agent/docs/rpc.md | 2 +- .../examples/extensions/custom-header.ts | 2 +- .../examples/extensions/doom-overlay/index.ts | 2 +- .../examples/extensions/handoff.ts | 2 +- .../examples/extensions/interactive-shell.ts | 2 +- .../coding-agent/examples/extensions/qna.ts | 2 +- .../examples/extensions/question.ts | 2 +- .../examples/extensions/questionnaire.ts | 2 +- .../coding-agent/examples/extensions/snake.ts | 2 +- .../examples/extensions/space-invaders.ts | 2 +- .../examples/extensions/summarize.ts | 2 +- .../examples/extensions/tic-tac-toe.ts | 2 +- .../coding-agent/examples/extensions/todo.ts | 2 +- .../coding-agent/examples/extensions/tools.ts | 5 +++ .../coding-agent/src/core/agent-session.ts | 8 +++- .../coding-agent/src/core/extensions/index.ts | 1 + .../src/core/extensions/runner.ts | 9 ++++- .../coding-agent/src/core/extensions/types.ts | 6 ++- .../src/modes/interactive/interactive-mode.ts | 2 + packages/coding-agent/src/modes/print-mode.ts | 1 + .../coding-agent/src/modes/rpc/rpc-mode.ts | 1 + .../test/extensions-runner.test.ts | 39 ++++++++++++++++++- .../test/trigger-compact-extension.test.ts | 1 + 25 files changed, 98 insertions(+), 25 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 6013abac..72c008b6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `ctx.mode` to extension contexts so extensions can distinguish TUI, RPC, JSON, and print mode. + ### 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)). diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index a30bbf52..64b1b87b 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -860,9 +860,13 @@ All handlers receive `ctx: ExtensionContext`. 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 -`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 @@ -2516,14 +2520,14 @@ const highlighted = highlightCode(code, lang, theme); ## Mode Behavior -| Mode | UI Methods | Notes | -|------|-----------|-------| -| Interactive | Full TUI | Normal operation | -| RPC (`--mode rpc`) | JSON protocol | Host handles UI, see [rpc.md](rpc.md) | -| JSON (`--mode json`) | No-op | Event stream to stdout, see [json.md](json.md) | -| Print (`-p`) | No-op | Extensions run but can't prompt | +| Mode | `ctx.mode` | `ctx.hasUI` | Notes | +|------|------------|-------------|-------| +| Interactive | `"tui"` | `true` | Full TUI with terminal rendering | +| RPC (`--mode rpc`) | `"rpc"` | `true` | Dialogs and notifications via JSON protocol; `custom()` returns `undefined`. See [rpc.md](rpc.md) | +| JSON (`--mode json`) | `"json"` | `false` | Event stream to stdout; UI methods are no-ops | +| 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 diff --git a/packages/coding-agent/docs/rpc.md b/packages/coding-agent/docs/rpc.md index 846162c3..9aa16ffc 100644 --- a/packages/coding-agent/docs/rpc.md +++ b/packages/coding-agent/docs/rpc.md @@ -1003,7 +1003,7 @@ Some `ExtensionUIContext` methods are not supported or degraded in RPC mode beca - `getTheme()` returns `undefined` - `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) diff --git a/packages/coding-agent/examples/extensions/custom-header.ts b/packages/coding-agent/examples/extensions/custom-header.ts index e3aeb656..600a85ce 100644 --- a/packages/coding-agent/examples/extensions/custom-header.ts +++ b/packages/coding-agent/examples/extensions/custom-header.ts @@ -47,7 +47,7 @@ function getPiMascot(theme: Theme): string[] { export default function (pi: ExtensionAPI) { // Set custom header immediately on load (if UI is available) pi.on("session_start", async (_event, ctx) => { - if (ctx.hasUI) { + if (ctx.mode === "tui") { ctx.ui.setHeader((_tui, theme) => { return { render(_width: number): string[] { diff --git a/packages/coding-agent/examples/extensions/doom-overlay/index.ts b/packages/coding-agent/examples/extensions/doom-overlay/index.ts index 9d9dad13..843d52da 100644 --- a/packages/coding-agent/examples/extensions/doom-overlay/index.ts +++ b/packages/coding-agent/examples/extensions/doom-overlay/index.ts @@ -23,7 +23,7 @@ export default function (pi: ExtensionAPI) { description: "Play DOOM as an overlay. Q to pause and exit.", handler: async (args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("DOOM requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/handoff.ts b/packages/coding-agent/examples/extensions/handoff.ts index 10581dc6..4af6661f 100644 --- a/packages/coding-agent/examples/extensions/handoff.ts +++ b/packages/coding-agent/examples/extensions/handoff.ts @@ -81,7 +81,7 @@ export default function (pi: ExtensionAPI) { pi.registerCommand("handoff", { description: "Transfer context to a new focused session", handler: async (args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("handoff requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/interactive-shell.ts b/packages/coding-agent/examples/extensions/interactive-shell.ts index 99a7f963..b1dd3746 100644 --- a/packages/coding-agent/examples/extensions/interactive-shell.ts +++ b/packages/coding-agent/examples/extensions/interactive-shell.ts @@ -146,7 +146,7 @@ export default function (pi: ExtensionAPI) { } // No UI available (print mode, RPC, etc.) - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { return { result: { output: "(interactive commands require TUI)", exitCode: 1, cancelled: false, truncated: false }, }; diff --git a/packages/coding-agent/examples/extensions/qna.ts b/packages/coding-agent/examples/extensions/qna.ts index 2675b5d7..70dbef7b 100644 --- a/packages/coding-agent/examples/extensions/qna.ts +++ b/packages/coding-agent/examples/extensions/qna.ts @@ -31,7 +31,7 @@ export default function (pi: ExtensionAPI) { pi.registerCommand("qna", { description: "Extract questions from last assistant message into editor", handler: async (_args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("qna requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/question.ts b/packages/coding-agent/examples/extensions/question.ts index 7aba1d70..eb1af51a 100644 --- a/packages/coding-agent/examples/extensions/question.ts +++ b/packages/coding-agent/examples/extensions/question.ts @@ -41,7 +41,7 @@ export default function question(pi: ExtensionAPI) { parameters: QuestionParams, async execute(_toolCallId, params, _signal, _onUpdate, ctx) { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { return { content: [{ type: "text", text: "Error: UI not available (running in non-interactive mode)" }], details: { diff --git a/packages/coding-agent/examples/extensions/questionnaire.ts b/packages/coding-agent/examples/extensions/questionnaire.ts index 34fb666c..653bd864 100644 --- a/packages/coding-agent/examples/extensions/questionnaire.ts +++ b/packages/coding-agent/examples/extensions/questionnaire.ts @@ -82,7 +82,7 @@ export default function questionnaire(pi: ExtensionAPI) { parameters: QuestionnaireParams, 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)"); } if (params.questions.length === 0) { diff --git a/packages/coding-agent/examples/extensions/snake.ts b/packages/coding-agent/examples/extensions/snake.ts index 29995b36..1b58edf9 100644 --- a/packages/coding-agent/examples/extensions/snake.ts +++ b/packages/coding-agent/examples/extensions/snake.ts @@ -311,7 +311,7 @@ export default function (pi: ExtensionAPI) { description: "Play Snake!", handler: async (_args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("Snake requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/space-invaders.ts b/packages/coding-agent/examples/extensions/space-invaders.ts index c91071b8..306c5ac5 100644 --- a/packages/coding-agent/examples/extensions/space-invaders.ts +++ b/packages/coding-agent/examples/extensions/space-invaders.ts @@ -529,7 +529,7 @@ export default function (pi: ExtensionAPI) { description: "Play Space Invaders!", handler: async (_args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("Space Invaders requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/summarize.ts b/packages/coding-agent/examples/extensions/summarize.ts index dd30c350..e6480974 100644 --- a/packages/coding-agent/examples/extensions/summarize.ts +++ b/packages/coding-agent/examples/extensions/summarize.ts @@ -115,7 +115,7 @@ const buildSummaryPrompt = (conversationText: string): string => ].join("\n"); const showSummaryUi = async (summary: string, ctx: ExtensionCommandContext) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { return; } diff --git a/packages/coding-agent/examples/extensions/tic-tac-toe.ts b/packages/coding-agent/examples/extensions/tic-tac-toe.ts index 8cbfdeab..8077a57f 100644 --- a/packages/coding-agent/examples/extensions/tic-tac-toe.ts +++ b/packages/coding-agent/examples/extensions/tic-tac-toe.ts @@ -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", handler: async (_args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("Tic-tac-toe requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/todo.ts b/packages/coding-agent/examples/extensions/todo.ts index 89b147c8..67cba742 100644 --- a/packages/coding-agent/examples/extensions/todo.ts +++ b/packages/coding-agent/examples/extensions/todo.ts @@ -284,7 +284,7 @@ export default function (pi: ExtensionAPI) { pi.registerCommand("todos", { description: "Show all todos on the current branch", handler: async (_args, ctx) => { - if (!ctx.hasUI) { + if (ctx.mode !== "tui") { ctx.ui.notify("/todos requires interactive mode", "error"); return; } diff --git a/packages/coding-agent/examples/extensions/tools.ts b/packages/coding-agent/examples/extensions/tools.ts index 5e617ab3..8db496a5 100644 --- a/packages/coding-agent/examples/extensions/tools.ts +++ b/packages/coding-agent/examples/extensions/tools.ts @@ -67,6 +67,11 @@ export default function toolsExtension(pi: ExtensionAPI) { pi.registerCommand("tools", { description: "Enable/disable tools", handler: async (_args, ctx) => { + if (ctx.mode !== "tui") { + ctx.ui.notify("/tools requires TUI mode", "error"); + return; + } + // Refresh tool list allTools = pi.getAllTools(); diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 608a6153..93916496 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -56,6 +56,7 @@ import { type ContextUsage, type ExtensionCommandContextActions, type ExtensionErrorListener, + type ExtensionMode, ExtensionRunner, type ExtensionUIContext, type InputSource, @@ -187,6 +188,7 @@ export interface AgentSessionConfig { export interface ExtensionBindings { uiContext?: ExtensionUIContext; + mode?: ExtensionMode; commandContextActions?: ExtensionCommandContextActions; abortHandler?: () => void; shutdownHandler?: ShutdownHandler; @@ -300,6 +302,7 @@ export class AgentSession { private _baseToolsOverride?: Record; private _sessionStartEvent: SessionStartEvent; private _extensionUIContext?: ExtensionUIContext; + private _extensionMode: ExtensionMode = "print"; private _extensionCommandContextActions?: ExtensionCommandContextActions; private _extensionAbortHandler?: () => void; private _extensionShutdownHandler?: ShutdownHandler; @@ -2064,6 +2067,9 @@ export class AgentSession { if (bindings.uiContext !== undefined) { this._extensionUIContext = bindings.uiContext; } + if (bindings.mode !== undefined) { + this._extensionMode = bindings.mode; + } if (bindings.commandContextActions !== undefined) { this._extensionCommandContextActions = bindings.commandContextActions; } @@ -2136,7 +2142,7 @@ export class AgentSession { } private _applyExtensionBindings(runner: ExtensionRunner): void { - runner.setUIContext(this._extensionUIContext); + runner.setUIContext(this._extensionUIContext, this._extensionMode); runner.bindCommandContext(this._extensionCommandContextActions); this._extensionErrorUnsubscriber?.(); diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index acbc50b1..a25fa37b 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -66,6 +66,7 @@ export type { ExtensionFactory, ExtensionFlag, ExtensionHandler, + ExtensionMode, // Runtime ExtensionRuntime, ExtensionShortcut, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 751e28a4..ed6f3b23 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -28,6 +28,7 @@ import type { ExtensionError, ExtensionEvent, ExtensionFlag, + ExtensionMode, ExtensionRuntime, ExtensionShortcut, ExtensionUIContext, @@ -225,6 +226,7 @@ export class ExtensionRunner { private extensions: Extension[]; private runtime: ExtensionRuntime; private uiContext: ExtensionUIContext; + private mode: ExtensionMode = "print"; private cwd: string; private sessionManager: SessionManager; private modelRegistry: ModelRegistry; @@ -354,8 +356,9 @@ export class ExtensionRunner { this.reloadHandler = async () => {}; } - setUIContext(uiContext?: ExtensionUIContext): void { + setUIContext(uiContext?: ExtensionUIContext, mode: ExtensionMode = "print"): void { this.uiContext = uiContext ?? noOpUIContext; + this.mode = mode; } getUIContext(): ExtensionUIContext { @@ -578,6 +581,10 @@ export class ExtensionRunner { runner.assertActive(); return runner.uiContext; }, + get mode() { + runner.assertActive(); + return runner.mode; + }, get hasUI() { runner.assertActive(); return runner.hasUI(); diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 7696c62e..db2a8f77 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -295,10 +295,14 @@ export interface CompactOptions { /** * Context passed to extension event handlers. */ +export type ExtensionMode = "tui" | "rpc" | "json" | "print"; + export interface ExtensionContext { /** UI methods for user interaction */ 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; /** Current working directory */ cwd: string; diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 987a62cb..b874495b 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -1508,6 +1508,7 @@ export class InteractiveMode { const uiContext = this.createExtensionUIContext(); await this.session.bindExtensions({ uiContext, + mode: "tui", abortHandler: () => { this.restoreQueuedMessagesToEditor({ abort: true }); }, @@ -1654,6 +1655,7 @@ export class InteractiveMode { // Create a context for shortcut handlers const createContext = (): ExtensionContext => ({ ui: this.createExtensionUIContext(), + mode: "tui", hasUI: true, cwd: this.sessionManager.getCwd(), sessionManager: this.sessionManager, diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index c9553c55..a5ee0235 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -71,6 +71,7 @@ export async function runPrintMode(runtimeHost: AgentSessionRuntime, options: Pr const rebindSession = async (): Promise => { session = runtimeHost.session; await session.bindExtensions({ + mode: mode === "json" ? "json" : "print", commandContextActions: { waitForIdle: () => session.agent.waitForIdle(), newSession: async (newSessionOptions) => runtimeHost.newSession(newSessionOptions), diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index aa09cbb3..9bb1f089 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -317,6 +317,7 @@ export async function runRpcMode(runtimeHost: AgentSessionRuntime): Promise session.agent.waitForIdle(), newSession: async (options) => runtimeHost.newSession(options), diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 2e7b0c01..0432ec39 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -9,7 +9,12 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { createExtensionRuntime, discoverAndLoadExtensions } from "../src/core/extensions/loader.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 { ModelRegistry } from "../src/core/model-registry.ts"; import { SessionManager } from "../src/core/session-manager.ts"; @@ -441,6 +446,38 @@ describe("ExtensionRunner", () => { controller.abort(); 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", () => { diff --git a/packages/coding-agent/test/trigger-compact-extension.test.ts b/packages/coding-agent/test/trigger-compact-extension.test.ts index a8bdbb2a..c114fac9 100644 --- a/packages/coding-agent/test/trigger-compact-extension.test.ts +++ b/packages/coding-agent/test/trigger-compact-extension.test.ts @@ -4,6 +4,7 @@ import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from ".. function createContext(tokens: number | null, compact = vi.fn()): ExtensionContext { return { + mode: "print", hasUI: false, ui: {} as ExtensionContext["ui"], cwd: process.cwd(),