@marimo-team/islands 0.16.5 → 0.17.1

package/src/components/chat/chat-panel.tsx

@@ -5,7 +5,7 @@ import { useChat } from "@ai-sdk/react";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { DefaultChatTransport, type ToolUIPart } from "ai";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
   AtSignIcon,
   BotMessageSquareIcon,
@@ -30,13 +30,19 @@ import {
 import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useModelChange } from "@/core/ai/config";
 import { AiModelId, type ProviderId } from "@/core/ai/ids/ids";
+import { useStagedAICellsActions } from "@/core/ai/staged-cells";
 import {
   activeChatAtom,
   type Chat,
   type ChatId,
   chatStateAtom,
 } from "@/core/ai/state";
-import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
+import type { ToolNotebookContext } from "@/core/ai/tools/base";
+import {
+  type CopilotMode,
+  FRONTEND_TOOL_REGISTRY,
+} from "@/core/ai/tools/registry";
+import { useCellActions } from "@/core/cells/cells";
 import { aiAtom, aiEnabledAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
@@ -45,7 +51,6 @@ import { useRuntimeManager } from "@/core/runtime/config";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
-
 import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
 import { PromptInput } from "../editor/ai/add-cell-with-ai";
@@ -301,7 +306,11 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
 
     const { saveModeChange } = useModelChange();
 
-    const modeOptions = [
+    const modeOptions: {
+      value: CopilotMode;
+      label: string;
+      subtitle: string;
+    }[] = [
       {
         value: "ask",
         label: "Ask",
@@ -313,6 +322,11 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
         label: "Manual",
         subtitle: "Pure chat, no tool usage",
       },
+      {
+        value: "agent",
+        label: "Agent (beta)",
+        subtitle: "Use AI with access to read and write tools",
+      },
     ];
 
     const isAttachmentSupported =
@@ -511,9 +525,20 @@ const ChatPanelBody = () => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const runtimeManager = useRuntimeManager();
-  const { invokeAiTool } = useRequestClient();
+  const { invokeAiTool, sendRun } = useRequestClient();
 
   const activeChatId = activeChat?.id;
+  const store = useStore();
+
+  const { addStagedCell } = useStagedAICellsActions();
+  const { createNewCell, prepareForRun } = useCellActions();
+  const toolContext: ToolNotebookContext = {
+    addStagedCell,
+    createNewCell,
+    prepareForRun,
+    sendRun,
+    store,
+  };
 
   const {
     messages,
@@ -555,6 +580,13 @@ const ChatPanelBody = () => {
       });
     },
     onToolCall: async ({ toolCall }) => {
+      // Dynamic tool calls will throw an error for toolName
+      // https://ai-sdk.dev/docs/ai-sdk-ui/chatbot-tool-usage#client-side-page
+      if (toolCall.dynamic) {
+        Logger.debug("Skipping dynamic tool call", toolCall);
+        return;
+      }
+
       await handleToolCall({
         invokeAiTool,
         addToolResult,
@@ -563,6 +595,7 @@ const ChatPanelBody = () => {
           toolCallId: toolCall.toolCallId,
           input: toolCall.input as Record<string, never>,
         },
+        toolContext,
       });
     },
     onError: (error) => {

package/src/components/chat/chat-utils.ts

@@ -2,6 +2,7 @@
 
 import type { components } from "@marimo-team/marimo-api";
 import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import type { ToolNotebookContext } from "@/core/ai/tools/base";
 import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
 import type {
   InvokeAiToolRequest,
@@ -106,8 +107,9 @@ interface AddToolResult {
 
 export async function handleToolCall({
   invokeAiTool,
-  addToolResult,
+  addToolResult, // Important that we don't await addToolResult to prevent potential deadlocks
   toolCall,
+  toolContext,
 }: {
   invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
   addToolResult: (result: AddToolResult) => Promise<void>;
@@ -116,6 +118,7 @@ export async function handleToolCall({
     toolCallId: string;
     input: Record<string, never>;
   };
+  toolContext: ToolNotebookContext;
 }) {
   try {
     if (FRONTEND_TOOL_REGISTRY.has(toolCall.toolName)) {
@@ -123,11 +126,12 @@ export async function handleToolCall({
       const response = await FRONTEND_TOOL_REGISTRY.invoke(
         toolCall.toolName,
         toolCall.input,
+        toolContext,
       );
       addToolResult({
         tool: toolCall.toolName,
         toolCallId: toolCall.toolCallId,
-        output: response,
+        output: response.result || response.error,
       });
     } else {
       // Invoke the backend/mcp tool
@@ -153,7 +157,7 @@ export async function handleToolCall({
 
 /**
  * Checks if we should send a message automatically based on the messages.
- * We only want to send a message if we have completed tool calls and there is no reply yet.
+ * We only want to send a message if all tool calls are completed and there is no reply yet.
  */
 export function hasPendingToolCalls(messages: UIMessage[]): boolean {
   if (messages.length === 0) {
@@ -177,7 +181,12 @@ export function hasPendingToolCalls(messages: UIMessage[]): boolean {
     part.type.startsWith("tool-"),
   ) as ToolUIPart[];
 
-  const hasCompletedToolCalls = toolParts.some(
+  // Guard against no tool parts
+  if (toolParts.length === 0) {
+    return false;
+  }
+
+  const allToolCallsCompleted = toolParts.every(
     (part) => part.state === "output-available",
   );
 
@@ -186,6 +195,8 @@ export function hasPendingToolCalls(messages: UIMessage[]): boolean {
   const hasTextContent =
     lastPart.type === "text" && lastPart.text?.trim().length > 0;
 
+  Logger.warn("All tool calls completed: %s", allToolCallsCompleted);
+
   // Only auto-send if we have completed tool calls and there is no reply yet
-  return hasCompletedToolCalls && !hasTextContent;
+  return allToolCallsCompleted && !hasTextContent;
 }

package/src/components/chat/markdown-renderer.css

@@ -33,6 +33,7 @@
 
   code {
     font-size: inherit !important;
+    white-space: pre-wrap;
   }
 
   li.task-list-item {

package/src/components/chat/tool-call-accordion.tsx

@@ -119,6 +119,8 @@ const ToolArgsRenderer: React.FC<{ input: unknown }> = ({ input }) => {
     return null;
   }
 
+  const isEmptyInput = isEmpty(input);
+
   const isObject =
     typeof input === "object" &&
     Object.keys(input as Record<string, unknown>).length > 0;
@@ -129,7 +131,11 @@ const ToolArgsRenderer: React.FC<{ input: unknown }> = ({ input }) => {
         Tool Request
       </h3>
       <pre className="bg-[var(--slate-2)] p-2 text-muted-foreground border border-[var(--slate-4)] rounded text-xs overflow-auto scrollbar-thin max-h-64">
-        {isObject ? JSON.stringify(input, null, 2) : String(input)}
+        {isEmptyInput
+          ? "{}"
+          : isObject
+            ? JSON.stringify(input, null, 2)
+            : String(input)}
       </pre>
     </div>
   );

package/src/components/data-table/__tests__/header-items.test.tsx

@@ -0,0 +1,117 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { SortingState } from "@tanstack/react-table";
+import { describe, expect, it, vi } from "vitest";
+
+describe("multi-column sorting logic", () => {
+  // Extract the core sorting logic to test in isolation
+  const handleSort = (
+    columnId: string,
+    desc: boolean,
+    sortingState: SortingState,
+    setSorting: (state: SortingState) => void,
+    clearSorting: () => void,
+  ) => {
+    const currentSort = sortingState.find((s) => s.id === columnId);
+
+    if (currentSort && currentSort.desc === desc) {
+      // Clicking the same sort again - remove it
+      clearSorting();
+    } else {
+      // New sort or different direction - move to end of stack
+      const otherSorts = sortingState.filter((s) => s.id !== columnId);
+      const newSort = { id: columnId, desc };
+      setSorting([...otherSorts, newSort]);
+    }
+  };
+
+  it("implements stack-based sorting: moves re-clicked column to end", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: false },
+    ];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age - should move age to end with desc=true
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+    ]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("removes sort when clicking same direction twice", () => {
+    const sortingState: SortingState = [{ id: "age", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Asc on age again - should remove the sort
+    handleSort("age", false, sortingState, setSorting, clearSorting);
+
+    expect(clearSorting).toHaveBeenCalled();
+    expect(setSorting).not.toHaveBeenCalled();
+  });
+
+  it("adds new column to end of stack", () => {
+    const sortingState: SortingState = [{ id: "name", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Asc on age - should add age to end
+    handleSort("age", false, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([
+      { id: "name", desc: false },
+      { id: "age", desc: false },
+    ]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("toggles sort direction when clicking opposite", () => {
+    const sortingState: SortingState = [{ id: "age", desc: false }];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age - should toggle to descending
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(setSorting).toHaveBeenCalledWith([{ id: "age", desc: true }]);
+    expect(clearSorting).not.toHaveBeenCalled();
+  });
+
+  it("correctly calculates priority numbers", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+      { id: "dept", desc: false },
+    ];
+
+    // Priority is index + 1
+    const nameSort = sortingState.find((s) => s.id === "name");
+    const namePriority = nameSort ? sortingState.indexOf(nameSort) + 1 : null;
+    expect(namePriority).toBe(1);
+
+    const deptSort = sortingState.find((s) => s.id === "dept");
+    const deptPriority = deptSort ? sortingState.indexOf(deptSort) + 1 : null;
+    expect(deptPriority).toBe(3);
+  });
+
+  it("handles removing column from middle of stack", () => {
+    const sortingState: SortingState = [
+      { id: "name", desc: false },
+      { id: "age", desc: true },
+      { id: "dept", desc: false },
+    ];
+    const setSorting = vi.fn();
+    const clearSorting = vi.fn();
+
+    // Click Desc on age again - should remove it
+    handleSort("age", true, sortingState, setSorting, clearSorting);
+
+    expect(clearSorting).toHaveBeenCalled();
+    // After removal, dept should move from priority 3 to priority 2
+  });
+});

package/src/components/data-table/charts/__tests__/storage.test.ts

@@ -2,8 +2,18 @@
 
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { asMock, SetupMocks } from "@/__mocks__/common";
+
+// Mock the storage module to use our mock storage
+vi.mock("@/utils/storage/storage", () => ({
+  availableStorage: {
+    getItem: vi.fn(),
+    setItem: vi.fn(),
+    removeItem: vi.fn(),
+  },
+}));
+
 import type { CellId } from "@/core/cells/ids";
+import { availableStorage } from "@/utils/storage/storage";
 import { ChartSchema } from "../schemas";
 import type { TabName } from "../storage";
 import { KEY, tabsStorageAtom } from "../storage";
@@ -12,7 +22,6 @@ import { ChartType } from "../types";
 describe("Chart Transforms Storage", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    SetupMocks.localStorage();
     // Reset the store before each test
     const store = getDefaultStore();
     store.set(tabsStorageAtom, new Map());
@@ -87,14 +96,11 @@ describe("Chart Transforms Storage", () => {
 
   describe("LocalStorage integration", () => {
     it("should call localStorage.setItem when setting data", () => {
-      // Reset the mock to ensure clean state
-      asMock(window.localStorage.setItem).mockReset();
-
       const store = getDefaultStore();
       const newMap = new Map();
       store.set(tabsStorageAtom, newMap);
 
-      expect(window.localStorage.setItem).toHaveBeenCalledWith(
+      expect(availableStorage.setItem).toHaveBeenCalledWith(
         KEY,
         JSON.stringify([]),
       );

package/src/components/data-table/charts/charts.tsx

@@ -48,6 +48,7 @@ const CHART_HEIGHT = 290;
 
 export interface TablePanelProps {
   cellId: CellId | null;
+  data: unknown[];
   dataTable: JSX.Element;
   displayHeader: boolean;
   getDataUrl?: GetDataUrl;
@@ -56,6 +57,7 @@ export interface TablePanelProps {
 
 export const TablePanel: React.FC<TablePanelProps> = ({
   cellId,
+  data,
   dataTable,
   getDataUrl,
   fieldTypes,
@@ -211,6 +213,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
         return (
           <TabsContent key={idx} value={tab.tabName} className="h-[400px] mt-1">
             <ChartPanel
+              tableData={data}
               chartConfig={tab.config}
               chartType={tab.chartType}
               saveChart={saveChart}
@@ -228,6 +231,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
 const CHART_PLACEHOLDER_CODE = "X and Y columns are not set";
 
 export const ChartPanel: React.FC<{
+  tableData: unknown[];
   chartConfig: ChartSchemaType | null;
   chartType: ChartType;
   saveChart: (formValues: ChartSchemaType) => void;
@@ -235,6 +239,7 @@ export const ChartPanel: React.FC<{
   getDataUrl?: GetDataUrl;
   fieldTypes?: FieldTypesWithExternalType | null;
 }> = ({
+  tableData,
   chartConfig,
   chartType,
   saveChart,
@@ -255,7 +260,7 @@ export const ChartPanel: React.FC<{
   const { ref: chartContainerRef } = useResizeObserver();
 
   const { data, isPending, error } = useAsyncData(async () => {
-    if (!getDataUrl) {
+    if (!getDataUrl || tableData.length === 0) {
       return [];
     }
 
@@ -276,7 +281,8 @@ export const ChartPanel: React.FC<{
       },
     );
     return chartData;
-  }, []);
+    // Re-run when the data table changes
+  }, [tableData]);
 
   const formValues = form.watch();
 

package/src/components/data-table/charts/storage.ts

@@ -5,7 +5,7 @@ import { capitalize } from "lodash-es";
 import { z } from "zod";
 import type { CellId } from "@/core/cells/ids";
 import { Logger } from "@/utils/Logger";
-import { NotebookScopedLocalStorage } from "@/utils/localStorage";
+import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 import type { TypedString } from "@/utils/typed";
 import { ChartSchema, type ChartSchemaType } from "./schemas";
 import type { ChartType } from "./types";

package/src/components/data-table/column-header.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 "use no memo";
 
-import type { Column } from "@tanstack/react-table";
+import type { Column, Table } from "@tanstack/react-table";
 import { capitalize } from "lodash-es";
 import { FilterIcon, MinusIcon, TextIcon, XIcon } from "lucide-react";
 import { useMemo, useRef, useState } from "react";
@@ -68,6 +68,7 @@ interface DataTableColumnHeaderProps<TData, TValue>
   column: Column<TData, TValue>;
   header: React.ReactNode;
   calculateTopKRows?: CalculateTopKRows;
+  table?: Table<TData>;
 }
 
 export const DataTableColumnHeader = <TData, TValue>({
@@ -75,6 +76,7 @@ export const DataTableColumnHeader = <TData, TValue>({
   header,
   className,
   calculateTopKRows,
+  table,
 }: DataTableColumnHeaderProps<TData, TValue>) => {
   const [isFilterValueOpen, setIsFilterValueOpen] = useState(false);
   const { locale } = useLocale();
@@ -117,7 +119,7 @@ export const DataTableColumnHeader = <TData, TValue>({
         </DropdownMenuTrigger>
         <DropdownMenuContent align="start">
           {renderDataType(column)}
-          {renderSorts(column)}
+          {renderSorts(column, table)}
           {renderCopyColumn(column)}
           {renderColumnPinning(column)}
           {renderColumnWrapping(column)}

package/src/components/data-table/column-wrapping/feature.ts

@@ -13,6 +13,8 @@ import type {
   ColumnWrappingTableState,
 } from "./types";
 
+export const COLUMN_WRAPPING_STYLES = "whitespace-pre-wrap min-w-[200px]";
+
 export const ColumnWrappingFeature: TableFeature = {
   getInitialState: (state): ColumnWrappingTableState => {
     return {

package/src/components/data-table/columns.tsx

@@ -21,6 +21,7 @@ import { Tooltip } from "../ui/tooltip";
 import { DataTableColumnHeader } from "./column-header";
 import type { ColumnChartSpecModel } from "./column-summary/chart-spec-model";
 import { TableColumnSummary } from "./column-summary/column-summary";
+import { COLUMN_WRAPPING_STYLES } from "./column-wrapping/feature";
 import { DatePopover } from "./date-popover";
 import type { FilterType } from "./filters";
 import { getMimeValues, MimeCell } from "./mime-cell";
@@ -165,7 +166,7 @@ export function generateColumns<T>({
         return row[key as keyof T];
       },
 
-      header: ({ column }) => {
+      header: ({ column, table }) => {
         const stats = chartSpecModel?.getColumnStats(key);
         const dtype = column.columnDef.meta?.dtype;
         const headerTitle = headerTooltip?.[key];
@@ -208,6 +209,7 @@ export function generateColumns<T>({
             header={headerWithTooltip}
             column={column}
             calculateTopKRows={calculateTopKRows}
+            table={table}
           />
         );
 
@@ -314,6 +316,7 @@ const PopoutColumn = ({
   rawStringValue,
   contentClassName,
   buttonText,
+  wrapped,
   children,
 }: {
   cellStyles?: string;
@@ -321,6 +324,7 @@ const PopoutColumn = ({
   rawStringValue: string;
   contentClassName?: string;
   buttonText?: string;
+  wrapped?: boolean;
   children: React.ReactNode;
 }) => {
   return (
@@ -335,7 +339,10 @@ const PopoutColumn = ({
           }}
         >
           <span
-            className="cursor-pointer hover:text-link"
+            className={cn(
+              "cursor-pointer hover:text-link",
+              wrapped && COLUMN_WRAPPING_STYLES,
+            )}
             title={rawStringValue}
           >
             {rawStringValue}
@@ -407,7 +414,7 @@ function getCellStyleClass(
     "truncate",
     justify === "center" && "text-center",
     justify === "right" && "text-right",
-    wrapped && "whitespace-pre-wrap min-w-[200px] break-words",
+    wrapped && `${COLUMN_WRAPPING_STYLES} break-words`,
   );
 }
 
@@ -516,6 +523,7 @@ export function renderCellValue<TData, TValue>({
         rawStringValue={stringValue}
         contentClassName="max-h-64 overflow-auto whitespace-pre-wrap break-words text-sm"
         buttonText="X"
+        wrapped={column.getColumnWrapping?.() === "wrap"}
       >
         <UrlDetector parts={parts} />
       </PopoutColumn>
@@ -557,6 +565,7 @@ export function renderCellValue<TData, TValue>({
         cellStyles={cellStyles}
         selectCell={selectCell}
         rawStringValue={rawStringValue}
+        wrapped={column.getColumnWrapping?.() === "wrap"}
       >
         <JsonOutput data={value} format="tree" className="max-h-64" />
       </PopoutColumn>

package/src/components/data-table/data-table.tsx

@@ -215,8 +215,13 @@ const DataTableInternal = <TData,>({
     manualPagination: manualPagination,
     getPaginationRowModel: getPaginationRowModel(),
     // sorting
-    ...(setSorting ? { onSortingChange: setSorting } : {}),
+    ...(setSorting
+      ? {
+          onSortingChange: setSorting,
+        }
+      : {}),
     manualSorting: manualSorting,
+    enableMultiSort: true,
     getSortedRowModel: getSortedRowModel(),
     // filtering
     manualFiltering: true,
@@ -302,7 +307,7 @@ const DataTableInternal = <TData,>({
         )}
         <Table className="relative" ref={tableRef}>
           {showLoadingBar && (
-            <div className="absolute top-0 left-0 h-[3px] w-1/2 bg-primary animate-slide" />
+            <thead className="absolute top-0 left-0 h-[3px] w-1/2 bg-primary animate-slide" />
           )}
           {renderTableHeader(table, Boolean(maxHeight))}
           <CellSelectionProvider>

package/src/components/data-table/header-items.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { PinLeftIcon, PinRightIcon } from "@radix-ui/react-icons";
-import type { Column } from "@tanstack/react-table";
+import type { Column, SortDirection, Table } from "@tanstack/react-table";
 import {
   AlignJustifyIcon,
   ArrowDownWideNarrowIcon,
@@ -163,27 +163,80 @@ export function renderCopyColumn<TData, TValue>(column: Column<TData, TValue>) {
 const AscIcon = ArrowUpNarrowWideIcon;
 const DescIcon = ArrowDownWideNarrowIcon;
 
-export function renderSorts<TData, TValue>(column: Column<TData, TValue>) {
+export function renderSorts<TData, TValue>(
+  column: Column<TData, TValue>,
+  table?: Table<TData>,
+) {
   if (!column.getCanSort()) {
     return null;
   }
 
+  const sortDirection = column.getIsSorted();
+  const sortingIndex = column.getSortIndex();
+
+  const sortingState = table?.getState().sorting;
+  const hasMultiSort = sortingState?.length && sortingState.length > 1;
+
+  const renderSortIndex = () => {
+    return (
+      <span className="ml-auto text-xs font-medium">{sortingIndex + 1}</span>
+    );
+  };
+
+  const renderClearSort = () => {
+    if (!sortDirection) {
+      return null;
+    }
+
+    if (!hasMultiSort) {
+      // render clear sort for this column
+      return (
+        <DropdownMenuItem onClick={() => column.clearSorting()}>
+          <ChevronsUpDown className="mo-dropdown-icon" />
+          Clear sort
+        </DropdownMenuItem>
+      );
+    }
+
+    // render clear sort for all columns
+    return (
+      <DropdownMenuItem onClick={() => table?.resetSorting()}>
+        <ChevronsUpDown className="mo-dropdown-icon" />
+        Clear all sorts
+      </DropdownMenuItem>
+    );
+  };
+
+  const toggleSort = (direction: SortDirection) => {
+    // Clear sort if clicking the same direction
+    if (sortDirection === direction) {
+      column.clearSorting();
+    } else {
+      // Toggle sort direction
+      const descending = direction === "desc";
+      column.toggleSorting(descending, true);
+    }
+  };
+
   return (
     <>
-      <DropdownMenuItem onClick={() => column.toggleSorting(false)}>
+      <DropdownMenuItem
+        onClick={() => toggleSort("asc")}
+        className={sortDirection === "asc" ? "bg-accent" : ""}
+      >
         <AscIcon className="mo-dropdown-icon" />
         Asc
+        {sortDirection === "asc" && renderSortIndex()}
       </DropdownMenuItem>
-      <DropdownMenuItem onClick={() => column.toggleSorting(true)}>
+      <DropdownMenuItem
+        onClick={() => toggleSort("desc")}
+        className={sortDirection === "desc" ? "bg-accent" : ""}
+      >
         <DescIcon className="mo-dropdown-icon" />
         Desc
+        {sortDirection === "desc" && renderSortIndex()}
       </DropdownMenuItem>
-      {column.getIsSorted() && (
-        <DropdownMenuItem onClick={() => column.clearSorting()}>
-          <ChevronsUpDown className="mo-dropdown-icon" />
-          Clear sort
-        </DropdownMenuItem>
-      )}
+      {renderClearSort()}
       <DropdownMenuSeparator />
     </>
   );