# Plexus Method Wrappers

## Problem

Schema computation is slow because `plugin_children()` calls `plugin_schema()` directly on each activation, bypassing any caching or interception. When `plexus.schema` is called:

1. `Activation::plugin_schema(&plexus)` is called
2. This calls `plugin_children()` which iterates all 11+ activations
3. Each `a.plugin_schema()` recomputes the full schema from scratch
4. No caching, no interception point

More broadly, we want a way to wrap any method with arbitrary computation (caching, logging, auth, transforms) without modifying the method implementation.

## Solution: Typed Method Wrappers

Instead of untyped hooks with `Value` params, use **typed wrappers** generated by the macro. Each method gets a wrapper API with compile-time type checking.

**Key insight**: The macro knows the full method signature at compile time, so it can generate type-safe wrapper infrastructure.

## Design

### Usage (What We Want)

```rust
// Wrap schema method with caching
cone.wrap_schema(|inner| {
    // Check cache first
    if let Some(cached) = cache.get("cone") {
        return cached.clone();
    }

    // Call the actual method
    let result = inner();

    // Cache and return
    cache.insert("cone", result.clone());
    result
});

// Wrap chat method with logging
cone.wrap_chat(|inner, identifier, prompt| {
    println!("Chat starting for {:?}", identifier);

    let stream = inner(identifier.clone(), prompt.clone());

    // Could transform the stream
    stream.inspect(|event| println!("Event: {:?}", event))
});

// Wrap with auth check
health.wrap_check(|inner| {
    if !is_authorized() {
        return HealthStatus::unauthorized();
    }
    inner()
});
```

### Generated Types

For each method, the macro generates type aliases based on the method signature:

```rust
// Given this method:
#[hub_method]
async fn schema(&self) -> PluginSchema { ... }

// Macro generates:
pub type ConeSchemaFn = Box<dyn Fn() -> PluginSchema + Send + Sync>;
pub type ConeSchemaWrapper = Box<dyn Fn(ConeSchemaFn) -> PluginSchema + Send + Sync>;
```

```rust
// Given this method:
#[hub_method]
async fn chat(&self, identifier: ConeIdentifier, prompt: String) -> impl Stream<Item = ChatEvent> { ... }

// Macro generates:
pub type ConeChatFn = Box<dyn Fn(ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync>;
pub type ConeChatWrapper = Box<dyn Fn(ConeChatFn, ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync>;
```

### Generated Wrapper Storage

```rust
// Generated struct to hold all wrappers for an activation
pub struct ConeWrappers {
    pub schema: Option<ConeSchemaWrapper>,
    pub chat: Option<ConeChatWrapper>,
    pub list: Option<ConeListWrapper>,
    pub get: Option<ConeGetWrapper>,
    // ... one per method
}

impl Default for ConeWrappers {
    fn default() -> Self {
        Self {
            schema: None,
            chat: None,
            list: None,
            get: None,
        }
    }
}
```

### Generated Registration Methods

```rust
impl Cone {
    /// Wrap the schema method with custom logic
    pub fn wrap_schema<F>(&mut self, wrapper: F) -> &mut Self
    where
        F: Fn(ConeSchemaFn) -> PluginSchema + Send + Sync + 'static
    {
        self.wrappers.schema = Some(Box::new(wrapper));
        self
    }

    /// Wrap the chat method with custom logic
    pub fn wrap_chat<F>(&mut self, wrapper: F) -> &mut Self
    where
        F: Fn(ConeChatFn, ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync + 'static
    {
        self.wrappers.chat = Some(Box::new(wrapper));
        self
    }
}
```

### Generated Dispatch with Wrappers

The `call()` method checks for wrappers and invokes them:

```rust
// Generated call() dispatch
async fn call(&self, method: &str, params: Value) -> Result<PlexusStream, PlexusError> {
    match method {
        "schema" => {
            // Create the "inner" function that calls the real method
            let inner: ConeSchemaFn = Box::new({
                let self_clone = self.clone();
                move || self_clone.schema_impl()
            });

            // Check for wrapper
            let result = if let Some(ref wrapper) = self.wrappers.schema {
                wrapper(inner)
            } else {
                inner()
            };

            Ok(wrap_as_stream(result))
        }

        "chat" => {
            let params: ChatParams = serde_json::from_value(params)?;

            let inner: ConeChatFn = Box::new({
                let self_clone = self.clone();
                move |id, prompt| self_clone.chat_impl(id, prompt)
            });

            let result = if let Some(ref wrapper) = self.wrappers.chat {
                wrapper(inner, params.identifier, params.prompt)
            } else {
                inner(params.identifier, params.prompt)
            };

            Ok(wrap_as_plexus_stream(result))
        }

        _ => Err(PlexusError::MethodNotFound { ... })
    }
}
```

## Macro Implementation Sketch

### Input (User Code)

```rust
#[hub_methods(namespace = "cone", version = "1.0.0", wrappers)]
impl Cone {
    #[hub_method]
    async fn schema(&self) -> PluginSchema {
        // actual implementation
    }

    #[hub_method]
    async fn chat(&self, identifier: ConeIdentifier, prompt: String)
        -> impl Stream<Item = ChatEvent>
    {
        // actual implementation
    }
}
```

### Output (Generated Code)

```rust
// ===== Type Aliases =====

pub type ConeSchemaFn = Box<dyn Fn() -> PluginSchema + Send + Sync>;
pub type ConeSchemaWrapper = Box<dyn Fn(ConeSchemaFn) -> PluginSchema + Send + Sync>;

pub type ConeChatFn = Box<dyn Fn(ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync>;
pub type ConeChatWrapper = Box<dyn Fn(ConeChatFn, ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync>;

// ===== Wrapper Storage =====

pub struct ConeWrappers {
    pub schema: Option<ConeSchemaWrapper>,
    pub chat: Option<ConeChatWrapper>,
}

impl Default for ConeWrappers {
    fn default() -> Self {
        Self { schema: None, chat: None }
    }
}

// ===== Registration Methods =====

impl Cone {
    pub fn wrap_schema<F>(&mut self, f: F) -> &mut Self
    where F: Fn(ConeSchemaFn) -> PluginSchema + Send + Sync + 'static
    {
        self.wrappers.schema = Some(Box::new(f));
        self
    }

    pub fn wrap_chat<F>(&mut self, f: F) -> &mut Self
    where F: Fn(ConeChatFn, ConeIdentifier, String) -> Pin<Box<dyn Stream<Item = ChatEvent> + Send>> + Send + Sync + 'static
    {
        self.wrappers.chat = Some(Box::new(f));
        self
    }
}

// ===== Activation Trait Impl =====

impl Activation for Cone {
    // ... other methods ...

    async fn call(&self, method: &str, params: Value) -> Result<PlexusStream, PlexusError> {
        match method {
            "schema" => {
                let self_clone = self.clone();
                let inner: ConeSchemaFn = Box::new(move || self_clone.schema_impl());

                let result = match &self.wrappers.schema {
                    Some(wrapper) => wrapper(inner),
                    None => inner(),
                };

                Ok(wrap_as_stream(result))
            }

            "chat" => {
                let p: ChatParams = serde_json::from_value(params)
                    .map_err(|e| PlexusError::InvalidParams(e.to_string()))?;

                let self_clone = self.clone();
                let inner: ConeChatFn = Box::new(move |id, prompt| {
                    Box::pin(self_clone.chat_impl(id, prompt))
                });

                let result = match &self.wrappers.chat {
                    Some(wrapper) => wrapper(inner, p.identifier, p.prompt),
                    None => inner(p.identifier, p.prompt),
                };

                Ok(wrap_as_plexus_stream(result))
            }

            _ => Err(PlexusError::MethodNotFound {
                activation: self.namespace().to_string(),
                method: method.to_string(),
            })
        }
    }
}
```

## Usage Examples

### Caching Wrapper

```rust
use std::sync::{Arc, RwLock};
use std::collections::HashMap;

let cache: Arc<RwLock<HashMap<String, PluginSchema>>> = Arc::new(RwLock::new(HashMap::new()));

let cache_clone = cache.clone();
cone.wrap_schema(move |inner| {
    // Check cache
    if let Some(schema) = cache_clone.read().unwrap().get("cone") {
        return schema.clone();
    }

    // Compute
    let schema = inner();

    // Store
    cache_clone.write().unwrap().insert("cone".to_string(), schema.clone());

    schema
});
```

### Logging Wrapper

```rust
cone.wrap_chat(|inner, identifier, prompt| {
    let start = std::time::Instant::now();
    println!("[CHAT] Starting for {:?}", identifier);

    let stream = inner(identifier, prompt);

    // Wrap stream to log completion
    Box::pin(stream.inspect(move |_| {
        println!("[CHAT] Event received, elapsed: {:?}", start.elapsed());
    }))
});
```

### Auth Wrapper

```rust
health.wrap_check(|inner| {
    if !current_user_authorized() {
        return HealthStatus {
            status: "unauthorized".into(),
            uptime_seconds: 0,
        };
    }
    inner()
});
```

### Retry Wrapper

```rust
cone.wrap_chat(|inner, identifier, prompt| {
    let mut attempts = 0;
    loop {
        attempts += 1;
        let result = inner(identifier.clone(), prompt.clone());
        // In practice, you'd check for errors in the stream
        if attempts >= 3 {
            return result;
        }
        // Could add backoff logic here
    }
});
```

## Files to Modify

1. **hub-macro/src/codegen/wrappers.rs** (new) - Generate wrapper types and registration methods
2. **hub-macro/src/codegen/activation.rs** - Add `wrappers` flag, generate dispatch with wrapper checks
3. **hub-macro/src/codegen/method_enum.rs** - May need updates for type extraction

## Considerations

### Wrapper Field on Struct

The activation struct needs a `wrappers` field:

```rust
pub struct Cone {
    // ... existing fields ...
    wrappers: ConeWrappers,
}
```

This could be:
- Added by the macro (requires the macro to modify the struct definition)
- Required by convention (user adds it manually)
- Stored externally in a registry

### Clone Requirements

For the `inner` closure to capture `self`, the activation needs to be `Clone` or wrapped in `Arc`.

### Async Wrappers

For async methods, the wrapper return type needs to handle futures:

```rust
pub type ConeSchemaWrapper = Box<dyn Fn(ConeSchemaFn) -> BoxFuture<'static, PluginSchema> + Send + Sync>;
```

### Multiple Wrappers (Chaining)

Currently supports one wrapper per method. Could extend to a `Vec` of wrappers that chain:

```rust
pub struct ConeWrappers {
    pub schema: Vec<ConeSchemaWrapper>,  // Chain of wrappers
}
```

Each wrapper receives the next wrapper's `inner`, forming a middleware chain.

## Alternative: Trait-Based Wrappers

Instead of closures, could use traits:

```rust
pub trait SchemaWrapper: Send + Sync {
    fn wrap(&self, inner: impl Fn() -> PluginSchema) -> PluginSchema;
}

struct CachingSchemaWrapper { cache: Arc<RwLock<HashMap<String, PluginSchema>>> }

impl SchemaWrapper for CachingSchemaWrapper {
    fn wrap(&self, inner: impl Fn() -> PluginSchema) -> PluginSchema {
        // caching logic
    }
}
```

This allows more complex wrapper implementations with state, but closures are more ergonomic for simple cases.
