@pontia/pi-client-plugin 0.2.0 → 0.2.2

package/src/index.ts

@@ -1,19 +1,61 @@
-import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
-import { defaultHookLogFile, loadTurnContext, type EnvLike, type LoadTurnContextResult, type TurnContext } from "./context.js";
+import { AsyncLocalStorage } from "node:async_hooks";
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { connectPi, RpcError, type PiConnection, type ControlInput } from "./control-socket.js";
+import { defaultHookLogFile, type EnvLike, type TurnContext } from "./context.js";
 import { appendDiagnostic, type DiagnosticEntry } from "./diagnostics.js";
-import { pontiaHomeFromEnv, resolvePontiaConnection } from "./discovery.js";
-import { buildSessionContextUsageUpdatedEvent, buildSessionExitedEvent, buildSessionMessageUpdatedEvent, buildSessionReadyEvent, buildTurnCompletedEvent, buildTurnFailedEvent, buildTurnInterruptedEvent, buildTurnOutputEvent, buildTurnStartedEvent, contextUsageFromPiHook, type InternalEvent, type PiTopologyContext, type PiTopologyEntryKind, type SessionMessageUpdatedReason } from "./events.js";
-import { asRecord, optionalString, parseJsonResponse } from "./internal-api.js";
-import { hasTmuxPaneEnvironment, isPontiaManagedTmuxPane, loadPontiaManagedRuntimeIdentity, type ManagedRuntimeIdentity } from "./managed-runtime.js";
-import { agentEndWasInterrupted, assistantDeltaFromEvent, assistantTextFromMessage, errorMessageFromAgentEnd, isTranscriptBoundaryMessageUpdate, lastAssistantTextFromMessages } from "./pi-message.js";
+import { pontiaHomeFromEnv } from "./discovery.js";
+import {
+  buildSessionContextUsageUpdatedEvent,
+  buildSessionExitedEvent,
+  buildSessionMessageUpdatedEvent,
+  buildSessionReadyEvent,
+  buildTurnCompletedEvent,
+  buildTurnFailedEvent,
+  buildTurnInterruptedEvent,
+  buildTurnOutputEvent,
+  buildTurnStartedEvent,
+  contextUsageFromPiHook,
+  type InternalEvent,
+  type PiTopologyContext,
+  type PiTopologyEntryKind,
+  type SessionMessageUpdatedReason,
+} from "./events.js";
+import { asRecord, optionalString } from "./values.js";
+import {
+  completeToolCallFromMessageUpdate,
+  LiveOutputPublisher,
+  type LiveOutputPublisherLike,
+} from "./live-output.js";
+import {
+  hasTmuxPaneEnvironment,
+  isPontiaManagedTmuxPane,
+  loadPontiaManagedRuntimeIdentity,
+  type ManagedRuntimeIdentity,
+} from "./managed-runtime.js";
+import {
+  agentEndWasInterrupted,
+  assistantDeltaFromEvent,
+  assistantTextFromMessage,
+  errorMessageFromAgentEnd,
+  isTranscriptBoundaryMessageUpdate,
+  lastAssistantTextFromMessages,
+} from "./pi-message.js";
 import { loadProfileSystemPrompt } from "./profile.js";
 import { EventReporter, type EventReportResult } from "./reporter.js";
-import { bindSession, loadExistingSessionContext, piSessionDetailsFromHookContext, type PiSessionDetails } from "./runtime-binding.js";
+import {
+  bindSession,
+  loadExistingSessionContext,
+  piSessionDetailsFromHookContext,
+  type PiSessionDetails,
+} from "./runtime-binding.js";
 import type { SessionContext } from "./session.js";
 import { isActiveRegisteredWorkspace } from "./workspace.js";
 
 interface ReporterLike {
-  report(context: { internalEventUrl: string }, event: InternalEvent): Promise<EventReportResult | boolean>;
+  report(
+    context: { runtimeInstanceId: string },
+    event: InternalEvent,
+  ): Promise<EventReportResult | boolean>;
 }
 
 function reportAccepted(result: EventReportResult | boolean): boolean {
@@ -28,18 +70,19 @@ function isPersistentTuiContext(ctx: unknown): boolean {
 
 export interface PontiaPiExtensionDependencies {
   env?: EnvLike;
-  loadContext?: (env: EnvLike, sessionContext?: SessionContext) => Promise<LoadTurnContextResult>;
   makeReporter?: (logFile: string) => ReporterLike;
   logDiagnostic?: (logFile: string, entry: DiagnosticEntry) => Promise<void>;
-  fetch?: typeof fetch;
   loadManagedRuntime?: (env: EnvLike) => Promise<ManagedRuntimeIdentity | undefined>;
   isManagedPane?: (env: EnvLike) => Promise<boolean>;
+  makeLiveOutputPublisher?: (context: TurnContext & { turnId: string }) => LiveOutputPublisherLike;
+  connectPi?: typeof connectPi;
 }
 
 interface ActiveTurnState {
   context: TurnContext & { turnId: string };
   logFile: string;
   reporter: ReporterLike;
+  liveOutput: LiveOutputPublisherLike;
   output: string;
   ended: boolean;
 }
@@ -68,13 +111,12 @@ const NON_MESSAGE_ENTRY_KINDS = new Set([
 function topologyEntryKind(entry: Record<string, unknown>): PiTopologyEntryKind {
   if (entry.type !== "message") {
     return typeof entry.type === "string" && NON_MESSAGE_ENTRY_KINDS.has(entry.type)
-      ? entry.type as PiTopologyEntryKind
+      ? (entry.type as PiTopologyEntryKind)
       : "other";
   }
   const message = entry.message;
-  const role = message && typeof message === "object"
-    ? (message as Record<string, unknown>).role
-    : undefined;
+  const role =
+    message && typeof message === "object" ? (message as Record<string, unknown>).role : undefined;
   if (role === "user") return "user_message";
   if (role === "assistant") return "assistant_message";
   if (role === "toolResult") return "tool_result_message";
@@ -103,7 +145,10 @@ function topologyContextFromHookContext(ctx: unknown): PiTopologyContext | undef
   }
 }
 
-export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPiExtensionDependencies = {}): void {
+export function createPontiaPiExtension(
+  pi: ExtensionAPI,
+  dependencies: PontiaPiExtensionDependencies = {},
+): void {
   const sourceEnv = dependencies.env ?? process.env;
   if (!pontiaHomeFromEnv(sourceEnv) || !hasTmuxPaneEnvironment(sourceEnv)) return;
   // Environment extensions may load PONTIA_HOME after extension registration.
@@ -111,12 +156,26 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
   const currentEnv = () => ({ ...sourceEnv, PONTIA_HOME: currentPontiaHome() });
   const currentHookLogFile = () => defaultHookLogFile(currentPontiaHome());
 
-  const contextLoader = dependencies.loadContext ?? ((contextEnv, sessionContext) => loadTurnContext(contextEnv, { sessionContext }));
-  const makeReporter = dependencies.makeReporter ?? ((logFile: string) => new EventReporter({ logFile }));
+  const rpcConnection: Pick<PiConnection, "request"> = {
+    request(method, params) {
+      if (!controlSocket) return Promise.reject(new Error("Pi connection is unavailable"));
+      return controlSocket.request(method, params);
+    },
+  };
+  const makeReporter =
+    dependencies.makeReporter ??
+    ((logFile: string) =>
+      new EventReporter({
+        logFile,
+        connection: rpcConnection,
+      }));
   const logDiagnostic = dependencies.logDiagnostic ?? appendDiagnostic;
-  const fetchImpl = dependencies.fetch ?? fetch;
   const loadManagedRuntime = dependencies.loadManagedRuntime ?? loadPontiaManagedRuntimeIdentity;
   const isManagedPane = dependencies.isManagedPane ?? isPontiaManagedTmuxPane;
+  const makeLiveOutputPublisher =
+    dependencies.makeLiveOutputPublisher ??
+    ((context: TurnContext & { turnId: string }) =>
+      new LiveOutputPublisher(context, { connection: rpcConnection }));
 
   let activeTurn: ActiveTurnState | undefined;
   let readyReported = false;
@@ -126,6 +185,169 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
   let managedPaneConfirmed = false;
   let lastContextUsageJson: string | undefined;
   let pendingPrompt: string | undefined;
+  let controlSocket: PiConnection | undefined;
+  let controlGeneration = 0;
+  let piContext: ExtensionContext | undefined;
+  let changingModel = false;
+  const directInput = new AsyncLocalStorage<{
+    submission: ControlInput;
+    sessionContext: SessionContext;
+    consumed: boolean;
+  }>();
+  let queuedControlInput: ControlInput | undefined;
+
+  const modelId = (model: { provider: string; id: string }) => `${model.provider}/${model.id}`;
+
+  async function reportModel(): Promise<void> {
+    const context = boundSessionContext;
+    const model = piContext?.model;
+    if (!context || !model || !readyReported || reportingDisabled) return;
+    const result = await makeReporter(currentHookLogFile()).report(context, {
+      session_id: context.sessionId,
+      type: "session.model_updated",
+      data: { model: modelId(model), runtime_instance_id: context.runtimeInstanceId },
+    });
+    if (!reportAccepted(result)) throw new Error("Pi model observation was not accepted");
+  }
+
+  function sendControlInput(submission: ControlInput): void {
+    if (!boundSessionContext) throw new Error("Pi session is not bound");
+    try {
+      directInput.run({ submission, sessionContext: boundSessionContext, consumed: false }, () =>
+        pi.sendUserMessage(submission.input),
+      );
+    } catch (error) {
+      throw new RpcError(
+        -32007,
+        `Pi input delivery uncertain: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  }
+
+  async function controlError(error: unknown): Promise<void> {
+    await logDiagnostic(currentHookLogFile(), {
+      level: "warn",
+      code: "pi_control_unavailable",
+      message: "Pi control channel failed",
+      details: error instanceof Error ? error.message : String(error),
+    });
+  }
+
+  async function closeControlSocket(): Promise<void> {
+    controlGeneration += 1;
+    queuedControlInput = undefined;
+    const socket = controlSocket;
+    controlSocket = undefined;
+    if (socket) await socket.close().catch(controlError);
+  }
+
+  async function registrationConnection(): Promise<PiConnection> {
+    if (controlSocket) return controlSocket;
+    const generation = controlGeneration;
+    const lifecycleContext = () => {
+      if (
+        generation !== controlGeneration ||
+        reportingDisabled ||
+        !boundSessionContext ||
+        !piContext ||
+        !readyReported
+      ) {
+        throw new Error("Pi lifecycle control session is no longer current");
+      }
+      return piContext;
+    };
+    const socket = await (dependencies.connectPi ?? connectPi)(
+      currentPontiaHome(),
+      (error) => {
+        void controlError(error).catch(() => {});
+      },
+      (submission) => {
+        if (generation !== controlGeneration || reportingDisabled || !boundSessionContext) {
+          throw new Error("Pi control session is no longer current");
+        }
+        if (queuedControlInput)
+          throw new RpcError(-32010, "Pi already has pending input; input was not submitted");
+        if (!piContext?.isIdle()) {
+          if (!activeTurn?.ended) throw new RpcError(-32010, "Pi is busy; input was not submitted");
+          queuedControlInput = submission;
+          return;
+        }
+        sendControlInput(submission);
+      },
+      {
+        listModels() {
+          if (
+            generation !== controlGeneration ||
+            reportingDisabled ||
+            !boundSessionContext ||
+            !piContext ||
+            !readyReported
+          ) {
+            throw new Error("Pi model control session is no longer current");
+          }
+          return piContext.modelRegistry.getAvailable().map((model) => ({
+            id: modelId(model),
+            name: model.name,
+            description: model.provider,
+          }));
+        },
+        async setModel(id) {
+          this.listModels();
+          if (changingModel) throw new Error("Pi already has a model change in progress");
+          const model = piContext!.modelRegistry
+            .getAvailable()
+            .find((model) => modelId(model) === id);
+          if (!model) throw new Error("The selected model is not available");
+          changingModel = true;
+          try {
+            if (!(await pi.setModel(model)))
+              throw new RpcError(-32006, "Pi could not authenticate the selected model");
+            this.listModels();
+            // Selecting the current model does not emit model_select; observe it again.
+            await reportModel();
+          } catch (error) {
+            if (error instanceof RpcError) throw error;
+            throw new RpcError(
+              -32007,
+              `Pi model change outcome is unknown: ${error instanceof Error ? error.message : String(error)}`,
+            );
+          } finally {
+            changingModel = false;
+          }
+        },
+        async onReconnect() {
+          if (generation !== controlGeneration)
+            throw new Error("Pi session changed during reconnect");
+          await reportModel();
+        },
+      },
+      (inboxMessageId) => {
+        if (generation !== controlGeneration || reportingDisabled || !boundSessionContext) {
+          throw new Error("Pi control session is no longer current");
+        }
+        // Command dispatch supplies the native context required by navigateTree.
+        pi.sendUserMessage(`/pontia-edit ${inboxMessageId}`, { expandPromptTemplates: true });
+      },
+      {
+        interrupt() {
+          lifecycleContext().abort();
+        },
+        shutdown() {
+          const ctx = lifecycleContext();
+          // Ordinary extension context: abort first, then exit after Pi settles.
+          // Unlike double Ctrl+C, this allows turn cleanup before session_shutdown.
+          ctx.abort();
+          ctx.shutdown();
+        },
+      },
+    );
+    if (generation !== controlGeneration) {
+      await socket.close();
+      throw new Error("Pi session changed during connection initialization");
+    }
+    controlSocket = socket;
+    return socket;
+  }
 
   async function confirmManagedPane(refresh = false): Promise<boolean> {
     if (managedPaneConfirmed && !refresh) return true;
@@ -137,13 +359,10 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
     if (boundSessionContext) return boundSessionContext;
     const runtimeIdentity = await loadManagedRuntime(currentEnv());
     if (!runtimeIdentity) return undefined;
-    const connection = await resolvePontiaConnection({ pontiaHome: currentPontiaHome(), fetch: fetchImpl });
-    if (!connection?.internalEventUrl) return undefined;
     return {
       sessionId: runtimeIdentity.sessionId,
       runtimeInstanceId: runtimeIdentity.runtimeInstanceId,
       clientType: "pi",
-      internalEventUrl: connection.internalEventUrl,
     };
   }
 
@@ -172,41 +391,18 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
         return;
       }
       const loaded = { logFile };
-      const connection = await resolvePontiaConnection({ pontiaHome: currentPontiaHome() });
-      if (!connection?.externalApiToken) {
-        await logDiagnostic(loaded.logFile, {
-          level: "error",
-          code: "branch_replay_stale_context",
-          message: "pontia-edit requires an authenticated Pontia connection",
-        });
-        return;
-      }
-      const resolveUrl = commandContext.internalEventUrl.replace(
-        /\/events\/?$/,
-        "/inbox/branch-replay/resolve",
-      );
 
       let replacementInput: string;
       let targetEntryId: string;
       try {
-        const response = await fetchImpl(resolveUrl, {
-          method: "POST",
-          headers: {
-            Authorization: `Bearer ${connection.externalApiToken}`,
-            "Content-Type": "application/json",
-          },
-          body: JSON.stringify({
-            inbox_message_id: inboxMessageId,
-            session_id: commandContext.sessionId,
-            runtime_instance_id: commandContext.runtimeInstanceId,
-            client_type: "pi",
-          }),
+        const connection = await registrationConnection();
+        const body = await connection.request("branch.resolve", {
+          inbox_message_id: inboxMessageId,
+          session_id: commandContext.sessionId,
+          runtime_instance_id: commandContext.runtimeInstanceId,
+          client_type: "pi",
         });
-        const body = await parseJsonResponse(response);
-        if (!response.ok) {
-          throw new Error(`${response.status} ${response.statusText}`);
-        }
-        const replay = asRecord(asRecord(asRecord(body)?.data)?.branch_replay);
+        const replay = asRecord(asRecord(body)?.branch_replay);
         replacementInput = optionalString(replay?.replacement_input) ?? "";
         targetEntryId = optionalString(replay?.target_entry_id) ?? "";
         if (
@@ -253,7 +449,21 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
 
       try {
         ctx.ui.setEditorText("");
-        await pi.sendUserMessage(replacementInput);
+        const current = await currentManagedSessionContext();
+        if (
+          current?.sessionId !== commandContext.sessionId ||
+          current?.runtimeInstanceId !== commandContext.runtimeInstanceId
+        ) {
+          throw new Error("branch replay Runtime is no longer current");
+        }
+        await directInput.run(
+          {
+            submission: { input: replacementInput, inboxMessageId },
+            sessionContext: current,
+            consumed: false,
+          },
+          () => pi.sendUserMessage(replacementInput),
+        );
       } catch (error) {
         await logDiagnostic(loaded.logFile, {
           level: "error",
@@ -267,11 +477,17 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
 
   async function scheduleMessageRefresh(reason: SessionMessageUpdatedReason): Promise<void> {
     if (!activeTurn || activeTurn.ended) return;
-    await activeTurn.reporter.report(activeTurn.context, buildSessionMessageUpdatedEvent(activeTurn.context, reason));
+    await activeTurn.reporter.report(
+      activeTurn.context,
+      buildSessionMessageUpdatedEvent(activeTurn.context, reason),
+    );
   }
 
   async function reportFinalMessageRefresh(state: ActiveTurnState): Promise<void> {
-    await state.reporter.report(state.context, buildSessionMessageUpdatedEvent(state.context, "final"));
+    await state.reporter.report(
+      state.context,
+      buildSessionMessageUpdatedEvent(state.context, "final"),
+    );
   }
 
   async function reportContextUsageFromHookEvent(event: unknown, ctx?: unknown): Promise<void> {
@@ -281,24 +497,23 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
     const usageJson = JSON.stringify(observation);
     if (usageJson === lastContextUsageJson) return;
     lastContextUsageJson = usageJson;
-    await activeTurn.reporter.report(activeTurn.context, buildSessionContextUsageUpdatedEvent(activeTurn.context, observation.context_usage, observation.model));
+    await activeTurn.reporter.report(
+      activeTurn.context,
+      buildSessionContextUsageUpdatedEvent(activeTurn.context, observation.context_usage),
+    );
   }
 
   pi.on("before_agent_start", async (event) => {
-    const eventRecord = event as unknown as Record<string, unknown>;
-    pendingPrompt = optionalString(eventRecord.prompt);
-    if (reportingDisabled || !await confirmManagedPane()) {
-      return { systemPrompt: typeof eventRecord.systemPrompt === "string" ? eventRecord.systemPrompt : "" };
-    }
-    const currentSystemPrompt = typeof eventRecord.systemPrompt === "string" ? eventRecord.systemPrompt : "";
+    pendingPrompt = event.prompt;
+    if (reportingDisabled || !(await confirmManagedPane())) return;
     try {
       const profilePrompt = await loadProfileSystemPrompt(
-        currentPontiaHome(),
-        fetchImpl,
+        rpcConnection,
         boundSessionContext?.sessionId,
       );
-      if (!profilePrompt) return { systemPrompt: currentSystemPrompt };
-      return { systemPrompt: `${currentSystemPrompt}\n\n${profilePrompt}` };
+      if (profilePrompt) {
+        event.systemPromptOptions.sections.pontia_execution_profile = profilePrompt;
+      }
     } catch (error) {
       await logDiagnostic(currentHookLogFile(), {
         level: "warn",
@@ -306,15 +521,30 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
         message: "failed to append pontia execution profile system prompt",
         details: error instanceof Error ? error.message : String(error),
       });
-      return { systemPrompt: currentSystemPrompt };
     }
   });
 
+  pi.on("model_select", async (_event, ctx) => {
+    piContext = ctx;
+    await reportModel().catch(controlError);
+  });
+
   pi.on("session_start", async (event, ctx) => {
     if (reportingDisabled) return;
     const reason = (event as unknown as Record<string, unknown> | undefined)?.reason;
-    if (readyReported && reason !== "fork" && reason !== "resume") return;
-    if (reason !== "startup" && reason !== "new" && reason !== "resume" && reason !== "fork") return;
+    if (readyReported && reason !== "fork" && reason !== "resume" && reason !== "new") return;
+    if (
+      reason !== "startup" &&
+      reason !== "reload" &&
+      reason !== "new" &&
+      reason !== "resume" &&
+      reason !== "fork"
+    )
+      return;
+    await closeControlSocket();
+    const parentSessionId = boundSessionContext?.sessionId;
+    boundSessionContext = undefined;
+    readyReported = false;
     if (!isPersistentTuiContext(ctx)) {
       reportingDisabled = true;
       return;
@@ -325,26 +555,28 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
       deferredManualSessionDetails = sessionDetails;
       const logFile = currentHookLogFile();
       let context: SessionContext | undefined;
-      const pontiaHome = currentPontiaHome();
       const env = currentEnv();
 
-      const workspaceActive = await isActiveRegisteredWorkspace(pontiaHome, fetchImpl, sessionDetails.clientCwd);
-      if (workspaceActive !== true) {
+      const workspaceActive = await isActiveRegisteredWorkspace(
+        await registrationConnection(),
+        sessionDetails.clientCwd,
+      );
+      if (!workspaceActive) {
+        await closeControlSocket();
         reportingDisabled = true;
         await logDiagnostic(logFile, {
           level: "info",
-          code: workspaceActive === false ? "workspace_not_active" : "workspace_check_unavailable",
-          message: workspaceActive === false
-            ? "current pi workspace is not an active registered pontia workspace; pontia reporting disabled"
-            : "could not verify active registered pontia workspace; pontia reporting disabled",
+          code: "workspace_not_active",
+          message:
+            "current pi workspace is not an active registered pontia workspace; pontia reporting disabled",
           details: { client_cwd: sessionDetails.clientCwd },
         });
         return;
       }
 
       if (reason === "fork") {
-        const parentSessionId = boundSessionContext?.sessionId;
         if (!parentSessionId) {
+          await closeControlSocket();
           await logDiagnostic(logFile, {
             level: "error",
             code: "missing_fork_parent_session",
@@ -352,21 +584,42 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
           });
           return;
         }
-        context = await bindSession(pontiaHome, env, fetchImpl, sessionDetails, { startKind: "fork", parentSessionId });
+        context = await bindSession(await registrationConnection(), env, sessionDetails, {
+          startKind: "fork",
+          parentSessionId,
+        });
         readyReported = false;
       } else {
-        const existingSession = await loadExistingSessionContext(pontiaHome, fetchImpl, sessionDetails);
+        const existingSession = await loadExistingSessionContext(
+          await registrationConnection(),
+          sessionDetails,
+        );
+        if (reason === "reload" && !existingSession) {
+          reportingDisabled = true;
+          deferredManualSessionDetails = undefined;
+          await closeControlSocket();
+          await logDiagnostic(logFile, {
+            level: "error",
+            code: "reload_session_context_missing",
+            message: "pi reload could not recover the existing pontia runtime binding",
+            details: { client_session_key: sessionDetails.clientSessionKey },
+          });
+          return;
+        }
         if (
-          existingSession
-          && ["idle", "busy", "interrupted"].includes(existingSession.sessionState)
+          reason !== "reload" &&
+          existingSession &&
+          ["idle", "busy", "interrupted"].includes(existingSession.sessionState)
         ) {
+          await closeControlSocket();
           reportingDisabled = true;
           boundSessionContext = undefined;
           readyReported = false;
           await logDiagnostic(logFile, {
             level: "info",
             code: "duplicate_active_client_session",
-            message: "native pi session is already bound to an active pontia session; duplicate TUI reporting disabled",
+            message:
+              "native pi session is already bound to an active pontia session; duplicate TUI reporting disabled",
             details: {
               client_session_key: sessionDetails.clientSessionKey,
               session_id: existingSession.sessionId,
@@ -375,22 +628,48 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
           });
           return;
         }
-        if (!existingSession && !await confirmManagedPane()) {
+        if (!existingSession && !(await confirmManagedPane())) {
+          await closeControlSocket();
           boundSessionContext = undefined;
           readyReported = false;
           return;
         }
 
-        context = await bindSession(pontiaHome, env, fetchImpl, sessionDetails, {
-          runtimeInstanceId: existingSession?.runtimeInstanceId,
+        context = await bindSession(await registrationConnection(), env, sessionDetails, {
+          runtimeInstanceId:
+            reason === "reload" || existingSession?.sessionState === "starting"
+              ? existingSession?.runtimeInstanceId
+              : undefined,
         });
-        if (reason === "resume" || reason === "new") readyReported = false;
+        if (
+          reason === "reload" &&
+          (context?.sessionId !== existingSession!.sessionId ||
+            context.runtimeInstanceId !== existingSession!.runtimeInstanceId)
+        ) {
+          reportingDisabled = true;
+          deferredManualSessionDetails = undefined;
+          throw new Error("pi reload runtime binding changed identity");
+        }
+        if (reason === "reload" || reason === "resume" || reason === "new") readyReported = false;
+      }
+      if (!context || !(await confirmManagedPane(true))) {
+        if (reason === "reload") {
+          reportingDisabled = true;
+          deferredManualSessionDetails = undefined;
+        }
+        await closeControlSocket();
+        return;
       }
-      if (!context || !await confirmManagedPane(true)) return;
 
       boundSessionContext = context;
-      readyReported = reportAccepted(await makeReporter(logFile).report(context, buildSessionReadyEvent(context)));
+      piContext = ctx;
+      controlSocket!.registered(context);
+      readyReported = reportAccepted(
+        await makeReporter(logFile).report(context, buildSessionReadyEvent(context)),
+      );
+      await reportModel().catch(controlError);
     } catch (error) {
+      await closeControlSocket();
       const logFile = currentHookLogFile();
       await logDiagnostic(logFile, {
         level: "error",
@@ -402,14 +681,18 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
   });
 
   pi.on("session_shutdown", async (event) => {
-    if (reportingDisabled || !await confirmManagedPane()) return;
-    const reason = (event as unknown as Record<string, unknown> | undefined)?.reason;
-    if (reason !== "quit" && reason !== "new" && reason !== "resume" && reason !== "fork") return;
-
+    piContext = undefined;
+    readyReported = false;
     try {
+      if (reportingDisabled || !(await confirmManagedPane())) return;
+      const reason = (event as unknown as Record<string, unknown> | undefined)?.reason;
+      if (reason !== "quit" && reason !== "new" && reason !== "resume" && reason !== "fork") return;
       const logFile = currentHookLogFile();
       if (!boundSessionContext) return;
-      await makeReporter(logFile).report(boundSessionContext, buildSessionExitedEvent(boundSessionContext, reason));
+      await makeReporter(logFile).report(
+        boundSessionContext,
+        buildSessionExitedEvent(boundSessionContext, reason),
+      );
     } catch (error) {
       const logFile = currentHookLogFile();
       await logDiagnostic(logFile, {
@@ -418,6 +701,8 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
         message: "failed to report pontia exited signal",
         details: error instanceof Error ? error.message : String(error),
       });
+    } finally {
+      await closeControlSocket();
     }
   });
 
@@ -426,75 +711,81 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
     const previousLeafId = leafIdFromHookContext(ctx);
     if (reportingDisabled) return;
     try {
-      const loaded = await contextLoader(currentEnv(), boundSessionContext);
-      let turnContext: TurnContext | undefined;
-      let logFile: string;
-      if (loaded.ok) {
-        turnContext = loaded.context;
-        logFile = loaded.logFile;
-      } else if (loaded.silent) {
-        if (!loaded.logFile) {
+      const dispatch = directInput.getStore();
+      const submission =
+        dispatch && !dispatch.consumed && dispatch.sessionContext === boundSessionContext
+          ? dispatch.submission
+          : undefined;
+      if (dispatch) dispatch.consumed = true;
+      const logFile = currentHookLogFile();
+      if (!boundSessionContext) {
+        const hookSessionDetails = piSessionDetailsFromHookContext(ctx);
+        const sessionDetails = hookSessionDetails.clientSessionKey
+          ? hookSessionDetails
+          : deferredManualSessionDetails;
+        if (!sessionDetails) {
           activeTurn = undefined;
           pendingPrompt = undefined;
           return;
         }
-        logFile = loaded.logFile;
-        if (!boundSessionContext) {
-          const hookSessionDetails = piSessionDetailsFromHookContext(ctx);
-          const sessionDetails = hookSessionDetails.clientSessionKey ? hookSessionDetails : deferredManualSessionDetails;
-          if (!sessionDetails) {
-            activeTurn = undefined;
-            pendingPrompt = undefined;
-            return;
-          }
-          const pontiaHome = currentPontiaHome();
-          const workspaceActive = await isActiveRegisteredWorkspace(pontiaHome, fetchImpl, sessionDetails.clientCwd);
-          if (workspaceActive !== true) {
-            reportingDisabled = true;
-            await logDiagnostic(logFile, {
-              level: "info",
-              code: workspaceActive === false ? "workspace_not_active" : "workspace_check_unavailable",
-              message: workspaceActive === false
-                ? "current pi workspace is not an active registered pontia workspace; pontia reporting disabled"
-                : "could not verify active registered pontia workspace; pontia reporting disabled",
-              details: { client_cwd: sessionDetails.clientCwd },
-            });
-            activeTurn = undefined;
-            pendingPrompt = undefined;
-            return;
-          }
-          boundSessionContext = await bindSession(pontiaHome, currentEnv(), fetchImpl, sessionDetails);
-          if (boundSessionContext && !readyReported) {
-            readyReported = reportAccepted(await makeReporter(logFile).report(boundSessionContext, buildSessionReadyEvent(boundSessionContext)));
-          }
-        }
-        if (boundSessionContext) {
-          turnContext = {
-            sessionId: boundSessionContext.sessionId,
-            runtimeInstanceId: boundSessionContext.runtimeInstanceId,
-            clientType: "pi",
-            internalEventUrl: boundSessionContext.internalEventUrl,
-            input: pendingPrompt,
-          };
-        } else {
+        const workspaceActive = await isActiveRegisteredWorkspace(
+          await registrationConnection(),
+          sessionDetails.clientCwd,
+        );
+        if (!workspaceActive) {
+          await closeControlSocket();
+          reportingDisabled = true;
+          await logDiagnostic(logFile, {
+            level: "info",
+            code: "workspace_not_active",
+            message:
+              "current pi workspace is not an active registered pontia workspace; pontia reporting disabled",
+            details: { client_cwd: sessionDetails.clientCwd },
+          });
           activeTurn = undefined;
           pendingPrompt = undefined;
           return;
         }
-      } else {
+        boundSessionContext = await bindSession(
+          await registrationConnection(),
+          currentEnv(),
+          sessionDetails,
+        );
+        if (boundSessionContext && !readyReported) {
+          piContext = ctx;
+          controlSocket!.registered(boundSessionContext);
+          readyReported = reportAccepted(
+            await makeReporter(logFile).report(
+              boundSessionContext,
+              buildSessionReadyEvent(boundSessionContext),
+            ),
+          );
+          await reportModel().catch(controlError);
+        }
+      }
+      if (!boundSessionContext) {
         activeTurn = undefined;
-        if (!loaded.silent && ctx?.hasUI) ctx.ui.notify(`pontia: ${loaded.reason}`, "warning");
+        pendingPrompt = undefined;
         return;
       }
+      const turnContext: TurnContext = {
+        sessionId: boundSessionContext.sessionId,
+        runtimeInstanceId: boundSessionContext.runtimeInstanceId,
+        clientType: "pi",
+        ...(submission ?? { input: pendingPrompt }),
+      };
 
       pendingPrompt = undefined;
-      if (!await confirmManagedPane(true)) {
+      if (!(await confirmManagedPane(true))) {
         activeTurn = undefined;
         return;
       }
       const reporter = makeReporter(logFile);
       lastContextUsageJson = undefined;
-      const started = await reporter.report(turnContext, buildTurnStartedEvent(turnContext, previousLeafId, topologyContext));
+      const started = await reporter.report(
+        turnContext,
+        buildTurnStartedEvent(turnContext, previousLeafId, topologyContext),
+      );
       const canonicalTurnId = typeof started === "boolean" ? turnContext.turnId : started.turnId;
       if (!reportAccepted(started) || !canonicalTurnId) {
         activeTurn = undefined;
@@ -505,14 +796,17 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
         });
         return;
       }
+      const activeContext = { ...turnContext, turnId: canonicalTurnId };
       activeTurn = {
-        context: { ...turnContext, turnId: canonicalTurnId },
+        context: activeContext,
         logFile,
         reporter,
+        liveOutput: makeLiveOutputPublisher(activeContext),
         output: "",
         ended: false,
       };
     } catch (error) {
+      if (!boundSessionContext) await closeControlSocket();
       pendingPrompt = undefined;
       activeTurn = undefined;
       const logFile = currentHookLogFile();
@@ -529,66 +823,99 @@ export function createPontiaPiExtension(pi: ExtensionAPI, dependencies: PontiaPi
     if (!activeTurn || activeTurn.ended) return;
     await reportContextUsageFromHookEvent(event, ctx);
 
-    const fullText = assistantTextFromMessage((event as unknown as Record<string, unknown> | undefined)?.message);
+    const delta = assistantDeltaFromEvent(event);
+    if (delta) activeTurn.liveOutput.appendText(delta);
+    const toolCall = completeToolCallFromMessageUpdate(event);
+    if (toolCall) activeTurn.liveOutput.appendToolCall(toolCall);
+
+    const fullText = assistantTextFromMessage(
+      (event as unknown as Record<string, unknown> | undefined)?.message,
+    );
     if (fullText) {
       activeTurn.output = fullText;
-    } else {
-      const delta = assistantDeltaFromEvent(event);
-      if (delta) activeTurn.output += delta;
+    } else if (delta) {
+      activeTurn.output += delta;
     }
 
     if (isTranscriptBoundaryMessageUpdate(event)) await scheduleMessageRefresh("update");
   });
 
-  pi.on("tool_execution_start", async () => {
-    await scheduleMessageRefresh("update");
+  pi.on("tool_execution_start", async (event) => {
+    if (!event.parentToolCallId) await scheduleMessageRefresh("update");
   });
 
-  pi.on("tool_execution_end", async () => {
-    await scheduleMessageRefresh("update");
+  pi.on("tool_execution_end", async (event) => {
+    if (!event.parentToolCallId) await scheduleMessageRefresh("update");
   });
 
   pi.on("message_end", async (event, ctx) => {
     if (!activeTurn || activeTurn.ended) return;
     await reportContextUsageFromHookEvent(event, ctx);
-    const fullText = assistantTextFromMessage((event as unknown as Record<string, unknown> | undefined)?.message);
+    const fullText = assistantTextFromMessage(
+      (event as unknown as Record<string, unknown> | undefined)?.message,
+    );
     if (fullText) activeTurn.output = fullText;
     await scheduleMessageRefresh("append");
   });
 
+  pi.on("agent_settled", (_event, ctx) => {
+    if (!queuedControlInput || !ctx.isIdle()) return;
+    const submission = queuedControlInput;
+    queuedControlInput = undefined;
+    sendControlInput(submission);
+  });
+
   pi.on("agent_end", async (event, ctx) => {
     if (!activeTurn || activeTurn.ended) return;
     await reportContextUsageFromHookEvent(event, ctx);
     activeTurn.ended = true;
 
     const state = activeTurn;
-    const terminalLeafId = leafIdFromHookContext(ctx);
-    if (agentEndWasInterrupted(event, ctx?.signal)) {
-      await state.reporter.report(state.context, buildTurnInterruptedEvent(state.context, terminalLeafId));
-      await reportFinalMessageRefresh(state);
-      return;
-    }
+    try {
+      const terminalLeafId = leafIdFromHookContext(ctx);
+      if (agentEndWasInterrupted(event, ctx?.signal)) {
+        await state.reporter.report(
+          state.context,
+          buildTurnInterruptedEvent(state.context, terminalLeafId),
+        );
+        await reportFinalMessageRefresh(state);
+        return;
+      }
 
-    const failureMessage = errorMessageFromAgentEnd(event);
-    if (failureMessage) {
-      await state.reporter.report(state.context, buildTurnFailedEvent(state.context, failureMessage, terminalLeafId));
-      await reportFinalMessageRefresh(state);
-      return;
-    }
+      const failureMessage = errorMessageFromAgentEnd(event);
+      if (failureMessage) {
+        await state.reporter.report(
+          state.context,
+          buildTurnFailedEvent(state.context, failureMessage, terminalLeafId),
+        );
+        await reportFinalMessageRefresh(state);
+        return;
+      }
 
-    if (!state.output) {
-      const finalText = lastAssistantTextFromMessages((event as unknown as Record<string, unknown> | undefined)?.messages);
-      if (finalText) state.output = finalText;
-    }
+      if (!state.output) {
+        const finalText = lastAssistantTextFromMessages(
+          (event as unknown as Record<string, unknown> | undefined)?.messages,
+        );
+        if (finalText) state.output = finalText;
+      }
 
-    const output = state.output.trim();
-    if (output.length > 0) {
-      const outputResult = await state.reporter.report(state.context, buildTurnOutputEvent(state.context, output));
-      if (!reportAccepted(outputResult)) return;
-    }
+      const output = state.output.trim();
+      if (output.length > 0) {
+        const outputResult = await state.reporter.report(
+          state.context,
+          buildTurnOutputEvent(state.context, output),
+        );
+        if (!reportAccepted(outputResult)) return;
+      }
 
-    await state.reporter.report(state.context, buildTurnCompletedEvent(state.context, terminalLeafId));
-    await reportFinalMessageRefresh(state);
+      await state.reporter.report(
+        state.context,
+        buildTurnCompletedEvent(state.context, terminalLeafId),
+      );
+      await reportFinalMessageRefresh(state);
+    } finally {
+      await state.liveOutput.close();
+    }
   });
 }