wip websockets (#8)

* wip websockets

* wip

* wip 2

* yayy

* yay 2

.gitignore

@@ -15,4 +15,5 @@ server/secrets.json
 **/yarn.lock
 .env.prod
 .env
+
 client/graphql.schema.json

client/codegen.yml

@@ -3,6 +3,7 @@ schema: "../server/src/graphql/schema.gql"
 documents:
   - "src/graphql/queries.gql"
   - "src/graphql/fragments.gql"
+  - "src/graphql/subscriptions.gql"
 generates:
   src/graphql/generated/graphql.ts:
     plugins:

client/package.json

@@ -12,7 +12,7 @@
     "pre": "yarn lint && yarn build"
   },
   "dependencies": {
-    "@apollo/client": "^3.11.4",
+    "@apollo/client": "^3.11.8",
     "@clerk/clerk-react": "^5.4.2",
     "@googlemaps/js-api-loader": "^1.16.6",
     "@radix-ui/react-checkbox": "^1.0.4",

client/src/App.tsx

@@ -8,7 +8,6 @@ import { ApolloClientProvider } from "./providers/ApolloClientProvider";
 import MainLayout from "@/components/layouts/MainLayout";
 import AuthWrapper from "./components/custom/AuthWrapper";
 
-
 const App = () => {
   return (
     <BrowserRouter>
@@ -32,6 +31,7 @@ const Providers = (props: PropsWithChildren<NonNullable<unknown>>) => {
         <ApolloClientProvider url={import.meta.env.VITE_GRAPHQL_URL}>
           <TooltipProvider>
             <MainLayout>
+              {/* Remove the WebSocketHandler component from here */}
               {props.children}
             </MainLayout>
           </TooltipProvider>

client/src/graphql/queries.gql

@@ -4,14 +4,14 @@ query FetchUser {
   }
 }
 
-query FetchSwarm($id: ID!) {
-  fetchSwarm(id: $id) {
+query FetchSwarm($swarmId: ID!) {
+  fetchSwarm(swarmId: $swarmId) {
     ...SwarmWithData
   }
 }
 
-query FetchChat($id: ID!) {
-  fetchChat(id: $id) {
+query FetchChat($chatId: ID!) {
+  fetchChat(chatId: $chatId) {
     ...ChatWithData
   }
 }

client/src/graphql/subscriptions.gql

@@ -0,0 +1,5 @@
+subscription NewMessage($chatId: ID!) {
+  messageSent(chatId: $chatId) {
+    ...Message
+  }
+}

client/src/providers/ApolloClientProvider.tsx

@@ -1,13 +1,10 @@
 import { PropsWithChildren, useCallback, useMemo } from 'react'
-
-import { ApolloClient, ApolloProvider, createHttpLink, from, InMemoryCache } from '@apollo/client'
+import { ApolloClient, ApolloProvider, createHttpLink, from, InMemoryCache, split } from '@apollo/client'
 import { onError } from '@apollo/client/link/error'
 import { setContext } from '@apollo/client/link/context'
-import { GraphQLWsLink } from '@apollo/client/link/subscriptions';
-import { createClient } from 'graphql-ws';
-import { split } from '@apollo/client';
-import { getMainDefinition } from '@apollo/client/utilities';
-
+import { getMainDefinition } from '@apollo/client/utilities'
+import { GraphQLWsLink } from '@apollo/client/link/subscriptions'
+import { createClient } from 'graphql-ws'
 import { useAuth } from '@clerk/clerk-react'
 
 interface Props {
@@ -16,12 +13,25 @@ interface Props {
 
 export const ApolloClientProvider = ({ children, url }: PropsWithChildren<Props>) => {
   const { getToken, isSignedIn } = useAuth()
+
   const httpLink = useMemo(() => createHttpLink({ uri: url, credentials: 'include' }), [url])
 
+  const wsLink = new GraphQLWsLink(createClient({
+    url: url.replace(/^http/, 'ws'),
+    connectionParams: async () => {
+      if (isSignedIn) {
+        const token = await getToken()
+        return {
+          Authorization: `Bearer ${token}`,
+        }
+      }
+      return {}
+    }
+  }))
+
   const getConnectionContext = useCallback(async () => {
     if (isSignedIn) {
       try {
-        // Fetch a new token for each request
         const token = await getToken()
         return {
           headers: {
@@ -44,45 +54,21 @@ export const ApolloClientProvider = ({ children, url }: PropsWithChildren<Props>
     }
   })
 
-  const wsLink = new GraphQLWsLink(createClient({
-    url: url.replace('http', 'ws'),
-    connectionParams: async () => {
-      if (isSignedIn) {
-        const token = await getToken();
-        return { Authorization: `Bearer ${token}` };
-      }
-      return {};
-    },
-    retryAttempts: 5,
-    retryWait: (retries) => new Promise((resolve) => setTimeout(resolve, retries * 1000)),
-    shouldRetry: (error: unknown) => {
-      const err = error as { message: string };
-      return !err.message.includes('Unauthorized') && !err.message.includes('Invalid token');
-    },
-  }));
-
   const splitLink = split(
     ({ query }) => {
-      const definition = getMainDefinition(query);
+      const definition = getMainDefinition(query)
       return (
         definition.kind === 'OperationDefinition' &&
         definition.operation === 'subscription'
-      );
+      )
     },
     wsLink,
     httpLink,
-  );
+  )
 
   const client = new ApolloClient({
     link: from([authLink, errorLink, splitLink]),
-    uri: url,
-    cache: new InMemoryCache({
-      typePolicies: {
-        StripeQuery: { keyFields: [] },
-        NoticeQuery: { keyFields: [] },
-        User: { keyFields: [] },
-      },
-    }),
+    cache: new InMemoryCache(),
     connectToDevTools: true,
   })
 

client/src/views/Chat/ChatMessages.tsx

@@ -1,5 +1,11 @@
-import { useEffect, useRef, useState } from "react";
-import { MessageFragment, MessageRoleEnum, useFetchChatLazyQuery } from "@/graphql/generated/graphql";
+import { useEffect, useRef } from "react";
+import {
+  MessageRoleEnum,
+  NewMessageDocument,
+  useFetchChatQuery,
+  NewMessageSubscription,
+  FetchChatQuery,
+} from "@/graphql/generated/graphql";
 
 interface UserMessageProps {
   content: string;
@@ -32,29 +38,57 @@ interface ChatMessagesProps {
   isDialogMode?: boolean;
 }
 
-export function ChatMessages({selectedChatId, isDialogMode }: ChatMessagesProps) {
+export function ChatMessages({ selectedChatId, isDialogMode }: ChatMessagesProps) {
   const messagesEndRef = useRef<HTMLDivElement>(null)
-  const [messages, setMessages] = useState<MessageFragment[]>([])
-  const [fetchChat, { data: chatData }] = useFetchChatLazyQuery()
+
+  const { data, subscribeToMore } = useFetchChatQuery({
+    variables: { chatId: selectedChatId ?? '' },
+    skip: !selectedChatId,
+  });
 
   useEffect(() => {
     if (selectedChatId) {
-      fetchChat({ variables: { id: selectedChatId } })
-    }
-  }, [selectedChatId, fetchChat]);
+      const subscribeToNewMessages = subscribeToMore<NewMessageSubscription>({
+        document: NewMessageDocument,
+        variables: { chatId: selectedChatId ?? '' },
+        updateQuery: (prev: FetchChatQuery, { subscriptionData }) => {
+          if (!subscriptionData.data?.messageSent) return prev;
+          const newMessage = subscriptionData.data.messageSent;
 
-  useEffect(() => {
-    if (chatData?.fetchChat) {
-      setMessages(chatData.fetchChat.data?.messages ?? [])
+          return {
+            ...prev,
+            fetchChat: prev.fetchChat
+              ? {
+                  ...prev.fetchChat,
+                  data: prev.fetchChat.data
+                    ? {
+                        ...prev.fetchChat.data,
+                        messages: [
+                          ...(prev.fetchChat.data.messages?.filter(Boolean) || []),
+                          newMessage,
+                        ].filter((message): message is NonNullable<typeof message> => 
+                          message !== null && message !== undefined
+                        ),
+                      }
+                    : null,
+                }
+              : null,
+          };
+        },
+      });
+
+      return () => subscribeToNewMessages();
     }
-  }, [chatData])
+  }, [selectedChatId, subscribeToMore]);
 
   useEffect(() => {
     messagesEndRef.current?.scrollIntoView({ behavior: "smooth" })
-  }, [messages])
+  }, [data?.fetchChat?.data?.messages])
 
   const heightClass = isDialogMode ? "h-[calc(100vh-163px)]" : "h-[calc(100vh-200px)]"
 
+  const messages = data?.fetchChat?.data?.messages || [];
+
   return (
     <div className={`flex flex-col justify-start overflow-y-auto mt-10 gap-4 ${heightClass}`}>
       {messages.map((message) => (

client/src/views/HomePage.tsx

@@ -31,7 +31,7 @@ export default function HomePage() {
 
   useEffect(() => {
     if (selectedSwarmId) {
-      fetchSwarm({ variables: { id: selectedSwarmId } });
+      fetchSwarm({ variables: { swarmId: selectedSwarmId } });
     }
   }, [selectedSwarmId, fetchSwarm]);
 

docs/references.md

@@ -0,0 +1,8 @@
+https://www.cursor.com/blog/problems-2023
+
+https://www.cursor.com/blog/problems-2024
+
+https://www.cursor.com/blog/shadow-workspace
+
+https://honeycomb.sh/blog/swe-bench-technical-report
+

infra/docker-compose.yml

@@ -1,7 +1,8 @@
 version: '3.9'
 services:
-  postgres:
+  postgres_swarmstar:
     image: postgres:16-alpine
+    container_name: postgres_swarmstar
     ports:
       - "5434:5432"
     volumes:
@@ -10,7 +11,13 @@ services:
       POSTGRES_USER: user
       POSTGRES_PASSWORD: password
       POSTGRES_DB: swarmstarv2
+    networks:
+      - swarmstar_network
 
 volumes:
   swarmstarv2:
     driver: local
+
+networks:
+  swarmstar_network:
+    driver: bridge

infra/local-setup.md

@@ -2,4 +2,21 @@
 
    ```bash
    cd infra && docker-compose up -d
-   ```
+   ```
+
+2. Install dependencies
+
+   ```bash
+   yarn
+   ```
+
+3. Start the server
+
+   ```bash
+   yarn server
+   ```
+
+3. Start the client
+   ```bash
+   yarn client
+   ```

package.json

@@ -8,6 +8,7 @@
     "scripts": {
         "client": "cd client && yarn start",
         "server": "cd server && yarn start:dev",
+        "ws": "cd server && yarn start:dev:ws",
         "dev": "concurrently \"yarn client\" \"yarn server\"",
         "lint:python": "cd swarmstar && yarn lint:python",
         "generate": "yarn server:generate && yarn client:generate",

server/package.json

@@ -6,7 +6,7 @@
   "scripts": {
     "cli": "ts-node src/cli/index.ts",
     "generate-action-metadata": "ts-node src/cli/generateActionMetadata.ts",
-    "start:dev": "nodemon src/server.ts",
+    "start:dev": "nodemon src/graphql-server.ts",
     "start:dev:ws": "nodemon src/websocket-server.ts",
     "build": "tsc --project ./ && yarn copy-schemas",
     "copy-schemas": "npx cpx -v \"src/**/*.gql\" dist && cp package.json dist/package.json && cp yarn.lock dist/yarn.lock && mkdir -p dist/prisma && cp -r prisma/schema.prisma dist/prisma/schema.prisma",
@@ -58,6 +58,7 @@
     "dotenv": "^16.4.5",
     "eslint": "^9.9.0",
     "graphql": "^16.9.0",
+    "graphql-subscriptions": "^2.0.0",
     "inversify": "^6.0.2",
     "ioredis": "^5.4.1",
     "openai": "^4.56.0",

server/prisma/schema.prisma

@@ -81,11 +81,11 @@ model Chat {
 
   messages Message[]
 
-  actionNodeId String?     @db.Uuid
-  actionNode   ActionNode? @relation(fields: [actionNodeId], references: [id])
+  actionNodeId String     @db.Uuid
+  actionNode   ActionNode @relation(fields: [actionNodeId], references: [id])
 
-  Swarm   Swarm?  @relation(fields: [swarmId], references: [id])
-  swarmId String? @db.Uuid
+  swarm   Swarm  @relation(fields: [swarmId], references: [id])
+  swarmId String @db.Uuid
 
   createdAt DateTime @default(now())
   updatedAt DateTime @updatedAt