use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use aion_core::{ActivityEvent, ActivityEventKind, MessageRole};
use chrono::Utc;
use uuid::Uuid;
use crate::context::ActivityContext;
use crate::process::{CommandOutputObserver, CommandStream};
const STDOUT_PRODUCER: &str = "command stdout";
const STDERR_PRODUCER: &str = "command stderr";
const fn producer_label(stream: CommandStream) -> &'static str {
match stream {
CommandStream::Stdout => STDOUT_PRODUCER,
CommandStream::Stderr => STDERR_PRODUCER,
}
}
#[derive(Debug)]
pub struct CommandTranscript<'context> {
context: &'context ActivityContext,
next_seq: AtomicU64,
seam_failure_reported: AtomicBool,
}
impl<'context> CommandTranscript<'context> {
#[must_use]
pub fn new(context: &'context ActivityContext) -> Self {
Self {
context,
next_seq: AtomicU64::new(0),
seam_failure_reported: AtomicBool::new(false),
}
}
fn event(&self, stream: CommandStream, line: &str) -> ActivityEvent {
ActivityEvent {
workflow_id: self.context.workflow_id().clone(),
run_id: self.context.run_id().clone(),
activity_id: self.context.activity_id().clone(),
attempt: self.context.attempt(),
agent_id: Uuid::nil(),
agent_role: producer_label(stream).to_owned(),
emitted_at: Utc::now(),
worker_seq: self.next_seq.fetch_add(1, Ordering::Relaxed),
store_seq: None,
ephemeral: false,
kind: ActivityEventKind::Message {
role: MessageRole::Tool,
text: line.to_owned(),
},
}
}
}
impl CommandOutputObserver for CommandTranscript<'_> {
fn on_line(&self, stream: CommandStream, line: &str) {
let event = self.event(stream, line);
if let Err(error) = self.context.emit_event(event) {
if !self.seam_failure_reported.swap(true, Ordering::Relaxed) {
tracing::warn!(
%error,
stream = stream.name(),
activity_id = %self.context.activity_id(),
attempt = self.context.attempt(),
"command transcript: the activity's event seam is closed; this command's \
remaining output is not streamed to the transcript"
);
}
}
}
}
#[cfg(test)]
mod tests {
use aion_core::{ActivityEventKind, ActivityId, MessageRole, RunId, WorkflowId};
use tokio::sync::mpsc;
use uuid::Uuid;
use super::{CommandTranscript, STDERR_PRODUCER, STDOUT_PRODUCER};
use crate::context::ActivityContext;
use crate::process::{CommandOutputObserver, CommandStream};
type TestResult = Result<(), Box<dyn std::error::Error>>;
fn workflow_id() -> WorkflowId {
WorkflowId::new(Uuid::from_u128(0x5eed))
}
fn run_id() -> RunId {
RunId::new(Uuid::from_u128(0x5eed_0001))
}
#[tokio::test]
async fn an_observed_line_becomes_a_transcript_message_keyed_to_the_activity() -> TestResult {
let (sender, mut events) = mpsc::unbounded_channel();
let (context, cancellation) = ActivityContext::with_transcript(
workflow_id(),
run_id(),
ActivityId::from_sequence_position(4),
2,
sender,
);
drop(cancellation);
let transcript = CommandTranscript::new(&context);
transcript.on_line(CommandStream::Stdout, "building");
let event = events.recv().await.ok_or("the line must be emitted")?;
assert_eq!(event.workflow_id, workflow_id());
assert_eq!(event.run_id, run_id());
assert_eq!(event.activity_id, ActivityId::from_sequence_position(4));
assert_eq!(event.attempt, 2);
assert_eq!(event.agent_role, STDOUT_PRODUCER);
assert_eq!(event.agent_id, Uuid::nil());
assert_eq!(event.worker_seq, 0);
assert_eq!(event.store_seq, None);
assert!(
!event.ephemeral,
"command output is retained transcript, never a WS-only delta"
);
let ActivityEventKind::Message { role, text } = event.kind else {
return Err("an output line must be a Message".into());
};
assert_eq!(role, MessageRole::Tool);
assert_eq!(text, "building");
Ok(())
}
#[tokio::test]
async fn both_streams_are_labelled_and_share_one_monotonic_sequence() -> TestResult {
let (sender, mut events) = mpsc::unbounded_channel();
let (context, cancellation) = ActivityContext::with_transcript(
workflow_id(),
run_id(),
ActivityId::from_sequence_position(1),
0,
sender,
);
drop(cancellation);
let transcript = CommandTranscript::new(&context);
transcript.on_line(CommandStream::Stdout, "out");
transcript.on_line(CommandStream::Stderr, "err");
transcript.on_line(CommandStream::Stdout, "out again");
let mut observed = Vec::new();
for _ in 0..3u8 {
let event = events.recv().await.ok_or("every line must be emitted")?;
let ActivityEventKind::Message { text, .. } = event.kind else {
return Err("an output line must be a Message".into());
};
observed.push((event.agent_role, text, event.worker_seq));
}
assert_eq!(
observed,
vec![
(STDOUT_PRODUCER.to_owned(), "out".to_owned(), 0),
(STDERR_PRODUCER.to_owned(), "err".to_owned(), 1),
(STDOUT_PRODUCER.to_owned(), "out again".to_owned(), 2),
]
);
Ok(())
}
#[tokio::test]
async fn a_context_without_a_seam_publishes_nothing() {
let (context, cancellation) = ActivityContext::new(
workflow_id(),
run_id(),
ActivityId::from_sequence_position(1),
0,
);
drop(cancellation);
let transcript = CommandTranscript::new(&context);
transcript.on_line(CommandStream::Stdout, "ignored");
}
#[tokio::test]
async fn a_closed_seam_does_not_fail_the_command() -> TestResult {
let (sender, events) = mpsc::unbounded_channel();
let (context, cancellation) = ActivityContext::with_transcript(
workflow_id(),
run_id(),
ActivityId::from_sequence_position(1),
0,
sender,
);
drop(cancellation);
drop(events);
let transcript = CommandTranscript::new(&context);
transcript.on_line(CommandStream::Stdout, "into the void");
transcript.on_line(CommandStream::Stderr, "also lost");
Ok(())
}
}