# AgentShield WASM Setup for Next.js Edge Runtime

## The Challenge

Next.js Edge Runtime (used by middleware) has strict requirements for WebAssembly:

- ✅ Supports `WebAssembly.instantiate()`
- ❌ Cannot use `WebAssembly.compile()` with buffers
- ❌ Cannot dynamically load WASM from node_modules
- ✅ Requires static imports with `?module` suffix

## Solution: Manual WASM Setup

Since npm packages cannot reliably provide WASM files that work in Edge Runtime, we provide a setup
process that copies the necessary files to your Next.js project.

## Quick Start

### Step 1: Install AgentShield

```bash
npm install @kya-os/agentshield-nextjs
# or
pnpm add @kya-os/agentshield-nextjs
```

### Step 2: Run Setup Script

```bash
npx agentshield-setup-edge
```

This will:

1. Copy the WASM file to your project root
2. Create TypeScript definitions
3. Generate an Edge-compatible loader

### Step 3: Update your middleware.ts

```typescript
// middleware.ts
import { NextResponse } from 'next/server';
import type { NextRequest } from 'next/server';
import { createAgentShieldMiddleware } from './lib/agentshield-edge';

const agentShield = createAgentShieldMiddleware({
  enableWasm: true, // Enable WASM for 95%+ confidence
  onAgentDetected: agent => {
    console.log('AI Agent detected:', agent);
  },
});

export async function middleware(request: NextRequest) {
  // Initialize on first request
  await agentShield.init();

  // Check for AI agents
  const result = await agentShield.detect(request);

  if (result.isAgent && result.confidence > 0.85) {
    // Handle AI agent traffic
    console.log(`Detected ${result.agent} with ${result.confidence * 100}% confidence`);

    // Optional: Block or redirect
    // return NextResponse.redirect(new URL('/api-access-denied', request.url));
  }

  return NextResponse.next();
}

export const config = {
  matcher: ['/((?!_next/static|_next/image|favicon.ico).*)'],
};
```

## Manual Setup (Alternative)

If you prefer to set up manually or the script doesn't work:

### 1. Copy WASM File

```bash
# Create wasm directory in your project root
mkdir -p wasm

# Copy the WASM file from node_modules
cp node_modules/@kya-os/agentshield-nextjs/dist/wasm/agentshield_wasm_bg.wasm ./wasm/
```

### 2. Create TypeScript Definitions

Create `wasm/agentshield.d.ts`:

```typescript
// wasm/agentshield.d.ts
declare module '*.wasm?module' {
  const wasmModule: WebAssembly.Module;
  export default wasmModule;
}

export interface AgentShieldWasm {
  detect_agent(metadata: string): string;
  get_version(): string;
  Memory: WebAssembly.Memory;
}
```

### 3. Create Edge Loader

Create `lib/agentshield-edge.ts`:

```typescript
// lib/agentshield-edge.ts
import type { NextRequest } from 'next/server';
import wasmModule from '../wasm/agentshield_wasm_bg.wasm?module';

let wasmInstance: WebAssembly.Instance | null = null;
let wasmMemory: WebAssembly.Memory | null = null;

interface WasmExports {
  detect_agent(metadataPtr: number, metadataLen: number): number;
  get_version(): number;
  __wbindgen_malloc(size: number): number;
  __wbindgen_free(ptr: number, size: number): void;
  __wbindgen_realloc(ptr: number, oldSize: number, newSize: number): number;
  memory: WebAssembly.Memory;
}

async function initWasm(): Promise<void> {
  if (wasmInstance) return;

  // Import the WASM module (Edge Runtime compatible)
  wasmMemory = new WebAssembly.Memory({ initial: 17, maximum: 256 });

  const imports = {
    wbg: {
      __wbindgen_throw: (ptr: number, len: number) => {
        throw new Error(readString(ptr, len));
      },
    },
    env: {
      memory: wasmMemory,
    },
  };

  wasmInstance = await WebAssembly.instantiate(wasmModule, imports);
}

function readString(ptr: number, len: number): string {
  if (!wasmInstance || !wasmMemory) throw new Error('WASM not initialized');
  const memory = new Uint8Array(wasmMemory.buffer);
  const bytes = memory.slice(ptr, ptr + len);
  return new TextDecoder().decode(bytes);
}

function writeString(str: string): [number, number] {
  if (!wasmInstance) throw new Error('WASM not initialized');
  const exports = wasmInstance.exports as unknown as WasmExports;

  const encoded = new TextEncoder().encode(str);
  const ptr = exports.__wbindgen_malloc(encoded.length);

  const memory = new Uint8Array(exports.memory.buffer);
  memory.set(encoded, ptr);

  return [ptr, encoded.length];
}

export interface DetectionResult {
  isAgent: boolean;
  confidence: number;
  agent?: string;
  verificationMethod: 'cryptographic' | 'pattern';
}

export function createAgentShieldMiddleware(options: {
  enableWasm?: boolean;
  onAgentDetected?: (result: DetectionResult) => void;
}) {
  let initialized = false;

  return {
    async init(): Promise<void> {
      if (initialized) return;

      if (options.enableWasm !== false) {
        try {
          await initWasm();
          initialized = true;
          console.log('✅ AgentShield WASM initialized in Edge Runtime');
        } catch (error) {
          console.warn('⚠️ WASM initialization failed, falling back to pattern detection:', error);
          initialized = true;
        }
      } else {
        initialized = true;
      }
    },

    async detect(request: NextRequest): Promise<DetectionResult> {
      const metadata = {
        userAgent: request.headers.get('user-agent') || '',
        ipAddress: request.ip || request.headers.get('x-forwarded-for') || '',
        headers: Object.fromEntries(request.headers.entries()),
      };

      if (wasmInstance && options.enableWasm !== false) {
        try {
          const exports = wasmInstance.exports as unknown as WasmExports;
          const [ptr, len] = writeString(JSON.stringify(metadata));

          const resultPtr = exports.detect_agent(ptr, len);
          const resultLen = 1024; // Assume max result size
          const resultStr = readString(resultPtr, resultLen);

          exports.__wbindgen_free(ptr, len);
          exports.__wbindgen_free(resultPtr, resultLen);

          const result = JSON.parse(resultStr);

          if (options.onAgentDetected && result.isAgent) {
            options.onAgentDetected(result);
          }

          return {
            ...result,
            verificationMethod: 'cryptographic',
          };
        } catch (error) {
          console.error('WASM detection failed:', error);
          // Fall through to pattern detection
        }
      }

      // Fallback: Pattern-based detection (85% confidence)
      const userAgent = metadata.userAgent.toLowerCase();
      const aiAgentPatterns = [
        { pattern: /chatgpt/i, name: 'ChatGPT' },
        { pattern: /claude/i, name: 'Claude' },
        { pattern: /anthropic/i, name: 'Anthropic' },
        { pattern: /openai/i, name: 'OpenAI' },
        { pattern: /gpt-/i, name: 'GPT' },
        { pattern: /copilot/i, name: 'GitHub Copilot' },
        { pattern: /bard/i, name: 'Google Bard' },
        { pattern: /gemini/i, name: 'Google Gemini' },
      ];

      for (const { pattern, name } of aiAgentPatterns) {
        if (pattern.test(userAgent)) {
          const result: DetectionResult = {
            isAgent: true,
            confidence: 0.85,
            agent: name,
            verificationMethod: 'pattern',
          };

          if (options.onAgentDetected) {
            options.onAgentDetected(result);
          }

          return result;
        }
      }

      return {
        isAgent: false,
        confidence: 0.95,
        verificationMethod: 'pattern',
      };
    },
  };
}
```

## How It Works

1. **Static Import**: The WASM file is imported at build time using `?module` suffix
2. **Edge Compatible**: Uses `WebAssembly.instantiate()` with the imported module
3. **Graceful Fallback**: Falls back to pattern detection if WASM fails
4. **Type Safety**: Full TypeScript support with proper typing

## Verification Methods

- **Cryptographic (95-100% confidence)**: Uses WASM for cryptographic verification
- **Pattern (85% confidence)**: Fallback pattern matching for known AI agents

## Troubleshooting

### WASM not loading?

1. Ensure the WASM file exists at `wasm/agentshield_wasm_bg.wasm`
2. Check Next.js config doesn't exclude `.wasm` files
3. Verify the import path is correct

### TypeScript errors?

Add to your `tsconfig.json`:

```json
{
  "compilerOptions": {
    "types": ["@types/node"],
    "moduleResolution": "bundler"
  }
}
```

### Build errors?

Update your `next.config.js`:

```javascript
/** @type {import('next').NextConfig} */
const nextConfig = {
  webpack: config => {
    // Handle WASM imports
    config.module.rules.push({
      test: /\.wasm$/,
      type: 'asset/resource',
    });

    return config;
  },
  experimental: {
    // Ensure Edge Runtime can access WASM
    serverComponentsExternalPackages: ['@kya-os/agentshield-nextjs'],
  },
};

module.exports = nextConfig;
```

## Why Manual Setup?

The Edge Runtime's security model prevents dynamic code evaluation. While this adds a setup step, it
ensures:

1. **Reliability**: Works consistently across all deployment platforms
2. **Security**: No dynamic code evaluation
3. **Performance**: WASM is loaded at build time, not runtime
4. **Vercel Compatibility**: Works perfectly with Vercel's Edge Runtime

## Support

- GitHub Issues:
  [github.com/kya-os/agentshield/issues](https://github.com/kya-os/agentshield/issues)
- Documentation: [agentshield.ai/docs](https://agentshield.ai/docs)