From 0c6d2b43b2d061c104e827bba939e09531dcdf33 Mon Sep 17 00:00:00 2001
From: Justin Brooks <justin@wandb.com>
Date: Fri, 15 Nov 2024 18:07:58 -0500
Subject: [PATCH] add more logging to pty

---
 apps/pty-proxy/src/controller/index.ts       | 11 +++++-----
 apps/pty-proxy/src/controller/user-socket.ts |  1 -
 apps/pty-proxy/src/index.ts                  |  4 +++-
 apps/pty-proxy/src/routing.ts                | 23 +++++++++++++++++++-
 apps/pty-proxy/src/server.ts                 |  2 +-
 5 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/apps/pty-proxy/src/controller/index.ts b/apps/pty-proxy/src/controller/index.ts
index 0e42fd57..3daa02ed 100644
--- a/apps/pty-proxy/src/controller/index.ts
+++ b/apps/pty-proxy/src/controller/index.ts
@@ -14,9 +14,9 @@ const onConnect = async (ws: WebSocket, request: IncomingMessage) => {
   if (agent != null) {
     logger.info("Agent connected");
     if (agent.resource?.id == null) {
-      logger.error("Agent resource id is null");
-      ws.close();
-      throw new Error("Agent resource id is null");
+      logger.error("Agent resource ID is null");
+      ws.close(1008, "Agent resource ID is null");
+      throw new Error("Agent resource ID is null");
     }
 
     agents.set(agent.resource.id, agent);
@@ -32,8 +32,9 @@ const onConnect = async (ws: WebSocket, request: IncomingMessage) => {
     return;
   }
 
-  logger.warn("Connection rejected - neither agent nor user");
-  ws.close();
+  const msg = "Neither agent nor user";
+  logger.warn(msg);
+  ws.close(1008, msg);
 };
 
 const wss = new WebSocketServer({ noServer: true });
diff --git a/apps/pty-proxy/src/controller/user-socket.ts b/apps/pty-proxy/src/controller/user-socket.ts
index c48ac107..7cd9dd09 100644
--- a/apps/pty-proxy/src/controller/user-socket.ts
+++ b/apps/pty-proxy/src/controller/user-socket.ts
@@ -62,7 +62,6 @@ export class UserSocket {
 
           logger.info("Found agent for session create", {
             targetId: data.targetId,
-            agentName: agent.resource?.name,
           });
 
           createSessionSocket(data.sessionId);
diff --git a/apps/pty-proxy/src/index.ts b/apps/pty-proxy/src/index.ts
index cd85eaf0..f3a61a18 100644
--- a/apps/pty-proxy/src/index.ts
+++ b/apps/pty-proxy/src/index.ts
@@ -1,11 +1,13 @@
 import ms from "ms";
 
+import { logger } from "@ctrlplane/logger";
+
 import { env } from "./config.js";
 import { addSocket } from "./routing.js";
 import { app } from "./server.js";
 
 const server = addSocket(app).listen(env.PORT, () => {
-  console.log(`Server is running on port ${env.PORT}`);
+  logger.info(`Server is running on port ${env.PORT}`);
 });
 
 const onCloseSignal = () => {
diff --git a/apps/pty-proxy/src/routing.ts b/apps/pty-proxy/src/routing.ts
index 684dd4e7..4abd049f 100644
--- a/apps/pty-proxy/src/routing.ts
+++ b/apps/pty-proxy/src/routing.ts
@@ -1,28 +1,49 @@
 import { createServer } from "node:http";
 import type { Express } from "express";
 
+import { logger } from "@ctrlplane/logger";
+
 import { controllerOnUpgrade } from "./controller/index.js";
 import { sessionOnUpgrade } from "./sessions/index.js";
 
 export const addSocket = (expressApp: Express) => {
   const server = createServer(expressApp);
+  logger.info("Created HTTP server for WebSocket upgrades");
+
   server.on("upgrade", (request, socket, head) => {
+    logger.debug("Received WebSocket upgrade request", {
+      url: request.url,
+      headers: request.headers,
+    });
+
     if (request.url == null) {
+      logger.warn("WebSocket upgrade rejected - no URL provided");
       socket.destroy();
       return;
     }
 
     const { pathname } = new URL(request.url, "ws://base.ws");
+    logger.debug("Processing WebSocket upgrade path", { pathname });
+
     if (pathname.startsWith("/api/v1/resources/proxy/session")) {
+      logger.info("Upgrading WebSocket connection for session proxy", {
+        pathname,
+      });
       sessionOnUpgrade(request, socket, head);
       return;
     }
 
     if (pathname.startsWith("/api/v1/resources/proxy/controller")) {
+      logger.info("Upgrading WebSocket connection for controller", {
+        pathname,
+      });
       controllerOnUpgrade(request, socket, head);
       return;
     }
-    socket.destroy();
+
+    logger.warn("WebSocket upgrade rejected - invalid path", { pathname });
+    socket.destroy(new Error("Incorrect path for WebSocket upgrade"));
   });
+
   return server;
 };
diff --git a/apps/pty-proxy/src/server.ts b/apps/pty-proxy/src/server.ts
index e14d739f..4910702c 100644
--- a/apps/pty-proxy/src/server.ts
+++ b/apps/pty-proxy/src/server.ts
@@ -17,7 +17,7 @@ app.use(express.json());
 app.use(cookieParser());
 
 // Health check endpoint
-app.get("/api/proxy/health", (_req, res) => {
+app.get("/api/v1/resources/proxy/health", (_req, res) => {
   res.status(200).json({ status: "ok" });
 });