@marimo-team/frontend 0.15.5 → 0.16.0

package/src/components/editor/ai/completion-utils.ts

@@ -10,6 +10,7 @@ import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import type { FileUIPart } from "ai";
 import { getAIContextRegistry } from "@/core/ai/context/context";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
+import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
@@ -145,3 +146,94 @@ export function addContextCompletion(
     startCompletion(inputRef.current.view);
   }
 }
+
+export interface AiCompletion {
+  language: LanguageAdapterType;
+  code: string;
+}
+
+/**
+ * Extracts code blocks (delimited by triple backticks) and their language ("python", "sql", "markdown").
+ * Defaults to "python" if no language is specified or no code blocks are found.
+ * Returns an array of AiCompletion objects.
+ */
+export function codeToCells(code: string): AiCompletion[] {
+  if (code.trim().length === 0) {
+    return [];
+  }
+
+  // If there are no backticks, assume code is in 1 cell and python
+  if (!code.includes("```")) {
+    return [{ language: "python", code: code }];
+  }
+
+  // If code has opening backticks, get the code after it
+  const cells: AiCompletion[] = [];
+  let start = 0;
+
+  let openIndex = code.indexOf("```", start);
+  while (openIndex !== -1) {
+    const newlineIndex = code.indexOf("\n", openIndex);
+    if (newlineIndex === -1) {
+      // If there's no newline after opening backticks, treat everything after as code
+      const remaining = code.slice(openIndex + 3);
+      const firstSpace = remaining.indexOf(" ");
+      const language =
+        firstSpace === -1 ? remaining : remaining.slice(0, firstSpace);
+      const finalLanguage =
+        language === "markdown"
+          ? "markdown"
+          : language === "sql"
+            ? "sql"
+            : "python";
+      // Extract code after the language identifier
+      const codeContent =
+        firstSpace === -1 ? "" : remaining.slice(firstSpace + 1);
+      if (codeContent.trim()) {
+        cells.push({ language: finalLanguage, code: codeContent.trim() });
+      }
+      break;
+    }
+
+    let language = code.slice(openIndex + 3, newlineIndex).trim() || "";
+    language =
+      language === "markdown"
+        ? "markdown"
+        : language === "sql"
+          ? "sql"
+          : "python";
+    const codeStart = newlineIndex + 1;
+
+    const closeIndex = code.indexOf("```", codeStart);
+    if (closeIndex === -1) {
+      // If there's no closing backticks, treat everything after the opening as code
+      const codeContent = code.slice(codeStart).replace(/\n+$/, "");
+      if (codeContent.trim()) {
+        cells.push({
+          language: language as LanguageAdapterType,
+          code: codeContent,
+        });
+      }
+      break;
+    }
+
+    // Remove trailing newlines
+    const codeContent = code.slice(codeStart, closeIndex).replace(/\n+$/, "");
+    if (codeContent.trim()) {
+      cells.push({
+        language: language as LanguageAdapterType,
+        code: codeContent,
+      });
+    }
+
+    start = closeIndex + 3;
+    openIndex = code.indexOf("```", start);
+  }
+
+  // If no cells found, assume code is in 1 cell and python
+  if (cells.length === 0) {
+    cells.push({ language: "python", code: code });
+  }
+
+  return cells;
+}

package/src/components/editor/ai/transport/chat-transport.tsx

@@ -0,0 +1,36 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import {
+  DefaultChatTransport,
+  type HttpChatTransportInitOptions,
+  type UIMessage,
+  type UIMessageChunk,
+} from "ai";
+
+/**
+ * Thin wrapper around the DefaultChatTransport that calls a callback when a chunk is received.
+ */
+export class StreamingChunkTransport<
+  UI_MESSAGE extends UIMessage,
+> extends DefaultChatTransport<UI_MESSAGE> {
+  constructor(
+    options: HttpChatTransportInitOptions<UI_MESSAGE> = {},
+    private onChunkReceived: (chunk: UIMessageChunk) => void,
+  ) {
+    super(options);
+  }
+
+  protected override processResponseStream(
+    stream: ReadableStream<Uint8Array>,
+  ): ReadableStream<UIMessageChunk> {
+    const onChunkReceived = this.onChunkReceived;
+    return super.processResponseStream(stream).pipeThrough(
+      new TransformStream<UIMessageChunk, UIMessageChunk>({
+        async transform(chunk, controller) {
+          onChunkReceived(chunk);
+          controller.enqueue(chunk);
+        },
+      }),
+    );
+  }
+}

package/src/components/editor/cell/StagedAICell.tsx

@@ -0,0 +1,51 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useAtomValue, useStore } from "jotai";
+import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
+import type { CellId } from "@/core/cells/ids";
+import { cn } from "@/utils/cn";
+import { CompletionActionsCellFooter } from "../ai/completion-handlers";
+
+export const StagedAICellBackground: React.FC<{
+  cellId: CellId;
+  className?: string;
+}> = ({ cellId, className }) => {
+  const stagedAICells = useAtomValue(stagedAICellsAtom);
+
+  if (!stagedAICells.has(cellId)) {
+    return null;
+  }
+
+  return <div className={cn("mo-ai-generated-cell", className)} />;
+};
+
+export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
+  cellId,
+}) => {
+  const store = useStore();
+  const stagedAICells = useAtomValue(stagedAICellsAtom);
+  const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
+
+  if (!stagedAICells.has(cellId)) {
+    return null;
+  }
+
+  const handleAcceptCompletion = () => {
+    removeStagedCell(cellId);
+  };
+
+  const handleDeclineCompletion = () => {
+    deleteStagedCell(cellId);
+  };
+
+  return (
+    <div className="flex items-center justify-end gap-1.5 w-full pb-1 pt-2">
+      <CompletionActionsCellFooter
+        isLoading={false}
+        onAccept={handleAcceptCompletion}
+        onDecline={handleDeclineCompletion}
+        size="xs"
+      />
+    </div>
+  );
+};

package/src/components/editor/cell/cell-actions.tsx

@@ -56,7 +56,8 @@ const CellActionsDropdownInternal = (
   ref: React.Ref<CellActionsDropdownHandle>,
 ) => {
   const [open, setOpen] = useState(false);
-  const actions = useCellActionButtons({ cell: props });
+  const closePopover = () => setOpen(false);
+  const actions = useCellActionButtons({ cell: props, closePopover });
 
   // store the last focused element so we can restore it when the popover closes
   const restoreFocus = useRestoreFocus();

package/src/components/terminal/__tests__/state.test.ts

@@ -0,0 +1,207 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { exportedForTesting } from "../state";
+
+const { reducer, initialState } = exportedForTesting;
+
+describe("terminal state", () => {
+  describe("initialState", () => {
+    it("should return initial state with empty pendingCommands and isReady false", () => {
+      const state = initialState();
+      expect(state).toEqual({
+        pendingCommands: [],
+        isReady: false,
+      });
+    });
+  });
+
+  describe("reducer", () => {
+    it("should add a command to pendingCommands", () => {
+      const state = initialState();
+      const text = "ls -la";
+
+      const newState = reducer(state, { type: "addCommand", payload: text });
+
+      expect(newState.pendingCommands).toHaveLength(1);
+      expect(newState.pendingCommands[0]).toMatchObject({
+        text: "ls -la",
+        timestamp: expect.any(Number),
+      });
+      expect(newState.pendingCommands[0].id).toBeDefined();
+      expect(newState.isReady).toBe(false);
+    });
+
+    it("should add multiple commands to pendingCommands", () => {
+      const state = initialState();
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, { type: "addCommand", payload: "cd /home" });
+      newState = reducer(newState, { type: "addCommand", payload: "pwd" });
+
+      expect(newState.pendingCommands).toHaveLength(3);
+      expect(newState.pendingCommands[0].text).toBe("ls -la");
+      expect(newState.pendingCommands[1].text).toBe("cd /home");
+      expect(newState.pendingCommands[2].text).toBe("pwd");
+    });
+
+    it("should remove a command by id", () => {
+      const state = initialState();
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, { type: "addCommand", payload: "cd /home" });
+      newState = reducer(newState, { type: "addCommand", payload: "pwd" });
+
+      const commandToRemove = newState.pendingCommands[1];
+      newState = reducer(newState, {
+        type: "removeCommand",
+        payload: commandToRemove.id,
+      });
+
+      expect(newState.pendingCommands).toHaveLength(2);
+      expect(newState.pendingCommands[0].text).toBe("ls -la");
+      expect(newState.pendingCommands[1].text).toBe("pwd");
+    });
+
+    it("should not remove anything if command id does not exist", () => {
+      const state = initialState();
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, { type: "addCommand", payload: "cd /home" });
+
+      const originalLength = newState.pendingCommands.length;
+      newState = reducer(newState, {
+        type: "removeCommand",
+        payload: "non-existent-id",
+      });
+
+      expect(newState.pendingCommands).toHaveLength(originalLength);
+    });
+
+    it("should set isReady to true", () => {
+      const state = initialState();
+
+      const newState = reducer(state, { type: "setReady", payload: true });
+
+      expect(newState.isReady).toBe(true);
+      expect(newState.pendingCommands).toEqual([]);
+    });
+
+    it("should set isReady to false", () => {
+      const state = { ...initialState(), isReady: true };
+
+      const newState = reducer(state, { type: "setReady", payload: false });
+
+      expect(newState.isReady).toBe(false);
+    });
+
+    it("should clear all pending commands", () => {
+      const state = initialState();
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, { type: "addCommand", payload: "cd /home" });
+      newState = reducer(newState, { type: "addCommand", payload: "pwd" });
+
+      expect(newState.pendingCommands).toHaveLength(3);
+
+      newState = reducer(newState, {
+        type: "clearCommands",
+        payload: undefined,
+      });
+
+      expect(newState.pendingCommands).toHaveLength(0);
+      expect(newState.isReady).toBe(false);
+    });
+
+    it("should preserve other state when adding commands", () => {
+      const state = { ...initialState(), isReady: true };
+
+      const newState = reducer(state, {
+        type: "addCommand",
+        payload: "ls -la",
+      });
+
+      expect(newState.isReady).toBe(true);
+      expect(newState.pendingCommands).toHaveLength(1);
+    });
+
+    it("should preserve other state when removing commands", () => {
+      const state = { ...initialState(), isReady: true };
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, {
+        type: "removeCommand",
+        payload: newState.pendingCommands[0].id,
+      });
+
+      expect(newState.isReady).toBe(true);
+      expect(newState.pendingCommands).toHaveLength(0);
+    });
+
+    it("should preserve other state when setting ready", () => {
+      const state = initialState();
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, { type: "addCommand", payload: "cd /home" });
+      newState = reducer(newState, { type: "setReady", payload: true });
+
+      expect(newState.isReady).toBe(true);
+      expect(newState.pendingCommands).toHaveLength(2);
+    });
+
+    it("should preserve other state when clearing commands", () => {
+      const state = { ...initialState(), isReady: true };
+
+      let newState = reducer(state, { type: "addCommand", payload: "ls -la" });
+      newState = reducer(newState, {
+        type: "clearCommands",
+        payload: undefined,
+      });
+
+      expect(newState.isReady).toBe(true);
+      expect(newState.pendingCommands).toHaveLength(0);
+    });
+  });
+
+  describe("command properties", () => {
+    it("should generate unique ids for commands", () => {
+      const state = initialState();
+
+      let newState = reducer(state, {
+        type: "addCommand",
+        payload: "command1",
+      });
+      newState = reducer(newState, { type: "addCommand", payload: "command2" });
+
+      const ids = newState.pendingCommands.map((cmd) => cmd.id);
+      expect(ids[0]).not.toBe(ids[1]);
+      expect(ids[0]).toBeDefined();
+      expect(ids[1]).toBeDefined();
+    });
+
+    it("should set timestamp when adding commands", () => {
+      const state = initialState();
+      const beforeTime = Date.now();
+
+      const newState = reducer(state, {
+        type: "addCommand",
+        payload: "ls -la",
+      });
+
+      const afterTime = Date.now();
+      const command = newState.pendingCommands[0];
+
+      expect(command.timestamp).toBeGreaterThanOrEqual(beforeTime);
+      expect(command.timestamp).toBeLessThanOrEqual(afterTime);
+    });
+
+    it("should preserve text when adding commands", () => {
+      const state = initialState();
+      const text = "echo 'Hello World'";
+
+      const newState = reducer(state, { type: "addCommand", payload: text });
+
+      expect(newState.pendingCommands[0].text).toBe(text);
+    });
+  });
+});

package/src/components/terminal/hooks.ts

@@ -0,0 +1,41 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useChromeActions } from "../editor/chrome/state";
+import { useTerminalActions } from "./state";
+
+/**
+ * Hook for sending commands to the terminal programmatically.
+ * This will:
+ * 1. Open the terminal if it's not already open
+ * 2. Wait for the terminal to be connected
+ * 3. Send the command text to the terminal
+ *
+ * @example
+ * ```tsx
+ * function CopyButton({ command }: { command: string }) {
+ *   const { sendCommand } = useTerminalCommands();
+ *
+ *   return (
+ *     <button onClick={() => sendCommand(command)}>
+ *       Copy to Terminal
+ *     </button>
+ *   );
+ * }
+ * ```
+ */
+export function useTerminalCommands() {
+  const { addCommand } = useTerminalActions();
+  const { setIsTerminalOpen } = useChromeActions();
+
+  const sendCommand = (text: string) => {
+    // First, ensure the terminal is open
+    setIsTerminalOpen(true);
+
+    // Add the command to the queue
+    addCommand(text);
+  };
+
+  return {
+    sendCommand,
+  };
+}

package/src/components/terminal/state.ts

@@ -0,0 +1,75 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { useAtomValue } from "jotai";
+import { createReducerAndAtoms } from "@/utils/createReducer";
+import { generateUUID } from "@/utils/uuid";
+
+export interface TerminalCommand {
+  id: string;
+  text: string;
+  timestamp: number;
+}
+
+export interface TerminalState {
+  pendingCommands: TerminalCommand[];
+  isReady: boolean;
+}
+
+function initialState(): TerminalState {
+  return {
+    pendingCommands: [],
+    isReady: false,
+  };
+}
+
+const {
+  reducer,
+  createActions,
+  valueAtom: terminalStateAtom,
+  useActions,
+} = createReducerAndAtoms(initialState, {
+  addCommand: (state, text: string) => {
+    const command: TerminalCommand = {
+      id: generateUUID(),
+      text,
+      timestamp: Date.now(),
+    };
+
+    return {
+      ...state,
+      pendingCommands: [...state.pendingCommands, command],
+    };
+  },
+  removeCommand: (state, commandId: string) => ({
+    ...state,
+    pendingCommands: state.pendingCommands.filter(
+      (cmd) => cmd.id !== commandId,
+    ),
+  }),
+  setReady: (state, isReady: boolean) => ({
+    ...state,
+    isReady,
+  }),
+  clearCommands: (state) => ({
+    ...state,
+    pendingCommands: [],
+  }),
+});
+
+/**
+ * React hook to get the terminal state.
+ */
+export const useTerminalState = () => useAtomValue(terminalStateAtom);
+
+/**
+ * React hook to get the terminal actions.
+ */
+export function useTerminalActions() {
+  return useActions();
+}
+
+export const exportedForTesting = {
+  reducer,
+  createActions,
+  initialState,
+};