wren node: Train in the panel, as a run of wren's training flow

POST/GET /flows/{name}/nodes/{id}/train submit wren_train on the node's
targets, examples and tests and read its latest run (run.for_node,
cause "train"); "Use this model" writes `model` into the draft.

frontend/src/client/schemas.gen.ts

@@ -4505,6 +4505,61 @@ export const TokenRequestSchema = {
     title: 'TokenRequest'
 } as const;
 
+export const TrainingSchema = {
+    properties: {
+        id: {
+            type: 'string',
+            title: 'Id'
+        },
+        status: {
+            type: 'string',
+            title: 'Status'
+        },
+        error: {
+            type: 'string',
+            title: 'Error',
+            default: ''
+        },
+        result: {
+            additionalProperties: true,
+            type: 'object',
+            title: 'Result',
+            default: {}
+        },
+        model: {
+            type: 'string',
+            title: 'Model',
+            default: ''
+        },
+        threshold: {
+            anyOf: [
+                {
+                    type: 'number'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Threshold'
+        },
+        test_exact_match: {
+            anyOf: [
+                {
+                    type: 'number'
+                },
+                {
+                    type: 'null'
+                }
+            ],
+            title: 'Test Exact Match'
+        }
+    },
+    type: 'object',
+    required: ['id', 'status'],
+    title: 'Training',
+    description: "A wren node's latest training run, as its panel shows it."
+} as const;
+
 export const TriggerRequestSchema = {
     properties: {
         values: {

frontend/src/client/types.gen.ts

@@ -1471,6 +1471,21 @@ export type TokenRequest = {
     labels?: Array<(string)>;
 };
 
+/**
+ * A wren node's latest training run, as its panel shows it.
+ */
+export type Training = {
+    id: string;
+    status: string;
+    error?: string;
+    result?: {
+        [key: string]: unknown;
+    };
+    model?: string;
+    threshold?: (number | null);
+    test_exact_match?: (number | null);
+};
+
 export type TriggerRequest = {
     values?: {
         [key: string]: unknown;
@@ -2112,6 +2127,20 @@ export type FlowsAcknowledgeNodeErrorData = {
 
 export type FlowsAcknowledgeNodeErrorResponse = (Message);
 
+export type FlowsTrainNodeData = {
+    name: string;
+    nodeId: string;
+};
+
+export type FlowsTrainNodeResponse = (Training);
+
+export type FlowsReadTrainingData = {
+    name: string;
+    nodeId: string;
+};
+
+export type FlowsReadTrainingResponse = ((Training | null));
+
 export type FlowsReadFlowStateData = {
     name: string;
 };

frontend/src/components/Flow/NodePanel.tsx

@@ -1,4 +1,5 @@
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"
+import { Link } from "@tanstack/react-router"
 import {
   Activity,
   CornerDownRight,
@@ -33,6 +34,8 @@ import {
   wiresAs,
 } from "@/components/Devices/ports"
 import { devicesQueryOptions } from "@/components/Devices/queries"
+import { isLive, shortId } from "@/components/Runs/queries"
+import { RunStatusBadge } from "@/components/Runs/RunStatus"
 import { Button } from "@/components/ui/button"
 import {
   Command,
@@ -56,6 +59,7 @@ import useAuth from "@/hooks/useAuth"
 import useCustomToast from "@/hooks/useCustomToast"
 import { inCodeEditor, useShortcuts } from "@/lib/shortcuts"
 import { cn } from "@/lib/utils"
+import { handleError } from "@/utils"
 import { qualify } from "./deriveEdges"
 import { DtypeValue } from "./FlowBoundary"
 import {
@@ -1593,6 +1597,119 @@ function SharingSection({
   )
 }
 
+/** A fraction as a percentage, or a dash when there is none. */
+const percent = (value: unknown) =>
+  typeof value === "number"
+    ? value.toLocaleString(undefined, {
+        style: "percent",
+        maximumFractionDigits: 1,
+      })
+    : "–"
+
+/**
+ * Train a new model for a wren node and offer it as the node's `model`. Using
+ * it is an edit like any other, so it waits in the draft until publishing.
+ */
+function TrainingSection({
+  flow,
+  node,
+  onUse,
+}: {
+  flow: string
+  node: NodeDef_Input
+  onUse: (model: string) => void
+}) {
+  const queryClient = useQueryClient()
+  // Why it would not start, such as wren's training flow not being registered.
+  const [problem, setProblem] = useState<string | null>(null)
+  const queryKey = ["wren-training", flow, node.id]
+  const { data: training } = useQuery({
+    queryKey,
+    queryFn: () => FlowsService.readTraining({ name: flow, nodeId: node.id }),
+    // Training takes minutes, so only a run still going is read again.
+    refetchInterval: (query) =>
+      query.state.data && isLive(query.state.data.status) ? 5_000 : false,
+  })
+  const train = useMutation({
+    mutationFn: () => FlowsService.trainNode({ name: flow, nodeId: node.id }),
+    onMutate: () => setProblem(null),
+    onSuccess: () => queryClient.invalidateQueries({ queryKey }),
+    onError: handleError.bind(setProblem),
+  })
+  const running = training ? isLive(training.status) : false
+  const model = training?.model
+  const inUse = Boolean(model) && node.params?.model === model
+
+  return (
+    <div className="grid gap-2" data-testid="wren-training">
+      <div className="flex items-center justify-between gap-3">
+        <div className="flex items-center gap-1">
+          <span className={SECTION}>Training</span>
+          <InfoTip label="Training">
+            Trains a new model on the targets wired here and on the examples and
+            tests, as a run of wren's training flow. The node keeps its model
+            until you use the new one and publish.
+          </InfoTip>
+        </div>
+        <Button
+          variant="outline"
+          size="sm"
+          className="h-8"
+          disabled={train.isPending || running}
+          onClick={() => train.mutate()}
+          data-testid="wren-train"
+        >
+          Train
+        </Button>
+      </div>
+      {problem ? <p className="text-xs text-destructive">{problem}</p> : null}
+      {training ? (
+        <div className="flex items-center gap-2">
+          <RunStatusBadge
+            run={{ status: training.status, status_reason: "", labels: [] }}
+          />
+          <Link
+            to="/runs/$id"
+            params={{ id: training.id }}
+            className="font-mono text-sm hover:underline"
+          >
+            {shortId(training.id)}
+          </Link>
+        </div>
+      ) : null}
+      {training?.error ? (
+        <p className="text-xs text-destructive">{training.error}</p>
+      ) : null}
+      {training && model ? (
+        <>
+          <dl className="grid gap-1 text-sm">
+            {[
+              ["Exact match", percent(training.result?.exact_match)],
+              ["Coverage", percent(training.result?.coverage)],
+              ["Threshold", training.threshold?.toFixed(2) ?? "–"],
+              ["Test exact match", percent(training.test_exact_match)],
+            ].map(([label, value]) => (
+              <div key={label} className="flex justify-between gap-3">
+                <dt className="text-muted-foreground">{label}</dt>
+                <dd className="font-mono">{value}</dd>
+              </div>
+            ))}
+          </dl>
+          <Button
+            size="sm"
+            className="h-8 justify-self-start"
+            disabled={inUse}
+            onClick={() => onUse(model)}
+            data-testid="wren-use-model"
+          >
+            {inUse ? "In use" : "Use this model"}
+          </Button>
+        </>
+      ) : null}
+    </div>
+  )
+}
+
 /** The value a fresh output returns, so the scaffold runs as written. */
 const PLACEHOLDER: Record<DType, string> = {
   float: "0.0",
@@ -2037,6 +2154,15 @@ function PanelBody({
         {nodeType?.plugin && user?.is_superuser ? (
           <StoredStateSection nodeType={nodeType.type} />
         ) : null}
+        {node.type === "wren" ? (
+          <TrainingSection
+            flow={flow}
+            node={node}
+            onUse={(model) =>
+              editNode({ ...node, params: { ...node.params, model } })
+            }
+          />
+        ) : null}
         {hasSource ? (
           <div className="grid gap-1.5">
             <div className="flex items-center gap-1">