fix(providers): replay only after reconnect succeeds

Author: TheGreatAxios

src/tui/command-surfaces.test.ts

@@ -1840,6 +1840,37 @@ describe("mcp surface", () => {
 });
 
 describe("host-routed surfaces", () => {
+  test("scoped add-provider forwards its completion callback", async () => {
+    await withShell((shell) => {
+      let received:
+        | {
+            initialKind?: string;
+            initialProfile?: string;
+            onComplete?: (connected: boolean) => void;
+          }
+        | undefined;
+      const completed: boolean[] = [];
+      expect(
+        openCommandSurface(
+          shell,
+          "add-provider",
+          {
+            notify: () => undefined,
+            openAddProvider: (opts) => {
+              received = opts;
+            },
+          },
+          { kind: "xai", profile: "default-2" },
+          (connected) => completed.push(connected),
+        ),
+      ).toBe(true);
+      expect(received?.initialKind).toBe("xai");
+      expect(received?.initialProfile).toBe("default-2");
+      received?.onComplete?.(true);
+      expect(completed).toEqual([true]);
+    });
+  });
+
   test.each([
     {
       surface: "models" as const,

src/tui/command-surfaces.ts

@@ -238,12 +238,14 @@ export interface CommandSurfaceDeps {
   /**
    * Opens the host's add-provider selector (owned by the product host).
    * `/connect` omits returnToModels. `/connect <kind> [profile]` pre-scopes
-   * via initialKind/initialProfile (kind row focused, profile to the flow).
+   * via initialKind/initialProfile (kind row focused, profile to the flow),
+   * with completion reported after that interactive flow settles.
    */
   readonly openAddProvider?: (opts?: {
     returnToModels?: boolean;
     initialKind?: string;
     initialProfile?: string;
+    onComplete?: (connected: boolean) => void;
   }) => void;
   /** Fallback channel for surfaces with no live data source. */
   readonly notify: (text: string) => void;
@@ -1666,6 +1668,7 @@ export function openCommandSurface(
   kind: CommandSurfaceKind,
   deps: CommandSurfaceDeps,
   connectScope?: ReconnectScope,
+  onConnectComplete?: (connected: boolean) => void,
 ): boolean {
   switch (kind) {
     case "help":
@@ -1697,6 +1700,9 @@ export function openCommandSurface(
           ? {
               initialKind: connectScope.kind,
               initialProfile: connectScope.profile,
+              ...(onConnectComplete !== undefined
+                ? { onComplete: onConnectComplete }
+                : {}),
             }
           : undefined,
       );

src/tui/product-host.test.ts

@@ -1106,6 +1106,36 @@ describe("flat type-to-filter model picker", () => {
     }
   });
 
+  test("scoped add-provider forwards completion with the reconnect request", async () => {
+    const completed: boolean[] = [];
+    let finish: ((connected: boolean) => void) | undefined;
+    const { harness, host } = await mountPicker({
+      onConnectProvider: (_name, req) => {
+        expect(req?.kind).toBe("xai");
+        expect(req?.profile).toBe("default-2");
+        finish = req?.onComplete;
+      },
+      addProviderChoices: () => [
+        { id: "openai", label: "OpenAI", hint: "", accountCount: 0 },
+        { id: "xai", label: "xAI", hint: "", accountCount: 1 },
+      ],
+    });
+    try {
+      host.openAddProvider?.({
+        initialKind: "xai",
+        initialProfile: "default-2",
+        onComplete: (connected) => completed.push(connected),
+      });
+      await harness.renderOnce();
+      acceptOverlaySelection(host.shell);
+      finish?.(true);
+      expect(completed).toEqual([true]);
+    } finally {
+      host.dispose();
+      harness.destroy();
+    }
+  });
+
   test("typed /connect then Enter opens add-provider and Esc leaves overlay null", async () => {
     const queued: { open?: () => void } = {};
     const { harness, host } = await mountPicker({

src/tui/product-host.ts

@@ -130,6 +130,8 @@ export interface ProductHostAddProviderChoice {
 export interface ProductHostConnectRequest {
   readonly kind?: string;
   readonly profile?: string;
+  /** Reports whether the interactive connect flow completed successfully. */
+  readonly onComplete?: (connected: boolean) => void;
 }
 
 export interface ProductHostConfig {
@@ -268,6 +270,7 @@ export interface ProductHost {
     returnToModels?: boolean;
     initialKind?: string;
     initialProfile?: string;
+    onComplete?: (connected: boolean) => void;
   }) => void;
   /** Swap the picker's rows/descriptions in place (e.g. after a provider connects). */
   readonly setModels?: (
@@ -598,6 +601,7 @@ export async function mountProductHost(
         returnToModels?: boolean;
         initialKind?: string;
         initialProfile?: string;
+        onComplete?: (connected: boolean) => void;
       }) => void)
     | undefined;
   if (config.onModelSelect) {
@@ -620,6 +624,7 @@ export async function mountProductHost(
             returnToModels?: boolean;
             initialKind?: string;
             initialProfile?: string;
+            onComplete?: (connected: boolean) => void;
           }): void => {
             const rows = addProviderChoices();
             const scopedIndex =
@@ -644,7 +649,13 @@ export async function mountProductHost(
                   opts?.initialProfile !== undefined &&
                   (opts.initialKind === undefined || opts.initialKind === id)
                 ) {
-                  onConnect(id, { kind: id, profile: opts.initialProfile });
+                  onConnect(id, {
+                    kind: id,
+                    profile: opts.initialProfile,
+                    ...(opts.onComplete !== undefined
+                      ? { onComplete: opts.onComplete }
+                      : {}),
+                  });
                   return;
                 }
                 onConnect(id);

src/tui/runner/host.ts

@@ -180,6 +180,7 @@ export type RunnerHost = ProductHost & {
   readonly openSurface: (
     kind: CommandSurfaceKind,
     connectScope?: ReconnectScope,
+    onConnectComplete?: (connected: boolean) => void,
   ) => boolean;
   /**
    * Recompute the models-first catalog from fresh recent/favorite refs and
@@ -511,8 +512,14 @@ export async function mountRunnerHost(
   return {
     ...host,
     dispose,
-    openSurface: (kind, connectScope) =>
-      openCommandSurface(host.shell, kind, surfaceDeps, connectScope),
+    openSurface: (kind, connectScope, onConnectComplete) =>
+      openCommandSurface(
+        host.shell,
+        kind,
+        surfaceDeps,
+        connectScope,
+        onConnectComplete,
+      ),
     refreshModels,
     refreshCostContext: pushCostContext,
     openCredentialRecovery,

src/tui/runner/index.ts

@@ -257,7 +257,8 @@ export async function runTUI(initialConfig: Config): Promise<number> {
     state.presentReconnectRecovery = createReconnectRecoveryPresenter({
       recovery: state.reconnectRecovery,
       openDialog: (dialog) => host.openReconnectRecovery(dialog),
-      openReconnect: (scope) => host.openSurface("add-provider", scope),
+      openReconnect: (scope, onComplete) =>
+        host.openSurface("add-provider", scope, onComplete),
       readDirector: () => services.directorHolder.instance,
       deliverContinuation: (message) => liveAgent(state).deliver(message),
     });

src/tui/runner/reconnect-recovery.test.ts

@@ -157,7 +157,7 @@ describe("reconnect recovery accept gate", () => {
 });
 
 describe("reconnect recovery selection", () => {
-  test("reconnects the failed scope and replays the turn once", () => {
+  test("reconnects the failed scope and replays only after success", () => {
     const state = createReconnectRecoveryState();
     const attempt = state.begin(operatorMessage(), "xai");
     state.observe(attempt, credentialRetry());
@@ -167,12 +167,14 @@ describe("reconnect recovery selection", () => {
     const seen: string[] = [];
     let armed: number | null = null;
     const delivered: InboundMessage[] = [];
+    let completeReconnect: ((connected: boolean) => void) | undefined;
     const outcome = applyReconnectRecoverySelection({
       state,
       generation: pending.generation,
       selectedId: reconnectRecoveryItemId(pending.scope),
-      reconnect: (scope) => {
+      reconnect: (scope, onComplete) => {
         seen.push(`${scope.kind}/${scope.profile}`);
+        completeReconnect = onComplete;
       },
       armContinuation: (generation) => {
         armed = generation;
@@ -184,10 +186,45 @@ describe("reconnect recovery selection", () => {
         delivered.push(message);
       },
     });
-    expect(outcome).toBe("continued");
+    expect(outcome).toBe("reconnected");
     expect(seen).toEqual(["xai/default-2"]);
+    expect(armed).toBeNull();
+    expect(delivered).toEqual([]);
+    completeReconnect?.(true);
     expect(armed ?? -1).toBe(pending.generation);
     expect(delivered).toHaveLength(1);
+    completeReconnect?.(true);
+    expect(delivered).toHaveLength(1);
+  });
+
+  test("cancel, failure, and committed attempts never replay", () => {
+    for (const connected of [false, true]) {
+      const state = createReconnectRecoveryState();
+      const attempt = state.begin(operatorMessage(), "xai");
+      state.observe(attempt, credentialRetry());
+      if (connected) {
+        state.observe(attempt, { type: "inference.tool-call", data: {} });
+      }
+      state.observe(attempt, credentialFailure("xai/default-2"));
+      const pending = required(state.settle(attempt), "pending recovery");
+      let completeReconnect: ((result: boolean) => void) | undefined;
+      applyReconnectRecoverySelection({
+        state,
+        generation: pending.generation,
+        selectedId: reconnectRecoveryItemId(pending.scope),
+        reconnect: (_scope, onComplete) => {
+          completeReconnect = onComplete;
+        },
+        armContinuation: () => {
+          throw new Error("should not arm");
+        },
+        cancelContinuation: () => undefined,
+        deliverContinuation: () => {
+          throw new Error("should not deliver");
+        },
+      });
+      completeReconnect?.(connected);
+    }
   });
 
   test("reconnect failure and stale generations never arm a continuation", () => {
@@ -254,6 +291,7 @@ describe("reconnect recovery presenter (runTUI wiring)", () => {
     };
     const dialogs: ReconnectDialog[] = [];
     const reconnected: ReconnectScope[] = [];
+    let completeReconnect: ((connected: boolean) => void) | undefined;
     const armed: number[] = [];
     const cancelled: number[] = [];
     const delivered: InboundMessage[] = [];
@@ -263,8 +301,9 @@ describe("reconnect recovery presenter (runTUI wiring)", () => {
         dialogs.push(next);
         return dialogOpens;
       },
-      openReconnect: (scope) => {
+      openReconnect: (scope, onComplete?: (connected: boolean) => void) => {
         reconnected.push(scope);
+        completeReconnect = onComplete;
         return options?.reconnectOpens ?? true;
       },
       readDirector: () =>
@@ -290,17 +329,22 @@ describe("reconnect recovery presenter (runTUI wiring)", () => {
       pending,
       dialog,
       reconnected,
+      completeReconnect: (connected: boolean) => completeReconnect?.(connected),
       armed,
       cancelled,
       delivered,
     };
   }
 
-  test("accept re-keys the failed scope and replays the turn once", () => {
+  test("accept waits for a successful re-key before replaying the turn once", () => {
     const wired = wirePresenter();
     expect(wired.dialog.scope).toEqual(wired.pending.scope);
     wired.dialog.onAccept(reconnectRecoveryItemId(wired.pending.scope));
     expect(wired.reconnected).toEqual([wired.pending.scope]);
+    expect(wired.armed).toEqual([]);
+    expect(wired.delivered).toEqual([]);
+
+    wired.completeReconnect(true);
     expect(wired.armed).toEqual([wired.pending.generation]);
     expect(wired.cancelled).toEqual([]);
     // The continuation message carries a build-time timestamp, so pin the
@@ -314,6 +358,9 @@ describe("reconnect recovery presenter (runTUI wiring)", () => {
     expect(continuation.headers.interchangeCorrelationId).toBe(
       String(wired.pending.generation),
     );
+    wired.completeReconnect(true);
+    expect(wired.armed).toEqual([wired.pending.generation]);
+    expect(wired.delivered).toHaveLength(1);
     // The generation is consumed: a late cancel is a no-op.
     expect(wired.recovery.cancel(wired.pending.generation)).toBe(false);
   });

src/tui/runner/reconnect-recovery.ts

@@ -194,29 +194,36 @@ export function applyReconnectRecoverySelection(args: {
   state: ReturnType<typeof createReconnectRecoveryState>;
   generation: number;
   selectedId: string;
-  reconnect: (scope: ReconnectScope) => void;
+  reconnect: (
+    scope: ReconnectScope,
+    onComplete: (connected: boolean) => void,
+  ) => void;
   armContinuation: (generation: number) => void;
   cancelContinuation: (generation: number) => void;
   deliverContinuation: (message: InboundMessage) => void;
-}): "stale" | "invalid" | "reconnect-failed" | "reconnected" | "continued" {
+}): "stale" | "invalid" | "reconnect-failed" | "reconnected" {
   const acceptance = args.state.accept(args.generation, args.selectedId);
   if (acceptance.kind !== "accepted") return acceptance.kind;
+  let completed = false;
+  const onComplete = (connected: boolean): void => {
+    if (completed) return;
+    completed = true;
+    if (!connected || !acceptance.replay) return;
+    args.armContinuation(acceptance.generation);
+    try {
+      args.deliverContinuation(
+        buildCredentialRecoveryContinuationMessage(acceptance.generation),
+      );
+    } catch {
+      args.cancelContinuation(acceptance.generation);
+    }
+  };
   try {
-    args.reconnect(acceptance.scope);
+    args.reconnect(acceptance.scope, onComplete);
   } catch {
     return "reconnect-failed";
   }
-  if (!acceptance.replay) return "reconnected";
-  args.armContinuation(acceptance.generation);
-  try {
-    args.deliverContinuation(
-      buildCredentialRecoveryContinuationMessage(acceptance.generation),
-    );
-  } catch {
-    args.cancelContinuation(acceptance.generation);
-    return "reconnected";
-  }
-  return "continued";
+  return "reconnected";
 }
 
 /**
@@ -236,7 +243,10 @@ export function createReconnectRecoveryPresenter(args: {
     onAccept: (id: string) => void;
     onCancel: () => void;
   }) => boolean;
-  openReconnect: (scope: ReconnectScope) => boolean;
+  openReconnect: (
+    scope: ReconnectScope,
+    onComplete: (connected: boolean) => void,
+  ) => boolean;
   readDirector: () =>
     | {
         armCredentialRecoveryContinuation: (generation: number) => void;
@@ -261,8 +271,8 @@ export function createReconnectRecoveryPresenter(args: {
           state: args.recovery,
           generation: pending.generation,
           selectedId: id,
-          reconnect: (scope) => {
-            if (!args.openReconnect(scope)) {
+          reconnect: (scope, onComplete) => {
+            if (!args.openReconnect(scope, onComplete)) {
               throw new Error(
                 `reconnect surface unavailable for ${scope.kind}/${scope.profile}`,
               );