Code Mode

Code Mode enables MCP clients to generate and execute structured queries (GraphQL, SQL) with a validation pipeline that ensures security and provides human-readable explanations.

Overview

Code Mode acts as LLM guard rails - protecting against risks specific to LLM-generated code:

Note: Code Mode is NOT for user authorization (handled by OAuth + backend systems).

Architecture

┌─────────────────────────────────────────────────────────────────────────────┐
│                           Code Mode Flow                                     │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                              │
│  MCP Client (LLM)                         MCP Server                         │
│  ═══════════════                          ═══════════                        │
│                                                                              │
│  1. User: "Show orders over $1000"                                           │
│           │                                                                  │
│           ▼                                                                  │
│  2. LLM generates GraphQL:                                                   │
│     query { orders(minAmount: 1000) { id total } }                           │
│           │                                                                  │
│           ▼                                                                  │
│  3. validate_query(query) ──────────────►  Policy Evaluation                 │
│           │                                  │                               │
│           │                                  ▼                               │
│           │                               ALLOW/DENY                         │
│           │                                  │                               │
│           ◄──────────────────────────────────┘                               │
│  4. Result:                                                                  │
│     - Explanation: "Reads orders over $1000"                                 │
│     - Risk: LOW                                                              │
│     - Token: <approval_token>                                                │
│           │                                                                  │
│           ▼                                                                  │
│  5. Show to user, get approval                                               │
│           │                                                                  │
│           ▼                                                                  │
│  6. execute_query(query, token) ─────────►  Verify token, execute            │
│           │                                                                  │
│           ◄──────────────────────────────── Return results                   │
│                                                                              │
└─────────────────────────────────────────────────────────────────────────────┘

Policy Evaluation Modes

Code Mode supports two policy evaluation approaches:

1. Local Cedar (Current Implementation)

Uses the cedar-policy crate for in-process evaluation:

• Latency: ~2-5ms
• Cost: $0 per request
• Configuration: TOML-based policies

[code_mode]
enabled = true
token_secret = "${CODE_MODE_SECRET}"

[code_mode.cedar_policies]
read_only_mode = true
block_introspection = true
max_depth = 10
max_fields = 100

2. Amazon Verified Permissions (Recommended for Production)

Uses AVP for managed policy storage and evaluation:

• Latency: ~10-50ms (imperceptible in Code Mode workflow)
• Cost: $1.50 per 10K requests
• Benefits: Managed service, immediate policy updates, audit trail

See AVP_ARCHITECTURE.md for the full AVP design.

Configuration

Basic Configuration (Read-Only API)

[code_mode]
enabled = true
allow_mutations = false
token_secret = "${CODE_MODE_SECRET}"

[code_mode.cedar_policies]
read_only_mode = true
block_introspection = true
max_depth = 10
max_fields = 100
block_subscriptions = true

With Specific Mutations Allowed

[code_mode]
enabled = true
allow_mutations = true
allowed_mutations = ["createOrder", "updateOrderStatus"]
blocked_mutations = ["deleteOrder"]
token_secret = "${CODE_MODE_SECRET}"

With Sensitive Data Controls

[code_mode]
enabled = true
blocked_fields = ["User.ssn", "User.password", "CreditCard.number"]
token_secret = "${CODE_MODE_SECRET}"

MCP Primitives

Code Mode adds the following MCP primitives to enabled servers:

Resource: graphql://schema

The GraphQL schema for query generation context.

Prompt: start_code_mode

Initiates Code Mode with schema and instructions for the LLM.

Tool: validate_query

Validates a GraphQL query and returns:

• Human-readable explanation
• Risk assessment
• Approval token (if valid)
• Policy violations (if invalid)

Tool: execute_query

Executes a validated query using an approval token.

Security

Approval Tokens

• HMAC-signed with server secret
• Bound to exact query hash (different query = invalid token)
• Short TTL (default: 5 minutes)
• Include context hash for TOCTOU protection

Policy Hierarchy

1. Baseline policies - Block subscriptions, enforce limits
2. Server policies - Generated from TOML config
3. Custom policies - Added via UI/API (AVP mode)

Feature Flags

Usage (Rust)

use mcp_server_common::code_mode::{
    CodeModeConfig, ValidationPipeline, ValidationContext
};

// With Cedar policies (local evaluation)
#[cfg(feature = "cedar")]
{
    use mcp_server_common::code_mode::CedarPolicyConfig;

    let config = CodeModeConfig::with_cedar_policies(CedarPolicyConfig::default());
    let pipeline = ValidationPipeline::new(config, b"secret-key".to_vec());
}

// Validate a query
let context = ValidationContext::new("user-123", "session-456", "schema-hash", "perms-hash");
let result = pipeline.validate_graphql_query("query { users { id name } }", &context)?;

if result.is_valid {
    println!("Explanation: {}", result.explanation);
    println!("Risk: {:?}", result.risk_level);
    // Show to user, get approval, then execute with token
}

Documentation

• AVP_ARCHITECTURE.md - Full AVP-based architecture design
• CEDAR_SECURITY_ARCHITECTURE.md - Cedar policy design details
• POLICY_ENGINE_DESIGN.md - Policy engine design
• SECURITY_SERVICES_DESIGN.md - Security services design

Roadmap

1. Current: Local Cedar evaluation with TOML config
2. Next: AVP integration for managed policy storage
3. Future: SQL Code Mode, REST Code Mode, Policy UI