Add freebuff session-end banner and drain-window handling

When the server flips the seat into `draining` (past `expires_at`) or
past the hard cutoff, the chat input is replaced by a "Session ended"
banner. While an agent is still streaming under the grace window, Enter
is disabled and the banner shows "Agent is wrapping up. Rejoin the wait
room after it's finished." — Esc still interrupts. Once idle, Enter
re-POSTs /session to drop back into the waiting room.

Adds a small countdown to the far right of the status bar (muted, turning
soft warning in the final minute — no red) and schedules the next poll
just after expires_at / gracePeriodEndsAt so the draining transition
shows up promptly instead of stalling at 0 for a full interval.

Moves getFreebuffInstanceId onto the session hook's module handle and
deletes the now-vestigial freebuff-session-store.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: jahooma

cli/src/app.tsx

@@ -375,6 +375,10 @@ const AuthedSurface = ({
   //   'none'   → server lost our row; hook is about to re-POST
   // Falling through to <Chat> on 'none' would leave the user unable to send
   // any free-mode request until the next poll cycle.
+  //
+  // 'draining' and 'ended' deliberately fall through to <Chat>: the agent
+  // may still be finishing work under the server-side grace period, and the
+  // chat surface itself swaps the input box for the session-ended banner.
   if (
     IS_FREEBUFF &&
     (session === null ||
@@ -401,6 +405,7 @@ const AuthedSurface = ({
       initialMode={initialMode}
       gitRoot={gitRoot}
       onSwitchToGitRoot={onSwitchToGitRoot}
+      freebuffSession={session}
     />
   )
 }

cli/src/chat.tsx

@@ -21,7 +21,7 @@ import { ReviewScreen } from './components/review-screen'
 import { MessageWithAgents } from './components/message-with-agents'
 import { areCreditsRestored } from './components/out-of-credits-banner'
 import { PendingBashMessage } from './components/pending-bash-message'
-import { useHasActiveFreebuffSession } from './components/freebuff-session-countdown'
+import { SessionEndedBanner } from './components/session-ended-banner'
 import { StatusBar } from './components/status-bar'
 import { TopBanner } from './components/top-banner'
 import { getSlashCommandsWithSkills } from './data/slash-commands'
@@ -84,6 +84,7 @@ import { computeInputLayoutMetrics } from './utils/text-layout'
 import type { CommandResult } from './commands/command-registry'
 import type { MultilineInputHandle } from './components/multiline-input'
 import type { MatchedSlashCommand } from './hooks/use-suggestion-engine'
+import type { FreebuffSessionResponse } from './types/freebuff-session'
 import type { User } from './utils/auth'
 import type { AgentMode } from './utils/constants'
 import type { FileTreeNode } from '@codebuff/common/util/file'
@@ -106,6 +107,7 @@ export const Chat = ({
   initialMode,
   gitRoot,
   onSwitchToGitRoot,
+  freebuffSession,
 }: {
   headerContent: React.ReactNode
   initialPrompt: string | null
@@ -121,6 +123,7 @@ export const Chat = ({
   initialMode?: AgentMode
   gitRoot?: string | null
   onSwitchToGitRoot?: () => void
+  freebuffSession: FreebuffSessionResponse | null
 }) => {
   const [forceFileOnlyMentions, setForceFileOnlyMentions] = useState(false)
 
@@ -1338,7 +1341,12 @@ export const Chat = ({
     return ` ${segments.join('   ')} `
   }, [queuePreviewTitle, pausedQueueText])
 
-  const hasActiveFreebuffSession = useHasActiveFreebuffSession()
+  const hasActiveFreebuffSession =
+    IS_FREEBUFF && freebuffSession?.status === 'active'
+  const isFreebuffSessionOver =
+    IS_FREEBUFF &&
+    (freebuffSession?.status === 'draining' ||
+      freebuffSession?.status === 'ended')
   const shouldShowStatusLine =
     !feedbackMode &&
     (hasStatusIndicatorContent ||
@@ -1447,6 +1455,7 @@ export const Chat = ({
             scrollToLatest={scrollToLatest}
             statusIndicatorState={statusIndicatorState}
             onStop={chatKeyboardHandlers.onInterruptStream}
+            freebuffSession={freebuffSession}
           />
         )}
 
@@ -1466,11 +1475,18 @@ export const Chat = ({
         )}
 
         {reviewMode ? (
+          // Review takes precedence over the session-ended banner: during a
+          // draining session the agent may still be asking to run tools, and
+          // those approvals must be reachable for the run to finish.
           <ReviewScreen
             onSelectOption={handleReviewOptionSelect}
             onCustom={handleReviewCustom}
             onCancel={handleCloseReviewScreen}
           />
+        ) : isFreebuffSessionOver ? (
+          <SessionEndedBanner
+            isStreaming={isStreaming || isWaitingForResponse}
+          />
         ) : (
           <ChatInputBar
             inputValue={inputValue}

cli/src/components/freebuff-session-countdown.tsx

@@ -1,11 +1,11 @@
 import React, { useEffect, useState } from 'react'
 
 import { useTheme } from '../hooks/use-theme'
-import { useFreebuffSessionStore } from '../state/freebuff-session-store'
 import { IS_FREEBUFF } from '../utils/constants'
 
-const LOW_THRESHOLD_MS = 5 * 60_000
-const CRITICAL_THRESHOLD_MS = 60_000
+import type { FreebuffSessionResponse } from '../types/freebuff-session'
+
+const LOW_THRESHOLD_MS = 60_000
 
 const formatRemaining = (ms: number): string => {
   if (ms <= 0) return 'expiring…'
@@ -24,9 +24,10 @@ const formatRemaining = (ms: number): string => {
  * surprised when their seat is released. Returns null in non-freebuff
  * builds or when no active session exists — safe to always mount.
  */
-export const FreebuffSessionCountdown: React.FC = () => {
+export const FreebuffSessionCountdown: React.FC<{
+  session: FreebuffSessionResponse | null
+}> = ({ session }) => {
   const theme = useTheme()
-  const session = useFreebuffSessionStore((s) => s.session)
   const expiresAtMs =
     session?.status === 'active' ? Date.parse(session.expiresAt) : null
 
@@ -40,21 +41,9 @@ export const FreebuffSessionCountdown: React.FC = () => {
   if (!IS_FREEBUFF || !expiresAtMs) return null
 
   const remainingMs = expiresAtMs - now
-  const color =
-    remainingMs < CRITICAL_THRESHOLD_MS
-      ? theme.error
-      : remainingMs < LOW_THRESHOLD_MS
-        ? theme.warning
-        : theme.muted
+  // Muted until the final minute, then a soft warning — deliberately not
+  // `theme.error` so the countdown reads informational, not alarming.
+  const color = remainingMs < LOW_THRESHOLD_MS ? theme.warning : theme.muted
 
   return <span fg={color}>{formatRemaining(remainingMs)}</span>
 }
-
-/** True when the freebuff session countdown will render non-null content.
- *  Used by the chat surface to keep the status bar visible while a
- *  session is active, even when there's no streaming/queue activity. */
-export const useHasActiveFreebuffSession = (): boolean => {
-  return useFreebuffSessionStore(
-    (s) => IS_FREEBUFF && s.session?.status === 'active',
-  )
-}

cli/src/components/session-ended-banner.tsx

@@ -0,0 +1,101 @@
+import { TextAttributes } from '@opentui/core'
+import { useKeyboard } from '@opentui/react'
+import React, { useCallback, useState } from 'react'
+
+import { Button } from './button'
+import { refreshFreebuffSession } from '../hooks/use-freebuff-session'
+import { useTheme } from '../hooks/use-theme'
+import { useChatStore } from '../state/chat-store'
+import { BORDER_CHARS } from '../utils/ui-constants'
+
+import type { KeyEvent } from '@opentui/core'
+
+interface SessionEndedBannerProps {
+  /** True while an agent request is still streaming under the server-side
+   *  grace window. Swaps the Enter-to-rejoin affordance for a "let it
+   *  finish" hint so the user doesn't abort their in-flight work. */
+  isStreaming: boolean
+}
+
+/**
+ * Replaces the chat input when the freebuff session has ended (client state
+ * `draining` or `ended`). Captures Enter to re-queue the user; Esc keeps
+ * falling through to the global stream-interrupt handler so in-flight work
+ * can be cancelled.
+ */
+export const SessionEndedBanner: React.FC<SessionEndedBannerProps> = ({
+  isStreaming,
+}) => {
+  const theme = useTheme()
+  const [rejoining, setRejoining] = useState(false)
+
+  // While a request is still streaming, rejoin is disabled: it would
+  // unmount <Chat> and abort the in-flight agent run. The promise is "we
+  // let the agent finish" — honoring that means Enter does nothing until
+  // the stream ends or the user hits Esc.
+  const canRejoin = !isStreaming && !rejoining
+  const rejoin = useCallback(() => {
+    if (!canRejoin) return
+    setRejoining(true)
+    // Once the POST lands, the hook flips status to 'queued' and app.tsx
+    // swaps us into <WaitingRoomScreen>, unmounting this banner. No need to
+    // clear `rejoining` on success — the component will be gone.
+    refreshFreebuffSession()
+      .then(() => {
+        // Wipe the prior conversation so the next admitted session starts
+        // with empty history instead of continuing the one that just ended.
+        useChatStore.getState().reset()
+      })
+      .catch(() => setRejoining(false))
+  }, [canRejoin])
+
+  useKeyboard(
+    useCallback(
+      (key: KeyEvent) => {
+        if (!canRejoin) return
+        if (key.name === 'return' || key.name === 'enter') {
+          key.preventDefault?.()
+          rejoin()
+        }
+      },
+      [rejoin, canRejoin],
+    ),
+  )
+
+  return (
+    <box
+      title="Session ended"
+      titleAlignment="center"
+      style={{
+        width: '100%',
+        borderStyle: 'single',
+        borderColor: theme.muted,
+        customBorderChars: BORDER_CHARS,
+        paddingLeft: 1,
+        paddingRight: 1,
+        paddingTop: 0,
+        paddingBottom: 0,
+        flexDirection: 'column',
+        gap: 0,
+      }}
+    >
+      <text style={{ fg: theme.foreground, wrapMode: 'word' }}>
+        Your freebuff session has ended.
+      </text>
+      {isStreaming ? (
+        <text style={{ fg: theme.muted, wrapMode: 'word' }}>
+          Agent is wrapping up. Rejoin the wait room after it's finished.
+        </text>
+      ) : (
+        <Button onClick={rejoin}>
+          <text
+            style={{ fg: rejoining ? theme.muted : theme.primary }}
+            attributes={TextAttributes.BOLD}
+          >
+            {rejoining ? 'Rejoining…' : '[Enter] Rejoin waiting room'}
+          </text>
+        </Button>
+      )}
+    </box>
+  )
+}

cli/src/components/status-bar.tsx

@@ -7,6 +7,7 @@ import { StopButton } from './stop-button'
 import { useTheme } from '../hooks/use-theme'
 import { formatElapsedTime } from '../utils/format-elapsed-time'
 
+import type { FreebuffSessionResponse } from '../types/freebuff-session'
 import type { StatusIndicatorState } from '../utils/status-indicator-state'
 
 
@@ -18,6 +19,7 @@ interface StatusBarProps {
   scrollToLatest: () => void
   statusIndicatorState: StatusIndicatorState
   onStop?: () => void
+  freebuffSession: FreebuffSessionResponse | null
 }
 
 export const StatusBar = ({
@@ -26,6 +28,7 @@ export const StatusBar = ({
   scrollToLatest,
   statusIndicatorState,
   onStop,
+  freebuffSession,
 }: StatusBarProps) => {
   const theme = useTheme()
   const [elapsedSeconds, setElapsedSeconds] = useState(0)
@@ -170,12 +173,12 @@ export const StatusBar = ({
         }}
       >
         <text style={{ wrapMode: 'none' }}>{elapsedTimeContent}</text>
-        <text style={{ wrapMode: 'none' }}>
-          <FreebuffSessionCountdown />
-        </text>
         {onStop && (statusIndicatorState.kind === 'waiting' || statusIndicatorState.kind === 'streaming') && (
           <StopButton onClick={onStop} />
         )}
+        <text style={{ wrapMode: 'none' }}>
+          <FreebuffSessionCountdown session={freebuffSession} />
+        </text>
       </box>
     </box>
   )

cli/src/hooks/helpers/__tests__/send-message.test.ts

@@ -1581,7 +1581,7 @@ describe('freebuff gate errors', () => {
     expect(messages[0].userError).toContain('Another freebuff CLI took over')
   })
 
-  test('handleRunError maps 410 session_expired to the rejoining message', () => {
+  test('handleRunError suppresses the inline error for 410 session_expired (ended banner takes over)', () => {
     const messages = baseMessage()
     const updater = makeUpdater(messages)
     handleRunError({
@@ -1594,10 +1594,13 @@ describe('freebuff gate errors', () => {
       updateChainInProgress: () => {},
     })
     updater.flush()
-    expect(messages[0].userError).toContain('no longer active')
+    // New contract: the gate handler flips the session store into `ended`
+    // and the session-ended banner is the user-facing signal, so we do NOT
+    // also surface an inline userError inside the chat transcript.
+    expect(messages[0].userError).toBeUndefined()
   })
 
-  test('handleRunError maps 428 waiting_room_required to the rejoining message', () => {
+  test('handleRunError suppresses the inline error for 428 waiting_room_required (ended banner takes over)', () => {
     const messages = baseMessage()
     const updater = makeUpdater(messages)
     handleRunError({
@@ -1610,7 +1613,7 @@ describe('freebuff gate errors', () => {
       updateChainInProgress: () => {},
     })
     updater.flush()
-    expect(messages[0].userError).toContain('no longer active')
+    expect(messages[0].userError).toBeUndefined()
   })
 
   test('handleRunError maps 429 waiting_room_queued to the still-queued message', () => {
@@ -1679,6 +1682,10 @@ describe('freebuff gate errors', () => {
       setHasReceivedPlanResponse: () => {},
     })
     updater.flush()
-    expect(messages[0].userError).toContain('no longer active')
+    // 410 is now handled by the ended banner, not an inline error. The
+    // assertion here just confirms routing happened via the gate handler
+    // (which swallows the userError) rather than the generic error path
+    // (which would set a userError from the message).
+    expect(messages[0].userError).toBeUndefined()
   })
 })

cli/src/hooks/helpers/send-message.ts

@@ -1,6 +1,7 @@
 import { getErrorObject } from '@codebuff/common/util/error'
 
 import {
+  markFreebuffSessionEnded,
   markFreebuffSessionSuperseded,
   refreshFreebuffSession,
 } from '../use-freebuff-session'
@@ -507,14 +508,14 @@ function handleFreebuffGateError(
   updater: BatchedMessageUpdater,
 ) {
   switch (kind) {
-    case 'waiting_room_required':
     case 'session_expired':
-      updater.setError(
-        'Your freebuff session is no longer active. Rejoining the waiting room…',
-      )
-      // Re-POST asynchronously; UI flips back to the waiting room as soon as
-      // the store picks up status: 'queued'.
-      refreshFreebuffSession().catch(() => {})
+    case 'waiting_room_required':
+      // Our seat is gone mid-chat. Flip to the client-only `ended` state
+      // instead of auto re-queuing: the Chat surface stays mounted so any
+      // in-flight agent work can finish under the server-side grace period,
+      // and the session-ended banner prompts the user to press Enter when
+      // they're ready to rejoin the waiting room.
+      markFreebuffSessionEnded()
       return
     case 'waiting_room_queued':
       updater.setError(