pi-grok-agent 0.1.0

package/src/client.ts

@@ -0,0 +1,75 @@
+// WebSocket transport to the Grok gateway: one JSON-RPC message per frame, bearer secret in the header.
+import type { AnyMessage, Stream } from '@agentclientprotocol/sdk';
+import WebSocket from 'ws';
+
+export type ConnectionOptions = { url: string; secret: string };
+
+export function validateEndpoint(value: string): string {
+  const url = new URL(value);
+  if (!['ws:', 'wss:'].includes(url.protocol) || url.username || url.password || url.search || url.hash) {
+    throw new Error('Use a ws:// or wss:// endpoint without credentials, query, or fragment.');
+  }
+  if (url.protocol === 'ws:' && !['localhost', '127.0.0.1', '[::1]'].includes(url.hostname)) {
+    throw new Error('Use wss:// for a non-loopback server.');
+  }
+  return url.toString();
+}
+
+/** One JSON-RPC message per WebSocket frame; no Pi filesystem/terminal bridge. */
+export async function openSocket(options: ConnectionOptions, signal?: AbortSignal): Promise<{ stream: Stream; close(): void; closed: Promise<void> }> {
+  signal?.throwIfAborted();
+  const url = validateEndpoint(options.url);
+  if (!options.secret.trim()) throw new Error('Grok WebSocket secret is missing.');
+  const socket = new WebSocket(url, {
+    headers: { Authorization: `Bearer ${options.secret}` },
+    handshakeTimeout: 10_000,
+    maxPayload: 16 * 1024 * 1024,
+    followRedirects: false,
+  });
+  let input: ReadableStreamDefaultController<AnyMessage>;
+  let ended = false;
+  const end = (error?: Error) => {
+    if (ended) return;
+    ended = true;
+    if (error) input.error(error); else input.close();
+  };
+  const readable = new ReadableStream<AnyMessage>({
+    start(controller) { input = controller; },
+    cancel() { ended = true; socket.terminate(); }
+  });
+  socket.on('message', (data) => {
+    if (ended) return;
+    try { input.enqueue(JSON.parse(data.toString())); }
+    catch { end(new Error('Invalid JSON from Grok WebSocket.')); socket.terminate(); }
+  });
+  let resolveClosed!: () => void;
+  const closed = new Promise<void>((resolve) => { resolveClosed = resolve; });
+  socket.on('close', () => { end(); resolveClosed(); });
+  socket.on('error', () => end(new Error('Grok WebSocket connection failed. Check the endpoint, server, and secret.')));
+  await new Promise<void>((resolve, reject) => {
+    const cleanup = () => { signal?.removeEventListener('abort', abort); socket.off('open', opened); socket.off('error', failed); socket.off('close', closed); };
+    const opened = () => { cleanup(); resolve(); };
+    const failed = () => { cleanup(); socket.terminate(); reject(new Error('Grok WebSocket handshake failed. Check the endpoint, server, and secret.')); };
+    const closed = () => { cleanup(); reject(new Error('Grok WebSocket closed during connection.')); };
+    const abort = () => { cleanup(); socket.terminate(); reject(new Error('Grok connection cancelled.')); };
+    socket.once('open', opened); socket.once('error', failed); socket.once('close', closed);
+    signal?.addEventListener('abort', abort, { once: true });
+    if (signal?.aborted) abort();
+  });
+  return {
+    stream: {
+      readable,
+      writable: new WritableStream<AnyMessage>({
+        write(message) {
+          return new Promise<void>((resolve, reject) => {
+            socket.send(JSON.stringify(message), (error) => error ? reject(new Error('Grok WebSocket send failed.')) : resolve());
+          });
+        },
+        close() { socket.close(); },
+        abort() { socket.terminate(); },
+      }),
+    },
+    close() { socket.terminate(); },
+    closed,
+  };
+}

package/src/config.ts

@@ -0,0 +1,125 @@
+import { readFile } from 'node:fs/promises';
+import { homedir } from 'node:os';
+import { isAbsolute, join } from 'node:path';
+import { validateEndpoint } from './client.ts';
+
+export const agentDir = process.env.PI_CODING_AGENT_DIR || join(homedir(), '.pi', 'agent');
+export const configPath = join(agentDir, 'grok-ws.json');
+export const defaultSecretFile = join(agentDir, 'grok-ws.secret');
+
+/** Which Pi tools the `grok` model provider offers to Grok in addition to Grok's own harness tools. */
+export type PiToolPolicy = 'none' | 'extensions' | 'all' | string[];
+export const PI_CORE_TOOLS = new Set(['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls']);
+
+/**
+ * How the `grok` model provider answers Grok's native permission prompts when Pi has no UI (`-p`, Fabric workers).
+ * `dialog` (default): ask in Pi's UI; without a UI, deny. `deny`: always reject. `allow`: always allow once.
+ * `reads`: allow read-class prompts, deny the rest. Interactive Pi always shows the dialog.
+ */
+export type HeadlessPermissionPolicy = 'dialog' | 'deny' | 'reads' | 'allow';
+/**
+ * Grok-side permission mode for sessions the provider creates, independent of Pi's `/grok perms` gate.
+ * `default`: Grok's normal rules. `auto`: Grok's auto permission mode (`_meta.autoMode`). A third value is accepted but not documented.
+ */
+export type GrokMode = 'default' | 'auto' | 'yolo';
+
+/**
+ * Gateway guard tiers, in milliseconds. Grok fails OPEN when a client hook times out and waits forever on a
+ * permission prompt, so the gateway answers on Pi's behalf when Pi cannot. Each value is a deadline after which
+ * the gateway answers fail-closed (deny / continue / reject) unless Pi has answered.
+ */
+export type GuardSettings = {
+  /** No `pi/gate-ack` from Pi within this window: Pi is hung or gone. Default 5000. */
+  ackMs?: number;
+  /** Acked without dialog or check: a policy answer is expected promptly. Default 15000. */
+  policyMs?: number;
+  /** Acked with `check: true` (post_tool_use / stop running a command). Default 590000. */
+  checkBudgetMs?: number;
+  /** Acked with `dialog: true` (a human is deciding). Default 600000. */
+  dialogMs?: number;
+};
+/** Grok's own client-hook deadline cap (`MAX_HOOK_TIMEOUT_SECS`); every guard tier for hooks must stay below it. */
+export const GROK_HOOK_CAP_MS = 600_000;
+/** Grok's deadline for pre_tool_use as registered by this package (`CLIENT_HOOKS`); ackMs and policyMs must stay below it. */
+export const GATE_REGISTRATION_MS = 30_000;
+const GUARD_DEFAULTS: Required<GuardSettings> = { ackMs: 5_000, policyMs: 15_000, checkBudgetMs: 590_000, dialogMs: 600_000 };
+
+export function resolveGuard(settings: GuardSettings | undefined, env: NodeJS.ProcessEnv = process.env): Required<GuardSettings> {
+  const pick = (key: keyof GuardSettings, envName: string): number => {
+    const raw = env[envName] ?? settings?.[key];
+    if (raw === undefined || raw === '') return GUARD_DEFAULTS[key];
+    const n = Number(raw);
+    if (!Number.isFinite(n) || n <= 0) throw new Error(`guard.${key} must be a positive number of milliseconds (got ${String(raw)}).`);
+    return Math.floor(n);
+  };
+  const guard = { ackMs: pick('ackMs', 'PI_GROK_ACK_MS'), policyMs: pick('policyMs', 'PI_GROK_POLICY_MS'), checkBudgetMs: pick('checkBudgetMs', 'PI_GROK_CHECK_BUDGET_MS'), dialogMs: pick('dialogMs', 'PI_GROK_DIALOG_MS') };
+  // A tier at or past Grok's own deadline would let Grok fail open first, which defeats the guard.
+  if (guard.ackMs >= GATE_REGISTRATION_MS || guard.policyMs >= GATE_REGISTRATION_MS) throw new Error('guard.ackMs and guard.policyMs must be below the pre_tool_use registration deadline (' + GATE_REGISTRATION_MS + ' ms).');
+  if (guard.checkBudgetMs >= GROK_HOOK_CAP_MS) throw new Error('guard.checkBudgetMs must be below Grok\'s hook cap (' + GROK_HOOK_CAP_MS + ' ms).');
+  if (guard.ackMs > guard.policyMs) throw new Error('guard.ackMs must not exceed guard.policyMs.');
+  return guard;
+}
+
+/** Optional overrides for the Grok-side hook layers. Regexes match Grok tool names. */
+export type HookSettings = {
+  /** Extra Grok tools to deny at pre_tool_use, beyond the capability mirror. */
+  denyGrokTools?: string[];
+  /** Grok tools to allow even when the capability mirror would deny them. */
+  allowGrokTools?: string[];
+  /**
+   * MCP servers (plugin or configured) whose tools count as read-only in a read-only Pi session.
+   * Grok 1.0.41 drops MCP `annotations` (readOnlyHint), so per-tool hints are only honored from `_meta`.
+   */
+  mcpReadOnlyServers?: string[];
+  /** Command run after a Grok edit; `{file}` is replaced. Non-zero exit output becomes additionalContext. Default: built-in syntax checks. */
+  postEditCheck?: string;
+  /** Command run when Grok wants to end its turn; non-zero exit blocks the stop with the output as reason. */
+  stopCheck?: string;
+};
+
+/** The gateway's shared secret, or undefined when the file does not exist yet. Other read errors propagate. */
+export async function readSecretFile(secretFile: string): Promise<string | undefined> {
+  try { return (await readFile(secretFile, 'utf8')).trim(); }
+  catch (error) { if ((error as NodeJS.ErrnoException).code === 'ENOENT') return undefined; throw error; }
+}
+
+export async function readConfig() {
+  let settings: { url?: string; secretFile?: string; piTools?: PiToolPolicy; hooks?: HookSettings; headlessPermissions?: HeadlessPermissionPolicy; guard?: GuardSettings; mediaDir?: string; grokMode?: GrokMode; autoStartGateway?: boolean } = {};
+  try { settings = JSON.parse(await readFile(configPath, 'utf8')); }
+  catch (error) { if ((error as NodeJS.ErrnoException).code !== 'ENOENT') throw error; }
+  const url = validateEndpoint(process.env.GROK_ACP_URL || settings.url || 'ws://127.0.0.1:2419/ws');
+  const configuredFile = settings.secretFile || defaultSecretFile;
+  const secretFile = configuredFile.startsWith('~/') ? join(homedir(), configuredFile.slice(2)) : configuredFile;
+  if (!isAbsolute(secretFile)) throw new Error('grok-ws secretFile must be absolute or start with ~/.');
+  // A missing secret file is not a load error: Pi exits on any extension load failure, and the gateway that creates
+  // the file may not have run yet. The connection reads the file again when it opens (`readSecretFile`).
+  const secret = process.env.GROK_AGENT_SECRET || (await readSecretFile(secretFile)) || '';
+  const piTools: PiToolPolicy = process.env.PI_GROK_PI_TOOLS ? parsePolicy(process.env.PI_GROK_PI_TOOLS) : settings.piTools ?? 'extensions';
+  const hooks: HookSettings = { ...settings.hooks };
+  if (process.env.PI_GROK_STOP_CHECK) hooks.stopCheck = process.env.PI_GROK_STOP_CHECK;
+  if (process.env.PI_GROK_POST_EDIT_CHECK) hooks.postEditCheck = process.env.PI_GROK_POST_EDIT_CHECK;
+  if (process.env.PI_GROK_DENY_TOOLS) hooks.denyGrokTools = process.env.PI_GROK_DENY_TOOLS.split(',').map((s) => s.trim()).filter(Boolean);
+  const headlessPermissions = (process.env.PI_GROK_HEADLESS_PERMISSIONS as HeadlessPermissionPolicy | undefined) ?? settings.headlessPermissions ?? 'dialog';
+  if (!['dialog', 'deny', 'reads', 'allow'].includes(headlessPermissions)) throw new Error(`headlessPermissions must be dialog, deny, reads, or allow (got ${headlessPermissions}).`);
+  const guard = resolveGuard(settings.guard);
+  // Where Pi copies Grok's generated media. Relative paths resolve against the Pi session cwd. Empty string disables the copy.
+  const mediaDir = process.env.PI_GROK_MEDIA_DIR ?? settings.mediaDir ?? '.pi/grok-images';
+  const grokMode = (process.env.PI_GROK_GROK_MODE as GrokMode | undefined) ?? settings.grokMode ?? 'default';
+  if (!['default', 'auto', 'yolo'].includes(grokMode)) throw new Error(`grokMode must be default, auto, or yolo (got ${grokMode}).`);
+  // Start the bundled gateway when nothing listens on the endpoint. PI_GROK_AUTOSTART=0 or autoStartGateway: false turns it off.
+  const autoStartGateway = process.env.PI_GROK_AUTOSTART ? !['0', 'false', 'no', 'off'].includes(process.env.PI_GROK_AUTOSTART.toLowerCase()) : settings.autoStartGateway ?? true;
+  return { url, secret, secretFile, piTools, hooks, headlessPermissions, guard, mediaDir, grokMode, autoStartGateway };
+}
+
+function parsePolicy(value: string): PiToolPolicy {
+  if (value === 'none' || value === 'extensions' || value === 'all') return value;
+  return value.split(',').map((s) => s.trim()).filter(Boolean);
+}
+
+export function selectPiTools<T extends { name: string }>(tools: T[], policy: PiToolPolicy): T[] {
+  if (policy === 'none') return [];
+  if (policy === 'all') return tools;
+  if (policy === 'extensions') return tools.filter((t) => !PI_CORE_TOOLS.has(t.name));
+  const allowed = new Set(policy);
+  return tools.filter((t) => allowed.has(t.name));
+}

package/src/launch.ts

@@ -0,0 +1,55 @@
+// Start the gateway that ships with this package when nothing listens on its loopback endpoint, so a user needs only
+// `pi install npm:pi-grok-agent`. The gateway runs detached and outlives Pi: every Pi process on the machine shares it.
+// Two Pi processes that start one at the same time are safe: the gateway binds its port before it touches a leader,
+// so the loser exits with EADDRINUSE and both connect to the winner.
+import { spawn } from 'node:child_process';
+import { existsSync, mkdirSync, openSync } from 'node:fs';
+import { createConnection } from 'node:net';
+import { homedir } from 'node:os';
+import { basename, dirname, join, sep } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+const packageRoot = join(dirname(fileURLToPath(import.meta.url)), '..');
+
+/** The gateway entry point for this install. Node refuses to strip TypeScript under node_modules, so an npm install runs the compiled copy. */
+export function gatewayEntry(root = packageRoot): string {
+  const source = join(root, 'scripts', 'server.ts');
+  const compiled = join(root, 'dist', 'scripts', 'server.js');
+  if (root.includes(`${sep}node_modules${sep}`) || !existsSync(source)) return compiled;
+  return source;
+}
+
+/** True when something accepts TCP connections on the endpoint's host and port. */
+export function endpointListening(endpoint: string, timeoutMs = 1000): Promise<boolean> {
+  const url = new URL(endpoint);
+  const host = url.hostname === '[::1]' ? '::1' : url.hostname;
+  return new Promise((resolve) => {
+    const socket = createConnection({ host, port: Number(url.port || 80) });
+    const done = (value: boolean) => { socket.destroy(); resolve(value); };
+    socket.setTimeout(timeoutMs, () => done(false));
+    socket.once('connect', () => done(true));
+    socket.once('error', () => done(false));
+  });
+}
+
+/** Start the gateway detached, logging to `<agentDir>/grok-ws.log`, and wait until its port accepts connections. Returns its pid. */
+export async function launchGateway(endpoint: string, logDir: string, timeoutMs = 30_000): Promise<number | undefined> {
+  const entry = gatewayEntry();
+  if (!existsSync(entry)) throw new Error(`Grok gateway not found at ${entry}. Run npm run build in the package, or start the gateway yourself.`);
+  mkdirSync(logDir, { recursive: true, mode: 0o700 });
+  const log = openSync(join(logDir, 'grok-ws.log'), 'a', 0o600);
+  // Inside Pi, execPath is the Node binary that runs Pi. A single-file Pi build would be Pi itself; use `node` then.
+  const nodeBinary = /^node(\.exe)?$/.test(basename(process.execPath)) ? process.execPath : 'node';
+  const child = spawn(nodeBinary, [entry], { detached: true, stdio: ['ignore', log, log], cwd: homedir(), env: process.env });
+  child.unref();
+  // Exited without a port: usually another Pi's launch won the port a moment earlier, so keep waiting a few seconds for it.
+  let exitedAt: number | undefined;
+  child.once('exit', () => { exitedAt = Date.now(); });
+  child.once('error', () => { exitedAt = Date.now(); });
+  const deadline = Date.now() + timeoutMs;
+  while (Date.now() < deadline && (exitedAt === undefined || Date.now() - exitedAt < 5000)) {
+    if (await endpointListening(endpoint, 500)) return exitedAt === undefined ? child.pid : undefined;
+    await new Promise((r) => setTimeout(r, 200));
+  }
+  throw new Error(`Grok gateway did not start within ${Math.round(timeoutMs / 1000)}s. See ${join(logDir, 'grok-ws.log')}.`);
+}

package/src/login.ts

@@ -0,0 +1,44 @@
+// `/grok login`: run Grok Build's own device-code sign-in and report the URL and code to Pi. Grok stores the
+// credential in ~/.grok/auth.json, as a `grok login` in a terminal would; Pi stores nothing.
+import { spawn } from 'node:child_process';
+
+export type DeviceCode = { url: string; code: string };
+
+/** Pull the verification URL and user code from `grok login --device-auth` output. */
+export function parseDeviceCode(output: string): DeviceCode | undefined {
+  const plain = output.replace(/\x1b\[[0-9;]*m/g, '');
+  const url = plain.match(/https:\/\/\S+user_code=([A-Z0-9-]+)/);
+  if (!url) return undefined;
+  return { url: url[0], code: url[1] };
+}
+
+/**
+ * Start the sign-in. `onCode` fires once with the URL and code; the promise settles when Grok's login exits.
+ * Grok may also open the URL in a browser itself. The abort signal stops the login.
+ */
+export function grokLogin(onCode: (device: DeviceCode) => void, options: { binary?: string; signal?: AbortSignal; timeoutMs?: number } = {}): Promise<void> {
+  const binary = options.binary ?? process.env.PI_GROK_BINARY ?? 'grok';
+  return new Promise((resolve, reject) => {
+    const child = spawn(binary, ['login', '--device-auth'], { stdio: ['ignore', 'pipe', 'pipe'], env: process.env });
+    let output = ''; let reported = false;
+    const read = (chunk: Buffer) => {
+      output += chunk.toString();
+      if (reported) return;
+      const device = parseDeviceCode(output);
+      if (device) { reported = true; onCode(device); }
+    };
+    child.stdout.on('data', read);
+    child.stderr.on('data', read);
+    const stop = () => child.kill('SIGTERM');
+    const timer = setTimeout(stop, options.timeoutMs ?? 15 * 60_000);
+    options.signal?.addEventListener('abort', stop, { once: true });
+    child.once('error', (error) => { clearTimeout(timer); reject(new Error(`Could not run ${binary} login: ${error.message}`)); });
+    child.once('exit', (code, signal) => {
+      clearTimeout(timer);
+      options.signal?.removeEventListener('abort', stop);
+      if (code === 0) { resolve(); return; }
+      const tail = output.replace(/\x1b\[[0-9;]*m/g, '').trim().split('\n').slice(-2).join(' ').trim();
+      reject(new Error(`Grok login ${signal ? 'stopped' : `failed (exit ${code})`}${tail ? `: ${tail}` : ''}`));
+    });
+  });
+}

package/src/model/connection.ts

@@ -0,0 +1,239 @@
+// One WebSocket to Grok, many sessions. Grok keeps its full native harness (its own tools,
+// permissions, subagents). Pi tools are offered additively as an HTTP MCP server that the gateway
+// fronts at /mcp/<serverId>; the gateway relays each MCP message back over this socket as an
+// _x.ai/mcp/sdk_call request, routed here by serverId. The stock leader never sees that traffic.
+import { client, type ClientConnection, type InitializeResponse, type NewSessionResponse, type LoadSessionResponse, type SessionNotification, type RequestPermissionRequest, type RequestPermissionResponse } from '@agentclientprotocol/sdk';
+import { openSocket, type ConnectionOptions } from '../client.ts';
+import { readSecretFile } from '../config.ts';
+import { endpointListening, launchGateway } from '../launch.ts';
+import { GATE_REGISTRATION_MS } from '../config.ts';
+
+export type McpToolDefinition = { name: string; description: string; inputSchema: Record<string, unknown> };
+export type McpToolResult = { content: { type: 'text'; text: string }[] | { type: 'image'; data: string; mimeType: string }[] | ({ type: 'text'; text: string } | { type: 'image'; data: string; mimeType: string })[]; isError?: boolean };
+export type SdkCall = { method: string; id: unknown; params?: any };
+
+export interface SessionHandlers {
+  onUpdate(notification: SessionNotification): void;
+  /** Answer one MCP JSON-RPC message from Grok. Return the JSON-RPC `result` or throw for an error. */
+  onMcp(message: SdkCall): Promise<unknown>;
+  onPermission?(request: RequestPermissionRequest): Promise<RequestPermissionResponse>;
+  /**
+   * Blocking client hook (`_x.ai/hooks/run`): pre_tool_use, post_tool_use, stop. `gate.dialog()` tells the gateway a
+   * human is deciding, so it waits the dialog window instead of the short policy window before answering for Pi.
+   */
+  onHookRun?(payload: any, gate?: { dialog(): void }): Promise<Record<string, unknown>>;
+  /** Passive client hook notification (`_x.ai/hooks/event`). */
+  onHookEvent?(payload: any): void;
+  /** Grok's ask_user_question (`_x.ai/ask_user_question`). Default: cancelled. */
+  onQuestion?(request: any): Promise<Record<string, unknown>>;
+  /** Grok's extension session notifications (`_x.ai/session_notification`): turn_completed carries full token usage. */
+  onSessionExt?(update: any): void;
+}
+
+/**
+ * Client hook registration sent in session/new. Grok caps timeouts at 600 s and fails OPEN on expiry.
+ * The gate gets a short deadline; the gateway denies on Pi's behalf well before it (PI_GROK_GATE_DENY_MS).
+ */
+export const CLIENT_HOOKS = {
+  PreToolUse: [{ hookCallbackIds: ['pi-pre'], timeout: GATE_REGISTRATION_MS / 1000 }],
+  PostToolUse: [{ hookCallbackIds: ['pi-post'], timeout: 600 }],
+  PostToolUseFailure: [{ hookCallbackIds: ['pi-post-failure'], timeout: 60 }],
+  Stop: [{ hookCallbackIds: ['pi-stop'], timeout: 600 }],
+};
+
+type SdkCallParams = { serverId: string; sessionId?: string; message: SdkCall };
+
+export class GrokModelConnection {
+  private socket?: Awaited<ReturnType<typeof openSocket>>;
+  private connection?: ClientConnection;
+  private initialized?: InitializeResponse;
+  private readonly sessions = new Map<string, SessionHandlers>();
+  private readonly servers = new Map<string, string>(); // serverId -> sessionId
+  private opening?: Promise<void>;
+  private closed = false;
+  private readonly options: ConnectionOptions & { secretFile?: string; autoStart?: { logDir: string } };
+
+  /**
+   * `secret` may be empty when the gateway has not created its file yet; `secretFile` is read again on each open.
+   * With `autoStart`, an open that finds nothing listening on a loopback `ws://` endpoint starts the bundled gateway.
+   */
+  constructor(options: ConnectionOptions & { secretFile?: string; autoStart?: { logDir: string } }) { this.options = { ...options }; }
+
+  /** Pid of the gateway this connection started, if any. */
+  launchedGateway?: number;
+
+  private async ensureGateway() {
+    if (!this.options.autoStart || !this.options.url.startsWith('ws://')) return;
+    if (await endpointListening(this.options.url)) return;
+    this.launchedGateway = await launchGateway(this.options.url, this.options.autoStart.logDir);
+  }
+
+  /** Fill in the secret from its file on first open, so Pi loads and the gateway may start after it. */
+  private async resolveSecret() {
+    if (this.options.secret.trim() || !this.options.secretFile) return;
+    const secret = await readSecretFile(this.options.secretFile);
+    if (!secret) throw new Error(`Grok gateway secret not found at ${this.options.secretFile}. Start the gateway once (pi-grok-gateway, or npm run server in the clone); it creates the file. Then send the message again.`);
+    this.options.secret = secret;
+  }
+
+  /** Reasons the last connection ended, for status and error text. */
+  lastDrop?: string;
+
+  get isOpen() { return !!this.connection && !this.closed; }
+
+  /** The socket went away underneath us (gateway restart, network). Sessions stay in the map so attach() can session/load them. */
+  private markDropped(reason: string) {
+    if (!this.connection && !this.socket) return;
+    this.lastDrop = reason;
+    this.connection = undefined;
+    this.socket = undefined;
+    this.initialized = undefined;
+    this.generation++;
+  }
+
+  /** Increments on every drop; a session attached under an older generation must session/load again. */
+  generation = 0;
+
+  async open(signal?: AbortSignal) {
+    if (this.connection) return;
+    if (this.opening) return this.opening;
+    this.closed = false;
+    this.opening = (async () => {
+      await this.ensureGateway();
+      await this.resolveSecret();
+      const socket = await openSocket(this.options, signal);
+      this.socket = socket;
+      void socket.closed.then(() => { if (this.socket === socket) this.markDropped('Grok WebSocket closed'); });
+      this.connection = client({ name: 'pi-grok-model' })
+        .onNotification('session/update', ({ params }) => {
+          this.sessions.get(params.sessionId)?.onUpdate(params);
+        })
+        .onRequest('session/request_permission', async ({ params }) => {
+          const handler = this.sessions.get(params.sessionId)?.onPermission;
+          this.ack(`perm:${params.toolCall?.toolCallId ?? ''}`, { dialog: this.hasUI });
+          return handler ? handler(params) : { outcome: { outcome: 'cancelled' } };
+        })
+        .onRequest('_x.ai/ask_user_question', (raw) => raw as any, async ({ params }) => {
+          const handler = this.sessions.get(params.sessionId ?? params.session_id)?.onQuestion;
+          this.ack(`ask:${params.toolCallId ?? params.tool_call_id ?? ''}`, { dialog: this.hasUI });
+          return handler ? handler(params) : { outcome: 'cancelled' };
+        })
+        .onRequest('_x.ai/hooks/run', (raw) => raw as any, async ({ params }) => {
+          const handler = this.sessions.get(params.sessionId ?? params.session_id)?.onHookRun;
+          const event = String(params.hookEventName ?? '');
+          const key = event === 'stop' ? `stop:${params.sessionId ?? params.session_id ?? ''}` : `${event}:${params.toolUseId ?? ''}`;
+          this.ack(key, { check: event === 'post_tool_use' || event === 'stop' });
+          if (handler) return handler(params, { dialog: () => this.ack(key, { dialog: true }) });
+          // No Pi session owns this Grok session (detached by /new or shutdown while a turn was still running).
+          // Pi's capability gate is gone, so tool use is denied rather than left to Grok's own permission mode.
+          if (event === 'pre_tool_use') return { decision: 'deny', reason: 'Pi detached from this Grok session; tool use is denied until a Pi session owns it again.' };
+          return { decision: 'continue' };
+        })
+        .onNotification('_x.ai/session_notification', (raw) => raw as any, ({ params }) => {
+          this.sessions.get(params.sessionId ?? params.session_id)?.onSessionExt?.(params.update);
+        })
+        .onNotification('_x.ai/hooks/event', (raw) => raw as any, ({ params }) => {
+          this.sessions.get(params.sessionId ?? params.session_id)?.onHookEvent?.(params);
+        })
+        .onRequest('_x.ai/mcp/sdk_call', (raw) => raw as SdkCallParams, async ({ params }) => {
+          const sessionId = params.sessionId ?? this.servers.get(params.serverId);
+          const handlers = sessionId ? this.sessions.get(sessionId) : undefined;
+          const { message } = params;
+          if (!handlers) return { jsonrpc: '2.0', id: message.id, error: { code: -32001, message: `no session for server ${params.serverId}` } };
+          try {
+            const result = await handlers.onMcp(message);
+            return { jsonrpc: '2.0', id: message.id, result };
+          } catch (error) {
+            return { jsonrpc: '2.0', id: message.id, error: { code: -32603, message: error instanceof Error ? error.message : String(error) } };
+          }
+        })
+        .connect(socket.stream);
+      const agent = this.connection.agent;
+      this.initialized = await agent.request('initialize', {
+        protocolVersion: 1,
+        clientInfo: { name: 'pi-grok-model', version: '0.1.0' },
+        clientCapabilities: { fs: { readTextFile: false, writeTextFile: false }, terminal: false },
+        _meta: { 'x.ai/mcp/sdk': true },
+      });
+      if ((this.initialized.authMethods ?? []).some((m) => m.id === 'cached_token')) {
+        await agent.request('authenticate', { methodId: 'cached_token' });
+      } else {
+        // Grok offers `cached_token` only with a stored login. Drop this connection so the turn after a login
+        // initializes again and sees the new credential.
+        this.socket?.close();
+        this.markDropped('Grok Build is not signed in');
+        throw new Error('Grok Build is not signed in. Run /grok login, approve the code in your browser, then send the message again.');
+      }
+    })().finally(() => { this.opening = undefined; });
+    return this.opening;
+  }
+
+  get agent() {
+    if (!this.connection) throw new Error('Grok model connection is not open.');
+    return this.connection.agent;
+  }
+
+  /** Whether a human can answer dialogs; the gateway extends permission deadlines when true. */
+  hasUI = false;
+
+  /**
+   * Tell the gateway Pi is alive and what it is doing with a reverse request (`pi/gate-ack`).
+   * The gateway consumes this; it never reaches Grok. Missing acks make the gateway fail closed. A later ack for
+   * the same key moves the request to that tier's deadline (for example a hook that turns into a dialog).
+   */
+  private ack(key: string, state: { dialog?: boolean; check?: boolean }) {
+    void this.connection?.agent.notify('pi/gate-ack', { key, ...state }).catch(() => {});
+  }
+
+  /** HTTP origin of the gateway that fronts this WebSocket (ws://host:port -> http://host:port). */
+  get mcpBaseUrl() { return new URL(this.options.url).origin.replace(/^ws/, 'http'); }
+
+  /** Create or load a Grok session with Grok's native harness intact. `offerPiTools` adds the Pi-hosted MCP server. */
+  async attachSession(input: { sessionId?: string; cwd: string; serverId: string; serverName: string; rules?: string; handlers: SessionHandlers; toolTimeoutMs?: number; offerPiTools: boolean; hooks?: boolean; grokMode?: 'default' | 'auto' | 'yolo' }) {
+    const _meta: Record<string, unknown> = { yoloMode: input.grokMode === 'yolo', ...(input.grokMode === 'auto' ? { autoMode: true } : {}) };
+    if (input.hooks !== false) _meta['x.ai/hooks'] = CLIENT_HOOKS;
+    const mcpServers: unknown[] = [];
+    if (input.offerPiTools) {
+      mcpServers.push({ type: 'http', name: input.serverName, url: `${this.mcpBaseUrl}/mcp/${input.serverId}`, headers: [] });
+      _meta.mcpConfig = { [input.serverName]: { toolTimeoutMs: input.toolTimeoutMs ?? 6 * 60 * 60 * 1000 } };
+    }
+    if (input.rules) _meta.rules = input.rules;
+    const params = { cwd: input.cwd, mcpServers, _meta };
+    this.servers.set(input.serverId, input.sessionId ?? '');
+    let session: NewSessionResponse | LoadSessionResponse;
+    let sessionId: string;
+    if (input.sessionId) {
+      this.sessions.set(input.sessionId, input.handlers);
+      session = await this.agent.request<LoadSessionResponse>('session/load', { ...params, sessionId: input.sessionId });
+      sessionId = input.sessionId;
+    } else {
+      const created = await this.agent.request<NewSessionResponse>('session/new', params);
+      session = created;
+      sessionId = created.sessionId;
+      this.sessions.set(sessionId, input.handlers);
+    }
+    this.servers.set(input.serverId, sessionId);
+    return { sessionId, response: session };
+  }
+
+  detachSession(sessionId: string, serverId: string) {
+    this.sessions.delete(sessionId);
+    this.servers.delete(serverId);
+  }
+
+  /** Close the socket but keep this connection usable: the next `open()` reconnects and sessions `session/load` themselves. */
+  drop(reason = 'reconnect requested') {
+    const socket = this.socket;
+    this.markDropped(reason);
+    socket?.close();
+  }
+
+  async close() {
+    this.closed = true;
+    this.sessions.clear();
+    this.servers.clear();
+    this.connection = undefined;
+    this.socket?.close();
+    this.socket = undefined;
+  }
+}