rax-flow-core 0.2.0 → 2.0.1

package/src/plugins/plugin-system.ts

@@ -1,23 +1,112 @@
-import { AgentInput, AgentOutput, WorkflowNode } from "../types/contracts.js";
+import { AgentInput, AgentOutput, WorkflowNode, Intent, RoutedPlan } from "../types/contracts.js";
+
+export interface CacheHookParams {
+  key: string;
+  intent: Intent;
+  prompt: string;
+  cached: boolean;
+  cachedValue?: unknown;
+}
+
+export interface MemoryHookParams {
+  taskId: string;
+  nodeId?: string;
+  action: "read" | "write" | "search";
+  data?: unknown;
+  result?: unknown;
+}
+
+export interface FallbackHookParams {
+  taskId: string;
+  nodeId: string;
+  agent: string;
+  originalProvider: string;
+  fallbackProvider: string;
+  reason: string;
+  attempt: number;
+}
+
+export interface EscalationHookParams {
+  taskId: string;
+  nodeId: string;
+  agent: string;
+  fromTier: string;
+  toTier: string;
+  reason: string;
+  confidence: number;
+}
+
+export interface RoutingHookParams {
+  prompt: string;
+  intent: Intent;
+  route: RoutedPlan;
+  changedBy?: string;
+}
+
+export interface WorkflowHookParams {
+  taskId: string;
+  workflow: { id: string; nodes: WorkflowNode[] };
+  phase: "planned" | "executing" | "completed" | "failed";
+}
 
 export interface RaxPlugin {
   name: string;
+  priority?: number;
   setup?(): Promise<void> | void;
+  teardown?(): Promise<void> | void;
   beforeNode?(params: { node: WorkflowNode; input: AgentInput }): Promise<void> | void;
   afterNode?(params: { node: WorkflowNode; input: AgentInput; output: AgentOutput }): Promise<void> | void;
+  beforeCache?(params: CacheHookParams): Promise<CacheHookParams | void> | CacheHookParams | void;
+  afterCache?(params: CacheHookParams): Promise<void> | void;
+  beforeMemory?(params: MemoryHookParams): Promise<MemoryHookParams | void> | MemoryHookParams | void;
+  afterMemory?(params: MemoryHookParams): Promise<void> | void;
+  beforeFallback?(params: FallbackHookParams): Promise<FallbackHookParams | void> | FallbackHookParams | void;
+  afterFallback?(params: FallbackHookParams): Promise<void> | void;
+  onEscalation?(params: EscalationHookParams): Promise<void> | void;
+  onRouting?(params: RoutingHookParams): Promise<RoutingHookParams | void> | RoutingHookParams | void;
+  onWorkflowPhase?(params: WorkflowHookParams): Promise<void> | void;
+  onError?(params: { taskId: string; nodeId?: string; error: Error; phase: string }): Promise<void> | void;
 }
 
 export class PluginSystem {
   private readonly plugins: RaxPlugin[] = [];
+  private initialized = false;
 
   register(plugin: RaxPlugin): void {
     this.plugins.push(plugin);
+    this.plugins.sort((a, b) => (a.priority ?? 100) - (b.priority ?? 100));
+  }
+
+  unregister(pluginName: string): boolean {
+    const idx = this.plugins.findIndex(p => p.name === pluginName);
+    if (idx >= 0) {
+      this.plugins.splice(idx, 1);
+      return true;
+    }
+    return false;
   }
 
   async initialize(): Promise<void> {
+    if (this.initialized) return;
     for (const plugin of this.plugins) {
       await plugin.setup?.();
     }
+    this.initialized = true;
+  }
+
+  async teardown(): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.teardown?.();
+    }
+    this.initialized = false;
+  }
+
+  getPlugin(name: string): RaxPlugin | undefined {
+    return this.plugins.find(p => p.name === name);
+  }
+
+  listPlugins(): RaxPlugin[] {
+    return [...this.plugins];
   }
 
   async runBeforeNode(params: { node: WorkflowNode; input: AgentInput }): Promise<void> {
@@ -31,4 +120,76 @@ export class PluginSystem {
       await plugin.afterNode?.(params);
     }
   }
+
+  async runBeforeCache(params: CacheHookParams): Promise<CacheHookParams> {
+    let result = params;
+    for (const plugin of this.plugins) {
+      const modified = await plugin.beforeCache?.(result);
+      if (modified) result = modified;
+    }
+    return result;
+  }
+
+  async runAfterCache(params: CacheHookParams): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.afterCache?.(params);
+    }
+  }
+
+  async runBeforeMemory(params: MemoryHookParams): Promise<MemoryHookParams> {
+    let result = params;
+    for (const plugin of this.plugins) {
+      const modified = await plugin.beforeMemory?.(result);
+      if (modified) result = modified;
+    }
+    return result;
+  }
+
+  async runAfterMemory(params: MemoryHookParams): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.afterMemory?.(params);
+    }
+  }
+
+  async runBeforeFallback(params: FallbackHookParams): Promise<FallbackHookParams> {
+    let result = params;
+    for (const plugin of this.plugins) {
+      const modified = await plugin.beforeFallback?.(result);
+      if (modified) result = modified;
+    }
+    return result;
+  }
+
+  async runAfterFallback(params: FallbackHookParams): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.afterFallback?.(params);
+    }
+  }
+
+  async runOnEscalation(params: EscalationHookParams): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.onEscalation?.(params);
+    }
+  }
+
+  async runOnRouting(params: RoutingHookParams): Promise<RoutingHookParams> {
+    let result = params;
+    for (const plugin of this.plugins) {
+      const modified = await plugin.onRouting?.(result);
+      if (modified) result = modified;
+    }
+    return result;
+  }
+
+  async runOnWorkflowPhase(params: WorkflowHookParams): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.onWorkflowPhase?.(params);
+    }
+  }
+
+  async runOnError(params: { taskId: string; nodeId?: string; error: Error; phase: string }): Promise<void> {
+    for (const plugin of this.plugins) {
+      await plugin.onError?.(params);
+    }
+  }
 }