rax-flow-core 0.2.0 → 2.0.1

package/src/governance/policy-engine.ts

@@ -0,0 +1,582 @@
+import { Policy } from "../plugins/governance-plugin.js";
+import { PolicyViolation, PolicyDecision, AgentInput, AgentOutput } from "../types/contracts.js";
+
+export interface GovernanceRule {
+  id: string;
+  name: string;
+  description: string;
+  enabled: boolean;
+  severity: "low" | "medium" | "high" | "critical";
+  scope: "input" | "output" | "both";
+  evaluate: (context: GovernanceContext) => Promise<PolicyDecision>;
+}
+
+export interface GovernanceContext {
+  taskId: string;
+  nodeId?: string;
+  agent?: string;
+  input?: AgentInput;
+  output?: AgentOutput;
+  metadata?: Record<string, unknown>;
+  costUsd?: number;
+  tokensUsed?: number;
+  latencyMs?: number;
+}
+
+export interface GovernanceConfig {
+  maxCostPerRun: number;
+  maxTokensPerRun: number;
+  maxLatencyMs: number;
+  allowedIntents?: string[];
+  blockedPatterns?: string[];
+  requireApprovalFor?: string[];
+  rateLimitPerMinute?: number;
+}
+
+export interface GovernanceStats {
+  totalEvaluations: number;
+  violationsByType: Record<string, number>;
+  violationsBySeverity: Record<string, number>;
+  blockedCount: number;
+  allowedCount: number;
+}
+
+const DEFAULT_CONFIG: GovernanceConfig = {
+  maxCostPerRun: 1.0,
+  maxTokensPerRun: 100000,
+  maxLatencyMs: 60000
+};
+
+export class GovernanceEngine {
+  private rules: Map<string, GovernanceRule> = new Map();
+  private config: GovernanceConfig;
+  private stats: GovernanceStats = {
+    totalEvaluations: 0,
+    violationsByType: {},
+    violationsBySeverity: {},
+    blockedCount: 0,
+    allowedCount: 0
+  };
+  private rateLimitTracker: Map<string, number[]> = new Map();
+
+  constructor(config: Partial<GovernanceConfig> = {}) {
+    this.config = { ...DEFAULT_CONFIG, ...config };
+    this.initializeDefaultRules();
+  }
+
+  private initializeDefaultRules(): void {
+    this.addRule({
+      id: "cost_limit",
+      name: "Cost Limit",
+      description: "Ensures execution cost stays within limits",
+      enabled: true,
+      severity: "high",
+      scope: "output",
+      evaluate: async (ctx) => this.evaluateCostLimit(ctx)
+    });
+
+    this.addRule({
+      id: "token_limit",
+      name: "Token Limit",
+      description: "Ensures token usage stays within limits",
+      enabled: true,
+      severity: "medium",
+      scope: "output",
+      evaluate: async (ctx) => this.evaluateTokenLimit(ctx)
+    });
+
+    this.addRule({
+      id: "latency_limit",
+      name: "Latency Limit",
+      description: "Ensures execution time stays within limits",
+      enabled: true,
+      severity: "medium",
+      scope: "output",
+      evaluate: async (ctx) => this.evaluateLatencyLimit(ctx)
+    });
+
+    this.addRule({
+      id: "blocked_content",
+      name: "Blocked Content",
+      description: "Blocks prompts containing disallowed patterns",
+      enabled: true,
+      severity: "critical",
+      scope: "input",
+      evaluate: async (ctx) => this.evaluateBlockedContent(ctx)
+    });
+
+    this.addRule({
+      id: "intent_whitelist",
+      name: "Intent Whitelist",
+      description: "Only allows specified intents",
+      enabled: false,
+      severity: "high",
+      scope: "input",
+      evaluate: async (ctx) => this.evaluateIntentWhitelist(ctx)
+    });
+
+    this.addRule({
+      id: "approval_required",
+      name: "Approval Required",
+      description: "Requires human approval for certain operations",
+      enabled: true,
+      severity: "medium",
+      scope: "both",
+      evaluate: async (ctx) => this.evaluateApprovalRequired(ctx)
+    });
+
+    this.addRule({
+      id: "rate_limit",
+      name: "Rate Limit",
+      description: "Limits request frequency",
+      enabled: true,
+      severity: "high",
+      scope: "input",
+      evaluate: async (ctx) => this.evaluateRateLimit(ctx)
+    });
+  }
+
+  addRule(rule: GovernanceRule): void {
+    this.rules.set(rule.id, rule);
+  }
+
+  removeRule(ruleId: string): boolean {
+    return this.rules.delete(ruleId);
+  }
+
+  enableRule(ruleId: string): boolean {
+    const rule = this.rules.get(ruleId);
+    if (rule) {
+      rule.enabled = true;
+      return true;
+    }
+    return false;
+  }
+
+  disableRule(ruleId: string): boolean {
+    const rule = this.rules.get(ruleId);
+    if (rule) {
+      rule.enabled = false;
+      return true;
+    }
+    return false;
+  }
+
+  async evaluateInput(context: GovernanceContext): Promise<PolicyDecision> {
+    const allViolations: PolicyViolation[] = [];
+
+    for (const rule of this.rules.values()) {
+      if (!rule.enabled || (rule.scope !== "input" && rule.scope !== "both")) continue;
+
+      const decision = await rule.evaluate(context);
+      this.updateStats(rule.id, decision);
+
+      if (!decision.allowed) {
+        allViolations.push(...decision.violations);
+      }
+    }
+
+    return {
+      allowed: allViolations.length === 0 || !allViolations.some(v => v.severity === "critical"),
+      violations: allViolations
+    };
+  }
+
+  async evaluateOutput(context: GovernanceContext): Promise<PolicyDecision> {
+    const allViolations: PolicyViolation[] = [];
+
+    for (const rule of this.rules.values()) {
+      if (!rule.enabled || (rule.scope !== "output" && rule.scope !== "both")) continue;
+
+      const decision = await rule.evaluate(context);
+      this.updateStats(rule.id, decision);
+
+      if (!decision.allowed) {
+        allViolations.push(...decision.violations);
+      }
+    }
+
+    return {
+      allowed: !allViolations.some(v => v.severity === "critical" || v.severity === "high"),
+      violations: allViolations
+    };
+  }
+
+  async evaluate(context: GovernanceContext, phase: "input" | "output"): Promise<PolicyDecision> {
+    this.stats.totalEvaluations++;
+
+    if (phase === "input") {
+      return this.evaluateInput(context);
+    }
+    return this.evaluateOutput(context);
+  }
+
+  private async evaluateCostLimit(context: GovernanceContext): Promise<PolicyDecision> {
+    const cost = context.costUsd ?? context.output?.costUsd ?? 0;
+
+    if (cost > this.config.maxCostPerRun) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: "cost_limit",
+          severity: "high",
+          message: `Cost $${cost.toFixed(4)} exceeds limit $${this.config.maxCostPerRun}`
+        }]
+      };
+    }
+
+    if (cost > this.config.maxCostPerRun * 0.8) {
+      return {
+        allowed: true,
+        violations: [{
+          policy: "cost_limit",
+          severity: "low",
+          message: `Cost approaching limit: $${cost.toFixed(4)} / $${this.config.maxCostPerRun}`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  private async evaluateTokenLimit(context: GovernanceContext): Promise<PolicyDecision> {
+    const tokens = context.tokensUsed ?? context.output?.usage?.totalTokens ?? 0;
+
+    if (tokens > this.config.maxTokensPerRun) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: "token_limit",
+          severity: "medium",
+          message: `Token usage ${tokens} exceeds limit ${this.config.maxTokensPerRun}`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  private async evaluateLatencyLimit(context: GovernanceContext): Promise<PolicyDecision> {
+    const latency = context.latencyMs ?? 0;
+
+    if (latency > this.config.maxLatencyMs) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: "latency_limit",
+          severity: "medium",
+          message: `Latency ${latency}ms exceeds limit ${this.config.maxLatencyMs}ms`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  private async evaluateBlockedContent(context: GovernanceContext): Promise<PolicyDecision> {
+    const violations: PolicyViolation[] = [];
+    const patterns = this.config.blockedPatterns ?? [];
+
+    if (!context.input?.userPrompt) {
+      return { allowed: true, violations: [] };
+    }
+
+    const text = context.input.userPrompt.toLowerCase();
+
+    for (const pattern of patterns) {
+      if (text.includes(pattern.toLowerCase())) {
+        violations.push({
+          policy: "blocked_content",
+          severity: "critical",
+          message: `Prompt contains blocked pattern: ${pattern}`
+        });
+      }
+    }
+
+    const dangerousPatterns = [
+      "ignore previous instructions",
+      "disregard all",
+      "override safety",
+      "bypass restrictions"
+    ];
+
+    for (const pattern of dangerousPatterns) {
+      if (text.includes(pattern)) {
+        violations.push({
+          policy: "blocked_content",
+          severity: "critical",
+          message: `Potential prompt injection detected: ${pattern}`
+        });
+      }
+    }
+
+    return {
+      allowed: violations.length === 0,
+      violations
+    };
+  }
+
+  private async evaluateIntentWhitelist(context: GovernanceContext): Promise<PolicyDecision> {
+    if (!this.config.allowedIntents || this.config.allowedIntents.length === 0) {
+      return { allowed: true, violations: [] };
+    }
+
+    const intent = context.metadata?.intent as string | undefined;
+    if (!intent) {
+      return { allowed: true, violations: [] };
+    }
+
+    if (!this.config.allowedIntents.includes(intent)) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: "intent_whitelist",
+          severity: "high",
+          message: `Intent '${intent}' is not in allowed list`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  private async evaluateApprovalRequired(context: GovernanceContext): Promise<PolicyDecision> {
+    const approvalFor = this.config.requireApprovalFor ?? [];
+    const agent = context.agent ?? context.metadata?.agent as string | undefined;
+
+    if (!agent || !approvalFor.includes(agent)) {
+      return { allowed: true, violations: [] };
+    }
+
+    if (context.metadata?.approved === true) {
+      return { allowed: true, violations: [] };
+    }
+
+    return {
+      allowed: false,
+      violations: [{
+        policy: "approval_required",
+        severity: "medium",
+        message: `Agent '${agent}' requires human approval before execution`
+      }]
+    };
+  }
+
+  private async evaluateRateLimit(context: GovernanceContext): Promise<PolicyDecision> {
+    if (!this.config.rateLimitPerMinute) {
+      return { allowed: true, violations: [] };
+    }
+
+    const key = context.metadata?.userId as string ?? "anonymous";
+    const now = Date.now();
+    const windowStart = now - 60000;
+
+    let requests = this.rateLimitTracker.get(key) ?? [];
+    requests = requests.filter(t => t > windowStart);
+
+    if (requests.length >= this.config.rateLimitPerMinute) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: "rate_limit",
+          severity: "high",
+          message: `Rate limit exceeded: ${requests.length} requests in the last minute`
+        }]
+      };
+    }
+
+    requests.push(now);
+    this.rateLimitTracker.set(key, requests);
+
+    return { allowed: true, violations: [] };
+  }
+
+  private updateStats(ruleId: string, decision: PolicyDecision): void {
+    if (decision.allowed) {
+      this.stats.allowedCount++;
+    } else {
+      this.stats.blockedCount++;
+    }
+
+    for (const violation of decision.violations) {
+      this.stats.violationsByType[ruleId] = (this.stats.violationsByType[ruleId] ?? 0) + 1;
+      this.stats.violationsBySeverity[violation.severity] = (this.stats.violationsBySeverity[violation.severity] ?? 0) + 1;
+    }
+  }
+
+  getStats(): GovernanceStats {
+    return { ...this.stats };
+  }
+
+  getConfig(): GovernanceConfig {
+    return { ...this.config };
+  }
+
+  updateConfig(updates: Partial<GovernanceConfig>): void {
+    this.config = { ...this.config, ...updates };
+  }
+
+  getRules(): GovernanceRule[] {
+    return Array.from(this.rules.values());
+  }
+
+  getRule(id: string): GovernanceRule | undefined {
+    return this.rules.get(id);
+  }
+
+  resetStats(): void {
+    this.stats = {
+      totalEvaluations: 0,
+      violationsByType: {},
+      violationsBySeverity: {},
+      blockedCount: 0,
+      allowedCount: 0
+    };
+    this.rateLimitTracker.clear();
+  }
+}
+
+export interface CostPolicyConfig {
+  maxPerRequest: number;
+  maxPerSession: number;
+  maxPerDay: number;
+  alertThreshold: number;
+}
+
+export class CostPolicy implements Policy {
+  name = "CostPolicy";
+  private sessionCosts: Map<string, number> = new Map();
+  private dailyCosts: Map<string, number> = new Map();
+
+  constructor(private config: CostPolicyConfig) {}
+
+  async validateInput(input: AgentInput): Promise<PolicyDecision> {
+    const sessionId = input.context?.sessionId as string | undefined;
+    const today = new Date().toISOString().split("T")[0];
+
+    const sessionCost = this.sessionCosts.get(sessionId ?? "default") ?? 0;
+    const dailyCost = this.dailyCosts.get(today) ?? 0;
+
+    const violations: PolicyViolation[] = [];
+
+    if (sessionCost >= this.config.maxPerSession) {
+      violations.push({
+        policy: this.name,
+        severity: "critical",
+        message: `Session cost limit reached: $${sessionCost.toFixed(4)}`
+      });
+    }
+
+    if (dailyCost >= this.config.maxPerDay) {
+      violations.push({
+        policy: this.name,
+        severity: "critical",
+        message: `Daily cost limit reached: $${dailyCost.toFixed(4)}`
+      });
+    }
+
+    return {
+      allowed: violations.length === 0,
+      violations
+    };
+  }
+
+  async validateOutput(output: AgentOutput): Promise<PolicyDecision> {
+    const cost = output.costUsd ?? 0;
+
+    if (cost > this.config.maxPerRequest) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: this.name,
+          severity: "high",
+          message: `Request cost $${cost.toFixed(4)} exceeds limit $${this.config.maxPerRequest}`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  recordCost(sessionId: string, cost: number): void {
+    const today = new Date().toISOString().split("T")[0];
+    this.sessionCosts.set(sessionId, (this.sessionCosts.get(sessionId) ?? 0) + cost);
+    this.dailyCosts.set(today, (this.dailyCosts.get(today) ?? 0) + cost);
+  }
+
+  getSessionCost(sessionId: string): number {
+    return this.sessionCosts.get(sessionId) ?? 0;
+  }
+
+  getDailyCost(): number {
+    const today = new Date().toISOString().split("T")[0];
+    return this.dailyCosts.get(today) ?? 0;
+  }
+}
+
+export interface TokenPolicyConfig {
+  maxPerRequest: number;
+  maxPerSession: number;
+  reserveForResponse: number;
+}
+
+export class TokenPolicy implements Policy {
+  name = "TokenPolicy";
+  private sessionTokens: Map<string, number> = new Map();
+
+  constructor(private config: TokenPolicyConfig) {}
+
+  async validateInput(input: AgentInput): Promise<PolicyDecision> {
+    const sessionId = input.context?.sessionId as string | undefined;
+    const estimatedTokens = this.estimateTokens(input.userPrompt);
+    const sessionTokens = this.sessionTokens.get(sessionId ?? "default") ?? 0;
+
+    const violations: PolicyViolation[] = [];
+
+    if (estimatedTokens > this.config.maxPerRequest) {
+      violations.push({
+        policy: this.name,
+        severity: "high",
+        message: `Request would exceed token limit: ${estimatedTokens} > ${this.config.maxPerRequest}`
+      });
+    }
+
+    if (sessionTokens + estimatedTokens > this.config.maxPerSession) {
+      violations.push({
+        policy: this.name,
+        severity: "high",
+        message: `Session token limit would be exceeded`
+      });
+    }
+
+    return {
+      allowed: violations.length === 0,
+      violations
+    };
+  }
+
+  async validateOutput(output: AgentOutput): Promise<PolicyDecision> {
+    const tokens = output.usage?.totalTokens ?? 0;
+
+    if (tokens > this.config.maxPerRequest) {
+      return {
+        allowed: false,
+        violations: [{
+          policy: this.name,
+          severity: "medium",
+          message: `Token usage ${tokens} exceeds limit ${this.config.maxPerRequest}`
+        }]
+      };
+    }
+
+    return { allowed: true, violations: [] };
+  }
+
+  recordTokens(sessionId: string, tokens: number): void {
+    this.sessionTokens.set(sessionId, (this.sessionTokens.get(sessionId) ?? 0) + tokens);
+  }
+
+  private estimateTokens(text: string): number {
+    return Math.ceil(text.split(/\s+/).length * 1.3);
+  }
+}