@smooai/smooth-extension-sdk 0.3.0 → 0.5.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@smooai/smooth-extension-sdk",
-  "version": "0.3.0",
+  "version": "0.5.0",
   "description": "TypeScript SDK for building Smooth Extension Protocol (SEP) extensions: `defineExtension`, `defineTool`, a stdio JSON-RPC transport, an in-process test host, and a conformance runner. Extensions are subprocesses speaking JSON-RPC 2.0 ndjson to any SEP host (smooth-operator-core and its polyglot servers).",
   "license": "MIT",
   "type": "module",

package/src/conformance.ts

@@ -115,6 +115,8 @@ export async function runConformance(opts: RunConformanceOptions): Promise<Confo
         await check('initialize', method.INITIALIZE, initParams(fixtures), 'methods/initialize.schema.json#/$defs/Result');
         await check('ping', method.PING, {}, 'methods/ping.schema.json#/$defs/Result');
         await check('tool/execute', method.TOOL_EXECUTE, fixtures.tool_execute_params!.instance, 'methods/tool-execute.schema.json#/$defs/Result');
+        await check('command/execute', method.COMMAND_EXECUTE, fixtures.command_execute_params!.instance, 'methods/command-execute.schema.json#/$defs/Result');
+        await check('command/complete', method.COMMAND_COMPLETE, fixtures.command_complete_params!.instance, 'methods/command-complete.schema.json#/$defs/Result');
         await check('shutdown', method.SHUTDOWN, {}, 'methods/shutdown.schema.json#/$defs/Result');
     } finally {
         peer.close();

package/src/extension.ts

@@ -12,7 +12,28 @@
  */
 import { Peer } from './jsonrpc.js';
 import { PROTOCOL_VERSION, method } from './protocol.js';
-import type { Context, EventParams, HookOutcome, HookParams, InitializeParams, InitializeResult, ToolExecuteParams, ToolExecuteResult, ToolUpdateParams } from './protocol.js';
+import type {
+    CommandCompleteParams,
+    CommandCompleteResult,
+    CommandExecuteParams,
+    CommandExecuteResult,
+    CommandRegistration,
+    Completion,
+    Context,
+    DeliverAs,
+    EventParams,
+    HookOutcome,
+    HookParams,
+    InitializeParams,
+    InitializeResult,
+    ShortcutRegistration,
+    ToolExecuteParams,
+    ToolExecuteResult,
+    ToolUpdateParams,
+    UiKind,
+    UiRequestParams,
+    UiRequestResult,
+} from './protocol.js';
 import { toJsonSchema, type ParameterSchema } from './schema.js';
 import { stdioTransport, type Transport } from './transport.js';
 
@@ -20,6 +41,74 @@ import { stdioTransport, type Transport } from './transport.js';
  *  is a fire-and-forget observe event. Kept in sync with the engine's HookType. */
 const HOOK_NAMES = new Set(['tool_call', 'tool_result', 'before_agent_start', 'message_end', 'context', 'before_provider_request', 'input', 'user_bash']);
 
+/**
+ * The `ui/request` surface handed to tools (and to event handlers via
+ * `smooth.ui`). Each call is an ext→host request; the frontend renders it and
+ * replies. `select`/`confirm`/`input` return an answer (or `{ cancelled: true }`
+ * if dismissed); `notify`/`setStatus`/`setWidget`/`setTitle` resolve empty. A
+ * headless or uncapable host rejects with an `RpcError` of code -32001 (NoUI) —
+ * gate with `hasUI(kind)` to avoid it.
+ */
+export interface UiApi {
+    select(prompt: string, options: string[]): Promise<UiRequestResult>;
+    confirm(prompt: string): Promise<UiRequestResult>;
+    input(prompt: string, opts?: { default?: string }): Promise<UiRequestResult>;
+    notify(message: string, level?: 'info' | 'warn' | 'error'): Promise<void>;
+    setStatus(status: string): Promise<void>;
+    setWidget(widget: Record<string, unknown>): Promise<void>;
+    setTitle(title: string): Promise<void>;
+}
+
+/** Build a [`UiApi`] that speaks `ui/request` over `peer`. */
+function makeUi(peer: Peer): UiApi {
+    const req = (params: UiRequestParams) => peer.request<UiRequestResult>(method.UI_REQUEST, params);
+    return {
+        select: (prompt, options) => req({ kind: 'select', prompt, options }),
+        confirm: (prompt) => req({ kind: 'confirm', prompt }),
+        input: (prompt, opts) => req({ kind: 'input', prompt, ...(opts?.default !== undefined ? { default: opts.default } : {}) }),
+        notify: async (message, level) => {
+            await req({ kind: 'notify', message, ...(level ? { level } : {}) });
+        },
+        setStatus: async (status) => {
+            await req({ kind: 'set_status', status });
+        },
+        setWidget: async (widget) => {
+            await req({ kind: 'set_widget', widget });
+        },
+        setTitle: async (title) => {
+            await req({ kind: 'set_title', title });
+        },
+    };
+}
+
+/**
+ * Session-mutating ext→host actions. Available only from a COMMAND-tier context
+ * (command handlers) — the host rejects them from an event-tier context with
+ * -32003 ContextViolation. `sendMessage` posts a message, `sendUserMessage`
+ * delivers a user message (steer/follow_up/next_turn), `appendEntry` persists an
+ * LLM-invisible transcript entry.
+ */
+export interface SessionApi {
+    sendMessage(text: string, opts?: { role?: 'user' | 'assistant' }): Promise<void>;
+    sendUserMessage(text: string, opts?: { deliverAs?: DeliverAs }): Promise<void>;
+    appendEntry(entry: Record<string, unknown>): Promise<void>;
+}
+
+/** Build a [`SessionApi`] bound to `context` (must be command-tier) over `peer`. */
+function makeSession(peer: Peer, context: Context): SessionApi {
+    return {
+        sendMessage: async (text, opts) => {
+            await peer.request(method.SESSION_SEND_MESSAGE, { context, text, ...(opts?.role ? { role: opts.role } : {}) });
+        },
+        sendUserMessage: async (text, opts) => {
+            await peer.request(method.SESSION_SEND_USER_MESSAGE, { context, text, ...(opts?.deliverAs ? { deliver_as: opts.deliverAs } : {}) });
+        },
+        appendEntry: async (entry) => {
+            await peer.request(method.SESSION_APPEND_ENTRY, { context, entry });
+        },
+    };
+}
+
 /** Progress + cancellation handed to a tool while it runs. */
 export interface ToolContext {
     /** Correlates `onUpdate` calls with this execution. */
@@ -30,6 +119,10 @@ export interface ToolContext {
     signal: AbortSignal;
     /** Stream a progress notification back to the host. */
     onUpdate(update: Omit<ToolUpdateParams, 'call_id'>): void;
+    /** Ask the frontend to render a dialog/widget. See [`UiApi`]. */
+    ui: UiApi;
+    /** True if the host's frontend can render this `ui/request` kind. */
+    hasUI(kind: UiKind): boolean;
 }
 
 /** What a tool's `execute` may return: a full result or just its `content`. */
@@ -48,6 +141,48 @@ export function defineTool<TArgs = Record<string, unknown>>(def: ToolDef<TArgs>)
     return def;
 }
 
+/** What a command handler receives: the command-tier context plus the session,
+ *  ui, and args bound to it. Session actions are valid because a command runs at
+ *  command tier. */
+export interface CommandContext {
+    /** The dispatch context (command tier). */
+    context: Context;
+    /** Free-form arguments parsed from the invocation. */
+    args: Record<string, unknown> | undefined;
+    /** Session-mutating actions, bound to this command's context. */
+    session: SessionApi;
+    /** Ask the frontend to render a dialog/widget. See [`UiApi`]. */
+    ui: UiApi;
+    /** True if the host's frontend can render this `ui/request` kind. */
+    hasUI(kind: UiKind): boolean;
+    /** Structured log line into host tracing. */
+    log(level: 'debug' | 'info' | 'warn' | 'error', message: string, fields?: Record<string, unknown>): void;
+}
+
+/** What a command's `execute` may return: text to surface, a full result, or
+ *  nothing. */
+export type CommandReturn = CommandExecuteResult | string | void;
+
+export interface CommandDef {
+    name: string;
+    description: string;
+    execute(ctx: CommandContext): Promise<CommandReturn> | CommandReturn;
+    /** Optional argument autocomplete: given the partial text, return candidates. */
+    complete?(partial: string, context: Context): Promise<Completion[]> | Completion[];
+}
+
+/** Identity helper for `registerCommand` call sites. */
+export function defineCommand(def: CommandDef): CommandDef {
+    return def;
+}
+
+/** A CLI/slash flag the extension declares. The host delivers its parsed value
+ *  in `initialize`; read it with `smooth.getFlag(name)`. */
+export interface FlagDef {
+    name: string;
+    description?: string;
+}
+
 /** A hook handler's friendly return: veto the operation, or replace its input
  *  with a patch (shallow-merged onto the input). Returning nothing = continue. */
 export type HookResult = { block: true; reason?: string } | { patch: Record<string, unknown> };
@@ -62,8 +197,23 @@ export interface SmoothApi {
     name: string;
     version: string;
     registerTool(tool: ToolDef<any>): void;
+    /** Register a slash-command surfaced in the host's `/` palette. */
+    registerCommand(command: CommandDef): void;
+    /** Declare a CLI/slash flag; read its delivered value with [`getFlag`]. */
+    registerFlag(flag: FlagDef): void;
+    /** Bind a keyboard shortcut (TUI frontends) to a registered command. */
+    registerShortcut(shortcut: ShortcutRegistration): void;
     on(event: string, handler: EventHandler): void;
     log(level: 'debug' | 'info' | 'warn' | 'error', message: string, fields?: Record<string, unknown>): void;
+    /** The parsed value the host delivered for a declared flag (undefined before
+     * `initialize`, or when the host has no CLI surface). */
+    getFlag(name: string): unknown;
+    /** True if the host's frontend can render this `ui/request` kind. Only
+     * meaningful after `initialize` — returns false before the handshake. */
+    hasUI(kind: UiKind): boolean;
+    /** The `ui/request` surface. Available after the extension connects; throws
+     * if read before then. Gate with [`hasUI`]. */
+    readonly ui: UiApi;
 }
 
 export type ExtensionSetup = (smooth: SmoothApi) => void;
@@ -75,11 +225,18 @@ export interface ConnectHandle {
 
 export class Extension {
     private readonly tools = new Map<string, ToolDef<any>>();
+    private readonly commands = new Map<string, CommandDef>();
+    private readonly flagDefs = new Map<string, FlagDef>();
+    private readonly shortcuts: ShortcutRegistration[] = [];
     private readonly events = new Map<string, EventHandler[]>();
     private name = 'extension';
     private version = '0.0.0';
     /** Set once connected so `log()` before connect is a safe no-op. */
     private live?: Peer;
+    /** UI kinds the host declared answerable at `initialize`. */
+    private hostUiCaps: string[] = [];
+    /** Flag values the host delivered at `initialize` (name → value). */
+    private flagValues: Record<string, unknown> = {};
 
     constructor(setup: ExtensionSetup) {
         const api: SmoothApi = {
@@ -98,6 +255,15 @@ export class Extension {
             registerTool: (tool) => {
                 this.tools.set(tool.name, tool);
             },
+            registerCommand: (command) => {
+                this.commands.set(command.name, command);
+            },
+            registerFlag: (flag) => {
+                this.flagDefs.set(flag.name, flag);
+            },
+            registerShortcut: (shortcut) => {
+                this.shortcuts.push(shortcut);
+            },
             on: (event, handler) => {
                 const list = this.events.get(event) ?? [];
                 list.push(handler);
@@ -106,6 +272,12 @@ export class Extension {
             log: (level, message, fields) => {
                 this.live?.notify(method.LOG, { level, message, ...(fields ? { fields } : {}) });
             },
+            getFlag: (name) => self.flagValues[name],
+            hasUI: (kind) => self.hostUiCaps.includes(kind),
+            get ui() {
+                if (!self.live) throw new Error('smooth.ui is only available after the extension connects');
+                return makeUi(self.live);
+            },
         };
         // `self` alias so the getter/setter pair above closes over the instance.
         const self = this;
@@ -125,6 +297,8 @@ export class Extension {
         });
         peer.setRequestHandler(method.TOOL_EXECUTE, (params, signal) => this.executeTool(params as ToolExecuteParams, peer, signal));
         peer.setRequestHandler(method.HOOK, (params) => this.dispatchHook(params as HookParams));
+        peer.setRequestHandler(method.COMMAND_EXECUTE, (params) => this.executeCommand(params as CommandExecuteParams, peer));
+        peer.setRequestHandler(method.COMMAND_COMPLETE, (params) => this.completeCommand(params as CommandCompleteParams));
         peer.setNotificationHandler(method.EVENT, (params) => this.dispatchEvent(params as EventParams));
 
         transport.start((frame) => peer.receive(frame));
@@ -145,23 +319,56 @@ export class Extension {
         });
     }
 
-    private initialize(_params: InitializeParams): InitializeResult {
+    private initialize(params: InitializeParams): InitializeResult {
+        this.hostUiCaps = params.ui_capabilities ?? [];
+        this.flagValues = params.flags ?? {};
         const tools = [...this.tools.values()].map((t) => ({
             name: t.name,
             description: t.description,
             parameters: toJsonSchema(t.parameters),
             ...(t.deferred ? { deferred: true } : {}),
         }));
+        const commands: CommandRegistration[] = [...this.commands.values()].map((c) => ({ name: c.name, description: c.description }));
+        const flags = [...this.flagDefs.keys()];
         // Only observe events go in `subscriptions` — hook names are intercepts
         // the host always calls, not events it filters by subscription.
         const subscriptions = [...this.events.keys()].filter((name) => !HOOK_NAMES.has(name));
         return {
             protocol_version: PROTOCOL_VERSION,
             extension: { name: this.name, version: this.version },
-            registrations: { tools, subscriptions },
+            registrations: {
+                tools,
+                ...(commands.length ? { commands } : {}),
+                ...(flags.length ? { flags } : {}),
+                ...(this.shortcuts.length ? { shortcuts: this.shortcuts } : {}),
+                subscriptions,
+            },
         };
     }
 
+    private async executeCommand(params: CommandExecuteParams, peer: Peer): Promise<CommandExecuteResult> {
+        const command = this.commands.get(params.command);
+        if (!command) return { content: `unknown command: ${params.command}` };
+        const ctx: CommandContext = {
+            context: params.context,
+            args: params.arguments,
+            session: makeSession(peer, params.context),
+            ui: makeUi(peer),
+            hasUI: (kind) => this.hostUiCaps.includes(kind),
+            log: (level, message, fields) => peer.notify(method.LOG, { level, message, ...(fields ? { fields } : {}) }),
+        };
+        const out = await command.execute(ctx);
+        if (out === undefined) return {};
+        return typeof out === 'string' ? { content: out } : out;
+    }
+
+    private async completeCommand(params: CommandCompleteParams): Promise<CommandCompleteResult> {
+        const command = this.commands.get(params.command);
+        if (!command?.complete) return { completions: [] };
+        const completions = await command.complete(params.partial ?? '', params.context);
+        return { completions };
+    }
+
     private async executeTool(params: ToolExecuteParams, peer: Peer, signal: AbortSignal): Promise<ToolExecuteResult> {
         const tool = this.tools.get(params.tool);
         if (!tool) return { content: `unknown tool: ${params.tool}`, is_error: true };
@@ -170,6 +377,8 @@ export class Extension {
             context: params.context,
             signal,
             onUpdate: (update) => peer.notify(method.TOOL_UPDATE, { call_id: params.call_id, ...update }),
+            ui: makeUi(peer),
+            hasUI: (kind) => this.hostUiCaps.includes(kind),
         };
         const out = await tool.execute(params.arguments, ctx);
         return typeof out === 'string' ? { content: out } : out;

package/src/index.ts

@@ -8,10 +8,25 @@
  * `createTestHost`, and gate it against the shared fixtures with
  * `runConformance`.
  */
-export { defineExtension, defineTool, Extension } from './extension.js';
-export type { ExtensionSetup, SmoothApi, ToolDef, ToolContext, ToolReturn, EventHandler, HookResult, ConnectHandle } from './extension.js';
+export { defineExtension, defineTool, defineCommand, Extension } from './extension.js';
+export type {
+    ExtensionSetup,
+    SmoothApi,
+    ToolDef,
+    ToolContext,
+    ToolReturn,
+    CommandDef,
+    CommandContext,
+    CommandReturn,
+    FlagDef,
+    SessionApi,
+    EventHandler,
+    HookResult,
+    ConnectHandle,
+    UiApi,
+} from './extension.js';
 export { createTestHost } from './test-host.js';
-export type { TestHost, CallToolOptions } from './test-host.js';
+export type { TestHost, CallToolOptions, CreateTestHostOptions, UiResponder } from './test-host.js';
 export { runConformance, DEFAULT_SPEC_DIR } from './conformance.js';
 export type { ConformanceReport, ConformanceStep, RunConformanceOptions } from './conformance.js';
 export { toJsonSchema } from './schema.js';
@@ -33,4 +48,18 @@ export type {
     EventParams,
     HookParams,
     HookOutcome,
+    UiKind,
+    UiRequestParams,
+    UiRequestResult,
+    ShortcutRegistration,
+    CommandRegistration,
+    CommandExecuteParams,
+    CommandExecuteResult,
+    CommandCompleteParams,
+    CommandCompleteResult,
+    Completion,
+    DeliverAs,
+    SessionSendMessageParams,
+    SessionSendUserMessageParams,
+    SessionAppendEntryParams,
 } from './protocol.js';

package/src/protocol.ts

@@ -18,9 +18,15 @@ export const method = {
     HOOK: 'hook',
     TOOL_EXECUTE: 'tool/execute',
     TOOL_UPDATE: 'tool/update',
+    UI_REQUEST: 'ui/request',
     REGISTRY_UPDATE: 'registry/update',
     LOG: 'log',
     CANCEL: '$/cancel',
+    COMMAND_EXECUTE: 'command/execute',
+    COMMAND_COMPLETE: 'command/complete',
+    SESSION_SEND_MESSAGE: 'session/send_message',
+    SESSION_SEND_USER_MESSAGE: 'session/send_user_message',
+    SESSION_APPEND_ENTRY: 'session/append_entry',
 } as const;
 
 /** JSON-RPC + SEP error codes (see spec/extension/envelope.md). */
@@ -61,6 +67,8 @@ export interface InitializeParams {
     session?: { id?: string };
     mode: 'tui' | 'web' | 'widget' | 'cli' | 'headless';
     ui_capabilities?: string[];
+    /** Parsed values for the flags the extension declares (name → value). */
+    flags?: Record<string, unknown>;
     capabilities_enabled?: Record<string, boolean>;
 }
 
@@ -77,10 +85,19 @@ export interface CommandRegistration {
     description: string;
 }
 
+export interface ShortcutRegistration {
+    /** A human-typed chord, e.g. `ctrl+p`; the frontend parses it. */
+    key: string;
+    /** The registered command this chord invokes (no leading `/`). */
+    command: string;
+    description?: string;
+}
+
 export interface Registrations {
     tools?: ToolRegistration[];
     commands?: CommandRegistration[];
     flags?: string[];
+    shortcuts?: ShortcutRegistration[];
     subscriptions?: string[];
 }
 
@@ -130,3 +147,76 @@ export type HookOutcome =
     | { action: 'continue' }
     | { action: 'block'; reason?: string }
     | { action: 'modify'; patch: Record<string, unknown> };
+
+/** The seven `ui/request` kinds (snake_case wire names). */
+export type UiKind = 'select' | 'confirm' | 'input' | 'notify' | 'set_status' | 'set_widget' | 'set_title';
+
+/** Params of `ui/request` (ext → host), discriminated by `kind`. */
+export type UiRequestParams =
+    | { kind: 'select'; prompt: string; options: string[] }
+    | { kind: 'confirm'; prompt: string }
+    | { kind: 'input'; prompt: string; default?: string }
+    | { kind: 'notify'; message: string; level?: 'info' | 'warn' | 'error' }
+    | { kind: 'set_status'; status: string }
+    | { kind: 'set_widget'; widget: Record<string, unknown> }
+    | { kind: 'set_title'; title: string };
+
+/**
+ * Reply to a `ui/request`. Which field is set depends on the request `kind`:
+ * `select` → `value`, `confirm` → `confirmed`, `input` → `text`; the rest are
+ * empty. Any kind may set `cancelled` if the user dismissed the UI.
+ */
+export interface UiRequestResult {
+    value?: string;
+    confirmed?: boolean;
+    text?: string;
+    cancelled?: boolean;
+}
+
+/** Host → ext `command/execute`: run a registered slash-command (command tier). */
+export interface CommandExecuteParams {
+    command: string;
+    context: Context;
+    arguments?: Record<string, unknown>;
+}
+
+export interface CommandExecuteResult {
+    content?: string;
+}
+
+/** Host → ext `command/complete`: argument autocomplete for a slash-command. */
+export interface CommandCompleteParams {
+    command: string;
+    context: Context;
+    partial?: string;
+}
+
+export interface Completion {
+    value: string;
+    description?: string;
+}
+
+export interface CommandCompleteResult {
+    completions: Completion[];
+}
+
+/** How a `session/send_user_message` is delivered relative to the current turn. */
+export type DeliverAs = 'steer' | 'follow_up' | 'next_turn';
+
+/** Ext → host `session/send_message` params (command tier — carries `context`). */
+export interface SessionSendMessageParams {
+    context: Context;
+    text: string;
+    role?: 'user' | 'assistant';
+}
+
+export interface SessionSendUserMessageParams {
+    context: Context;
+    text: string;
+    deliver_as?: DeliverAs;
+}
+
+export interface SessionAppendEntryParams {
+    context: Context;
+    entry: Record<string, unknown>;
+}

package/src/test-host.ts

@@ -5,14 +5,39 @@
  * (with progress + cancellation), events, ping and shutdown directly against a
  * `defineExtension(...)` object.
  */
-import { Peer } from './jsonrpc.js';
-import { PROTOCOL_VERSION, method } from './protocol.js';
-import type { Context, HookOutcome, InitializeParams, InitializeResult, ToolExecuteResult, ToolUpdateParams } from './protocol.js';
+import { Peer, RpcError } from './jsonrpc.js';
+import { PROTOCOL_VERSION, errorCode, method } from './protocol.js';
+import type {
+    CommandExecuteResult,
+    Context,
+    HookOutcome,
+    InitializeParams,
+    InitializeResult,
+    ToolExecuteResult,
+    ToolUpdateParams,
+    UiRequestParams,
+    UiRequestResult,
+} from './protocol.js';
 import type { Extension } from './extension.js';
 import { linkedPair } from './transport.js';
 
 let callSeq = 0;
 
+/** Answers the extension's `ui/request` calls. Return a result, or throw an
+ * `RpcError` (e.g. code -32001) to simulate a headless/uncapable frontend. */
+export type UiResponder = (params: UiRequestParams) => UiRequestResult | Promise<UiRequestResult>;
+
+/** A recorded ext→host `session/*` request the test can assert on. */
+export interface SessionCall {
+    method: string;
+    params: Record<string, unknown>;
+}
+
+export interface CreateTestHostOptions {
+    /** Answers `ui/request`. Default: reject every call with -32001 NoUI. */
+    onUiRequest?: UiResponder;
+}
+
 export interface CallToolOptions {
     /** Receives each `tool/update` the extension streams for this call. */
     onUpdate?: (update: ToolUpdateParams) => void;
@@ -27,15 +52,21 @@ export interface TestHost {
     callTool(tool: string, args: Record<string, unknown>, opts?: CallToolOptions): Promise<ToolExecuteResult>;
     /** Drive a `hook` request and get back the extension's folded outcome. */
     callHook(hook: string, input: Record<string, unknown>, context?: Context): Promise<HookOutcome>;
+    /** Dispatch a `command/execute` with a command-tier context by default. */
+    runCommand(command: string, args?: Record<string, unknown>, context?: Context): Promise<CommandExecuteResult>;
+    /** Dispatch a `command/complete` for argument autocomplete. */
+    completeCommand(command: string, partial: string, context?: Context): Promise<{ completions: { value: string; description?: string }[] }>;
     ping(): Promise<Record<string, unknown>>;
     sendEvent(event: string, payload?: Record<string, unknown>, context?: Context): void;
+    /** Every `session/*` request the extension made, in order — for assertions. */
+    readonly sessionCalls: SessionCall[];
     shutdown(): Promise<void>;
     close(): void;
 }
 
 const DEFAULT_CONTEXT: Context = { token: 'test-epoch', tier: 'command' };
 
-export function createTestHost(extension: Extension): TestHost {
+export function createTestHost(extension: Extension, options: CreateTestHostOptions = {}): TestHost {
     const [hostT, extT] = linkedPair();
     const extHandle = extension.connect(extT);
     /** call_id → the caller's onUpdate, so streamed progress reaches the test. */
@@ -46,9 +77,32 @@ export function createTestHost(extension: Extension): TestHost {
         const p = params as ToolUpdateParams;
         updateSinks.get(p.call_id)?.(p);
     });
+    // Answer ext→host `ui/request`. Default mimics a headless frontend (NoUI).
+    host.setRequestHandler(method.UI_REQUEST, async (params) => {
+        if (!options.onUiRequest) throw new RpcError(errorCode.NoUI, 'no UI available (headless test host)');
+        return options.onUiRequest(params as UiRequestParams);
+    });
     // Extension notifications the host just observes in tests.
     host.setNotificationHandler(method.LOG, () => {});
     host.setNotificationHandler(method.REGISTRY_UPDATE, () => {});
+
+    // Service ext→host `session/*` requests, enforcing the same command-tier
+    // guard the real host does (event-tier → -32003) so a demo's session calls
+    // are exercised realistically. Every call is recorded for assertions.
+    const sessionCalls: SessionCall[] = [];
+    const sessionHandler = (params: unknown) => {
+        const p = (params ?? {}) as Record<string, unknown>;
+        const tier = (p.context as { tier?: string } | undefined)?.tier;
+        if (tier !== 'command') throw new RpcError(errorCode.ContextViolation, 'session action requires a command-tier context');
+        return p;
+    };
+    for (const m of [method.SESSION_SEND_MESSAGE, method.SESSION_SEND_USER_MESSAGE, method.SESSION_APPEND_ENTRY]) {
+        host.setRequestHandler(m, (params) => {
+            const recorded = sessionHandler(params);
+            sessionCalls.push({ method: m, params: recorded });
+            return {};
+        });
+    }
     hostT.start((frame) => host.receive(frame));
 
     return {
@@ -79,6 +133,17 @@ export function createTestHost(extension: Extension): TestHost {
         callHook(hook, input, context) {
             return host.request<HookOutcome>(method.HOOK, { hook, input, context: context ?? DEFAULT_CONTEXT });
         },
+        runCommand(command, args, context) {
+            return host.request<CommandExecuteResult>(method.COMMAND_EXECUTE, {
+                command,
+                context: context ?? DEFAULT_CONTEXT,
+                ...(args ? { arguments: args } : {}),
+            });
+        },
+        completeCommand(command, partial, context) {
+            return host.request(method.COMMAND_COMPLETE, { command, context: context ?? DEFAULT_CONTEXT, partial });
+        },
+        sessionCalls,
         ping() {
             return host.request<Record<string, unknown>>(method.PING, {});
         },