@aegis-firewall/sdk

TypeScript/JavaScript SDK for Aegis LLM Inference Firewall.

Installation

npm install @aegis-firewall/sdk

Quick Start

import { AegisClient, AegisFirewall } from '@aegis-firewall/sdk';

// Create client
const client = new AegisClient({
  baseUrl: 'http://localhost:8080',
  apiKey: 'your-api-key',
});

// Create firewall wrapper
const firewall = new AegisFirewall(client);

// Add user message
const userSpan = await firewall.addUserMessage('Send email to [email protected]');

// Execute tool call with taint tracking
const result = await firewall.safeExecuteToolCall(
  'tool.send_email',
  { to: '[email protected]', subject: 'Hello' },
  [userSpan.id]
);

Features

• ✅ Type-safe API - Full TypeScript support
• ✅ Taint tracking - Automatic span management
• ✅ Policy builder - Programmatic policy creation
• ✅ Approval workflows - Human-in-the-loop support
• ✅ Trace verification - CBOR bundle validation
• ✅ Error handling - Rich error types
• ✅ Retry logic - Automatic retries with backoff

API Reference

AegisClient

Main client for interacting with Aegis gateway.

const client = new AegisClient({
  baseUrl: 'http://localhost:8080',
  apiKey: 'optional-api-key',
  timeout: 30000,
  retries: 3,
  enableTracing: false,
});

Methods

• validateToolCall(toolCall) - Validate tool call against policy
• createSpan(span) - Create new span
• getSpan(spanId) - Get span by ID
• listDecisions(sessionId?) - List policy decisions
• getApproval(approvalId) - Get approval request
• listApprovals() - List pending approvals
• approve(approvalId) - Approve request
• deny(approvalId, reason) - Deny request
• getTrace(sessionId) - Get trace session
• verifyTrace(bundleData) - Verify trace bundle
• health() - Health check

AegisFirewall

High-level wrapper for easy integration.

const firewall = new AegisFirewall(client);

Methods

• addUserMessage(content) - Add user message
• addDocument(content, sensitivity?) - Add RAG document
• addSystemPrompt(content) - Add system prompt
• executeToolCall(capability, args, sourceSpanIds) - Execute tool call
• safeExecuteToolCall(capability, args, sourceSpanIds, onApprovalRequired?) - Safe execution with error handling
• getSessionSpans() - Get all spans in session
• clearSession() - Clear session

PolicyBuilder

Build policies programmatically.

const policy = new PolicyBuilder()
  .allow('tool.read_calendar', TrustLevel.UntrustedUser)
  .allow('tool.send_email', TrustLevel.UntrustedUser, true)
  .deny('payments.*', TrustLevel.UntrustedDocument)
  .blockUntrustedDocuments()
  .blockSecretData()
  .requireApprovalForHighRisk()
  .autoDenyAfter(30)
  .build();

// Export as .aegis format
const aegisFormat = new PolicyBuilder()
  .allow('tool.send_email', TrustLevel.UntrustedUser, true)
  .toAegisFormat();

Examples

Basic Tool Call Validation

const result = await client.validateToolCall({
  capability: 'tool.send_email',
  arguments: {
    to: '[email protected]',
    subject: 'Hello',
    body: 'Test message',
  },
});

if (result.allowed) {
  // Execute tool
} else {
  console.error('Blocked:', result.reason);
}

RAG with Document Isolation

// Add user query
const userSpan = await firewall.addUserMessage('What is our pricing?');

// Add retrieved documents (automatically marked as untrusted)
const doc = await firewall.addDocument('Pricing: $99/month');

// Tool call sourced only from documents will be BLOCKED
const result = await firewall.safeExecuteToolCall(
  'tool.send_email',
  { to: '[email protected]' },
  [doc.id] // Only document source - blocked by taint tracking!
);

Approval Workflow

const result = await firewall.safeExecuteToolCall(
  'tool.execute_code',
  { code: 'print("hello")' },
  [userSpan.id],
  async (approvalId) => {
    console.log('Waiting for approval:', approvalId);
    
    // Poll for approval status
    const approval = await client.getApproval(approvalId);
    
    // Or handle in separate workflow
    await notifyAdmin(approvalId);
  }
);

Trace Verification

import { TraceManager } from '@aegis-firewall/sdk';

// Export trace
const trace = await client.getTrace('session-123');
await TraceManager.exportToFile(trace, './trace.cbor');

// Import and verify
const bundle = await TraceManager.importFromFile('./trace.cbor');
const verification = await TraceManager.verifyBundle(bundle);

if (!verification.valid) {
  console.error('Verification failed:', verification.errors);
}

Error Handling

The SDK provides rich error types:

import {
  CapabilityBlockedError,
  ApprovalRequiredError,
  TaintViolationError,
  NetworkError,
} from '@aegis-firewall/sdk';

try {
  await client.validateToolCall(toolCall);
} catch (error) {
  if (error instanceof CapabilityBlockedError) {
    console.log('Capability:', error.capability);
    console.log('Reason:', error.reason);
  } else if (error instanceof ApprovalRequiredError) {
    console.log('Approval ID:', error.approvalId);
    await handleApproval(error.approvalId);
  } else if (error instanceof NetworkError) {
    console.log('Status:', error.statusCode);
  }
}

TypeScript Support

Full TypeScript definitions included:

import {
  TrustLevel,
  SensitivityLevel,
  RiskLevel,
  RunMode,
  Span,
  ToolCall,
  PolicyDecision,
} from '@aegis-firewall/sdk';

License

MIT