use std::collections::BTreeMap;
use std::future::{pending, ready};
use std::mem;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, PoisonError, mpsc};
use std::time::{Duration, Instant};
use agentcat_api::models::PublishEventRequest;
use serde_json::{Map, Value};
use tokio::runtime::Handle;
use tokio::sync::Semaphore;
use tokio::task::JoinHandle;
use super::json::object;
use super::vectors;
use crate::config::{HookError, Options, UserIdentity};
use crate::delivery::{LineWriter, Publisher};
use crate::event::Event;
use crate::event::assemble::EventJob;
use crate::handles::is_valid_session_id;
use crate::hooks::PendingHooks;
use crate::pipeline::process::Redactors;
use crate::sink::BoxFuture;
use crate::wire;
pub(crate) fn cases() -> Vec<Map<String, Value>> {
vectors::cases("queue")
}
pub(crate) fn case(name: &str) -> Map<String, Value> {
cases()
.into_iter()
.find(|case| vectors::name(case) == name)
.unwrap_or_else(|| panic!("queue.json has no case {name}"))
}
pub(crate) type Task<T> = JoinHandle<Result<Option<T>, HookError>>;
pub(crate) fn returns<T: Send + 'static>(value: Option<T>) -> Task<T> {
tokio::spawn(async move { Ok(value) })
}
pub(crate) fn throws<T: Send + 'static>(message: &str) -> Task<T> {
let message = message.to_owned();
tokio::spawn(async move { Err(message.into()) })
}
pub(crate) fn hangs<T: Send + 'static>() -> Task<T> {
tokio::spawn(pending())
}
fn outcome_task<T: Send + 'static>(outcome: &Value, value: fn(&Value) -> T) -> Task<T> {
if outcome.get("hangs").is_some() {
return hangs();
}
if let Some(message) = outcome.get("throws") {
return throws(message.as_str().expect("a message"));
}
let returned = &outcome["returns"];
returns((!returned.is_null()).then(|| value(returned)))
}
fn identity_of(value: &Value) -> UserIdentity {
let keys = value.as_object().expect("an identify result is an object");
for key in keys.keys() {
assert!(
["user_id", "user_name", "user_data"].contains(&key.as_str()),
"the identify result has the unknown key {key:?}"
);
}
let mut identity = UserIdentity::new(value["user_id"].as_str().expect("a user_id"));
if let Some(name) = value.get("user_name").and_then(Value::as_str) {
identity = identity.user_name(name);
}
if let Some(Value::Object(data)) = value.get("user_data") {
identity = identity.user_data(data.clone());
}
identity
}
pub(crate) fn pending_of(name: &str, hooks: &Map<String, Value>) -> PendingHooks {
for key in hooks.keys() {
assert!(
[
"resolve_session_id",
"identify",
"event_tags",
"event_properties"
]
.contains(&key.as_str()),
"queue.json {name}: the hooks have the unknown key {key:?}"
);
}
PendingHooks {
resolve_session_id: hooks.get("resolve_session_id").map(|outcome| {
outcome_task(outcome, |value| {
value.as_str().expect("a string").to_owned()
})
}),
identify: hooks
.get("identify")
.map(|outcome| outcome_task(outcome, identity_of)),
event_tags: hooks
.get("event_tags")
.map(|outcome| outcome_task(outcome, |value| object(value.clone()))),
event_properties: hooks
.get("event_properties")
.map(|outcome| outcome_task(outcome, |value| object(value.clone()))),
}
}
fn redaction_stages(case: &Map<String, Value>) -> (bool, bool) {
match case
.get("custom")
.and_then(|custom| custom["form"].as_str())
{
None => (true, true),
Some("string") => (false, false),
Some(_) => (false, true),
}
}
pub(crate) fn job_of(case: &Map<String, Value>, runtime: Option<&Handle>) -> EventJob {
let hooks = case.get("hooks").and_then(Value::as_object).map(|hooks| {
let _entered = runtime.expect("a runtime for the hooks").enter();
pending_of(vectors::name(case), hooks)
});
let (string_redaction, event_redaction) = redaction_stages(case);
EventJob {
event: object(case["event"].clone()),
hooks: hooks.unwrap_or_default(),
hook_mode: case.get("hook_mode") == Some(&Value::Bool(true)),
runtime: runtime.cloned(),
string_redaction,
event_redaction,
backtrace: None,
}
}
pub(crate) fn redactors_of(case: &Map<String, Value>) -> Redactors {
let options = options_of(case);
Redactors {
string: options.string_redactor().cloned(),
event: options.event_redactor().cloned(),
}
}
pub(crate) fn options_of(case: &Map<String, Value>) -> Options {
let mut options = Options::default();
if let Some(spec) = case.get("redact_event") {
let spec = spec.clone();
options = options.redact_event(move |event: Event| {
let result: Result<Option<Event>, HookError> = if spec.get("throws").is_some() {
Err(spec["throws"].as_str().expect("a message").into())
} else if let Some(Value::Object(fields)) = spec.get("merge") {
let mut merged = event.into_object();
merged.extend(fields.clone());
Ok(Some(Event::from_object(&merged)))
} else {
Ok(None)
};
ready(result)
});
}
if let Some(spec) = case.get("redact_strings") {
let replace = spec
.get("replace")
.and_then(Value::as_str)
.map(str::to_owned);
let throw_at = spec
.get("throw_at")
.and_then(Value::as_u64)
.map(|n| usize::try_from(n).expect("a call number"));
let calls = Arc::new(AtomicUsize::new(0));
options = options.redact_sensitive_information(move |text: String| {
let call = calls.fetch_add(1, Ordering::SeqCst) + 1;
let result: Result<String, HookError> = match (&replace, throw_at) {
(Some(secret), _) => Ok(text.replace(secret.as_str(), "[REDACTED]")),
(None, Some(n)) if n == call => Err("redactor failed".into()),
(None, _) => Ok(text),
};
ready(result)
});
}
options
}
const DELIVERY_LINES: [&str; 2] = ["Successfully sent event ", "Failed to send event "];
const CUSTOM_BUILDER_LINES: [&str; 3] = [
"Published custom event ",
"publish_custom_event: ",
"Dropping invalid tag: ",
];
const ERROR_PREFIXES: [&str; 8] = [
"resolve_session_id hook error: ",
"Error: identify hook returned an error - ",
"event_tags callback error: ",
"event_properties callback error: ",
"Failed to resolve pending hook results: ",
"Failed to redact event (event-level hook): ",
"Failed to redact event: ",
"Failed to truncate event: ",
];
pub(crate) fn expected_line_as_rust(line: &str) -> String {
for prefix in ERROR_PREFIXES {
let message = line
.strip_prefix(prefix)
.and_then(|rest| rest.strip_prefix("Error: "));
if let Some(message) = message {
return format!("{prefix}{message}");
}
}
line.to_owned()
}
fn queue_lines(case: &Map<String, Value>) -> impl Iterator<Item = &str> {
let custom = case.contains_key("custom");
case["logs"]
.as_array()
.expect("logs")
.iter()
.map(|line| line.as_str().expect("a line"))
.filter(move |line| {
!custom
|| !CUSTOM_BUILDER_LINES
.iter()
.any(|start| line.starts_with(start))
})
}
pub(crate) fn expected_logs(case: &Map<String, Value>) -> Vec<String> {
queue_lines(case)
.filter(|line| !DELIVERY_LINES.iter().any(|start| line.starts_with(start)))
.map(expected_line_as_rust)
.collect()
}
const DELIVERY_FAILED: &str = "Failed to deliver AgentCat event after retries: ";
const ATTEMPT_ERROR: &str = "[Error: ";
pub(crate) fn expected_delivery_line_as_rust(line: &str) -> String {
if let Some(rendered) = line.strip_prefix(DELIVERY_FAILED) {
return format!("{DELIVERY_FAILED}{}", error_message_of_rendered(rendered));
}
if line.starts_with(DELIVERY_LINES[1])
&& let Some(at) = line.find(ATTEMPT_ERROR)
&& let Some(rendered) = line[at + ATTEMPT_ERROR.len()..].strip_suffix(']')
{
let head = &line[..at + ATTEMPT_ERROR.len()];
return format!("{head}{}]", error_message_of_rendered(rendered));
}
line.to_owned()
}
fn error_message_of_rendered(rendered: &str) -> String {
if let Ok(Value::Object(error)) = serde_json::from_str::<Value>(rendered)
&& let Some(Value::String(message)) = error.get("message")
{
return message.clone();
}
rendered
.strip_prefix("Error: ")
.unwrap_or(rendered)
.to_owned()
}
pub(crate) fn worker_logs(case: &Map<String, Value>) -> Vec<String> {
queue_lines(case)
.map(|line| expected_delivery_line_as_rust(&expected_line_as_rust(line)))
.collect()
}
pub(crate) fn full_logs(case: &Map<String, Value>) -> Vec<String> {
case["logs"]
.as_array()
.expect("logs")
.iter()
.map(|line| {
expected_delivery_line_as_rust(&expected_line_as_rust(line.as_str().expect("a line")))
})
.collect()
}
const EVENT_ID_BODY: usize = 27;
pub(crate) fn mask_event_ids(line: &str) -> String {
let mut masked = String::with_capacity(line.len());
let mut rest = line;
while let Some(at) = rest.find("evt_") {
let body = &rest[at + 4..];
let is_id = body.len() >= EVENT_ID_BODY
&& body.as_bytes()[..EVENT_ID_BODY]
.iter()
.all(u8::is_ascii_alphanumeric);
masked.push_str(&rest[..at]);
if is_id {
masked.push_str("<evt>");
rest = &body[EVENT_ID_BODY..];
} else {
masked.push_str("evt_");
rest = body;
}
}
masked.push_str(rest);
masked
}
const MODEL_FIELDS: [(&str, &str); 28] = [
("id", "id"),
("project_id", "projectId"),
("session_id", "sessionId"),
("timestamp", "timestamp"),
("duration", "duration"),
("input_tokens", "inputTokens"),
("output_tokens", "outputTokens"),
("event_type", "eventType"),
("resource_name", "resourceName"),
("parameters", "parameters"),
("response", "response"),
("user_intent", "userIntent"),
("is_error", "isError"),
("error", "error"),
("identify_actor_given_id", "identifyActorGivenId"),
("identify_actor_name", "identifyActorName"),
("identify_data", "identifyData"),
("ip_address", "ipAddress"),
("sdk_language", "sdkLanguage"),
("agentcat_version", "agentcatVersion"),
("server_name", "serverName"),
("server_version", "serverVersion"),
("client_name", "clientName"),
("client_version", "clientVersion"),
("actor_id", "actorId"),
("event_id", "eventId"),
("tags", "tags"),
("properties", "properties"),
];
fn sort_tags(model: &mut Map<String, Value>) {
if let Some(Value::Object(tags)) = model.get_mut("tags") {
let sorted: BTreeMap<String, Value> = mem::take(tags).into_iter().collect();
*tags = sorted.into_iter().collect();
}
}
pub(crate) fn to_request_model(request: &PublishEventRequest) -> Value {
let mut wire = object(serde_json::to_value(request).expect("a request is JSON"));
let mut model = Map::new();
for (field, name) in MODEL_FIELDS {
if let Some(value) = wire.remove(field) {
model.insert(name.to_owned(), value);
}
}
assert!(
wire.is_empty(),
"fields the fixture's model does not name: {wire:?}"
);
let id = model["id"].as_str().expect("an id");
let body = id.strip_prefix("evt_").expect("an evt_ ID");
assert!(
body.len() == 27 && body.bytes().all(|b| b.is_ascii_alphanumeric()),
"{id} is not an evt_ ID"
);
model.insert("id".to_owned(), Value::from("<evt>"));
if model
.get("tags")
.and_then(|tags| tags.get("agentcat_session_id_source"))
== Some(&Value::from("minted"))
{
let session = model["sessionId"].as_str().expect("a session");
assert!(is_valid_session_id(session), "{session}");
model.insert("sessionId".to_owned(), Value::from("<minted>"));
}
if let Some(Value::String(timestamp)) = model.get_mut("timestamp") {
let time = wire::parse_iso_millis(timestamp).expect("an RFC 3339 timestamp");
*timestamp = wire::iso_millis(time);
}
sort_tags(&mut model);
Value::Object(model)
}
pub(crate) fn expected_requests(case: &Map<String, Value>) -> Vec<Value> {
case["requests"]
.as_array()
.expect("requests")
.iter()
.map(|request| {
let mut request = object(request.clone());
match request.get_mut("userIntent") {
Some(Value::Null) => {
request.remove("userIntent");
}
Some(intent) if !intent.is_string() => {
*intent = Value::String(wire::stringify(intent));
}
_ => {}
}
sort_tags(&mut request);
Value::Object(request)
})
.collect()
}
pub(crate) const MOCK_FAILURE: &str = "mock failure";
#[derive(Debug, Clone)]
pub(crate) struct Attempt {
pub(crate) at: Instant,
pub(crate) request: PublishEventRequest,
}
#[derive(Debug, Clone)]
pub(crate) struct Gate(Arc<Semaphore>);
impl Gate {
pub(crate) fn open(&self) {
self.0.close();
}
}
#[derive(Default)]
struct Recorded {
failures: Mutex<usize>,
gate: Option<Gate>,
delay: Option<Duration>,
starts: Mutex<Option<mpsc::Sender<()>>>,
attempts: Mutex<Vec<Attempt>>,
running: AtomicUsize,
most_running: AtomicUsize,
finished: AtomicUsize,
}
#[derive(Clone, Default)]
pub(crate) struct RecordingPublisher(Arc<Recorded>);
struct Running(Arc<Recorded>);
impl Drop for Running {
fn drop(&mut self) {
self.0.running.fetch_sub(1, Ordering::SeqCst);
}
}
impl RecordingPublisher {
pub(crate) fn new() -> Self {
Self::default()
}
pub(crate) fn failing(attempts: usize) -> Self {
Self::build(attempts, None, None)
}
pub(crate) fn gated() -> (Self, Gate) {
let gate = Gate(Arc::new(Semaphore::new(0)));
(Self::build(0, Some(gate.clone()), None), gate)
}
pub(crate) fn slow(delay: Duration) -> Self {
Self::build(0, None, Some(delay))
}
fn build(failures: usize, gate: Option<Gate>, delay: Option<Duration>) -> Self {
Self(Arc::new(Recorded {
failures: Mutex::new(failures),
gate,
delay,
..Recorded::default()
}))
}
pub(crate) fn starts(&self) -> mpsc::Receiver<()> {
let (sender, receiver) = mpsc::channel();
*self.0.starts.lock().unwrap_or_else(PoisonError::into_inner) = Some(sender);
receiver
}
pub(crate) fn attempts(&self) -> Vec<Attempt> {
self.0
.attempts
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
pub(crate) fn requests(&self) -> Vec<PublishEventRequest> {
self.attempts()
.into_iter()
.map(|attempt| attempt.request)
.collect()
}
pub(crate) fn resource_names(&self) -> Vec<String> {
self.requests()
.into_iter()
.map(|request| request.resource_name.unwrap_or_default())
.collect()
}
pub(crate) fn most_running(&self) -> usize {
self.0.most_running.load(Ordering::SeqCst)
}
pub(crate) fn finished(&self) -> usize {
self.0.finished.load(Ordering::SeqCst)
}
}
impl Publisher for RecordingPublisher {
fn publish(&self, request: PublishEventRequest) -> BoxFuture<'static, Result<(), String>> {
let recorded = Arc::clone(&self.0);
Box::pin(async move {
recorded
.attempts
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push(Attempt {
at: Instant::now(),
request,
});
let fail = {
let mut left = recorded
.failures
.lock()
.unwrap_or_else(PoisonError::into_inner);
let fail = *left > 0;
*left = left.saturating_sub(1);
fail
};
let now = recorded.running.fetch_add(1, Ordering::SeqCst) + 1;
recorded.most_running.fetch_max(now, Ordering::SeqCst);
let _running = Running(Arc::clone(&recorded));
if let Some(starts) = &*recorded
.starts
.lock()
.unwrap_or_else(PoisonError::into_inner)
{
let _ = starts.send(());
}
if let Some(Gate(gate)) = &recorded.gate {
let _ = gate.acquire().await;
}
if let Some(delay) = recorded.delay {
tokio::time::sleep(delay).await;
}
recorded.finished.fetch_add(1, Ordering::SeqCst);
if fail {
Err(MOCK_FAILURE.to_owned())
} else {
Ok(())
}
})
}
}
pub(crate) struct DropSignal(mpsc::Sender<()>);
impl DropSignal {
pub(crate) fn new() -> (Self, mpsc::Receiver<()>) {
let (sender, receiver) = mpsc::channel();
(Self(sender), receiver)
}
}
impl Drop for DropSignal {
fn drop(&mut self) {
let _ = self.0.send(());
}
}
type Entry = (String, Option<String>);
#[derive(Debug, Clone, Default)]
pub(crate) struct CapturedLines(Arc<(Mutex<Vec<Entry>>, Condvar)>);
impl CapturedLines {
pub(crate) fn writer(&self) -> LineWriter {
let lines = self.clone();
Arc::new(move |line: &str| {
let thread = std::thread::current().name().map(str::to_owned);
let (entries, added) = &*lines.0;
entries
.lock()
.unwrap_or_else(PoisonError::into_inner)
.push((line.to_owned(), thread));
added.notify_all();
})
}
pub(crate) fn wait_for(&self, count: usize, timeout: Duration) -> bool {
let (entries, added) = &*self.0;
let entries = entries.lock().unwrap_or_else(PoisonError::into_inner);
let (entries, _) = added
.wait_timeout_while(entries, timeout, |entries| entries.len() < count)
.unwrap_or_else(PoisonError::into_inner);
entries.len() >= count
}
pub(crate) fn lines(&self) -> Vec<String> {
self.entries().into_iter().map(|(line, _)| line).collect()
}
pub(crate) fn entries(&self) -> Vec<Entry> {
self.0
.0
.lock()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
}
mod tests {
use serde_json::{Value, json};
use super::{expected_requests, object};
#[test]
fn expected_user_intent_follows_the_wire() {
let case = object(json!({
"requests": [
{"id": "<evt>", "userIntent": null},
{"id": "<evt>", "userIntent": ["a", 1]},
{"id": "<evt>", "userIntent": {"why": true}},
{"id": "<evt>", "userIntent": "plain"},
{"id": "<evt>"},
],
}));
let intents: Vec<Option<Value>> = expected_requests(&case)
.iter()
.map(|request| request.get("userIntent").cloned())
.collect();
assert_eq!(
intents,
[
None,
Some(json!("[\"a\",1]")),
Some(json!("{\"why\":true}")),
Some(json!("plain")),
None,
]
);
}
}