feat(status-ui): improve bridge recovery and desktop controls

ui/src/App.tsx

@@ -1,4 +1,11 @@
-import { useEffect, useLayoutEffect, useMemo, useRef, useState } from "react";
+import {
+	type KeyboardEvent as ReactKeyboardEvent,
+	useEffect,
+	useLayoutEffect,
+	useMemo,
+	useRef,
+	useState,
+} from "react";
 import { listen } from "@tauri-apps/api/event";
 import { invoke } from "@tauri-apps/api/core";
 import { getCurrentWindow } from "@tauri-apps/api/window";
@@ -8,6 +15,7 @@ type Agent = { id: string; worktreePath: string; state: string };
 type Directory = {
 	worktreePath: string;
 	state: string;
+	isHome: boolean;
 	agentId?: string;
 };
 type Session = {
@@ -122,7 +130,6 @@ const thinkingLevels = [
 	"xhigh",
 	"max",
 ];
-const rememberedKey = "pi-status-ui.remembered-worktrees";
 const selectedPathKey = "pi-status-ui.selected-worktree";
 const controlCommands: Command[] = [
 	{
@@ -198,12 +205,12 @@ function currentTodos(messages: Message[]) {
 	return [];
 }
 
-function rememberedPaths() {
-	try {
-		return JSON.parse(localStorage.getItem(rememberedKey) ?? "[]") as string[];
-	} catch {
-		return [];
-	}
+function focusableExtensionControls(container: HTMLElement) {
+	return [
+		...container.querySelectorAll<HTMLElement>(
+			'button:not(:disabled), input:not(:disabled), select:not(:disabled), textarea:not(:disabled), [href], [tabindex]:not([tabindex="-1"])',
+		),
+	];
 }
 
 function App() {
@@ -221,6 +228,7 @@ function App() {
 	const [commands, setCommands] = useState<Command[]>([]);
 	const [message, setMessage] = useState("");
 	const [folderPath, setFolderPath] = useState("");
+	const [addingDirectory, setAddingDirectory] = useState(false);
 	const [status, setStatus] = useState("Connecting to Pi Status Bridge…");
 	const [view, setView] = useState<"conversation" | "settings">("conversation");
 	const [commandFollowUp, setCommandFollowUp] = useState<FollowUpCommand>();
@@ -237,8 +245,13 @@ function App() {
 		Record<string, ExtensionWidget>
 	>({});
 	const transcriptRef = useRef<HTMLElement>(null);
+	const extensionRef = useRef<HTMLElement>(null);
+	const directoryTabRefs = useRef<Record<string, HTMLButtonElement | null>>({});
 	const pendingSubmissionId = useRef(0);
 	const loadGeneration = useRef(0);
+	const refreshGeneration = useRef(0);
+	const bridgeRetryTimer = useRef<number | undefined>(undefined);
+	const bridgeRetryAttempt = useRef(0);
 	const lastEventSequence = useRef<Record<string, number>>({});
 	const [workProgress, setWorkProgress] = useState<WorkProgress>({
 		phase: "idle",
@@ -247,6 +260,13 @@ function App() {
 	});
 
 	const selected = agents.find((agent) => agent.id === selectedId);
+	const selectedDirectoryIndex = directories.findIndex(
+		(directory) => directory.agentId === selectedId,
+	);
+	const selectedDirectoryTabId =
+		selectedDirectoryIndex >= 0
+			? `directory-tab-${selectedDirectoryIndex}`
+			: undefined;
 	const filteredCommands = useMemo(() => {
 		if (!message.startsWith("/")) return [];
 		const uniqueCommands = [...controlCommands, ...commands].filter(
@@ -282,6 +302,16 @@ function App() {
 			const loadedCommands = unpack(snapshot.commands);
 			const loadedModels = unpack(snapshot.models);
 			setState(loadedState as AgentState);
+			setDirectories((current) =>
+				current.map((directory) =>
+					directory.agentId === agentId
+						? {
+								...directory,
+								state: loadedState.isStreaming ? "streaming" : "idle",
+							}
+						: directory,
+				),
+			);
 			setStats(loadedStats as SessionStats);
 			setMessages((loadedTranscript.messages as Message[] | undefined) ?? []);
 			setCommands((loadedCommands.commands as Command[] | undefined) ?? []);
@@ -304,17 +334,36 @@ function App() {
 				await invoke("subscribe_agent", { agentId, cursor: 0 });
 		} catch (error) {
 			setStatus(String(error));
+			scheduleBridgeReconnect();
 		}
 	}
 
+	function scheduleBridgeReconnect(preferredPath?: string) {
+		if (bridgeRetryTimer.current !== undefined) return;
+		const delay = Math.min(250 * 2 ** bridgeRetryAttempt.current, 2_000);
+		bridgeRetryAttempt.current += 1;
+		setStatus(`Bridge unavailable; retrying in ${delay}ms…`);
+		bridgeRetryTimer.current = window.setTimeout(() => {
+			bridgeRetryTimer.current = undefined;
+			void refreshAgents(preferredPath);
+		}, delay);
+	}
+
 	async function refreshAgents(
 		preferredPath = localStorage.getItem(selectedPathKey) ?? undefined,
 	) {
+		const generation = ++refreshGeneration.current;
 		try {
 			const [agentResponse, directoryResponse] = await Promise.all([
 				invoke<{ agents: Agent[] }>("list_agents"),
 				invoke<{ directories: Directory[] }>("list_directories"),
 			]);
+			if (generation !== refreshGeneration.current) return;
+			if (bridgeRetryTimer.current !== undefined) {
+				window.clearTimeout(bridgeRetryTimer.current);
+				bridgeRetryTimer.current = undefined;
+			}
+			bridgeRetryAttempt.current = 0;
 			setAgents(agentResponse.agents);
 			setDirectories(directoryResponse.directories);
 			const next =
@@ -329,7 +378,8 @@ function App() {
 				await loadAgent(next.id, true);
 			} else setStatus("No Pi agents available");
 		} catch (error) {
-			setStatus(String(error));
+			if (generation !== refreshGeneration.current) return;
+			scheduleBridgeReconnect(preferredPath);
 		}
 	}
 
@@ -346,6 +396,10 @@ function App() {
 
 	useEffect(() => {
 		void refreshAgents();
+		return () => {
+			if (bridgeRetryTimer.current !== undefined)
+				window.clearTimeout(bridgeRetryTimer.current);
+		};
 	}, []);
 
 	useEffect(() => {
@@ -541,6 +595,18 @@ function App() {
 			if (bridgeEvent.type === "extension_ui_request") {
 				handleExtension(bridgeEvent.data?.event as Extension | undefined);
 			} else {
+				if (
+					bridgeEvent.type === "agent_state" &&
+					typeof bridgeEvent.data?.state === "string"
+				) {
+					setDirectories((current) =>
+						current.map((directory) =>
+							directory.agentId === bridgeEvent.agentId
+								? { ...directory, state: bridgeEvent.data?.state as string }
+								: directory,
+						),
+					);
+				}
 				updateWorkProgress(bridgeEvent);
 				if (selectedId) void loadAgent(selectedId);
 			}
@@ -551,7 +617,7 @@ function App() {
 			setStatus(
 				event.payload.message ?? "Bridge connection lost; reconnecting…",
 			);
-			window.setTimeout(() => void refreshAgents(), 500);
+			scheduleBridgeReconnect();
 		}).then((stop) => {
 			unlistenError = stop;
 		});
@@ -580,7 +646,58 @@ function App() {
 		else await activateWorktree(directory.worktreePath);
 	}
 
-	async function startNewSession() {
+	function handleDirectoryTabKeyDown(
+		event: ReactKeyboardEvent<HTMLButtonElement>,
+		index: number,
+	) {
+		let nextIndex: number | undefined;
+		if (event.key === "ArrowRight")
+			nextIndex = (index + 1) % directories.length;
+		else if (event.key === "ArrowLeft")
+			nextIndex = (index - 1 + directories.length) % directories.length;
+		else if (event.key === "Home") nextIndex = 0;
+		else if (event.key === "End") nextIndex = directories.length - 1;
+		if (nextIndex === undefined) return;
+		event.preventDefault();
+		const nextDirectory = directories[nextIndex];
+		directoryTabRefs.current[nextDirectory.worktreePath]?.focus();
+	}
+
+	async function forgetDirectory(directory: Directory) {
+		if (directory.isHome) return;
+		if (directory.state === "streaming") {
+			setStatus("Wait for Pi to finish before forgetting this directory");
+			return;
+		}
+		if (
+			!window.confirm(
+				`Forget ${directory.worktreePath}? Its saved sessions will remain available when you add it again.`,
+			)
+		)
+			return;
+		try {
+			const closingIndex = directories.findIndex(
+				(candidate) => candidate.worktreePath === directory.worktreePath,
+			);
+			const nextFocusPath =
+				directories[closingIndex + 1]?.worktreePath ??
+				directories[closingIndex - 1]?.worktreePath;
+			await invoke("forget_directory", {
+				worktreePath: directory.worktreePath,
+			});
+			const wasSelected = directory.agentId === selectedId;
+			if (wasSelected) localStorage.removeItem(selectedPathKey);
+			await refreshAgents(wasSelected ? "" : selected?.worktreePath);
+			window.requestAnimationFrame(() => {
+				if (nextFocusPath) directoryTabRefs.current[nextFocusPath]?.focus();
+			});
+			setStatus(`Forgot ${worktreeLabel(directory.worktreePath)}`);
+		} catch (error) {
+			setStatus(String(error));
+		}
+	}
+
+	async function startNewSession(successMessage = "Started a new Pi session") {
 		if (!selectedId) return;
 		try {
 			const response = await invoke<{ data?: { cancelled?: boolean } }>(
@@ -594,12 +711,34 @@ function App() {
 			await loadAgent(selectedId, true);
 			await loadSessions(selectedId);
 			setCommandFollowUp(undefined);
-			setStatus("Started a new Pi session");
+			setStatus(successMessage);
 		} catch (error) {
 			setStatus(String(error));
 		}
 	}
 
+	async function closeCurrentSession() {
+		if (
+			!window.confirm(
+				"Close the current session and start a new one? The current history remains available through /resume.",
+			)
+		)
+			return;
+		await startNewSession(
+			"Closed the current session. Use /resume to reopen it.",
+		);
+	}
+
+	async function restartPi() {
+		if (
+			!window.confirm(
+				"Restart Pi for the selected directory? Any in-progress response will be interrupted.",
+			)
+		)
+			return;
+		await invokeAgent("restart");
+	}
+
 	async function switchSession(session: Session) {
 		if (!selectedId || session.isCurrent) return;
 		try {
@@ -649,10 +788,9 @@ function App() {
 			const result = await invoke<{ agent: Agent }>("select_worktree", {
 				worktreePath,
 			});
-			const paths = [...new Set([...rememberedPaths(), worktreePath])];
-			localStorage.setItem(rememberedKey, JSON.stringify(paths));
 			localStorage.setItem(selectedPathKey, worktreePath);
 			setFolderPath("");
+			setAddingDirectory(false);
 			setSelectedId(result.agent.id);
 			await refreshAgents(worktreePath);
 			setView("conversation");
@@ -724,17 +862,25 @@ function App() {
 
 	useEffect(() => {
 		const handleKeydown = (event: KeyboardEvent) => {
-			if (event.key === "Escape") {
-				event.preventDefault();
-				event.stopPropagation();
-				void getCurrentWindow()
-					.hide()
-					.catch((error) => setStatus(String(error)));
+			if (event.key !== "Escape") return;
+			if (pendingExtension) return;
+			event.preventDefault();
+			event.stopPropagation();
+			if (commandFollowUp) {
+				setCommandFollowUp(undefined);
+				return;
 			}
+			if (addingDirectory) {
+				setAddingDirectory(false);
+				return;
+			}
+			void getCurrentWindow()
+				.hide()
+				.catch((error) => setStatus(String(error)));
 		};
 		window.addEventListener("keydown", handleKeydown, true);
 		return () => window.removeEventListener("keydown", handleKeydown, true);
-	}, []);
+	}, [addingDirectory, commandFollowUp, pendingExtension]);
 
 	async function respondToExtension(response: Record<string, unknown>) {
 		if (!selectedId || !pendingExtension?.id) return;
@@ -751,6 +897,37 @@ function App() {
 		}
 	}
 
+	useEffect(() => {
+		if (!pendingExtension) return;
+		const dialog = extensionRef.current;
+		if (!dialog) return;
+		const controls = focusableExtensionControls(dialog);
+		const initial =
+			dialog.querySelector<HTMLElement>("[data-initial-focus]") ?? controls[0];
+		initial?.focus();
+		const handleKeydown = (event: KeyboardEvent) => {
+			if (event.key === "Escape") {
+				event.preventDefault();
+				void respondToExtension({ cancelled: true });
+				return;
+			}
+			if (event.key !== "Tab") return;
+			const focusable = focusableExtensionControls(dialog);
+			const first = focusable[0];
+			const last = focusable[focusable.length - 1];
+			if (!first || !last) return;
+			if (event.shiftKey && document.activeElement === first) {
+				event.preventDefault();
+				last.focus();
+			} else if (!event.shiftKey && document.activeElement === last) {
+				event.preventDefault();
+				first.focus();
+			}
+		};
+		window.addEventListener("keydown", handleKeydown);
+		return () => window.removeEventListener("keydown", handleKeydown);
+	}, [pendingExtension]);
+
 	const extensionMethod = pendingExtension?.method ?? "";
 	const workLabel =
 		workProgress.phase === "recovering"
@@ -789,10 +966,13 @@ function App() {
 							{workLabel}
 						</span>
 					</div>
-					<div className="status-row" role="status">
-						<span title={status}>{status}</span>
+					<div className="status-row">
+						<span role="status" aria-atomic="true" title={status}>
+							{status}
+						</span>
 						{contextWindow && (
 							<span
+								aria-hidden="true"
 								className="status-metric"
 								title={`Current context: ${formatTokens(contextTokens)} / ${formatTokens(contextWindow)} tokens`}
 							>
@@ -802,6 +982,7 @@ function App() {
 						)}
 						{stats.tokens?.total !== undefined && (
 							<span
+								aria-hidden="true"
 								className="status-metric"
 								title={`Session tokens — input: ${formatTokens(stats.tokens.input)}, output: ${formatTokens(stats.tokens.output)}, cache read: ${formatTokens(stats.tokens.cacheRead)}, cache write: ${formatTokens(stats.tokens.cacheWrite)}`}
 							>
@@ -819,6 +1000,8 @@ function App() {
 					</div>
 				</div>
 				<button
+					aria-controls="directory-workspace"
+					aria-expanded={view === "settings"}
 					className="quiet compact-button"
 					onClick={() =>
 						setView(view === "conversation" ? "settings" : "conversation")
@@ -828,23 +1011,86 @@ function App() {
 				</button>
 			</header>
 
-			<section className="agents" aria-label="Managed directories">
-				<span className="agent-caption">Directories</span>
-				{directories.map((directory) => (
-					<button
-						className={
-							directory.agentId === selectedId ? "agent selected" : "agent"
-						}
-						key={directory.worktreePath}
-						title={directory.worktreePath}
-						onClick={() => void chooseDirectory(directory)}
+			<section className="directory-tabs" aria-label="Managed directories">
+				<div
+					aria-label="Managed directories"
+					className="directory-tab-list"
+					role="tablist"
+				>
+					<span className="agent-caption">Directories</span>
+					{directories.map((directory, index) => {
+						const isSelected = directory.agentId === selectedId;
+						return (
+							<div
+								className="directory-tab"
+								key={directory.worktreePath}
+								role="presentation"
+							>
+								<button
+									aria-controls="directory-workspace"
+									aria-label={`${worktreeLabel(directory.worktreePath) ?? directory.worktreePath}, ${directory.state}${directory.isHome ? ", home directory" : ""}`}
+									aria-selected={directory.agentId === selectedId}
+									className={
+										isSelected
+											? "directory-tab-button selected"
+											: "directory-tab-button"
+									}
+									id={`directory-tab-${index}`}
+									onClick={() => void chooseDirectory(directory)}
+									onKeyDown={(event) => handleDirectoryTabKeyDown(event, index)}
+									ref={(element) => {
+										directoryTabRefs.current[directory.worktreePath] = element;
+									}}
+									role="tab"
+									tabIndex={isSelected || (!selectedId && index === 0) ? 0 : -1}
+									title={directory.worktreePath}
+								>
+									<span>
+										{worktreeLabel(directory.worktreePath) ??
+											directory.worktreePath}
+									</span>
+									<small>{directory.state}</small>
+								</button>
+								{!directory.isHome && (
+									<button
+										aria-label={`Forget ${directory.worktreePath}`}
+										className="directory-close"
+										disabled={directory.state === "streaming"}
+										title={`Forget ${directory.worktreePath}`}
+										onClick={() => void forgetDirectory(directory)}
+									>
+										×
+									</button>
+								)}
+							</div>
+						);
+					})}
+				</div>
+				<button
+					aria-expanded={addingDirectory}
+					className="quiet directory-add-toggle"
+					onClick={() => setAddingDirectory((current) => !current)}
+				>
+					{addingDirectory ? "Cancel" : "Add directory"}
+				</button>
+				{addingDirectory && (
+					<form
+						className="directory-add"
+						onSubmit={(event) => {
+							event.preventDefault();
+							void addFolder();
+						}}
 					>
-						<span>
-							{worktreeLabel(directory.worktreePath) ?? directory.worktreePath}
-						</span>
-						<small>{directory.state}</small>
-					</button>
-				))}
+						<input
+							autoFocus
+							value={folderPath}
+							onChange={(event) => setFolderPath(event.currentTarget.value)}
+							placeholder="/absolute/path/to/project"
+							aria-label="Directory path"
+						/>
+						<button type="submit">Add</button>
+					</form>
+				)}
 			</section>
 
 			{extensionNotifications.length > 0 && (
@@ -853,6 +1099,7 @@ function App() {
 						<div
 							className={`notification ${notice.notifyType ?? "info"}`}
 							key={`${notice.id ?? "notice"}-${index}`}
+							role={notice.notifyType === "error" ? "alert" : "status"}
 						>
 							{notice.message ?? "Extension notification"}
 						</div>
@@ -861,32 +1108,13 @@ function App() {
 			)}
 
 			{view === "settings" ? (
-				<section className="settings">
-					<h2>Agents and session</h2>
-					<div className="folder-form">
-						<input
-							value={folderPath}
-							onChange={(event) => setFolderPath(event.currentTarget.value)}
-							placeholder="/absolute/path/to/project"
-						/>
-						<button onClick={() => void addFolder()}>Add folder</button>
-					</div>
-					<div className="remembered">
-						<p className="muted">Remembered folders</p>
-						{rememberedPaths().length ? (
-							rememberedPaths().map((path) => (
-								<button
-									className="quiet"
-									key={path}
-									onClick={() => void activateWorktree(path)}
-								>
-									{path}
-								</button>
-							))
-						) : (
-							<p className="muted">Home only</p>
-						)}
-					</div>
+				<section
+					aria-labelledby={selectedDirectoryTabId}
+					className="settings"
+					id="directory-workspace"
+					role="tabpanel"
+				>
+					<h2>Session controls</h2>
 					<p className="muted">
 						{state.sessionName ?? state.sessionId ?? "Session"} ·{" "}
 						{state.messageCount ?? 0} messages ·{" "}
@@ -904,17 +1132,29 @@ function App() {
 										"Open a directory to view its sessions"}
 								</p>
 							</div>
-							<button
-								disabled={!selectedId || state.isStreaming}
-								onClick={() => void startNewSession()}
-							>
-								New session
-							</button>
+							<div className="session-actions">
+								<button
+									aria-label="Close current session and start a new session"
+									className="quiet"
+									disabled={!selectedId || state.isStreaming}
+									onClick={() => void closeCurrentSession()}
+								>
+									Close current
+								</button>
+								<button
+									disabled={!selectedId || state.isStreaming}
+									onClick={() => void startNewSession()}
+								>
+									New session
+								</button>
+							</div>
 						</div>
 						{sessions.length ? (
 							<div className="session-list">
 								{sessions.map((session) => (
 									<button
+										aria-current={session.isCurrent ? "page" : undefined}
+										aria-label={`${session.name ?? session.firstMessage ?? "Untitled session"}${session.isCurrent ? ", current session" : `, ${session.messageCount} messages`}`}
 										className={
 											session.isCurrent ? "session current" : "session quiet"
 										}
@@ -944,8 +1184,9 @@ function App() {
 							Retry
 						</button>
 						<button
+							aria-label="Restart Pi for the selected directory"
 							className="quiet danger"
-							onClick={() => void invokeAgent("restart")}
+							onClick={() => void restartPi()}
 						>
 							Restart Pi
 						</button>
@@ -953,7 +1194,12 @@ function App() {
 				</section>
 			) : (
 				<>
-					<div className="workflow">
+					<div
+						aria-labelledby={selectedDirectoryTabId}
+						className="workflow"
+						id="directory-workspace"
+						role="tabpanel"
+					>
 						<div className="workflow-main">
 							<section
 								className={`transcript ${pendingExtension ? "with-extension" : ""} ${workProgress.phase}`}
@@ -1227,61 +1473,79 @@ function App() {
 			)}
 
 			{pendingExtension && (
-				<section className="extension" aria-label="Pi extension request">
-					<h2>{pendingExtension.title ?? "Pi extension request"}</h2>
-					<p>{pendingExtension.message ?? "Choose a response."}</p>
-					{extensionMethod === "select" &&
-						pendingExtension.options?.map((option) => (
-							<button
-								key={option}
-								onClick={() => void respondToExtension({ value: option })}
-							>
-								{option}
-							</button>
-						))}
-					{extensionMethod === "confirm" && (
-						<div className="composer-actions">
-							<button
-								onClick={() => void respondToExtension({ confirmed: true })}
-							>
-								Confirm
-							</button>
-							<button
-								className="quiet"
-								onClick={() => void respondToExtension({ confirmed: false })}
-							>
-								Decline
-							</button>
-						</div>
-					)}
-					{(extensionMethod === "input" || extensionMethod === "editor") && (
-						<>
-							<textarea
-								value={extensionValue}
-								onChange={(event) =>
-									setExtensionValue(event.currentTarget.value)
-								}
-								placeholder={
-									pendingExtension.placeholder ?? "Extension response"
-								}
-								rows={4}
-							/>
-							<button
-								onClick={() =>
-									void respondToExtension({ value: extensionValue })
-								}
-							>
-								Submit
-							</button>
-						</>
-					)}
-					<button
-						className="quiet"
-						onClick={() => void respondToExtension({ cancelled: true })}
+				<div className="extension-backdrop">
+					<section
+						className="extension"
+						ref={extensionRef}
+						role="dialog"
+						aria-modal="true"
+						aria-labelledby="extension-title"
+						aria-describedby="extension-description"
 					>
-						Cancel request
-					</button>
-				</section>
+						<h2 id="extension-title">
+							{pendingExtension.title ?? "Pi extension request"}
+						</h2>
+						<p id="extension-description">
+							{pendingExtension.message ?? "Choose a response."}
+						</p>
+						{extensionMethod === "select" &&
+							pendingExtension.options?.map((option) => (
+								<button
+									key={option}
+									onClick={() => void respondToExtension({ value: option })}
+								>
+									{option}
+								</button>
+							))}
+						{extensionMethod === "confirm" && (
+							<div className="composer-actions">
+								<button
+									onClick={() => void respondToExtension({ confirmed: true })}
+								>
+									Confirm
+								</button>
+								<button
+									className="quiet"
+									onClick={() => void respondToExtension({ confirmed: false })}
+								>
+									Decline
+								</button>
+							</div>
+						)}
+						{(extensionMethod === "input" || extensionMethod === "editor") && (
+							<>
+								<label className="sr-only" htmlFor="extension-value">
+									Extension response
+								</label>
+								<textarea
+									id="extension-value"
+									value={extensionValue}
+									onChange={(event) =>
+										setExtensionValue(event.currentTarget.value)
+									}
+									placeholder={
+										pendingExtension.placeholder ?? "Extension response"
+									}
+									rows={4}
+								/>
+								<button
+									onClick={() =>
+										void respondToExtension({ value: extensionValue })
+									}
+								>
+									Submit
+								</button>
+							</>
+						)}
+						<button
+							data-initial-focus
+							className="quiet"
+							onClick={() => void respondToExtension({ cancelled: true })}
+						>
+							Cancel request
+						</button>
+					</section>
+				</div>
 			)}
 		</main>
 	);