From 6a6c4829cabc45cef1e972200ba38488bc743f6b Mon Sep 17 00:00:00 2001
From: Pedro Farias <feliperodriguesf62@gmail.com>
Date: Sat, 17 Jan 2026 13:08:40 -0300
Subject: [PATCH] feat: server reliability and observability improvements

---
 src/config.js |  2 +-
 src/index.js  | 20 +++++++++++++++++++-
 src/server.js | 36 ++++++++++++++++++++++++++++--------
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/src/config.js b/src/config.js
index 04f2238..3aa3109 100644
--- a/src/config.js
+++ b/src/config.js
@@ -60,7 +60,7 @@ function loadConfig() {
         if (process.env.DEBUG === 'true') config.debug = true;
 
     } catch (error) {
-        console.error('[Config] Error loading config:', error);
+        logger.error('[Config] Error loading config:', error);
     }
 }
 
diff --git a/src/index.js b/src/index.js
index ce740b5..f17e9f6 100644
--- a/src/index.js
+++ b/src/index.js
@@ -34,7 +34,7 @@ const PORT = process.env.PORT || DEFAULT_PORT;
 const HOME_DIR = os.homedir();
 const CONFIG_DIR = path.join(HOME_DIR, '.antigravity-claude-proxy');
 
-app.listen(PORT, () => {
+const server = app.listen(PORT, () => {
     // Clear console for a clean start
     console.clear();
 
@@ -105,3 +105,21 @@ ${border}    ${align4(`export ANTHROPIC_BASE_URL=http://localhost:${PORT}`)}${bo
         logger.warn('Running in DEBUG mode - verbose logs enabled');
     }
 });
+
+// Graceful shutdown
+const shutdown = () => {
+    logger.info('Shutting down server...');
+    server.close(() => {
+        logger.success('Server stopped');
+        process.exit(0);
+    });
+
+    // Force close if it takes too long
+    setTimeout(() => {
+        logger.error('Could not close connections in time, forcefully shutting down');
+        process.exit(1);
+    }, 10000);
+};
+
+process.on('SIGTERM', shutdown);
+process.on('SIGINT', shutdown);
\ No newline at end of file
diff --git a/src/server.js b/src/server.js
index 43e2ccb..b5b935f 100644
--- a/src/server.js
+++ b/src/server.js
@@ -28,6 +28,9 @@ const FALLBACK_ENABLED = args.includes('--fallback') || process.env.FALLBACK ===
 
 const app = express();
 
+// Disable x-powered-by header for security
+app.disable('x-powered-by');
+
 // Initialize account manager (will be fully initialized on first request or startup)
 const accountManager = new AccountManager();
 
@@ -150,14 +153,31 @@ function parseError(error) {
 
 // Request logging middleware
 app.use((req, res, next) => {
-    // Skip logging for event logging batch unless in debug mode
-    if (req.path === '/api/event_logging/batch') {
-        if (logger.isDebugEnabled) {
-             logger.debug(`[${req.method}] ${req.path}`);
+    const start = Date.now();
+
+    // Log response on finish
+    res.on('finish', () => {
+        const duration = Date.now() - start;
+        const status = res.statusCode;
+        const logMsg = `[${req.method}] ${req.path} ${status} (${duration}ms)`;
+
+        // Skip standard logging for event logging batch unless in debug mode
+        if (req.path === '/api/event_logging/batch') {
+            if (logger.isDebugEnabled) {
+                logger.debug(logMsg);
+            }
+        } else {
+            // Colorize status code
+            if (status >= 500) {
+                logger.error(logMsg);
+            } else if (status >= 400) {
+                logger.warn(logMsg);
+            } else {
+                logger.info(logMsg);
+            }
         }
-    } else {
-        logger.info(`[${req.method}] ${req.path}`);
-    }
+    });
+
     next();
 });
 
@@ -802,4 +822,4 @@ app.use('*', (req, res) => {
     });
 });
 
-export default app;
+export default app;
\ No newline at end of file