Merge branch 'feat/tui-environment-background-connect-start'

Author: 7Sageer

.changeset/environment-startup-ux.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": minor
+---
+
+Starting with `--environment` no longer freezes the interface while the remote environment connects and reports an unknown environment id immediately; subagent cards now show the environment they run in.

apps/kimi-code/src/tui/components/messages/agent-group.ts

@@ -304,6 +304,7 @@ function formatStats(snap: ToolCallSubagentSnapshot): string {
   const parts: string[] = [];
   if (snap.model !== undefined) parts.push(snap.model);
   if (snap.effort !== undefined) parts.push(snap.effort);
+  if (snap.environment !== undefined) parts.push(`env ${snap.environment}`);
   parts.push(`${String(snap.toolCount)} tool${snap.toolCount === 1 ? '' : 's'}`);
   if (snap.elapsedSeconds !== undefined) parts.push(formatElapsed(snap.elapsedSeconds));
   if (snap.tokens > 0) parts.push(formatTokens(snap.tokens));

apps/kimi-code/src/tui/components/messages/tool-call.ts

@@ -114,6 +114,8 @@ export interface ToolCallSubagentSnapshot {
   readonly model?: string;
   /** Thinking effort, present only for concrete levels (on/off hidden). */
   readonly effort?: string;
+  /** Remote environment the subagent runs in (the Agent tool's `environment` arg), absent for local. */
+  readonly environment?: string;
   readonly phase: SubagentPhase | undefined;
   readonly toolCount: number;
   readonly elapsedSeconds: number | undefined;
@@ -158,6 +160,11 @@ function str(v: unknown): string {
   return typeof v === 'string' ? v : '';
 }
 
+function subagentEnvironment(args: Readonly<Record<string, unknown>>): string | undefined {
+  const raw = str(args['environment']).trim();
+  return raw.length > 0 && raw !== 'local' ? raw : undefined;
+}
+
 function formatSubagentContextTokens(contextTokens: number | undefined): string | undefined {
   if (contextTokens === undefined || contextTokens <= 0) return undefined;
   return `${formatTokenCount(contextTokens)} tok`;
@@ -1135,6 +1142,7 @@ export class ToolCallComponent extends Container {
       agentName: this.subagentAgentName,
       model: this.subagentModel,
       effort: this.subagentEffort,
+      environment: subagentEnvironment(this.toolCall.args),
       phase: derivedPhase,
       toolCount: finished,
       elapsedSeconds: this.getSubagentElapsedSeconds(),
@@ -2098,6 +2106,8 @@ export class ToolCallComponent extends Container {
     const parts: string[] = [];
     if (this.subagentModel !== undefined) parts.push(this.subagentModel);
     if (this.subagentEffort !== undefined) parts.push(this.subagentEffort);
+    const environment = subagentEnvironment(this.toolCall.args);
+    if (environment !== undefined) parts.push(`env ${environment}`);
     parts.push(`${String(this.subToolActivities.size)} tool${this.subToolActivities.size === 1 ? '' : 's'}`);
     const elapsed = this.getSubagentElapsedSeconds();
     if (elapsed !== undefined) parts.push(formatElapsed(elapsed));

apps/kimi-code/src/tui/kimi-tui.ts

@@ -15,10 +15,12 @@ import type {
   PluginCommandDef,
   PromptPart,
   Session,
+  SessionEnvironmentType,
   SkillSummary,
   TokenUsage,
   TurnEndedEvent,
   TurnStartedEvent,
+  WorkspaceEnvironmentDeclarationInfo,
   WorkspaceTrustInfo,
 } from '@moonshot-ai/kimi-code-sdk';
 import { isTelemetryDisabledByEnv } from '@moonshot-ai/kimi-telemetry';
@@ -313,6 +315,12 @@ interface SendMessageOptions {
    * it — the queue item owns the raw ids and re-leases at dequeue.
    */
   readonly lease?: StagingLease;
+  /**
+   * Set when the caller already appended the user transcript entry and ran
+   * beginSessionRequest (the pre-session announce in sendNormalUserInput):
+   * dispatch skips both and goes straight to the prompt.
+   */
+  readonly announced?: boolean;
 }
 
 /** How long the one-shot "moved to background" footer hint stays visible. */
@@ -340,6 +348,8 @@ export class KimiTUI {
   state: TUIState;
   /** In-flight lazy session creation (v2 engine), shared by concurrent first-use triggers. */
   private ensureSessionPromise: Promise<Session | undefined> | null = null;
+  /** Type of the validated `--environment` startup declaration, used for the synthetic connecting footer slot. */
+  private startupEnvironmentType: SessionEnvironmentType | undefined;
   private readonly cacheHint = new CacheHintController(this);
   /** Staged prompt media lifecycle (daemon uploads + cache copies) — see StagingLeaseTracker. */
   private readonly staging: StagingLeaseTracker;
@@ -670,7 +680,20 @@ export class KimiTUI {
       }
 
       startupTrace('initMainTui:begin');
-      const shouldReplayHistory = await this.initMainTui();
+      let shouldReplayHistory: boolean;
+      try {
+        shouldReplayHistory = await this.initMainTui();
+      } catch (error) {
+        // The event loop runs here only when the trust prompt already started
+        // it; a startup failure then (e.g. an unknown --environment id failing
+        // fast) must not leak raw mode into the user's shell — mirror the
+        // migration branch's cleanup.
+        if (trustPromptStartedLoop) {
+          this.disposeTerminalTracking();
+          this.state.ui.stop();
+        }
+        throw error;
+      }
       startupTrace('initMainTui:end');
       // Debug-only input→render latency overlay (KIMI_TUI_INPUT_LATENCY=1).
       if (process.env['KIMI_TUI_INPUT_LATENCY']) installInputLatencyProbe(this.state.ui);
@@ -978,6 +1001,9 @@ export class KimiTUI {
         // time (model, permission, plan mode, thinking effort, context cap).
         await this.hydrateLazyConfigDefaults();
         this.appendStartupNotice(SESSIONLESS_STARTUP_NOTICE);
+        if (startup.environment !== undefined && startup.environment !== 'local') {
+          await this.prepareStartupEnvironment(startup.environment);
+        }
       }
       if (session !== undefined && shouldReplayHistory) {
         await this.applyStartupModesToResumedSession(session);
@@ -1448,9 +1474,55 @@ export class KimiTUI {
     }
     let session = this.session;
     if (session === undefined) {
+      // The lazy create can block on a remote environment connect
+      // (--environment startup). Announce the prompt and enter the waiting
+      // state before awaiting it: the UI keeps rendering (and the footer
+      // shows the connecting spinner), and the prompt is not lost when
+      // creation fails.
+      const announced =
+        this.state.appState.streamingPhase === 'idle' &&
+        !this.deferUserMessages &&
+        !this.state.appState.isCompacting;
+      if (announced) {
+        this.appendTranscriptEntry({
+          id: nextTranscriptId(),
+          kind: 'user',
+          turnId: undefined,
+          renderMode: 'plain',
+          content: text,
+          imageAttachmentIds:
+            extraction.imageAttachmentIds.length > 0
+              ? [...extraction.imageAttachmentIds]
+              : undefined,
+        });
+        this.beginSessionRequest();
+      }
       session = await this.ensureSession();
       if (session === undefined) {
         this.staging.release(stagingLease);
+        if (announced) {
+          // Creation failed with the error already on screen; unwind only the
+          // waiting state — the announced prompt stays in the transcript.
+          this.setAppState({ streamingPhase: 'idle' });
+          this.resetLivePane();
+        }
+        return;
+      }
+      if (announced) {
+        if (extraction.hasMedia) {
+          this.sendMessageInternal(session, text, {
+            hasMedia: true,
+            parts: extraction.parts,
+            imageAttachmentIds: extraction.imageAttachmentIds,
+            videoAttachmentIds: extraction.videoAttachmentIds,
+            lease: stagingLease,
+            announced: true,
+          });
+        } else {
+          this.sendMessageInternal(session, text, { announced: true });
+        }
+        this.updateQueueDisplay();
+        this.state.ui.requestRender();
         return;
       }
     }
@@ -1794,14 +1866,16 @@ export class KimiTUI {
       options?.imageAttachmentIds !== undefined && options.imageAttachmentIds.length > 0
         ? options.imageAttachmentIds
         : undefined;
-    this.appendTranscriptEntry({
-      id: nextTranscriptId(),
-      kind: 'user',
-      turnId: undefined,
-      renderMode: 'plain',
-      content: input,
-      imageAttachmentIds,
-    });
+    if (options?.announced !== true) {
+      this.appendTranscriptEntry({
+        id: nextTranscriptId(),
+        kind: 'user',
+        turnId: undefined,
+        renderMode: 'plain',
+        content: input,
+        imageAttachmentIds,
+      });
+    }
     // A goal-active steer is buffered into the running goal turn — no new
     // turn.started will fire for handleTurnStarted to claim the lease — so
     // bind it to that turn here. The turn context must be read BEFORE
@@ -1812,7 +1886,9 @@ export class KimiTUI {
       this.state.appState.streamingPhase === 'idle' || this.state.appState.streamingPhase === 'shell'
         ? undefined
         : this.streamingUI.getTurnContext().turnId;
-    this.beginSessionRequest();
+    if (options?.announced !== true) {
+      this.beginSessionRequest();
+    }
 
     // Compression captions for pasted images are authored here — not at
     // extraction — because only now is the session (and its media-originals
@@ -2213,6 +2289,63 @@ export class KimiTUI {
     return this.session;
   }
 
+  /**
+   * `--environment <id>` startup binding: fail fast on an id the merged
+   * [environments] declarations do not know — the engine only rejects it at
+   * createSession, which the lazy startup defers to the first prompt (losing
+   * that prompt). Then pre-create the session in the background so the remote
+   * connect overlaps startup instead of blocking the first prompt.
+   */
+  private async prepareStartupEnvironment(environmentId: string): Promise<void> {
+    let declarations: readonly WorkspaceEnvironmentDeclarationInfo[] | undefined;
+    try {
+      declarations = await this.harness.listEnvironmentDeclarations(this.state.appState.workDir);
+    } catch {
+      // Declaration resolution failed (e.g. unreadable config): skip the early
+      // check — the background create surfaces the engine's own error.
+      declarations = undefined;
+    }
+    if (declarations !== undefined) {
+      const declared = declarations.find((entry) => entry.id === environmentId);
+      if (declared === undefined) {
+        throw new Error(`environment "${environmentId}" is not declared in [environments]`);
+      }
+      if (declared.defaultCwd === undefined) {
+        throw new Error(`environment "${environmentId}" does not set defaultCwd in [environments]`);
+      }
+      this.startupEnvironmentType = declared.type;
+    }
+    void this.ensureSession();
+  }
+
+  /**
+   * While the `--environment` startup session is being created there is no
+   * binding to sync from yet; a synthetic connecting slot drives the same
+   * footer spinner the registry-backed slot shows once refreshEnvironmentSlot
+   * takes over.
+   */
+  private markStartupEnvironmentConnecting(): void {
+    const environmentId = this.options.startup.environment;
+    if (environmentId === undefined || environmentId === 'local') return;
+    const current = this.state.appState.environment;
+    if (current?.environmentId === environmentId && current.status === 'connecting') return;
+    this.setAppState({
+      environment: {
+        environmentId,
+        type: this.startupEnvironmentType ?? 'command',
+        status: 'connecting',
+      },
+    });
+  }
+
+  private clearStartupEnvironmentConnecting(): void {
+    const environmentId = this.options.startup.environment;
+    const current = this.state.appState.environment;
+    if (environmentId === undefined || current?.environmentId !== environmentId) return;
+    if (current.status !== 'connecting') return;
+    this.setAppState({ environment: undefined });
+  }
+
   /**
    * Seed appState with the config defaults the v2 engine would apply at
    * createSession time (model, permission, plan mode, thinking effort,
@@ -2357,22 +2490,25 @@ export class KimiTUI {
   }
 
   private async lazyCreateSession(): Promise<Session | undefined> {
+    this.markStartupEnvironmentConnecting();
     let session: Session;
     try {
       session = await this.createSessionFromCurrentState(true);
     } catch (error) {
+      this.clearStartupEnvironmentConnecting();
       const msg = formatErrorMessage(error);
       this.showError(`Failed to start a session: ${msg}`);
       return undefined;
     }
-    this.resetSessionRuntime();
+    this.resetSessionRuntime(true);
     await this.setSession(session);
     this.setAppState({ sessionId: session.id });
     try {
       await this.activateRuntime();
       await this.syncRuntimeState(session);
     } catch (error) {
       this.sessionEventHandler.startSubscription();
+      this.clearStartupEnvironmentConnecting();
       const msg = formatErrorMessage(error);
       this.showError(`Post-create setup failed: ${msg}`);
       return undefined;
@@ -2673,12 +2809,15 @@ export class KimiTUI {
     this.state.terminal.setTitle(label);
   }
 
-  resetSessionRuntime(): void {
+  resetSessionRuntime(preserveQueue = false): void {
     this.aborted = false;
     this.cacheHint.resetRuntime();
     this.surveyController.reset();
     this.streamingUI.discardPending();
-    this.clearQueuedMessages();
+    // The lazy first creation keeps input queued while it was in flight (a
+    // bash command behind the announced first prompt belongs to the session
+    // being created); every other reset discards the old session's backlog.
+    if (!preserveQueue) this.clearQueuedMessages();
     this.state.swarmModeEntry = undefined;
     this.streamingUI.resetToolCallState();
     this.streamingUI.resetToolUi();

apps/kimi-code/test/tui/components/messages/agent-group.test.ts

@@ -111,6 +111,55 @@ describe('AgentGroupComponent', () => {
     running.dispose();
   });
 
+  it('shows the environment badge in the row stats only for a non-local subagent', () => {
+    vi.useFakeTimers();
+    vi.setSystemTime(0);
+    const ui = stubTui();
+    const group = new AgentGroupComponent(ui);
+    const remote = new ToolCallComponent(
+      {
+        id: 'call_agent_env',
+        name: 'Agent',
+        args: { description: 'inspect project', environment: 'dev-box' },
+      },
+      undefined,
+      ui,
+    );
+    remote.onSubagentSpawned({
+      agentId: 'sub_call_agent_env',
+      agentName: 'explore',
+      runInBackground: false,
+    });
+    startAgent(remote, 'call_agent_env', 'explore');
+    const local = new ToolCallComponent(
+      {
+        id: 'call_agent_local',
+        name: 'Agent',
+        args: { description: 'write tests', environment: 'local' },
+      },
+      undefined,
+      ui,
+    );
+    local.onSubagentSpawned({
+      agentId: 'sub_call_agent_local',
+      agentName: 'coder',
+      runInBackground: false,
+    });
+    startAgent(local, 'call_agent_local', 'coder');
+
+    group.attach('call_agent_env', remote);
+    group.attach('call_agent_local', local);
+
+    const output = renderText(group);
+    expect(output).toContain('explore · inspect project · env dev-box · 0 tools');
+    expect(output).toContain('coder · write tests · 0 tools');
+    expect(output).not.toContain('env local');
+
+    group.dispose();
+    remote.dispose();
+    local.dispose();
+  });
+
   it('shows the Ctrl+B hint while agents are running and hides it once all are backgrounded', () => {
     vi.useFakeTimers();
     vi.setSystemTime(0);