AWS Lambda Coding Agent Sandbox: Safe Tool Isolation for LLM Agents

User → Orchestrator (Bedrock / custom) → Tool router
        ├─ ReadRepoLambda     (S3 GetObject only)
        ├─ ApplyPatchLambda   (S3 read/write session prefix only)
        ├─ RunTestsLambda     (no AWS APIs; VPC deny egress)
        └─ OpenPRLambda       (GitHub App token from Secrets Manager, scoped)

// apply-patch/handler.ts
import { S3Client, GetObjectCommand, PutObjectCommand } from "@aws-sdk/client-s3";
import { createHash } from "crypto";
import { promises as fs } from "fs";
import path from "path";
import { execFile } from "child_process";
import { promisify } from "util";

const execFileAsync = promisify(execFile);
const s3 = new S3Client({});
const BUCKET = process.env.WORKSPACE_BUCKET!;
const MAX_DIFF_BYTES = 200_000; // ✅ hard cap — illustrative

type PatchEvent = {
  tenantId: string;
  sessionId: string;
  relativePath: string;
  unifiedDiff: string;
};

export const handler = async (event: PatchEvent) => {
  assertIds(event.tenantId, event.sessionId);
  if (Buffer.byteLength(event.unifiedDiff, "utf8") > MAX_DIFF_BYTES) {
    return { ok: false, error: "diff_too_large" };
  }
  if (event.relativePath.includes("..") || path.isAbsolute(event.relativePath)) {
    return { ok: false, error: "invalid_path" }; // ❌ path traversal
  }

  const prefix = `tenants/${event.tenantId}/sessions/${event.sessionId}/`;
  const key = `${prefix}${event.relativePath}`;
  const local = `/tmp/work/${event.relativePath}`;
  await fs.mkdir(path.dirname(local), { recursive: true });

  const obj = await s3.send(new GetObjectCommand({ Bucket: BUCKET, Key: key }));
  const body = await obj.Body!.transformToByteArray();
  await fs.writeFile(local, body);

  // Apply with `patch` binary — no shell metacharacters
  const diffPath = "/tmp/change.diff";
  await fs.writeFile(diffPath, event.unifiedDiff);
  try {
    // ✅ execFile argv array — never exec(`patch < ${user}`)
    await execFileAsync("patch", ["-p0", "--batch", "-i", diffPath], {
      cwd: "/tmp/work",
      timeout: 10_000,
      maxBuffer: 2_000_000,
    });
  } catch (e: any) {
    return { ok: false, error: "patch_failed", detail: String(e?.stderr || e).slice(0, 2000) };
  }

  const patched = await fs.readFile(local);
  await s3.send(
    new PutObjectCommand({
      Bucket: BUCKET,
      Key: key,
      Body: patched,
      Metadata: {
        sha256: createHash("sha256").update(patched).digest("hex"),
        session: event.sessionId,
      },
    })
  );

  return { ok: true, bytes: patched.length };
};

function assertIds(tenantId: string, sessionId: string) {
  const re = /^[a-zA-Z0-9_-]{8,64}$/;
  if (!re.test(tenantId) || !re.test(sessionId)) throw new Error("bad_ids");
}

{
  "Effect": "Allow",
  "Action": ["s3:GetObject", "s3:PutObject"],
  "Resource": "arn:aws:s3:::agent-workspaces/tenants/${aws:PrincipalTag/tenantId}/sessions/*"
}

resource "aws_security_group" "test_sandbox" {
  name   = "agent-test-sandbox"
  vpc_id = var.vpc_id

  # ❌ No egress 0.0.0.0/0
  # ✅ Optional: allow HTTPS only to VPC endpoint SGs for CodeArtifact/S3
  egress {
    from_port       = 443
    to_port         = 443
    protocol        = "tcp"
    security_groups = [var.codeartifact_vpce_sg]
  }
}

resource "aws_lambda_function" "run_tests" {
  function_name = "agent-run-tests"
  runtime       = "nodejs20.x"
  handler       = "index.handler"
  role          = aws_iam_role.tests_no_aws_apis.arn
  timeout       = 120          # illustrative wall clock
  memory_size   = 2048         # tests often need CPU ~ proportional to memory
  reserved_concurrent_executions = 20  # blast-radius cap

  vpc_config {
    subnet_ids         = var.private_subnet_ids
    security_group_ids = [aws_security_group.test_sandbox.id]
  }

  environment {
    variables = {
      # ❌ No GITHUB_TOKEN here
      MAX_STDOUT_BYTES = "100000"
    }
  }
}

// run-tests — return truncated structured output to the LLM
export const handler = async (evt: { sessionPath: string; cmd: "test" | "lint" }) => {
  const allowed = {
    test: ["npm", ["test", "--", "--runInBand"]],
    lint: ["npm", ["run", "lint"]],
  } as const;
  const [bin, args] = allowed[evt.cmd];

  try {
    const { stdout, stderr } = await execFileAsync(bin, args, {
      cwd: evt.sessionPath,
      timeout: 90_000,
      maxBuffer: 100_000,
      env: { PATH: "/usr/local/bin:/usr/bin", HOME: "/tmp", CI: "1" }, // ✅ minimal env
    });
    return {
      ok: true,
      exitCode: 0,
      stdout: stdout.slice(0, 100_000),
      stderr: stderr.slice(0, 20_000),
    };
  } catch (e: any) {
    return {
      ok: false,
      exitCode: e.code ?? 1,
      stdout: String(e.stdout || "").slice(0, 100_000),
      stderr: String(e.stderr || e.message).slice(0, 20_000),
    };
  }
};

# open_pr/handler.py — illustrative
import os, time, jwt, urllib.request, json, boto3

sm = boto3.client("secretsmanager")

def handler(event, _ctx):
    # event: {tenantId, sessionId, branch, title, body}
    secret = json.loads(sm.get_secret_value(SecretId=os.environ["GH_APP_SECRET_ARN"])["SecretString"])
    token = mint_installation_token(secret, event["installation_id"])
    # ✅ token lives only in this invocation's memory
    # ❌ never return token to the model / never log Authorization headers
    pr = create_pull_request(token, event)
    return {"ok": True, "prUrl": pr["html_url"], "number": pr["number"]}

def mint_installation_token(secret, installation_id: int) -> str:
    now = int(time.time())
    payload = {"iat": now - 60, "exp": now + 540, "iss": secret["app_id"]}
    app_jwt = jwt.encode(payload, secret["private_key"], algorithm="RS256")
    req = urllib.request.Request(
        f"https://api.github.com/app/installations/{installation_id}/access_tokens",
        method="POST",
        headers={
            "Authorization": f"Bearer {app_jwt}",
            "Accept": "application/vnd.github+json",
        },
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        return json.load(resp)["token"]

{
  "name": "run_checks",
  "description": "Run lint or unit tests in the session sandbox",
  "parameters": {
    "type": "object",
    "additionalProperties": false,
    "required": ["cmd"],
    "properties": {
      "cmd": { "type": "string", "enum": ["lint", "test"] }
    }
  }
}