some last changes

node_modules/.vite/results.json

@@ -1 +1 @@
-{"version":"3.0.9","results":[[":src/store/__tests__/openApiStore.test.ts",{"duration":13.622202999999956,"failed":false}],[":integration/__tests__/proxy.test.ts",{"duration":1053.5150509999999,"failed":false}],[":src/middleware/__tests__/harRecorder.test.ts",{"duration":6.968161000000009,"failed":false}],[":src/__tests__/cli.test.ts",{"duration":4.447293999999999,"failed":false}]]}
+{"version":"3.0.9","results":[[":src/store/__tests__/openApiStore.test.ts",{"duration":14.241412999999994,"failed":false}],[":integration/__tests__/proxy.test.ts",{"duration":1061.098798,"failed":false}],[":src/middleware/__tests__/harRecorder.test.ts",{"duration":7.512958000000026,"failed":false}],[":src/__tests__/cli.test.ts",{"duration":3.995672000000013,"failed":false}]]}

src/cli.ts

@@ -6,7 +6,8 @@ import { startServers } from './server.js';
 
 const program = new Command();
 
-console.log('Starting Arbiter...');
+// Use console.info for startup messages
+console.info('Starting Arbiter...');
 
 program
   .name('arbiter')
@@ -24,11 +25,11 @@ const options = program.opts();
 
 // Start the servers
 startServers({
-  target: options.target,
-  proxyPort: parseInt(options.port),
-  docsPort: parseInt(options.docsPort),
-  verbose: options.verbose,
-}).catch((error) => {
-  console.error(chalk.red('Failed to start servers:'), error);
+  target: options.target as string,
+  proxyPort: parseInt(options.port as string, 10),
+  docsPort: parseInt(options.docsPort as string, 10),
+  verbose: options.verbose as boolean,
+}).catch((error: Error) => {
+  console.error(chalk.red('Failed to start servers:'), error.message);
   process.exit(1);
 });

src/middleware/apiDocGenerator.ts

@@ -15,23 +15,31 @@ export function apiDocGenerator(store: OpenAPIStore): (c: Context, next: Next) =
     const endTime = Date.now();
     const responseTime = endTime - startTime;
 
-    // Get request details
-    const url = new URL(c.req.url);
-    const queryParams: Record<string, string> = {};
-    for (const [key, value] of url.searchParams.entries()) {
-      queryParams[key] = value;
-    }
-
-    // Get request headers
-    const requestHeaders: Record<string, string> = {};
-    for (const [key, value] of Object.entries(c.req.header())) {
-      if (typeof value === 'string') {
-        requestHeaders[key] = value;
-      }
-    }
-
-    // Record the endpoint in OpenAPI format
+    // Record the request/response in OpenAPI format
     try {
+      const url = new URL(c.req.url);
+      const queryParams: Record<string, string> = {};
+      for (const [key, value] of url.searchParams.entries()) {
+        queryParams[key] = value;
+      }
+
+      // Get request headers
+      const requestHeaders: Record<string, string> = {};
+      for (const [key, value] of Object.entries(c.req.header())) {
+        if (typeof value === 'string') {
+          requestHeaders[key] = value;
+        }
+      }
+
+      // Get response headers
+      const responseHeaders: Record<string, string> = {};
+      if (c.res) {
+        for (const [key, value] of c.res.headers.entries()) {
+          responseHeaders[key] = value;
+        }
+      }
+
+      // Record the endpoint
       store.recordEndpoint(
         c.req.path,
         c.req.method.toLowerCase(),
@@ -39,15 +47,17 @@ export function apiDocGenerator(store: OpenAPIStore): (c: Context, next: Next) =
           query: queryParams,
           headers: requestHeaders,
           contentType: c.req.header('content-type') || 'application/json',
+          body: undefined, // We'll need to handle body parsing if needed
         },
         {
-          status: c.res.status,
-          contentType: c.res.headers.get('content-type') || 'application/json',
-          headers: Object.fromEntries(c.res.headers.entries()),
+          status: c.res?.status || 500,
+          headers: responseHeaders,
+          contentType: c.res?.headers.get('content-type') || 'application/json',
+          body: c.res ? await c.res.clone().text() : '',
         }
       );
     } catch (error) {
-      console.error('Error recording endpoint:', error);
+      console.error('Error recording OpenAPI entry:', error);
     }
   };
 }

src/middleware/harRecorder.ts

@@ -1,6 +1,5 @@
 import type { Context, Next } from 'hono';
 import type { OpenAPIStore } from '../store/openApiStore.js';
-import { SecurityInfo } from '../store/openApiStore.js';
 
 export function harRecorder(store: OpenAPIStore): (c: Context, next: Next) => Promise<void> {
   return async (c: Context, next: Next): Promise<void> => {

src/server.ts

@@ -66,12 +66,10 @@ export async function startServers(
 
   // Configure proxy server middleware
   proxyApp.use('*', async (c, next) => {
-    await harRecorder(openApiStore)(c);
-    await next();
+    await harRecorder(openApiStore)(c, next);
   });
   proxyApp.use('*', async (c, next) => {
-    await apiDocGenerator(openApiStore)(c);
-    await next();
+    await apiDocGenerator(openApiStore)(c, next);
   });
 
   // Documentation endpoints
@@ -392,6 +390,24 @@ export async function startServers(
   console.log(chalk.cyan(`  OpenAPI YAML: http://localhost:${availableDocsPort}/openapi.yaml`));
   console.log('\n' + chalk.yellow('Press Ctrl+C to stop'));
 
+  // Handle graceful shutdown
+  const shutdown = async (signal: string): Promise<void> => {
+    console.info(`Received ${signal}, shutting down...`);
+    await Promise.all([
+      proxyServer.close(),
+      docsServer.close(),
+    ]);
+    process.exit(0);
+  };
+
+  process.on('SIGTERM', () => {
+    void shutdown('SIGTERM');
+  });
+
+  process.on('SIGINT', () => {
+    void shutdown('SIGINT');
+  });
+
   return { proxyServer, docsServer };
 }