From 4a71bb3f912600840fddcdf36d38f5fe10d2a296 Mon Sep 17 00:00:00 2001
From: Declan Brady <declan.brady@scale.com>
Date: Mon, 4 May 2026 11:00:16 -0400
Subject: [PATCH] Add custom UI

---
 agentex-ui/app/custom/page.tsx                |  30 +++
 agentex-ui/components/custom/config-panel.tsx | 214 ++++++++++++++++++
 .../components/custom/custom-chat-panel.tsx   |  75 ++++++
 .../components/custom/custom-page-root.tsx    | 118 ++++++++++
 .../components/custom/custom-prompt-input.tsx | 153 +++++++++++++
 agentex-ui/components/custom/tag-input.tsx    |  90 ++++++++
 6 files changed, 680 insertions(+)
 create mode 100644 agentex-ui/app/custom/page.tsx
 create mode 100644 agentex-ui/components/custom/config-panel.tsx
 create mode 100644 agentex-ui/components/custom/custom-chat-panel.tsx
 create mode 100644 agentex-ui/components/custom/custom-page-root.tsx
 create mode 100644 agentex-ui/components/custom/custom-prompt-input.tsx
 create mode 100644 agentex-ui/components/custom/tag-input.tsx

diff --git a/agentex-ui/app/custom/page.tsx b/agentex-ui/app/custom/page.tsx
new file mode 100644
index 0000000..87840be
--- /dev/null
+++ b/agentex-ui/app/custom/page.tsx
@@ -0,0 +1,30 @@
+import { connection } from 'next/server';
+
+import { CustomPageRoot } from '@/components/custom/custom-page-root';
+import { AgentexProvider } from '@/components/providers';
+
+export default async function CustomPage() {
+  await connection();
+
+  const sgpAppURL = process.env.NEXT_PUBLIC_SGP_APP_URL ?? '';
+  const agentexAPIBaseURL =
+    process.env.NEXT_PUBLIC_AGENTEX_API_BASE_URL ?? 'http://localhost:5003';
+
+  if (!agentexAPIBaseURL) {
+    return (
+      <div role="alert">
+        <p>Missing some configs</p>
+        <pre>{JSON.stringify({ sgpAppURL, agentexAPIBaseURL }, null, 2)}</pre>
+      </div>
+    );
+  }
+
+  return (
+    <AgentexProvider
+      sgpAppURL={sgpAppURL}
+      agentexAPIBaseURL={agentexAPIBaseURL}
+    >
+      <CustomPageRoot />
+    </AgentexProvider>
+  );
+}
diff --git a/agentex-ui/components/custom/config-panel.tsx b/agentex-ui/components/custom/config-panel.tsx
new file mode 100644
index 0000000..0b8f50a
--- /dev/null
+++ b/agentex-ui/components/custom/config-panel.tsx
@@ -0,0 +1,214 @@
+'use client';
+
+import { useCallback, useEffect } from 'react';
+
+import { RotateCcw } from 'lucide-react';
+import { useForm } from 'react-hook-form';
+
+import { TagInput } from '@/components/custom/tag-input';
+import { Button } from '@/components/ui/button';
+import {
+  Form,
+  FormControl,
+  FormField,
+  FormItem,
+  FormLabel,
+} from '@/components/ui/form';
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from '@/components/ui/select';
+import { Textarea } from '@/components/ui/textarea';
+
+export type GoldenAgentConfig = {
+  harness: string;
+  model: string;
+  system_prompt: string;
+  allowed_tools: string[];
+};
+
+const SUPPORTED_HARNESSES = [
+  'sandbox-claude',
+  'claude-code',
+  'codex',
+  'agentex',
+] as const;
+
+const DEFAULT_MODEL_BY_HARNESS: Record<string, string> = {
+  'sandbox-claude': 'claude-opus-4-6',
+  'claude-code': 'claude-opus-4-6',
+  codex: 'gpt-5.4',
+  agentex: 'gpt-5.4',
+};
+
+export const DEFAULT_CONFIG: GoldenAgentConfig = {
+  harness: 'sandbox-claude',
+  model: 'claude-opus-4-6',
+  system_prompt:
+    "You are a developer for Scale AI's SGP team. Your primary sources of truth are scaleapi/packages/egp-api-backend and scaleapi/packages/egp-annotation. Use these packages to do the development you need to do. Make sure to read the READMEs and CLAUDE.mds in those directories to get the context you need to address any issues. Your main goal is to take in the information from the prompt, use your sources of truth to come up with a course of action to address it, and then make a PR with the solution.",
+  allowed_tools: [
+    'Read',
+    'Write',
+    'Edit',
+    'Glob',
+    'Grep',
+    'Bash',
+    'WebSearch',
+    'WebFetch',
+    'List',
+  ],
+};
+
+type ConfigPanelProps = {
+  disabled: boolean;
+  onConfigChange: (config: GoldenAgentConfig) => void;
+  onReset: () => void;
+};
+
+export function ConfigPanel({
+  disabled,
+  onConfigChange,
+  onReset,
+}: ConfigPanelProps) {
+  const form = useForm<GoldenAgentConfig>({
+    defaultValues: DEFAULT_CONFIG,
+  });
+
+  const harness = form.watch('harness');
+
+  useEffect(() => {
+    const subscription = form.watch(values => {
+      onConfigChange(values as GoldenAgentConfig);
+    });
+    return () => subscription.unsubscribe();
+  }, [form, onConfigChange]);
+
+  const handleHarnessChange = useCallback(
+    (value: string) => {
+      form.setValue('harness', value);
+      form.setValue('model', DEFAULT_MODEL_BY_HARNESS[value] ?? '');
+    },
+    [form]
+  );
+
+  const handleReset = useCallback(() => {
+    form.reset(DEFAULT_CONFIG);
+    onReset();
+  }, [form, onReset]);
+
+  return (
+    <div className="flex h-full flex-col">
+      <div className="border-b px-4 py-3">
+        <div className="flex items-center justify-between">
+          <h2 className="text-sm font-semibold">Configuration</h2>
+          {disabled && (
+            <Button
+              variant="ghost"
+              size="sm"
+              onClick={handleReset}
+              className="gap-1.5 text-xs"
+            >
+              <RotateCcw className="size-3" />
+              New Chat
+            </Button>
+          )}
+        </div>
+        <p className="text-muted-foreground mt-0.5 text-xs">golden-agent</p>
+      </div>
+
+      <Form {...form}>
+        <form className="flex-1 space-y-4 overflow-y-auto p-4">
+          <FormField
+            control={form.control}
+            name="harness"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel className="text-xs">Harness</FormLabel>
+                <Select
+                  value={field.value}
+                  onValueChange={handleHarnessChange}
+                  disabled={disabled}
+                >
+                  <FormControl>
+                    <SelectTrigger className="w-full">
+                      <SelectValue />
+                    </SelectTrigger>
+                  </FormControl>
+                  <SelectContent>
+                    {SUPPORTED_HARNESSES.map(h => (
+                      <SelectItem key={h} value={h}>
+                        {h}
+                      </SelectItem>
+                    ))}
+                  </SelectContent>
+                </Select>
+              </FormItem>
+            )}
+          />
+
+          <FormField
+            control={form.control}
+            name="model"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel className="text-xs">Model</FormLabel>
+                <FormControl>
+                  <input
+                    type="text"
+                    value={field.value}
+                    onChange={field.onChange}
+                    disabled={disabled}
+                    placeholder={
+                      DEFAULT_MODEL_BY_HARNESS[harness] ?? 'Enter model name'
+                    }
+                    className="border-input placeholder:text-muted-foreground flex h-9 w-full rounded-md border bg-transparent px-3 py-1 text-sm shadow-xs outline-none focus-visible:ring-[1px] focus-visible:ring-[#756BA2] disabled:cursor-not-allowed disabled:opacity-50"
+                  />
+                </FormControl>
+              </FormItem>
+            )}
+          />
+
+          <FormField
+            control={form.control}
+            name="system_prompt"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel className="text-xs">System Prompt</FormLabel>
+                <FormControl>
+                  <Textarea
+                    value={field.value}
+                    onChange={field.onChange}
+                    disabled={disabled}
+                    placeholder="Enter system instructions..."
+                    className="min-h-32 resize-y"
+                  />
+                </FormControl>
+              </FormItem>
+            )}
+          />
+
+          <FormField
+            control={form.control}
+            name="allowed_tools"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel className="text-xs">Allowed Tools</FormLabel>
+                <FormControl>
+                  <TagInput
+                    value={field.value}
+                    onChange={field.onChange}
+                    disabled={disabled}
+                    placeholder="Add tool name and press Enter"
+                  />
+                </FormControl>
+              </FormItem>
+            )}
+          />
+        </form>
+      </Form>
+    </div>
+  );
+}
diff --git a/agentex-ui/components/custom/custom-chat-panel.tsx b/agentex-ui/components/custom/custom-chat-panel.tsx
new file mode 100644
index 0000000..f7455a2
--- /dev/null
+++ b/agentex-ui/components/custom/custom-chat-panel.tsx
@@ -0,0 +1,75 @@
+'use client';
+
+import { useRef } from 'react';
+
+import { motion } from 'framer-motion';
+
+import { useAgentexClient } from '@/components/providers';
+import { TaskMessages } from '@/components/task-messages/task-messages';
+import { CopyButton } from '@/components/ui/copy-button';
+import { useTaskSubscription } from '@/hooks/use-task-subscription';
+
+const GOLDEN_AGENT_NAME = 'golden-agent';
+
+type CustomChatPanelProps = {
+  taskId: string | null;
+};
+
+export function CustomChatPanel({ taskId }: CustomChatPanelProps) {
+  const { agentexClient } = useAgentexClient();
+  const headerRef = useRef<HTMLDivElement>(null);
+  const scrollContainerRef = useRef<HTMLDivElement>(null);
+
+  useTaskSubscription({
+    agentexClient,
+    taskId: taskId ?? '',
+    agentName: GOLDEN_AGENT_NAME,
+    enabled: !!taskId,
+  });
+
+  if (!taskId) {
+    return (
+      <div className="flex flex-1 items-center justify-center">
+        <p className="text-muted-foreground text-sm">
+          Configure the agent and send a message to start.
+        </p>
+      </div>
+    );
+  }
+
+  return (
+    <motion.div
+      key={taskId}
+      ref={scrollContainerRef}
+      className="relative flex-1 flex-col overflow-x-hidden overflow-y-auto"
+      initial={{ opacity: 0 }}
+      animate={{ opacity: 1 }}
+      transition={{ duration: 0.25, ease: 'easeInOut' }}
+    >
+      <div
+        ref={headerRef}
+        className="bg-background/80 sticky top-0 z-10 flex items-center gap-2 border-b px-4 py-2 backdrop-blur-sm"
+      >
+        <span className="text-muted-foreground text-xs font-medium">
+          golden-agent
+        </span>
+        <span className="text-muted-foreground text-xs">|</span>
+        <span className="text-muted-foreground font-mono text-xs">
+          {taskId.slice(0, 8)}...
+        </span>
+        <CopyButton content={taskId} tooltip="Copy task ID" />
+      </div>
+
+      <div className="flex w-full flex-col items-center px-4 sm:px-6 md:px-8">
+        <div className="w-full max-w-3xl">
+          <TaskMessages
+            taskId={taskId}
+            headerRef={headerRef}
+            scrollContainerRef={scrollContainerRef}
+            agentNameOverride={GOLDEN_AGENT_NAME}
+          />
+        </div>
+      </div>
+    </motion.div>
+  );
+}
diff --git a/agentex-ui/components/custom/custom-page-root.tsx b/agentex-ui/components/custom/custom-page-root.tsx
new file mode 100644
index 0000000..19f3f17
--- /dev/null
+++ b/agentex-ui/components/custom/custom-page-root.tsx
@@ -0,0 +1,118 @@
+'use client';
+
+import { useCallback, useState } from 'react';
+
+import { ToastContainer } from 'react-toastify';
+
+import { useAgentexClient } from '@/components/providers';
+import { ResizableSidebar } from '@/components/ui/resizable-sidebar';
+import { useAgents } from '@/hooks/use-agents';
+
+import { ConfigPanel, DEFAULT_CONFIG } from './config-panel';
+import { CustomChatPanel } from './custom-chat-panel';
+import { CustomPromptInput } from './custom-prompt-input';
+
+import type { GoldenAgentConfig } from './config-panel';
+
+const GOLDEN_AGENT_NAME = 'golden-agent';
+
+export function CustomPageRoot() {
+  const { agentexClient } = useAgentexClient();
+  const { data: agents = [], isLoading } = useAgents(agentexClient);
+
+  const [config, setConfig] = useState<GoldenAgentConfig>(DEFAULT_CONFIG);
+  const [taskId, setTaskId] = useState<string | null>(null);
+  const [isConfigLocked, setIsConfigLocked] = useState(false);
+  const [prompt, setPrompt] = useState('');
+
+  const goldenAgent = agents.find(a => a.name === GOLDEN_AGENT_NAME);
+
+  const handleTaskCreated = useCallback((newTaskId: string) => {
+    setTaskId(newTaskId);
+    setIsConfigLocked(true);
+  }, []);
+
+  const handleReset = useCallback(() => {
+    setTaskId(null);
+    setIsConfigLocked(false);
+    setPrompt('');
+  }, []);
+
+  const handleConfigChange = useCallback(
+    (newConfig: GoldenAgentConfig) => {
+      if (!isConfigLocked) {
+        setConfig(newConfig);
+      }
+    },
+    [isConfigLocked]
+  );
+
+  if (isLoading) {
+    return (
+      <div className="fixed inset-0 flex items-center justify-center">
+        <p className="text-muted-foreground text-sm">Loading agents...</p>
+      </div>
+    );
+  }
+
+  if (!goldenAgent) {
+    return (
+      <div className="fixed inset-0 flex items-center justify-center">
+        <div className="text-center">
+          <p className="text-sm font-medium">golden-agent not found</p>
+          <p className="text-muted-foreground mt-1 text-xs">
+            Make sure the golden-agent is deployed and has status Ready.
+          </p>
+        </div>
+      </div>
+    );
+  }
+
+  if (goldenAgent.status !== 'Ready') {
+    return (
+      <div className="fixed inset-0 flex items-center justify-center">
+        <div className="text-center">
+          <p className="text-sm font-medium">golden-agent is not ready</p>
+          <p className="text-muted-foreground mt-1 text-xs">
+            Current status: {goldenAgent.status}
+          </p>
+        </div>
+      </div>
+    );
+  }
+
+  return (
+    <>
+      <div className="fixed inset-0 flex w-full">
+        <ResizableSidebar
+          side="left"
+          storageKey="custom-config-width"
+          defaultWidth={400}
+          minWidth={300}
+          maxWidth={600}
+        >
+          <ConfigPanel
+            disabled={isConfigLocked}
+            onConfigChange={handleConfigChange}
+            onReset={handleReset}
+          />
+        </ResizableSidebar>
+
+        <main className="flex min-w-0 flex-1 flex-col">
+          <CustomChatPanel taskId={taskId} />
+
+          <div className="flex w-full justify-center border-t px-4 py-4">
+            <CustomPromptInput
+              taskId={taskId}
+              config={config}
+              prompt={prompt}
+              setPrompt={setPrompt}
+              onTaskCreated={handleTaskCreated}
+            />
+          </div>
+        </main>
+      </div>
+      <ToastContainer />
+    </>
+  );
+}
diff --git a/agentex-ui/components/custom/custom-prompt-input.tsx b/agentex-ui/components/custom/custom-prompt-input.tsx
new file mode 100644
index 0000000..593c4d0
--- /dev/null
+++ b/agentex-ui/components/custom/custom-prompt-input.tsx
@@ -0,0 +1,153 @@
+'use client';
+
+import { useCallback, useMemo, useRef } from 'react';
+
+import { ArrowUp } from 'lucide-react';
+
+import { useAgentexClient } from '@/components/providers';
+import { IconButton } from '@/components/ui/icon-button';
+import { toast } from '@/components/ui/toast';
+import { useCreateTask } from '@/hooks/use-create-task';
+import { useSendMessage } from '@/hooks/use-task-messages';
+import { useTask } from '@/hooks/use-tasks';
+import { TaskStatusEnum } from '@/lib/types';
+
+import type { GoldenAgentConfig } from './config-panel';
+import type { DataContent, TextContent } from 'agentex/resources';
+
+const GOLDEN_AGENT_NAME = 'golden-agent';
+
+type CustomPromptInputProps = {
+  taskId: string | null;
+  config: GoldenAgentConfig;
+  prompt: string;
+  setPrompt: (prompt: string) => void;
+  onTaskCreated: (taskId: string) => void;
+};
+
+export function CustomPromptInput({
+  taskId,
+  config,
+  prompt,
+  setPrompt,
+  onTaskCreated,
+}: CustomPromptInputProps) {
+  const { agentexClient } = useAgentexClient();
+  const createTaskMutation = useCreateTask({ agentexClient });
+  const sendMessageMutation = useSendMessage({ agentexClient });
+  const { data: task } = useTask({
+    agentexClient,
+    taskId: taskId ?? '',
+  });
+  const inputRef = useRef<HTMLInputElement>(null);
+
+  const isTaskTerminal = useMemo(() => {
+    if (!taskId || !task) return false;
+    return task.status != null && task.status !== TaskStatusEnum.RUNNING;
+  }, [taskId, task]);
+
+  const isDisabled = isTaskTerminal;
+
+  const handleSendPrompt = useCallback(async () => {
+    if (!prompt.trim()) {
+      toast.error('Please enter a prompt');
+      return;
+    }
+
+    const currentPrompt = prompt;
+    setPrompt('');
+
+    let currentTaskId = taskId;
+
+    if (!currentTaskId) {
+      const params: Record<string, unknown> = {
+        system_prompt: config.system_prompt || null,
+        allowed_tools: config.allowed_tools,
+        harness: config.harness,
+        model: config.model,
+      };
+
+      const createdTask = await createTaskMutation.mutateAsync({
+        agentName: GOLDEN_AGENT_NAME,
+        params,
+      });
+      currentTaskId = createdTask.id;
+      onTaskCreated(currentTaskId);
+
+      const content: DataContent = {
+        type: 'data',
+        author: 'user',
+        data: {
+          system_prompt: config.system_prompt || null,
+          allowed_tools: config.allowed_tools,
+          harness: config.harness,
+          model: config.model,
+          message: currentPrompt,
+        },
+      };
+
+      await sendMessageMutation.mutateAsync({
+        taskId: currentTaskId,
+        agentName: GOLDEN_AGENT_NAME,
+        content,
+      });
+    } else {
+      const content: TextContent = {
+        type: 'text',
+        author: 'user',
+        format: 'plain',
+        attachments: [],
+        content: currentPrompt,
+      };
+
+      await sendMessageMutation.mutateAsync({
+        taskId: currentTaskId,
+        agentName: GOLDEN_AGENT_NAME,
+        content,
+      });
+    }
+  }, [
+    prompt,
+    taskId,
+    config,
+    setPrompt,
+    createTaskMutation,
+    sendMessageMutation,
+    onTaskCreated,
+  ]);
+
+  return (
+    <div className="flex w-full max-w-3xl flex-col gap-2">
+      <div
+        className={`border-input dark:bg-input ${isDisabled ? 'bg-muted scale-90 cursor-not-allowed' : 'scale-100'} flex w-full items-center justify-between rounded-4xl border py-2 pr-2 pl-6 shadow-sm transition-transform duration-300`}
+      >
+        <input
+          ref={inputRef}
+          type="text"
+          value={prompt}
+          onChange={e => setPrompt(e.target.value)}
+          onKeyDown={e => {
+            if (e.key === 'Enter' && !isDisabled && prompt.trim()) {
+              handleSendPrompt();
+            }
+          }}
+          disabled={isDisabled}
+          placeholder={
+            isTaskTerminal
+              ? `Task ${task?.status?.toLowerCase() ?? 'ended'}`
+              : 'Enter your prompt'
+          }
+          className="mr-2 flex-1 outline-none focus:ring-0 focus:outline-none"
+          style={{ backgroundColor: 'inherit', cursor: 'inherit' }}
+        />
+        <IconButton
+          className="pointer-events-auto size-10 rounded-full"
+          onClick={handleSendPrompt}
+          disabled={isDisabled || !prompt.trim()}
+          icon={ArrowUp}
+          aria-label="Send Prompt"
+        />
+      </div>
+    </div>
+  );
+}
diff --git a/agentex-ui/components/custom/tag-input.tsx b/agentex-ui/components/custom/tag-input.tsx
new file mode 100644
index 0000000..0d9bb8b
--- /dev/null
+++ b/agentex-ui/components/custom/tag-input.tsx
@@ -0,0 +1,90 @@
+'use client';
+
+import { useCallback, useState, type KeyboardEvent } from 'react';
+
+import { X } from 'lucide-react';
+
+import { Badge } from '@/components/ui/badge';
+import { cn } from '@/lib/utils';
+
+type TagInputProps = {
+  value: string[];
+  onChange: (tags: string[]) => void;
+  disabled?: boolean;
+  placeholder?: string;
+  className?: string;
+};
+
+export function TagInput({
+  value,
+  onChange,
+  disabled = false,
+  placeholder = 'Type and press Enter',
+  className,
+}: TagInputProps) {
+  const [inputValue, setInputValue] = useState('');
+
+  const addTag = useCallback(
+    (tag: string) => {
+      const trimmed = tag.trim();
+      if (trimmed && !value.includes(trimmed)) {
+        onChange([...value, trimmed]);
+      }
+      setInputValue('');
+    },
+    [value, onChange]
+  );
+
+  const removeTag = useCallback(
+    (tagToRemove: string) => {
+      onChange(value.filter(tag => tag !== tagToRemove));
+    },
+    [value, onChange]
+  );
+
+  const handleKeyDown = useCallback(
+    (e: KeyboardEvent<HTMLInputElement>) => {
+      if (e.key === 'Enter') {
+        e.preventDefault();
+        addTag(inputValue);
+      } else if (e.key === 'Backspace' && !inputValue && value.length > 0) {
+        removeTag(value[value.length - 1]!);
+      }
+    },
+    [inputValue, value, addTag, removeTag]
+  );
+
+  return (
+    <div
+      className={cn(
+        'border-input flex min-h-9 flex-wrap items-center gap-1.5 rounded-md border px-3 py-1.5 shadow-xs',
+        disabled && 'cursor-not-allowed opacity-50',
+        className
+      )}
+    >
+      {value.map(tag => (
+        <Badge key={tag} variant="secondary" className="gap-1 pr-1">
+          {tag}
+          {!disabled && (
+            <button
+              type="button"
+              onClick={() => removeTag(tag)}
+              className="hover:text-destructive rounded-sm"
+            >
+              <X className="size-3" />
+            </button>
+          )}
+        </Badge>
+      ))}
+      <input
+        type="text"
+        value={inputValue}
+        onChange={e => setInputValue(e.target.value)}
+        onKeyDown={handleKeyDown}
+        disabled={disabled}
+        placeholder={value.length === 0 ? placeholder : ''}
+        className="placeholder:text-muted-foreground min-w-[120px] flex-1 bg-transparent text-sm outline-none disabled:cursor-not-allowed"
+      />
+    </div>
+  );
+}