Save endpoint with the wsConnections record instead of environment variables

Author: Perlkonig

api/sockets/authHandler.ts

@@ -52,6 +52,9 @@ export const handler = async (event: WebSocketEvent) => {
       return { statusCode: 400, body: "Missing token" };
   }
 
+  const domain = event.requestContext.domainName;
+  const stage = event.requestContext.stage;
+
   try {
     // console.log(`Verifying JWT: ${token}`);
     const payload = await verifier.verify(token);
@@ -72,6 +75,7 @@ export const handler = async (event: WebSocketEvent) => {
                 connectionId,
                 userId,
                 invisible,
+                endpoint: `https://${domain}/${stage}`,
 
                 // Optional TTL for auto-cleanup
                 ttl: Math.floor(Date.now() / 1000) + 3600,

api/sockets/messageHandler.ts

@@ -53,22 +53,13 @@ async function processRecord(record: SQSRecord) {
     return;
   }
 
-  const { verb, payload, exclude, domainName, stage } = body;
-
-  if (!domainName || !stage) {
-    console.error("Missing domainName or stage in SQS message");
-    return;
-  }
+  const { verb, payload, exclude } = body;
 
   if (!["chat", "game", "test", "connections"].includes(verb)) {
     console.warn("Unsupported verb:", verb);
     return;
   }
 
-  const apigw = new ApiGatewayManagementApiClient({
-    endpoint: `https://${domainName}/${stage}`,
-  });
-
   // Query all active connections
   const result = await ddbDocClient.send(
     new QueryCommand({
@@ -83,7 +74,10 @@ async function processRecord(record: SQSRecord) {
   const now = Math.floor(Date.now() / 1000);
 
   for (const conn of result.Items ?? []) {
-    console.log(`Processing connection: ${JSON.stringify(conn)}`);
+    const endpoint = conn.endpoint.S!;
+    const apigw = new ApiGatewayManagementApiClient({
+        endpoint,
+    });
     const connectionId = conn.sk.S!;
     const userId = conn.userId.S!;
     const ttl = conn.ttl?.N ? parseInt(conn.ttl.N) : null;

lib/wsBroadcast.ts

@@ -4,8 +4,6 @@ const REGION = "us-east-1";
 const sqsClient = new SQSClient({ region: REGION });
 
 type WsMsgBody = {
-  domainName: string;
-  stage: string;
   verb: string;
   payload?: any;
   exclude?: string[];
@@ -14,8 +12,6 @@ type WsMsgBody = {
 export async function wsBroadcast (verb: string, payload: any, exclude?: string[]): Promise<SendMessageCommandOutput> {
     // construct message
     const body: WsMsgBody = {
-        domainName: process.env.WEBSOCKET_DOMAIN!,
-        stage: process.env.WEBSOCKET_STAGE!,
         verb,
         payload,
         exclude,