use std::io;
use std::path::Path;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::atomic::AtomicBool;
use rhema_agent_tools::{
Approval, ExpandPastes, Overlay, ShellTools, Skill, revive_dangling_calls,
};
use rhema_core::{
Agent, AgentError, ChatParams, CommitHook, EventSink, Message, ModelOutput, Role, Session,
StopReason, ToolContext, ToolResolver, ToolSpec, run_turn,
};
use crate::config::{CodingAgentConfig, Steering};
const ENVIRONMENT_STATE_KEY: &str = "environment_report";
pub fn clear_environment_stamp(session: &mut Session) {
if let Some(state) = session.driver_state.as_object_mut() {
state.remove(ENVIRONMENT_STATE_KEY);
}
}
#[derive(Debug, PartialEq, Eq, thiserror::Error)]
pub enum CodingAgentError {
#[error(transparent)]
Agent(#[from] AgentError),
#[error("workspace escape: {0}")]
WorkspaceEscape(std::path::PathBuf),
#[error("system prompt template: {0}")]
Template(String),
}
impl From<minijinja::Error> for CodingAgentError {
fn from(err: minijinja::Error) -> Self {
let line = err
.line()
.map(|line| format!(", line {line}"))
.unwrap_or_default();
Self::Template(format!("{err}{line}"))
}
}
pub fn stamped_system(
base: Option<&str>,
workspace: &Path,
extra: &[ToolSpec],
) -> Result<String, CodingAgentError> {
let home = rhema_config::Home::discover().ok();
let agents = home.as_ref().map(|h| h.agents_skills());
let global = home.as_ref().map(|h| h.skills());
let skills =
rhema_agent_tools::skills::discover(agents.as_deref(), global.as_deref(), workspace);
stamped_system_with_sections(base, workspace, extra, &skills, None, None, &[], None, None)
}
pub fn stamped_system_with_skills(
base: Option<&str>,
workspace: &Path,
extra: &[ToolSpec],
skills: &[Skill],
) -> Result<String, CodingAgentError> {
stamped_system_with_sections(base, workspace, extra, skills, None, None, &[], None, None)
}
#[expect(
clippy::too_many_arguments,
reason = "one parameter per prompt section keeps the render's inputs flat"
)]
pub fn stamped_system_with_sections(
base: Option<&str>,
workspace: &Path,
extra: &[ToolSpec],
skills: &[Skill],
environment: Option<&str>,
user: Option<&str>,
mcp: &[rhema_agent_tools::prompt::McpServerLine],
vision: Option<&str>,
jobs: Option<&str>,
) -> Result<String, CodingAgentError> {
let mut tools = ShellTools::builtin_specs();
tools.extend_from_slice(extra);
rhema_agent_tools::render_system_with_sections(
base.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
workspace,
skills,
&tools,
environment,
user,
mcp,
vision,
jobs,
)
.map_err(CodingAgentError::from)
}
pub struct CodingAgent {
agent: Agent,
config: CodingAgentConfig,
session: Session,
tools: Arc<ShellTools>,
resolver: Arc<dyn ToolResolver>,
transcript: Option<Arc<Mutex<rhema_transcript::Transcript>>>,
turns: u64,
}
impl CodingAgent {
pub fn new(
provider: Arc<dyn rhema_core::ModelProvider>,
config: CodingAgentConfig,
) -> Result<Self, CodingAgentError> {
let provider: Arc<dyn rhema_core::ModelProvider> = Arc::new(ExpandPastes::new(provider));
let tools = config.tools.clone().unwrap_or_else(|| {
Arc::new(ShellTools::with_env(
config.workspace.clone(),
&config.shell_env,
))
});
let resolver: Arc<dyn ToolResolver> = match config.extra_resolver.clone() {
Some(extra) => Arc::new(Overlay::new(tools.clone(), extra)),
None => tools.clone(),
};
let mut extra: Vec<ToolSpec> = config
.extra_resolver
.as_ref()
.map(|resolver| resolver.specs())
.unwrap_or_default();
let (provider, resolver, transcript) = match config.transcript {
Some(settings) => {
extra.push(rhema_transcript::Zoom::spec());
let (provider, zoom, mirror) = rhema_transcript::attach_with_summarizer(
provider,
settings,
config.summary_worker.clone(),
);
let resolver = Arc::new(Overlay::new(resolver, Arc::new(zoom)));
(
provider,
Arc::clone(&resolver) as Arc<dyn ToolResolver>,
Some(mirror),
)
}
None => (provider, resolver, None),
};
let mut skills = match &config.skills {
Some(skills) => skills.clone(),
None => {
let home = config
.home
.clone()
.or_else(|| rhema_config::Home::discover().ok());
let agents = home.as_ref().map(|h| h.agents_skills());
let global = home.as_ref().map(|h| h.skills());
rhema_agent_tools::skills::discover(
agents.as_deref(),
global.as_deref(),
tools.workspace(),
)
}
};
skills.retain(|skill| !config.disabled_skills.contains(&skill.name));
let mcp: Vec<rhema_agent_tools::prompt::McpServerLine> = config
.mcp_servers
.iter()
.filter(|line| !config.disabled_mcp_servers.contains(&line.server))
.cloned()
.collect();
let system = match config.system.clone() {
Some(system) => system,
None => stamped_system_with_sections(
config.base_system.as_deref(),
tools.workspace(),
&extra,
&skills,
config.environment.as_deref(),
config.user_preferences.as_deref(),
&mcp,
config.vision.as_deref(),
config.jobs.as_deref(),
)?,
};
let skills_bytes = match config.system {
Some(_) => 0,
None => {
let mut bare = ShellTools::builtin_specs();
bare.extend_from_slice(&extra);
rhema_agent_tools::render_system_with_sections(
config
.base_system
.as_deref()
.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
tools.workspace(),
&[],
&bare,
config.environment.as_deref(),
config.user_preferences.as_deref(),
&mcp,
config.vision.as_deref(),
config.jobs.as_deref(),
)
.map(|without| system.len().saturating_sub(without.len()))
.unwrap_or(0)
}
};
if let Some(mirror) = &transcript {
let mut mirror = mirror.lock().expect("transcript lock poisoned");
mirror.set_skills_bytes(skills_bytes);
mirror.set_skills(
skills
.iter()
.map(|skill| rhema_transcript::SkillText {
name: skill.name.clone(),
body: skill.body.clone(),
})
.collect(),
);
}
let agent = Agent::builder()
.provider(provider)
.resolver(resolver.clone())
.system(system)
.params(ChatParams {
thinking: config.thinking,
..ChatParams::new(config.model.clone())
})
.build()?;
let mut built = Self {
agent,
config,
session: Session::default(),
tools,
resolver,
transcript,
turns: 0,
};
built.stamp_environment();
Ok(built)
}
pub fn from_session(
provider: Arc<dyn rhema_core::ModelProvider>,
mut config: CodingAgentConfig,
session: Session,
) -> Result<Self, CodingAgentError> {
if let Some(stamped) = session
.driver_state
.get(ENVIRONMENT_STATE_KEY)
.and_then(|value| value.as_str())
&& config.environment.is_some()
{
config.environment = Some(stamped.to_owned());
}
let mut agent = Self::new(provider, config)?;
agent.session = session;
agent.stamp_environment();
Ok(agent)
}
fn stamp_environment(&mut self) {
let Some(report) = self.config.environment.clone() else {
return;
};
let state = &mut self.session.driver_state;
if !state.is_object() {
if !state.is_null() {
return;
}
*state = serde_json::json!({});
}
if state.get(ENVIRONMENT_STATE_KEY).is_none() {
state[ENVIRONMENT_STATE_KEY] = serde_json::json!(report);
}
}
pub fn set_commit_hook(&mut self, hook: Option<CommitHook>) {
self.agent.set_commit_hook(hook);
}
pub fn session(&self) -> &Session {
&self.session
}
pub fn system(&self) -> Option<&str> {
self.agent.system()
}
pub fn transcript_stats(&self) -> Option<rhema_transcript::TranscriptStats> {
self.transcript
.as_ref()
.map(|mirror| mirror.lock().expect("transcript lock poisoned").stats())
}
pub fn carry_transcript_stats(&mut self, saved: &rhema_transcript::TranscriptStats) {
if let Some(mirror) = &self.transcript {
mirror
.lock()
.expect("transcript lock poisoned")
.carry(saved);
}
}
pub fn transcript_sections(&self) -> Vec<rhema_transcript::Section> {
match &self.transcript {
Some(mirror) => mirror
.lock()
.expect("transcript lock poisoned")
.saved_sections(),
None => Vec::new(),
}
}
pub fn adopt_transcript_sections(&mut self, sections: Vec<rhema_transcript::Section>) {
if let Some(mirror) = &self.transcript {
mirror
.lock()
.expect("transcript lock poisoned")
.adopt_sections(sections);
}
}
pub fn reset_transcript_stats(&mut self) {
if let Some(mirror) = &self.transcript {
mirror
.lock()
.expect("transcript lock poisoned")
.reset_stats();
}
}
pub fn session_mut(&mut self) -> &mut Session {
&mut self.session
}
pub async fn turn(
&mut self,
input: impl Into<String>,
sink: &mut dyn EventSink,
approval: Arc<dyn Approval>,
) -> Result<ModelOutput, CodingAgentError> {
self.turn_with_cancel(input, sink, approval, Arc::new(AtomicBool::new(false)))
.await
}
pub async fn turn_with_cancel(
&mut self,
input: impl Into<String>,
sink: &mut dyn EventSink,
approval: Arc<dyn Approval>,
cancelled: Arc<AtomicBool>,
) -> Result<ModelOutput, CodingAgentError> {
self.turn_message_with_cancel(Message::text(Role::User, input), sink, approval, cancelled)
.await
}
pub async fn turn_message(
&mut self,
message: Message,
sink: &mut dyn EventSink,
approval: Arc<dyn Approval>,
) -> Result<ModelOutput, CodingAgentError> {
self.turn_message_with_cancel(message, sink, approval, Arc::new(AtomicBool::new(false)))
.await
}
pub async fn turn_message_with_cancel(
&mut self,
message: Message,
sink: &mut dyn EventSink,
approval: Arc<dyn Approval>,
cancelled: Arc<AtomicBool>,
) -> Result<ModelOutput, CodingAgentError> {
self.turns += 1;
let ctx = ToolContext {
cwd: self.tools.workspace().to_path_buf(),
run_id: format!("turn-{}", self.turns),
cancelled,
};
self.tools.install_approval(approval);
revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
let result = match self.config.steering.clone() {
Some(steering) => {
self.steps(Some(message), sink, &ctx, Some(steering.as_ref()))
.await
}
None => {
run_turn(
&mut self.agent,
&mut self.session,
message,
sink,
&ctx,
self.config.max_turn_steps,
)
.await
}
};
self.tools.clear_approval();
Ok(result?)
}
pub async fn retry_turn_with_cancel(
&mut self,
sink: &mut dyn EventSink,
approval: Arc<dyn Approval>,
cancelled: Arc<AtomicBool>,
) -> Result<ModelOutput, CodingAgentError> {
self.turns += 1;
let ctx = ToolContext {
cwd: self.tools.workspace().to_path_buf(),
run_id: format!("turn-{}", self.turns),
cancelled,
};
self.tools.install_approval(approval);
revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
let steering = self.config.steering.clone();
let result = self.steps(None, sink, &ctx, steering.as_deref()).await;
self.tools.clear_approval();
Ok(result?)
}
async fn steps(
&mut self,
input: Option<Message>,
sink: &mut dyn EventSink,
ctx: &ToolContext,
steering: Option<&dyn Steering>,
) -> Result<ModelOutput, AgentError> {
let mut announced: Vec<String> = Vec::new();
if let Some(steering) = steering {
self.append_notifications(steering, &mut announced);
}
if let Some(input) = input {
self.session.history.push(input);
}
for _ in 0..self.config.max_turn_steps {
let step = self.agent.step(&mut self.session, sink, ctx).await?;
if step.output.stop_reason != StopReason::ToolUse {
return Ok(step.output);
}
if let Some(steering) = steering {
self.append_notifications(steering, &mut announced);
}
}
Err(AgentError::MaxIterations(self.config.max_turn_steps))
}
fn append_notifications(&mut self, steering: &dyn Steering, announced: &mut Vec<String>) {
let lines: Vec<String> = steering
.drain()
.into_iter()
.filter(|line| !announced.contains(line))
.collect();
if lines.is_empty() {
return;
}
announced.extend(lines.iter().cloned());
self.session
.history
.push(Message::text(Role::User, lines.join("\n")));
}
pub fn save_session_to(&self, path: &impl AsRef<Path>) -> io::Result<()> {
let json = serde_json::to_string_pretty(&self.session)
.map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
rhema_config::privacy::private_write_atomic(path.as_ref(), json.as_bytes())
}
pub fn load_session_from(path: &impl AsRef<Path>) -> io::Result<Session> {
let raw = std::fs::read_to_string(path)?;
serde_json::from_str(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
}
}
#[cfg(test)]
mod paste_tests {
use super::*;
use rhema_agent_tools::AllowAll;
use rhema_core::{
Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
};
use std::sync::Mutex;
struct Recording(Mutex<Vec<String>>);
#[async_trait]
impl ModelProvider for Recording {
async fn complete(
&self,
req: ModelRequest,
_sink: &mut dyn EventSink,
) -> Result<ModelOutput, ProviderError> {
let mut user = String::new();
for message in &req.messages {
if message.role != Role::User {
continue;
}
for content in &message.content {
if let Content::Text(text) = content {
user.push_str(text);
}
}
}
self.0.lock().unwrap().push(user);
Ok(ModelOutput {
message: Message::text(Role::Assistant, "done"),
stop_reason: StopReason::EndTurn,
usage: Usage::default(),
})
}
}
struct NoopSink;
impl EventSink for NoopSink {
fn emit(&mut self, _event: Event<'_>) {}
}
#[tokio::test]
async fn the_provider_sees_pastes_expanded_and_the_session_keeps_the_chips() {
let payload = "l1\nl2\nl3\nl4";
let embedded = rhema_core::paste::embed(
&format!("look at {} this", rhema_core::paste::MARK),
&[payload.to_owned()],
);
let recording = Arc::new(Recording(Mutex::new(Vec::new())));
let mut agent = CodingAgent::new(
recording.clone(),
CodingAgentConfig::new("m", std::env::temp_dir()),
)
.unwrap();
agent
.turn(&embedded, &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
assert_eq!(
agent.session().history[0],
Message::text(Role::User, embedded.clone())
);
let seen = recording.0.lock().unwrap();
assert_eq!(seen.len(), 1);
assert_eq!(seen[0], format!("look at {payload} this"));
}
struct RecordingBlocks(Mutex<Vec<Vec<Content>>>);
#[async_trait]
impl ModelProvider for RecordingBlocks {
async fn complete(
&self,
req: ModelRequest,
_sink: &mut dyn EventSink,
) -> Result<ModelOutput, ProviderError> {
let mut blocks = Vec::new();
for message in &req.messages {
if message.role == Role::User {
blocks.extend(message.content.iter().cloned());
}
}
self.0.lock().unwrap().push(blocks);
Ok(ModelOutput {
message: Message::text(Role::Assistant, "done"),
stop_reason: StopReason::EndTurn,
usage: Usage::default(),
})
}
}
fn blocks_agent(recording: &Arc<RecordingBlocks>) -> CodingAgent {
CodingAgent::new(
recording.clone(),
CodingAgentConfig::new("m", std::env::temp_dir()),
)
.unwrap()
}
#[tokio::test]
async fn turn_message_delivers_text_then_image_blocks_to_the_provider() {
let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
let mut agent = blocks_agent(&recording);
let message = Message::text_with_images(
Role::User,
"what is in this picture?",
vec![
rhema_core::Image::new(rhema_core::ImageMediaType::Png, "aGVsbG8="),
rhema_core::Image::new(rhema_core::ImageMediaType::Jpeg, "d29ybGQ="),
],
);
agent
.turn_message(message, &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let seen = recording.0.lock().unwrap();
assert_eq!(seen.len(), 1);
assert_eq!(
seen[0],
vec![
Content::Text("what is in this picture?".into()),
Content::Image(rhema_core::Image::new(
rhema_core::ImageMediaType::Png,
"aGVsbG8="
)),
Content::Image(rhema_core::Image::new(
rhema_core::ImageMediaType::Jpeg,
"d29ybGQ="
)),
]
);
assert_eq!(agent.session().history[0].content, seen[0]);
}
#[tokio::test]
async fn the_text_entry_still_sends_exactly_one_text_block() {
let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
let mut agent = blocks_agent(&recording);
agent
.turn("plain words", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let seen = recording.0.lock().unwrap();
assert_eq!(seen.len(), 1);
assert_eq!(seen[0], vec![Content::Text("plain words".into())]);
}
}
#[cfg(test)]
mod extra_tools_tests {
use super::*;
use rhema_agent_tools::AllowAll;
use rhema_core::{
Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
ToolResult, Usage, async_trait,
};
use std::sync::Mutex;
use serde_json::json;
struct NoopSink;
impl EventSink for NoopSink {
fn emit(&mut self, _event: Event<'_>) {}
}
struct Echo;
#[async_trait]
impl rhema_core::ToolResolver for Echo {
fn specs(&self) -> Vec<rhema_core::ToolSpec> {
vec![rhema_core::ToolSpec::new(
"echo",
"Repeats the text argument back. The extension's canary.",
json!({
"type": "object",
"properties": { "text": { "type": "string" } },
"required": ["text"]
}),
)]
}
async fn resolve(&self, call: ToolCall, _ctx: &ToolContext) -> ToolResult {
let text = call.args.get("text").and_then(|v| v.as_str()).unwrap_or("");
ToolResult::success(call.id, format!("echo: {text}"))
}
}
struct Scripted(Mutex<Vec<Vec<String>>>);
#[async_trait]
impl ModelProvider for Scripted {
async fn complete(
&self,
req: ModelRequest,
_sink: &mut dyn EventSink,
) -> Result<ModelOutput, ProviderError> {
let round = self.0.lock().unwrap().len();
self.0
.lock()
.unwrap()
.push(req.tools.iter().map(|t| t.name.clone()).collect());
if round == 0 {
Ok(ModelOutput {
message: Message {
role: Role::Assistant,
content: vec![Content::ToolCall(ToolCall {
id: "c1".into(),
name: "echo".into(),
args: json!({"text": "through the overlay"}),
..ToolCall::default()
})],
},
stop_reason: StopReason::ToolUse,
usage: Usage::default(),
})
} else {
Ok(ModelOutput {
message: Message::text(Role::Assistant, "final"),
stop_reason: StopReason::EndTurn,
usage: Usage::default(),
})
}
}
}
#[tokio::test]
async fn the_extra_resolver_reaches_the_request_and_resolves() {
let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
let config = CodingAgentConfig {
extra_resolver: Some(Arc::new(Echo)),
..CodingAgentConfig::new("m", std::env::temp_dir())
};
let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
let output = agent
.turn("use the tool", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let text: String = output
.message
.content
.iter()
.filter_map(|c| match c {
Content::Text(text) => Some(text.as_str()),
_ => None,
})
.collect();
assert_eq!(text, "final");
let rounds = scripted.0.lock().unwrap();
assert_eq!(rounds.len(), 2);
assert!(
rounds[0].contains(&"echo".to_owned()) && rounds[0].contains(&"bash".to_owned()),
"the tools array carries both: {:?}",
rounds[0]
);
let prompt = agent
.system()
.expect("the agent was built with a system prompt");
assert!(prompt.contains("echo"), "{prompt}");
assert!(
prompt.contains("canary"),
"the extension's description is in the prompt: {prompt}"
);
}
#[tokio::test]
async fn the_default_advertises_zoom_and_none_skips_it() {
let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
let mut agent = CodingAgent::new(
scripted.clone(),
CodingAgentConfig::new("m", std::env::temp_dir()),
)
.unwrap();
agent
.turn("hi", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
{
let rounds = scripted.0.lock().unwrap();
assert_eq!(rounds[0], vec!["bash".to_owned(), "zoom".to_owned()]);
}
let prompt = agent
.system()
.expect("the agent was built with a system prompt");
assert!(prompt.contains("zoom"), "{prompt}");
assert!(
!agent
.system()
.is_some_and(|prompt| prompt.contains("canary"))
);
let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
let config = CodingAgentConfig {
transcript: None,
..CodingAgentConfig::new("m", std::env::temp_dir())
};
let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
agent
.turn("hi", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let rounds = scripted.0.lock().unwrap();
assert_eq!(rounds[0], vec!["bash".to_owned()]);
let prompt = agent
.system()
.expect("the agent was built with a system prompt");
assert!(!prompt.contains("zoom"), "{prompt}");
assert!(
!agent
.system()
.is_some_and(|prompt| prompt.contains("canary"))
);
}
#[test]
fn pre_collected_skills_reach_the_prompt_without_a_filesystem() {
let workspace = std::env::temp_dir();
let skill = Skill {
name: "deploy".into(),
description: "Triggers on deploy tasks.".into(),
license: None,
path: workspace.join(".rhema/skills/deploy/SKILL.md"),
body: "the deploy rules".into(),
};
let prompt = stamped_system_with_skills(None, &workspace, &[], &[skill]).unwrap();
assert!(prompt.contains("deploy"), "{prompt}");
assert!(prompt.contains("Triggers on deploy tasks."), "{prompt}");
let without = stamped_system(None, &workspace, &[]).unwrap();
assert!(!without.contains("Triggers on deploy tasks."), "{without}");
}
#[test]
fn stamped_system_carries_the_extra_specs() {
let workspace = std::env::temp_dir();
let spec = rhema_core::ToolSpec::new(
"echo",
"Repeats text back. The extension's canary.",
json!({"type": "object"}),
);
let with = stamped_system(None, &workspace, std::slice::from_ref(&spec)).unwrap();
assert!(with.contains("echo") && with.contains("canary"), "{with}");
let without = stamped_system(None, &workspace, &[]).unwrap();
assert!(!without.contains("canary"), "{without}");
}
#[test]
fn the_optional_sections_reach_the_prompt_where_the_template_places_them() {
let workspace = std::env::temp_dir();
let prompt = stamped_system_with_sections(
None,
&workspace,
&[],
&[],
Some("## Environment\n\n- OS: test os"),
Some("prefer concise answers"),
&[],
Some("pipe bytes into `see`"),
Some("`job run` starts it"),
)
.unwrap();
assert!(prompt.contains("- OS: test os"), "{prompt}");
assert!(prompt.contains("prefer concise answers"), "{prompt}");
assert!(prompt.contains("pipe bytes into `see`"), "{prompt}");
assert!(prompt.contains("`job run` starts it"), "{prompt}");
let bare = stamped_system_with_skills(None, &workspace, &[], &[]).unwrap();
assert!(!bare.contains("- OS: test os"), "{bare}");
assert!(!bare.contains("`job run` starts it"), "{bare}");
assert!(!bare.contains("prefer concise answers"), "{bare}");
assert!(!bare.contains("pipe bytes into `see`"), "{bare}");
}
}
#[cfg(test)]
mod environment_stamp_tests {
use super::*;
use rhema_agent_tools::AllowAll;
use rhema_core::{
Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
};
struct Ender;
#[async_trait]
impl ModelProvider for Ender {
async fn complete(
&self,
_req: ModelRequest,
_sink: &mut dyn EventSink,
) -> Result<ModelOutput, ProviderError> {
Ok(ModelOutput {
message: Message::text(Role::Assistant, "done"),
stop_reason: StopReason::EndTurn,
usage: Usage::default(),
})
}
}
struct NoopSink;
impl EventSink for NoopSink {
fn emit(&mut self, _event: Event<'_>) {}
}
fn config(report: &str) -> CodingAgentConfig {
CodingAgentConfig {
environment: Some(report.to_owned()),
..CodingAgentConfig::new("m", std::env::temp_dir())
}
}
#[test]
fn a_fresh_conversation_stamps_at_construction() {
let agent = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
assert_eq!(
agent
.session()
.driver_state
.get(ENVIRONMENT_STATE_KEY)
.and_then(|v| v.as_str()),
Some("REPORT-A")
);
}
#[tokio::test]
async fn a_resumed_conversation_keeps_its_stamped_report() {
let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
first
.turn("hi", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let session = first.session().clone();
let second =
CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
let prompt = second.system().expect("the agent has a prompt");
assert!(prompt.contains("REPORT-A"), "{prompt}");
assert!(!prompt.contains("REPORT-B"), "{prompt}");
assert_eq!(
second
.session()
.driver_state
.get(ENVIRONMENT_STATE_KEY)
.and_then(|v| v.as_str()),
Some("REPORT-A"),
"the stamp is idempotent across rebuilds"
);
}
#[tokio::test]
async fn a_pre_report_session_stamps_the_fresh_render_once() {
let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
first
.turn("hi", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let session = Session {
driver_state: serde_json::Value::Null,
..first.session().clone()
};
let second =
CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
let prompt = second.system().expect("the agent has a prompt");
assert!(prompt.contains("REPORT-B"), "{prompt}");
assert_eq!(
second
.session()
.driver_state
.get(ENVIRONMENT_STATE_KEY)
.and_then(|v| v.as_str()),
Some("REPORT-B"),
"stamped once on arrival, carried forward"
);
}
#[test]
fn driver_state_shapes_coexist_or_stay_untouched() {
let session = Session {
driver_state: serde_json::json!({ "someone-elses": "shape" }),
..Session::default()
};
let agent =
CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), session).unwrap();
let state = &agent.session().driver_state;
assert_eq!(
state.get("someone-elses"),
Some(&serde_json::json!("shape")),
"the other driver's key stands"
);
assert_eq!(
state.get(ENVIRONMENT_STATE_KEY).and_then(|v| v.as_str()),
Some("REPORT-A"),
"the stamp coexists beside it"
);
let opaque = Session {
driver_state: serde_json::json!("a driver's opaque shape"),
..Session::default()
};
let agent = CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), opaque).unwrap();
assert_eq!(
agent.session().driver_state,
serde_json::json!("a driver's opaque shape"),
"a non-object shape is never clobbered"
);
assert!(
agent
.system()
.is_some_and(|prompt| prompt.contains("REPORT-A")),
"the section still renders for this conversation"
);
}
}
#[cfg(test)]
mod transcript_tests {
use super::*;
use rhema_agent_tools::AllowAll;
use rhema_core::{
Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
ToolResult, Usage, async_trait,
};
use std::sync::Mutex;
struct Recording(Mutex<Vec<Vec<Message>>>);
#[async_trait]
impl ModelProvider for Recording {
async fn complete(
&self,
req: ModelRequest,
_sink: &mut dyn EventSink,
) -> Result<ModelOutput, ProviderError> {
self.0.lock().unwrap().push(req.messages);
Ok(ModelOutput {
message: Message::text(Role::Assistant, "done"),
stop_reason: StopReason::EndTurn,
usage: Usage::default(),
})
}
}
struct NoopSink;
impl EventSink for NoopSink {
fn emit(&mut self, _event: Event<'_>) {}
}
fn big_history(rounds: usize) -> Session {
let mut history = vec![Message::text(Role::User, "fix the build")];
for i in 0..rounds {
history.push(Message {
role: Role::Assistant,
content: vec![
Content::Text(format!("round {i}")),
Content::ToolCall(ToolCall {
id: format!("c{i}"),
name: "bash".into(),
args: serde_json::json!({ "command": "make" }),
provider_data: serde_json::Value::Null,
}),
],
});
history.push(Message {
role: Role::Tool,
content: vec![Content::ToolResult(ToolResult::success(
format!("c{i}"),
"x".repeat(2048),
))],
});
}
Session {
history,
..Session::default()
}
}
#[tokio::test]
async fn off_the_history_rides_verbatim() {
let recording = Arc::new(Recording(Mutex::new(Vec::new())));
let config = CodingAgentConfig {
transcript: None,
..CodingAgentConfig::new("m", std::env::temp_dir())
};
let mut agent =
CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
agent
.turn("next", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let sent = recording.0.lock().unwrap();
assert_eq!(sent.len(), 1);
assert_eq!(sent[0].len(), 5 * 2 + 2);
assert_eq!(sent[0][0], Message::text(Role::User, "fix the build"));
}
#[tokio::test]
async fn on_by_default_a_big_history_is_viewed_and_the_session_stays_whole() {
let recording = Arc::new(Recording(Mutex::new(Vec::new())));
let config = CodingAgentConfig {
transcript: Some(rhema_transcript::TranscriptSettings {
max_view_bytes: 4_000,
..rhema_transcript::TranscriptSettings::default()
}),
..CodingAgentConfig::new("m", std::env::temp_dir())
};
let mut agent =
CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
let prompt = agent
.system()
.expect("the agent was built with a system prompt");
assert!(prompt.contains("zoom"), "{prompt}");
agent
.turn("next", &mut NoopSink, Arc::new(AllowAll))
.await
.unwrap();
let sent = recording.0.lock().unwrap();
let messages = &sent[0];
assert!(
messages.len() < 5 * 2 + 2,
"the view is shorter: {} messages",
messages.len()
);
let rhema_core::Content::Text(text) = &messages[0].content[0] else {
panic!("a text view message");
};
assert!(text.starts_with("<transcript>"), "{text}");
assert_eq!(messages.last().unwrap(), &Message::text(Role::User, "next"));
assert_eq!(agent.session().history.len(), 5 * 2 + 3);
}
}