pe-core 0.1.0

Core types for Potential Expectations — messages, channels, state, traits
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
//! 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"
        );
    }
}