Co-authored-by: aider (openai/andrew/openrouter/deepseek/deepseek-chat-v3.1) <aider@aider.chat>
138 lines
4.1 KiB
Rust
138 lines
4.1 KiB
Rust
use std::time::Instant;
|
|
|
|
use crate::meta_plugin::{MetaPlugin, MetaPluginType};
|
|
|
|
#[derive(Debug, Clone, Default)]
|
|
pub struct ReadRateMetaPlugin {
|
|
start_time: Option<Instant>,
|
|
bytes_read: u64,
|
|
is_finalized: bool,
|
|
outputs: std::collections::HashMap<String, serde_yaml::Value>,
|
|
options: std::collections::HashMap<String, serde_yaml::Value>,
|
|
}
|
|
|
|
impl ReadRateMetaPlugin {
|
|
pub fn new(
|
|
_options: Option<std::collections::HashMap<String, serde_yaml::Value>>,
|
|
outputs: Option<std::collections::HashMap<String, serde_yaml::Value>>,
|
|
) -> ReadRateMetaPlugin {
|
|
// Start with default options
|
|
let mut final_options = std::collections::HashMap::new();
|
|
if let Some(opts) = _options {
|
|
for (key, value) in opts {
|
|
final_options.insert(key, value);
|
|
}
|
|
}
|
|
|
|
// Start with default outputs
|
|
let mut final_outputs = std::collections::HashMap::new();
|
|
let default_outputs = Self::default().default_outputs();
|
|
for output_name in default_outputs {
|
|
final_outputs.insert(output_name.clone(), serde_yaml::Value::String(output_name));
|
|
}
|
|
if let Some(outs) = outputs {
|
|
for (key, value) in outs {
|
|
final_outputs.insert(key, value);
|
|
}
|
|
}
|
|
|
|
ReadRateMetaPlugin {
|
|
start_time: None,
|
|
bytes_read: 0,
|
|
is_finalized: false,
|
|
outputs: final_outputs,
|
|
options: final_options,
|
|
}
|
|
}
|
|
|
|
}
|
|
|
|
impl MetaPlugin for ReadRateMetaPlugin {
|
|
fn is_finalized(&self) -> bool {
|
|
self.is_finalized
|
|
}
|
|
|
|
fn set_finalized(&mut self, finalized: bool) {
|
|
self.is_finalized = finalized;
|
|
}
|
|
|
|
fn finalize(&mut self) -> crate::meta_plugin::MetaPluginResponse {
|
|
// If already finalized, don't process again
|
|
if self.is_finalized {
|
|
return crate::meta_plugin::MetaPluginResponse {
|
|
metadata: Vec::new(),
|
|
is_finalized: true,
|
|
};
|
|
}
|
|
|
|
let mut metadata = Vec::new();
|
|
|
|
if let Some(start_time) = self.start_time {
|
|
let duration = start_time.elapsed();
|
|
let rate = if duration.as_secs_f64() > 0.0 {
|
|
format!("{:.2} KB/s", (self.bytes_read as f64 / 1024.0) / duration.as_secs_f64())
|
|
} else {
|
|
"N/A".to_string()
|
|
};
|
|
|
|
// Use process_metadata_outputs to handle output mapping
|
|
if let Some(meta_data) = crate::meta_plugin::process_metadata_outputs(
|
|
"read_rate",
|
|
serde_yaml::Value::String(rate),
|
|
&self.outputs
|
|
) {
|
|
metadata.push(meta_data);
|
|
}
|
|
}
|
|
|
|
// Mark as finalized
|
|
self.is_finalized = true;
|
|
|
|
crate::meta_plugin::MetaPluginResponse {
|
|
metadata,
|
|
is_finalized: true,
|
|
}
|
|
}
|
|
|
|
fn update(&mut self, data: &[u8]) -> crate::meta_plugin::MetaPluginResponse {
|
|
// If already finalized, don't process more data
|
|
if self.is_finalized {
|
|
return crate::meta_plugin::MetaPluginResponse {
|
|
metadata: Vec::new(),
|
|
is_finalized: true,
|
|
};
|
|
}
|
|
|
|
if self.start_time.is_none() {
|
|
self.start_time = Some(Instant::now());
|
|
}
|
|
self.bytes_read += data.len() as u64;
|
|
crate::meta_plugin::MetaPluginResponse {
|
|
metadata: Vec::new(),
|
|
is_finalized: false,
|
|
}
|
|
}
|
|
|
|
fn meta_type(&self) -> MetaPluginType {
|
|
MetaPluginType::ReadRate
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|