pe_core/

lobe.rs

//! Lobe — the cognitive processing unit.
//!
//! A lobe is like a brain region: it has a specialized function, runs in
//! parallel with other lobes, and feeds its output into synthesis. The
//! library provides the trait — users implement domain-specific lobes.
//!
//! Lobes compile down to `NodeFn<CognitiveState>` via the [`LobeNode`](crate::lobe_node::LobeNode)
//! bridge (in `lobe_node` module), so the cognitive graph uses the same
//! Pregel engine as the outer execution graph.

use crate::cognitive_memory::WorkingNote;
use crate::cognitive_signal::CognitiveSignal;
use crate::error::PeError;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::future::Future;
use std::path::PathBuf;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;

/// A cognitive processing unit — one lens on the agent's reasoning.
///
/// Lobes are the building blocks of the cognitive architecture.
/// Each lobe receives lean, restricted context (not the full conversation),
/// processes it through its specialized function, and returns structured output.
///
/// # Implementing a Lobe
///
/// ```ignore
/// struct CriticLobe { model: Arc<dyn LlmProvider> }
///
/// impl Lobe for CriticLobe {
///     fn name(&self) -> &str { "critic" }
///     fn should_activate(&self, _context: &LobeContext) -> bool { true }
///     fn priority(&self) -> u32 { 20 }
///     fn budget(&self) -> LobeBudget { LobeBudget::default() }
///     fn output_format(&self) -> LobeOutputFormat { LobeOutputFormat::FreeText }
///     fn process(&self, input: &LobeInput) -> LobeFuture {
///         let model = self.model.clone();
///         let prompt = input.input.clone();
///         Box::pin(async move {
///             // Call LLM with critic-focused prompt
///             Ok(LobeOutput::new("Risks identified: ...".into(), 0.85))
///         })
///     }
/// }
/// ```
pub trait Lobe: Send + Sync {
    /// Unique name for this lobe (used as key in stream_outputs).
    fn name(&self) -> &str;

    /// Whether this lobe should run for the current context.
    /// Returning `false` skips the lobe entirely (zero cost).
    /// Receives a lean [`LobeContext`], not the full `CognitiveState`.
    fn should_activate(&self, context: &LobeContext) -> bool;

    /// Execution priority — lower values run first in synthesis ordering.
    /// Does NOT affect parallel execution (all lobes run concurrently).
    fn priority(&self) -> u32;

    /// Resource limits for this lobe's execution.
    fn budget(&self) -> LobeBudget;

    /// What kind of output this lobe produces.
    fn output_format(&self) -> LobeOutputFormat;

    /// Process the input and produce output. This is the lobe's main work.
    fn process(&self, input: &LobeInput) -> LobeFuture;
}

/// Future type for lobe processing — async, Send, returns Result.
pub type LobeFuture = Pin<Box<dyn Future<Output = Result<LobeOutput, PeError>> + Send>>;

/// Lean, restricted input for a lobe.
///
/// Each lobe sees only what it needs — not the full conversation.
/// This keeps cognitive processing cheap and focused.
#[derive(Clone)]
pub struct LobeInput {
    /// The task/prompt being processed.
    pub input: String,

    /// Context slice relevant to this lobe.
    pub context: LobeContext,

    /// Recent working notes (for lobes that build on prior observations).
    pub notes: Vec<WorkingNote>,

    /// Optional runtime-owned services available to this lobe.
    ///
    /// When present, the lobe can request accountable runtime services
    /// such as bounded inspection through the same runtime substrate as
    /// the outer agent.
    pub runtime_services: Option<Arc<dyn LobeRuntimeServices>>,
}

impl std::fmt::Debug for LobeInput {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LobeInput")
            .field("input", &self.input)
            .field("context", &self.context)
            .field("notes", &self.notes)
            .field("has_runtime_services", &self.runtime_services.is_some())
            .finish()
    }
}

/// Restricted context for a lobe — not the full conversation.
///
/// Use [`from_cognitive_state`](LobeContext::from_cognitive_state) to build
/// from a full `CognitiveState`, or construct directly for testing.
#[derive(Debug, Clone, Default)]
pub struct LobeContext {
    /// Agent's self-description (from SelfModel.self_context, summarized).
    pub self_summary: Option<String>,

    /// Recent error history (for lobes that learn from failures).
    pub recent_errors: Vec<String>,

    /// Current confidence level (from matrix C value).
    pub confidence: f64,

    /// Current plan (if one exists).
    pub current_plan: Option<String>,

    /// Arbitrary metadata from the outer graph.
    pub metadata: HashMap<String, serde_json::Value>,
}

impl LobeContext {
    /// Build a lean context from a full [`CognitiveState`](crate::cognitive::CognitiveState).
    ///
    /// This is the canonical conversion — used by both [`LobeNode`](crate::lobe_node::LobeNode)
    /// and [`LobeRegistry`](crate::lobe_registry::LobeRegistry) to avoid duplication.
    pub fn from_cognitive_state(state: &crate::cognitive::CognitiveState) -> Self {
        let mut metadata = HashMap::new();
        metadata.insert(
            "working_notes_count".into(),
            serde_json::Value::from(state.working_notes.len()),
        );
        metadata.insert(
            "failure_records_count".into(),
            serde_json::Value::from(state.failure_records.len()),
        );
        Self {
            self_summary: None, // Populated by the cognitive graph runner from SelfModel
            recent_errors: state.error_history.clone(),
            confidence: state.confidence,
            current_plan: state.current_plan.clone(),
            metadata,
        }
    }
}

/// Output from a lobe's processing.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct LobeOutput {
    /// Name of the lobe that produced this output (for synthesizer labeling).
    #[serde(default)]
    pub lobe_name: String,

    /// The lobe's output content.
    pub content: String,

    /// How confident the lobe is in its output (0.0 to 1.0).
    pub confidence: f64,

    /// Signals to emit to the outer graph.
    #[serde(default)]
    pub signals: Vec<CognitiveSignal>,

    /// Lobe-specific metadata.
    #[serde(default)]
    pub metadata: HashMap<String, serde_json::Value>,
}

impl LobeOutput {
    /// Create a new lobe output with content and confidence.
    pub fn new(content: impl Into<String>, confidence: f64) -> Self {
        Self {
            lobe_name: String::new(),
            content: content.into(),
            confidence,
            signals: Vec::new(),
            metadata: HashMap::new(),
        }
    }

    /// Set the lobe name on this output.
    #[must_use]
    pub fn with_lobe_name(mut self, name: impl Into<String>) -> Self {
        self.lobe_name = name.into();
        self
    }

    /// Add a signal to emit.
    #[must_use]
    pub fn with_signal(mut self, signal: CognitiveSignal) -> Self {
        self.signals.push(signal);
        self
    }
}

/// Resource limits for a single lobe.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LobeBudget {
    /// Max tokens this lobe can consume.
    pub max_tokens: u32,

    /// Max wall-clock time for this lobe.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub max_duration: Option<Duration>,

    /// Whether this lobe streams LLM tokens incrementally.
    ///
    /// Default is `false` because lobes run in parallel — streaming multiple
    /// lobes concurrently causes SSE I/O thrash with no user-visible benefit
    /// (the synthesizer waits for all lobes anyway). Enable selectively for
    /// lobes that need token-level progress reporting.
    #[serde(default)]
    pub streaming: bool,
}

impl Default for LobeBudget {
    fn default() -> Self {
        Self {
            max_tokens: 500,
            max_duration: Some(Duration::from_secs(5)),
            streaming: false,
        }
    }
}

/// What kind of output a lobe produces.
///
/// Helps the synthesizer interpret outputs correctly.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[non_exhaustive]
pub enum LobeOutputFormat {
    /// Unstructured text (most common).
    FreeText,
    /// JSON-structured output.
    Structured,
    /// Numeric score (0.0 to 1.0).
    Score,
    /// Yes/no decision.
    Boolean,
    /// User-defined format.
    Custom(String),
}

/// When a lobe activates — used by the cognitive graph scheduler.
///
/// Currently informational. The `Lobe::should_activate()` method is the
/// runtime activation gate. This enum is available for future scheduling
/// optimizations (e.g., skip `should_activate()` call for `AlwaysOn` lobes).
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[non_exhaustive]
pub enum LobeActivation {
    /// Runs every cognitive cycle.
    AlwaysOn,
    /// Runs only when `should_activate()` returns true.
    OnDemand,
    /// Runs based on external trigger (e.g., low confidence).
    Conditional,
}

/// Runtime-owned services available to a lobe.
///
/// Defined in `pe-core` so lobes can depend on the interface without
/// depending on `pe-runtime`.
pub trait LobeRuntimeServices: Send + Sync {
    /// Run a bounded runtime-owned inspection.
    fn inspect(&self, request: LobeInspectionRequest) -> Result<LobeInspectionResult, PeError>;
}

/// Factory for creating source-attributed runtime service handles per lobe.
pub trait LobeRuntimeServiceFactory: Send + Sync {
    /// Build runtime services for the named lobe.
    fn for_lobe(&self, lobe_name: &str) -> Arc<dyn LobeRuntimeServices>;
}

/// Lobe-facing request for bounded runtime-owned inspection.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LobeInspectionRequest {
    pub root: PathBuf,
    pub allowed_roots: Vec<PathBuf>,
    pub max_files: Option<u64>,
    pub max_bytes: Option<u64>,
    pub max_depth: Option<usize>,
    pub include_contents: bool,
    pub include_extensions: Vec<String>,
    pub exclude_names: Vec<String>,
    pub exclude_path_prefixes: Vec<PathBuf>,
    pub max_preview_bytes_per_file: Option<u64>,
    pub skip_hidden: bool,
}

impl LobeInspectionRequest {
    /// Create a new request rooted at `root`.
    pub fn new(root: impl Into<PathBuf>) -> Self {
        Self {
            root: root.into(),
            allowed_roots: Vec::new(),
            max_files: None,
            max_bytes: None,
            max_depth: None,
            include_contents: false,
            include_extensions: Vec::new(),
            exclude_names: Vec::new(),
            exclude_path_prefixes: Vec::new(),
            max_preview_bytes_per_file: None,
            skip_hidden: false,
        }
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_allowed_roots(mut self, roots: impl IntoIterator<Item = PathBuf>) -> Self {
        self.allowed_roots = roots.into_iter().collect();
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_contents(mut self, include_contents: bool) -> Self {
        self.include_contents = include_contents;
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_extensions<I, S>(mut self, extensions: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.include_extensions = extensions
            .into_iter()
            .map(|ext| ext.as_ref().trim_start_matches('.').to_ascii_lowercase())
            .filter(|ext| !ext.is_empty())
            .collect();
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_excluded_names<I, S>(mut self, names: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: AsRef<str>,
    {
        self.exclude_names = names
            .into_iter()
            .map(|name| name.as_ref().trim().to_string())
            .filter(|name| !name.is_empty())
            .collect();
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_excluded_path_prefixes<I, P>(mut self, prefixes: I) -> Self
    where
        I: IntoIterator<Item = P>,
        P: Into<PathBuf>,
    {
        self.exclude_path_prefixes = prefixes.into_iter().map(Into::into).collect();
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_max_preview_bytes_per_file(mut self, max_bytes: u64) -> Self {
        self.max_preview_bytes_per_file = Some(max_bytes);
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_skip_hidden(mut self, skip_hidden: bool) -> Self {
        self.skip_hidden = skip_hidden;
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_max_files(mut self, max_files: u64) -> Self {
        self.max_files = Some(max_files);
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_max_bytes(mut self, max_bytes: u64) -> Self {
        self.max_bytes = Some(max_bytes);
        self
    }

    #[must_use = "builder methods return the modified builder"]
    pub fn with_max_depth(mut self, max_depth: usize) -> Self {
        self.max_depth = Some(max_depth);
        self
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LobeInspectionEntry {
    pub path: PathBuf,
    pub kind: LobeInspectionEntryKind,
    pub depth: usize,
    pub size_bytes: Option<u64>,
    pub content_preview: Option<String>,
    pub content_truncated: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum LobeInspectionEntryKind {
    Directory,
    File,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LobeInspectionResult {
    pub root: PathBuf,
    pub max_files: Option<u64>,
    pub max_bytes: Option<u64>,
    pub max_depth: Option<usize>,
    pub entries: Vec<LobeInspectionEntry>,
    pub files_seen: u64,
    pub bytes_read: u64,
    pub truncated: bool,
    pub truncation_reason: Option<String>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::cognitive_signal::CognitiveSignal;

    #[test]
    fn test_lobe_output_creation() {
        let output = LobeOutput::new("analysis result", 0.92);
        assert_eq!(output.content, "analysis result");
        assert!((output.confidence - 0.92).abs() < f64::EPSILON);
        assert!(output.signals.is_empty());
    }

    #[test]
    fn test_lobe_output_with_signal() {
        let output =
            LobeOutput::new("risky", 0.3).with_signal(CognitiveSignal::ProceedWithCaution {
                concern: "low confidence".into(),
            });
        assert_eq!(output.signals.len(), 1);
        assert!(output.signals[0].is_cautionary());
    }

    #[test]
    fn test_lobe_budget_defaults() {
        let budget = LobeBudget::default();
        assert_eq!(budget.max_tokens, 500);
        assert_eq!(budget.max_duration, Some(Duration::from_secs(5)));
    }

    #[test]
    fn test_lobe_output_format_variants() {
        let formats = vec![
            LobeOutputFormat::FreeText,
            LobeOutputFormat::Score,
            LobeOutputFormat::Boolean,
            LobeOutputFormat::Custom("risk_matrix".into()),
        ];
        for fmt in &formats {
            let json = serde_json::to_string(fmt).unwrap();
            let back: LobeOutputFormat = serde_json::from_str(&json).unwrap();
            assert_eq!(&back, fmt);
        }
    }

    #[test]
    fn test_lobe_input_construction() {
        let input = LobeInput {
            input: "analyze this code".into(),
            context: LobeContext {
                confidence: 0.7,
                current_plan: Some("review then test".into()),
                ..Default::default()
            },
            notes: vec![],
            runtime_services: None,
        };
        assert_eq!(input.input, "analyze this code");
        assert!((input.context.confidence - 0.7).abs() < f64::EPSILON);
    }

    #[test]
    fn test_lobe_budget_streaming_default_false() {
        let budget = LobeBudget::default();
        assert!(
            !budget.streaming,
            "default streaming must be false — parallel lobes + SSE = I/O thrash"
        );
    }

    #[test]
    fn test_lobe_budget_streaming_serialization() {
        // Round-trip with streaming = true
        let budget_on = LobeBudget {
            streaming: true,
            ..Default::default()
        };
        let json = serde_json::to_string(&budget_on).unwrap();
        let back: LobeBudget = serde_json::from_str(&json).unwrap();
        assert!(back.streaming);

        // Round-trip with streaming = false
        let budget_off = LobeBudget::default();
        let json = serde_json::to_string(&budget_off).unwrap();
        let back: LobeBudget = serde_json::from_str(&json).unwrap();
        assert!(!back.streaming);

        // Deserialization from JSON missing the field defaults to false
        let json_no_field = r#"{"max_tokens":500}"#;
        let back: LobeBudget = serde_json::from_str(json_no_field).unwrap();
        assert!(!back.streaming);
    }

    #[test]
    fn test_lobe_context_metadata_counts() {
        use crate::cognitive::CognitiveState;
        use crate::cognitive_memory::{NoteCategory, WorkingNote};
        use crate::self_model::FailureRecord;

        let state = CognitiveState {
            working_notes: vec![
                WorkingNote::new("note 1", NoteCategory::Discovery),
                WorkingNote::new("note 2", NoteCategory::Concern),
                WorkingNote::new("note 3", NoteCategory::Reflection),
            ],
            failure_records: vec![
                FailureRecord::new("db", "ALTER TABLE"),
                FailureRecord::new("api", "POST /users"),
            ],
            ..Default::default()
        };

        let ctx = LobeContext::from_cognitive_state(&state);
        assert_eq!(
            ctx.metadata.get("working_notes_count"),
            Some(&serde_json::json!(3)),
            "metadata must contain working_notes_count from CognitiveState"
        );
        assert_eq!(
            ctx.metadata.get("failure_records_count"),
            Some(&serde_json::json!(2)),
            "metadata must contain failure_records_count from CognitiveState"
        );
    }
}