use std::io;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use rhema_agent_tools::{Approval, ExpandPastes, ShellTools, revive_dangling_calls};
use rhema_core::{
Agent, AgentError, ChatParams, CommitHook, EventSink, Message, ModelOutput, Role, Session,
ToolContext, run_turn,
};
use crate::config::CodingAgentConfig;
#[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) -> Result<String, CodingAgentError> {
let global = rhema_config::Config::skills_dir().ok();
let skills = rhema_agent_tools::skills::discover(global.as_deref(), workspace);
rhema_agent_tools::render_system(
base.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
workspace,
&skills,
&rhema_agent_tools::ShellTools::builtin_specs(),
)
.map_err(CodingAgentError::from)
}
pub struct CodingAgent {
agent: Agent,
config: CodingAgentConfig,
session: Session,
tools: Arc<ShellTools>,
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 = Arc::new(ShellTools::new(config.workspace.clone()));
let system = match config.system.clone() {
Some(system) => system,
None => stamped_system(config.base_system.as_deref(), tools.workspace())?,
};
let agent = Agent::builder()
.provider(provider)
.resolver(tools.clone())
.system(system)
.params(ChatParams {
thinking: config.thinking,
..ChatParams::new(config.model.clone())
})
.build()?;
Ok(Self {
agent,
config,
session: Session::default(),
tools,
turns: 0,
})
}
pub fn from_session(
provider: Arc<dyn rhema_core::ModelProvider>,
config: CodingAgentConfig,
session: Session,
) -> Result<Self, CodingAgentError> {
let mut agent = Self::new(provider, config)?;
agent.session = session;
Ok(agent)
}
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) -> &str {
self.agent.system()
}
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.turns += 1;
let ctx = ToolContext {
cwd: self.tools.workspace().to_path_buf(),
run_id: format!("turn-{}", self.turns),
cancelled: Arc::new(AtomicBool::new(false)),
};
revive_dangling_calls(&mut self.session, self.tools.as_ref(), &ctx).await;
self.tools.install_approval(approval);
let result = run_turn(
&mut self.agent,
&mut self.session,
Message::text(Role::User, input),
sink,
&ctx,
self.config.max_turn_steps,
)
.await;
self.tools.clear_approval();
Ok(result?)
}
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))?;
std::fs::write(path, json)
}
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"));
}
}