diff --git a/components/Notebook/AgentChat/AgentChatPanel.tsx b/components/Notebook/AgentChat/AgentChatPanel.tsx index c6927234e..d9e4f8301 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,10 @@ 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 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. const noteIsEmpty = useEditorIsEmpty(editor); @@ -277,13 +289,25 @@ 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, }); + // A selectable tier must never submit its first turn without an authoritative + // model. Cached budget and catalog data remain usable through refresh failures. + const budgetSendDisabled = + researchAI.budget === null || + researchAI.isSubmissionBlocked() || + (hasModelSelection && modelSelection.model === null); // ---- drafts (per chat, surviving switches and failed sends) ---- const draftsRef = useRef(new Map()); @@ -345,11 +369,34 @@ 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. + if ( + !open && + researchAI.budgetStatus !== 'loading' && + researchAI.budget?.tier !== 'blocked' && + accessDenied + ) { onUnavailable(); } - }, [list.access, chatState.access, onUnavailable]); + }, [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 @@ -386,7 +433,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 @@ -411,6 +458,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 }); @@ -435,8 +484,13 @@ export function AgentChatPanel({ list, chatState, modelSelection.request, + modelSelection.adoptConversation, + noteId, updateDraft, isCurrentTarget, + budgetSendDisabled, + creatingChat, + queuedMessage, ]); // Fire the queued first message once the freshly created chat is live. @@ -928,24 +982,21 @@ 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. 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) { if (list.access === 'loading') return ; if (list.access === 'error') { @@ -1208,18 +1271,52 @@ export function AgentChatPanel({ busy={composerBusy} canStop={canStop} disabled={composerDisabled} + sendDisabled={budgetSendDisabled} notice={notice} + footer={ + <> + { + 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/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..3cf65fd1c --- /dev/null +++ b/components/Notebook/AgentChat/CreditMeter.tsx @@ -0,0 +1,67 @@ +'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); + let balanceLabel = 'Credits unavailable'; + if (limit === null) { + balanceLabel = 'Unlimited credits'; + } else if (remaining !== null) { + balanceLabel = `${formatCredits(remaining)} credits remaining`; + } + return ( +

+
+ + {balanceLabel} + + +
+ {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..ca1cf8613 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; @@ -64,52 +48,49 @@ 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; } /** - * 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. + const [choice, setChoice] = useState<{ key: string; preference: StoredPreference }>({ + key: conversationKey, + preference: {}, + }); + const preference = choice.key === conversationKey ? choice.preference : {}; useEffect(() => { - setPreference(readPreference()); - setHydrated(true); - }, []); - - 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; + // 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, + [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,42 +106,55 @@ 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?.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]); + + const adoptConversation = useCallback((key: string, generation: GenerationRequest) => { + const { model: ref, ...options } = generation; + setChoice({ key, preference: { ref, ...options } }); + }, []); return { status, + multiplierExplanation: modelMultiplierExplanation(catalog), models, model, - pinned: pinnedRef != null, + pinned: locked, effortPinned, options, selectModel, setOptions, + adoptConversation, request, }; } 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..ed4ec1bb5 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(() => { @@ -479,39 +476,55 @@ 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' }; + 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'); } return { ok: true }; } catch (err) { - const outcome = sendFailureOutcome(err); - // 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 === '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] + [noteId, chatId, fetchChat, handleSendFailure, getSnapshot, isSubmissionBlocked] ); const cancel = useCallback(async () => { @@ -523,8 +536,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..141d1b257 --- /dev/null +++ b/store/researchAI.ts @@ -0,0 +1,137 @@ +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 catalogQueued = 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 = (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' })) + // 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; + }; + + 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/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..a418280bb 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,24 @@ 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; + const comparison = baseline ? ` relative to ${baseline}` : ''; + return `Estimated credit usage${comparison}. 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..70e22e9f5 --- /dev/null +++ b/types/researchAI.ts @@ -0,0 +1,57 @@ +/** 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'; +} + +/** Display credits with comma grouping and exactly two decimal places. */ +export function formatCredits(value: string): string { + return Number(value).toLocaleString('en-US', { + minimumFractionDigits: 2, + maximumFractionDigits: 2, + }); +} + +export function formatBudgetReset(resetsAt: string): string { + return new Intl.DateTimeFormat(undefined, { hour: 'numeric', minute: '2-digit' }).format( + new Date(resetsAt) + ); +}