Add VPC Direct Egress support for Cloud Run

cloud_run_client.py:
- Add vpc_network/vpc_subnet/vpc_egress fields to CloudRunClientConfig
- Configure VPC Direct Egress via run_v2.VpcAccess on job creation
- Include VPC fields in config hash
- Upload large commands to GCS when they exceed env var limits

claude_code_client.py:
- Add vpc_network/vpc_subnet/vpc_egress fields to ClaudeCodeClientConfig
- Pass VPC config through to CloudRunClientConfig

cache_manager.py:
- Fix FileBasedCacheManager self-eviction bug

README.md:
- Document VPC egress firewall setup and configuration

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: ElleNajt

safetytooling/apis/inference/cache_manager.py

@@ -1,7 +1,7 @@
 import logging
 import os
 import sys
-from collections import OrderedDict, deque
+from collections import deque
 from itertools import chain
 from pathlib import Path
 from typing import List, Tuple, Union
@@ -137,54 +137,56 @@ def save_embeddings(self, params: EmbeddingParams, response: EmbeddingResponseBa
 
 
 class FileBasedCacheManager(BaseCacheManager):
-    """File-based cache with an LRU-evicted in-memory layer."""
+    """Original file-based cache manager implementation."""
 
     def __init__(self, cache_dir: Path, num_bins: int = 20, max_mem_usage_mb: float = 5_000):
         super().__init__(cache_dir, num_bins)
-        self.in_memory_cache: OrderedDict[Path, dict] = OrderedDict()
-        self.sizes: dict[Path, float] = {}
-        self.total_usage_mb = 0.0
+        self.in_memory_cache = {}
+        self.sizes = {}  # Track the size of each cache file in memory
+        self.total_usage_mb = 0
         self.max_mem_usage_mb = max_mem_usage_mb
 
-    def _evict_lru(self):
-        """Evict the least-recently-used bin from memory."""
-        lru_key = next(iter(self.in_memory_cache))
-        del self.in_memory_cache[lru_key]
-        self.total_usage_mb -= self.sizes.pop(lru_key)
-        LOGGER.info(f"Evicted LRU entry {lru_key} from mem cache. Total usage is now {self.total_usage_mb:.1f} MB.")
+    def remove_entry(self, cache_file: Path):
+        self.in_memory_cache.pop(cache_file)
 
-    def add_entry(self, cache_file: Path, contents: dict) -> bool:
-        """Add or replace a bin in the in-memory cache, evicting LRU entries if needed.
-
-        Returns False if the single entry is larger than the entire cache limit.
-        """
-        size_mb = total_size(contents)
-
-        if self.max_mem_usage_mb is not None and size_mb > self.max_mem_usage_mb:
-            LOGGER.warning(f"Entry {cache_file} ({size_mb:.1f} MB) exceeds cache limit ({self.max_mem_usage_mb} MB).")
-            return False
-
-        # Remove old version first. This prevents self-eviction (the eviction
-        # loop below can only see OTHER bins) and prevents double-counting.
-        if cache_file in self.sizes:
-            del self.in_memory_cache[cache_file]
-            self.total_usage_mb -= self.sizes.pop(cache_file)
-
-        # Evict least-recently-used bins until there's room
         if self.max_mem_usage_mb is not None:
-            while self.in_memory_cache and self.total_usage_mb + size_mb > self.max_mem_usage_mb:
-                self._evict_lru()
+            size = self.sizes.pop(cache_file)
+            self.total_usage_mb -= size
+            LOGGER.info(f"Removed entry from mem cache. Freed {size} MB.")
 
-        # Insert at the back (most-recently-used position)
+    def add_entry(self, cache_file: Path, contents: dict):
         self.in_memory_cache[cache_file] = contents
-        self.sizes[cache_file] = size_mb
-        self.total_usage_mb += size_mb
-        return True
 
-    def touch(self, cache_file: Path):
-        """Mark a bin as recently used (moves it to the back of the LRU queue)."""
-        if cache_file in self.in_memory_cache:
-            self.in_memory_cache.move_to_end(cache_file)
+        if self.max_mem_usage_mb is not None:
+            size = total_size(contents)
+            if self.total_usage_mb + size > self.max_mem_usage_mb:
+                space_available = self.free_space_for(size)
+                if not space_available:
+                    return False
+            self.sizes[cache_file] = size
+            self.total_usage_mb += size
+
+    def free_space_for(self, needed_space_mb: float):
+        if self.max_mem_usage_mb is None:
+            return True
+
+        if needed_space_mb > self.max_mem_usage_mb:
+            LOGGER.warning(
+                f"Needed space {needed_space_mb} MB is greater than max mem usage {self.max_mem_usage_mb} MB. "
+                "This is not possible."
+            )
+            return False
+        LOGGER.info(f"Evicting entry from mem cache to free up {needed_space_mb} MB")
+        while self.total_usage_mb > self.max_mem_usage_mb - needed_space_mb:
+            # Find the entry with the smallest size
+            try:
+                smallest_entry = min(self.sizes.items(), key=lambda x: x[1])
+            except ValueError:
+                LOGGER.warning("No entries in mem cache to evict")
+                return True
+            self.remove_entry(smallest_entry[0])
+            LOGGER.info(f"Evicted entry from mem cache. Total usage is now {self.total_usage_mb} MB.")
+        return True
 
     def get_cache_file(self, prompt: Prompt, params: LLMParams) -> tuple[Path, str]:
         # Use the SHA-1 hash of the prompt for the dictionary key
@@ -209,7 +211,6 @@ def maybe_load_cache(self, prompt: Prompt, params: LLMParams):
                 self.add_entry(cache_file, contents)
         else:
             contents = self.in_memory_cache[cache_file]
-            self.touch(cache_file)
 
         data = contents.get(prompt_hash, None)
         return None if data is None else LLMCache.model_validate_json(data)

safetytooling/infra/cloud_run/README.md

@@ -177,12 +177,106 @@ client = ClaudeCodeClient(
 - **Without this, Claude could take over your entire GCP project** - don't skip this step!
 
 **What this doesn't limit:**
-- Outbound network access (Claude could exfiltrate data to external URLs)
+- Outbound network access (see Egress Firewall below)
 - Anthropic API usage (Claude could use your API key for other purposes)
 
 For the "yolo Claude" use case, the main risks are data exfiltration and API key abuse.
 Containers are ephemeral (destroyed after job), so there's no persistence risk.
 
+## Egress Firewall (Recommended)
+
+By default, containers can make outbound requests to any host. To restrict egress (e.g., only allow `api.anthropic.com` and Google APIs), use VPC Direct Egress with Cloud NGFW firewall rules.
+
+**How it works:** When `vpc_network` is set, all container traffic routes through a VPC where a Cloud NGFW firewall policy controls access by domain name (FQDN rules). This covers both IPv4 and IPv6.
+
+**Usage:**
+
+```python
+client = ClaudeCodeClient(
+    project_id="my-project",
+    gcs_bucket="my-bucket",
+    api_key_secret="anthropic-api-key-USERNAME",
+    service_account="claude-runner@my-project.iam.gserviceaccount.com",
+    vpc_network="my-egress-vpc",       # VPC with NGFW firewall policy
+    vpc_subnet="my-egress-subnet",     # Subnet in the VPC
+    vpc_egress="all-traffic",          # Route all traffic through VPC
+)
+```
+
+**One-time GCP setup:**
+
+1. **VPC + Subnet** (with Private Google Access for Google APIs):
+   ```bash
+   gcloud compute networks create egress-firewall-vpc --subnet-mode=custom
+   gcloud compute networks subnets create egress-firewall-subnet \
+       --network=egress-firewall-vpc --region=us-central1 \
+       --range=10.100.0.0/24 --enable-private-ip-google-access
+   ```
+
+2. **Cloud Router + NAT** (required for internet access from VPC):
+   ```bash
+   gcloud compute routers create egress-firewall-router \
+       --network=egress-firewall-vpc --region=us-central1
+   gcloud compute routers nats create egress-firewall-nat \
+       --router=egress-firewall-router --region=us-central1 \
+       --auto-allocate-nat-external-ips \
+       --endpoint-types=ENDPOINT_TYPE_VM,ENDPOINT_TYPE_MANAGED_PROXY_LB \
+       --nat-all-subnet-ip-ranges
+   ```
+   Note: `ENDPOINT_TYPE_MANAGED_PROXY_LB` is required — Cloud Run Direct VPC Egress uses managed proxy load balancers internally.
+
+3. **Cloud NGFW firewall policy** with FQDN rules:
+   ```bash
+   # Create policy and associate with VPC
+   gcloud compute network-firewall-policies create egress-firewall-policy --global
+   gcloud compute network-firewall-policies associations create \
+       --firewall-policy=egress-firewall-policy --network=egress-firewall-vpc --global-firewall-policy
+
+   # Allow DNS
+   gcloud compute network-firewall-policies rules create 100 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=allow \
+       --dest-ip-ranges=0.0.0.0/0 --layer4-configs=udp:53,tcp:53 --global-firewall-policy
+
+   # Allow metadata server
+   gcloud compute network-firewall-policies rules create 200 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=allow \
+       --dest-ip-ranges=169.254.169.254/32 --layer4-configs=all --global-firewall-policy
+
+   # Allow Google APIs (list each subdomain — wildcards not supported)
+   gcloud compute network-firewall-policies rules create 250 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=allow \
+       --dest-fqdns=storage.googleapis.com,oauth2.googleapis.com,www.googleapis.com,\
+   secretmanager.googleapis.com,accounts.googleapis.com,cloudresourcemanager.googleapis.com,\
+   run.googleapis.com,logging.googleapis.com,gcr.io,iamcredentials.googleapis.com \
+       --layer4-configs=tcp:443 --global-firewall-policy
+
+   # Allow Private Google Access VIPs
+   gcloud compute network-firewall-policies rules create 300 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=allow \
+       --dest-ip-ranges=199.36.153.0/24 --layer4-configs=tcp:443 --global-firewall-policy
+
+   # Allow your API provider (e.g., Anthropic)
+   gcloud compute network-firewall-policies rules create 400 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=allow \
+       --dest-fqdns=api.anthropic.com --layer4-configs=tcp:443 --global-firewall-policy
+
+   # Deny everything else (IPv4 + IPv6)
+   gcloud compute network-firewall-policies rules create 10000 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=deny \
+       --dest-ip-ranges=0.0.0.0/0 --layer4-configs=all --global-firewall-policy
+   gcloud compute network-firewall-policies rules create 10001 \
+       --firewall-policy=egress-firewall-policy --direction=EGRESS --action=deny \
+       --dest-ip-ranges=::/0 --layer4-configs=all --global-firewall-policy
+   ```
+
+**Costs:** ~$32/month for Cloud NAT gateway + ~$0.018/GB for NGFW FQDN rule evaluation. The NAT gateway runs 24/7 regardless of job activity.
+
+**Key facts:**
+- FQDN rules don't support wildcards — must list each Google API subdomain individually
+- ~20s cold start penalty on first outbound connection (NAT port allocation)
+- IPv6 is fully blocked at the VPC level (deny `::/0`)
+- Cloud NGFW Standard tier pricing applies for FQDN rule traffic
+
 ## How It Works
 
 ```
@@ -255,6 +349,9 @@ ClaudeCodeClientConfig(
     memory: str = "2Gi",   # Up to 32Gi
     skip_permissions: bool = True,  # --dangerously-skip-permissions
     image: str = DEFAULT_CLAUDE_CODE_IMAGE,  # Pre-built image with Claude Code
+    vpc_network: str = None,       # VPC for egress firewall (see Egress Firewall section)
+    vpc_subnet: str = None,        # Subnet in the VPC (required when vpc_network is set)
+    vpc_egress: str = "all-traffic",  # "all-traffic" or "private-ranges-only"
 )
 ```
 
@@ -333,6 +430,9 @@ CloudRunClientConfig(
     env: dict = {},        # Environment variables
     secrets: dict = {},    # Secret Manager secrets as env vars
     service_account: str = None,  # Restricted service account (see Security Hardening)
+    vpc_network: str = None,       # VPC for egress firewall
+    vpc_subnet: str = None,        # Subnet in the VPC
+    vpc_egress: str = None,        # "all-traffic" or "private-ranges-only"
 )
 ```
 

safetytooling/infra/cloud_run/claude_code_client.py

@@ -87,6 +87,12 @@ class ClaudeCodeClientConfig:
                         SECURITY: Use a restricted service account to limit container access.
                         See README for setup instructions.
                         Format: "name@project.iam.gserviceaccount.com"
+        vpc_network: VPC network name for Direct VPC Egress. When set with vpc_egress="all-traffic",
+                     all outbound traffic routes through the VPC where Cloud NGFW firewall policies
+                     control access. This covers both IPv4 and IPv6. Requires a Cloud NAT gateway
+                     with ENDPOINT_TYPE_MANAGED_PROXY_LB on the VPC for internet access.
+        vpc_subnet: VPC subnet name (required when vpc_network is set).
+        vpc_egress: VPC egress setting - "all-traffic" or "private-ranges-only" (default: "all-traffic").
     """
 
     project_id: str
@@ -102,6 +108,9 @@ class ClaudeCodeClientConfig:
     image: str = DEFAULT_CLAUDE_CODE_IMAGE
     api_key_secret: str | None = None
     service_account: str | None = None
+    vpc_network: str | None = None
+    vpc_subnet: str | None = None
+    vpc_egress: str = "all-traffic"
 
 
 # Instructions prepended to task when output_instructions=True
@@ -405,6 +414,9 @@ def __init__(
             env={},
             secrets=secrets,
             service_account=self.config.service_account,
+            vpc_network=self.config.vpc_network,
+            vpc_subnet=self.config.vpc_subnet,
+            vpc_egress=self.config.vpc_egress if self.config.vpc_network else None,
         )
         self._cloud_run = CloudRunClient(cloud_run_config)
 

safetytooling/infra/cloud_run/cloud_run_client.py

@@ -102,6 +102,9 @@ class CloudRunClientConfig:
     env: dict[str, str] = field(default_factory=dict)
     secrets: dict[str, str] = field(default_factory=dict)
     service_account: str | None = None
+    vpc_network: str | None = None
+    vpc_subnet: str | None = None
+    vpc_egress: str | None = None  # "all-traffic" or "private-ranges-only"
 
 
 @dataclass(frozen=True)
@@ -496,15 +499,27 @@ def _get_or_create_job(self, timeout: int) -> str:
             if self.config.service_account:
                 job.template.template.service_account = self.config.service_account
 
+            if self.config.vpc_network:
+                vpc_access = run_v2.VpcAccess(
+                    network_interfaces=[
+                        run_v2.VpcAccess.NetworkInterface(
+                            network=self.config.vpc_network,
+                            subnetwork=self.config.vpc_subnet,
+                        )
+                    ],
+                )
+                if self.config.vpc_egress == "all-traffic":
+                    vpc_access.egress = run_v2.VpcAccess.VpcEgress.ALL_TRAFFIC
+                job.template.template.vpc_access = vpc_access
+
             parent = f"projects/{self.config.project_id}/locations/{self.config.region}"
-            request = CreateJobRequest(parent=parent, job=job, job_id=job_id)
 
+            request = CreateJobRequest(parent=parent, job=job, job_id=job_id)
             try:
                 operation = self._jobs_client.create_job(request=request)
                 created_job = operation.result()
                 job_name = created_job.name
             except Exception as e:
-                # Job might already exist (from previous process/session)
                 if "already exists" in str(e).lower():
                     job_name = f"{parent}/jobs/{job_id}"
                 else:
@@ -513,6 +528,19 @@ def _get_or_create_job(self, timeout: int) -> str:
             self._job_cache[config_hash] = job_name
             return job_name
 
+    _GCS_COMMANDS_PREFIX: ClassVar[str] = "cloudrun-commands"
+    _COMMAND_SIZE_LIMIT: ClassVar[int] = 30000  # Leave headroom below 32768 env var limit
+
+    def _upload_command_to_gcs(self, command: str) -> str:
+        """Upload a large command to GCS and return its path."""
+        cmd_hash = hashlib.sha256(command.encode()).hexdigest()[:16]
+        gcs_path = f"{self._GCS_COMMANDS_PREFIX}/{cmd_hash}.sh"
+        bucket = self._storage_client.bucket(self.config.gcs_bucket)
+        blob = bucket.blob(gcs_path)
+        if not blob.exists():
+            blob.upload_from_string(command, content_type="text/plain")
+        return gcs_path
+
     def _run_job_execution(
         self,
         job_name: str,
@@ -524,7 +552,17 @@ def _run_job_execution(
         """Run an execution of an existing job with specific inputs/outputs/command.
 
         Uses RunJobRequest.Overrides to pass per-execution environment variables.
+        If the command exceeds the env var size limit, it's uploaded to GCS and
+        a small bootstrap script downloads and evals it.
         """
+        # If command is too large for an env var, stash it in GCS
+        if len(command.encode()) > self._COMMAND_SIZE_LIMIT:
+            gcs_path = self._upload_command_to_gcs(command)
+            command = (
+                f'gcloud storage cp "gs://{self.config.gcs_bucket}/{gcs_path}" /tmp/large_command.sh '
+                f"&& bash /tmp/large_command.sh"
+            )
+
         # Build env var overrides for this execution
         env_overrides = [
             run_v2.EnvVar(name="OUTPUT_GCS_PATH", value=output_gcs_path),
@@ -634,6 +672,9 @@ def _compute_config_hash(self) -> str:
             self.config.memory,
             self.config.service_account or "",
             self.config.gcs_bucket,
+            self.config.vpc_network or "",
+            self.config.vpc_subnet or "",
+            self.config.vpc_egress or "",
         ]
         # Add sorted env vars
         for k, v in sorted(self.config.env.items()):