feat(datasets): edit an episode's task in the dataset browser

frontend/src/components/dataset/EpisodeList.tsx

@@ -1,13 +1,16 @@
 import React, { useEffect, useRef } from "react";
+import { Pencil } from "lucide-react";
 import { EpisodeSummary, formatDuration } from "@/lib/datasetApi";
 
 interface EpisodeListProps {
   episodes: EpisodeSummary[];
   selected: number | null;
   onSelect: (episodeIndex: number) => void;
+  /** Shows an edit button on the selected episode's task when set. */
+  onEditTask?: (episodeIndex: number) => void;
 }
 
-const EpisodeList: React.FC<EpisodeListProps> = ({ episodes, selected, onSelect }) => {
+const EpisodeList: React.FC<EpisodeListProps> = ({ episodes, selected, onSelect, onEditTask }) => {
   const activeRef = useRef<HTMLButtonElement | null>(null);
 
   // Follow the selection. Paging with Previous/Next, or opening a link straight
@@ -39,29 +42,43 @@ const EpisodeList: React.FC<EpisodeListProps> = ({ episodes, selected, onSelect
         {episodes.map((ep) => {
           const active = ep.episode_index === selected;
           return (
-            <button
+            <div
               key={ep.episode_index}
-              ref={active ? activeRef : undefined}
-              type="button"
-              onClick={() => onSelect(ep.episode_index)}
-              className={`flex w-full items-center gap-2 border-b border-gray-900 px-3 py-2 text-left transition last:border-b-0 ${
+              className={`flex items-center border-b border-gray-900 transition last:border-b-0 ${
                 active ? "bg-gray-800/80" : "hover:bg-gray-900"
-              } focus:outline-none focus-visible:ring-1 focus-visible:ring-inset focus-visible:ring-orange-500`}
+              }`}
             >
-              <span
-                className={`w-8 text-xs tabular-nums ${
-                  active ? "font-semibold text-orange-400" : "text-gray-500"
-                }`}
+              <button
+                ref={active ? activeRef : undefined}
+                type="button"
+                onClick={() => onSelect(ep.episode_index)}
+                className="flex min-w-0 flex-1 items-center gap-2 px-3 py-2 text-left focus:outline-none focus-visible:ring-1 focus-visible:ring-inset focus-visible:ring-orange-500"
               >
-                {ep.episode_index}
-              </span>
-              <span className="flex-1 truncate text-xs text-gray-300">
-                {ep.tasks.length > 0 ? ep.tasks.join(", ") : "—"}
-              </span>
-              <span className="text-[11px] tabular-nums text-gray-500">
-                {formatDuration(ep.duration_s)}
-              </span>
-            </button>
+                <span
+                  className={`w-8 text-xs tabular-nums ${
+                    active ? "font-semibold text-orange-400" : "text-gray-500"
+                  }`}
+                >
+                  {ep.episode_index}
+                </span>
+                <span className="flex-1 truncate text-xs text-gray-300">
+                  {ep.tasks.length > 0 ? ep.tasks.join(", ") : "—"}
+                </span>
+                <span className="text-[11px] tabular-nums text-gray-500">
+                  {formatDuration(ep.duration_s)}
+                </span>
+              </button>
+              {active && onEditTask && (
+                <button
+                  type="button"
+                  onClick={() => onEditTask(ep.episode_index)}
+                  aria-label={`Edit the task of episode ${ep.episode_index}`}
+                  className="mr-2 rounded p-1 text-gray-500 hover:bg-gray-700 hover:text-gray-200 focus:outline-none focus-visible:ring-1 focus-visible:ring-orange-500"
+                >
+                  <Pencil className="h-3 w-3" />
+                </button>
+              )}
+            </div>
           );
         })}
       </div>

frontend/src/components/dataset/EpisodeTaskDialog.test.tsx

@@ -0,0 +1,61 @@
+// @vitest-environment jsdom
+import React, { act } from "react";
+import { createRoot } from "react-dom/client";
+import { expect, test } from "vitest";
+
+import EpisodeTaskDialog from "./EpisodeTaskDialog";
+
+test("saves the task for the browsed episode or for every episode", async () => {
+  Object.assign(globalThis, { IS_REACT_ACT_ENVIRONMENT: true });
+  const host = document.createElement("div");
+  document.body.append(host);
+  const root = createRoot(host);
+  const saved: [string, number | null][] = [];
+
+  try {
+    await act(async () => {
+      root.render(
+        <EpisodeTaskDialog
+          repoId="local/demo"
+          episodeIndex={3}
+          currentTask="pick the cube"
+          totalEpisodes={5}
+          open
+          onOpenChange={() => {}}
+          onSave={async (task, episodeIndex) => {
+            saved.push([task, episodeIndex]);
+          }}
+        />,
+      );
+    });
+
+    const input = document.getElementById("episode-task") as HTMLInputElement;
+    const saveButton = () =>
+      [...document.body.querySelectorAll("button")].find(
+        (b) => b.textContent === "Save task",
+      ) as HTMLButtonElement;
+    const type = async (value: string) => {
+      const setter = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")!.set!;
+      await act(async () => {
+        setter.call(input, value);
+        input.dispatchEvent(new Event("input", { bubbles: true }));
+      });
+    };
+
+    expect(input.value).toBe("pick the cube");
+
+    await type("  ");
+    expect(saveButton().disabled).toBe(true);
+
+    await type(" pick the bowl ");
+    await act(async () => saveButton().click());
+    expect(saved).toEqual([["pick the bowl", 3]]);
+
+    await act(async () => (document.getElementById("episode-task-all") as HTMLButtonElement).click());
+    await act(async () => saveButton().click());
+    expect(saved.at(-1)).toEqual(["pick the bowl", null]);
+  } finally {
+    await act(async () => root.unmount());
+    host.remove();
+  }
+});

frontend/src/components/dataset/EpisodeTaskDialog.tsx

@@ -0,0 +1,130 @@
+import { useEffect, useState } from "react";
+import { Loader2 } from "lucide-react";
+
+import { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import {
+  Dialog,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+} from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import { Label } from "@/components/ui/label";
+
+interface EpisodeTaskDialogProps {
+  repoId: string;
+  episodeIndex: number;
+  currentTask: string;
+  totalEpisodes: number;
+  open: boolean;
+  onOpenChange: (open: boolean) => void;
+  /** `episodeIndex` is null when the task applies to every episode. */
+  onSave: (task: string, episodeIndex: number | null) => Promise<void>;
+}
+
+export default function EpisodeTaskDialog({
+  repoId,
+  episodeIndex,
+  currentTask,
+  totalEpisodes,
+  open,
+  onOpenChange,
+  onSave,
+}: EpisodeTaskDialogProps) {
+  const [task, setTask] = useState("");
+  const [allEpisodes, setAllEpisodes] = useState(false);
+  const [saving, setSaving] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+
+  useEffect(() => {
+    if (open) {
+      setTask(currentTask);
+      setAllEpisodes(false);
+      setSaving(false);
+      setError(null);
+    }
+  }, [open, currentTask]);
+
+  const canSave = task.trim() !== "" && !saving;
+
+  const handleOpenChange = (nextOpen: boolean) => {
+    if (saving) return;
+    onOpenChange(nextOpen);
+  };
+
+  const handleSave = async () => {
+    if (!canSave) return;
+    setSaving(true);
+    try {
+      await onSave(task.trim(), allEpisodes ? null : episodeIndex);
+      onOpenChange(false);
+    } catch (caughtError) {
+      setError(caughtError instanceof Error ? caughtError.message : "Could not save the task.");
+    } finally {
+      setSaving(false);
+    }
+  };
+
+  return (
+    <Dialog open={open} onOpenChange={handleOpenChange}>
+      <DialogContent className="border-gray-700 bg-gray-900 text-gray-300 sm:max-w-lg">
+        <DialogHeader>
+          <DialogTitle className="text-white">Edit task</DialogTitle>
+          <DialogDescription>
+            The task is the instruction a policy is trained on, saved with each episode of{" "}
+            <code>{repoId}</code>.
+          </DialogDescription>
+        </DialogHeader>
+
+        <div className="space-y-3">
+          <div className="space-y-2">
+            <Label htmlFor="episode-task" className="text-gray-300">
+              Task
+            </Label>
+            <Input
+              id="episode-task"
+              value={task}
+              onChange={(event) => {
+                setError(null);
+                setTask(event.target.value);
+              }}
+              onKeyDown={(event) => {
+                if (event.key === "Enter") void handleSave();
+              }}
+              disabled={saving}
+              className="border-gray-700 bg-gray-800 text-white"
+            />
+          </div>
+          {totalEpisodes > 1 && (
+            <div className="flex items-center gap-2">
+              <Checkbox
+                id="episode-task-all"
+                checked={allEpisodes}
+                onCheckedChange={(checked) => setAllEpisodes(checked === true)}
+                disabled={saving}
+                className="border-gray-500"
+              />
+              <Label htmlFor="episode-task-all" className="text-sm text-gray-300">
+                Apply to all {totalEpisodes} episodes, not only episode {episodeIndex}
+              </Label>
+            </div>
+          )}
+          {error && <p className="text-sm text-red-400">{error}</p>}
+        </div>
+
+        <DialogFooter>
+          <Button variant="outline" onClick={() => handleOpenChange(false)} disabled={saving}>
+            Cancel
+          </Button>
+          <Button onClick={handleSave} disabled={!canSave}>
+            {saving && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
+            {saving ? "Saving…" : "Save task"}
+          </Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
+  );
+}

frontend/src/hooks/useEpisodes.ts

@@ -1,4 +1,4 @@
-import { useEffect, useState } from "react";
+import { useCallback, useEffect, useState } from "react";
 import { useApi } from "@/contexts/ApiContext";
 import {
   EpisodeDetail,
@@ -13,6 +13,7 @@ export const useEpisodes = (repoId: string | null) => {
   const [data, setData] = useState<EpisodeListResponse | null>(null);
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
+  const [version, setVersion] = useState(0);
 
   useEffect(() => {
     setData(null);
@@ -44,11 +45,13 @@ export const useEpisodes = (repoId: string | null) => {
         if (!controller.signal.aborted) setLoading(false);
       });
     return () => controller.abort();
-  }, [baseUrl, fetchWithHeaders, repoId]);
+  }, [baseUrl, fetchWithHeaders, repoId, version]);
+
+  const refresh = useCallback(() => setVersion((v) => v + 1), []);
 
   // Effects run after render: do not let the new dataset auto-select an
   // episode from the old list before this effect clears it.
-  return { data: data?.repo_id === repoId ? data : null, loading, error };
+  return { data: data?.repo_id === repoId ? data : null, loading, error, refresh };
 };
 
 /** One episode's detail. Either arg null → idle. */

frontend/src/lib/replayApi.ts

@@ -55,3 +55,25 @@ export async function mergeLocalDatasets(
     action: "Merge datasets",
   });
 }
+
+export interface EpisodeTaskResult {
+  success: true;
+  repo_id: string;
+  task: string;
+  updated_episodes: number;
+}
+
+/** Set the task of one episode, or of every episode when `episodeIndex` is null. */
+export async function setEpisodeTask(
+  baseUrl: string,
+  fetcher: Fetcher,
+  repoId: string,
+  task: string,
+  episodeIndex: number | null,
+): Promise<EpisodeTaskResult> {
+  return apiRequest<EpisodeTaskResult>(baseUrl, fetcher, "/datasets/episode-task", {
+    method: "POST",
+    body: { repo_id: repoId, task, episode_index: episodeIndex },
+    action: "Edit task",
+  });
+}