feat: add Google Vertex AI support with service account auth

2026-03-31 18:00:40 +02:00 · 2026-03-03 02:42:46 +05:30
parent 3ec491b30b
commit 6a76df2f4c
7 changed files with 141 additions and 5 deletions
--- a/.env.example
+++ b/.env.example
@@ -48,6 +48,23 @@ ANTHROPIC_API_KEY=your-api-key-here
 # AWS_REGION=us-east-1
 # AWS_BEARER_TOKEN_BEDROCK=your-bearer-token

+# =============================================================================
+# OPTION 4: Google Vertex AI
+# =============================================================================
+# https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-partner-models
+# Requires a GCP service account with roles/aiplatform.user.
+# Download the SA key JSON from GCP Console (IAM > Service Accounts > Keys).
+# Requires the model tier overrides above to be set with Vertex AI model IDs.
+# Example Vertex AI model IDs:
+#   ANTHROPIC_SMALL_MODEL=claude-haiku-4-5@20251001
+#   ANTHROPIC_MEDIUM_MODEL=claude-sonnet-4-6
+#   ANTHROPIC_LARGE_MODEL=claude-opus-4-6
+
+# CLAUDE_CODE_USE_VERTEX=1
+# CLOUD_ML_REGION=us-east5
+# ANTHROPIC_VERTEX_PROJECT_ID=your-gcp-project-id
+# GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json
+
 # =============================================================================
 # Available Models
 # =============================================================================
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,6 @@
 node_modules/
 .env
 audit-logs/
+credentials/
 dist/
 repos/
--- a/README.md
+++ b/README.md
@@ -88,6 +88,7 @@ Shannon is available in two editions:
  - [Workspaces and Resuming](#workspaces-and-resuming)
  - [Configuration (Optional)](#configuration-optional)
  - [AWS Bedrock](#aws-bedrock)
+  - [Google Vertex AI](#google-vertex-ai)
  - [[EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)](#experimental---unsupported-router-mode-alternative-providers)
  - [Output and Results](#output-and-results)
 - [Sample Reports](#-sample-reports)
@@ -109,6 +110,7 @@ Shannon is available in two editions:
  - **Anthropic API key** (recommended) - Get from [Anthropic Console](https://console.anthropic.com)
  - **Claude Code OAuth token**
  - **AWS Bedrock** - Route through Amazon Bedrock with AWS credentials (see [AWS Bedrock](#aws-bedrock))
+  - **Google Vertex AI** - Route through Google Cloud Vertex AI (see [Google Vertex AI](#google-vertex-ai))
  - **[EXPERIMENTAL - UNSUPPORTED] Alternative providers via Router Mode** - OpenAI or Google Gemini via OpenRouter (see [Router Mode](#experimental---unsupported-router-mode-alternative-providers))

 ### Quick Start
@@ -377,6 +379,43 @@ ANTHROPIC_LARGE_MODEL=us.anthropic.claude-opus-4-6

 Shannon uses three model tiers: **small** (`claude-haiku-4-5-20251001`) for summarization, **medium** (`claude-sonnet-4-6`) for security analysis, and **large** (`claude-opus-4-6`) for deep reasoning. Set `ANTHROPIC_SMALL_MODEL`, `ANTHROPIC_MEDIUM_MODEL`, and `ANTHROPIC_LARGE_MODEL` to the Bedrock model IDs for your region.

+### Google Vertex AI
+
+Shannon also supports [Google Vertex AI](https://cloud.google.com/vertex-ai) instead of using an Anthropic API key.
+
+#### Quick Setup
+
+1. Create a service account with the `roles/aiplatform.user` role in the [GCP Console](https://console.cloud.google.com/iam-admin/serviceaccounts), then download a JSON key file.
+
+2. Place the key file in the `./credentials/` directory:
+
+```bash
+mkdir -p ./credentials
+cp /path/to/your-sa-key.json ./credentials/gcp-sa-key.json
+```
+
+3. Add your GCP configuration to `.env`:
+
+```bash
+CLAUDE_CODE_USE_VERTEX=1
+CLOUD_ML_REGION=us-east5
+ANTHROPIC_VERTEX_PROJECT_ID=your-gcp-project-id
+GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json
+
+# Set models with Vertex AI model IDs
+ANTHROPIC_SMALL_MODEL=claude-haiku-4-5@20251001
+ANTHROPIC_MEDIUM_MODEL=claude-sonnet-4-6
+ANTHROPIC_LARGE_MODEL=claude-opus-4-6
+```
+
+4. Run Shannon as usual:
+
+```bash
+./shannon start URL=https://example.com REPO=repo-name
+```
+
+Set `CLOUD_ML_REGION=global` for global endpoints, or a specific region like `us-east5`. Some models may not be available on global endpoints — see the [Vertex AI Model Garden](https://console.cloud.google.com/vertex-ai/model-garden) for region availability.
+
 ### [EXPERIMENTAL - UNSUPPORTED] Router Mode (Alternative Providers)

 Shannon can experimentally route requests through alternative AI providers using claude-code-router. This mode is not officially supported and is intended primarily for:
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -27,6 +27,10 @@ services:
      - CLAUDE_CODE_USE_BEDROCK=${CLAUDE_CODE_USE_BEDROCK:-}
      - AWS_REGION=${AWS_REGION:-}
      - AWS_BEARER_TOKEN_BEDROCK=${AWS_BEARER_TOKEN_BEDROCK:-}
+      - CLAUDE_CODE_USE_VERTEX=${CLAUDE_CODE_USE_VERTEX:-}
+      - CLOUD_ML_REGION=${CLOUD_ML_REGION:-}
+      - ANTHROPIC_VERTEX_PROJECT_ID=${ANTHROPIC_VERTEX_PROJECT_ID:-}
+      - GOOGLE_APPLICATION_CREDENTIALS=${GOOGLE_APPLICATION_CREDENTIALS:-}
      - ANTHROPIC_SMALL_MODEL=${ANTHROPIC_SMALL_MODEL:-}
      - ANTHROPIC_MEDIUM_MODEL=${ANTHROPIC_MEDIUM_MODEL:-}
      - ANTHROPIC_LARGE_MODEL=${ANTHROPIC_LARGE_MODEL:-}
@@ -39,6 +43,7 @@ services:
      - ./prompts:/app/prompts
      - ./audit-logs:/app/audit-logs
      - ${OUTPUT_DIR:-./audit-logs}:/app/output
+      - ./credentials:/app/credentials:ro
      - ./repos:/repos
      - ${BENCHMARKS_BASE:-.}:/benchmarks
    shm_size: 2gb
--- a/27
+++ b/27
@@ -156,12 +156,37 @@ cmd_start() {
        echo "ERROR: Bedrock mode requires the following env vars in .env:$MISSING"
        exit 1
      fi
+    elif [ "$CLAUDE_CODE_USE_VERTEX" = "1" ]; then
+      # Vertex AI mode — validate required GCP credentials
+      MISSING=""
+      [ -z "$CLOUD_ML_REGION" ] && MISSING="$MISSING CLOUD_ML_REGION"
+      [ -z "$ANTHROPIC_VERTEX_PROJECT_ID" ] && MISSING="$MISSING ANTHROPIC_VERTEX_PROJECT_ID"
+      [ -z "$ANTHROPIC_SMALL_MODEL" ] && MISSING="$MISSING ANTHROPIC_SMALL_MODEL"
+      [ -z "$ANTHROPIC_MEDIUM_MODEL" ] && MISSING="$MISSING ANTHROPIC_MEDIUM_MODEL"
+      [ -z "$ANTHROPIC_LARGE_MODEL" ] && MISSING="$MISSING ANTHROPIC_LARGE_MODEL"
+      if [ -n "$MISSING" ]; then
+        echo "ERROR: Vertex AI mode requires the following env vars in .env:$MISSING"
+        exit 1
+      fi
+      # Validate service account key file (must be inside ./credentials/ for Docker mount)
+      if [ -z "$GOOGLE_APPLICATION_CREDENTIALS" ]; then
+        echo "ERROR: Vertex AI mode requires GOOGLE_APPLICATION_CREDENTIALS in .env"
+        echo "       Place your service account key in ./credentials/ and set:"
+        echo "       GOOGLE_APPLICATION_CREDENTIALS=./credentials/gcp-sa-key.json"
+        exit 1
+      fi
+      if [ ! -f "$GOOGLE_APPLICATION_CREDENTIALS" ]; then
+        echo "ERROR: Service account key file not found: $GOOGLE_APPLICATION_CREDENTIALS"
+        echo "       Download a key from the GCP Console (IAM > Service Accounts > Keys)"
+        exit 1
+      fi
    elif [ "$ROUTER" = "true" ] && { [ -n "$OPENAI_API_KEY" ] || [ -n "$OPENROUTER_API_KEY" ]; }; then
      # Router mode with alternative provider - set a placeholder for SDK init
      export ANTHROPIC_API_KEY="router-mode"
    else
      echo "ERROR: Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env"
      echo "       (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock,"
+      echo "        CLAUDE_CODE_USE_VERTEX=1 for Google Vertex AI,"
      echo "        or ROUTER=true with OPENAI_API_KEY or OPENROUTER_API_KEY)"
      exit 1
    fi
@@ -222,7 +247,7 @@ cmd_start() {
  fi

  # Ensure audit-logs directory exists with write permissions for container user (UID 1001)
-  mkdir -p ./audit-logs
+  mkdir -p ./audit-logs ./credentials
  chmod 777 ./audit-logs

  # Ensure repo deliverables directory is writable by container user (UID 1001)
--- a/src/ai/claude-executor.ts
+++ b/src/ai/claude-executor.ts
@@ -235,6 +235,10 @@ export async function runClaudePrompt(
    'CLAUDE_CODE_USE_BEDROCK',
    'AWS_REGION',
    'AWS_BEARER_TOKEN_BEDROCK',
+    'CLAUDE_CODE_USE_VERTEX',
+    'CLOUD_ML_REGION',
+    'ANTHROPIC_VERTEX_PROJECT_ID',
+    'GOOGLE_APPLICATION_CREDENTIALS',
    'ANTHROPIC_SMALL_MODEL',
    'ANTHROPIC_MEDIUM_MODEL',
    'ANTHROPIC_LARGE_MODEL',
--- a/src/services/preflight.ts
+++ b/src/services/preflight.ts
@@ -14,7 +14,7 @@
 * Checks run sequentially, cheapest first:
 * 1. Repository path exists and contains .git
 * 2. Config file parses and validates (if provided)
- * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, or router mode)
+ * 3. Credentials validate via Claude Agent SDK query (API key, OAuth, Bedrock, Vertex AI, or router mode)
 */

 import fs from 'fs/promises';
@@ -185,11 +185,56 @@ async function validateCredentials(
    return ok(undefined);
  }

-  // 3. Check that at least one credential is present
+  // 3. Vertex AI mode — validate required GCP credentials are present
+  if (process.env.CLAUDE_CODE_USE_VERTEX === '1') {
+    const required = ['CLOUD_ML_REGION', 'ANTHROPIC_VERTEX_PROJECT_ID', 'ANTHROPIC_SMALL_MODEL', 'ANTHROPIC_MEDIUM_MODEL', 'ANTHROPIC_LARGE_MODEL'];
+    const missing = required.filter(v => !process.env[v]);
+    if (missing.length > 0) {
+      return err(
+        new PentestError(
+          `Vertex AI mode requires the following env vars in .env: ${missing.join(', ')}`,
+          'config',
+          false,
+          { missing },
+          ErrorCode.AUTH_FAILED
+        )
+      );
+    }
+    // Validate service account credentials file is accessible
+    const credPath = process.env.GOOGLE_APPLICATION_CREDENTIALS;
+    if (!credPath) {
+      return err(
+        new PentestError(
+          'Vertex AI mode requires GOOGLE_APPLICATION_CREDENTIALS pointing to a service account key JSON file',
+          'config',
+          false,
+          {},
+          ErrorCode.AUTH_FAILED
+        )
+      );
+    }
+    try {
+      await fs.access(credPath);
+    } catch {
+      return err(
+        new PentestError(
+          `Service account key file not found at: ${credPath}`,
+          'config',
+          false,
+          { credPath },
+          ErrorCode.AUTH_FAILED
+        )
+      );
+    }
+    logger.info('Vertex AI credentials OK');
+    return ok(undefined);
+  }
+
+  // 4. Check that at least one credential is present
  if (!process.env.ANTHROPIC_API_KEY && !process.env.CLAUDE_CODE_OAUTH_TOKEN) {
    return err(
      new PentestError(
-        'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock)',
+        'No API credentials found. Set ANTHROPIC_API_KEY or CLAUDE_CODE_OAUTH_TOKEN in .env (or use CLAUDE_CODE_USE_BEDROCK=1 for AWS Bedrock, or CLAUDE_CODE_USE_VERTEX=1 for Google Vertex AI)',
        'config',
        false,
        {},
@@ -198,7 +243,7 @@ async function validateCredentials(
    );
  }

-  // 4. Validate via SDK query
+  // 5. Validate via SDK query
  const authType = process.env.CLAUDE_CODE_OAUTH_TOKEN ? 'OAuth token' : 'API key';
  logger.info(`Validating ${authType} via SDK...`);