From 54a6e63b4cb50f8f93fb78c13e29b6ceab7409bb Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Fri, 4 Sep 2026 11:44:58 -0400 Subject: [PATCH 1/6] Add notebook AI credit budgets and tier-aware model controls --- .../Notebook/AgentChat/AgentChatPanel.tsx | 107 +++-- .../Notebook/AgentChat/ChatComposer.tsx | 7 +- components/Notebook/AgentChat/CreditMeter.tsx | 65 ++++ .../Notebook/AgentChat/ModelControls.tsx | 38 +- hooks/useAgentModelSelection.ts | 106 +++-- hooks/useAgentModels.ts | 56 +-- hooks/useNotebookChat.ts | 83 ++-- hooks/useResearchAI.ts | 77 ++++ services/notebookChat.service.ts | 70 +++- store/researchAI.ts | 125 ++++++ tests/notebook-ai.test.cjs | 365 ++++++++++++++++++ types/notebookChat.ts | 2 +- types/notebookModels.ts | 42 +- types/researchAI.ts | 54 +++ 14 files changed, 1008 insertions(+), 189 deletions(-) create mode 100644 components/Notebook/AgentChat/CreditMeter.tsx create mode 100644 hooks/useResearchAI.ts create mode 100644 store/researchAI.ts create mode 100644 tests/notebook-ai.test.cjs create mode 100644 types/researchAI.ts diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index c6927234e..59dbf0d54 100644 --- a/components/Notebook/AgentChat/AgentChatPanel.tsx +++ b/components/Notebook/AgentChat/AgentChatPanel.tsx @@ -30,6 +30,9 @@ import { ChatPicker } from './ChatPicker'; import { ChatPresets } from './ChatPresets'; import { ChatSources, collectChatSources } from './ChatSources'; import { ChatTranscript } from './ChatTranscript'; +import { CreditMeter } from './CreditMeter'; +import { useResearchAI } from '@/hooks/useResearchAI'; +import { canSelectAIModel } from '@/types/researchAI'; import { ModelControls } from './ModelControls'; import { Logo } from '@/components/ui/Logo'; import { @@ -52,13 +55,18 @@ interface QueuedMessage { /** Pixels per arrow key press while the resize divider has focus. */ const RESIZE_KEY_STEP = 24; -function noticeFromOutcome(outcome: SendOutcome & { ok: false }): ComposerNotice { +function noticeFromOutcome(outcome: SendOutcome & { ok: false }): ComposerNotice | null { switch (outcome.reason) { + case 'usage_limit': + // The shared meter owns this notice, including when the allowance resets. + return null; + case 'account_busy': case 'busy': return { tone: 'warning', text: outcome.detail ?? 'The assistant is still working on a previous message.', }; + case 'model_not_allowed': case 'invalid': return { tone: 'error', text: outcome.detail ?? 'That message can’t be sent.' }; case 'not_found': @@ -220,8 +228,8 @@ interface AgentChatPanelProps { /** * The notebook AI assistant panel: chat picker, transcript with live turn * progress, and composer. Stays mounted while the notebook is open so chat - * selection and drafts survive closing the panel; all network activity is - * gated on `open`. + * selection and drafts survive closing the panel. Chat requests are gated on + * `open`; user-wide allowances load with the notebook. */ export function AgentChatPanel({ noteId, @@ -237,6 +245,11 @@ export function AgentChatPanel({ onReviewChange, }: AgentChatPanelProps) { const { editor, currentNote } = useNotebookContext(); + // This panel stays mounted even when closed: load allowances on notebook open. + const researchAI = useResearchAI(true); + const canSelectModel = + canSelectAIModel(researchAI.budget?.tier) && researchAI.budgetStatus === 'ok'; + const budgetSendDisabled = researchAI.budgetStatus !== 'ok' || researchAI.isSubmissionBlocked(); // Decide which writing preset the empty chat screen offers, and what it // calls the document: the notebook holds RFPs as well as proposals. const noteIsEmpty = useEditorIsEmpty(editor); @@ -277,9 +290,15 @@ export function AgentChatPanel({ // ---- model selection ---- // The catalog loads with the panel. A chat that has already run a turn is // locked to the model it started on, and reports it here; until then the - // browser-level preference decides. + // API default decides. const modelSelection = useAgentModelSelection({ - enabled: open, + enabled: false, + canSelect: canSelectModel, + conversationKey: `${noteId}:${selectedChatId ?? 'new'}`, + locked: + (chatState.chat?.executions.length ?? 0) > 0 || + (chatState.chat?.messages.length ?? 0) > 0 || + chatState.pendingSend !== null, pinnedRef: chatState.pinnedModelRef, effortPinned: chatState.latestExecution != null, pinnedEffort: chatState.latestExecution?.effort ?? null, @@ -346,10 +365,24 @@ export function AgentChatPanel({ // ---- server-side access gate ---- useEffect(() => { - if (list.access === 'hidden' || chatState.access === 'unauthorized') { + // Leave a visible restriction until the user closes the panel. A blocked + // account keeps the entry point so its unavailable state remains reachable. + if ( + !open && + researchAI.budgetStatus !== 'loading' && + researchAI.budget?.tier !== 'blocked' && + (list.access === 'hidden' || chatState.access === 'unauthorized') + ) { onUnavailable(); } - }, [list.access, chatState.access, onUnavailable]); + }, [ + open, + list.access, + chatState.access, + onUnavailable, + researchAI.budgetStatus, + researchAI.budget?.tier, + ]); // ---- keep the listing fresh as the open chat evolves ---- // Derived titles land after the first turn, previews/spinners change as @@ -386,7 +419,7 @@ export function AgentChatPanel({ const handleSend = useCallback(async () => { const text = draft.trim(); - if (!text) return; + if (!text || budgetSendDisabled || chatState.isBusy || creatingChat || queuedMessage) return; setNotice(null); const target = targetRef.current; // Captured before the awaits: the turn runs on what was selected when the @@ -437,6 +470,9 @@ export function AgentChatPanel({ modelSelection.request, updateDraft, isCurrentTarget, + budgetSendDisabled, + creatingChat, + queuedMessage, ]); // Fire the queued first message once the freshly created chat is live. @@ -928,16 +964,13 @@ export function AgentChatPanel({ // ---- derived composer state ---- // Sending before the catalog lands would run the turn on the server default - // and pin the conversation to it, silently losing the user's chosen model - // with no way back. Busy rather than disabled: the draft stays editable, only - // send waits. A catalog that fails resolves to `unavailable`, which sends on - // the server default by design. + // and pin the conversation to it. Keep the draft editable while send waits. const composerBusy = chatState.isBusy || chatState.isFinishing || creatingChat || queuedMessage != null || - modelSelection.status === 'loading'; + (canSelectModel && modelSelection.status === 'loading'); // Stop is only offered once something cancellable exists server-side. While // the message POST is still in flight or the chat is being created, cancel // would no-op and the turn would start anyway. @@ -957,6 +990,20 @@ export function AgentChatPanel({ ); const renderBody = () => { + if ( + researchAI.budget?.tier === 'blocked' || + list.access === 'hidden' || + chatState.access === 'unauthorized' + ) { + return ( +
+ You do not have access to the research assistant for this notebook. +
+ ); + } if (selectedChatId == null) { if (list.access === 'loading') return ; if (list.access === 'error') { @@ -1208,19 +1255,33 @@ export function AgentChatPanel({ busy={composerBusy} canStop={canStop} disabled={composerDisabled} + sendDisabled={budgetSendDisabled} notice={notice} - toolbar={ - { + void researchAI.refreshBudget(true); + }} /> } + toolbar={ + canSelectModel && ( + + ) + } /> ); diff --git a/components/Notebook/AgentChat/ChatComposer.tsx b/components/Notebook/AgentChat/ChatComposer.tsx index c4bf4a799..8d3f73a74 100644 --- a/components/Notebook/AgentChat/ChatComposer.tsx +++ b/components/Notebook/AgentChat/ChatComposer.tsx @@ -25,6 +25,8 @@ interface ChatComposerProps { readonly canStop: boolean; /** Hard-disable everything (chat unavailable). */ readonly disabled: boolean; + readonly sendDisabled?: boolean; + readonly footer?: ReactNode; readonly notice: ComposerNotice | null; readonly placeholder?: string; /** @@ -54,6 +56,8 @@ export function ChatComposer({ busy, canStop, disabled, + sendDisabled = false, + footer, notice, placeholder = 'Ask the assistant…', textareaRef, @@ -67,7 +71,7 @@ export function ChatComposer({ textarea.style.height = `${Math.min(textarea.scrollHeight, 160)}px`; }, [value]); - const canSend = !disabled && !busy && value.trim().length > 0; + const canSend = !disabled && !sendDisabled && !busy && value.trim().length > 0; const handleKeyDown = (event: KeyboardEvent) => { if (event.key === 'Enter' && !event.shiftKey) { @@ -143,6 +147,7 @@ export function ChatComposer({ )} + {footer} {value.length >= COUNTER_THRESHOLD && (

{value.length.toLocaleString()} / {MAX_CHAT_MESSAGE_LENGTH.toLocaleString()} diff --git a/components/Notebook/AgentChat/CreditMeter.tsx b/components/Notebook/AgentChat/CreditMeter.tsx new file mode 100644 index 000000000..964b020b5 --- /dev/null +++ b/components/Notebook/AgentChat/CreditMeter.tsx @@ -0,0 +1,65 @@ +'use client'; + +import type { ResearchAIState } from '@/store/researchAI'; +import { formatBudgetReset, formatCredits, isBudgetExhausted } from '@/types/researchAI'; + +export function CreditMeter({ + budget, + budgetStatus, + limitResetAt, + onRefresh, +}: Pick & { onRefresh: () => void }) { + if (budget?.tier === 'blocked') { + return ( +

+ Research AI is unavailable for this account. +

+ ); + } + if (!budget) { + return ( +

+ {budgetStatus === 'loading' ? 'Loading AI credits…' : 'Couldn’t load AI credits.'} + {budgetStatus === 'unavailable' && ( + + )} +

+ ); + } + const exhausted = isBudgetExhausted(budget) || limitResetAt !== null; + const { remaining, daily_limit: limit } = budget.credits; + const reset = formatBudgetReset(budget.resets_at); + return ( +
+
+ + {limit === null + ? 'Unlimited credits' + : remaining === null + ? 'Credits unavailable' + : `${formatCredits(remaining)} credits remaining`} + + +
+ {exhausted && ( +

+ Daily AI usage limit reached. Available again at {reset}. +

+ )} + {budgetStatus === 'unavailable' && ( +

+ Credits may be out of date.{' '} + +

+ )} +
+ ); +} diff --git a/components/Notebook/AgentChat/ModelControls.tsx b/components/Notebook/AgentChat/ModelControls.tsx index c442e8758..00c1d0908 100644 --- a/components/Notebook/AgentChat/ModelControls.tsx +++ b/components/Notebook/AgentChat/ModelControls.tsx @@ -11,6 +11,7 @@ import { clampTemperature, EFFORT_LABELS, formatTemperature, + formatModelMultiplier, summarizeGenerationOptions, TEMPERATURE_MAX, TEMPERATURE_MIN, @@ -35,6 +36,7 @@ interface ModelControlsProps { readonly onSelectModel: (ref: string) => void; readonly onChangeOptions: (options: GenerationOptions) => void; readonly disabled: boolean; + readonly multiplierExplanation: string; } type OpenMenu = 'model' | 'effort' | null; @@ -60,6 +62,7 @@ export function ModelControls({ onSelectModel, onChangeOptions, disabled, + multiplierExplanation, }: ModelControlsProps) { const [openMenu, setOpenMenu] = useState(null); const containerRef = useRef(null); @@ -133,7 +136,7 @@ export function ModelControls({ {model.label} - {hasEffortMenu && ( + {hasEffortMenu && model.allowed && ( toggle('effort')} open={openMenu === 'effort'} @@ -160,17 +163,20 @@ export function ModelControls({ {openMenu === 'model' && !disabled && !pinned && (
- {models.map((option) => ( - { - setOpenMenu(null); - onSelectModel(option.ref); - }} - /> - ))} + {models + .filter((option) => option.allowed) + .map((option) => ( + { + setOpenMenu(null); + onSelectModel(option.ref); + }} + /> + ))} {models.length === 0 && (

No models are available.

)} @@ -178,7 +184,7 @@ export function ModelControls({
)} - {openMenu === 'effort' && !disabled && ( + {openMenu === 'effort' && !disabled && model.allowed && (
{effortLocked ? ( @@ -326,7 +332,9 @@ function ModelRow({ model, selected, onSelect, + multiplierExplanation, }: { + readonly multiplierExplanation: string; readonly model: AgentModel; readonly selected: boolean; readonly onSelect: () => void; @@ -334,6 +342,7 @@ function ModelRow({ return ( ); } diff --git a/hooks/useAgentModelSelection.ts b/hooks/useAgentModelSelection.ts index fd1b8da5c..0a635e5fc 100644 --- a/hooks/useAgentModelSelection.ts +++ b/hooks/useAgentModelSelection.ts @@ -4,6 +4,7 @@ import { useCallback, useEffect, useMemo, useState } from 'react'; import { useAgentModels, type AgentModelsStatus } from '@/hooks/useAgentModels'; import { findModel, + modelMultiplierExplanation, normalizeGenerationOptions, unknownModel, type AgentModel, @@ -12,37 +13,19 @@ import { type GenerationRequest, } from '@/types/notebookModels'; -const STORAGE_KEY = 'notebook:agent-model'; - -/** Stable empty list so consumers can depend on `models` by identity. */ const NO_MODELS: AgentModel[] = []; - -/** - * The user's raw choices, kept exactly as they made them. Values a given - * model can't take are dropped on the way out rather than on the way in, so - * an effort survives a detour through a model that doesn't offer it. - */ interface StoredPreference extends GenerationOptions { ref?: string; } -function readPreference(): StoredPreference { - try { - const raw = window.localStorage.getItem(STORAGE_KEY); - if (!raw) return {}; - const parsed: unknown = JSON.parse(raw); - return parsed != null && typeof parsed === 'object' ? (parsed as StoredPreference) : {}; - } catch { - // Unparseable, or storage denied — fall back to the server's defaults. - return {}; - } -} - export interface UseAgentModelSelectionOptions { readonly enabled: boolean; + readonly canSelect: boolean; + readonly conversationKey: string; + readonly locked: boolean; /** * The model the open chat's first turn ran on. A conversation keeps its - * model for life, so this — when set — outranks the user's standing choice. + * model for life, so this — when set — outranks the new-chat default. */ readonly pinnedRef: string | null; /** Any recorded turn locks effort, including legacy turns without a model. */ @@ -53,6 +36,7 @@ export interface UseAgentModelSelectionOptions { export interface AgentModelSelection { readonly status: AgentModelsStatus; + readonly multiplierExplanation: string; readonly models: AgentModel[]; /** The model the next turn runs on, or null while there is no catalog. */ readonly model: AgentModel | null; @@ -69,47 +53,38 @@ export interface AgentModelSelection { } /** - * Which model the next turn runs on, and how. - * - * The model choice is a browser-level preference — the last one picked is the - * one a new chat starts on — while a chat already under way reports its own - * pin, which wins. Effort is also fixed after the first turn. Existing chats - * omit it so the server inherits its saved value, regardless of this browser's - * preference. Independent thinking and temperature controls remain per-turn. + * New chats start with the API default and keep choices only for that chat. + * Model and effort lock after the first turn; thinking and temperature remain + * configurable when the saved model/effort combination supports them. */ export function useAgentModelSelection({ enabled, + canSelect, + conversationKey, + locked, pinnedRef, effortPinned, pinnedEffort, }: UseAgentModelSelectionOptions): AgentModelSelection { const { status, catalog } = useAgentModels(enabled); - const [preference, setPreference] = useState({}); - const [hydrated, setHydrated] = useState(false); - - // Read after mount, never during initialization: localStorage is unavailable - // on the server and a differing first client render would hydrate-mismatch. - useEffect(() => { - setPreference(readPreference()); - setHydrated(true); - }, []); - + const [choice, setChoice] = useState<{ key: string; preference: StoredPreference }>({ + key: conversationKey, + preference: {}, + }); + const preference = choice.key === conversationKey ? choice.preference : {}; useEffect(() => { - if (!hydrated) return; - try { - window.localStorage.setItem(STORAGE_KEY, JSON.stringify(preference)); - } catch { - // A blocked or full store just means the choice lasts this session. - } - }, [hydrated, preference]); - - const models = catalog?.models ?? NO_MODELS; + setChoice({ key: conversationKey, preference: {} }); + }, [conversationKey]); + const models = useMemo( + () => catalog?.models.filter((model) => model.allowed) ?? NO_MODELS, + [catalog] + ); const model = useMemo(() => { if (catalog == null) return null; // A pinned ref is named even when the catalog no longer carries it, so a // chat on a retired model still says what it is running. - if (pinnedRef) return findModel(models, pinnedRef) ?? unknownModel(pinnedRef); + if (pinnedRef) return findModel(catalog.models, pinnedRef) ?? unknownModel(pinnedRef); return ( findModel(models, preference.ref ?? null) ?? findModel(models, catalog.default) ?? @@ -125,38 +100,45 @@ export function useAgentModelSelection({ const selectModel = useCallback( (ref: string) => { - if (pinnedRef != null) return; - setPreference((current) => ({ ...current, ref })); + if (!canSelect || locked || !models.some((model) => model.ref === ref)) return; + setChoice((current) => ({ + key: conversationKey, + preference: { ...(current.key === conversationKey ? current.preference : {}), ref }, + })); }, - [pinnedRef] + [canSelect, locked, models, conversationKey] ); const setOptions = useCallback( (next: GenerationOptions) => { - // Locked controls must not overwrite the preference for the next new chat. + if (!canSelect) return; + // A saved effort stays visible, but must never overwrite this chat's + // choice or leak into its later requests. const { effort, ...perTurn } = next; - setPreference((current) => ({ ...current, ...(effortPinned ? perTurn : next) })); + const patch = effortPinned ? perTurn : next; + setChoice((current) => ({ + key: conversationKey, + preference: { ...(current.key === conversationKey ? current.preference : {}), ...patch }, + })); }, - [effortPinned] + [canSelect, conversationKey, effortPinned] ); const request = useMemo(() => { - if (model == null) return {}; - // Saved effort is shown in the controls, but never sent again. Omitting - // pinned values lets the server inherit its own record even if it changed - // since this client last fetched the chat. + if (!canSelect || model == null || !model.allowed) return {}; const { effort, ...perTurn } = options; return { - ...(pinnedRef == null && { model: model.ref }), + ...(!locked && { model: model.ref }), ...(effortPinned ? perTurn : options), }; - }, [model, options, pinnedRef, effortPinned]); + }, [canSelect, model, options, locked, effortPinned]); return { status, + multiplierExplanation: modelMultiplierExplanation(catalog), models, model, - pinned: pinnedRef != null, + pinned: locked, effortPinned, options, selectModel, diff --git a/hooks/useAgentModels.ts b/hooks/useAgentModels.ts index ed4d376bc..99912a649 100644 --- a/hooks/useAgentModels.ts +++ b/hooks/useAgentModels.ts @@ -1,64 +1,16 @@ 'use client'; -import { useEffect, useState } from 'react'; -import { AgentModelService } from '@/services/agentModel.service'; +import { useResearchAI } from '@/hooks/useResearchAI'; import type { AgentModelCatalog } from '@/types/notebookModels'; -/** - * `loading` until the first fetch settles; `unavailable` for every failure — - * the gate, a network blip, a backend without the endpoint. All of them mean - * the same thing to the UI: no picker, and turns run on the server's default. - */ export type AgentModelsStatus = 'loading' | 'ok' | 'unavailable'; - export interface UseAgentModelsResult { readonly status: AgentModelsStatus; readonly catalog: AgentModelCatalog | null; } -// The catalog is per-deployment, not per-note or per-user-action: one fetch -// serves every panel this session. Only successes are cached, so a transient -// failure is retried by the next mount rather than disabling the picker for -// the rest of the session. -let cachedCatalog: AgentModelCatalog | null = null; -let inFlight: Promise | null = null; - -function loadCatalog(): Promise { - if (cachedCatalog) return Promise.resolve(cachedCatalog); - inFlight ??= AgentModelService.listModels() - .then((catalog) => { - cachedCatalog = catalog; - return catalog; - }) - .finally(() => { - inFlight = null; - }); - return inFlight; -} - -/** The models this user may select. Fetched once, when something needs it. */ +/** Availability is user-specific and refreshed with the shared AI budget. */ export function useAgentModels(enabled: boolean): UseAgentModelsResult { - const [result, setResult] = useState(() => - cachedCatalog ? { status: 'ok', catalog: cachedCatalog } : { status: 'loading', catalog: null } - ); - - useEffect(() => { - if (!enabled || result.catalog != null) return; - let cancelled = false; - loadCatalog().then( - (catalog) => { - if (!cancelled) setResult({ status: 'ok', catalog }); - }, - () => { - // Deliberately terminal for this mount: the deps can't change back, so - // a failing endpoint is asked once, not once per render. - if (!cancelled) setResult({ status: 'unavailable', catalog: null }); - } - ); - return () => { - cancelled = true; - }; - }, [enabled, result.catalog]); - - return result; + const { catalog, catalogStatus } = useResearchAI(enabled); + return { catalog, status: catalogStatus }; } diff --git a/hooks/useNotebookChat.ts b/hooks/useNotebookChat.ts index e6fd132f4..51a08cc38 100644 --- a/hooks/useNotebookChat.ts +++ b/hooks/useNotebookChat.ts @@ -4,7 +4,9 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { debounce, type DebouncedFunc } from 'lodash-es'; import { NotebookChatService, - chatErrorDetail, + chatErrorBody, + sendFailureOutcome, + type SendOutcome, chatErrorStatus, } from '@/services/notebookChat.service'; import { useNotebookChatSocket, type ChatSocketStatus } from '@/hooks/useNotebookChatSocket'; @@ -21,6 +23,8 @@ import { type NotebookChat, type NotebookChatListItem, } from '@/types/notebookChat'; +import { useResearchAI } from '@/hooks/useResearchAI'; +import { canSelectAIModel } from '@/types/researchAI'; import type { GenerationRequest } from '@/types/notebookModels'; /** Fallback poll cadence while a turn runs; the socket nudge usually wins. */ @@ -40,31 +44,7 @@ const STREAM_CHAR_CAPS: Record = { export type ChatAccess = 'loading' | 'ok' | 'not_found' | 'unauthorized' | 'error'; -export type SendOutcome = - | { ok: true } - | { - ok: false; - reason: 'busy' | 'invalid' | 'not_found' | 'unauthorized' | 'error'; - detail?: string; - }; - -/** Maps a failed send POST to its outcome; the state side-effects stay in `send`. */ -function sendFailureOutcome(err: unknown): Extract { - const detail = chatErrorDetail(err); - switch (chatErrorStatus(err)) { - case 409: - return { ok: false, reason: 'busy', detail }; - case 400: - return { ok: false, reason: 'invalid', detail }; - case 401: - case 403: - return { ok: false, reason: 'unauthorized', detail }; - case 404: - return { ok: false, reason: 'not_found', detail }; - default: - return { ok: false, reason: 'error', detail }; - } -} +export type { SendOutcome } from '@/services/notebookChat.service'; export interface PendingSend { text: string; @@ -281,6 +261,8 @@ export function useNotebookChat({ enabled, initialChat = null, }: UseNotebookChatOptions): UseNotebookChatResult { + const { refreshBudget, refreshCatalog, recordLimit, getSnapshot, isSubmissionBlocked } = + useResearchAI(); const [chat, setChat] = useState(null); const [access, setAccess] = useState('loading'); const [pendingSend, setPendingSend] = useState(null); @@ -316,6 +298,19 @@ export function useNotebookChat({ try { const data = await NotebookChatService.getChat(noteId, chatId, { live }); if (seq !== seqRef.current) return; + const previous = chatRef.current; + const settled = data.executions.filter( + (execution) => + !isActiveExecutionStatus(execution.status) && + previous?.executions.find((cached) => cached.id === execution.id)?.status !== + execution.status + ); + for (const execution of settled) { + if (execution.error?.code === 'usage_limit_exceeded') { + recordLimit(undefined, execution.finished_at ?? execution.started_at); + } + } + if (settled.length > 0) void refreshBudget(true); setChat((prev) => { const merged = mergeLiveChat(live ? prev : null, data); chatRef.current = merged; @@ -337,7 +332,7 @@ export function useNotebookChat({ } } }, - [noteId, chatId] + [noteId, chatId, refreshBudget, recordLimit] ); // Reset + initial load whenever the target chat changes or the panel opens. @@ -408,12 +403,13 @@ export function useNotebookChat({ const timer = setInterval(() => { if (inFlight) return; inFlight = true; + void refreshBudget(); fetchChat('live').finally(() => { inFlight = false; }); }, POLL_INTERVAL_MS); return () => clearInterval(timer); - }, [enabled, access, isBusy, isFinishing, fetchChat]); + }, [enabled, access, isBusy, isFinishing, fetchChat, refreshBudget]); // Debounced lifecycle nudge / stream-gap repair → live refetch. const nudgeRef = useRef void> | null>(null); @@ -447,6 +443,7 @@ export function useNotebookChat({ (event: ChatSocketEvent) => { if (event.conversation_id !== chatId) return; if (!isChatStreamSocketEvent(event)) { + void refreshBudget(['turn_finished', 'turn_failed', 'turn_cancelled'].includes(event.kind)); nudgeRef.current?.(); return; } @@ -463,7 +460,7 @@ export function useNotebookChat({ setChat(applied.chat); } }, - [chatId, repairStream] + [chatId, repairStream, refreshBudget] ); const handleSocketReconnect = useCallback(() => { @@ -482,10 +479,17 @@ export function useNotebookChat({ const send = useCallback( async (text: string, generation?: GenerationRequest): Promise => { if (noteId == null || chatId == null) return { ok: false, reason: 'error' }; + if (getSnapshot().budget?.tier === 'blocked') return { ok: false, reason: 'unauthorized' }; + if (isSubmissionBlocked()) return { ok: false, reason: 'usage_limit' }; const epoch = epochRef.current; setPendingSend({ text, executionId: null }); try { - const response = await NotebookChatService.sendMessage(noteId, chatId, text, generation); + const response = await NotebookChatService.sendMessage( + noteId, + chatId, + text, + canSelectAIModel(getSnapshot().budget?.tier) ? generation : undefined + ); if (epoch === epochRef.current) { setPendingSend({ text, executionId: response.execution_id }); fetchChat('live'); @@ -493,12 +497,15 @@ export function useNotebookChat({ return { ok: true }; } catch (err) { const outcome = sendFailureOutcome(err); + if (outcome.reason === 'usage_limit') recordLimit(chatErrorBody(err)); + else void refreshBudget(true); + if (outcome.reason === 'model_not_allowed') void refreshCatalog(); // The outcome is still reported either way, but a continuation for a // chat that is no longer selected must not mutate the current one. if (epoch === epochRef.current) { setPendingSend(null); // Raced an active turn — refetch so the busy state renders truthfully. - if (outcome.reason === 'busy') fetchChat('live'); + if (outcome.reason === 'busy' || outcome.reason === 'account_busy') fetchChat('live'); if (outcome.reason === 'not_found') setAccess('not_found'); // Session expired or permission revoked mid-chat: mirror what a // failed GET does so the access gate reacts instead of the composer @@ -511,7 +518,16 @@ export function useNotebookChat({ return outcome; } }, - [noteId, chatId, fetchChat] + [ + noteId, + chatId, + fetchChat, + refreshBudget, + refreshCatalog, + recordLimit, + getSnapshot, + isSubmissionBlocked, + ] ); const cancel = useCallback(async () => { @@ -523,8 +539,9 @@ export function useNotebookChat({ } catch { // Fall through: the refetch below renders whatever actually happened. } + void refreshBudget(true); if (epoch === epochRef.current) fetchChat('live'); - }, [noteId, chatId, fetchChat]); + }, [noteId, chatId, fetchChat, refreshBudget]); const rename = useCallback( async (title: string): Promise => { diff --git a/hooks/useResearchAI.ts b/hooks/useResearchAI.ts new file mode 100644 index 000000000..5c8d57d3f --- /dev/null +++ b/hooks/useResearchAI.ts @@ -0,0 +1,77 @@ +'use client'; + +import { useEffect, useMemo, useSyncExternalStore } from 'react'; +import { useSession } from 'next-auth/react'; +import { ApiClient } from '@/services/client'; +import { AgentModelService } from '@/services/agentModel.service'; +import { createResearchAIStore, INITIAL_RESEARCH_AI_STATE } from '@/store/researchAI'; + +const stores = new Map>(); +const createStore = () => + createResearchAIStore({ + budget: () => ApiClient.get('/api/research_ai/usage-budget/'), + catalog: () => AgentModelService.listModels(), + }); + +/** Session-scoped memory only: never persist allowances or share them between accounts. */ +export function useResearchAI(enabled = false) { + const { data: session } = useSession(); + const token = session?.authToken; + const store = useMemo(() => { + if (!token || typeof window === 'undefined') return createStore(); + let current = stores.get(token); + if (!current) { + current = createStore(); + stores.set(token, current); + } + return current; + }, [token]); + const state = useSyncExternalStore( + store.subscribe, + store.getSnapshot, + () => INITIAL_RESEARCH_AI_STATE + ); + + useEffect(() => { + if (!enabled || !token) return; + const refresh = () => { + void store.refreshBudget(true); + void store.refreshCatalog(); + }; + refresh(); + window.addEventListener('focus', refresh); + return () => window.removeEventListener('focus', refresh); + }, [enabled, token, store]); + + // Refresh at reset even when the composer is idle. Failed/stale reset reads + // retry with a bounded delay rather than leaving Send disabled all day. + useEffect(() => { + if (!enabled || !token) return; + const resetsAt = state.budget?.resets_at ?? state.limitResetAt; + if (!resetsAt) return; + let timer: ReturnType; + let cancelled = false; + const schedule = () => { + timer = setTimeout( + async () => { + await store.refreshBudget(true); + if (!cancelled) schedule(); + }, + Math.min( + 2_147_483_647, + Math.max( + 5_000, + Date.parse(store.getSnapshot().budget?.resets_at ?? resetsAt) - Date.now() + 1000 + ) + ) + ); + }; + schedule(); + return () => { + cancelled = true; + clearTimeout(timer); + }; + }, [enabled, token, store, state.budget?.resets_at, state.limitResetAt]); + + return { ...state, ...store }; +} diff --git a/services/notebookChat.service.ts b/services/notebookChat.service.ts index 0a2940b2d..2f6bde375 100644 --- a/services/notebookChat.service.ts +++ b/services/notebookChat.service.ts @@ -93,9 +93,77 @@ export function chatErrorStatus(error: unknown): number | undefined { */ export function chatErrorDetail(error: unknown): string | undefined { if (error instanceof ApiError) { - const detail = (error.errors as Record | undefined)?.detail; + const fields = error.errors as Record | undefined; + const detail = fields?.detail; if (typeof detail === 'string' && detail.length > 0) return detail; + if (fields) { + const messages = Object.entries(fields).flatMap(([field, value]) => + Array.isArray(value) + ? value + .filter((entry): entry is string => typeof entry === 'string') + .map((entry) => `${field}: ${entry}`) + : [] + ); + if (messages.length) return messages.join(' '); + } return error.message; } return error instanceof Error ? error.message : undefined; } + +export function chatErrorCode(error: unknown): string | undefined { + const code = chatErrorBody(error)?.code; + return typeof code === 'string' ? code : undefined; +} + +export function chatErrorBody(error: unknown): Record | undefined { + return error instanceof ApiError + ? (error.errors as Record | undefined) + : undefined; +} + +export type SendOutcome = + | { ok: true } + | { + ok: false; + reason: + | 'usage_limit' + | 'account_busy' + | 'model_not_allowed' + | 'busy' + | 'invalid' + | 'not_found' + | 'unauthorized' + | 'error'; + detail?: string; + }; + +/** Maps a failed send POST to its outcome; the state side-effects stay in `send`. */ +export function sendFailureOutcome(err: unknown): Extract { + const detail = chatErrorDetail(err); + const code = chatErrorCode(err); + switch (chatErrorStatus(err)) { + case 429: + return code === 'usage_limit_exceeded' + ? { ok: false, reason: 'usage_limit', detail: 'Daily AI usage limit reached.' } + : { ok: false, reason: 'error', detail }; + + case 409: + return code === 'usage_work_in_progress' + ? { ok: false, reason: 'account_busy', detail: 'Another AI request is still running.' } + : { ok: false, reason: 'busy', detail: 'This conversation already has an active turn.' }; + case 400: + return { + ok: false, + reason: code === 'model_not_allowed' ? 'model_not_allowed' : 'invalid', + detail, + }; + case 401: + case 403: + return { ok: false, reason: 'unauthorized', detail }; + case 404: + return { ok: false, reason: 'not_found', detail }; + default: + return { ok: false, reason: 'error', detail }; + } +} diff --git a/store/researchAI.ts b/store/researchAI.ts new file mode 100644 index 000000000..d8d215b47 --- /dev/null +++ b/store/researchAI.ts @@ -0,0 +1,125 @@ +import type { AgentModelCatalog } from '@/types/notebookModels'; +import { isBudgetExhausted, isResearchAIBudget, type ResearchAIBudget } from '@/types/researchAI'; + +export interface ResearchAIState { + budget: ResearchAIBudget | null; + budgetStatus: 'loading' | 'ok' | 'unavailable'; + catalog: AgentModelCatalog | null; + catalogStatus: 'loading' | 'ok' | 'unavailable'; + /** A provider may reject its next call while some recorded credits remain. */ + limitResetAt: string | null; +} + +export const INITIAL_RESEARCH_AI_STATE: ResearchAIState = { + budget: null, + budgetStatus: 'loading', + catalog: null, + catalogStatus: 'loading', + limitResetAt: null, +}; + +/** One store per authenticated session, shared across notes and AI workflows. */ +export function createResearchAIStore(loaders: { + budget: () => Promise; + catalog: () => Promise; +}) { + let state = INITIAL_RESEARCH_AI_STATE; + const listeners = new Set<() => void>(); + let budgetFlight: Promise | null = null; + let catalogFlight: Promise | null = null; + let budgetQueued = false; + let budgetRevision = 0; + let lastBudgetFetch = 0; + + const update = (patch: Partial) => { + state = { ...state, ...patch }; + listeners.forEach((listener) => listener()); + }; + const acceptBudget = (budget: ResearchAIBudget) => { + update({ + budget, + budgetStatus: 'ok', + limitResetAt: + state.limitResetAt && Date.parse(budget.resets_at) <= Date.parse(state.limitResetAt) + ? state.limitResetAt + : null, + }); + }; + + const refreshBudget = (force = false): Promise => { + if (budgetFlight) { + // A terminal event may follow the snapshot of the currently running GET. + if (force) budgetQueued = true; + return budgetFlight; + } + if (!force && Date.now() - lastBudgetFetch < 15_000) return Promise.resolve(); + lastBudgetFetch = Date.now(); + const revision = budgetRevision; + budgetFlight = loaders + .budget() + .then((value) => { + if (revision !== budgetRevision) return; + if (!isResearchAIBudget(value)) throw new Error('Credit budget unavailable'); + acceptBudget(value); + }) + .catch(() => { + if (revision === budgetRevision) update({ budgetStatus: 'unavailable' }); + }) + .finally(() => { + budgetFlight = null; + if (budgetQueued) { + budgetQueued = false; + void refreshBudget(true); + } + }); + return budgetFlight; + }; + + const refreshCatalog = (): Promise => { + if (catalogFlight) return catalogFlight; + catalogFlight = loaders + .catalog() + .then((catalog) => update({ catalog, catalogStatus: 'ok' })) + .catch(() => update({ catalog: null, catalogStatus: 'unavailable' })) + .finally(() => { + catalogFlight = null; + }); + return catalogFlight; + }; + + return { + getSnapshot: () => state, + subscribe: (listener: () => void) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }, + refreshBudget, + refreshCatalog, + /** 429 budget is top-level, and must outrank an older in-flight GET. */ + recordLimit: (value?: unknown, occurredAt?: string | null) => { + // Reopening yesterday's failed conversation must not exhaust today's allowance. + if ( + occurredAt && + Number.isFinite(Date.parse(occurredAt)) && + new Date(occurredAt).toISOString().slice(0, 10) !== new Date().toISOString().slice(0, 10) + ) + return; + budgetRevision += 1; + if (isResearchAIBudget(value)) acceptBudget(value); + const nextReset = new Date(); + nextReset.setUTCHours(24, 0, 0, 0); + const knownReset = state.budget?.resets_at; + update({ + limitResetAt: + knownReset && Date.parse(knownReset) > Date.now() ? knownReset : nextReset.toISOString(), + }); + void refreshBudget(true); + }, + isSubmissionBlocked: () => + state.budget?.tier === 'blocked' || + isBudgetExhausted(state.budget) || + state.limitResetAt !== null, + }; +} diff --git a/tests/notebook-ai.test.cjs b/tests/notebook-ai.test.cjs new file mode 100644 index 000000000..2fd308a92 --- /dev/null +++ b/tests/notebook-ai.test.cjs @@ -0,0 +1,365 @@ +/* eslint-disable @typescript-eslint/no-require-imports -- Standalone CommonJS Node test runner. */ +// Run with: node --test tests/notebook-ai.test.cjs +// Uses the repository's TypeScript compiler and Node's test runner; no extra test dependencies. +const { test } = require('node:test'); +const assert = require('node:assert/strict'); +const { readFileSync, existsSync } = require('node:fs'); +const path = require('node:path'); +const ts = require('typescript'); +const root = path.resolve(__dirname, '..'); +const cache = new Map(); +const api = {}; +function load(relative) { + const filename = path.resolve(root, relative); + if (cache.has(filename)) return cache.get(filename).exports; + const compiledModule = { exports: {} }; + cache.set(filename, compiledModule); + const source = ts.transpileModule(readFileSync(filename, 'utf8'), { + compilerOptions: { + module: ts.ModuleKind.CommonJS, + target: ts.ScriptTarget.ES2022, + jsx: ts.JsxEmit.ReactJSX, + }, + }).outputText; + const localRequire = (specifier) => { + if ( + specifier === '@/services/client' || + (specifier === './client' && relative.startsWith('services/')) + ) + return { ApiClient: api }; + if (specifier === '@/hooks/useAgentModels') + return { useAgentModels: () => ({ status: 'ok', catalog }) }; + if (!specifier.startsWith('@/') && !specifier.startsWith('.')) return require(specifier); + const base = specifier.startsWith('@/') + ? path.join(root, specifier.slice(2)) + : path.resolve(path.dirname(filename), specifier); + const target = [base + '.ts', base + '.tsx', path.join(base, 'index.ts')].find(existsSync); + return load(path.relative(root, target)); + }; + new Function('require', 'module', 'exports', source)( + localRequire, + compiledModule, + compiledModule.exports + ); + return compiledModule.exports; +} +const budgetTypes = load('types/researchAI.ts'); +const models = load('types/notebookModels.ts'); +const { createResearchAIStore } = load('store/researchAI.ts'); +const { ApiError } = load('services/types/api.ts'); +const service = load('services/notebookChat.service.ts'); +const { renderToStaticMarkup } = require('react-dom/server'); +const { createElement } = require('react'); +const { CreditMeter } = load('components/Notebook/AgentChat/CreditMeter.tsx'); +const { ChatComposer } = load('components/Notebook/AgentChat/ChatComposer.tsx'); +const { ModelControls } = load('components/Notebook/AgentChat/ModelControls.tsx'); +const tomorrow = new Date(); +tomorrow.setUTCHours(24, 0, 0, 0); +const budget = (overrides = {}) => ({ + tier: 'default', + credits: { daily_limit: '250', used: '1.65', remaining: '248.35' }, + turns_used: 2, + turn_cap: 10, + resets_at: tomorrow.toISOString(), + ...overrides, +}); +const catalog = models.toAgentModelCatalog({ + default: 'openrouter:test', + credit_pricing: { + multiplier_base_model: 'openrouter:base', + multiplier_basis: 'equal_input_output_tokens', + multiplier_is_estimate: true, + }, + models: [ + { + ref: 'openrouter:test', + label: 'Flash', + allowed: true, + multiplier: '0.03', + capabilities: { effort: ['low', 'high'], thinking: [], temperature: false }, + }, + { ref: 'openrouter:base', label: 'Baseline', allowed: true, multiplier: '1' }, + { ref: 'openrouter:unpriced', allowed: false, multiplier: null }, + ], +}); +const storeWith = (getBudget = async () => budget()) => + createResearchAIStore({ budget: getBudget, catalog: async () => catalog }); +const flush = () => new Promise((resolve) => setImmediate(resolve)); + +test('credit exhaustion and provider-call cap are independent; null means unlimited', () => { + assert.equal(budgetTypes.isBudgetExhausted(budget()), false); + assert.equal( + budgetTypes.isBudgetExhausted( + budget({ credits: { daily_limit: '250', used: '250', remaining: '0.00' } }) + ), + true + ); + assert.equal(budgetTypes.isBudgetExhausted(budget({ turns_used: 10 })), true); + assert.equal( + budgetTypes.isBudgetExhausted( + budget({ + credits: { daily_limit: null, used: '999', remaining: null }, + turn_cap: null, + turns_used: 999, + }) + ), + false + ); + assert.equal(budgetTypes.isResearchAIBudget({ tier: 'default', remaining: '100' }), false); +}); + +test('tiers and catalog permissions are authoritative; fractions are not rounded to zero', () => { + assert.equal(budgetTypes.canSelectAIModel('default'), false); + assert.equal(budgetTypes.canSelectAIModel('blocked'), false); + assert.equal(budgetTypes.canSelectAIModel('invited'), true); + assert.equal(budgetTypes.canSelectAIModel('privileged'), true); + assert.equal(catalog.models[2].allowed, false); + assert.equal(models.formatModelMultiplier('0.03'), '0.03×'); + assert.equal(models.formatModelMultiplier('3.75'), '3.75×'); + assert.equal(models.formatModelMultiplier('0.001'), '<0.01×'); + assert.equal(models.formatModelMultiplier(null), 'Pricing unavailable'); + assert.equal(budgetTypes.formatCredits('0.00010'), '0.0001'); + assert.match(models.modelMultiplierExplanation(catalog), /relative to Baseline/); + assert.deepEqual( + models.normalizeGenerationOptions(catalog.models[0], { + effort: 'high', + thinking: 'adaptive', + temperature: 1, + }), + { effort: 'high' } + ); +}); + +test('meter shows fractional credits, daily cap exhaustion, local reset and unlimited balances', () => { + const render = (b) => + renderToStaticMarkup( + createElement(CreditMeter, { + budget: b, + budgetStatus: 'ok', + limitResetAt: null, + onRefresh() {}, + }) + ); + assert.match(render(budget()), /248.35 credits remaining/); + assert.match(render(budget()), /250 daily credits/); + assert.match(render(budget()), /Resets at/); + const capped = render(budget({ turns_used: 10 })); + assert.match(capped, /Daily AI usage limit reached/); + assert.doesNotMatch(capped, /Out of credits|messages remaining/); + assert.match( + render(budget({ credits: { daily_limit: null, remaining: null, used: '1' } })), + /Unlimited credits/ + ); + const previousTZ = process.env.TZ; + process.env.TZ = 'America/New_York'; + assert.match(budgetTypes.formatBudgetReset('2026-09-05T00:00:00Z'), /8:00/); + if (previousTZ === undefined) delete process.env.TZ; + else process.env.TZ = previousTZ; +}); + +test('exhaustion disables Send but retains editable draft and Stop', () => { + const props = { + value: 'Keep my unsent question', + onChange() {}, + onSend() {}, + onStop() {}, + busy: false, + canStop: false, + disabled: false, + sendDisabled: true, + notice: null, + textareaRef: { current: null }, + }; + const html = renderToStaticMarkup(createElement(ChatComposer, props)); + assert.match(html, /Keep my unsent question/); + assert.doesNotMatch(html, /]* disabled=""/); + assert.match(html, /]*disabled=""[^>]*title="Send message"/); + const running = renderToStaticMarkup( + createElement(ChatComposer, { ...props, busy: true, canStop: true }) + ); + assert.match(running, /title="Stop the assistant"/); + assert.doesNotMatch(running, /]* disabled=""/); +}); + +test('pinned model control is disabled and explains how to switch', () => { + const html = renderToStaticMarkup( + createElement(ModelControls, { + models: catalog.models, + model: catalog.models[0], + pinned: true, + options: {}, + onSelectModel() {}, + onChangeOptions() {}, + disabled: false, + multiplierExplanation: models.modelMultiplierExplanation(catalog), + }) + ); + assert.match(html, /]*disabled=""/); + assert.match(html, /Start a new chat to switch models/); +}); + +test('shared subscribers see one fetch; progress refreshes throttle and sessions stay isolated', async () => { + let calls = 0; + const store = storeWith(async () => { + calls++; + return budget(); + }); + const snapshots = []; + const unsubscribe = store.subscribe(() => snapshots.push(store.getSnapshot())); + await Promise.all([store.refreshBudget(), store.refreshBudget()]); + await store.refreshBudget(); + assert.equal(calls, 1); + assert.equal(snapshots.at(-1).budget.credits.remaining, '248.35'); + assert.equal(storeWith().getSnapshot().budget, null); + unsubscribe(); +}); + +test('429 budget wins over an older GET, with a follow-up refresh after settlement', async () => { + let resolveOld; + let calls = 0; + const store = storeWith(() => + ++calls === 1 + ? new Promise((resolve) => { + resolveOld = resolve; + }) + : Promise.resolve(budget()) + ); + const pending = store.refreshBudget(true); + store.recordLimit(budget({ turns_used: 10 })); + assert.equal(store.getSnapshot().budget.turns_used, 10); + resolveOld(budget()); + await pending; + await flush(); + assert.equal(calls, 2); + assert.equal(store.isSubmissionBlocked(), true); +}); + +test('a post-202 limit stays blocked with credits remaining and recovers after reset', async () => { + let current = budget(); + const store = storeWith(async () => current); + await store.refreshBudget(); + store.recordLimit(); + await flush(); + assert.equal(store.isSubmissionBlocked(), true); + const nextReset = new Date(Date.parse(current.resets_at) + 86400000).toISOString(); + current = budget({ resets_at: nextReset, turns_used: 0 }); + await store.refreshBudget(true); + assert.equal(store.isSubmissionBlocked(), false); +}); + +test('opening an old failed turn does not block today, and cancellation refresh never refunds', async () => { + const store = storeWith(); + await store.refreshBudget(); + store.recordLimit(undefined, '2020-01-01T12:00:00Z'); + assert.equal(store.isSubmissionBlocked(), false); + await store.refreshBudget(true); + assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); +}); + +test('budget failure keeps the recorded balance and catalog refresh removes withdrawn choices', async () => { + let fail = false; + let available = catalog; + const store = createResearchAIStore({ + budget: async () => { + if (fail) throw Error('offline'); + return budget(); + }, + catalog: async () => available, + }); + await store.refreshBudget(); + fail = true; + await store.refreshBudget(true); + assert.equal(store.getSnapshot().budgetStatus, 'unavailable'); + assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); + await store.refreshCatalog(); + available = { ...catalog, models: [] }; + await store.refreshCatalog(); + assert.equal(store.getSnapshot().catalog.models.length, 0); +}); + +test('errors preserve structured codes, top-level budgets and ordinary field validation', () => { + const limit = new ApiError('Request failed', 429, { ...budget(), code: 'usage_limit_exceeded' }); + assert.equal(service.chatErrorCode(limit), 'usage_limit_exceeded'); + assert.equal(service.chatErrorBody(limit).credits.remaining, '248.35'); + assert.equal( + service.chatErrorDetail(new ApiError('Request failed', 400, { message: ['Too long.'] })), + 'message: Too long.' + ); + assert.equal( + service.chatErrorDetail(new ApiError('Request failed', 400, { detail: 'Model unavailable.' })), + 'Model unavailable.' + ); + assert.equal( + service.chatErrorCode(new ApiError('Request failed', 409, { code: 'usage_work_in_progress' })), + 'usage_work_in_progress' + ); +}); + +test('default-tier sends only message; selected model is omitted once the conversation is locked', async () => { + const { useAgentModelSelection } = load('hooks/useAgentModelSelection.ts'); + let selection; + function Probe(props) { + selection = useAgentModelSelection({ + enabled: false, + conversationKey: 'new', + pinnedRef: null, + locked: false, + ...props, + }); + return null; + } + let sent; + api.post = async (url, body) => { + sent = { url, body }; + return { execution_id: 42 }; + }; + renderToStaticMarkup(createElement(Probe, { canSelect: false })); + await service.NotebookChatService.sendMessage( + 1, + 2, + 'Summarize this notebook.', + selection.request + ); + assert.deepEqual(sent.body, { message: 'Summarize this notebook.' }); + renderToStaticMarkup(createElement(Probe, { canSelect: true })); + assert.equal(selection.model.ref, catalog.default); + assert.equal(selection.request.model, catalog.default); + renderToStaticMarkup( + createElement(Probe, { canSelect: true, locked: true, pinnedRef: 'openrouter:base' }) + ); + assert.equal(selection.model.ref, 'openrouter:base'); + assert.equal(selection.request.model, undefined); +}); + +test('immediate resubmission after cancellation handles account-wide 409 without classifying it as exhaustion', async () => { + api.post = async (url) => { + if (url.endsWith('/cancel/')) return { cancelled: true, execution_id: 42 }; + throw new ApiError('Request failed', 409, { code: 'usage_work_in_progress' }); + }; + await service.NotebookChatService.cancelTurn(1, 2); + await assert.rejects( + service.NotebookChatService.sendMessage(1, 2, 'Preserved draft'), + (error) => { + const outcome = service.sendFailureOutcome(error); + assert.equal(outcome.reason, 'account_busy'); + assert.equal(outcome.detail, 'Another AI request is still running.'); + return true; + } + ); + assert.equal(service.sendFailureOutcome(new ApiError('busy', 409)).reason, 'busy'); + assert.equal(service.sendFailureOutcome(new ApiError('forbidden', 403)).reason, 'unauthorized'); + assert.equal( + service.sendFailureOutcome( + new ApiError('invalid', 400, { code: 'model_not_allowed', detail: 'Choose another model.' }) + ).reason, + 'model_not_allowed' + ); + assert.equal( + service.sendFailureOutcome(new ApiError('invalid', 400, { message: ['Too long.'] })).reason, + 'invalid' + ); + assert.equal( + service.sendFailureOutcome(new ApiError('limit', 429, { code: 'usage_limit_exceeded' })).reason, + 'usage_limit' + ); +}); diff --git a/types/notebookChat.ts b/types/notebookChat.ts index 36d9703f6..afb2a6723 100644 --- a/types/notebookChat.ts +++ b/types/notebookChat.ts @@ -156,7 +156,7 @@ export interface ChatExecution { /** Heartbeat, stamped on every durable write. */ last_activity_at: string | null; iterations: number; - max_iterations: number; + max_iterations: number | null; /** True while the turn succeeded but its answer hasn't landed in `messages` yet. */ assistant_message_pending: boolean; error: ChatExecutionError | null; diff --git a/types/notebookModels.ts b/types/notebookModels.ts index f0e38ce14..c136b6b38 100644 --- a/types/notebookModels.ts +++ b/types/notebookModels.ts @@ -2,7 +2,7 @@ * Types for the agent model catalog (`GET /api/research_ai/models/`) and the * per-turn generation controls a selected model accepts. * - * Wire shapes stay snake_case-free but verbatim, like `types/notebookChat.ts`. + * Wire shapes stay snake_case and verbatim, like `types/notebookChat.ts`. * * The catalog says *what* each model accepts (its `capabilities`); the rules * below say which *combinations* the backend will take. They mirror @@ -58,9 +58,18 @@ export interface AgentModel { readonly description: string; readonly provider: string; readonly capabilities: AgentModelCapabilities; + readonly allowed: boolean; + readonly multiplier: string | null; +} + +export interface CreditPricing { + readonly multiplier_base_model: string; + readonly multiplier_basis: string; + readonly multiplier_is_estimate: boolean; } export interface AgentModelCatalog { + readonly credit_pricing?: CreditPricing; /** The ref that runs when a request names no model. */ readonly default: string; /** Server-ordered: strongest first within each family. */ @@ -89,9 +98,12 @@ export interface GenerationRequest extends GenerationOptions { /** Raw catalog response — `capabilities` arrives as open-ended strings. */ export interface AgentModelCatalogResponse { - default?: string; + default?: string | null; + credit_pricing?: CreditPricing; models?: Array<{ ref?: string; + allowed?: boolean; + multiplier?: string | null; label?: string; description?: string; provider?: string; @@ -134,6 +146,8 @@ export function toAgentModelCatalog(response: AgentModelCatalogResponse): AgentM if (!model?.ref) continue; models.push({ ref: model.ref, + allowed: model.allowed === true, + multiplier: model.multiplier ?? null, label: model.label?.trim() || modelIdOf(model.ref), description: model.description ?? '', provider: model.provider || providerOf(model.ref), @@ -144,7 +158,7 @@ export function toAgentModelCatalog(response: AgentModelCatalogResponse): AgentM }, }); } - return { default: response.default ?? '', models }; + return { default: response.default ?? '', models, credit_pricing: response.credit_pricing }; } export function findModel(models: AgentModel[], ref: string | null): AgentModel | null { @@ -160,6 +174,8 @@ export function findModel(models: AgentModel[], ref: string | null): AgentModel export function unknownModel(ref: string): AgentModel { return { ref, + allowed: false, + multiplier: null, label: modelIdOf(ref) || ref, description: '', provider: providerOf(ref), @@ -285,3 +301,23 @@ export function summarizeGenerationOptions(options: GenerationOptions): string[] if (options.temperature != null) parts.push(`Temp ${formatTemperature(options.temperature)}`); return parts; } + +/** Comparisons, never a per-message charge. Null pricing is not free. */ +export function formatModelMultiplier(value: string | null): string { + if (value === null || !Number.isFinite(Number(value)) || Number(value) <= 0) { + return 'Pricing unavailable'; + } + const multiplier = Number(value); + return multiplier < 0.01 + ? '<0.01×' + : `${new Intl.NumberFormat(undefined, { maximumFractionDigits: 2 }).format(multiplier)}×`; +} + +export function modelMultiplierExplanation(catalog: AgentModelCatalog | null): string { + const pricing = catalog?.credit_pricing; + const baseline = pricing + ? (findModel(catalog?.models ?? [], pricing.multiplier_base_model)?.label ?? + modelIdOf(pricing.multiplier_base_model)) + : null; + return `Estimated credit usage${baseline ? ` relative to ${baseline}` : ''}. Actual usage varies with input and output length, caching, and searches.`; +} diff --git a/types/researchAI.ts b/types/researchAI.ts new file mode 100644 index 000000000..481f90a8e --- /dev/null +++ b/types/researchAI.ts @@ -0,0 +1,54 @@ +/** User-wide Research AI allowances. Decimal strings are never dollar amounts. */ +export interface ResearchAIBudget { + tier: 'default' | 'invited' | 'privileged' | 'blocked'; + credits: { + daily_limit: string | null; + used: string; + remaining: string | null; + }; + turns_used: number; + turn_cap: number | null; + resets_at: string; +} + +export function isResearchAIBudget(value: unknown): value is ResearchAIBudget { + if (!value || typeof value !== 'object') return false; + const budget = value as ResearchAIBudget; + const decimal = (v: unknown) => typeof v === 'string' && /^-?\d+(\.\d+)?$/.test(v); + return ( + ['default', 'invited', 'privileged', 'blocked'].includes(budget.tier) && + budget.credits != null && + (budget.credits.daily_limit === null || decimal(budget.credits.daily_limit)) && + decimal(budget.credits.used) && + (budget.credits.remaining === null || decimal(budget.credits.remaining)) && + Number.isInteger(budget.turns_used) && + (budget.turn_cap === null || Number.isInteger(budget.turn_cap)) && + typeof budget.resets_at === 'string' && + Number.isFinite(Date.parse(budget.resets_at)) + ); +} + +export function isBudgetExhausted(budget: ResearchAIBudget | null): boolean { + if (!budget) return false; + return ( + (budget.credits.daily_limit !== null && + budget.credits.remaining !== null && + Number(budget.credits.remaining) <= 0) || + (budget.turn_cap !== null && budget.turns_used >= budget.turn_cap) + ); +} + +export function canSelectAIModel(tier: ResearchAIBudget['tier'] | undefined): boolean { + return tier === 'invited' || tier === 'privileged'; +} + +/** Keep the API's fractional precision, including balances smaller than 0.01. */ +export function formatCredits(value: string): string { + return value.includes('.') ? value.replace(/0+$/, '').replace(/\.$/, '') : value; +} + +export function formatBudgetReset(resetsAt: string): string { + return new Intl.DateTimeFormat(undefined, { hour: 'numeric', minute: '2-digit' }).format( + new Date(resetsAt) + ); +} From a4858916b25c48ab18dfe943df4b9700310423c9 Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Fri, 4 Sep 2026 11:50:20 -0400 Subject: [PATCH 2/6] Format AI credits with commas and two decimal places --- tests/notebook-ai.test.cjs | 6 ++++-- types/researchAI.ts | 7 +++++-- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/notebook-ai.test.cjs b/tests/notebook-ai.test.cjs index 2fd308a92..3a923f44e 100644 --- a/tests/notebook-ai.test.cjs +++ b/tests/notebook-ai.test.cjs @@ -118,7 +118,9 @@ test('tiers and catalog permissions are authoritative; fractions are not rounded assert.equal(models.formatModelMultiplier('3.75'), '3.75×'); assert.equal(models.formatModelMultiplier('0.001'), '<0.01×'); assert.equal(models.formatModelMultiplier(null), 'Pricing unavailable'); - assert.equal(budgetTypes.formatCredits('0.00010'), '0.0001'); + assert.equal(budgetTypes.formatCredits('0.00010'), '0.00'); + assert.equal(budgetTypes.formatCredits('12345.6'), '12,345.60'); + assert.equal(budgetTypes.formatCredits('250'), '250.00'); assert.match(models.modelMultiplierExplanation(catalog), /relative to Baseline/); assert.deepEqual( models.normalizeGenerationOptions(catalog.models[0], { @@ -141,7 +143,7 @@ test('meter shows fractional credits, daily cap exhaustion, local reset and unli }) ); assert.match(render(budget()), /248.35 credits remaining/); - assert.match(render(budget()), /250 daily credits/); + assert.match(render(budget()), /250\.00 daily credits/); assert.match(render(budget()), /Resets at/); const capped = render(budget({ turns_used: 10 })); assert.match(capped, /Daily AI usage limit reached/); diff --git a/types/researchAI.ts b/types/researchAI.ts index 481f90a8e..70e22e9f5 100644 --- a/types/researchAI.ts +++ b/types/researchAI.ts @@ -42,9 +42,12 @@ export function canSelectAIModel(tier: ResearchAIBudget['tier'] | undefined): bo return tier === 'invited' || tier === 'privileged'; } -/** Keep the API's fractional precision, including balances smaller than 0.01. */ +/** Display credits with comma grouping and exactly two decimal places. */ export function formatCredits(value: string): string { - return value.includes('.') ? value.replace(/0+$/, '').replace(/\.$/, '') : value; + return Number(value).toLocaleString('en-US', { + minimumFractionDigits: 2, + maximumFractionDigits: 2, + }); } export function formatBudgetReset(resetsAt: string): string { From b2cca5a7970cc4754bc337486716a1bbb14fcf30 Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Thu, 10 Sep 2026 11:18:35 -0400 Subject: [PATCH 3/6] Keep model catalog authoritative across refresh races --- .../Notebook/AgentChat/AgentChatPanel.tsx | 46 +++++++++++++++---- hooks/useNotebookChat.ts | 2 +- store/researchAI.ts | 18 ++++++-- tests/notebook-ai.test.cjs | 44 ++++++++++++++++-- 4 files changed, 91 insertions(+), 19 deletions(-) diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index 59dbf0d54..9730f6459 100644 --- a/components/Notebook/AgentChat/AgentChatPanel.tsx +++ b/components/Notebook/AgentChat/AgentChatPanel.tsx @@ -247,9 +247,9 @@ export function AgentChatPanel({ const { editor, currentNote } = useNotebookContext(); // This panel stays mounted even when closed: load allowances on notebook open. const researchAI = useResearchAI(true); - const canSelectModel = + const hasModelSelection = canSelectAIModel(researchAI.budget?.tier) && researchAI.budgetStatus === 'ok'; - const budgetSendDisabled = researchAI.budgetStatus !== 'ok' || researchAI.isSubmissionBlocked(); + const canSelectModel = hasModelSelection && researchAI.catalog !== null; // Decide which writing preset the empty chat screen offers, and what it // calls the document: the notebook holds RFPs as well as proposals. const noteIsEmpty = useEditorIsEmpty(editor); @@ -303,6 +303,12 @@ export function AgentChatPanel({ effortPinned: chatState.latestExecution != null, pinnedEffort: chatState.latestExecution?.effort ?? null, }); + // A selectable tier must never submit its first turn without an authoritative + // model. A cached catalog remains usable through a transient refresh failure. + const budgetSendDisabled = + researchAI.budgetStatus !== 'ok' || + researchAI.isSubmissionBlocked() || + (hasModelSelection && modelSelection.model === null); // ---- drafts (per chat, surviving switches and failed sends) ---- const draftsRef = useRef(new Map()); @@ -1258,14 +1264,34 @@ export function AgentChatPanel({ sendDisabled={budgetSendDisabled} notice={notice} footer={ - { - void researchAI.refreshBudget(true); - }} - /> + <> + { + void researchAI.refreshBudget(true); + }} + /> + {hasModelSelection && researchAI.catalog === null && ( +

+ {researchAI.catalogStatus === 'loading' + ? 'Loading available AI models…' + : 'Couldn’t load available AI models.'} + {researchAI.catalogStatus === 'unavailable' && ( + + )} +

+ )} + } toolbar={ canSelectModel && ( diff --git a/hooks/useNotebookChat.ts b/hooks/useNotebookChat.ts index 51a08cc38..77a3ccc96 100644 --- a/hooks/useNotebookChat.ts +++ b/hooks/useNotebookChat.ts @@ -499,7 +499,7 @@ export function useNotebookChat({ const outcome = sendFailureOutcome(err); if (outcome.reason === 'usage_limit') recordLimit(chatErrorBody(err)); else void refreshBudget(true); - if (outcome.reason === 'model_not_allowed') void refreshCatalog(); + if (outcome.reason === 'model_not_allowed') void refreshCatalog(true); // The outcome is still reported either way, but a continuation for a // chat that is no longer selected must not mutate the current one. if (epoch === epochRef.current) { diff --git a/store/researchAI.ts b/store/researchAI.ts index d8d215b47..141d1b257 100644 --- a/store/researchAI.ts +++ b/store/researchAI.ts @@ -28,6 +28,7 @@ export function createResearchAIStore(loaders: { let budgetFlight: Promise | null = null; let catalogFlight: Promise | null = null; let budgetQueued = false; + let catalogQueued = false; let budgetRevision = 0; let lastBudgetFetch = 0; @@ -75,14 +76,25 @@ export function createResearchAIStore(loaders: { return budgetFlight; }; - const refreshCatalog = (): Promise => { - if (catalogFlight) return catalogFlight; + const refreshCatalog = (force = false): Promise => { + if (catalogFlight) { + // A model rejection can arrive after the active GET was snapshotted. + // Queue one authoritative follow-up rather than accepting stale access. + if (force) catalogQueued = true; + return catalogFlight; + } catalogFlight = loaders .catalog() .then((catalog) => update({ catalog, catalogStatus: 'ok' })) - .catch(() => update({ catalog: null, catalogStatus: 'unavailable' })) + // Keep the last successful catalog through transient failures so an + // already-selected model cannot silently turn into the server default. + .catch(() => update({ catalogStatus: 'unavailable' })) .finally(() => { catalogFlight = null; + if (catalogQueued) { + catalogQueued = false; + void refreshCatalog(true); + } }); return catalogFlight; }; diff --git a/tests/notebook-ai.test.cjs b/tests/notebook-ai.test.cjs index 3a923f44e..2397b4f8d 100644 --- a/tests/notebook-ai.test.cjs +++ b/tests/notebook-ai.test.cjs @@ -258,27 +258,61 @@ test('opening an old failed turn does not block today, and cancellation refresh assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); }); -test('budget failure keeps the recorded balance and catalog refresh removes withdrawn choices', async () => { - let fail = false; +test('transient failures retain the recorded budget and last successful catalog', async () => { + let failBudget = false; + let failCatalog = false; let available = catalog; const store = createResearchAIStore({ budget: async () => { - if (fail) throw Error('offline'); + if (failBudget) throw Error('offline'); return budget(); }, - catalog: async () => available, + catalog: async () => { + if (failCatalog) throw Error('offline'); + return available; + }, }); await store.refreshBudget(); - fail = true; + failBudget = true; await store.refreshBudget(true); assert.equal(store.getSnapshot().budgetStatus, 'unavailable'); assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); await store.refreshCatalog(); + failCatalog = true; + await store.refreshCatalog(); + assert.equal(store.getSnapshot().catalogStatus, 'unavailable'); + assert.equal(store.getSnapshot().catalog.models[0].ref, catalog.models[0].ref); + failCatalog = false; available = { ...catalog, models: [] }; await store.refreshCatalog(); assert.equal(store.getSnapshot().catalog.models.length, 0); }); +test('an authoritative catalog refresh queues behind an active fetch', async () => { + let resolveFirst; + let calls = 0; + const withdrawn = { ...catalog, models: catalog.models.slice(1) }; + const store = createResearchAIStore({ + budget: async () => budget(), + catalog: () => { + calls += 1; + if (calls === 1) { + return new Promise((resolve) => { + resolveFirst = resolve; + }); + } + return Promise.resolve(withdrawn); + }, + }); + const pending = store.refreshCatalog(); + store.refreshCatalog(true); + resolveFirst(catalog); + await pending; + await flush(); + assert.equal(calls, 2); + assert.deepEqual(store.getSnapshot().catalog.models, withdrawn.models); +}); + test('errors preserve structured codes, top-level budgets and ordinary field validation', () => { const limit = new ApiError('Request failed', 429, { ...budget(), code: 'usage_limit_exceeded' }); assert.equal(service.chatErrorCode(limit), 'usage_limit_exceeded'); From 898cf04025988395b2baa9a4ff9281b8653cec56 Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Thu, 10 Sep 2026 12:06:01 -0400 Subject: [PATCH 4/6] Preserve new chat settings on creation and remove PR tests --- .../Notebook/AgentChat/AgentChatPanel.tsx | 4 + hooks/useAgentModelSelection.ts | 14 +- tests/notebook-ai.test.cjs | 401 ------------------ 3 files changed, 17 insertions(+), 402 deletions(-) delete mode 100644 tests/notebook-ai.test.cjs diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index 9730f6459..83be02501 100644 --- a/components/Notebook/AgentChat/AgentChatPanel.tsx +++ b/components/Notebook/AgentChat/AgentChatPanel.tsx @@ -450,6 +450,8 @@ export function AgentChatPanel({ return; } draftsRef.current.delete('new'); + // A rejected first attempt must retry with the same model and settings. + modelSelection.adoptConversation(`${noteId}:${created.conversation_id}`, generation); setInitialChat(created); setSelectedChatId(created.conversation_id); setQueuedMessage({ text, generation }); @@ -474,6 +476,8 @@ export function AgentChatPanel({ list, chatState, modelSelection.request, + modelSelection.adoptConversation, + noteId, updateDraft, isCurrentTarget, budgetSendDisabled, diff --git a/hooks/useAgentModelSelection.ts b/hooks/useAgentModelSelection.ts index 0a635e5fc..8aabb10e7 100644 --- a/hooks/useAgentModelSelection.ts +++ b/hooks/useAgentModelSelection.ts @@ -48,6 +48,8 @@ export interface AgentModelSelection { readonly selectModel: (ref: string) => void; /** Patch: pass a field as `undefined` to hand it back to the server. */ readonly setOptions: (options: GenerationOptions) => void; + /** Carry the first send's choices onto the conversation the server just created. */ + readonly adoptConversation: (key: string, generation: GenerationRequest) => void; /** Generation fields for a send, ready to spread into the request body. */ readonly request: GenerationRequest; } @@ -73,7 +75,11 @@ export function useAgentModelSelection({ }); const preference = choice.key === conversationKey ? choice.preference : {}; useEffect(() => { - setChoice({ key: conversationKey, preference: {} }); + // Creation already transferred the captured choices to the assigned ID. + // Only an ordinary chat switch should reset this hook's selection. + setChoice((current) => + current.key === conversationKey ? current : { key: conversationKey, preference: {} } + ); }, [conversationKey]); const models = useMemo( () => catalog?.models.filter((model) => model.allowed) ?? NO_MODELS, @@ -133,6 +139,11 @@ export function useAgentModelSelection({ }; }, [canSelect, model, options, locked, effortPinned]); + const adoptConversation = useCallback((key: string, generation: GenerationRequest) => { + const { model: ref, ...options } = generation; + setChoice({ key, preference: { ref, ...options } }); + }, []); + return { status, multiplierExplanation: modelMultiplierExplanation(catalog), @@ -143,6 +154,7 @@ export function useAgentModelSelection({ options, selectModel, setOptions, + adoptConversation, request, }; } diff --git a/tests/notebook-ai.test.cjs b/tests/notebook-ai.test.cjs deleted file mode 100644 index 2397b4f8d..000000000 --- a/tests/notebook-ai.test.cjs +++ /dev/null @@ -1,401 +0,0 @@ -/* eslint-disable @typescript-eslint/no-require-imports -- Standalone CommonJS Node test runner. */ -// Run with: node --test tests/notebook-ai.test.cjs -// Uses the repository's TypeScript compiler and Node's test runner; no extra test dependencies. -const { test } = require('node:test'); -const assert = require('node:assert/strict'); -const { readFileSync, existsSync } = require('node:fs'); -const path = require('node:path'); -const ts = require('typescript'); -const root = path.resolve(__dirname, '..'); -const cache = new Map(); -const api = {}; -function load(relative) { - const filename = path.resolve(root, relative); - if (cache.has(filename)) return cache.get(filename).exports; - const compiledModule = { exports: {} }; - cache.set(filename, compiledModule); - const source = ts.transpileModule(readFileSync(filename, 'utf8'), { - compilerOptions: { - module: ts.ModuleKind.CommonJS, - target: ts.ScriptTarget.ES2022, - jsx: ts.JsxEmit.ReactJSX, - }, - }).outputText; - const localRequire = (specifier) => { - if ( - specifier === '@/services/client' || - (specifier === './client' && relative.startsWith('services/')) - ) - return { ApiClient: api }; - if (specifier === '@/hooks/useAgentModels') - return { useAgentModels: () => ({ status: 'ok', catalog }) }; - if (!specifier.startsWith('@/') && !specifier.startsWith('.')) return require(specifier); - const base = specifier.startsWith('@/') - ? path.join(root, specifier.slice(2)) - : path.resolve(path.dirname(filename), specifier); - const target = [base + '.ts', base + '.tsx', path.join(base, 'index.ts')].find(existsSync); - return load(path.relative(root, target)); - }; - new Function('require', 'module', 'exports', source)( - localRequire, - compiledModule, - compiledModule.exports - ); - return compiledModule.exports; -} -const budgetTypes = load('types/researchAI.ts'); -const models = load('types/notebookModels.ts'); -const { createResearchAIStore } = load('store/researchAI.ts'); -const { ApiError } = load('services/types/api.ts'); -const service = load('services/notebookChat.service.ts'); -const { renderToStaticMarkup } = require('react-dom/server'); -const { createElement } = require('react'); -const { CreditMeter } = load('components/Notebook/AgentChat/CreditMeter.tsx'); -const { ChatComposer } = load('components/Notebook/AgentChat/ChatComposer.tsx'); -const { ModelControls } = load('components/Notebook/AgentChat/ModelControls.tsx'); -const tomorrow = new Date(); -tomorrow.setUTCHours(24, 0, 0, 0); -const budget = (overrides = {}) => ({ - tier: 'default', - credits: { daily_limit: '250', used: '1.65', remaining: '248.35' }, - turns_used: 2, - turn_cap: 10, - resets_at: tomorrow.toISOString(), - ...overrides, -}); -const catalog = models.toAgentModelCatalog({ - default: 'openrouter:test', - credit_pricing: { - multiplier_base_model: 'openrouter:base', - multiplier_basis: 'equal_input_output_tokens', - multiplier_is_estimate: true, - }, - models: [ - { - ref: 'openrouter:test', - label: 'Flash', - allowed: true, - multiplier: '0.03', - capabilities: { effort: ['low', 'high'], thinking: [], temperature: false }, - }, - { ref: 'openrouter:base', label: 'Baseline', allowed: true, multiplier: '1' }, - { ref: 'openrouter:unpriced', allowed: false, multiplier: null }, - ], -}); -const storeWith = (getBudget = async () => budget()) => - createResearchAIStore({ budget: getBudget, catalog: async () => catalog }); -const flush = () => new Promise((resolve) => setImmediate(resolve)); - -test('credit exhaustion and provider-call cap are independent; null means unlimited', () => { - assert.equal(budgetTypes.isBudgetExhausted(budget()), false); - assert.equal( - budgetTypes.isBudgetExhausted( - budget({ credits: { daily_limit: '250', used: '250', remaining: '0.00' } }) - ), - true - ); - assert.equal(budgetTypes.isBudgetExhausted(budget({ turns_used: 10 })), true); - assert.equal( - budgetTypes.isBudgetExhausted( - budget({ - credits: { daily_limit: null, used: '999', remaining: null }, - turn_cap: null, - turns_used: 999, - }) - ), - false - ); - assert.equal(budgetTypes.isResearchAIBudget({ tier: 'default', remaining: '100' }), false); -}); - -test('tiers and catalog permissions are authoritative; fractions are not rounded to zero', () => { - assert.equal(budgetTypes.canSelectAIModel('default'), false); - assert.equal(budgetTypes.canSelectAIModel('blocked'), false); - assert.equal(budgetTypes.canSelectAIModel('invited'), true); - assert.equal(budgetTypes.canSelectAIModel('privileged'), true); - assert.equal(catalog.models[2].allowed, false); - assert.equal(models.formatModelMultiplier('0.03'), '0.03×'); - assert.equal(models.formatModelMultiplier('3.75'), '3.75×'); - assert.equal(models.formatModelMultiplier('0.001'), '<0.01×'); - assert.equal(models.formatModelMultiplier(null), 'Pricing unavailable'); - assert.equal(budgetTypes.formatCredits('0.00010'), '0.00'); - assert.equal(budgetTypes.formatCredits('12345.6'), '12,345.60'); - assert.equal(budgetTypes.formatCredits('250'), '250.00'); - assert.match(models.modelMultiplierExplanation(catalog), /relative to Baseline/); - assert.deepEqual( - models.normalizeGenerationOptions(catalog.models[0], { - effort: 'high', - thinking: 'adaptive', - temperature: 1, - }), - { effort: 'high' } - ); -}); - -test('meter shows fractional credits, daily cap exhaustion, local reset and unlimited balances', () => { - const render = (b) => - renderToStaticMarkup( - createElement(CreditMeter, { - budget: b, - budgetStatus: 'ok', - limitResetAt: null, - onRefresh() {}, - }) - ); - assert.match(render(budget()), /248.35 credits remaining/); - assert.match(render(budget()), /250\.00 daily credits/); - assert.match(render(budget()), /Resets at/); - const capped = render(budget({ turns_used: 10 })); - assert.match(capped, /Daily AI usage limit reached/); - assert.doesNotMatch(capped, /Out of credits|messages remaining/); - assert.match( - render(budget({ credits: { daily_limit: null, remaining: null, used: '1' } })), - /Unlimited credits/ - ); - const previousTZ = process.env.TZ; - process.env.TZ = 'America/New_York'; - assert.match(budgetTypes.formatBudgetReset('2026-09-05T00:00:00Z'), /8:00/); - if (previousTZ === undefined) delete process.env.TZ; - else process.env.TZ = previousTZ; -}); - -test('exhaustion disables Send but retains editable draft and Stop', () => { - const props = { - value: 'Keep my unsent question', - onChange() {}, - onSend() {}, - onStop() {}, - busy: false, - canStop: false, - disabled: false, - sendDisabled: true, - notice: null, - textareaRef: { current: null }, - }; - const html = renderToStaticMarkup(createElement(ChatComposer, props)); - assert.match(html, /Keep my unsent question/); - assert.doesNotMatch(html, /]* disabled=""/); - assert.match(html, /]*disabled=""[^>]*title="Send message"/); - const running = renderToStaticMarkup( - createElement(ChatComposer, { ...props, busy: true, canStop: true }) - ); - assert.match(running, /title="Stop the assistant"/); - assert.doesNotMatch(running, /]* disabled=""/); -}); - -test('pinned model control is disabled and explains how to switch', () => { - const html = renderToStaticMarkup( - createElement(ModelControls, { - models: catalog.models, - model: catalog.models[0], - pinned: true, - options: {}, - onSelectModel() {}, - onChangeOptions() {}, - disabled: false, - multiplierExplanation: models.modelMultiplierExplanation(catalog), - }) - ); - assert.match(html, /]*disabled=""/); - assert.match(html, /Start a new chat to switch models/); -}); - -test('shared subscribers see one fetch; progress refreshes throttle and sessions stay isolated', async () => { - let calls = 0; - const store = storeWith(async () => { - calls++; - return budget(); - }); - const snapshots = []; - const unsubscribe = store.subscribe(() => snapshots.push(store.getSnapshot())); - await Promise.all([store.refreshBudget(), store.refreshBudget()]); - await store.refreshBudget(); - assert.equal(calls, 1); - assert.equal(snapshots.at(-1).budget.credits.remaining, '248.35'); - assert.equal(storeWith().getSnapshot().budget, null); - unsubscribe(); -}); - -test('429 budget wins over an older GET, with a follow-up refresh after settlement', async () => { - let resolveOld; - let calls = 0; - const store = storeWith(() => - ++calls === 1 - ? new Promise((resolve) => { - resolveOld = resolve; - }) - : Promise.resolve(budget()) - ); - const pending = store.refreshBudget(true); - store.recordLimit(budget({ turns_used: 10 })); - assert.equal(store.getSnapshot().budget.turns_used, 10); - resolveOld(budget()); - await pending; - await flush(); - assert.equal(calls, 2); - assert.equal(store.isSubmissionBlocked(), true); -}); - -test('a post-202 limit stays blocked with credits remaining and recovers after reset', async () => { - let current = budget(); - const store = storeWith(async () => current); - await store.refreshBudget(); - store.recordLimit(); - await flush(); - assert.equal(store.isSubmissionBlocked(), true); - const nextReset = new Date(Date.parse(current.resets_at) + 86400000).toISOString(); - current = budget({ resets_at: nextReset, turns_used: 0 }); - await store.refreshBudget(true); - assert.equal(store.isSubmissionBlocked(), false); -}); - -test('opening an old failed turn does not block today, and cancellation refresh never refunds', async () => { - const store = storeWith(); - await store.refreshBudget(); - store.recordLimit(undefined, '2020-01-01T12:00:00Z'); - assert.equal(store.isSubmissionBlocked(), false); - await store.refreshBudget(true); - assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); -}); - -test('transient failures retain the recorded budget and last successful catalog', async () => { - let failBudget = false; - let failCatalog = false; - let available = catalog; - const store = createResearchAIStore({ - budget: async () => { - if (failBudget) throw Error('offline'); - return budget(); - }, - catalog: async () => { - if (failCatalog) throw Error('offline'); - return available; - }, - }); - await store.refreshBudget(); - failBudget = true; - await store.refreshBudget(true); - assert.equal(store.getSnapshot().budgetStatus, 'unavailable'); - assert.equal(store.getSnapshot().budget.credits.remaining, '248.35'); - await store.refreshCatalog(); - failCatalog = true; - await store.refreshCatalog(); - assert.equal(store.getSnapshot().catalogStatus, 'unavailable'); - assert.equal(store.getSnapshot().catalog.models[0].ref, catalog.models[0].ref); - failCatalog = false; - available = { ...catalog, models: [] }; - await store.refreshCatalog(); - assert.equal(store.getSnapshot().catalog.models.length, 0); -}); - -test('an authoritative catalog refresh queues behind an active fetch', async () => { - let resolveFirst; - let calls = 0; - const withdrawn = { ...catalog, models: catalog.models.slice(1) }; - const store = createResearchAIStore({ - budget: async () => budget(), - catalog: () => { - calls += 1; - if (calls === 1) { - return new Promise((resolve) => { - resolveFirst = resolve; - }); - } - return Promise.resolve(withdrawn); - }, - }); - const pending = store.refreshCatalog(); - store.refreshCatalog(true); - resolveFirst(catalog); - await pending; - await flush(); - assert.equal(calls, 2); - assert.deepEqual(store.getSnapshot().catalog.models, withdrawn.models); -}); - -test('errors preserve structured codes, top-level budgets and ordinary field validation', () => { - const limit = new ApiError('Request failed', 429, { ...budget(), code: 'usage_limit_exceeded' }); - assert.equal(service.chatErrorCode(limit), 'usage_limit_exceeded'); - assert.equal(service.chatErrorBody(limit).credits.remaining, '248.35'); - assert.equal( - service.chatErrorDetail(new ApiError('Request failed', 400, { message: ['Too long.'] })), - 'message: Too long.' - ); - assert.equal( - service.chatErrorDetail(new ApiError('Request failed', 400, { detail: 'Model unavailable.' })), - 'Model unavailable.' - ); - assert.equal( - service.chatErrorCode(new ApiError('Request failed', 409, { code: 'usage_work_in_progress' })), - 'usage_work_in_progress' - ); -}); - -test('default-tier sends only message; selected model is omitted once the conversation is locked', async () => { - const { useAgentModelSelection } = load('hooks/useAgentModelSelection.ts'); - let selection; - function Probe(props) { - selection = useAgentModelSelection({ - enabled: false, - conversationKey: 'new', - pinnedRef: null, - locked: false, - ...props, - }); - return null; - } - let sent; - api.post = async (url, body) => { - sent = { url, body }; - return { execution_id: 42 }; - }; - renderToStaticMarkup(createElement(Probe, { canSelect: false })); - await service.NotebookChatService.sendMessage( - 1, - 2, - 'Summarize this notebook.', - selection.request - ); - assert.deepEqual(sent.body, { message: 'Summarize this notebook.' }); - renderToStaticMarkup(createElement(Probe, { canSelect: true })); - assert.equal(selection.model.ref, catalog.default); - assert.equal(selection.request.model, catalog.default); - renderToStaticMarkup( - createElement(Probe, { canSelect: true, locked: true, pinnedRef: 'openrouter:base' }) - ); - assert.equal(selection.model.ref, 'openrouter:base'); - assert.equal(selection.request.model, undefined); -}); - -test('immediate resubmission after cancellation handles account-wide 409 without classifying it as exhaustion', async () => { - api.post = async (url) => { - if (url.endsWith('/cancel/')) return { cancelled: true, execution_id: 42 }; - throw new ApiError('Request failed', 409, { code: 'usage_work_in_progress' }); - }; - await service.NotebookChatService.cancelTurn(1, 2); - await assert.rejects( - service.NotebookChatService.sendMessage(1, 2, 'Preserved draft'), - (error) => { - const outcome = service.sendFailureOutcome(error); - assert.equal(outcome.reason, 'account_busy'); - assert.equal(outcome.detail, 'Another AI request is still running.'); - return true; - } - ); - assert.equal(service.sendFailureOutcome(new ApiError('busy', 409)).reason, 'busy'); - assert.equal(service.sendFailureOutcome(new ApiError('forbidden', 403)).reason, 'unauthorized'); - assert.equal( - service.sendFailureOutcome( - new ApiError('invalid', 400, { code: 'model_not_allowed', detail: 'Choose another model.' }) - ).reason, - 'model_not_allowed' - ); - assert.equal( - service.sendFailureOutcome(new ApiError('invalid', 400, { message: ['Too long.'] })).reason, - 'invalid' - ); - assert.equal( - service.sendFailureOutcome(new ApiError('limit', 429, { code: 'usage_limit_exceeded' })).reason, - 'usage_limit' - ); -}); From e182e06550a4a0afd2173be1533b5c04ab55f10b Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Thu, 10 Sep 2026 12:52:41 -0400 Subject: [PATCH 5/6] Preserve notebook access denials and address SonarCloud findings --- .../Notebook/AgentChat/AgentChatPanel.tsx | 43 ++++++++------ components/Notebook/AgentChat/CreditMeter.tsx | 24 ++++---- hooks/useAgentModelSelection.ts | 2 +- hooks/useNotebookChat.ts | 57 +++++++++---------- types/notebookModels.ts | 3 +- 5 files changed, 68 insertions(+), 61 deletions(-) diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index 83be02501..d853ca999 100644 --- a/components/Notebook/AgentChat/AgentChatPanel.tsx +++ b/components/Notebook/AgentChat/AgentChatPanel.tsx @@ -370,6 +370,22 @@ export function AgentChatPanel({ }, [noteId]); // ---- server-side access gate ---- + const [deniedNoteId, setDeniedNoteId] = useState(null); + const accessNoteRef = useRef(noteId); + useEffect(() => { + if (accessNoteRef.current !== noteId) { + accessNoteRef.current = noteId; + setDeniedNoteId(null); + // The chat hooks reset after a note switch; their current access values + // can still belong to the previous note on this render. + return; + } + if (list.access === 'hidden' || chatState.access === 'unauthorized') { + setDeniedNoteId(noteId); + } + }, [noteId, list.access, chatState.access]); + const accessDenied = deniedNoteId === noteId; + useEffect(() => { // Leave a visible restriction until the user closes the panel. A blocked // account keeps the entry point so its unavailable state remains reachable. @@ -377,18 +393,11 @@ export function AgentChatPanel({ !open && researchAI.budgetStatus !== 'loading' && researchAI.budget?.tier !== 'blocked' && - (list.access === 'hidden' || chatState.access === 'unauthorized') + accessDenied ) { onUnavailable(); } - }, [ - open, - list.access, - chatState.access, - onUnavailable, - researchAI.budgetStatus, - researchAI.budget?.tier, - ]); + }, [open, accessDenied, onUnavailable, researchAI.budgetStatus, researchAI.budget?.tier]); // ---- keep the listing fresh as the open chat evolves ---- // Derived titles land after the first turn, previews/spinners change as @@ -987,8 +996,8 @@ export function AgentChatPanel({ const turnActive = chatState.latestExecution != null && isActiveExecutionStatus(chatState.latestExecution.status); const canStop = turnActive || chatState.pendingSend?.executionId != null; - const composerDisabled = - selectedChatId == null ? list.access !== 'ok' : chatState.access !== 'ok'; + const chatAccessible = selectedChatId == null ? list.access === 'ok' : chatState.access === 'ok'; + const composerDisabled = accessDenied || !chatAccessible; const emptyState = ( { if ( researchAI.budget?.tier === 'blocked' || + accessDenied || list.access === 'hidden' || chatState.access === 'unauthorized' ) { return ( -
+ You do not have access to the research assistant for this notebook. -
+ ); } if (selectedChatId == null) { @@ -1278,7 +1285,7 @@ export function AgentChatPanel({ }} /> {hasModelSelection && researchAI.catalog === null && ( -

+ {researchAI.catalogStatus === 'loading' ? 'Loading available AI models…' : 'Couldn’t load available AI models.'} @@ -1293,7 +1300,7 @@ export function AgentChatPanel({ Retry )} -

+ )} } diff --git a/components/Notebook/AgentChat/CreditMeter.tsx b/components/Notebook/AgentChat/CreditMeter.tsx index 964b020b5..3cf65fd1c 100644 --- a/components/Notebook/AgentChat/CreditMeter.tsx +++ b/components/Notebook/AgentChat/CreditMeter.tsx @@ -11,46 +11,48 @@ export function CreditMeter({ }: Pick & { onRefresh: () => void }) { if (budget?.tier === 'blocked') { return ( -

+ Research AI is unavailable for this account. -

+ ); } if (!budget) { return ( -

+ {budgetStatus === 'loading' ? 'Loading AI credits…' : 'Couldn’t load AI credits.'} {budgetStatus === 'unavailable' && ( )} -

+ ); } const exhausted = isBudgetExhausted(budget) || limitResetAt !== null; const { remaining, daily_limit: limit } = budget.credits; const reset = formatBudgetReset(budget.resets_at); + let balanceLabel = 'Credits unavailable'; + if (limit === null) { + balanceLabel = 'Unlimited credits'; + } else if (remaining !== null) { + balanceLabel = `${formatCredits(remaining)} credits remaining`; + } return (
- {limit === null - ? 'Unlimited credits' - : remaining === null - ? 'Credits unavailable' - : `${formatCredits(remaining)} credits remaining`} + {balanceLabel}
{exhausted && ( -

+ Daily AI usage limit reached. Available again at {reset}. -

+ )} {budgetStatus === 'unavailable' && (

diff --git a/hooks/useAgentModelSelection.ts b/hooks/useAgentModelSelection.ts index 8aabb10e7..ca1cf8613 100644 --- a/hooks/useAgentModelSelection.ts +++ b/hooks/useAgentModelSelection.ts @@ -131,7 +131,7 @@ export function useAgentModelSelection({ ); const request = useMemo(() => { - if (!canSelect || model == null || !model.allowed) return {}; + if (!canSelect || !model?.allowed) return {}; const { effort, ...perTurn } = options; return { ...(!locked && { model: model.ref }), diff --git a/hooks/useNotebookChat.ts b/hooks/useNotebookChat.ts index 77a3ccc96..ed4ec1bb5 100644 --- a/hooks/useNotebookChat.ts +++ b/hooks/useNotebookChat.ts @@ -476,6 +476,31 @@ export function useNotebookChat({ onReconnect: handleSocketReconnect, }); + const handleSendFailure = useCallback( + (err: unknown, epoch: number): SendOutcome => { + const outcome = sendFailureOutcome(err); + if (outcome.reason === 'usage_limit') recordLimit(chatErrorBody(err)); + else void refreshBudget(true); + if (outcome.reason === 'model_not_allowed') void refreshCatalog(true); + + // Refresh account allowances even after switching chats, but only update + // the transcript and access state for the chat that sent the message. + if (epoch !== epochRef.current) return outcome; + setPendingSend(null); + // Raced an active turn — refetch so the busy state renders truthfully. + if (outcome.reason === 'busy' || outcome.reason === 'account_busy') fetchChat('live'); + if (outcome.reason === 'not_found') setAccess('not_found'); + // Session expired or permission revoked mid-chat: mirror a failed GET + // so the access gate reacts instead of showing generic composer errors. + if (outcome.reason === 'unauthorized') { + setChat(null); + setAccess('unauthorized'); + } + return outcome; + }, + [fetchChat, refreshBudget, refreshCatalog, recordLimit] + ); + const send = useCallback( async (text: string, generation?: GenerationRequest): Promise => { if (noteId == null || chatId == null) return { ok: false, reason: 'error' }; @@ -496,38 +521,10 @@ export function useNotebookChat({ } return { ok: true }; } catch (err) { - const outcome = sendFailureOutcome(err); - if (outcome.reason === 'usage_limit') recordLimit(chatErrorBody(err)); - else void refreshBudget(true); - if (outcome.reason === 'model_not_allowed') void refreshCatalog(true); - // The outcome is still reported either way, but a continuation for a - // chat that is no longer selected must not mutate the current one. - if (epoch === epochRef.current) { - setPendingSend(null); - // Raced an active turn — refetch so the busy state renders truthfully. - if (outcome.reason === 'busy' || outcome.reason === 'account_busy') fetchChat('live'); - if (outcome.reason === 'not_found') setAccess('not_found'); - // Session expired or permission revoked mid-chat: mirror what a - // failed GET does so the access gate reacts instead of the composer - // showing generic errors forever. - if (outcome.reason === 'unauthorized') { - setChat(null); - setAccess('unauthorized'); - } - } - return outcome; + return handleSendFailure(err, epoch); } }, - [ - noteId, - chatId, - fetchChat, - refreshBudget, - refreshCatalog, - recordLimit, - getSnapshot, - isSubmissionBlocked, - ] + [noteId, chatId, fetchChat, handleSendFailure, getSnapshot, isSubmissionBlocked] ); const cancel = useCallback(async () => { diff --git a/types/notebookModels.ts b/types/notebookModels.ts index c136b6b38..a418280bb 100644 --- a/types/notebookModels.ts +++ b/types/notebookModels.ts @@ -319,5 +319,6 @@ export function modelMultiplierExplanation(catalog: AgentModelCatalog | null): s ? (findModel(catalog?.models ?? [], pricing.multiplier_base_model)?.label ?? modelIdOf(pricing.multiplier_base_model)) : null; - return `Estimated credit usage${baseline ? ` relative to ${baseline}` : ''}. Actual usage varies with input and output length, caching, and searches.`; + const comparison = baseline ? ` relative to ${baseline}` : ''; + return `Estimated credit usage${comparison}. Actual usage varies with input and output length, caching, and searches.`; } From 9e62cf305c9a22aab08b01ad8989470cf0de19b0 Mon Sep 17 00:00:00 2001 From: Taki Koutsomitis Date: Thu, 10 Sep 2026 13:10:33 -0400 Subject: [PATCH 6/6] Keep cached AI budgets usable after refresh failures --- components/Notebook/AgentChat/AgentChatPanel.tsx | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index d853ca999..d9e4f8301 100644 --- a/components/Notebook/AgentChat/AgentChatPanel.tsx +++ b/components/Notebook/AgentChat/AgentChatPanel.tsx @@ -247,8 +247,7 @@ export function AgentChatPanel({ const { editor, currentNote } = useNotebookContext(); // This panel stays mounted even when closed: load allowances on notebook open. const researchAI = useResearchAI(true); - const hasModelSelection = - canSelectAIModel(researchAI.budget?.tier) && researchAI.budgetStatus === 'ok'; + const hasModelSelection = canSelectAIModel(researchAI.budget?.tier); const canSelectModel = hasModelSelection && researchAI.catalog !== null; // Decide which writing preset the empty chat screen offers, and what it // calls the document: the notebook holds RFPs as well as proposals. @@ -304,9 +303,9 @@ export function AgentChatPanel({ pinnedEffort: chatState.latestExecution?.effort ?? null, }); // A selectable tier must never submit its first turn without an authoritative - // model. A cached catalog remains usable through a transient refresh failure. + // model. Cached budget and catalog data remain usable through refresh failures. const budgetSendDisabled = - researchAI.budgetStatus !== 'ok' || + researchAI.budget === null || researchAI.isSubmissionBlocked() || (hasModelSelection && modelSelection.model === null);