fix(tui): prewarm agent runtime before first send

vincentkoc · vincentkoc · commit aaa2bb980b2a · 2026-05-26T18:20:51.000+01:00
diff --git a/src/tui/embedded-backend.ts b/src/tui/embedded-backend.ts
@@ -1,8 +1,13 @@
 import { randomUUID } from "node:crypto";
 import { agentCommandFromIngress } from "../agents/agent-command.js";
-import { resolveSessionAgentId } from "../agents/agent-scope.js";
+import {
+  resolveAgentDir,
+  resolveAgentWorkspaceDir,
+  resolveSessionAgentId,
+} from "../agents/agent-scope.js";
 import { ensureContextWindowCacheLoaded } from "../agents/context.js";
 import { DEFAULT_PROVIDER } from "../agents/defaults.js";
+import { prewarmAgentHarnessRuntime } from "../agents/harness/prewarm.js";
 import {
   buildAllowedModelSet,
   buildConfiguredModelCatalog,
@@ -385,7 +390,6 @@ export class EmbeddedTuiBackend implements TuiBackend {
         catalog,
       });
     }
-
     return {
       sessionKey: opts.sessionKey,
       sessionId,
@@ -396,6 +400,28 @@ export class EmbeddedTuiBackend implements TuiBackend {
     };
   }
 
+  async prewarmAgentRuntime(opts: { sessionKey: string }) {
+    const { cfg, entry } = loadSessionEntry(opts.sessionKey);
+    const sessionId = entry?.sessionId;
+    const sessionAgentId = resolveSessionAgentId({ sessionKey: opts.sessionKey, config: cfg });
+    const resolvedSessionModel = resolveSessionModelRef(cfg, entry, sessionAgentId);
+    const agentRuntimePrewarm = await prewarmAgentHarnessRuntime({
+      cfg,
+      provider: resolvedSessionModel.provider,
+      modelId: resolvedSessionModel.model,
+      agentId: sessionAgentId,
+      sessionKey: opts.sessionKey,
+      sessionId,
+      sessionFile: entry?.sessionFile,
+      agentDir: resolveAgentDir(cfg, sessionAgentId),
+      workspaceDir: entry?.spawnedWorkspaceDir ?? resolveAgentWorkspaceDir(cfg, sessionAgentId),
+      authProfileId: entry?.authProfileOverride,
+      authProfileIdSource: entry?.authProfileOverrideSource,
+      reason: "tui-startup",
+    });
+    return { agentRuntimePrewarm };
+  }
+
   async listSessions(opts?: Parameters<TuiBackend["listSessions"]>[0]): Promise<TuiSessionList> {
     const cfg = getRuntimeConfig();
     const { storePath, store } = loadCombinedSessionStoreForGateway(cfg);
diff --git a/src/tui/tui-backend.ts b/src/tui/tui-backend.ts
@@ -114,7 +114,11 @@ export type TuiBackend = {
     sessionKey: string;
     runId: string;
   }) => Promise<{ ok: boolean; aborted: boolean }>;
-  loadHistory: (opts: { sessionKey: string; limit?: number }) => Promise<unknown>;
+  loadHistory: (opts: {
+    sessionKey: string;
+    limit?: number;
+  }) => Promise<unknown>;
+  prewarmAgentRuntime?: (opts: { sessionKey: string }) => Promise<unknown>;
   listSessions: (opts?: SessionsListParams) => Promise<TuiSessionList>;
   listAgents: () => Promise<TuiAgentsList>;
   patchSession: (opts: SessionsPatchParams) => Promise<SessionsPatchResult>;
diff --git a/src/tui/tui-session-actions.test.ts b/src/tui/tui-session-actions.test.ts
@@ -492,4 +492,43 @@ describe("tui session actions", () => {
     expect(state.currentSessionId).toBe("session-main");
     expect(rememberSessionKey).toHaveBeenCalledWith("agent:main:main");
   });
+
+  it("requests runtime prewarm while loading the first history snapshot", async () => {
+    const loadHistory = vi.fn().mockResolvedValue({
+      sessionId: "session-main",
+      messages: [],
+    });
+    const prewarmAgentRuntime = vi.fn().mockResolvedValue({
+      agentRuntimePrewarm: { status: "warmed" },
+    });
+    const setActivityStatus = vi.fn();
+    const state = createBaseState({ historyLoaded: false });
+
+    const { loadHistory: runLoadHistory } = createTestSessionActions({
+      client: {
+        listSessions: vi.fn().mockResolvedValue({
+          ts: Date.now(),
+          path: "/tmp/sessions.json",
+          count: 0,
+          defaults: {},
+          sessions: [],
+        }),
+        loadHistory,
+        prewarmAgentRuntime,
+      } as unknown as TuiBackend,
+      state,
+      setActivityStatus,
+    });
+
+    await runLoadHistory();
+
+    expect(loadHistory).toHaveBeenCalledWith({
+      sessionKey: "agent:main:main",
+      limit: 200,
+    });
+    expect(prewarmAgentRuntime).toHaveBeenCalledWith({ sessionKey: "agent:main:main" });
+    expect(setActivityStatus).toHaveBeenNthCalledWith(1, "warming runtime");
+    expect(setActivityStatus).toHaveBeenLastCalledWith("idle");
+  });
+
 });
diff --git a/src/tui/tui-session-actions.ts b/src/tui/tui-session-actions.ts
@@ -295,6 +295,16 @@ export function createSessionActions(context: SessionActionContext) {
   };
 
   const loadHistory = async () => {
+    const shouldPrewarmRuntime = !state.historyLoaded;
+    if (shouldPrewarmRuntime) {
+      setActivityStatus("warming runtime");
+    }
+    const prewarmPromise =
+      shouldPrewarmRuntime && client.prewarmAgentRuntime
+        ? client
+            .prewarmAgentRuntime({ sessionKey: state.currentSessionKey })
+            .catch((err) => ({ agentRuntimePrewarm: { status: "failed", error: String(err) } }))
+        : undefined;
     try {
       const history = await client.loadHistory({
         sessionKey: state.currentSessionKey,
@@ -367,11 +377,23 @@ export function createSessionActions(context: SessionActionContext) {
         }
       }
       state.historyLoaded = true;
+      tui.requestRender();
+      const prewarm = (await prewarmPromise) as
+        | { agentRuntimePrewarm?: { status?: string; error?: string } }
+        | undefined;
+      if (prewarm?.agentRuntimePrewarm?.status === "failed") {
+        chatLog.addSystem(
+          `runtime prewarm failed: ${prewarm.agentRuntimePrewarm.error ?? "unknown"}`,
+        );
+      }
       void rememberSessionKey?.(state.currentSessionKey);
     } catch (err) {
       chatLog.addSystem(`history failed: ${String(err)}`);
     }
     await refreshSessionInfo();
+    if (shouldPrewarmRuntime) {
+      setActivityStatus("idle");
+    }
     tui.requestRender();
   };
 
diff --git a/src/tui/tui.test.ts b/src/tui/tui.test.ts
@@ -118,12 +118,25 @@ describe("canSubmitTuiChatMessage", () => {
       }),
     ).toBe(false);
   });
+
+  it("blocks the first submit while runtime prewarm is active", () => {
+    expect(
+      canSubmitTuiChatMessage({
+        local: false,
+        activityStatus: "warming runtime",
+      }),
+    ).toBe(false);
+  });
 });
 
 describe("isTuiBusyActivityStatus", () => {
   it("treats finishing context as a visible busy status", () => {
     expect(isTuiBusyActivityStatus("finishing context")).toBe(true);
   });
+
+  it("treats runtime warmup as a visible busy status", () => {
+    expect(isTuiBusyActivityStatus("warming runtime")).toBe(true);
+  });
 });
 
 describe("resolveTuiShutdownHardExitMs", () => {
diff --git a/src/tui/tui.ts b/src/tui/tui.ts
@@ -374,6 +374,9 @@ export function canSubmitTuiChatMessage(params: {
   pendingOptimisticUserMessage?: boolean;
 }): boolean {
   const pending = Boolean(params.pendingChatRunId) || params.pendingOptimisticUserMessage === true;
+  if (params.activityStatus === "warming runtime") {
+    return false;
+  }
   if (params.activeChatRunId) {
     return params.local === true && params.activityStatus === "finishing context" && !pending;
   }
@@ -385,6 +388,7 @@ const TUI_BUSY_ACTIVITY_STATUSES = new Set([
   "waiting",
   "streaming",
   "running",
+  "warming runtime",
   "finishing context",
 ]);
 
