added gateway (#27)

* added types for ai agent execution

* added session life cycle management

* added gateway.ts

---------

Co-authored-by: vivek1504 <vivekjadhav2743@gmail.com>

Author: vivek1504

src/exec/gateway.ts

@@ -0,0 +1,66 @@
+import { runtimeStore } from "../runtime/store.js";
+import { readVsockResponse } from "../runtime/protocol.js";
+import { getVmSocket } from "../runtime/transport.js";
+import { getSession, createSession, touchSession } from "./session.js";
+import { gatewayLogger } from "../utils/logger.js";
+import crypto from "crypto";
+import type { Vm } from "../types/types.js";
+
+import { createVm } from "../runtime/vm-manager.js";
+import { Deque } from "../runtime/deque.js";
+
+export async function ensureSession(sessionId: string): Promise<Vm> {
+  let session = getSession(sessionId);
+  if (session?.state === "active") {
+    const fn = runtimeStore.functions.get(sessionId);
+    if (fn && fn.vms.length > 0) return fn.vms[0]!;
+  }
+
+  session = session || createSession(sessionId);
+
+  let fn = runtimeStore.functions.get(sessionId);
+  if (!fn) {
+    fn = {
+      functionId: sessionId,
+      queue: new Deque(),
+      vms: [],
+      readyVms: new Set(),
+      weight: 1,
+      inflightCount: 0,
+      deficit: 0,
+      pendingCreations: 0,
+    };
+    runtimeStore.functions.set(sessionId, fn);
+  }
+
+  if (fn.vms.length === 0) {
+    await createVm(sessionId, fn, "__exec__");
+  }
+
+  session.state = "active";
+  return fn.vms[0]!;
+}
+
+export async function sendSessionMessage(
+  sessionId: string,
+  message: Record<string, any>,
+  onStream?: (chunk: any) => void,
+  timeout: number = 60000,
+): Promise<any> {
+  const vm = await ensureSession(sessionId);
+  touchSession(sessionId);
+
+  const id = message.id || crypto.randomUUID();
+  const fullMessage = { ...message, id };
+
+  const socket = await getVmSocket(vm);
+  socket.write(JSON.stringify(fullMessage) + "\n");
+
+  gatewayLogger.debug(
+    { sessionId, messageType: message.type, messageId: id },
+    "message sent to VM"
+  );
+
+  const result = await readVsockResponse(socket, timeout, onStream);
+  return { ...result, messageId: id };
+}

src/runtime/vm-manager.ts

@@ -11,6 +11,7 @@ import type { RuntimeFunction, Vm } from "../types/types.js";
 export async function createVm(
   functionId: string,
   fn: RuntimeFunction,
+  snapshotId?: string,
 ): Promise<Vm> {
   const instanceId = crypto.randomBytes(4).toString("hex");
   const apiSock = `/tmp/firecracker-${functionId}-${instanceId}.socket`;
@@ -38,7 +39,7 @@ export async function createVm(
     await waitForFirecrackerApiSocket(apiSock);
 
     const client = createFcClient(apiSock);
-    await restoreVm(client, functionId, vsock);
+    await restoreVm(client, snapshotId || functionId, vsock);
 
     const vm: Vm = {
       id: instanceId,

src/utils/logger.ts

@@ -33,7 +33,8 @@ export const vmManagerLogger = logger.child({ module: "vm-manager" });
 export const transportLogger = logger.child({ module: "transport" });
 export const protocolLogger = logger.child({ module: "protocol" });
 export const cleanupLogger = logger.child({ module: "cleanup" });
-export const sessionLogger = logger.child({module: "session"})
+export const sessionLogger = logger.child({ module: "session" })
+export const gatewayLogger = logger.child({module: "gateway"})
 
 export const httpLoggerOptions: Options = {
   logger: logger.child({ module: "http" }),