Test script for AI Agent Intent Mode SSE responses

test-intent-sse.ts

#!/usr/bin/env npx tsx
/**
 * Test Script for AI Agent Intent Mode SSE Responses
 *
 * Tests classification intent types with streaming enabled:
 * - siem_sweep: SIEM entity search (IP, domain, user, hash)
 * - context_sweep: Context/baseline search (hosts)
 * - ask_direct: AI knowledge queries
 * - siem_query: Query generation for SIEM searches
 * - ti_sweep: Threat Intelligence IOC lookups
 * - multi_sweep: Multi-entity sweep operations
 *
 * Usage:
 *   npx tsx scripts/test-intent-sse.ts [options]
 *
 * Options:
 *   --base-url <url>    API base URL (default: http://localhost:8000/api/v1/th)
 *   --intent <type>     Run only specific intent type
 *   --query <text>      Run a custom query
 *   --no-stream         Disable streaming (use JSON response)
 *   --verbose           Show full response payloads
 *   --help              Show help
 */

import { parseArgs } from "node:util";

// ============================================================================
// Types
// ============================================================================

interface TestQuery {
  name: string;
  query: string;
  expectedIntent: string;
  description: string;
}

interface SSEEvent {
  event: string;
  data: unknown;
  raw: string;
}

interface TestResult {
  name: string;
  query: string;
  expectedIntent: string;
  actualIntent: string | null;
  success: boolean;
  events: SSEEvent[];
  error: string | null;
  duration: number;
  classification: unknown | null;
  result: unknown | null;
}

// ============================================================================
// Test Queries by Intent Type
// ============================================================================

/**
 * Backend Intent Types (from /aiagent/query classification):
 *
 * SIEM/Entity Search:
 * - siem_sweep: SIEM entity search (IP, domain, user, hash)
 * - context_sweep: Context/baseline search (typically for hosts)
 *
 * Knowledge Queries:
 * - ask_direct: Direct AI knowledge query (security concepts, explanations)
 * - siem_query: Query generation for SIEM searches
 *
 * Threat Intelligence:
 * - ti_sweep: Threat Intelligence IOC lookup
 *
 * Deep Investigations (via /aiagent/hunt, not /aiagent/query):
 * - c2_investigation, user_investigation, host_investigation
 *
 * Note: The /aiagent/query endpoint handles most intents directly.
 * Deep investigations may route to sweep operations or require /hunt endpoint.
 */

const TEST_QUERIES: Record<string, TestQuery[]> = {
  // SIEM Sweep tests - fast entity search
  siem_sweep: [
    {
      name: "IP Sweep",
      query: "Check 8.8.8.8 in SIEM",
      expectedIntent: "siem_sweep",
      description: "Single IP entity sweep",
    },
    {
      name: "Domain Sweep",
      query: "Search for domain google.com activity",
      expectedIntent: "siem_sweep",
      description: "Domain entity sweep",
    },
    {
      name: "User Sweep",
      query: "Check user john.doe",
      expectedIntent: "siem_sweep",
      description: "User entity sweep",
    },
    {
      name: "Hash Sweep",
      query: "Check hash d41d8cd98f00b204e9800998ecf8427e",
      expectedIntent: "siem_sweep",
      description: "MD5 hash entity sweep",
    },
    {
      name: "IP Activity",
      query: "Search for 192.168.1.100 activity",
      expectedIntent: "siem_sweep",
      description: "IP activity search",
    },
  ],

  // Context Sweep tests - baseline/context search
  context_sweep: [
    {
      name: "Host Sweep",
      query: "Look up host WORKSTATION-01",
      expectedIntent: "context_sweep",
      description: "Host entity context sweep",
    },
    {
      name: "Host Investigation",
      query: "Deep investigation of host INFECTED-PC",
      expectedIntent: "context_sweep",
      description: "Host context lookup",
    },
    {
      name: "Host Context",
      query: "Check context for host SERVER-01",
      expectedIntent: "context_sweep",
      description: "Host baseline context",
    },
  ],

  // Ask Direct tests - AI knowledge queries
  ask_direct: [
    {
      name: "Security Concept",
      query: "What is a lateral movement attack?",
      expectedIntent: "ask_direct",
      description: "General security knowledge query",
    },
    {
      name: "MITRE ATT&CK",
      query: "Explain MITRE ATT&CK T1059",
      expectedIntent: "ask_direct",
      description: "MITRE technique explanation",
    },
    {
      name: "Threat Info",
      query: "What are signs of ransomware infection?",
      expectedIntent: "ask_direct",
      description: "Threat indicator information",
    },
    {
      name: "Security Best Practice",
      query: "How should I respond to a phishing incident?",
      expectedIntent: "ask_direct",
      description: "Incident response guidance",
    },
  ],

  // SIEM Query tests - query generation
  siem_query: [
    {
      name: "Failed Logins Query",
      query: "Generate query for failed login attempts",
      expectedIntent: "siem_query",
      description: "Generate failed login detection query",
    },
    {
      name: "Detection Query",
      query: "How do I detect credential dumping?",
      expectedIntent: "siem_query",
      description: "Detection guidance query",
    },
    {
      name: "Network Query",
      query: "Build query for network connections to port 4444",
      expectedIntent: "siem_query",
      description: "Suspicious port connection query",
    },
  ],

  // Threat Intelligence tests
  ti_sweep: [
    {
      name: "TI IP Lookup",
      query: "TI lookup 185.220.101.1",
      expectedIntent: "ti_sweep",
      description: "Threat intelligence lookup for IP",
    },
    {
      name: "VirusTotal Check",
      query: "Check VirusTotal for 45.33.32.156",
      expectedIntent: "ti_sweep",
      description: "VirusTotal reputation check",
    },
    {
      name: "Domain Reputation",
      query: "Check threat intelligence for evil-domain.com",
      expectedIntent: "ti_sweep",
      description: "Domain TI enrichment",
    },
    {
      name: "IOC Reputation",
      query: "What is the reputation of IP 8.8.4.4?",
      expectedIntent: "ti_sweep",
      description: "IOC reputation lookup",
    },
  ],

  // Multi-entity sweep tests
  multi_sweep: [
    {
      name: "Multiple IPs",
      query: "Check IPs 8.8.8.8, 1.1.1.1, and 9.9.9.9",
      expectedIntent: "siem_sweep",
      description: "Multi-entity IP sweep",
    },
    {
      name: "Multiple Users",
      query: "Search for users john.doe, jane.smith, admin",
      expectedIntent: "siem_sweep",
      description: "Multi-entity user sweep",
    },
    {
      name: "Multiple Domains",
      query: "Look up domains google.com, microsoft.com, apple.com",
      expectedIntent: "siem_sweep",
      description: "Multi-entity domain sweep",
    },
  ],

  // Investigation tests - these typically route to sweeps via /query
  investigation: [
    {
      name: "C2 Investigation",
      query: "Investigate potential C2 communication from 45.33.32.156",
      expectedIntent: "siem_sweep",
      description: "C2 investigation routes to SIEM sweep",
    },
    {
      name: "User Investigation",
      query: "Investigate user suspicious.admin for anomalies",
      expectedIntent: "siem_sweep",
      description: "User investigation routes to SIEM sweep",
    },
  ],
};

// ============================================================================
// SSE Parser
// ============================================================================

async function parseSSEStream(
  response: Response,
  onEvent: (event: SSEEvent) => void
): Promise<void> {
  const reader = response.body?.getReader();
  if (!reader) {
    throw new Error("No response body reader available");
  }

  const decoder = new TextDecoder();
  let buffer = "";

  try {
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });

      // Process complete SSE messages (separated by double newlines)
      const messages = buffer.split("\n\n");
      buffer = messages.pop() || ""; // Keep incomplete message in buffer

      for (const message of messages) {
        if (!message.trim()) continue;

        let eventType = "message";
        let data = "";

        for (const line of message.split("\n")) {
          if (line.startsWith("event:")) {
            eventType = line.slice(6).trim();
          } else if (line.startsWith("data:")) {
            data = line.slice(5).trim();
          }
        }

        if (data) {
          try {
            const parsedData = JSON.parse(data);
            onEvent({ event: eventType, data: parsedData, raw: data });
          } catch {
            onEvent({ event: eventType, data, raw: data });
          }
        }
      }
    }
  } finally {
    reader.releaseLock();
  }
}

// ============================================================================
// Test Runner
// ============================================================================

async function runTest(
  baseUrl: string,
  testQuery: TestQuery,
  stream: boolean,
  verbose: boolean
): Promise<TestResult> {
  const startTime = Date.now();
  const events: SSEEvent[] = [];
  let classification: unknown = null;
  let result: unknown = null;
  let actualIntent: string | null = null;
  let error: string | null = null;

  try {
    const payload = {
      query: testQuery.query,
      stream,
      compact: true,
    };

    const response = await fetch(`${baseUrl}/aiagent/query`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify(payload),
    });

    if (!response.ok) {
      throw new Error(`HTTP ${response.status}: ${response.statusText}`);
    }

    if (stream) {
      // Parse SSE stream
      await parseSSEStream(response, (event) => {
        events.push(event);

        if (verbose) {
          console.log(`  [${event.event}]`, JSON.stringify(event.data, null, 2).slice(0, 200));
        }

        // Extract classification and result from events
        const eventData = event.data as Record<string, unknown>;

        // phase_complete with phase="understanding" contains classification info
        if (event.event === "phase_complete" && eventData.phase === "understanding") {
          if (eventData.intent && !actualIntent) {
            actualIntent = eventData.intent as string;
          }
        }

        // result event contains full classification and T2Trace_output
        if (event.event === "result" && typeof event.data === "object") {
          result = event.data;
          if (eventData.classification && typeof eventData.classification === "object") {
            classification = eventData.classification;
            const classObj = eventData.classification as Record<string, unknown>;
            if (classObj.intent) {
              actualIntent = classObj.intent as string;
            }
          }
        }

        // complete event may also have classification
        if (event.event === "complete" && typeof event.data === "object") {
          if (!result && eventData.T2Trace_output) {
            result = event.data;
          }
          if (!classification && eventData.classification) {
            classification = eventData.classification;
            actualIntent = (eventData.classification as Record<string, unknown>).intent as string;
          }
        }

        // Handle errors
        if (event.event === "error") {
          error = typeof event.data === "object"
            ? (eventData.message as string) || JSON.stringify(event.data)
            : String(event.data);
        }
      });
    } else {
      // Parse JSON response
      const data = await response.json();
      classification = data.classification;
      result = data;
      actualIntent = data.classification?.intent;

      if (verbose) {
        console.log("  Response:", JSON.stringify(data, null, 2).slice(0, 500));
      }

      if (!data.success && !data.confirmation_required) {
        error = data.error || "Query failed";
      }
    }
  } catch (err) {
    error = err instanceof Error ? err.message : String(err);
  }

  const duration = Date.now() - startTime;

  // Determine success
  // For ambiguous tests, we expect confirmation_required or low confidence
  let success = false;
  if (testQuery.expectedIntent === "ambiguous") {
    const classData = classification as Record<string, unknown> | null;
    success =
      classData?.ambiguous === true ||
      (classData?.confidence !== undefined && (classData.confidence as number) < 0.7) ||
      (result as Record<string, unknown>)?.confirmation_required === true;
  } else {
    success = actualIntent === testQuery.expectedIntent && !error;
  }

  return {
    name: testQuery.name,
    query: testQuery.query,
    expectedIntent: testQuery.expectedIntent,
    actualIntent,
    success,
    events,
    error,
    duration,
    classification,
    result,
  };
}

// ============================================================================
// Main
// ============================================================================

async function main() {
  const { values } = parseArgs({
    options: {
      "base-url": { type: "string", default: "http://localhost:8000/api/v1/th" },
      intent: { type: "string" },
      query: { type: "string" },
      "no-stream": { type: "boolean", default: false },
      verbose: { type: "boolean", short: "v", default: false },
      help: { type: "boolean", short: "h", default: false },
    },
    allowPositionals: true,
  });

  if (values.help) {
    console.log(`
AI Agent Intent Mode SSE Test Script

Tests all classification intent types with streaming enabled.

Usage:
  npx tsx scripts/test-intent-sse.ts [options]

Options:
  --base-url <url>    API base URL (default: http://localhost:8000/api/v1/th)
  --intent <type>     Run only specific intent type (see below)
  --query <text>      Run a custom query
  --no-stream         Disable streaming (use JSON response)
  --verbose, -v       Show full response payloads
  --help, -h          Show help

Intent Types:
  siem_sweep     - SIEM entity search (IP, domain, user, hash)
  context_sweep  - Context/baseline search (hosts)
  ask_direct     - AI knowledge queries
  siem_query     - Query generation for SIEM searches
  ti_sweep       - Threat Intelligence IOC lookups
  multi_sweep    - Multi-entity sweep operations
  investigation  - Investigation queries (routes to sweep)

Examples:
  # Run all tests with streaming
  npx tsx scripts/test-intent-sse.ts

  # Run only SIEM sweep tests
  npx tsx scripts/test-intent-sse.ts --intent siem_sweep

  # Run only ask intent tests
  npx tsx scripts/test-intent-sse.ts --intent ask_intent

  # Run a custom query
  npx tsx scripts/test-intent-sse.ts --query "Check user admin"

  # Run against remote backend
  npx tsx scripts/test-intent-sse.ts --base-url https://lab207.cl4.ai/api/v1/th
`);
    process.exit(0);
  }

  const baseUrl = values["base-url"] as string;
  const stream = !values["no-stream"];
  const verbose = values.verbose as boolean;
  const intentFilter = values.intent as string | undefined;
  const customQuery = values.query as string | undefined;

  console.log("\n🔬 AI Agent Intent Mode SSE Test");
  console.log("================================");
  console.log(`Base URL: ${baseUrl}`);
  console.log(`Streaming: ${stream ? "enabled" : "disabled"}`);
  console.log(`Verbose: ${verbose ? "yes" : "no"}`);
  console.log("");

  // Custom query mode
  if (customQuery) {
    console.log(`\n📝 Custom Query: "${customQuery}"\n`);

    const testQuery: TestQuery = {
      name: "Custom Query",
      query: customQuery,
      expectedIntent: "unknown",
      description: "User-provided custom query",
    };

    const result = await runTest(baseUrl, testQuery, stream, true);

    console.log("\n📊 Result:");
    console.log(`  Intent: ${result.actualIntent}`);
    console.log(`  Duration: ${result.duration}ms`);
    console.log(`  Events: ${result.events.length}`);
    console.log(`  Error: ${result.error || "none"}`);

    if (result.classification) {
      console.log("\n📋 Classification:");
      console.log(JSON.stringify(result.classification, null, 2));
    }

    if (result.result && verbose) {
      console.log("\n📦 Full Result:");
      console.log(JSON.stringify(result.result, null, 2));
    }

    process.exit(result.error ? 1 : 0);
  }

  // Get tests to run
  const testCategories = intentFilter
    ? { [intentFilter]: TEST_QUERIES[intentFilter] || [] }
    : TEST_QUERIES;

  if (intentFilter && !TEST_QUERIES[intentFilter]) {
    console.error(`❌ Unknown intent type: ${intentFilter}`);
    console.log(`Available types: ${Object.keys(TEST_QUERIES).join(", ")}`);
    process.exit(1);
  }

  // Run tests
  const allResults: TestResult[] = [];

  for (const [category, queries] of Object.entries(testCategories)) {
    console.log(`\n📁 ${category.toUpperCase()} Intent Tests`);
    console.log("-".repeat(40));

    for (const testQuery of queries) {
      process.stdout.write(`  ▶ ${testQuery.name}... `);

      const result = await runTest(baseUrl, testQuery, stream, verbose);
      allResults.push(result);

      if (result.success) {
        console.log(`✅ (${result.actualIntent}, ${result.duration}ms, ${result.events.length} events)`);
      } else {
        console.log(`❌ (expected: ${result.expectedIntent}, got: ${result.actualIntent}, error: ${result.error})`);
      }

      // Small delay between tests
      await new Promise((resolve) => setTimeout(resolve, 500));
    }
  }

  // Summary
  const passed = allResults.filter((r) => r.success).length;
  const failed = allResults.filter((r) => !r.success).length;
  const total = allResults.length;

  console.log("\n" + "=".repeat(50));
  console.log("📊 SUMMARY");
  console.log("=".repeat(50));
  console.log(`Total: ${total}`);
  console.log(`Passed: ${passed} ✅`);
  console.log(`Failed: ${failed} ❌`);
  console.log(`Success Rate: ${((passed / total) * 100).toFixed(1)}%`);

  if (failed > 0) {
    console.log("\n❌ Failed Tests:");
    for (const result of allResults.filter((r) => !r.success)) {
      console.log(`  - ${result.name}: ${result.error || `expected ${result.expectedIntent}, got ${result.actualIntent}`}`);
    }
  }

  // Event type summary
  const eventCounts: Record<string, number> = {};
  for (const result of allResults) {
    for (const event of result.events) {
      eventCounts[event.event] = (eventCounts[event.event] || 0) + 1;
    }
  }

  if (Object.keys(eventCounts).length > 0) {
    console.log("\n📈 SSE Event Summary:");
    for (const [event, count] of Object.entries(eventCounts).sort((a, b) => b[1] - a[1])) {
      console.log(`  ${event}: ${count}`);
    }
  }

  console.log("");
  process.exit(failed > 0 ? 1 : 0);
}

main().catch((err) => {
  console.error("Fatal error:", err);
  process.exit(1);
});