gpustack-ui/src/pages/usage/token-tab/tables/models-table.tsx

import PluginExtraFields from '@/components/plugin-extra-fields';
import { TABLE_SORT_DIRECTIONS } from '@/config/settings';
import PageBox from '@/pages/_components/page-box';
import { useIntl } from '@umijs/max';
import { Table } from 'antd';
import { useEffect, useMemo, useRef, useState } from 'react';
import { BreakdownFilters, BreakdownItem } from '../../config/types';
import useModelsColumns from '../../hooks/use-models-columns';
import useQueryBreakdownList from '../../services/use-query-breakdown-list';
import getBreakdownRowKey from '../../utils/get-breakdown-row-key';

const Models: React.FC<{
  filters: BreakdownFilters;
  dateRange: { start_date: string; end_date: string };
  scope: string;
  pageResetKey?: number;
  refreshKey?: number;
}> = ({ filters, dateRange, scope, pageResetKey = 0, refreshKey = 0 }) => {
  const intl = useIntl();

  const { loading, dataSource, fetchData } = useQueryBreakdownList({
    key: 'modelsTableData'
  });
  const [queryParams, setQueryParams] = useState<{
    page: number;
    perPage: number;
    sort_by: string;
  }>({
    page: 1,
    perPage: 10,
    sort_by: '-total_tokens'
  });
  const pendingPageResetRef = useRef(false);

  const handleTableChange = (pagination: any, filters: any, sorter: any) => {
    const sort_by =
      sorter.order === 'descend' ? `-${sorter.field}` : sorter.field;
    setQueryParams((prev) => ({
      ...prev,
      page: 1,
      sort_by
    }));
  };

  const handlePageChange = (page: number, pageSize: number) => {
    setQueryParams((prev) => ({
      ...prev,
      page,
      perPage: pageSize
    }));
  };

  const columns = useModelsColumns();

  useEffect(() => {
    if (queryParams.page !== 1) {
      pendingPageResetRef.current = true;
      setQueryParams((prev) => ({
        ...prev,
        page: 1
      }));
    }
  }, [pageResetKey]);

  useEffect(() => {
    if (pendingPageResetRef.current && queryParams.page !== 1) {
      return;
    }
    pendingPageResetRef.current = false;

    fetchData({
      ...queryParams,
      group_by: ['route'],
      // Send the full filter set (route / user / api_key), not just the
      // table's own dimension, so the breakdown matches the trend chart.
      filters,
      scope: scope,
      ...dateRange
    });
  }, [
    dateRange.end_date,
    dateRange.start_date,
    filters,
    queryParams.page,
    queryParams.perPage,
    queryParams.sort_by,
    refreshKey,
    scope
  ]);

  // Route ids currently visible in the Models tab. Mirrored to the
  // plugin slot below so enterprise plugins can bulk-fetch per-route
  // data (e.g. the caller's quota / usage on each route) in one call
  // rather than firing N round-trips from each cell. `refreshToken`
  // bumps with every successful fetch so plugins re-pull even when the
  // id set hasn't changed.
  const pluginContext = useMemo(
    () => ({
      routeIds: (dataSource.dataList || [])
        .map((item) => item.route?.identity?.current?.route_id)
        .filter((id): id is number => id != null),
      refreshToken: refreshKey
    }),
    [dataSource.dataList, refreshKey]
  );

  return (
    <>
      <PageBox>
        <Table
          columns={columns}
          dataSource={dataSource.dataList}
          rowKey={(record: BreakdownItem) =>
            getBreakdownRowKey(record, 'models')
          }
          loading={{
            spinning: loading,
            size: 'middle'
          }}
          sortDirections={TABLE_SORT_DIRECTIONS}
          showSorterTooltip={false}
          onChange={handleTableChange}
          pagination={{
            size: 'middle',
            showSizeChanger: true,
            pageSize: queryParams.perPage,
            current: queryParams.page,
            total: dataSource.total,
            hideOnSinglePage: queryParams.perPage === 10,
            onChange: handlePageChange
          }}
        ></Table>
      </PageBox>
      {/* Page-level data lifecycle for plugin-contributed extra
          columns on this tab. Receives the visible route ids so the
          plugin can bulk-fetch per-route data in one call; renders
          nothing when no plugin is registered. */}
      <PluginExtraFields name="UsageModelsPageGlobal" context={pluginContext} />
    </>
  );
};

export default Models;