Skip to main content
AI agents present unique risks compared to simple chat completions. They make autonomous decisions, call tools, access external systems, and execute multi-step workflows. A single compromised step can cascade through the entire pipeline. Know Your AI gives you full visibility and control over agent behavior.

​
What can go wrong with agents?

RiskDescriptionExample
Goal driftAgent deviates from its intended purposeSupport agent starts giving medical advice
Excessive agencyAgent takes actions beyond its scopeRead-only agent attempts to write data
Tool abuseAgent misuses available toolsAgent calls deleteUser instead of getUser
Recursive hijackingInjected instructions propagate across agent stepsPoisoned tool output redirects the agent
Data leakageAgent outputs sensitive information from its contextAgent reveals API keys from system prompt
Cost spiralsAgent enters loops consuming tokensInfinite retry loop with a large model

​
Trace every agent step

Use the tracing system to build a complete picture of what your agent does: 
import * as KnowYourAI from '@know-your-ai/node';
import { firewallIntegration } from '@know-your-ai/firewall';
import { GoogleGenAI } from '@google/genai';

KnowYourAI.init({
  dsn: process.env.KNOW_YOUR_AI_DSN!,
  traceMode: true,
  integrations: [
    KnowYourAI.googleGenAIIntegration(),
    firewallIntegration({
      baseUrl: process.env.FIREWALL_URL!,
      apiKey: process.env.FIREWALL_API_KEY!,
      onInputViolation: 'block',
      onOutputViolation: 'callback',
      violationCallback: async (ctx) => {
        console.warn(`[Agent Safety] ${ctx.phase} violation: ${ctx.validation.risks.map(r => r.category).join(', ')}`);
      },
    }),
  ],
});

const genAI = new GoogleGenAI({ apiKey: process.env.GOOGLE_API_KEY! });
const ai = KnowYourAI.instrumentGoogleGenAIClient(genAI);

​
Example: traced agent with tool calls

await KnowYourAI.withTrace(
  { name: 'order-management-agent', userId: 'user-456', sessionId: 'sess-789' },
  async () => {
    // Agent orchestration
    const agent = KnowYourAI.startAgent('order-agent', {});
    agent.setAvailableTools([
      { type: 'function', name: 'getOrder', description: 'Retrieve order details' },
      { type: 'function', name: 'updateOrderStatus', description: 'Update order status' },
      { type: 'function', name: 'sendNotification', description: 'Send email notification' },
    ]);

    // Step 1: Understand the request
    const intent = await KnowYourAI.withGeneration('classify-intent', async (gen) => {
      gen.setModel('gemini-2.0-flash');
      const res = await ai.models.generateContent({
        model: 'gemini-2.0-flash',
        contents: 'User says: Where is my order #12345?',
      });
      return res.text;
    });

    // Step 2: Tool call — fetch order
    const toolSpan = agent.startTool('getOrder', {});
    toolSpan.setToolDetails('getOrder', 'call_001', { orderId: '12345' });
    const order = await fetchOrder('12345');
    toolSpan.setResult(order);
    toolSpan.end();

    agent.incrementIterations();

    // Step 3: Generate response
    await KnowYourAI.withGeneration('compose-response', async (gen) => {
      gen.setModel('gemini-2.0-flash');
      const res = await ai.models.generateContent({
        model: 'gemini-2.0-flash',
        contents: `Order: ${JSON.stringify(order)}\nIntent: ${intent}\nGenerate a helpful status update.`,
      });
      return res.text;
    });

    agent.setFinalAction('Provided order status to user');
    agent.end();
  }
);
In the tracing dashboard you’ll see: 
order-management-agent (trace)
└── order-agent (agent) — 3 tools available, 1 iteration
    ├── classify-intent (generation) — gemini-2.0-flash, 120ms
    ├── getOrder (tool) — orderId: 12345, 45ms
    └── compose-response (generation) — gemini-2.0-flash, 850ms

​
Intercept suspicious agent behavior with hooks

Use beforeRequest hooks to enforce agent policies — what the agent can and cannot do:

​
Block unauthorized models

const hookManager = KnowYourAI.getHookManager();

const ALLOWED_MODELS = ['gemini-2.0-flash', 'gemini-1.5-pro'];

hookManager.addBeforeRequestHook(async (ctx) => {
  if (!ALLOWED_MODELS.includes(ctx.model)) {
    return {
      action: 'block',
      reason: `Model ${ctx.model} is not allowed. Allowed: ${ALLOWED_MODELS.join(', ')}`,
    };
  }
  return { action: 'continue' };
}, 'model-allowlist');

​
Enforce token limits

hookManager.addBeforeRequestHook(async (ctx) => {
  if (ctx.requestParams?.maxTokens && ctx.requestParams.maxTokens > 4096) {
    return {
      action: 'modify',
      modified: {
        requestParams: { ...ctx.requestParams, maxTokens: 4096 },
      },
      reason: 'Capped maxTokens to 4096',
    };
  }
  return { action: 'continue' };
}, 'token-limit');

​
Block sensitive keywords in inputs

const BLOCKED_PATTERNS = [
  /ignore.*(?:previous|all).*instructions/i,
  /you are (?:now )?DAN/i,
  /(?:reveal|output|show).*system prompt/i,
  /(?:admin|root) password/i,
];

hookManager.addBeforeRequestHook(async (ctx) => {
  const inputText = ctx.input?.map(m =>
    typeof m.content === 'string' ? m.content : JSON.stringify(m.content)
  ).join(' ') || '';

  for (const pattern of BLOCKED_PATTERNS) {
    if (pattern.test(inputText)) {
      return {
        action: 'block',
        reason: `Input matches blocked pattern: ${pattern.source}`,
      };
    }
  }
  return { action: 'continue' };
}, 'input-filter');

​
Detect PII in outputs

hookManager.addAfterResponseHook(async (ctx) => {
  if (!ctx.output) return { action: 'continue' };

  // Simple PII patterns
  const PII_PATTERNS = {
    email: /[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}/g,
    phone: /\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g,
    ssn: /\b\d{3}-\d{2}-\d{4}\b/g,
    creditCard: /\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b/g,
  };

  const detectedPII: string[] = [];
  for (const [type, pattern] of Object.entries(PII_PATTERNS)) {
    if (pattern.test(ctx.output)) {
      detectedPII.push(type);
    }
  }

  if (detectedPII.length > 0) {
    console.warn(`⚠️ PII detected in output: ${detectedPII.join(', ')}`);
    // Log, alert, or block based on your policy
    return {
      action: 'block',
      reason: `PII detected in output: ${detectedPII.join(', ')}`,
    };
  }
  return { action: 'continue' };
}, 'pii-detector');

​
Monitor agent behavior with after-response hooks

Track what your agents are actually doing in production:

​
Log all tool calls

hookManager.addAfterResponseHook(async (ctx) => {
  if (ctx.toolCalls && ctx.toolCalls.length > 0) {
    for (const call of ctx.toolCalls) {
      console.log(`[Agent Audit] Tool call: ${call.name}`, {
        model: ctx.model,
        arguments: call.arguments,
        duration: ctx.duration,
      });

      // Alert on dangerous tool calls
      const SENSITIVE_TOOLS = ['deleteUser', 'dropTable', 'executeCommand', 'sendPayment'];
      if (SENSITIVE_TOOLS.includes(call.name)) {
        await alertSecurityTeam({
          type: 'sensitive_tool_call',
          tool: call.name,
          model: ctx.model,
          arguments: call.arguments,
        });
      }
    }
  }
  return { action: 'continue' };
}, 'tool-audit');

​
Track cost and alert on spikes

let rollingCost = 0;
const COST_ALERT_THRESHOLD = 1.0; // $1 per minute

hookManager.addAfterResponseHook(async (ctx) => {
  if (ctx.tokenUsage) {
    // Rough cost estimate
    const inputCost = (ctx.tokenUsage.inputTokens || 0) * 0.000001;
    const outputCost = (ctx.tokenUsage.outputTokens || 0) * 0.000004;
    rollingCost += inputCost + outputCost;

    if (rollingCost > COST_ALERT_THRESHOLD) {
      await alertTeam(`🚨 AI cost threshold exceeded: $${rollingCost.toFixed(4)}`);
      rollingCost = 0; // reset
    }
  }
  return { action: 'continue' };
}, 'cost-monitor');

​
Combine monitoring + firewall for full protection

The most robust setup uses both layers together: 
import * as KnowYourAI from '@know-your-ai/node';
import { firewallIntegration } from '@know-your-ai/firewall';
import { GoogleGenAI } from '@google/genai';

KnowYourAI.init({
  dsn: process.env.KNOW_YOUR_AI_DSN!,
  environment: 'production',
  traceMode: true,
  integrations: [
    KnowYourAI.googleGenAIIntegration(),
    firewallIntegration({
      baseUrl: process.env.FIREWALL_URL!,
      apiKey: process.env.FIREWALL_API_KEY!,
      onInputViolation: 'block',
      onOutputViolation: 'callback',
      riskThreshold: 0.7,
      violationCallback: async (ctx) => {
        await logViolation(ctx);
        if (ctx.validation.risks.some(r => r.score >= 0.9)) {
          await alertSecurityTeam(ctx);
        }
      },
    }),
  ],
});

// Custom hooks for agent-specific policies
const hooks = KnowYourAI.getHookManager();

// 1. Only allow approved models
hooks.addBeforeRequestHook(async (ctx) => {
  if (!['gemini-2.0-flash', 'gemini-1.5-pro'].includes(ctx.model)) {
    return { action: 'block', reason: `Unauthorized model: ${ctx.model}` };
  }
}, 'model-policy');

// 2. Cap token usage to prevent runaway costs
hooks.addBeforeRequestHook(async (ctx) => {
  if (ctx.requestParams?.maxTokens && ctx.requestParams.maxTokens > 8192) {
    return {
      action: 'modify',
      modified: { requestParams: { ...ctx.requestParams, maxTokens: 8192 } },
    };
  }
}, 'token-cap');

// 3. Audit all tool calls
hooks.addAfterResponseHook(async (ctx) => {
  if (ctx.toolCalls?.length) {
    console.log(`[Audit] ${ctx.model} called tools: ${ctx.toolCalls.map(t => t.name).join(', ')}`);
  }
}, 'tool-audit');

// Now all AI calls are monitored, firewalled, and policy-enforced
const genAI = new GoogleGenAI({ apiKey: process.env.GOOGLE_API_KEY! });
const client = KnowYourAI.instrumentGoogleGenAIClient(genAI);
This gives you:
LayerProtection
Firewall (input)Blocks jailbreaks, prompt injection, PII extraction
Custom hooks (before)Enforces model allowlist, caps token usage
AI modelProcesses the validated request
Firewall (output)Flags harmful, biased, or toxic content
Custom hooks (after)Audits tool calls, tracks cost, detects PII
MonitoringCaptures everything to the dashboard and traces