plugin_demo/
plugin_demo.rs

1// CPEX Plugin Demo
2//
3// Demonstrates how to:
4//   1. Define hook types and payloads
5//   2. Build plugins that implement HookHandler
6//   3. Create plugin factories for config-driven loading
7//   4. Load a YAML config with routing rules
8//   5. Invoke hooks with MetaExtension for route resolution
9//
10// Run with: cargo run --example plugin_demo
11
12use std::sync::Arc;
13
14use async_trait::async_trait;
15use cpex_core::context::PluginContext;
16use cpex_core::error::{PluginError, PluginViolation};
17use cpex_core::executor::PipelineResult;
18use cpex_core::factory::{PluginFactory, PluginInstance};
19use cpex_core::hooks::adapter::TypedHandlerAdapter;
20use cpex_core::hooks::payload::{Extensions, MetaExtension};
21use cpex_core::hooks::trait_def::{HookHandler, HookTypeDef, PluginResult};
22use cpex_core::manager::PluginManager;
23use cpex_core::plugin::{Plugin, PluginConfig};
24
25// ---------------------------------------------------------------------------
26// Step 1: Define a payload and hook type
27// ---------------------------------------------------------------------------
28
29/// The payload carried through the tool_pre_invoke hook.
30#[derive(Debug, Clone)]
31struct ToolInvokePayload {
32    tool_name: String,
33    user: String,
34    arguments: String,
35}
36cpex_core::impl_plugin_payload!(ToolInvokePayload);
37
38/// Hook type for tool_pre_invoke — runs before a tool executes.
39struct ToolPreInvoke;
40impl HookTypeDef for ToolPreInvoke {
41    type Payload = ToolInvokePayload;
42    type Result = PluginResult<ToolInvokePayload>;
43    const NAME: &'static str = "tool_pre_invoke";
44}
45
46/// Hook type for tool_post_invoke — runs after a tool executes.
47struct ToolPostInvoke;
48impl HookTypeDef for ToolPostInvoke {
49    type Payload = ToolInvokePayload;
50    type Result = PluginResult<ToolInvokePayload>;
51    const NAME: &'static str = "tool_post_invoke";
52}
53
54// ---------------------------------------------------------------------------
55// Step 2: Build plugins
56// ---------------------------------------------------------------------------
57
58/// Identity resolver — checks that a user is present.
59struct IdentityResolver {
60    cfg: PluginConfig,
61}
62
63#[async_trait]
64impl Plugin for IdentityResolver {
65    fn config(&self) -> &PluginConfig {
66        &self.cfg
67    }
68    async fn initialize(&self) -> Result<(), Box<PluginError>> {
69        println!("  [identity-resolver] initialized");
70        Ok(())
71    }
72    async fn shutdown(&self) -> Result<(), Box<PluginError>> {
73        println!("  [identity-resolver] shutdown");
74        Ok(())
75    }
76}
77
78impl HookHandler<ToolPreInvoke> for IdentityResolver {
79    async fn handle(
80        &self,
81        payload: &ToolInvokePayload,
82        _extensions: &Extensions,
83        _ctx: &mut PluginContext,
84    ) -> PluginResult<ToolInvokePayload> {
85        if payload.user.is_empty() {
86            println!("  [identity-resolver] DENIED: no user identity");
87            return PluginResult::deny(PluginViolation::new(
88                "no_identity",
89                "User identity is required",
90            ));
91        }
92        println!(
93            "  [identity-resolver] OK: user '{}' identified",
94            payload.user
95        );
96        PluginResult::allow()
97    }
98}
99
100impl HookHandler<ToolPostInvoke> for IdentityResolver {
101    async fn handle(
102        &self,
103        payload: &ToolInvokePayload,
104        _extensions: &Extensions,
105        _ctx: &mut PluginContext,
106    ) -> PluginResult<ToolInvokePayload> {
107        println!(
108            "  [identity-resolver] post-invoke: user '{}' completed '{}'",
109            payload.user, payload.tool_name
110        );
111        PluginResult::allow()
112    }
113}
114
115/// PII guard — blocks access to sensitive tools without clearance.
116struct PiiGuard {
117    cfg: PluginConfig,
118}
119
120#[async_trait]
121impl Plugin for PiiGuard {
122    fn config(&self) -> &PluginConfig {
123        &self.cfg
124    }
125    // initialize() and shutdown() use defaults — no setup needed
126}
127
128impl HookHandler<ToolPreInvoke> for PiiGuard {
129    async fn handle(
130        &self,
131        payload: &ToolInvokePayload,
132        _extensions: &Extensions,
133        ctx: &mut PluginContext,
134    ) -> PluginResult<ToolInvokePayload> {
135        // Check if the user has PII clearance (simulated via context)
136        let has_clearance = ctx
137            .get_global("pii_clearance")
138            .and_then(|v| v.as_bool())
139            .unwrap_or(false);
140
141        if !has_clearance {
142            println!(
143                "  [pii-guard] DENIED: user '{}' lacks PII clearance for '{}'",
144                payload.user, payload.tool_name
145            );
146            return PluginResult::deny(PluginViolation::new(
147                "pii_access_denied",
148                "PII clearance required",
149            ));
150        }
151
152        println!(
153            "  [pii-guard] OK: user '{}' has PII clearance",
154            payload.user
155        );
156        PluginResult::allow()
157    }
158}
159
160/// Audit logger — logs all tool invocations (fire-and-forget).
161struct AuditLogger {
162    cfg: PluginConfig,
163}
164
165#[async_trait]
166impl Plugin for AuditLogger {
167    fn config(&self) -> &PluginConfig {
168        &self.cfg
169    }
170    // initialize() and shutdown() use defaults — no setup needed
171}
172
173impl HookHandler<ToolPreInvoke> for AuditLogger {
174    async fn handle(
175        &self,
176        payload: &ToolInvokePayload,
177        _extensions: &Extensions,
178        _ctx: &mut PluginContext,
179    ) -> PluginResult<ToolInvokePayload> {
180        println!(
181            "  [audit-logger] LOG: user='{}' tool='{}' args='{}'",
182            payload.user, payload.tool_name, payload.arguments
183        );
184        PluginResult::allow()
185    }
186}
187
188impl HookHandler<ToolPostInvoke> for AuditLogger {
189    async fn handle(
190        &self,
191        payload: &ToolInvokePayload,
192        _extensions: &Extensions,
193        _ctx: &mut PluginContext,
194    ) -> PluginResult<ToolInvokePayload> {
195        println!(
196            "  [audit-logger] LOG: post-invoke user='{}' tool='{}'",
197            payload.user, payload.tool_name
198        );
199        PluginResult::allow()
200    }
201}
202
203// ---------------------------------------------------------------------------
204// Awaiting plugin example — RemoteAuthz
205// ---------------------------------------------------------------------------
206//
207// `HookHandler<H>` is async by design — `handle` is `async fn`.
208// Plugins that don't need to `.await` anything still write
209// `async fn handle` and return synchronously; this plugin shows the
210// other direction, where the body genuinely awaits per-invocation
211// work. The realistic version would call a remote authz service
212// (gRPC, HTTP, OPA, Cedarling, etc.); here we simulate the network
213// round-trip with a small `tokio::time::sleep` so the demo runs
214// offline.
215//
216// Key things this shows:
217//   1. Per-request latency state is *cached at init* — the handler
218//      consults the in-memory ACL and only "calls out" on a miss.
219//      Hot-path I/O is the most common source of latency regressions
220//      in plugins, so prefer initialize-time loading wherever you can.
221//   2. Registration uses the exact same factory pattern as any other
222//      plugin — `TypedHandlerAdapter::<H, _>` and the same
223//      `register_factory` call. There is no separate async path.
224struct RemoteAuthz {
225    cfg: PluginConfig,
226    /// ACL "fetched" at init. Populated in Plugin::initialize.
227    allowed_users: tokio::sync::RwLock<std::collections::HashSet<String>>,
228}
229
230#[async_trait]
231impl Plugin for RemoteAuthz {
232    fn config(&self) -> &PluginConfig {
233        &self.cfg
234    }
235    /// Pretend we're loading the ACL from a remote service. In a real
236    /// plugin this would be `client.fetch_acl().await`; we simulate
237    /// the round-trip with a small sleep so the demo runs offline.
238    async fn initialize(&self) -> Result<(), Box<PluginError>> {
239        tokio::time::sleep(std::time::Duration::from_millis(2)).await;
240        let mut acl = self.allowed_users.write().await;
241        acl.extend(["alice", "bob"].iter().map(|s| s.to_string()));
242        println!(
243            "  [remote-authz] initialized — ACL cached ({} users)",
244            acl.len()
245        );
246        Ok(())
247    }
248    async fn shutdown(&self) -> Result<(), Box<PluginError>> {
249        println!("  [remote-authz] shutdown");
250        Ok(())
251    }
252}
253
254impl HookHandler<ToolPreInvoke> for RemoteAuthz {
255    async fn handle(
256        &self,
257        payload: &ToolInvokePayload,
258        _extensions: &Extensions,
259        _ctx: &mut PluginContext,
260    ) -> PluginResult<ToolInvokePayload> {
261        // Cache hit path — fast.
262        let acl = self.allowed_users.read().await;
263        if acl.contains(&payload.user) {
264            println!(
265                "  [remote-authz] OK (cache hit): user '{}' allowed",
266                payload.user
267            );
268            return PluginResult::allow();
269        }
270        drop(acl); // release read lock before the fake remote call
271                   // Cache miss path — simulate a remote authz check. In a real
272                   // plugin this is where you'd `.await` a gRPC or HTTP call.
273                   // The latency cost is real and shows up on the request path.
274        tokio::time::sleep(std::time::Duration::from_millis(1)).await;
275        println!(
276            "  [remote-authz] DENIED (cache miss + remote check): user '{}'",
277            payload.user
278        );
279        PluginResult::deny(PluginViolation::new(
280            "remote_authz_denied",
281            format!("User '{}' not in remote ACL", payload.user),
282        ))
283    }
284}
285
286// ---------------------------------------------------------------------------
287// Step 3: Create plugin factories
288// ---------------------------------------------------------------------------
289
290struct IdentityFactory;
291impl PluginFactory for IdentityFactory {
292    fn create(&self, config: &PluginConfig) -> Result<PluginInstance, Box<PluginError>> {
293        let plugin = Arc::new(IdentityResolver {
294            cfg: config.clone(),
295        });
296        Ok(PluginInstance {
297            plugin: plugin.clone(),
298            handlers: vec![
299                (
300                    "tool_pre_invoke",
301                    Arc::new(TypedHandlerAdapter::<ToolPreInvoke, _>::new(plugin.clone())),
302                ),
303                (
304                    "tool_post_invoke",
305                    Arc::new(TypedHandlerAdapter::<ToolPostInvoke, _>::new(plugin)),
306                ),
307            ],
308        })
309    }
310}
311
312struct PiiGuardFactory;
313impl PluginFactory for PiiGuardFactory {
314    fn create(&self, config: &PluginConfig) -> Result<PluginInstance, Box<PluginError>> {
315        let plugin = Arc::new(PiiGuard {
316            cfg: config.clone(),
317        });
318        Ok(PluginInstance {
319            plugin: plugin.clone(),
320            handlers: vec![(
321                "tool_pre_invoke",
322                Arc::new(TypedHandlerAdapter::<ToolPreInvoke, _>::new(plugin)),
323            )],
324        })
325    }
326}
327
328struct AuditLoggerFactory;
329impl PluginFactory for AuditLoggerFactory {
330    fn create(&self, config: &PluginConfig) -> Result<PluginInstance, Box<PluginError>> {
331        let plugin = Arc::new(AuditLogger {
332            cfg: config.clone(),
333        });
334        Ok(PluginInstance {
335            plugin: plugin.clone(),
336            handlers: vec![
337                (
338                    "tool_pre_invoke",
339                    Arc::new(TypedHandlerAdapter::<ToolPreInvoke, _>::new(plugin.clone())),
340                ),
341                (
342                    "tool_post_invoke",
343                    Arc::new(TypedHandlerAdapter::<ToolPostInvoke, _>::new(plugin)),
344                ),
345            ],
346        })
347    }
348}
349
350/// Factory for the async plugin. Note the factory body is identical
351/// in shape to the sync factories above — `TypedHandlerAdapter` and
352/// the `register_factory` path don't care that the underlying handler
353/// is async. The framework hides the choice.
354struct RemoteAuthzFactory;
355impl PluginFactory for RemoteAuthzFactory {
356    fn create(&self, config: &PluginConfig) -> Result<PluginInstance, Box<PluginError>> {
357        let plugin = Arc::new(RemoteAuthz {
358            cfg: config.clone(),
359            allowed_users: tokio::sync::RwLock::new(std::collections::HashSet::new()),
360        });
361        Ok(PluginInstance {
362            plugin: plugin.clone(),
363            handlers: vec![(
364                "tool_pre_invoke",
365                Arc::new(TypedHandlerAdapter::<ToolPreInvoke, _>::new(plugin)),
366            )],
367        })
368    }
369}
370
371// ---------------------------------------------------------------------------
372// Step 4: Build extensions with MetaExtension for routing
373// ---------------------------------------------------------------------------
374
375fn make_tool_extensions(tool_name: &str, tags: &[&str]) -> Extensions {
376    Extensions {
377        meta: Some(Arc::new(MetaExtension {
378            entity_type: Some("tool".into()),
379            entity_name: Some(tool_name.into()),
380            tags: tags.iter().map(|s| s.to_string()).collect(),
381            ..Default::default()
382        })),
383        ..Default::default()
384    }
385}
386
387// ---------------------------------------------------------------------------
388// Helper to print results
389// ---------------------------------------------------------------------------
390
391fn print_result(_label: &str, result: &PipelineResult) {
392    if result.continue_processing {
393        println!("  Result: ALLOWED");
394    } else {
395        let violation = result.violation.as_ref().unwrap();
396        println!(
397            "  Result: DENIED by '{}' — {} [{}]",
398            violation.plugin_name.as_deref().unwrap_or("unknown"),
399            violation.reason,
400            violation.code,
401        );
402    }
403    println!();
404}
405
406// ---------------------------------------------------------------------------
407// Step 5: Main — load config, invoke hooks, see results
408// ---------------------------------------------------------------------------
409
410#[tokio::main]
411async fn main() {
412    println!("=== CPEX Plugin Demo ===\n");
413
414    // --- Load config from YAML file ---
415    let config_path = "crates/cpex-core/examples/plugin_demo.yaml";
416    println!("--- Loading config from {} ---\n", config_path);
417    let yaml = std::fs::read_to_string(config_path)
418        .unwrap_or_else(|e| panic!("Failed to read {}: {}", config_path, e));
419    let cpex_config = cpex_core::config::parse_config(&yaml).unwrap();
420
421    let mgr = PluginManager::default();
422    mgr.register_factory("builtin/identity", Box::new(IdentityFactory));
423    mgr.register_factory("builtin/pii", Box::new(PiiGuardFactory));
424    mgr.register_factory("builtin/audit", Box::new(AuditLoggerFactory));
425    mgr.register_factory("builtin/remote_authz", Box::new(RemoteAuthzFactory));
426    mgr.load_config(cpex_config).unwrap();
427
428    println!("\n--- Initializing plugins ---\n");
429    mgr.initialize().await.unwrap();
430
431    println!("\nPlugins loaded: {}", mgr.plugin_count());
432    println!(
433        "Hooks registered: tool_pre_invoke={}, tool_post_invoke={}\n",
434        mgr.has_hooks_for("tool_pre_invoke"),
435        mgr.has_hooks_for("tool_post_invoke"),
436    );
437
438    // --- Scenario 1: PII tool without clearance ---
439    println!("=== Scenario 1: get_compensation (PII tool, no clearance) ===\n");
440    let payload = ToolInvokePayload {
441        tool_name: "get_compensation".into(),
442        user: "alice".into(),
443        arguments: "employee_id=42".into(),
444    };
445    let ext = make_tool_extensions("get_compensation", &[]);
446    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
447    print_result("get_compensation (no clearance)", &result);
448    // Wait for any fire-and-forget tasks
449    bg.wait_for_background_tasks().await;
450
451    // --- Scenario 2: PII tool with clearance ---
452    println!("=== Scenario 2: get_compensation (PII tool, with clearance) ===\n");
453    let payload = ToolInvokePayload {
454        tool_name: "get_compensation".into(),
455        user: "alice".into(),
456        arguments: "employee_id=42".into(),
457    };
458    let ext = make_tool_extensions("get_compensation", &[]);
459    // Simulate clearance by pre-populating global_state
460    // (In production, an earlier hook would set this from a token claim)
461    let mut ctx_table = cpex_core::context::PluginContextTable::new();
462    ctx_table
463        .global_state
464        .insert("pii_clearance".into(), serde_json::Value::Bool(true));
465    let (result, bg) = mgr
466        .invoke::<ToolPreInvoke>(payload, ext, Some(ctx_table))
467        .await;
468    print_result("get_compensation (with clearance)", &result);
469    bg.wait_for_background_tasks().await;
470
471    // Now call post-invoke — threads the context table from pre-invoke
472    println!("  --- post-invoke for get_compensation ---\n");
473    let payload = ToolInvokePayload {
474        tool_name: "get_compensation".into(),
475        user: "alice".into(),
476        arguments: "employee_id=42".into(),
477    };
478    let ext = make_tool_extensions("get_compensation", &[]);
479    let (post_result, bg) = mgr
480        .invoke::<ToolPostInvoke>(payload, ext, Some(result.context_table))
481        .await;
482    print_result("get_compensation post-invoke", &post_result);
483    bg.wait_for_background_tasks().await;
484
485    // --- Scenario 3: Non-PII tool ---
486    println!("=== Scenario 3: list_departments (non-PII tool) ===\n");
487    let payload = ToolInvokePayload {
488        tool_name: "list_departments".into(),
489        user: "bob".into(),
490        arguments: "".into(),
491    };
492    let ext = make_tool_extensions("list_departments", &[]);
493    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
494    print_result("list_departments", &result);
495    bg.wait_for_background_tasks().await;
496
497    // --- Scenario 4: Unknown tool (wildcard route) ---
498    println!("=== Scenario 4: some_other_tool (wildcard route) ===\n");
499    let payload = ToolInvokePayload {
500        tool_name: "some_other_tool".into(),
501        user: "charlie".into(),
502        arguments: "foo=bar".into(),
503    };
504    let ext = make_tool_extensions("some_other_tool", &[]);
505    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
506    print_result("some_other_tool (wildcard)", &result);
507    bg.wait_for_background_tasks().await;
508
509    // --- Scenario 5: Awaiting plugin — cache hit ---
510    // RemoteAuthz's `handle` is `async fn` and reads from a tokio
511    // RwLock. Its initialize() pre-loaded an ACL containing "alice"
512    // and "bob"; this call exercises the cache-hit fast path.
513    println!("=== Scenario 5: query_external_data (async plugin, cache hit) ===\n");
514    let payload = ToolInvokePayload {
515        tool_name: "query_external_data".into(),
516        user: "alice".into(),
517        arguments: "dataset=sales".into(),
518    };
519    let ext = make_tool_extensions("query_external_data", &[]);
520    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
521    print_result("query_external_data (alice — in ACL)", &result);
522    bg.wait_for_background_tasks().await;
523
524    // --- Scenario 6: Awaiting plugin — cache miss path with .await ---
525    // "charlie" is not in the cached ACL, so RemoteAuthz takes the
526    // cache-miss branch and `.await`s a simulated remote call before
527    // denying.
528    println!("=== Scenario 6: query_external_data (async plugin, cache miss) ===\n");
529    let payload = ToolInvokePayload {
530        tool_name: "query_external_data".into(),
531        user: "charlie".into(),
532        arguments: "dataset=sales".into(),
533    };
534    let ext = make_tool_extensions("query_external_data", &[]);
535    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
536    print_result("query_external_data (charlie — not in ACL)", &result);
537    bg.wait_for_background_tasks().await;
538
539    // --- Scenario 7: No user identity ---
540    println!("=== Scenario 7: list_departments (no user identity) ===\n");
541    let payload = ToolInvokePayload {
542        tool_name: "list_departments".into(),
543        user: "".into(),
544        arguments: "".into(),
545    };
546    let ext = make_tool_extensions("list_departments", &[]);
547    let (result, bg) = mgr.invoke::<ToolPreInvoke>(payload, ext, None).await;
548    print_result("list_departments (no user)", &result);
549    bg.wait_for_background_tasks().await;
550
551    // --- Shutdown ---
552    println!("--- Shutting down ---\n");
553    mgr.shutdown().await;
554
555    println!("=== Demo complete ===");
556}
plugin_demo/plugin_demo.rs

plugin_demo/
plugin_demo.rs
