From 4e745c05e2f0569b477f5d659b8e15e7338e25a1 Mon Sep 17 00:00:00 2001
From: Piotr Karwatka <piotr.karwatka@gmail.com>
Date: Tue, 13 May 2025 12:44:39 +0200
Subject: [PATCH] feat: exposing container ports

---
 examples/nodejs-project-example.ts |  3 ++-
 src/code-execution-tool.ts         |  5 ++++-
 src/container-manager.ts           | 13 +++++++++++++
 src/types.ts                       |  1 +
 4 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/examples/nodejs-project-example.ts b/examples/nodejs-project-example.ts
index 06f485a..9b2834b 100644
--- a/examples/nodejs-project-example.ts
+++ b/examples/nodejs-project-example.ts
@@ -21,7 +21,7 @@ async function main() {
 2. A package.json with necessary dependencies
 3. A README.md with setup instructions
 4. A simple route that returns "Hello World"
-5. Include a console log message "Listening on port 3000"
+5. Include a console log message "Listening on port 1234"
 
 Please format your response as follows:
 FILE: filename
@@ -83,6 +83,7 @@ Separate each file with a blank line.`
         cwd: '/project'
       },
       dependencies: ['express'],
+      ports: [1234],
       streamOutput: {
         dependencyStdout: (data) => {
           console.log('Dependency stdout:', data);
diff --git a/src/code-execution-tool.ts b/src/code-execution-tool.ts
index 16f9b7f..5798d11 100644
--- a/src/code-execution-tool.ts
+++ b/src/code-execution-tool.ts
@@ -44,6 +44,7 @@ export function createCodeExecutionTool(config: CodeExecutionToolConfig = {}) {
       entryFile: z.string().describe('Path to the entry file relative to the mounted directory'),
       cwd: z.string().describe('Working directory path that should be mounted')
     }).optional().describe('Optional configuration for running an entire application'),
+    ports: z.array(z.number()).optional().describe('Ports to expose from the container to the host'),
     streamOutput: z.object({
       stdout: z.function().args(z.string()).optional(),
       stderr: z.function().args(z.string()).optional(),
@@ -67,6 +68,7 @@ export function createCodeExecutionTool(config: CodeExecutionToolConfig = {}) {
   //    sessionId,
       environment = {},
       runApp,
+      ports,
       streamOutput
     }: z.infer<typeof codeExecutionSchema>): Promise<CodeExecutionResult> => {
       const strategy = config.defaultStrategy ?? 'per_execution';
@@ -77,7 +79,8 @@ export function createCodeExecutionTool(config: CodeExecutionToolConfig = {}) {
         containerConfig: {
           image: getImageForLanguage(language),
           environment,
-          mounts: config.mounts
+          mounts: config.mounts,
+          ports
         }
       });
 
diff --git a/src/container-manager.ts b/src/container-manager.ts
index b265048..399bb55 100644
--- a/src/container-manager.ts
+++ b/src/container-manager.ts
@@ -112,6 +112,17 @@ export class ContainerManager {
       mountsWithWorkspace.push({ type: 'directory', source: workspaceDir, target: '/workspace' });
     }
 
+    // Handle port bindings if provided
+    const exposedPorts: Record<string, {}> = {};
+    const portBindings: Record<string, Array<{ HostPort: string }>> = {};
+    if (config.ports && config.ports.length > 0) {
+      for (const p of config.ports) {
+        const key = `${p}/tcp`;
+        exposedPorts[key] = {};
+        portBindings[key] = [{ HostPort: p.toString() }];
+      }
+    }
+
     const container = await this.docker.createContainer({
       name: containerName,
       Image: config.image,
@@ -122,6 +133,7 @@ export class ContainerManager {
         CpuPeriod: 100000,
         CpuQuota: 50000,
         NetworkMode: 'bridge',
+        PortBindings: Object.keys(portBindings).length > 0 ? portBindings : undefined,
         Mounts: mountsWithWorkspace.map(mount => ({
           Target: mount.target,
           Source: mount.source,
@@ -130,6 +142,7 @@ export class ContainerManager {
         }))
       },
       WorkingDir: '/workspace',
+      ExposedPorts: Object.keys(exposedPorts).length > 0 ? exposedPorts : undefined,
       Cmd: ['sh', '-c', 'mkdir -p /workspace && tail -f /dev/null']
     });
 
diff --git a/src/types.ts b/src/types.ts
index b24fe8a..6e38e95 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -32,6 +32,7 @@ export interface ContainerConfig {
   mounts?: MountOptions[];
   environment?: Record<string, string>;
   name?: string;
+  ports?: number[]; // Host ports to publish (TCP)
 }
 
 export interface ExecutionResult {