Files
keep/src/meta_plugin/mod.rs
Andrew Phillips fed3722ef9 fix: Resolve compilation errors by refactoring imports and type annotations
Co-authored-by: aider (openai/andrew/openrouter/deepseek/deepseek-chat-v3.1) <aider@aider.chat>
2025-09-03 09:39:22 -03:00

337 lines
12 KiB
Rust

use log::debug;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Mutex;
use once_cell::sync::Lazy;
pub mod exec;
pub mod digest;
pub mod magic;
pub mod text;
pub mod read_time;
pub mod read_rate;
pub mod hostname;
pub mod cwd;
pub mod user;
pub mod shell;
pub mod shell_pid;
pub mod keep_pid;
pub mod env;
/// Represents metadata to be stored
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MetaData {
pub name: String,
pub value: String,
}
/// Response from meta plugin operations
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct MetaPluginResponse {
pub metadata: Vec<MetaData>,
pub is_finalized: bool,
}
/// Base implementation for meta plugins to reduce boilerplate
#[derive(Debug, Clone, Default)]
pub struct BaseMetaPlugin {
pub outputs: std::collections::HashMap<String, serde_yaml::Value>,
pub options: std::collections::HashMap<String, serde_yaml::Value>,
pub is_finalized: bool,
}
impl BaseMetaPlugin {
pub fn new() -> Self {
Self::default()
}
pub fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
pub fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
pub fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
pub fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
/// Helper function to initialize plugin options and outputs
pub fn initialize_plugin(
&mut self,
default_outputs: &[&str],
options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) {
// Set default outputs
for output_name in default_outputs {
self.outputs.insert(output_name.to_string(), serde_yaml::Value::String(output_name.to_string()));
}
// Apply provided options and outputs
if let Some(opts) = options {
self.options.extend(opts);
}
if let Some(outs) = outputs {
self.outputs.extend(outs);
}
}
}
impl MetaPlugin for BaseMetaPlugin {
fn meta_type(&self) -> MetaPluginType {
// This is a base implementation, so we need to return something
// This might not be used, but we need to satisfy the trait
MetaPluginType::Text
}
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.outputs
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.outputs
}
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
&self.options
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
&mut self.options
}
}
#[derive(Debug, Eq, PartialEq, Clone, strum::EnumIter, strum::Display, strum::EnumString, Serialize, Deserialize)]
#[strum(serialize_all = "snake_case", ascii_case_insensitive)]
pub enum MetaPluginType {
MagicFile,
Cwd,
Text,
User,
Shell,
ShellPid,
KeepPid,
Digest,
ReadTime,
ReadRate,
Hostname,
Exec,
Env,
}
/// Central function to handle metadata output with name mapping
/// outputs: HashMap where key is internal name, value is either custom name or "false" to disable
pub fn process_metadata_outputs(internal_name: &str, value: serde_yaml::Value, outputs: &std::collections::HashMap<String, serde_yaml::Value>) -> Option<MetaData> {
// Check if this output is disabled
if let Some(mapping) = outputs.get(internal_name) {
// Check for null to disable the output
if mapping.is_null() {
debug!("META: Skipping disabled output (null): {}", internal_name);
return None;
}
// Check for boolean false to disable the output
if let Some(false_val) = mapping.as_bool() {
if !false_val {
debug!("META: Skipping disabled output: {}", internal_name);
return None;
}
}
if let Some(custom_name) = mapping.as_str() {
// Convert the value to a string representation
let value_str = match &value {
serde_yaml::Value::Null => "null".to_string(),
serde_yaml::Value::Bool(b) => b.to_string(),
serde_yaml::Value::Number(n) => n.to_string(),
serde_yaml::Value::String(s) => s.clone(),
serde_yaml::Value::Sequence(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
serde_yaml::Value::Mapping(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
serde_yaml::Value::Tagged(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
};
debug!("META: Processing metadata: internal_name={}, custom_name={}, value={}", internal_name, custom_name, value_str);
return Some(MetaData {
name: custom_name.to_string(),
value: value_str,
});
}
}
// Convert the value to a string representation
let value_str = match &value {
serde_yaml::Value::Null => "null".to_string(),
serde_yaml::Value::Bool(b) => b.to_string(),
serde_yaml::Value::Number(n) => n.to_string(),
serde_yaml::Value::String(s) => s.clone(),
serde_yaml::Value::Sequence(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
serde_yaml::Value::Mapping(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
serde_yaml::Value::Tagged(_) => serde_yaml::to_string(&value).unwrap_or_else(|_| "".to_string()),
};
// Default: use internal name as output name
debug!("META: Processing metadata: name={}, value={}", internal_name, value_str);
Some(MetaData {
name: internal_name.to_string(),
value: value_str,
})
}
pub trait MetaPlugin where Self: 'static {
fn meta_type(&self) -> MetaPluginType;
fn is_supported(&self) -> bool {
true
}
fn is_internal(&self) -> bool {
true
}
// Check if the plugin is already finalized
fn is_finalized(&self) -> bool {
false
}
// Set the finalized state (only for plugins that can track this)
fn set_finalized(&mut self, _finalized: bool) {}
// Update the meta plugin with new data
fn update(&mut self, _data: &[u8]) -> MetaPluginResponse {
// Default implementation does nothing
MetaPluginResponse {
metadata: Vec::new(),
is_finalized: false,
}
}
fn finalize(&mut self) -> MetaPluginResponse {
// Default implementation does nothing
MetaPluginResponse {
metadata: Vec::new(),
is_finalized: true,
}
}
// Get program information for display in status
fn program_info(&self) -> Option<(&str, Vec<&str>)> {
None
}
// Initialize the plugin
fn initialize(&mut self) -> MetaPluginResponse {
// Default implementation does nothing
MetaPluginResponse {
metadata: Vec::new(),
is_finalized: false,
}
}
// Access to outputs mapping with default implementation
fn outputs(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
use once_cell::sync::Lazy;
use std::collections::HashMap;
use std::sync::Mutex;
static EMPTY: Lazy<std::collections::HashMap<String, serde_yaml::Value>> =
Lazy::new(|| std::collections::HashMap::new());
&EMPTY
}
fn outputs_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
panic!("outputs_mut() not implemented for this plugin")
}
// Access to options mapping with default implementation
fn options(&self) -> &std::collections::HashMap<String, serde_yaml::Value> {
use once_cell::sync::Lazy;
static EMPTY: Lazy<std::collections::HashMap<String, serde_yaml::Value>> =
Lazy::new(|| std::collections::HashMap::new());
&EMPTY
}
fn options_mut(&mut self) -> &mut std::collections::HashMap<String, serde_yaml::Value> {
panic!("options_mut() not implemented for this plugin")
}
// Get the default output names this plugin can produce
fn default_outputs(&self) -> Vec<String> {
// Default implementation returns the meta type as a string
vec![self.meta_type().to_string()]
}
// Method to downcast to concrete type (for checking finalization state)
fn as_any_mut(&mut self) -> &mut dyn std::any::Any where Self: Sized {
self
}
}
/// Global registry for meta plugins
static META_PLUGIN_REGISTRY: Lazy<Mutex<HashMap<MetaPluginType, fn(Option<HashMap<String, serde_yaml::Value>>, Option<HashMap<String, serde_yaml::Value>>) -> Box<dyn MetaPlugin>>>> =
Lazy::new(|| Mutex::new(HashMap::new()));
/// Register a meta plugin with the global registry
pub fn register_meta_plugin(
meta_plugin_type: MetaPluginType,
constructor: fn(Option<HashMap<String, serde_yaml::Value>>, Option<HashMap<String, serde_yaml::Value>>) -> Box<dyn MetaPlugin>
) {
META_PLUGIN_REGISTRY.lock().unwrap().insert(meta_plugin_type, constructor);
}
pub fn get_meta_plugin(
meta_plugin_type: MetaPluginType,
options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
) -> Box<dyn MetaPlugin> {
let registry = META_PLUGIN_REGISTRY.lock().unwrap();
if let Some(constructor) = registry.get(&meta_plugin_type) {
return constructor(options, outputs);
}
// Fallback for exec plugin which needs special handling
if meta_plugin_type == MetaPluginType::Exec {
// For exec type, we need to parse the command from options
let mut program_name = String::new();
let mut args = Vec::new();
let mut meta_name = MetaPluginType::Exec.to_string();
let mut split_whitespace = true;
if let Some(opts) = &options {
if let Some(command_value) = opts.get("command") {
if let Some(command_str) = command_value.as_str() {
let parts: Vec<&str> = command_str.split_whitespace().collect();
if !parts.is_empty() {
program_name = parts[0].to_string();
args = parts[1..].iter().map(|s| s.to_string()).collect();
}
}
}
// Handle other options if needed
if let Some(split_value) = opts.get("split_whitespace") {
if let Some(split_bool) = split_value.as_bool() {
split_whitespace = split_bool;
}
}
if let Some(name_value) = opts.get("name") {
if let Some(name_str) = name_value.as_str() {
meta_name = name_str.to_string();
}
}
}
return Box::new(MetaPluginExec::new(&program_name,
args.iter().map(|s: &String| s.as_str()).collect(),
meta_name,
split_whitespace,
options,
outputs));
}
// Fallback for unknown plugins
panic!("Meta plugin {:?} not registered", meta_plugin_type);
}