opencode-auth-load-balancer 0.2.2 → 0.4.0

package/tui/auth-load-balancer-tui.view.tsx

@@ -11,8 +11,8 @@
  *   - sidebar_content — a panel next to Context / MCP / LSP listing ALL accounts
  *     per provider, sorted by the scheduler's score (highest = use next), showing
  *     the score + use-order, usage + reset countdowns, in-use marker, and state.
- *     Click any account row for a menu to Rename (prompt) or Delete (confirm) it,
- *     written straight to the pool file.
+ *     Click any account row for a menu to Rename (prompt), Disable/Enable
+ *     (toggle), Re-login (provider OAuth), or Delete (confirm) it.
  *
  * The ranking is computed by the SAME code the server scheduler uses — imported from a
  * byte-identical copy of src/scheduler/score-core.ts installed alongside this file
@@ -31,12 +31,19 @@ import {
 } from './auth-load-balancer-scoring'
 import {
   cfg,
+  clearReloginTargetInPool,
   compareAscii,
   deleteFromPool,
+  MANUAL_DISABLED_REASON,
   pct,
+  pickAuthMethodIndex,
   type PoolShape,
   readPool,
   renameInPool,
+  sessionAccountId,
+  sessionFallback,
+  setDisabledInPool,
+  setReloginTargetInPool,
   stateOf,
   tierResets,
   toScore,
@@ -57,16 +64,51 @@ function providerLabel(id: string): string {
 }
 const POLL_MS = 3000
 
+/**
+ * The current session id from the TUI route, or undefined on the home screen.
+ * Unlike the `sidebar_content` slot (handed a `session_id` prop), the always-
+ * visible `app_bottom` bar receives no slot props, so it reads the session from
+ * the route to scope its in-use account to the VIEWER's own session.
+ */
+function routeSessionId(api: TuiPluginApi): string | undefined {
+  const r = api.route.current
+  if (r.name !== 'session') return undefined
+  const sid = r.params?.sessionID
+  return typeof sid === 'string' ? sid : undefined
+}
+
 // Scoring knobs read from env exactly like the server — single-sourced in the shared
 // ./auth-load-balancer-scoring module (a byte copy of src/scheduler/score-core.ts),
 // shared with auth-load-balancer-tui.logic.ts (imported above as `cfg`) so it is
 // computed exactly once.
 
+/** One poll sample: the pool, plus the instant it was read. */
+interface PoolSample {
+  pool: PoolShape
+  now: number
+}
+
+/**
+ * Capture the pool AND the clock together, beside each other.
+ *
+ * The timestamp deliberately lives HERE rather than in each consuming memo.
+ * A memo that calls `Date.now()` itself is impure — it depends on something
+ * Solid cannot track, so its countdowns refresh only by the coincidence that
+ * the poll signal happens to change on the same tick — and it measures every
+ * `until()` against an instant slightly LATER than the snapshot it is
+ * describing. Sampling both at once makes the memos pure functions of one
+ * reactive input and keeps each countdown consistent with the pool state it
+ * is rendered from.
+ */
+function pollSample(): PoolSample {
+  return { pool: readPool(), now: Date.now() }
+}
+
 function usePool() {
-  const [pool, setPool] = createSignal<PoolShape>(readPool())
-  const timer = setInterval(() => setPool(readPool()), POLL_MS)
+  const [sample, setSample] = createSignal<PoolSample>(pollSample())
+  const timer = setInterval(() => setSample(pollSample()), POLL_MS)
   onCleanup(() => clearInterval(timer))
-  return pool
+  return sample
 }
 
 /** Shared theme-color accessor used identically by `BottomBar` and `SidebarPanel`. */
@@ -85,24 +127,32 @@ interface WindowDisplay {
 interface Chip extends WindowDisplay {
   name: string
   label: string
+  /** Raw model ids (requested → served) when this session runs on a fallback model, else undefined. */
+  fallback?: { from: string; to: string }
 }
 
 /** Always-visible bottom bar (app_bottom): the in-use account per provider. */
 function BottomBar(props: { api: TuiPluginApi }) {
-  const pool = usePool()
+  const sample = usePool()
   const color = themeColor(props.api)
   const chips = createMemo<Chip[]>(() => {
-    const p = pool()
+    const { pool: p, now } = sample()
     const accounts = p.accounts ?? []
-    const now = Date.now()
     const out: Chip[] = []
-    // `lastSelected` key order is whichever provider served first; sort by
-    // provider id (byte-deterministic `< / >`, no locale) so the bar lists
-    // providers in the same order as the sidebar, CLI, and status tool.
-    const selected = Object.entries(p.lastSelected ?? {}).sort(([x], [y]) =>
-      compareAscii(x, y),
+    const sid = routeSessionId(props.api)
+    // Show the account THIS session is using per provider — NOT the global
+    // `lastSelected` (whichever session served last across the whole pool). On
+    // the home screen (no session) fall back to `lastSelected` so the bar isn't
+    // blank. Providers byte-sorted (`< / >`, no locale) so the bar lists them in
+    // the same order as the sidebar, CLI, and status tool.
+    const providerIds = [...new Set(accounts.map((x) => x.providerID))].sort(
+      compareAscii,
     )
-    for (const [providerID, id] of selected) {
+    for (const providerID of providerIds) {
+      const id =
+        sessionAccountId(p, providerID, sid) ??
+        (sid ? undefined : p.lastSelected?.[providerID])
+      if (!id) continue
       const a = accounts.find((x) => x.id === id)
       if (!a) continue
       out.push({
@@ -112,6 +162,9 @@ function BottomBar(props: { api: TuiPluginApi }) {
         weeklyReset: until(a.usage?.weekly?.resetAt, now),
         hourlyPct: winPct(a.usage?.hourly, now),
         hourlyReset: until(a.usage?.hourly?.resetAt, now),
+        // Only THIS session's fallback (keyed by sid); the `lastSelected` home-screen
+        // fallback path has no session, so no downgrade to surface there.
+        fallback: sessionFallback(p, providerID, sid),
       })
     }
     return out
@@ -125,6 +178,11 @@ function BottomBar(props: { api: TuiPluginApi }) {
             <span style={{ fg: color().primary }}>
               {a.name} {a.label}
             </span>
+            {a.fallback ? (
+              <span style={{ fg: color().warning }}>
+                {`  fallback ${a.fallback.from}→${a.fallback.to}`}
+              </span>
+            ) : null}
             {`  wk ${a.weeklyPct} (${a.weeklyReset}) · 5h ${a.hourlyPct} (${a.hourlyReset})`}
           </text>
         )}
@@ -136,10 +194,12 @@ function BottomBar(props: { api: TuiPluginApi }) {
 interface Row extends WindowDisplay {
   id: string
   label: string
+  providerID: string
   current: boolean
   score: number | null
   rank: number | null
   state: string
+  manuallyDisabled: boolean
 }
 interface Group {
   provider: string
@@ -147,17 +207,23 @@ interface Group {
 }
 
 /** Sidebar panel: all accounts per provider, scored + sorted, click to rename. */
-function SidebarPanel(props: { api: TuiPluginApi }) {
-  const pool = usePool()
+function SidebarPanel(props: {
+  api: TuiPluginApi
+  sessionId: string | undefined
+}) {
+  const sample = usePool()
   const color = themeColor(props.api)
   const groups = createMemo<Group[]>(() => {
-    const p = pool()
+    const { pool: p, now } = sample()
     const accounts = p.accounts ?? []
-    const now = Date.now()
     const providerIds = [...new Set(accounts.map((a) => a.providerID))].sort(
       compareAscii,
     )
     return providerIds.map((providerID) => {
+      // The account the VIEWER's session is pinned to for this provider — used
+      // for the in-use (▶) marker instead of the global `lastSelected`, so the
+      // marker reflects THIS session, not whichever session served last.
+      const sessionAcct = sessionAccountId(p, providerID, props.sessionId)
       const ranked = accounts
         .filter((a) => a.providerID === providerID)
         .map((a) => {
@@ -180,7 +246,8 @@ function SidebarPanel(props: { api: TuiPluginApi }) {
         return {
           id: a.id,
           label: a.label,
-          current: p.lastSelected?.[providerID] === a.id,
+          providerID,
+          current: sessionAcct === a.id,
           score: available ? score : null,
           rank: available ? rank : null,
           // Reuse the windows `toScore(a)` already normalized instead of
@@ -196,6 +263,7 @@ function SidebarPanel(props: { api: TuiPluginApi }) {
           // `1e999` entry (Infinity via JSON.parse) would otherwise render its
           // tier annotation forever until the server heals the file.
           state: stateOf(sa, tierResets(a, now), now),
+          manuallyDisabled: a.disabledReason === MANUAL_DISABLED_REASON,
         }
       })
       return { provider: providerLabel(providerID), rows }
@@ -234,25 +302,187 @@ function SidebarPanel(props: { api: TuiPluginApi }) {
     )
   }
 
-  // Click an account -> a small menu so both Rename and Delete are reachable.
-  function openMenu(id: string, label: string): void {
-    dialog().replace(() =>
-      props.api.ui.DialogSelect({
-        title: label,
-        options: [
-          {
-            title: 'Rename',
-            value: 'rename',
-            onSelect: () => openRename(id, label),
+  async function openRelogin(
+    id: string,
+    label: string,
+    providerID: string,
+  ): Promise<void> {
+    try {
+      const methods = (await props.api.client.provider.auth()).data?.[
+        providerID
+      ]
+      const method = pickAuthMethodIndex(methods)
+      if (method === null) {
+        props.api.ui.toast({
+          variant: 'error',
+          message: `No OAuth login method is registered for ${providerID}.`,
+        })
+        return
+      }
+
+      const auth = (
+        await props.api.client.provider.oauth.authorize({
+          providerID,
+          method,
+        })
+      ).data
+      if (!auth) {
+        props.api.ui.toast({
+          variant: 'error',
+          message: `Unable to start OAuth login for ${providerID}.`,
+        })
+        return
+      }
+
+      // Timestamp intentionally left to the callee's default — see its doc.
+      setReloginTargetInPool(id, providerID)
+
+      const completeRelogin = async (code?: string): Promise<void> => {
+        try {
+          const completed = (
+            await props.api.client.provider.oauth.callback({
+              providerID,
+              method,
+              ...(code === undefined ? {} : { code }),
+            })
+          ).data
+          if (completed) {
+            props.api.ui.toast({
+              variant: 'success',
+              message: `Re-login complete for ${label}.`,
+            })
+            // The existing 3s usePool poll observes the consumed intent and
+            // cleared disabledReason; no manual state refresh is needed.
+          } else {
+            clearReloginTargetInPool()
+            props.api.ui.toast({
+              variant: 'error',
+              message: `Re-login failed for ${label}.`,
+            })
+          }
+          dialog().clear()
+        } catch {
+          clearReloginTargetInPool()
+          dialog().clear()
+          props.api.ui.toast({
+            variant: 'error',
+            message: `Re-login failed for ${label}.`,
+          })
+        }
+      }
+
+      if (auth.method === 'auto') {
+        await completeRelogin()
+        return
+      }
+
+      dialog().replace(() =>
+        props.api.ui.DialogPrompt({
+          title: `Re-login "${label}"`,
+          description: () => (
+            <box>
+              <text>{auth.instructions}</text>
+              <text>{auth.url}</text>
+            </box>
+          ),
+          placeholder: 'Paste the authorization code',
+          onConfirm: (code: string) => {
+            const trimmed = code.trim()
+            if (!trimmed) {
+              clearReloginTargetInPool()
+              dialog().clear()
+              return
+            }
+            void completeRelogin(trimmed)
           },
-          {
-            title: 'Delete — remove from pool',
-            value: 'delete',
-            onSelect: () => openDelete(id, label),
+          onCancel: () => {
+            clearReloginTargetInPool()
+            dialog().clear()
           },
-        ],
-      }),
-    )
+        }),
+      )
+    } catch {
+      clearReloginTargetInPool()
+      dialog().clear()
+      props.api.ui.toast({
+        variant: 'error',
+        message: `Re-login failed for ${label}.`,
+      })
+    }
+  }
+
+  // Click an account -> a small menu so Rename, Disable/Enable, Re-login, and Delete are
+  // all reachable. Deliberately NOT api.ui.DialogSelect: that always renders an
+  // auto-focused filter <input>, and a focused opentui input swallows the FIRST
+  // Esc (to blur itself) — so the menu needed TWO Esc presses to close. A plain
+  // clickable list has no input, so the dialog stack's own Esc binding closes it
+  // in ONE press. The menu is opened by a mouse click on the row, so mouse-driven
+  // options stay consistent (there is no keyboard path that opens it).
+  function openMenu(
+    id: string,
+    label: string,
+    providerID: string,
+    manuallyDisabled: boolean,
+  ): void {
+    const items: { title: string; run: () => void }[] = [
+      { title: 'Rename', run: () => openRename(id, label) },
+      {
+        // Reversible (just a scheduler skip), so no confirm step — flip the
+        // pool flag and close, unlike Delete which drops the tokens.
+        title: manuallyDisabled
+          ? 'Enable — include in selection'
+          : 'Disable — exclude from selection',
+        run: () => {
+          setDisabledInPool(id, !manuallyDisabled)
+          dialog().clear()
+        },
+      },
+      {
+        title: 'Re-login — re-authorize with the provider',
+        run: () => {
+          void openRelogin(id, label, providerID)
+        },
+      },
+      { title: 'Delete — remove from pool', run: () => openDelete(id, label) },
+    ]
+    dialog().replace(() => {
+      const c = color()
+      const [hovered, setHovered] = createSignal(-1)
+      return (
+        <box
+          gap={1}
+          paddingBottom={1}
+          paddingLeft={4}
+          paddingRight={4}
+          paddingTop={1}
+        >
+          <box flexDirection="row" justifyContent="space-between">
+            <text fg={c.text}>
+              <b>{label}</b>
+            </text>
+            <text fg={c.textMuted} onMouseUp={() => dialog().clear()}>
+              esc
+            </text>
+          </box>
+          <box>
+            <For each={items}>
+              {(item, i) => (
+                <box
+                  onMouseMove={() => setHovered(i())}
+                  onMouseUp={() => item.run()}
+                  paddingLeft={1}
+                  paddingRight={1}
+                >
+                  <text fg={hovered() === i() ? c.primary : c.text}>
+                    {item.title}
+                  </text>
+                </box>
+              )}
+            </For>
+          </box>
+        </box>
+      )
+    })
   }
 
   return (
@@ -262,7 +492,7 @@ function SidebarPanel(props: { api: TuiPluginApi }) {
           <b>Auth accounts</b>
           <span style={{ fg: color().textMuted }}>
             {' '}
-            (click: rename / delete)
+            (click: rename / disable / re-login / delete)
           </span>
         </text>
         <For each={groups()}>
@@ -271,9 +501,19 @@ function SidebarPanel(props: { api: TuiPluginApi }) {
               <text fg={color().textMuted}>{g.provider}</text>
               <For each={g.rows}>
                 {(r) => (
-                  <box onMouseUp={() => openMenu(r.id, r.label)}>
+                  <box
+                    onMouseUp={() =>
+                      openMenu(r.id, r.label, r.providerID, r.manuallyDisabled)
+                    }
+                  >
                     <text
-                      fg={r.current ? color().primary : color().text}
+                      fg={
+                        r.manuallyDisabled
+                          ? color().textMuted
+                          : r.current
+                            ? color().primary
+                            : color().text
+                      }
                       wrapMode="word"
                     >
                       {(r.current ? '▶ ' : '  ') +
@@ -313,9 +553,11 @@ function makeSlots(api: TuiPluginApi): TuiSlotPlugin {
       app_bottom() {
         return <BottomBar api={api} />
       },
-      // Panel in the session sidebar, next to Context / MCP / LSP.
-      sidebar_content() {
-        return <SidebarPanel api={api} />
+      // Panel in the session sidebar, next to Context / MCP / LSP. The host
+      // hands this slot the current `session_id`, so the in-use marker can scope
+      // to the VIEWER's session instead of the global last-selected account.
+      sidebar_content(_ctx, props) {
+        return <SidebarPanel api={api} sessionId={props.session_id} />
       },
     },
   }