1import { useAtomValue } from "@effect/atom-react";
2import { AuthSettingsWriteScope, type EnvironmentId, sessionGrantsScope } from "@t3tools/contracts";
3import { ChevronDownIcon, PlusIcon, RotateCcwIcon, XIcon } from "lucide-react";
4import * as Option from "effect/Option";
5import { AsyncResult, Atom } from "effect/reactivity";
6import { useRef, useState } from "react";
7
8import { cn } from "../../lib/utils";
9import { environmentPresentations } from "../../state/presentation";
10import { serverEnvironment } from "../../state/server";
11import { environmentSession, readEnvironmentScope } from "../../state/session";
12import { useAtomCommand } from "../../state/use-atom-command";
13import type { EnvironmentUsageStatus } from "../../state/usage";
14import { Button } from "../ui/button";
15import {
16 Dialog,
17 DialogDescription,
18 DialogFooter,
19 DialogHeader,
20 DialogPanel,
21 DialogPopup,
22 DialogTitle,
23} from "../ui/dialog";
24import { Input } from "../ui/input";
25import { Label } from "../ui/label";
26import { Menu, MenuCheckboxItem, MenuPopup, MenuSeparator, MenuTrigger } from "../ui/menu";
27import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "../ui/table";
28import { Tooltip, TooltipTrigger, TooltipPopup } from "../ui/tooltip";
29import { USAGE_PRICE_FIELDS } from "./usagePriceForm";
30import {
31 isEmptyUsagePriceDraft,
32 usageAliasCell,
33 usagePriceCell,
34 usagePriceTableChanges,
35 usagePriceTableErrors,
36 type UsagePriceDraft,
37 type UsagePriceField,
38} from "./usagePriceTable";
39import {
40 writeUsagePrices,
41 type UsagePriceTarget,
42 type UsagePriceWriteResult,
43} from "./usagePriceTargets";
44
45const priceTargetsAtom = Atom.make((get): readonly UsagePriceTarget[] =>
46 [...get(environmentPresentations.presentationsAtom)].map(([environmentId, environment]) => {
47 const settings = get(serverEnvironment.settingsValueAtom(environmentId));
48 const session = get(environmentSession.sessionStateAtom(environmentId));
49 const sessionData = Option.getOrNull(AsyncResult.value(session));
50 const access =
51 session._tag === "Failure"
52 ? "denied"
53 : sessionData === null
54 ? session.waiting
55 ? "pending"
56 : "denied"
57 : sessionGrantsScope(sessionData, AuthSettingsWriteScope)
58 ? "granted"
59 : "denied";
60 return {
61 environmentId,
62 label: environment.entry.target.label,
63 prices: settings?.usagePriceOverrides ?? null,
64 aliases:
65 environment.serverConfig?.environment.capabilities.usageModelAliases === true
66 ? (settings?.usageModelAliases ?? null)
67 : null,
68 unavailable:
69 environment.connection.phase !== "connected"
70 ? "Offline"
71 : settings === null
72 ? "Prices not loaded"
73 : environment.serverConfig?.environment.capabilities.usagePriceOverrides !== true
74 ? "Update server to edit prices"
75 : access === "pending"
76 ? "Checking permissions…"
77 : access === "denied"
78 ? "Read-only access"
79 : null,
80 };
81 }),
82);
83
84type SaveAttempt = {
85 readonly drafts: readonly UsagePriceDraft[];
86 readonly destinations: readonly {
87 readonly environmentId: EnvironmentId;
88 readonly label: string;
89 }[];
90 readonly results: ReadonlyMap<
91 EnvironmentId,
92 UsagePriceWriteResult | { readonly status: "saving" }
93 >;
94};
95
96/**
97 * Edits custom model prices and mappings. A mapped model's usage counts as its
98 * target model. `initialModel` opens with a new row for that model.
99 */
100export function UsagePriceOverrides({
101 usage,
102 initialSelectedEnvironmentIds,
103 initialModel,
104 onOpenChange,
105}: {
106 readonly usage: readonly EnvironmentUsageStatus[];
107 readonly initialSelectedEnvironmentIds: ReadonlySet<EnvironmentId> | null;
108 readonly initialModel?: string | undefined;
109 readonly onOpenChange: (open: boolean) => void;
110}) {
111 const environments = useAtomValue(priceTargetsAtom);
112 const [selectedIds, setSelectedIds] = useState(initialSelectedEnvironmentIds);
113 const selected = environments.filter(
114 (environment) => selectedIds === null || selectedIds.has(environment.environmentId),
115 );
116 // A model that already has a custom price or mapping is edited in its existing row.
117 const [drafts, setDrafts] = useState<readonly UsagePriceDraft[]>(() =>
118 initialModel === undefined ||
119 selected.some(
120 (environment) =>
121 Object.hasOwn(environment.prices ?? {}, initialModel) ||
122 Object.hasOwn(environment.aliases ?? {}, initialModel),
123 )
124 ? []
125 : [{ id: "new:initial", model: initialModel, isNew: true, values: {} }],
126 );
127 const [pending, setPending] = useState(false);
128 const [attempt, setAttempt] = useState<SaveAttempt | null>(null);
129 const focusRowRef = useRef<string | null>(null);
130 const nextRowId = useRef(0);
131 const updateSettings = useAtomCommand(serverEnvironment.updateSettings, { reportFailure: false });
132 const customModels = [
133 ...new Set(
134 selected.flatMap((environment) => [
135 ...Object.keys(environment.prices ?? {}),
136 ...Object.keys(environment.aliases ?? {}),
137 ]),
138 ),
139 ].sort();
140 const models = [
141 ...new Set([
142 ...customModels,
143 ...usage
144 .filter((environment) =>
145 selected.some((target) => target.environmentId === environment.environmentId),
146 )
147 .flatMap((environment) => environment.summary?.buckets.map((bucket) => bucket.model) ?? []),
148 ]),
149 ].sort();
150 // Keep new rows in place while successful environments publish their updated settings.
151 const newModels = new Set(
152 drafts.filter((draft) => draft.isNew).map((draft) => draft.model.trim()),
153 );
154 const rows: readonly UsagePriceDraft[] = [
155 ...customModels
156 .filter((model) => attempt === null || !newModels.has(model))
157 .map(
158 (model) =>
159 drafts.find((draft) => draft.id === `model:${model}`) ?? {
160 id: `model:${model}`,
161 model,
162 isNew: false,
163 values: {},
164 },
165 ),
166 ...drafts.filter((draft) => draft.isNew),
167 ];
168 const stagedDrafts = drafts.filter((draft) => !isEmptyUsagePriceDraft(draft));
169 const errors = usagePriceTableErrors(selected, stagedDrafts);
170 for (const draft of stagedDrafts) {
171 if (!draft.isNew) continue;
172 if (draft.model.trim() === "") errors.set(draft.id, "Enter a model ID.");
173 else if (
174 attempt === null &&
175 (customModels.includes(draft.model.trim()) ||
176 drafts.some((other) => other.id !== draft.id && other.model.trim() === draft.model.trim()))
177 )
178 errors.set(draft.id, "This model already has a row. Edit its prices there.");
179 }
180 const failedDestinations =
181 attempt?.destinations.filter(
182 (destination) => attempt.results.get(destination.environmentId)?.status === "failed",
183 ) ?? [];
184 const locked = pending || failedDestinations.length > 0;
185 const hasChanges = stagedDrafts.length > 0;
186 const destinationLabel =
187 selected.length === 1 ? selected[0]!.label : `${selected.length} environments`;
188 const selectionLabel = selectedIds === null ? "All environments" : destinationLabel;
189 const discard = () => {
190 setDrafts([]);
191 setAttempt(null);
192 };
193 const selectEnvironments = (ids: ReadonlySet<EnvironmentId> | null) => {
194 setSelectedIds(ids);
195 setAttempt(null);
196 };
197 const updateDraft = (draft: UsagePriceDraft) => {
198 setDrafts((previous) =>
199 previous.some((entry) => entry.id === draft.id)
200 ? previous.map((entry) => (entry.id === draft.id ? draft : entry))
201 : [...previous, draft],
202 );
203 setAttempt(null);
204 };
205 // An existing row edited back to its saved values has nothing left to save.
206 const matchesSaved = (
207 row: UsagePriceDraft,
208 cell: { value: string; placeholder: string },
209 value: string,
210 ) =>
211 !row.isNew &&
212 cell.placeholder !== "Mixed" &&
213 cell.placeholder !== "Unavailable" &&
214 value === cell.value;
215 const updateOrDropDraft = (draft: UsagePriceDraft) => {
216 if (!draft.isNew && Object.keys(draft.values).length === 0 && draft.alias === undefined)
217 setDrafts((previous) => previous.filter((entry) => entry.id !== draft.id));
218 else updateDraft(draft);
219 };
220 const editCell = (row: UsagePriceDraft, field: UsagePriceField, value: string) => {
221 const values = { ...row.values, [field]: value };
222 const original = usagePriceCell(selected, row.model, field);
223 if (matchesSaved(row, original, value)) delete values[field];
224 updateOrDropDraft({ ...row, values });
225 };
226 const editAlias = (row: UsagePriceDraft, value: string) => {
227 const { alias: _alias, ...rest } = row;
228 const original = usageAliasCell(selected, row.model);
229 updateOrDropDraft(matchesSaved(row, original, value) ? rest : { ...rest, alias: value });
230 };
231 const save = async (retry = false) => {
232 if (pending) return;
233 const destinations = retry ? failedDestinations : selected;
234 const edits = retry && attempt ? attempt.drafts : stagedDrafts;
235 if (destinations.length === 0 || edits.length === 0) return;
236 const targets = destinations.map(
237 (destination) =>
238 environments.find(
239 (environment) => environment.environmentId === destination.environmentId,
240 ) ?? { ...destination, prices: null, aliases: null, unavailable: "Environment removed" },
241 );
242 const changes = new Map(
243 targets.map((target) => [target.environmentId, usagePriceTableChanges(target, edits)]),
244 );
245 setPending(true);
246 setAttempt((previous) => ({
247 drafts: edits,
248 destinations: retry && previous ? previous.destinations : targets,
249 results: new Map([
250 ...(retry && previous ? previous.results : []),
251 ...targets.map((target) => [target.environmentId, { status: "saving" } as const] as const),
252 ]),
253 }));
254 let failed = false;
255 await writeUsagePrices({
256 targets: targets.map((target) => ({
257 ...target,
258 unavailable:
259 target.unavailable ?? [...changes.get(target.environmentId)!.errors.values()][0] ?? null,
260 })),
261 changes: new Map([...changes].map(([id, plan]) => [id, plan.changes])),
262 write: async (input) =>
263 readEnvironmentScope(input.environmentId, AuthSettingsWriteScope)
264 ? updateSettings(input)
265 : { _tag: "Failure" },
266 onResult: (environmentId, result) => {
267 if (result.status === "failed") failed = true;
268 setAttempt((previous) =>
269 previous === null
270 ? null
271 : { ...previous, results: new Map(previous.results).set(environmentId, result) },
272 );
273 },
274 });
275 setPending(false);
276 if (!failed) setDrafts([]);
277 };
278
279 return (
280 <Dialog
281 open
282 onOpenChange={(open) => {
283 if (!pending) onOpenChange(open);
284 }}
285 >
286 <DialogPopup className="max-w-4xl">
287 <DialogHeader>
288 <DialogTitle>Custom model prices</DialogTitle>
289 <DialogDescription>
290 Prices and mappings apply to all past and future usage on the environments you select.
291 </DialogDescription>
292 </DialogHeader>
293 <DialogPanel>
294 <div className="flex flex-wrap items-center justify-between gap-2">
295 <div className="flex min-w-0 items-center gap-2">
296 <Label id="usage-prices-apply-label" className="shrink-0">
297 Apply to
298 </Label>
299 <Menu>
300 <MenuTrigger
301 render={<Button variant="outline" size="sm" className="min-w-0" />}
302 aria-labelledby="usage-prices-apply-label usage-prices-selection"
303 disabled={pending || hasChanges}
304 >
305 <span id="usage-prices-selection" className="truncate">
306 {selectionLabel}
307 </span>
308 <ChevronDownIcon className="size-3.5 shrink-0" aria-hidden />
309 </MenuTrigger>
310 <MenuPopup align="start">
311 <MenuCheckboxItem
312 checked={selectedIds === null}
313 closeOnClick={false}
314 onCheckedChange={(checked) => selectEnvironments(checked ? null : new Set())}
315 >
316 All environments
317 </MenuCheckboxItem>
318 <MenuSeparator />
319 {environments.map((environment) => (
320 <MenuCheckboxItem
321 key={environment.environmentId}
322 checked={selectedIds === null || selectedIds.has(environment.environmentId)}
323 closeOnClick={false}
324 onCheckedChange={(checked) => {
325 const next = new Set(selected.map((target) => target.environmentId));
326 if (checked) next.add(environment.environmentId);
327 else next.delete(environment.environmentId);
328 selectEnvironments(next.size === environments.length ? null : next);
329 }}
330 >
331 <span className="flex min-w-0 items-center gap-3">
332 <span className="min-w-0 flex-1 truncate">{environment.label}</span>
333 {environment.unavailable ? (
334 <span className="text-xs text-muted-foreground">
335 {environment.unavailable}
336 </span>
337 ) : null}
338 </span>
339 </MenuCheckboxItem>
340 ))}
341 </MenuPopup>
342 </Menu>
343 </div>
344 <span className="text-xs text-muted-foreground">USD / million tokens</span>
345 </div>
346 {selected.length === 0 ? (
347 <p className="text-sm text-muted-foreground">
348 {environments.length === 0
349 ? "Connect an environment to set model prices."
350 : "Select an environment to see and change its model prices."}
351 </p>
352 ) : (
353 <>
354 <div className="min-w-0 overflow-hidden rounded-lg border border-border">
355 <Table className="min-w-180 table-fixed" aria-label="Custom model prices">
356 <colgroup>
357 <col className="w-[22%]" />
358 <col className="w-[18%]" />
359 {USAGE_PRICE_FIELDS.map((field) => (
360 <col key={field.key} />
361 ))}
362 <col className="w-10" />
363 </colgroup>
364 <TableHeader>
365 <TableRow>
366 <TableHead>Model ID</TableHead>
367 <TableHead>Map to</TableHead>
368 {USAGE_PRICE_FIELDS.map((field) => (
369 <TableHead key={field.key}>{field.label}</TableHead>
370 ))}
371 <TableHead>
372 <Button
373 size="icon-xs"
374 variant="ghost"
375 aria-label="Add model price"
376 disabled={locked}
377 onClick={() => {
378 const id = `new:${nextRowId.current++}`;
379 focusRowRef.current = id;
380 updateDraft({ id, model: "", isNew: true, values: {} });
381 }}
382 >
383 <PlusIcon aria-hidden />
384 </Button>
385 </TableHead>
386 </TableRow>
387 </TableHeader>
388 <TableBody>
389 {rows.length === 0 ? (
390 <TableRow>
391 <TableCell colSpan={7} className="text-center whitespace-normal">
392 <p className="py-6 text-muted-foreground">
393 {selected.some((environment) => environment.prices === null)
394 ? "Some environment prices are unavailable."
395 : "No custom prices or mappings. Add a row to set one."}
396 </p>
397 </TableCell>
398 </TableRow>
399 ) : (
400 rows.map((row) => {
401 const aliasCell = usageAliasCell(selected, row.model);
402 const alias = (row.alias ?? aliasCell.value).trim();
403 return (
404 <TableRow key={row.id} data-row-id={row.id}>
405 <TableCell className="whitespace-normal">
406 {row.isNew ? (
407 <Input
408 size="compact"
409 value={row.model}
410 ref={(node) => {
411 if (node && focusRowRef.current === row.id) {
412 node.focus();
413 focusRowRef.current = null;
414 }
415 }}
416 aria-label="New model ID"
417 aria-invalid={
418 (row.model.trim() !== "" && errors.has(row.id)) || undefined
419 }
420 list="usage-price-models"
421 placeholder="Model ID"
422 autoComplete="off"
423 spellCheck={false}
424 disabled={locked}
425 onChange={(event) =>
426 updateDraft({ ...row, model: event.target.value })
427 }
428 />
429 ) : (
430 <span
431 className={cn(
432 "block break-all",
433 row.removed && "text-muted-foreground line-through",
434 )}
435 >
436 {row.model}
437 </span>
438 )}
439 {errors.has(row.id) &&
440 (row.model.trim() !== "" ||
441 Object.values(row.values).some((value) => value !== "")) ? (
442 <p role="alert" className="mt-1 text-xs text-destructive">
443 {errors.get(row.id)}
444 </p>
445 ) : null}
446 </TableCell>
447 {row.removed ? (
448 <TableCell colSpan={5}>
449 <span className="text-muted-foreground">
450 Automatic pricing after saving
451 </span>
452 </TableCell>
453 ) : (
454 <TableCell>
455 <Input
456 size="compact"
457 value={row.alias ?? aliasCell.value}
458 aria-label={`Map ${row.model || "new model"} to model`}
459 list="usage-alias-models"
460 placeholder={row.isNew ? "Optional" : aliasCell.placeholder}
461 autoComplete="off"
462 spellCheck={false}
463 disabled={locked}
464 onChange={(event) => editAlias(row, event.target.value)}
465 />
466 </TableCell>
467 )}
468 {row.removed ? null : alias !== "" ? (
469 <TableCell colSpan={4} className="whitespace-normal">
470 <span className="break-all text-muted-foreground">
471 Counted as {alias}
472 </span>
473 </TableCell>
474 ) : (
475 USAGE_PRICE_FIELDS.map((field) => {
476 const cell = usagePriceCell(selected, row.model, field.key);
477 return (
478 <TableCell key={field.key}>
479 <Input
480 size="compact"
481 inputMode="decimal"
482 aria-label={`${field.label} price for ${row.model || "new model"}`}
483 value={row.values[field.key] ?? (row.isNew ? "" : cell.value)}
484 placeholder={
485 row.isNew
486 ? field.optional
487 ? "Input rate"
488 : "0.00"
489 : cell.placeholder
490 }
491 autoComplete="off"
492 disabled={locked}
493 font="mono"
494 onChange={(event) =>
495 editCell(row, field.key, event.target.value)
496 }
497 />
498 </TableCell>
499 );
500 })
501 )}
502 <TableCell>
503 <Tooltip>
504 <TooltipTrigger
505 render={<Button size="icon-xs" variant="ghost" />}
506 disabled={locked}
507 aria-label={
508 row.removed
509 ? `Undo reset for ${row.model}`
510 : row.isNew
511 ? "Remove new model"
512 : `Reset price for ${row.model} to automatic`
513 }
514 onClick={() => {
515 if (row.isNew)
516 setDrafts((previous) =>
517 previous.filter((entry) => entry.id !== row.id),
518 );
519 else if (row.removed)
520 updateOrDropDraft({ ...row, removed: false });
521 else updateDraft({ ...row, removed: true });
522 }}
523 >
524 {row.isNew ? (
525 <XIcon aria-hidden />
526 ) : (
527 <RotateCcwIcon aria-hidden />
528 )}
529 </TooltipTrigger>
530 <TooltipPopup>
531 {row.removed
532 ? "Undo reset"
533 : row.isNew
534 ? "Remove row"
535 : "Reset to automatic"}
536 </TooltipPopup>
537 </Tooltip>
538 </TableCell>
539 </TableRow>
540 );
541 })
542 )}
543 </TableBody>
544 </Table>
545 </div>
546 <datalist id="usage-price-models">
547 {models
548 .filter((model) => !customModels.includes(model))
549 .map((model) => (
550 <option key={model} value={model} />
551 ))}
552 </datalist>
553 <datalist id="usage-alias-models">
554 {models
555 .filter((model) =>
556 selected.every(
557 (environment) => !Object.hasOwn(environment.aliases ?? {}, model),
558 ),
559 )
560 .map((model) => (
561 <option key={model} value={model} />
562 ))}
563 </datalist>
564 <p className="text-xs text-muted-foreground">
565 Blank cache rates use the input price. Enter 0 for free tokens. A mapped model’s
566 usage counts as the model it maps to.
567 {selected.length > 1
568 ? " Mixed cells keep each environment’s rate until you edit them."
569 : ""}
570 </p>
571 </>
572 )}
573 {attempt ? (
574 <div role="status" className="grid gap-1 text-xs">
575 {attempt.destinations.map((destination) => {
576 const result = attempt.results.get(destination.environmentId);
577 return (
578 <div key={destination.environmentId} className="flex justify-between gap-3">
579 <span className="truncate text-muted-foreground">{destination.label}</span>
580 <span
581 className={
582 result?.status === "failed"
583 ? "text-right text-destructive"
584 : "text-muted-foreground"
585 }
586 >
587 {result?.status === "failed"
588 ? `Not saved · ${result.error}`
589 : result?.status === "saved"
590 ? "Saved"
591 : "Saving…"}
592 </span>
593 </div>
594 );
595 })}
596 </div>
597 ) : null}
598 </DialogPanel>
599 <DialogFooter variant="bare" className="items-center sm:justify-between">
600 <span className="text-xs text-muted-foreground">
601 {hasChanges ? `Changes apply to ${destinationLabel}` : ""}
602 </span>
603 <div className="flex w-full flex-wrap justify-end gap-2 sm:w-auto">
604 {hasChanges ? (
605 <Button variant="ghost" disabled={pending} onClick={discard}>
606 {failedDestinations.length > 0 ? "Discard pending changes" : "Discard changes"}
607 </Button>
608 ) : null}
609 <Button
610 disabled={
611 pending ||
612 (!hasChanges && failedDestinations.length === 0) ||
613 (failedDestinations.length === 0 && (errors.size > 0 || selected.length === 0))
614 }
615 onClick={() => void save(failedDestinations.length > 0)}
616 >
617 {pending
618 ? "Saving…"
619 : failedDestinations.length > 0
620 ? "Retry failed saves"
621 : "Save changes"}
622 </Button>
623 </div>
624 </DialogFooter>
625 </DialogPopup>
626 </Dialog>
627 );
628}