create sampling response form

nathanArseneau · nathanArseneau · commit c3045495cb11 · 2025-04-05T18:24:33.000-04:00
diff --git a/client/src/components/DynamicJsonForm.tsx b/client/src/components/DynamicJsonForm.tsx
@@ -36,15 +36,17 @@ interface DynamicJsonFormProps {
   value: JsonValue;
   onChange: (value: JsonValue) => void;
   maxDepth?: number;
+  defaultIsJsonMode?: boolean;
 }
 
 const DynamicJsonForm = ({
   schema,
   value,
   onChange,
   maxDepth = 3,
+  defaultIsJsonMode = false,
 }: DynamicJsonFormProps) => {
-  const [isJsonMode, setIsJsonMode] = useState(false);
+  const [isJsonMode, setIsJsonMode] = useState(defaultIsJsonMode);
   const [jsonError, setJsonError] = useState<string>();
   // Store the raw JSON string to allow immediate feedback during typing
   // while deferring parsing until the user stops typing
@@ -370,11 +372,21 @@ const DynamicJsonForm = ({
     <div className="space-y-4">
       <div className="flex justify-end space-x-2">
         {isJsonMode && (
-          <Button variant="outline" size="sm" onClick={formatJson}>
+          <Button
+            type="button"
+            variant="outline"
+            size="sm"
+            onClick={formatJson}
+          >
             Format JSON
           </Button>
         )}
-        <Button variant="outline" size="sm" onClick={handleSwitchToFormMode}>
+        <Button
+          type="button"
+          variant="outline"
+          size="sm"
+          onClick={handleSwitchToFormMode}
+        >
           {isJsonMode ? "Switch to Form" : "Switch to JSON"}
         </Button>
       </div>
diff --git a/client/src/components/SamplingRequest.tsx b/client/src/components/SamplingRequest.tsx
@@ -0,0 +1,148 @@
+import { Button } from "@/components/ui/button";
+import JsonView from "./JsonView";
+import { useMemo, useState } from "react";
+import {
+  CreateMessageResult,
+  CreateMessageResultSchema,
+} from "@modelcontextprotocol/sdk/types.js";
+import { PendingRequest } from "./SamplingTab";
+import DynamicJsonForm, { JsonSchemaType, JsonValue } from "./DynamicJsonForm";
+
+import { useToast } from "@/hooks/use-toast";
+
+export type SamplingRequestProps = {
+  request: PendingRequest;
+  onApprove: (id: number, result: CreateMessageResult) => void;
+  onReject: (id: number) => void;
+};
+
+const SamplingRequest = ({
+  onApprove,
+  request,
+  onReject,
+}: SamplingRequestProps) => {
+  const { toast } = useToast();
+
+  const [messageResult, setMessageResult] = useState<JsonValue>({
+    model: "GPT-4o",
+    stopReason: "endTurn",
+    role: "assistant",
+    content: {
+      type: "text",
+      text: "",
+    },
+  });
+
+  const s = useMemo(() => {
+    const schema: JsonSchemaType = {
+      type: "object",
+      description: "Message result",
+      properties: {
+        model: {
+          type: "string",
+          default: "GPT-4o",
+          description: "model name",
+        },
+        stopReason: {
+          type: "string",
+          default: "endTurn",
+          description: "Stop reason",
+        },
+        role: {
+          type: "string",
+          default: "endTurn",
+          description: "Role of the model",
+        },
+        content: {
+          type: "object",
+          properties: {
+            type: {
+              type: "string",
+              default: "text",
+              description: "Type of content",
+            },
+          },
+        },
+      },
+    };
+
+    const contentType = (messageResult as any)?.content?.type;
+    if (contentType === "text" && schema.properties) {
+      schema.properties.content.properties = {
+        ...schema.properties.content.properties,
+        text: {
+          type: "string",
+          default: "",
+          description: "text content",
+        },
+      };
+    } else if (contentType === "image" && schema.properties) {
+      schema.properties.content.properties = {
+        ...schema.properties.content.properties,
+        data: {
+          type: "string",
+          default: "",
+          description: "Base64 encoded image data",
+        },
+        mimeType: {
+          type: "string",
+          default: "",
+          description: "Mime type of the image",
+        },
+      };
+    }
+
+    return schema;
+  }, [(messageResult as any)?.content?.type]);
+
+  const handleApprove = (id: number) => {
+    const validationResult = CreateMessageResultSchema.safeParse(messageResult);
+    if (!validationResult.success) {
+      toast({
+        title: "Error",
+        description: `There was an error validating the message result: ${validationResult.error.message}`,
+        variant: "destructive",
+      });
+      return;
+    }
+
+    onApprove(id, validationResult.data);
+  };
+
+  return (
+    <div
+      data-testid="sampling-request"
+      className="flex gap-4 p-4 border rounded-lg space-y-4"
+    >
+      <div className="flex-1 bg-gray-50 dark:bg-gray-800 dark:text-gray-100 p-2 rounded">
+        <JsonView data={JSON.stringify(request.request)} />
+      </div>
+      <form className="flex-1 space-y-4">
+        <div className="space-y-2">
+          <DynamicJsonForm
+            defaultIsJsonMode={true}
+            schema={s}
+            value={messageResult}
+            onChange={(newValue: JsonValue) => {
+              setMessageResult(newValue);
+            }}
+          />
+        </div>
+        <div className="flex space-x-2 mt-1">
+          <Button type="button" onClick={() => handleApprove(request.id)}>
+            Approve
+          </Button>
+          <Button
+            type="button"
+            variant="outline"
+            onClick={() => onReject(request.id)}
+          >
+            Reject
+          </Button>
+        </div>
+      </form>
+    </div>
+  );
+};
+
+export default SamplingRequest;
diff --git a/client/src/components/SamplingTab.tsx b/client/src/components/SamplingTab.tsx
@@ -1,11 +1,10 @@
 import { Alert, AlertDescription } from "@/components/ui/alert";
-import { Button } from "@/components/ui/button";
 import { TabsContent } from "@/components/ui/tabs";
 import {
   CreateMessageRequest,
   CreateMessageResult,
 } from "@modelcontextprotocol/sdk/types.js";
-import JsonView from "./JsonView";
+import SamplingRequest from "./SamplingRequest";
 
 export type PendingRequest = {
   id: number;
@@ -19,21 +18,8 @@ export type Props = {
 };
 
 const SamplingTab = ({ pendingRequests, onApprove, onReject }: Props) => {
-  const handleApprove = (id: number) => {
-    // For now, just return a stub response
-    onApprove(id, {
-      model: "stub-model",
-      stopReason: "endTurn",
-      role: "assistant",
-      content: {
-        type: "text",
-        text: "This is a stub response.",
-      },
-    });
-  };
-
   return (
-    <TabsContent value="sampling" className="h-96">
+    <TabsContent value="sampling" className="mh-96">
       <Alert>
         <AlertDescription>
           When the server requests LLM sampling, requests will appear here for
@@ -43,19 +29,12 @@ const SamplingTab = ({ pendingRequests, onApprove, onReject }: Props) => {
       <div className="mt-4 space-y-4">
         <h3 className="text-lg font-semibold">Recent Requests</h3>
         {pendingRequests.map((request) => (
-          <div key={request.id} className="p-4 border rounded-lg space-y-4">
-            <JsonView
-              className="bg-gray-50 dark:bg-gray-800 dark:text-gray-100 rounded"
-              data={JSON.stringify(request.request)}
-            />
-
-            <div className="flex space-x-2">
-              <Button onClick={() => handleApprove(request.id)}>Approve</Button>
-              <Button variant="outline" onClick={() => onReject(request.id)}>
-                Reject
-              </Button>
-            </div>
-          </div>
+          <SamplingRequest
+            key={request.id}
+            request={request}
+            onApprove={onApprove}
+            onReject={onReject}
+          />
         ))}
         {pendingRequests.length === 0 && (
           <p className="text-gray-500">No pending requests</p>
diff --git a/client/src/components/Sidebar.tsx b/client/src/components/Sidebar.tsx
@@ -460,7 +460,9 @@ const Sidebar = ({
                   </SelectTrigger>
                   <SelectContent>
                     {Object.values(LoggingLevelSchema.enum).map((level) => (
-                      <SelectItem value={level}>{level}</SelectItem>
+                      <SelectItem key={level} value={level}>
+                        {level}
+                      </SelectItem>
                     ))}
                   </SelectContent>
                 </Select>
diff --git a/client/src/components/__tests__/samplingRequest.test.tsx b/client/src/components/__tests__/samplingRequest.test.tsx
@@ -0,0 +1,73 @@
+import { render, screen, fireEvent } from "@testing-library/react";
+import SamplingRequest from "../SamplingRequest";
+import { PendingRequest } from "../SamplingTab";
+
+const mockRequest: PendingRequest = {
+  id: 1,
+  request: {
+    method: "sampling/createMessage",
+    params: {
+      messages: [
+        {
+          role: "user",
+          content: {
+            type: "text",
+            text: "What files are in the current directory?",
+          },
+        },
+      ],
+      systemPrompt: "You are a helpful file system assistant.",
+      includeContext: "thisServer",
+      maxTokens: 100,
+    },
+  },
+};
+
+describe("Form to handle sampling response", () => {
+  const mockOnApprove = jest.fn();
+  const mockOnReject = jest.fn();
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("should call onApprove with correct text content when Approve button is clicked", () => {
+    render(
+      <SamplingRequest
+        request={mockRequest}
+        onApprove={mockOnApprove}
+        onReject={mockOnReject}
+      />,
+    );
+
+    // Click the Approve button
+    fireEvent.click(screen.getByRole("button", { name: /approve/i }));
+
+    // Assert that onApprove is called with the correct arguments
+    expect(mockOnApprove).toHaveBeenCalledWith(mockRequest.id, {
+      model: "GPT-4o",
+      stopReason: "endTurn",
+      role: "assistant",
+      content: {
+        type: "text",
+        text: "",
+      },
+    });
+  });
+
+  it("should call onReject with correct request id when Reject button is clicked", () => {
+    render(
+      <SamplingRequest
+        request={mockRequest}
+        onApprove={mockOnApprove}
+        onReject={mockOnReject}
+      />,
+    );
+
+    // Click the Approve button
+    fireEvent.click(screen.getByRole("button", { name: /Reject/i }));
+
+    // Assert that onApprove is called with the correct arguments
+    expect(mockOnReject).toHaveBeenCalledWith(mockRequest.id);
+  });
+});
diff --git a/client/src/components/__tests__/samplingTab.test.tsx b/client/src/components/__tests__/samplingTab.test.tsx
@@ -0,0 +1,55 @@
+import { render, screen } from "@testing-library/react";
+import { Tabs } from "@/components/ui/tabs";
+import SamplingTab, { PendingRequest } from "../SamplingTab";
+
+describe("Sampling tab", () => {
+  const mockOnApprove = jest.fn();
+  const mockOnReject = jest.fn();
+
+  const renderSamplingTab = (pendingRequests: PendingRequest[]) =>
+    render(
+      <Tabs defaultValue="sampling">
+        <SamplingTab
+          pendingRequests={pendingRequests}
+          onApprove={mockOnApprove}
+          onReject={mockOnReject}
+        />
+      </Tabs>,
+    );
+
+  it("should render 'No pending requests' when there are no pending requests", () => {
+    renderSamplingTab([]);
+    expect(
+      screen.getByText(
+        "When the server requests LLM sampling, requests will appear here for approval.",
+      ),
+    ).toBeTruthy();
+    expect(screen.findByText("No pending requests")).toBeTruthy();
+  });
+
+  it("should render the correct number of requests", () => {
+    renderSamplingTab(
+      Array.from({ length: 5 }, (_, i) => ({
+        id: i,
+        request: {
+          method: "sampling/createMessage",
+          params: {
+            messages: [
+              {
+                role: "user",
+                content: {
+                  type: "text",
+                  text: "What files are in the current directory?",
+                },
+              },
+            ],
+            systemPrompt: "You are a helpful file system assistant.",
+            includeContext: "thisServer",
+            maxTokens: 100,
+          },
+        },
+      })),
+    );
+    expect(screen.getAllByTestId("sampling-request").length).toBe(5);
+  });
+});
