From 5da417470fd1487ca38cccc0af310bde2e734de7 Mon Sep 17 00:00:00 2001
From: Justin Spahr-Summers <justin@anthropic.com>
Date: Fri, 25 Oct 2024 14:47:08 +0100
Subject: [PATCH] Add tab and approval flow for server -> client sampling

---
 client/src/App.tsx                    | 50 +++++++++++++++++++++
 client/src/components/SamplingTab.tsx | 65 +++++++++++++++++++++++++++
 2 files changed, 115 insertions(+)
 create mode 100644 client/src/components/SamplingTab.tsx

diff --git a/client/src/App.tsx b/client/src/App.tsx
index c0c5b81..15cdff8 100644
--- a/client/src/App.tsx
+++ b/client/src/App.tsx
@@ -13,8 +13,12 @@ import {
   ProgressNotificationSchema,
   ServerNotification,
   EmptyResultSchema,
+  CreateMessageRequest,
+  CreateMessageResult,
+  CreateMessageRequestSchema,
 } from "mcp-typescript/types.js";
 import { useState, useRef, useEffect } from "react";
+
 import {
   Send,
   Terminal,
@@ -23,6 +27,7 @@ import {
   MessageSquare,
   Hammer,
   Play,
+  Hash,
 } from "lucide-react";
 import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { Input } from "@/components/ui/input";
@@ -45,6 +50,7 @@ import { AnyZodObject } from "zod";
 import HistoryAndNotifications from "./components/History";
 import "./App.css";
 import PingTab from "./components/PingTab";
+import SamplingTab, { PendingRequest } from "./components/SamplingTab";
 
 const App = () => {
   const [connectionStatus, setConnectionStatus] = useState<
@@ -77,6 +83,32 @@ const App = () => {
   const [mcpClient, setMcpClient] = useState<Client | null>(null);
   const [notifications, setNotifications] = useState<ServerNotification[]>([]);
 
+  const [pendingSampleRequests, setPendingSampleRequests] = useState<
+    Array<
+      PendingRequest & {
+        resolve: (result: CreateMessageResult) => void;
+        reject: (error: Error) => void;
+      }
+    >
+  >([]);
+  const nextRequestId = useRef(0);
+
+  const handleApproveSampling = (id: number, result: CreateMessageResult) => {
+    setPendingSampleRequests((prev) => {
+      const request = prev.find((r) => r.id === id);
+      request?.resolve(result);
+      return prev.filter((r) => r.id !== id);
+    });
+  };
+
+  const handleRejectSampling = (id: number) => {
+    setPendingSampleRequests((prev) => {
+      const request = prev.find((r) => r.id === id);
+      request?.reject(new Error("Sampling request rejected"));
+      return prev.filter((r) => r.id !== id);
+    });
+  };
+
   const [selectedResource, setSelectedResource] = useState<Resource | null>(
     null,
   );
@@ -229,6 +261,15 @@ const App = () => {
         },
       );
 
+      client.setRequestHandler(CreateMessageRequestSchema, (request) => {
+        return new Promise<CreateMessageResult>((resolve, reject) => {
+          setPendingSampleRequests((prev) => [
+            ...prev,
+            { id: nextRequestId.current++, request, resolve, reject },
+          ]);
+        });
+      });
+
       setMcpClient(client);
       setConnectionStatus("connected");
     } catch (e) {
@@ -314,6 +355,10 @@ const App = () => {
                     <Bell className="w-4 h-4 mr-2" />
                     Ping
                   </TabsTrigger>
+                  <TabsTrigger value="sampling">
+                    <Hash className="w-4 h-4 mr-2" />
+                    Sampling
+                  </TabsTrigger>
                 </TabsList>
 
                 <div className="w-full">
@@ -362,6 +407,11 @@ const App = () => {
                       );
                     }}
                   />
+                  <SamplingTab
+                    pendingRequests={pendingSampleRequests}
+                    onApprove={handleApproveSampling}
+                    onReject={handleRejectSampling}
+                  />
                 </div>
               </Tabs>
             ) : (
diff --git a/client/src/components/SamplingTab.tsx b/client/src/components/SamplingTab.tsx
new file mode 100644
index 0000000..372d3fd
--- /dev/null
+++ b/client/src/components/SamplingTab.tsx
@@ -0,0 +1,65 @@
+import { TabsContent } from "@/components/ui/tabs";
+import { Alert, AlertDescription } from "@/components/ui/alert";
+import { Button } from "@/components/ui/button";
+import {
+  CreateMessageRequest,
+  CreateMessageResult,
+} from "mcp-typescript/types.js";
+
+export type PendingRequest = {
+  id: number;
+  request: CreateMessageRequest;
+};
+
+export type Props = {
+  pendingRequests: PendingRequest[];
+  onApprove: (id: number, result: CreateMessageResult) => void;
+  onReject: (id: number) => void;
+};
+
+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">
+      <Alert>
+        <AlertDescription>
+          When the server requests LLM sampling, requests will appear here for
+          approval.
+        </AlertDescription>
+      </Alert>
+      <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">
+            <pre className="bg-gray-50 p-2 rounded">
+              {JSON.stringify(request.request, null, 2)}
+            </pre>
+            <div className="flex space-x-2">
+              <Button onClick={() => handleApprove(request.id)}>Approve</Button>
+              <Button variant="outline" onClick={() => onReject(request.id)}>
+                Reject
+              </Button>
+            </div>
+          </div>
+        ))}
+        {pendingRequests.length === 0 && (
+          <p className="text-gray-500">No pending requests</p>
+        )}
+      </div>
+    </TabsContent>
+  );
+};
+
+export default SamplingTab;