554 lines · 23.6 KB · TypeScript
RawDownload
import { useAtomValue } from "@effect/atom-react";import type { EnvironmentId } from "@t3tools/contracts";import { ChevronDownIcon, PlusIcon, RotateCcwIcon, XIcon } from "lucide-react";import * as Option from "effect/Option";import { AsyncResult, Atom } from "effect/unstable/reactivity";import { useRef, useState } from "react"; import { isElectron } from "../../env";import { cn } from "../../lib/utils";import { environmentPresentations } from "../../state/presentation";import { serverEnvironment } from "../../state/server";import { environmentSession } from "../../state/session";import { useAtomCommand } from "../../state/use-atom-command";import type { EnvironmentUsageStatus } from "../../state/usage";import {  resolvePrimaryOperateAccess,  resolveRemoteOperateAccess,} from "../settings/ProviderSettingsPanel.logic";import { Button } from "../ui/button";import {  Dialog,  DialogDescription,  DialogFooter,  DialogHeader,  DialogPanel,  DialogPopup,  DialogTitle,} from "../ui/dialog";import { Input } from "../ui/input";import { Label } from "../ui/label";import { Menu, MenuCheckboxItem, MenuPopup, MenuSeparator, MenuTrigger } from "../ui/menu";import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "../ui/table";import { Tooltip, TooltipTrigger, TooltipPopup } from "../ui/tooltip";import { USAGE_PRICE_FIELDS } from "./usagePriceForm";import {  isEmptyUsagePriceDraft,  usagePriceCell,  usagePriceTableChanges,  usagePriceTableErrors,  type UsagePriceDraft,  type UsagePriceField,} from "./usagePriceTable";import {  writeUsagePrices,  type UsagePriceTarget,  type UsagePriceWriteResult,} from "./usagePriceTargets"; const priceTargetsAtom = Atom.make((get): readonly UsagePriceTarget[] =>  [...get(environmentPresentations.presentationsAtom)].map(([environmentId, environment]) => {    const settings = get(serverEnvironment.settingsValueAtom(environmentId));    const session = get(environmentSession.sessionStateAtom(environmentId));    const sessionAccess = {      session: Option.getOrNull(AsyncResult.value(session)),      isPending: session.waiting,      hasError: session._tag === "Failure",    };    const isPrimary = environment.entry.target._tag === "PrimaryConnectionTarget";    const access = isPrimary      ? resolvePrimaryOperateAccess({ ...sessionAccess, isPrimary, hasDesktopBridge: isElectron })      : resolveRemoteOperateAccess(sessionAccess);    return {      environmentId,      label: environment.entry.target.label,      prices: settings?.usagePriceOverrides ?? null,      unavailable:        environment.connection.phase !== "connected"          ? "Offline"          : settings === null            ? "Prices not loaded"            : environment.serverConfig?.environment.capabilities.usagePriceOverrides !== true              ? "Update server to edit prices"              : access === "pending"                ? "Checking permissions…"                : access === "denied"                  ? "Read-only access"                  : null,    };  }),); type SaveAttempt = {  readonly drafts: readonly UsagePriceDraft[];  readonly destinations: readonly {    readonly environmentId: EnvironmentId;    readonly label: string;  }[];  readonly results: ReadonlyMap<    EnvironmentId,    UsagePriceWriteResult | { readonly status: "saving" }  >;}; export function UsagePriceOverrides({  usage,  initialSelectedEnvironmentIds,  onOpenChange,}: {  readonly usage: readonly EnvironmentUsageStatus[];  readonly initialSelectedEnvironmentIds: ReadonlySet<EnvironmentId> | null;  readonly onOpenChange: (open: boolean) => void;}) {  const environments = useAtomValue(priceTargetsAtom);  const [selectedIds, setSelectedIds] = useState(initialSelectedEnvironmentIds);  const selected = environments.filter(    (environment) => selectedIds === null || selectedIds.has(environment.environmentId),  );  const [drafts, setDrafts] = useState<readonly UsagePriceDraft[]>([]);  const [pending, setPending] = useState(false);  const [attempt, setAttempt] = useState<SaveAttempt | null>(null);  const focusRowRef = useRef<string | null>(null);  const nextRowId = useRef(0);  const updateSettings = useAtomCommand(serverEnvironment.updateSettings, { reportFailure: false });  const customModels = [    ...new Set(selected.flatMap((environment) => Object.keys(environment.prices ?? {}))),  ].sort();  const models = [    ...new Set([      ...customModels,      ...usage        .filter((environment) =>          selected.some((target) => target.environmentId === environment.environmentId),        )        .flatMap((environment) => environment.summary?.buckets.map((bucket) => bucket.model) ?? []),    ]),  ].sort();  // Keep new rows in place while successful environments publish their updated settings.  const newModels = new Set(    drafts.filter((draft) => draft.isNew).map((draft) => draft.model.trim()),  );  const rows: readonly UsagePriceDraft[] = [    ...customModels      .filter((model) => attempt === null || !newModels.has(model))      .map(        (model) =>          drafts.find((draft) => draft.id === `model:${model}`) ?? {            id: `model:${model}`,            model,            isNew: false,            values: {},          },      ),    ...drafts.filter((draft) => draft.isNew),  ];  const stagedDrafts = drafts.filter((draft) => !isEmptyUsagePriceDraft(draft));  const errors = usagePriceTableErrors(selected, stagedDrafts);  for (const draft of stagedDrafts) {    if (!draft.isNew) continue;    if (draft.model.trim() === "") errors.set(draft.id, "Enter a model ID.");    else if (      attempt === null &&      (customModels.includes(draft.model.trim()) ||        drafts.some((other) => other.id !== draft.id && other.model.trim() === draft.model.trim()))    )      errors.set(draft.id, "This model already has a row. Edit its prices there.");  }  const failedDestinations =    attempt?.destinations.filter(      (destination) => attempt.results.get(destination.environmentId)?.status === "failed",    ) ?? [];  const locked = pending || failedDestinations.length > 0;  const hasChanges = stagedDrafts.length > 0;  const destinationLabel =    selected.length === 1 ? selected[0]!.label : `${selected.length} environments`;  const selectionLabel = selectedIds === null ? "All environments" : destinationLabel;  const discard = () => {    setDrafts([]);    setAttempt(null);  };  const selectEnvironments = (ids: ReadonlySet<EnvironmentId> | null) => {    setSelectedIds(ids);    setAttempt(null);  };  const updateDraft = (draft: UsagePriceDraft) => {    setDrafts((previous) =>      previous.some((entry) => entry.id === draft.id)        ? previous.map((entry) => (entry.id === draft.id ? draft : entry))        : [...previous, draft],    );    setAttempt(null);  };  const editCell = (row: UsagePriceDraft, field: UsagePriceField, value: string) => {    const values = { ...row.values, [field]: value };    const original = usagePriceCell(selected, row.model, field);    if (      !row.isNew &&      original.placeholder !== "Mixed" &&      original.placeholder !== "Unavailable" &&      value === original.value    )      delete values[field];    if (!row.isNew && Object.keys(values).length === 0)      setDrafts((previous) => previous.filter((entry) => entry.id !== row.id));    else updateDraft({ ...row, values });  };  const save = async (retry = false) => {    if (pending) return;    const destinations = retry ? failedDestinations : selected;    const edits = retry && attempt ? attempt.drafts : stagedDrafts;    if (destinations.length === 0 || edits.length === 0) return;    const targets = destinations.map(      (destination) =>        environments.find(          (environment) => environment.environmentId === destination.environmentId,        ) ?? { ...destination, prices: null, unavailable: "Environment removed" },    );    const changes = new Map(      targets.map((target) => [target.environmentId, usagePriceTableChanges(target, edits)]),    );    setPending(true);    setAttempt((previous) => ({      drafts: edits,      destinations: retry && previous ? previous.destinations : targets,      results: new Map([        ...(retry && previous ? previous.results : []),        ...targets.map((target) => [target.environmentId, { status: "saving" } as const] as const),      ]),    }));    let failed = false;    await writeUsagePrices({      targets: targets.map((target) => ({        ...target,        unavailable:          target.unavailable ?? [...changes.get(target.environmentId)!.errors.values()][0] ?? null,      })),      changes: new Map([...changes].map(([id, plan]) => [id, plan.changes])),      write: updateSettings,      onResult: (environmentId, result) => {        if (result.status === "failed") failed = true;        setAttempt((previous) =>          previous === null            ? null            : { ...previous, results: new Map(previous.results).set(environmentId, result) },        );      },    });    setPending(false);    if (!failed) setDrafts([]);  };   return (    <Dialog      open      onOpenChange={(open) => {        if (!pending) onOpenChange(open);      }}    >      <DialogPopup className="max-w-3xl">        <DialogHeader>          <DialogTitle>Custom model prices</DialogTitle>          <DialogDescription>            Prices apply to all past and future usage on the environments you select.          </DialogDescription>        </DialogHeader>        <DialogPanel className="grid gap-4">          <div className="flex flex-wrap items-center justify-between gap-2">            <div className="flex min-w-0 items-center gap-2">              <Label id="usage-prices-apply-label" className="shrink-0">                Apply to              </Label>              <Menu>                <MenuTrigger                  render={<Button variant="outline" size="sm" className="min-w-0" />}                  aria-labelledby="usage-prices-apply-label usage-prices-selection"                  disabled={pending || hasChanges}                >                  <span id="usage-prices-selection" className="truncate">                    {selectionLabel}                  </span>                  <ChevronDownIcon className="size-3.5 shrink-0" aria-hidden />                </MenuTrigger>                <MenuPopup align="start" className="w-80 max-w-[calc(100vw-2rem)]">                  <MenuCheckboxItem                    checked={selectedIds === null}                    closeOnClick={false}                    onCheckedChange={(checked) => selectEnvironments(checked ? null : new Set())}                  >                    All environments                  </MenuCheckboxItem>                  <MenuSeparator />                  {environments.map((environment) => (                    <MenuCheckboxItem                      key={environment.environmentId}                      checked={selectedIds === null || selectedIds.has(environment.environmentId)}                      closeOnClick={false}                      onCheckedChange={(checked) => {                        const next = new Set(selected.map((target) => target.environmentId));                        if (checked) next.add(environment.environmentId);                        else next.delete(environment.environmentId);                        selectEnvironments(next.size === environments.length ? null : next);                      }}                    >                      <span className="flex min-w-0 items-center gap-3">                        <span className="min-w-0 flex-1 truncate">{environment.label}</span>                        {environment.unavailable ? (                          <span className="text-xs text-muted-foreground">                            {environment.unavailable}                          </span>                        ) : null}                      </span>                    </MenuCheckboxItem>                  ))}                </MenuPopup>              </Menu>            </div>            <span className="text-xs text-muted-foreground">USD / million tokens</span>          </div>          {selected.length === 0 ? (            <p className="text-sm text-muted-foreground">              {environments.length === 0                ? "Connect an environment to set model prices."                : "Select an environment to see and change its model prices."}            </p>          ) : (            <>              <div className="min-w-0 overflow-hidden rounded-lg border border-border">                <Table className="min-w-160 table-fixed" aria-label="Custom model prices">                  <colgroup>                    <col className="w-[34%]" />                    {USAGE_PRICE_FIELDS.map((field) => (                      <col key={field.key} />                    ))}                    <col className="w-10" />                  </colgroup>                  <TableHeader>                    <TableRow className="hover:bg-transparent">                      <TableHead className="pl-3">Model ID</TableHead>                      {USAGE_PRICE_FIELDS.map((field) => (                        <TableHead key={field.key}>{field.label}</TableHead>                      ))}                      <TableHead className="px-1">                        <Button                          size="icon-xs"                          variant="ghost"                          aria-label="Add model price"                          disabled={locked}                          onClick={() => {                            const id = `new:${nextRowId.current++}`;                            focusRowRef.current = id;                            updateDraft({ id, model: "", isNew: true, values: {} });                          }}                        >                          <PlusIcon aria-hidden />                        </Button>                      </TableHead>                    </TableRow>                  </TableHeader>                  <TableBody>                    {rows.length === 0 ? (                      <TableRow className="hover:bg-transparent">                        <TableCell                          colSpan={6}                          className="py-8 text-center whitespace-normal text-muted-foreground"                        >                          {selected.some((environment) => environment.prices === null)                            ? "Some environment prices are unavailable."                            : "No custom prices. Add a row to override automatic pricing."}                        </TableCell>                      </TableRow>                    ) : (                      rows.map((row) => (                        <TableRow                          key={row.id}                          data-row-id={row.id}                          className="hover:bg-transparent"                        >                          <TableCell className="pl-3 whitespace-normal">                            {row.isNew ? (                              <Input                                size="compact"                                value={row.model}                                ref={(node) => {                                  if (node && focusRowRef.current === row.id) {                                    node.focus();                                    focusRowRef.current = null;                                  }                                }}                                aria-label="New model ID"                                aria-invalid={                                  (row.model.trim() !== "" && errors.has(row.id)) || undefined                                }                                list="usage-price-models"                                placeholder="Model ID"                                autoComplete="off"                                spellCheck={false}                                disabled={locked}                                onChange={(event) =>                                  updateDraft({ ...row, model: event.target.value })                                }                              />                            ) : (                              <span                                className={cn(                                  "block break-all",                                  row.removed && "text-muted-foreground line-through",                                )}                              >                                {row.model}                              </span>                            )}                            {errors.has(row.id) &&                            (row.model.trim() !== "" ||                              Object.values(row.values).some((value) => value !== "")) ? (                              <p role="alert" className="mt-1 text-xs text-destructive">                                {errors.get(row.id)}                              </p>                            ) : null}                          </TableCell>                          {row.removed ? (                            <TableCell colSpan={4} className="text-muted-foreground">                              Automatic pricing after saving                            </TableCell>                          ) : (                            USAGE_PRICE_FIELDS.map((field) => {                              const cell = usagePriceCell(selected, row.model, field.key);                              return (                                <TableCell key={field.key} className="px-1">                                  <Input                                    size="compact"                                    inputMode="decimal"                                    aria-label={`${field.label} price for ${row.model || "new model"}`}                                    value={row.values[field.key] ?? (row.isNew ? "" : cell.value)}                                    placeholder={                                      row.isNew                                        ? field.optional                                          ? "Input rate"                                          : "0.00"                                        : cell.placeholder                                    }                                    autoComplete="off"                                    disabled={locked}                                    className="tabular-nums"                                    onChange={(event) =>                                      editCell(row, field.key, event.target.value)                                    }                                  />                                </TableCell>                              );                            })                          )}                          <TableCell className="px-1">                            <Tooltip>                              <TooltipTrigger                                render={<Button size="icon-xs" variant="ghost" />}                                disabled={locked}                                aria-label={                                  row.removed                                    ? `Undo reset for ${row.model}`                                    : row.isNew                                      ? "Remove new model"                                      : `Reset price for ${row.model} to automatic`                                }                                onClick={() => {                                  if (row.isNew)                                    setDrafts((previous) =>                                      previous.filter((entry) => entry.id !== row.id),                                    );                                  else if (row.removed) {                                    if (Object.keys(row.values).length === 0)                                      setDrafts((previous) =>                                        previous.filter((entry) => entry.id !== row.id),                                      );                                    else updateDraft({ ...row, removed: false });                                  } else updateDraft({ ...row, removed: true });                                }}                              >                                {row.isNew ? <XIcon aria-hidden /> : <RotateCcwIcon aria-hidden />}                              </TooltipTrigger>                              <TooltipPopup>                                {row.removed                                  ? "Undo reset"                                  : row.isNew                                    ? "Remove row"                                    : "Reset to automatic"}                              </TooltipPopup>                            </Tooltip>                          </TableCell>                        </TableRow>                      ))                    )}                  </TableBody>                </Table>              </div>              <datalist id="usage-price-models">                {models                  .filter((model) => !customModels.includes(model))                  .map((model) => (                    <option key={model} value={model} />                  ))}              </datalist>              <p className="text-xs text-muted-foreground">                Blank cache rates use the input price. Enter 0 for free tokens.                {selected.length > 1                  ? " Mixed cells keep each environment’s rate until you edit them."                  : ""}              </p>            </>          )}          {attempt ? (            <div role="status" className="grid gap-1 text-xs">

Lines 1–500 of 554.