use std::sync::{Arc, Mutex};
use rpi_ai::types::{AssistantMessage, Content, Message, UserContent};
use rpi_ai::Context;
use rpi_ai::{ContextPatch, Model, ProviderHooks, SimpleStreamOptions, SimpleStreamOptionsPatch};
use rpi_plugin_sdk::EventTag;
use crate::loader::ExtensionSession;
use crate::prompt_transform::emit_before_agent_start;
use crate::registry::RegistrySnapshot;
use crate::translate::dispatch_data_event;
use crate::PluginKeepalive;
pub struct ExtensionProviderHooks {
snapshot: Arc<RegistrySnapshot>,
_keepalive: Arc<PluginKeepalive>,
last_before_agent_start: Mutex<Option<i64>>,
last_forced_prompt: Mutex<Option<(i64, Option<String>)>>,
}
impl ExtensionProviderHooks {
pub fn from_session(session: &ExtensionSession) -> Option<Self> {
let snapshot = session.snapshot_arc()?;
let observes_events = [
EventTag::BeforeProviderRequest,
EventTag::BeforeProviderHeaders,
EventTag::AfterProviderResponse,
EventTag::BeforeAgentStart,
]
.iter()
.any(|t| !snapshot.handlers_for(*t).is_empty());
let transforms_prompt = !snapshot.before_agent_start().is_empty();
if !observes_events && !transforms_prompt {
return None;
}
Some(Self {
snapshot,
_keepalive: session.keepalive(),
last_before_agent_start: Mutex::new(None),
last_forced_prompt: Mutex::new(None),
})
}
}
impl ProviderHooks for ExtensionProviderHooks {
fn transform_context(&self, _model: &Model, ctx: &Context) -> Option<ContextPatch> {
if self.snapshot.before_agent_start().is_empty() {
return None;
}
let (timestamp, prompt, image_count) = latest_user_prompt(&ctx.messages)?;
let mut cache = self.last_forced_prompt.lock().unwrap();
if cache.as_ref().map(|(ts, _)| *ts) != Some(timestamp) {
let event = before_agent_start_event(&prompt, image_count, ctx);
let forced = emit_before_agent_start(
ctx.system_prompt.as_deref().unwrap_or_default(),
&event.to_string(),
&self.snapshot,
);
*cache = Some((timestamp, forced));
}
cache
.as_ref()
.and_then(|(_, forced)| forced.clone())
.map(|system_prompt| ContextPatch {
system_prompt: Some(system_prompt),
})
}
fn before_request(
&self,
model: &Model,
ctx: &Context,
opts: &SimpleStreamOptions,
) -> Option<SimpleStreamOptionsPatch> {
if let Some((timestamp, prompt, image_count)) = latest_user_prompt(&ctx.messages) {
let mut last = self.last_before_agent_start.lock().unwrap();
if *last != Some(timestamp) {
*last = Some(timestamp);
drop(last);
let event = before_agent_start_event(&prompt, image_count, ctx);
dispatch_data_event(
&self.snapshot,
EventTag::BeforeAgentStart,
&event.to_string(),
);
}
}
let request = serde_json::json!({
"model": model.id,
"provider": model.provider,
"baseUrl": model.base_url,
"reasoning": model.reasoning,
"apiKey": opts.api_key,
"timeoutMs": opts.timeout.map(|d| d.as_millis() as u64),
"headers": opts.headers,
"metadata": opts.metadata,
"maxTokens": opts.max_tokens,
"temperature": opts.temperature,
"messages": ctx.messages,
});
dispatch_data_event(
&self.snapshot,
EventTag::BeforeProviderRequest,
&request.to_string(),
);
let headers = serde_json::json!({ "headers": opts.headers });
dispatch_data_event(
&self.snapshot,
EventTag::BeforeProviderHeaders,
&headers.to_string(),
);
None
}
fn after_response(&self, _model: &Model, message: &AssistantMessage) {
let json = serde_json::to_value(message).unwrap_or(serde_json::json!({}));
dispatch_data_event(
&self.snapshot,
EventTag::AfterProviderResponse,
&json.to_string(),
);
}
}
fn before_agent_start_event(prompt: &str, image_count: usize, ctx: &Context) -> serde_json::Value {
serde_json::json!({
"prompt": prompt,
"imageCount": image_count,
"systemPrompt": ctx.system_prompt,
"sessionId": std::env::var(rpi_plugin_sdk::SESSION_ID_ENV).ok(),
"host": host_identity(),
})
}
fn host_identity() -> serde_json::Value {
serde_json::json!({
"name": rpi_plugin_sdk::HOST_NAME,
"version": env!("CARGO_PKG_VERSION"),
})
}
fn latest_user_prompt(messages: &[Message]) -> Option<(i64, String, usize)> {
messages.iter().rev().find_map(|message| match message {
Message::User(user) => Some((
user.timestamp,
user_content_text(&user.content),
user_content_image_count(&user.content),
)),
_ => None,
})
}
fn user_content_text(content: &UserContent) -> String {
match content {
UserContent::Text(text) => text.clone(),
UserContent::Blocks(blocks) => blocks
.iter()
.filter_map(|block| match block {
Content::Text(text) => Some(text.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("\n"),
}
}
fn user_content_image_count(content: &UserContent) -> usize {
match content {
UserContent::Text(_) => 0,
UserContent::Blocks(blocks) => blocks
.iter()
.filter(|block| matches!(block, Content::Image(_)))
.count(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use rpi_plugin_sdk::EventTag;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Mutex;
static PROVIDER_HITS: AtomicUsize = AtomicUsize::new(0);
static PROVIDER_LOCK: Mutex<()> = Mutex::new(());
static CAPTURED: Mutex<Vec<(EventTag, String)>> = Mutex::new(Vec::new());
rpi_plugin_sdk::event_handler! { fn capture_handler(ev, _ud) {
let payload = ev.data_text().unwrap_or_default();
CAPTURED.lock().unwrap().push((ev.tag(), payload));
0
} }
rpi_plugin_sdk::event_handler! { fn counting_provider_handler(ev, _ud) {
let s = ev.data_text().unwrap_or_default();
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&s) {
if v.get("model").is_some() {
assert!(v.get("model").is_some(), "request event must carry model");
}
PROVIDER_HITS.fetch_add(1, Ordering::SeqCst);
}
0
} }
#[test]
fn before_request_emits_before_agent_start_once_and_forwards_messages() {
let _guard = PROVIDER_LOCK.lock().unwrap();
CAPTURED.lock().unwrap().clear();
let mut registry = crate::registry::ExtensionRegistry::new();
registry.register_event_handler(
"test-plugin".to_string(),
EventTag::BeforeAgentStart,
capture_handler,
std::ptr::null_mut(),
);
registry.register_event_handler(
"test-plugin".to_string(),
EventTag::BeforeProviderRequest,
capture_handler,
std::ptr::null_mut(),
);
let snapshot = Arc::new(registry.snapshot());
let hooks = ExtensionProviderHooks {
snapshot: Arc::clone(&snapshot),
_keepalive: Arc::new(crate::PluginKeepalive::new(Vec::new(), None)),
last_before_agent_start: Mutex::new(None),
last_forced_prompt: Mutex::new(None),
};
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let ctx = rpi_ai::Context::new(vec![rpi_ai::types::Message::User(
rpi_ai::types::UserMessage::new(
rpi_ai::types::UserContent::Text("hello".to_string()),
42,
),
)]);
let opts = rpi_ai::SimpleStreamOptions::default();
hooks.before_request(&model, &ctx, &opts);
hooks.before_request(&model, &ctx, &opts);
let captured = CAPTURED.lock().unwrap().clone();
let agent_starts: Vec<_> = captured
.iter()
.filter(|(tag, _)| matches!(tag, EventTag::BeforeAgentStart))
.collect();
assert_eq!(
agent_starts.len(),
1,
"before_agent_start must fire once per prompt: {captured:?}"
);
let start_payload: serde_json::Value =
serde_json::from_str(&agent_starts[0].1).expect("before_agent_start payload is JSON");
assert_eq!(start_payload["prompt"], "hello");
assert_eq!(start_payload["imageCount"], 0);
assert_eq!(
start_payload.get("sessionId").map(|v| v.is_null()),
Some(std::env::var(rpi_plugin_sdk::SESSION_ID_ENV).is_err()),
"sessionId must mirror the host's session id env var: {}",
agent_starts[0].1
);
assert_eq!(start_payload["host"]["name"], rpi_plugin_sdk::HOST_NAME);
assert_eq!(
start_payload["host"]["version"],
env!("CARGO_PKG_VERSION"),
"the host reports the version of the host, not of the plugin SDK"
);
let requests: Vec<_> = captured
.iter()
.filter(|(tag, _)| matches!(tag, EventTag::BeforeProviderRequest))
.collect();
assert_eq!(requests.len(), 2, "every provider call is observed");
assert!(
requests[0].1.contains("\"messages\""),
"request payload must forward messages: {}",
requests[0].1
);
}
#[test]
fn provider_hooks_dispatch_to_subscribers() {
let _guard = PROVIDER_LOCK.lock().unwrap();
PROVIDER_HITS.store(0, Ordering::SeqCst);
let mut registry = crate::registry::ExtensionRegistry::new();
let handler: rpi_plugin_sdk::EventHandlerFn = counting_provider_handler;
registry.register_event_handler(
"test-plugin".to_string(),
EventTag::BeforeProviderRequest,
handler,
std::ptr::null_mut(),
);
registry.register_event_handler(
"test-plugin".to_string(),
EventTag::BeforeProviderHeaders,
handler,
std::ptr::null_mut(),
);
let snapshot = Arc::new(registry.snapshot());
assert!(dispatch_data_event(
&snapshot,
EventTag::BeforeProviderRequest,
r#"{"model":"m"}"#,
));
assert_eq!(PROVIDER_HITS.load(Ordering::SeqCst), 1);
let hooks = ExtensionProviderHooks {
snapshot: Arc::clone(&snapshot),
_keepalive: Arc::new(crate::PluginKeepalive::new(Vec::new(), None)),
last_before_agent_start: Mutex::new(None),
last_forced_prompt: Mutex::new(None),
};
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let ctx = rpi_ai::Context::default();
let opts = rpi_ai::SimpleStreamOptions::default();
let patch = hooks.before_request(&model, &ctx, &opts);
assert!(patch.is_none(), "observer-only v1: no patch returned");
assert_eq!(PROVIDER_HITS.load(Ordering::SeqCst), 3);
}
#[test]
fn provider_hooks_no_subscribers_is_noop() {
let registry = crate::registry::ExtensionRegistry::new();
let snapshot = Arc::new(registry.snapshot());
assert!(!dispatch_data_event(
&snapshot,
EventTag::AfterProviderResponse,
"{}",
));
}
fn append_marker(
envelope: &str,
out: &mut rpi_plugin_sdk::callbacks::Output<'_>,
marker: &str,
) -> i32 {
let current = serde_json::from_str::<serde_json::Value>(envelope)
.ok()
.and_then(|v| v.get("systemPrompt")?.as_str().map(str::to_string))
.unwrap_or_default();
let json = serde_json::json!({ "systemPrompt": format!("{current}{marker}") });
i32::from(!out.write(json.to_string()))
}
rpi_plugin_sdk::string_handler! { fn append_a(input, out, _data) { append_marker(input, out, " A") } }
rpi_plugin_sdk::string_handler! { fn append_b(input, out, _data) { append_marker(input, out, " B") } }
rpi_plugin_sdk::string_handler! { fn append_after(input, out, _data) { append_marker(input, out, " AFTER") } }
rpi_plugin_sdk::string_handler! { fn append_mark(input, out, _data) { append_marker(input, out, " MARK") } }
extern "C" fn declines(
_event_json: rpi_plugin_sdk::StbStringRef,
_out: *mut rpi_plugin_sdk::StbString,
_user_data: *mut std::ffi::c_void,
) -> i32 {
0
}
extern "C" fn fails(
_event_json: rpi_plugin_sdk::StbStringRef,
_out: *mut rpi_plugin_sdk::StbString,
_user_data: *mut std::ffi::c_void,
) -> i32 {
7
}
fn hooks_with(
handlers: Vec<crate::registry::BeforeAgentStartHandler>,
) -> ExtensionProviderHooks {
let mut registry = crate::registry::ExtensionRegistry::new();
for h in handlers {
registry.register_before_agent_start(h.handler, h.plugin_free_string, h.user_data);
}
ExtensionProviderHooks {
snapshot: Arc::new(registry.snapshot()),
_keepalive: Arc::new(crate::PluginKeepalive::new(Vec::new(), None)),
last_before_agent_start: Mutex::new(None),
last_forced_prompt: Mutex::new(None),
}
}
#[test]
fn prompt_transformers_chain_in_registration_order() {
let hooks = hooks_with(vec![
crate::registry::BeforeAgentStartHandler {
handler: append_a,
plugin_free_string: test_free,
user_data: std::ptr::null_mut(),
},
crate::registry::BeforeAgentStartHandler {
handler: append_b,
plugin_free_string: test_free,
user_data: std::ptr::null_mut(),
},
]);
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let mut ctx = rpi_ai::Context::new(vec![rpi_ai::types::Message::User(
rpi_ai::types::UserMessage::new(
rpi_ai::types::UserContent::Text("hello".to_string()),
1,
),
)]);
ctx.system_prompt = Some("base".to_string());
let patch = hooks
.transform_context(&model, &ctx)
.expect("both handlers changed the prompt");
assert_eq!(patch.system_prompt.as_deref(), Some("base A B"));
}
#[test]
fn a_failing_transformer_is_skipped_and_the_chain_continues() {
let hooks = hooks_with(vec![
crate::registry::BeforeAgentStartHandler {
handler: fails,
plugin_free_string: test_free,
user_data: std::ptr::null_mut(),
},
crate::registry::BeforeAgentStartHandler {
handler: append_after,
plugin_free_string: test_free,
user_data: std::ptr::null_mut(),
},
]);
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let mut ctx = rpi_ai::Context::new(vec![rpi_ai::types::Message::User(
rpi_ai::types::UserMessage::new(
rpi_ai::types::UserContent::Text("hello".to_string()),
1,
),
)]);
ctx.system_prompt = Some("base".to_string());
let patch = hooks.transform_context(&model, &ctx).expect("changed");
assert_eq!(
patch.system_prompt.as_deref(),
Some("base AFTER"),
"the failing handler must contribute nothing and not stop the next one"
);
}
#[test]
fn declining_transformers_leave_the_prompt_untouched() {
let hooks = hooks_with(vec![crate::registry::BeforeAgentStartHandler {
handler: declines,
plugin_free_string: test_free,
user_data: std::ptr::null_mut(),
}]);
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let ctx = rpi_ai::Context::new(vec![rpi_ai::types::Message::User(
rpi_ai::types::UserMessage::new(
rpi_ai::types::UserContent::Text("hello".to_string()),
1,
),
)]);
assert!(
hooks.transform_context(&model, &ctx).is_none(),
"no handler changed the prompt ⇒ no patch ⇒ harness leaves it alone"
);
}
#[test]
fn no_transformers_means_no_patch() {
let hooks = hooks_with(vec![]);
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let ctx = rpi_ai::Context::default();
assert!(hooks.transform_context(&model, &ctx).is_none());
}
use rpi_plugin_sdk::callbacks::free_string as test_free;
#[test]
fn the_observer_sees_the_prompt_after_the_transform() {
let _guard = PROVIDER_LOCK.lock().unwrap();
CAPTURED.lock().unwrap().clear();
let mut registry = crate::registry::ExtensionRegistry::new();
registry.register_event_handler(
"test-plugin".to_string(),
EventTag::BeforeAgentStart,
capture_handler,
std::ptr::null_mut(),
);
registry.register_before_agent_start(append_mark, test_free, std::ptr::null_mut());
let hooks = ExtensionProviderHooks {
snapshot: Arc::new(registry.snapshot()),
_keepalive: Arc::new(crate::PluginKeepalive::new(Vec::new(), None)),
last_before_agent_start: Mutex::new(None),
last_forced_prompt: Mutex::new(None),
};
let model = rpi_ai::Model::new(
"m",
"m",
rpi_ai::Api::AnthropicMessages,
"anthropic",
"https://api.anthropic.com",
);
let mut ctx = rpi_ai::Context::new(vec![rpi_ai::types::Message::User(
rpi_ai::types::UserMessage::new(
rpi_ai::types::UserContent::Text("hello".to_string()),
1,
),
)]);
ctx.system_prompt = Some("base".to_string());
let patch = hooks
.transform_context(&model, &ctx)
.expect("transformer changed the prompt");
patch.apply(&mut ctx);
hooks.before_request(&model, &ctx, &rpi_ai::SimpleStreamOptions::default());
let captured = CAPTURED.lock().unwrap().clone();
let start = captured
.iter()
.find(|(tag, _)| matches!(tag, EventTag::BeforeAgentStart))
.expect("BeforeAgentStart was dispatched");
let payload: serde_json::Value = serde_json::from_str(&start.1).expect("JSON");
assert_eq!(
payload["systemPrompt"], "base MARK",
"observers must be shown the prompt the provider will actually get"
);
}
}