more-compute 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

frontend/components/popups/SettingsPopup.tsx

@@ -47,7 +47,7 @@ const SettingsPopup: React.FC<{ onSettingsChange?: (settings: NotebookSettings)
   return (
     <div className="settings-container">
       <div style={{ fontSize: '12px', marginBottom: '8px', color: 'var(--mc-text-secondary, #6b7280)' }}>
-        See all themes at <a href="https://github.com/DannyMang/more-compute/blob/main/frontend/styling_README.md" target="_blank" rel="noopener noreferrer" style={{ color: 'var(--mc-link-color, #3b82f6)', textDecoration: 'underline' }}>styling_README.md</a>
+        See all settings + possible themes at <a href="https://github.com/DannyMang/more-compute/blob/main/frontend/settings.md" target="_blank" rel="noopener noreferrer" style={{ color: 'var(--mc-link-color, #3b82f6)', textDecoration: 'underline' }}>settings.md</a>
       </div>
       <textarea
         className="settings-editor"

frontend/contexts/ClaudeContext.tsx

@@ -0,0 +1,392 @@
+"use client";
+
+import React, {
+  createContext,
+  useContext,
+  useState,
+  useCallback,
+  useRef,
+  useEffect,
+} from "react";
+import type {
+  ClaudeMessage,
+  ProposedEdit,
+  ClaudeState,
+  ClaudeActions,
+  ClaudeModel,
+} from "@/types/claude";
+
+interface ClaudeContextType extends ClaudeState, ClaudeActions {}
+
+const ClaudeContext = createContext<ClaudeContextType | undefined>(undefined);
+
+export const useClaude = () => {
+  const context = useContext(ClaudeContext);
+  if (!context) {
+    throw new Error("useClaude must be used within a ClaudeProvider");
+  }
+  return context;
+};
+
+interface ClaudeProviderProps {
+  children: React.ReactNode;
+}
+
+export const ClaudeProvider: React.FC<ClaudeProviderProps> = ({ children }) => {
+  const [messages, setMessages] = useState<ClaudeMessage[]>([]);
+  const [pendingEdits, setPendingEdits] = useState<Map<number, ProposedEdit>>(
+    new Map()
+  );
+  const [isPanelOpen, setIsPanelOpen] = useState(false);
+  const [isLoading, setIsLoading] = useState(false);
+  const [isConfigured, setIsConfigured] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+  const [model, setModelState] = useState<ClaudeModel>("sonnet");
+
+  const wsRef = useRef<WebSocket | null>(null);
+  const currentMessageRef = useRef<string>("");
+  const currentMessageIdRef = useRef<string | null>(null);
+
+  // Check if Claude is configured on mount
+  useEffect(() => {
+    const checkConfig = async () => {
+      try {
+        const response = await fetch("/api/claude/config");
+        if (response.ok) {
+          const data = await response.json();
+          setIsConfigured(data.configured);
+        }
+      } catch (e) {
+        console.error("[Claude] Failed to check config:", e);
+      }
+    };
+    checkConfig();
+  }, []);
+
+  // Connect to WebSocket for Claude messages
+  useEffect(() => {
+    const connectWebSocket = () => {
+      const ws = new WebSocket("ws://127.0.0.1:3141/ws");
+
+      ws.onopen = () => {
+        // WebSocket connected
+      };
+
+      ws.onmessage = (event) => {
+        try {
+          const message = JSON.parse(event.data);
+          handleWebSocketMessage(message);
+        } catch (err) {
+          console.error("[Claude] Failed to parse message:", err);
+        }
+      };
+
+      ws.onerror = (error) => {
+        console.error("[Claude] WebSocket error:", error);
+      };
+
+      ws.onclose = () => {
+        wsRef.current = null;
+        // Reconnect after a delay
+        setTimeout(connectWebSocket, 2000);
+      };
+
+      wsRef.current = ws;
+    };
+
+    connectWebSocket();
+
+    return () => {
+      if (wsRef.current) {
+        wsRef.current.close();
+      }
+    };
+  }, []);
+
+  const handleWebSocketMessage = useCallback((message: any) => {
+    switch (message.type) {
+      case "claude_stream_start":
+        currentMessageIdRef.current = message.data.messageId;
+        currentMessageRef.current = "";
+        setIsLoading(true);
+        setError(null);
+
+        // Add a new assistant message that will be updated
+        setMessages((prev) => [
+          ...prev,
+          {
+            id: message.data.messageId,
+            role: "assistant",
+            content: "",
+            timestamp: Date.now(),
+            isStreaming: true,
+          },
+        ]);
+        break;
+
+      case "claude_stream_chunk":
+        if (message.data.messageId === currentMessageIdRef.current) {
+          currentMessageRef.current += message.data.chunk;
+
+          // Update the message with new content
+          setMessages((prev) =>
+            prev.map((msg) =>
+              msg.id === currentMessageIdRef.current
+                ? { ...msg, content: currentMessageRef.current }
+                : msg
+            )
+          );
+        }
+        break;
+
+      case "claude_stream_end":
+        if (message.data.messageId === currentMessageIdRef.current) {
+          const proposedEdits = message.data.proposedEdits || [];
+
+          // Update the message with final content and edits
+          setMessages((prev) =>
+            prev.map((msg) =>
+              msg.id === currentMessageIdRef.current
+                ? {
+                    ...msg,
+                    content: message.data.fullResponse,
+                    isStreaming: false,
+                    proposedEdits: proposedEdits.map((e: any) => ({
+                      id: e.id,
+                      cellIndex: e.cellIndex,
+                      originalCode: e.originalCode,
+                      newCode: e.newCode,
+                      explanation: e.explanation,
+                      status: "pending" as const,
+                    })),
+                  }
+                : msg
+            )
+          );
+
+          // Add pending edits to the map
+          const newPendingEdits = new Map<number, ProposedEdit>();
+          for (const edit of proposedEdits) {
+            newPendingEdits.set(edit.cellIndex, {
+              id: edit.id,
+              cellIndex: edit.cellIndex,
+              originalCode: edit.originalCode,
+              newCode: edit.newCode,
+              explanation: edit.explanation,
+              status: "pending",
+            });
+          }
+          setPendingEdits(newPendingEdits);
+
+          setIsLoading(false);
+          currentMessageIdRef.current = null;
+        }
+        break;
+
+      case "claude_error":
+        setError(message.data.error);
+        setIsLoading(false);
+        currentMessageIdRef.current = null;
+        break;
+
+      case "claude_edit_applied":
+        // Update the edit status
+        setPendingEdits((prev) => {
+          const next = new Map(prev);
+          const edit = Array.from(prev.values()).find(
+            (e) => e.id === message.data.editId
+          );
+          if (edit) {
+            next.delete(edit.cellIndex);
+          }
+          return next;
+        });
+
+        // Update message edit status
+        setMessages((prev) =>
+          prev.map((msg) => ({
+            ...msg,
+            proposedEdits: msg.proposedEdits?.map((e) =>
+              e.id === message.data.editId ? { ...e, status: "applied" as const } : e
+            ),
+          }))
+        );
+        break;
+
+      case "claude_edit_rejected":
+        // Update the edit status
+        setPendingEdits((prev) => {
+          const next = new Map(prev);
+          const edit = Array.from(prev.values()).find(
+            (e) => e.id === message.data.editId
+          );
+          if (edit) {
+            next.delete(edit.cellIndex);
+          }
+          return next;
+        });
+
+        // Update message edit status
+        setMessages((prev) =>
+          prev.map((msg) => ({
+            ...msg,
+            proposedEdits: msg.proposedEdits?.map((e) =>
+              e.id === message.data.editId ? { ...e, status: "rejected" as const } : e
+            ),
+          }))
+        );
+        break;
+    }
+  }, []);
+
+  const sendMessage = useCallback(
+    (message: string) => {
+      if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+        setError("Not connected to server");
+        return;
+      }
+
+      if (!message.trim()) {
+        return;
+      }
+
+      // Add user message to history
+      const userMessage: ClaudeMessage = {
+        id: `user-${Date.now()}`,
+        role: "user",
+        content: message,
+        timestamp: Date.now(),
+      };
+      setMessages((prev) => [...prev, userMessage]);
+
+      // Build history for context
+      const history = messages.map((msg) => ({
+        role: msg.role,
+        content: msg.content,
+      }));
+
+      // Send to backend
+      wsRef.current.send(
+        JSON.stringify({
+          type: "claude_message",
+          data: {
+            message,
+            history,
+            model,
+          },
+        })
+      );
+    },
+    [messages, model]
+  );
+
+  const applyEdit = useCallback((editId: string) => {
+    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+      setError("Not connected to server");
+      return;
+    }
+
+    // Find the edit
+    const edit = Array.from(pendingEdits.values()).find((e) => e.id === editId);
+    if (!edit) {
+      return;
+    }
+
+    wsRef.current.send(
+      JSON.stringify({
+        type: "claude_apply_edit",
+        data: {
+          editId,
+          cellIndex: edit.cellIndex,
+          newCode: edit.newCode,
+        },
+      })
+    );
+  }, [pendingEdits]);
+
+  const rejectEdit = useCallback((editId: string) => {
+    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
+      setError("Not connected to server");
+      return;
+    }
+
+    wsRef.current.send(
+      JSON.stringify({
+        type: "claude_reject_edit",
+        data: {
+          editId,
+        },
+      })
+    );
+  }, []);
+
+  const togglePanel = useCallback(() => {
+    setIsPanelOpen((prev) => !prev);
+  }, []);
+
+  const openPanel = useCallback(() => {
+    setIsPanelOpen(true);
+  }, []);
+
+  const closePanel = useCallback(() => {
+    setIsPanelOpen(false);
+  }, []);
+
+  const clearHistory = useCallback(() => {
+    setMessages([]);
+    setPendingEdits(new Map());
+    setError(null);
+  }, []);
+
+  const setModel = useCallback((newModel: ClaudeModel) => {
+    setModelState(newModel);
+  }, []);
+
+  const setApiKey = useCallback(async (apiKey: string): Promise<boolean> => {
+    try {
+      const response = await fetch("/api/claude/config", {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({ api_key: apiKey }),
+      });
+
+      if (response.ok) {
+        setIsConfigured(true);
+        setError(null);
+        return true;
+      } else {
+        const data = await response.json();
+        setError(data.detail || "Failed to save API key");
+        return false;
+      }
+    } catch (e) {
+      setError("Failed to connect to server");
+      return false;
+    }
+  }, []);
+
+  const value: ClaudeContextType = {
+    messages,
+    pendingEdits,
+    isPanelOpen,
+    isLoading,
+    isConfigured,
+    error,
+    model,
+    sendMessage,
+    applyEdit,
+    rejectEdit,
+    togglePanel,
+    openPanel,
+    closePanel,
+    clearHistory,
+    setApiKey,
+    setModel,
+  };
+
+  return (
+    <ClaudeContext.Provider value={value}>{children}</ClaudeContext.Provider>
+  );
+};

frontend/contexts/PodWebSocketContext.tsx

@@ -11,6 +11,8 @@ export interface GPUPod {
   region: string;
   costPerHour: number;
   sshConnection: string | null;
+  provider?: string;  // Provider name (e.g., "runpod", "lambda_labs", "prime_intellect")
+  gpuCount?: number;  // Number of GPUs
 }
 
 export type ConnectionState = 'provisioning' | 'deploying' | 'connected' | null;
@@ -99,7 +101,6 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
     const ws = new WebSocket(wsUrl);
 
     ws.onopen = () => {
-      console.log('[PodWebSocket] Connected');
       setIsConnected(true);
       reconnectAttemptsRef.current = 0;
     };
@@ -116,10 +117,19 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
 
           if (isFullyReady) {
             uiStatus = "running";
-          } else if (podData.status === "ACTIVE" || podData.status === "PROVISIONING" || podData.status === "PENDING") {
+          } else if (podData.status === "ACTIVE" || podData.status === "PROVISIONING" || podData.status === "PENDING" || podData.status === "STARTING") {
             uiStatus = "starting";
           }
 
+          // Check for auto-connect callback - do this OUTSIDE of setGpuPods to avoid closure issues
+          const callback = autoConnectCallbacksRef.current.get(podData.pod_id);
+          if (callback && isFullyReady) {
+            autoConnectCallbacksRef.current.delete(podData.pod_id);
+            setTimeout(() => {
+              callback(podData.pod_id);
+            }, 5000);
+          }
+
           // Update pod in the list
           setGpuPods((prevPods) => {
             const existingPodIndex = prevPods.findIndex((p) => p.id === podData.pod_id);
@@ -127,7 +137,6 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
             if (existingPodIndex >= 0) {
               // Update existing pod
               const updatedPods = [...prevPods];
-              const wasStarting = prevPods[existingPodIndex].status === "starting";
 
               updatedPods[existingPodIndex] = {
                 id: podData.pod_id,
@@ -137,22 +146,10 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
                 region: prevPods[existingPodIndex].region,
                 costPerHour: podData.price_hr,
                 sshConnection: podData.ssh_connection || null,
+                provider: podData.provider || prevPods[existingPodIndex].provider,
+                gpuCount: podData.gpu_count || prevPods[existingPodIndex].gpuCount,
               };
 
-              // Trigger auto-connect callback if pod just became ready
-              // Small delay to ensure SSH is accepting connections
-              if (wasStarting && isFullyReady) {
-                console.log(`[PodWebSocket] Pod ${podData.pod_id} is ready, scheduling auto-connect in 5s`);
-                const callback = autoConnectCallbacksRef.current.get(podData.pod_id);
-                if (callback) {
-                  setTimeout(() => {
-                    console.log(`[PodWebSocket] Triggering auto-connect for pod ${podData.pod_id}`);
-                    callback(podData.pod_id);
-                  }, 5000); // Short delay, backend now handles retries
-                  autoConnectCallbacksRef.current.delete(podData.pod_id);
-                }
-              }
-
               return updatedPods;
             } else {
               // Add new pod if not in list
@@ -166,6 +163,8 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
                   region: "Unknown",
                   costPerHour: podData.price_hr,
                   sshConnection: podData.ssh_connection || null,
+                  provider: podData.provider,
+                  gpuCount: podData.gpu_count,
                 },
               ];
             }
@@ -182,21 +181,17 @@ export const PodWebSocketProvider: React.FC<PodWebSocketProviderProps> = ({ chil
     };
 
     ws.onclose = () => {
-      console.log('[PodWebSocket] Disconnected');
       setIsConnected(false);
       wsRef.current = null;
 
       // Attempt to reconnect with exponential backoff
       if (reconnectAttemptsRef.current < maxReconnectAttempts) {
         const delay = Math.min(1000 * Math.pow(2, reconnectAttemptsRef.current), 30000);
-        console.log(`[PodWebSocket] Reconnecting in ${delay}ms (attempt ${reconnectAttemptsRef.current + 1}/${maxReconnectAttempts})`);
 
         reconnectTimeoutRef.current = setTimeout(() => {
           reconnectAttemptsRef.current++;
           connectWebSocket();
         }, delay);
-      } else {
-        console.log('[PodWebSocket] Max reconnect attempts reached');
       }
     };