1use std::io;
16use std::path::Path;
17use std::sync::Arc;
18use std::sync::Mutex;
19use std::sync::atomic::AtomicBool;
20
21use rhema_agent_tools::{
22 Approval, ExpandPastes, Overlay, ShellTools, Skill, revive_dangling_calls,
23};
24use rhema_core::{
25 Agent, AgentError, ChatParams, CommitHook, EventSink, Message, ModelOutput, Role, Session,
26 StopReason, ToolContext, ToolResolver, ToolSpec, run_turn,
27};
28
29use crate::config::{CodingAgentConfig, Steering};
30
31const ENVIRONMENT_STATE_KEY: &str = "environment_report";
36
37pub fn clear_environment_stamp(session: &mut Session) {
46 if let Some(state) = session.driver_state.as_object_mut() {
47 state.remove(ENVIRONMENT_STATE_KEY);
48 }
49}
50
51#[derive(Debug, PartialEq, Eq, thiserror::Error)]
53pub enum CodingAgentError {
54 #[error(transparent)]
57 Agent(#[from] AgentError),
58 #[error("workspace escape: {0}")]
63 WorkspaceEscape(std::path::PathBuf),
64 #[error("system prompt template: {0}")]
68 Template(String),
69}
70
71impl From<minijinja::Error> for CodingAgentError {
72 fn from(err: minijinja::Error) -> Self {
73 let line = err
74 .line()
75 .map(|line| format!(", line {line}"))
76 .unwrap_or_default();
77 Self::Template(format!("{err}{line}"))
78 }
79}
80
81pub fn stamped_system(
92 base: Option<&str>,
93 workspace: &Path,
94 extra: &[ToolSpec],
95) -> Result<String, CodingAgentError> {
96 let home = rhema_config::Home::discover().ok();
99 let agents = home.as_ref().map(|h| h.agents_skills());
100 let global = home.as_ref().map(|h| h.skills());
101 let skills =
102 rhema_agent_tools::skills::discover(agents.as_deref(), global.as_deref(), workspace);
103 stamped_system_with_sections(base, workspace, extra, &skills, None, None, &[], None)
104}
105
106pub fn stamped_system_with_skills(
113 base: Option<&str>,
114 workspace: &Path,
115 extra: &[ToolSpec],
116 skills: &[Skill],
117) -> Result<String, CodingAgentError> {
118 stamped_system_with_sections(base, workspace, extra, skills, None, None, &[], None)
119}
120
121#[expect(
128 clippy::too_many_arguments,
129 reason = "one parameter per prompt section keeps the render's inputs flat"
130)]
131pub fn stamped_system_with_sections(
132 base: Option<&str>,
133 workspace: &Path,
134 extra: &[ToolSpec],
135 skills: &[Skill],
136 environment: Option<&str>,
137 user: Option<&str>,
138 mcp: &[rhema_agent_tools::prompt::McpServerLine],
139 vision: Option<&str>,
140) -> Result<String, CodingAgentError> {
141 let mut tools = ShellTools::builtin_specs();
142 tools.extend_from_slice(extra);
143 rhema_agent_tools::render_system_with_sections(
144 base.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
145 workspace,
146 skills,
147 &tools,
148 environment,
149 user,
150 mcp,
151 vision,
152 )
153 .map_err(CodingAgentError::from)
154}
155
156pub struct CodingAgent {
159 agent: Agent,
160 config: CodingAgentConfig,
161 session: Session,
162 tools: Arc<ShellTools>,
163 resolver: Arc<dyn ToolResolver>,
168 transcript: Option<Arc<Mutex<rhema_transcript::Transcript>>>,
171 turns: u64,
172}
173
174impl CodingAgent {
175 pub fn new(
181 provider: Arc<dyn rhema_core::ModelProvider>,
182 config: CodingAgentConfig,
183 ) -> Result<Self, CodingAgentError> {
184 let provider: Arc<dyn rhema_core::ModelProvider> = Arc::new(ExpandPastes::new(provider));
188 let tools = config.tools.clone().unwrap_or_else(|| {
189 Arc::new(ShellTools::with_env(
190 config.workspace.clone(),
191 &config.shell_env,
192 ))
193 });
194 let resolver: Arc<dyn ToolResolver> = match config.extra_resolver.clone() {
195 Some(extra) => Arc::new(Overlay::new(tools.clone(), extra)),
196 None => tools.clone(),
197 };
198 let mut extra: Vec<ToolSpec> = config
202 .extra_resolver
203 .as_ref()
204 .map(|resolver| resolver.specs())
205 .unwrap_or_default();
206 let (provider, resolver, transcript) = match config.transcript {
207 Some(settings) => {
208 extra.push(rhema_transcript::Zoom::spec());
209 let (provider, zoom, mirror) = rhema_transcript::attach_with_summarizer(
210 provider,
211 settings,
212 config.summary_worker.clone(),
213 );
214 let resolver = Arc::new(Overlay::new(resolver, Arc::new(zoom)));
215 (
216 provider,
217 Arc::clone(&resolver) as Arc<dyn ToolResolver>,
218 Some(mirror),
219 )
220 }
221 None => (provider, resolver, None),
222 };
223 let mut skills = match &config.skills {
228 Some(skills) => skills.clone(),
229 None => {
230 let home = config
231 .home
232 .clone()
233 .or_else(|| rhema_config::Home::discover().ok());
234 let agents = home.as_ref().map(|h| h.agents_skills());
235 let global = home.as_ref().map(|h| h.skills());
236 rhema_agent_tools::skills::discover(
237 agents.as_deref(),
238 global.as_deref(),
239 tools.workspace(),
240 )
241 }
242 };
243 skills.retain(|skill| !config.disabled_skills.contains(&skill.name));
244 let mcp: Vec<rhema_agent_tools::prompt::McpServerLine> = config
249 .mcp_servers
250 .iter()
251 .filter(|line| !config.disabled_mcp_servers.contains(&line.server))
252 .cloned()
253 .collect();
254 let system = match config.system.clone() {
255 Some(system) => system,
258 None => stamped_system_with_sections(
259 config.base_system.as_deref(),
260 tools.workspace(),
261 &extra,
262 &skills,
263 config.environment.as_deref(),
264 config.user_preferences.as_deref(),
265 &mcp,
266 config.vision.as_deref(),
267 )?,
268 };
269 let skills_bytes = match config.system {
274 Some(_) => 0,
275 None => {
276 let mut bare = ShellTools::builtin_specs();
283 bare.extend_from_slice(&extra);
284 rhema_agent_tools::render_system_with_sections(
285 config
286 .base_system
287 .as_deref()
288 .unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
289 tools.workspace(),
290 &[],
291 &bare,
292 config.environment.as_deref(),
293 config.user_preferences.as_deref(),
294 &mcp,
295 config.vision.as_deref(),
296 )
297 .map(|without| system.len().saturating_sub(without.len()))
298 .unwrap_or(0)
299 }
300 };
301 if let Some(mirror) = &transcript {
302 let mut mirror = mirror.lock().expect("transcript lock poisoned");
303 mirror.set_skills_bytes(skills_bytes);
304 mirror.set_skills(
307 skills
308 .iter()
309 .map(|skill| rhema_transcript::SkillText {
310 name: skill.name.clone(),
311 body: skill.body.clone(),
312 })
313 .collect(),
314 );
315 }
316 let agent = Agent::builder()
317 .provider(provider)
318 .resolver(resolver.clone())
321 .system(system)
322 .params(ChatParams {
323 thinking: config.thinking,
324 ..ChatParams::new(config.model.clone())
325 })
326 .build()?;
327 let mut built = Self {
328 agent,
329 config,
330 session: Session::default(),
331 tools,
332 resolver,
333 transcript,
334 turns: 0,
335 };
336 built.stamp_environment();
337 Ok(built)
338 }
339
340 pub fn from_session(
353 provider: Arc<dyn rhema_core::ModelProvider>,
354 mut config: CodingAgentConfig,
355 session: Session,
356 ) -> Result<Self, CodingAgentError> {
357 if let Some(stamped) = session
358 .driver_state
359 .get(ENVIRONMENT_STATE_KEY)
360 .and_then(|value| value.as_str())
361 && config.environment.is_some()
362 {
363 config.environment = Some(stamped.to_owned());
364 }
365 let mut agent = Self::new(provider, config)?;
366 agent.session = session;
367 agent.stamp_environment();
368 Ok(agent)
369 }
370
371 fn stamp_environment(&mut self) {
377 let Some(report) = self.config.environment.clone() else {
378 return;
379 };
380 let state = &mut self.session.driver_state;
381 if !state.is_object() {
382 if !state.is_null() {
383 return;
386 }
387 *state = serde_json::json!({});
388 }
389 if state.get(ENVIRONMENT_STATE_KEY).is_none() {
390 state[ENVIRONMENT_STATE_KEY] = serde_json::json!(report);
391 }
392 }
393
394 pub fn set_commit_hook(&mut self, hook: Option<CommitHook>) {
401 self.agent.set_commit_hook(hook);
402 }
403
404 pub fn session(&self) -> &Session {
406 &self.session
407 }
408
409 pub fn system(&self) -> Option<&str> {
412 self.agent.system()
413 }
414
415 pub fn transcript_stats(&self) -> Option<rhema_transcript::TranscriptStats> {
419 self.transcript
420 .as_ref()
421 .map(|mirror| mirror.lock().expect("transcript lock poisoned").stats())
422 }
423
424 pub fn carry_transcript_stats(&mut self, saved: &rhema_transcript::TranscriptStats) {
429 if let Some(mirror) = &self.transcript {
430 mirror
431 .lock()
432 .expect("transcript lock poisoned")
433 .carry(saved);
434 }
435 }
436
437 pub fn transcript_sections(&self) -> Vec<rhema_transcript::Section> {
442 match &self.transcript {
443 Some(mirror) => mirror
444 .lock()
445 .expect("transcript lock poisoned")
446 .saved_sections(),
447 None => Vec::new(),
448 }
449 }
450
451 pub fn adopt_transcript_sections(&mut self, sections: Vec<rhema_transcript::Section>) {
456 if let Some(mirror) = &self.transcript {
457 mirror
458 .lock()
459 .expect("transcript lock poisoned")
460 .adopt_sections(sections);
461 }
462 }
463
464 pub fn reset_transcript_stats(&mut self) {
469 if let Some(mirror) = &self.transcript {
470 mirror
471 .lock()
472 .expect("transcript lock poisoned")
473 .reset_stats();
474 }
475 }
476
477 pub fn session_mut(&mut self) -> &mut Session {
480 &mut self.session
481 }
482
483 pub async fn turn(
490 &mut self,
491 input: impl Into<String>,
492 sink: &mut dyn EventSink,
493 approval: Arc<dyn Approval>,
494 ) -> Result<ModelOutput, CodingAgentError> {
495 self.turn_with_cancel(input, sink, approval, Arc::new(AtomicBool::new(false)))
496 .await
497 }
498
499 pub async fn turn_with_cancel(
506 &mut self,
507 input: impl Into<String>,
508 sink: &mut dyn EventSink,
509 approval: Arc<dyn Approval>,
510 cancelled: Arc<AtomicBool>,
511 ) -> Result<ModelOutput, CodingAgentError> {
512 self.turn_message_with_cancel(Message::text(Role::User, input), sink, approval, cancelled)
513 .await
514 }
515
516 pub async fn turn_message(
520 &mut self,
521 message: Message,
522 sink: &mut dyn EventSink,
523 approval: Arc<dyn Approval>,
524 ) -> Result<ModelOutput, CodingAgentError> {
525 self.turn_message_with_cancel(message, sink, approval, Arc::new(AtomicBool::new(false)))
526 .await
527 }
528
529 pub async fn turn_message_with_cancel(
532 &mut self,
533 message: Message,
534 sink: &mut dyn EventSink,
535 approval: Arc<dyn Approval>,
536 cancelled: Arc<AtomicBool>,
537 ) -> Result<ModelOutput, CodingAgentError> {
538 self.turns += 1;
539 let ctx = ToolContext {
540 cwd: self.tools.workspace().to_path_buf(),
541 run_id: format!("turn-{}", self.turns),
542 cancelled,
543 };
544 self.tools.install_approval(approval);
548 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
549 let result = match self.config.steering.clone() {
553 Some(steering) => {
554 self.steps(Some(message), sink, &ctx, Some(steering.as_ref()))
555 .await
556 }
557 None => {
558 run_turn(
559 &mut self.agent,
560 &mut self.session,
561 message,
562 sink,
563 &ctx,
564 self.config.max_turn_steps,
565 )
566 .await
567 }
568 };
569 self.tools.clear_approval();
570 Ok(result?)
571 }
572
573 pub async fn retry_turn_with_cancel(
584 &mut self,
585 sink: &mut dyn EventSink,
586 approval: Arc<dyn Approval>,
587 cancelled: Arc<AtomicBool>,
588 ) -> Result<ModelOutput, CodingAgentError> {
589 self.turns += 1;
590 let ctx = ToolContext {
591 cwd: self.tools.workspace().to_path_buf(),
592 run_id: format!("turn-{}", self.turns),
593 cancelled,
594 };
595 self.tools.install_approval(approval);
599 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
600 let steering = self.config.steering.clone();
601 let result = self.steps(None, sink, &ctx, steering.as_deref()).await;
602 self.tools.clear_approval();
603 Ok(result?)
604 }
605
606 async fn steps(
618 &mut self,
619 input: Option<Message>,
620 sink: &mut dyn EventSink,
621 ctx: &ToolContext,
622 steering: Option<&dyn Steering>,
623 ) -> Result<ModelOutput, AgentError> {
624 let mut announced: Vec<String> = Vec::new();
627 if let Some(steering) = steering {
628 self.append_notifications(steering, &mut announced);
629 }
630 if let Some(input) = input {
631 self.session.history.push(input);
632 }
633 for _ in 0..self.config.max_turn_steps {
634 let step = self.agent.step(&mut self.session, sink, ctx).await?;
635 if step.output.stop_reason != StopReason::ToolUse {
636 return Ok(step.output);
637 }
638 if let Some(steering) = steering {
639 self.append_notifications(steering, &mut announced);
640 }
641 }
642 Err(AgentError::MaxIterations(self.config.max_turn_steps))
643 }
644
645 fn append_notifications(&mut self, steering: &dyn Steering, announced: &mut Vec<String>) {
651 let lines: Vec<String> = steering
652 .drain()
653 .into_iter()
654 .filter(|line| !announced.contains(line))
655 .collect();
656 if lines.is_empty() {
657 return;
658 }
659 announced.extend(lines.iter().cloned());
660 self.session
661 .history
662 .push(Message::text(Role::User, lines.join("\n")));
663 }
664
665 pub fn save_session_to(&self, path: &impl AsRef<Path>) -> io::Result<()> {
672 let json = serde_json::to_string_pretty(&self.session)
673 .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
674 rhema_config::privacy::private_write_atomic(path.as_ref(), json.as_bytes())
675 }
676
677 pub fn load_session_from(path: &impl AsRef<Path>) -> io::Result<Session> {
679 let raw = std::fs::read_to_string(path)?;
680 serde_json::from_str(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
681 }
682}
683
684#[cfg(test)]
685mod paste_tests {
686 use super::*;
687 use rhema_agent_tools::AllowAll;
688 use rhema_core::{
689 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
690 };
691 use std::sync::Mutex;
692
693 struct Recording(Mutex<Vec<String>>);
695
696 #[async_trait]
697 impl ModelProvider for Recording {
698 async fn complete(
699 &self,
700 req: ModelRequest,
701 _sink: &mut dyn EventSink,
702 ) -> Result<ModelOutput, ProviderError> {
703 let mut user = String::new();
704 for message in &req.messages {
705 if message.role != Role::User {
706 continue;
707 }
708 for content in &message.content {
709 if let Content::Text(text) = content {
710 user.push_str(text);
711 }
712 }
713 }
714 self.0.lock().unwrap().push(user);
715 Ok(ModelOutput {
716 message: Message::text(Role::Assistant, "done"),
717 stop_reason: StopReason::EndTurn,
718 usage: Usage::default(),
719 })
720 }
721 }
722
723 struct NoopSink;
724
725 impl EventSink for NoopSink {
726 fn emit(&mut self, _event: Event<'_>) {}
727 }
728
729 #[tokio::test]
730 async fn the_provider_sees_pastes_expanded_and_the_session_keeps_the_chips() {
731 let payload = "l1\nl2\nl3\nl4";
732 let embedded = rhema_core::paste::embed(
733 &format!("look at {} this", rhema_core::paste::MARK),
734 &[payload.to_owned()],
735 );
736 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
737 let mut agent = CodingAgent::new(
738 recording.clone(),
739 CodingAgentConfig::new("m", std::env::temp_dir()),
740 )
741 .unwrap();
742 agent
743 .turn(&embedded, &mut NoopSink, Arc::new(AllowAll))
744 .await
745 .unwrap();
746
747 assert_eq!(
749 agent.session().history[0],
750 Message::text(Role::User, embedded.clone())
751 );
752 let seen = recording.0.lock().unwrap();
754 assert_eq!(seen.len(), 1);
755 assert_eq!(seen[0], format!("look at {payload} this"));
756 }
757
758 struct RecordingBlocks(Mutex<Vec<Vec<Content>>>);
761
762 #[async_trait]
763 impl ModelProvider for RecordingBlocks {
764 async fn complete(
765 &self,
766 req: ModelRequest,
767 _sink: &mut dyn EventSink,
768 ) -> Result<ModelOutput, ProviderError> {
769 let mut blocks = Vec::new();
770 for message in &req.messages {
771 if message.role == Role::User {
772 blocks.extend(message.content.iter().cloned());
773 }
774 }
775 self.0.lock().unwrap().push(blocks);
776 Ok(ModelOutput {
777 message: Message::text(Role::Assistant, "done"),
778 stop_reason: StopReason::EndTurn,
779 usage: Usage::default(),
780 })
781 }
782 }
783
784 fn blocks_agent(recording: &Arc<RecordingBlocks>) -> CodingAgent {
785 CodingAgent::new(
786 recording.clone(),
787 CodingAgentConfig::new("m", std::env::temp_dir()),
788 )
789 .unwrap()
790 }
791
792 #[tokio::test]
793 async fn turn_message_delivers_text_then_image_blocks_to_the_provider() {
794 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
795 let mut agent = blocks_agent(&recording);
796 let message = Message::text_with_images(
797 Role::User,
798 "what is in this picture?",
799 vec![
800 rhema_core::Image::new(rhema_core::ImageMediaType::Png, "aGVsbG8="),
801 rhema_core::Image::new(rhema_core::ImageMediaType::Jpeg, "d29ybGQ="),
802 ],
803 );
804 agent
805 .turn_message(message, &mut NoopSink, Arc::new(AllowAll))
806 .await
807 .unwrap();
808
809 let seen = recording.0.lock().unwrap();
810 assert_eq!(seen.len(), 1);
811 assert_eq!(
812 seen[0],
813 vec![
814 Content::Text("what is in this picture?".into()),
815 Content::Image(rhema_core::Image::new(
816 rhema_core::ImageMediaType::Png,
817 "aGVsbG8="
818 )),
819 Content::Image(rhema_core::Image::new(
820 rhema_core::ImageMediaType::Jpeg,
821 "d29ybGQ="
822 )),
823 ]
824 );
825 assert_eq!(agent.session().history[0].content, seen[0]);
828 }
829
830 #[tokio::test]
831 async fn the_text_entry_still_sends_exactly_one_text_block() {
832 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
833 let mut agent = blocks_agent(&recording);
834 agent
835 .turn("plain words", &mut NoopSink, Arc::new(AllowAll))
836 .await
837 .unwrap();
838
839 let seen = recording.0.lock().unwrap();
840 assert_eq!(seen.len(), 1);
841 assert_eq!(seen[0], vec![Content::Text("plain words".into())]);
842 }
843}
844
845#[cfg(test)]
846mod extra_tools_tests {
847 use super::*;
848 use rhema_agent_tools::AllowAll;
849 use rhema_core::{
850 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
851 ToolResult, Usage, async_trait,
852 };
853 use std::sync::Mutex;
854
855 use serde_json::json;
856
857 struct NoopSink;
858
859 impl EventSink for NoopSink {
860 fn emit(&mut self, _event: Event<'_>) {}
861 }
862
863 struct Echo;
865
866 #[async_trait]
867 impl rhema_core::ToolResolver for Echo {
868 fn specs(&self) -> Vec<rhema_core::ToolSpec> {
869 vec![rhema_core::ToolSpec::new(
870 "echo",
871 "Repeats the text argument back. The extension's canary.",
872 json!({
873 "type": "object",
874 "properties": { "text": { "type": "string" } },
875 "required": ["text"]
876 }),
877 )]
878 }
879
880 async fn resolve(&self, call: ToolCall, _ctx: &ToolContext) -> ToolResult {
881 let text = call.args.get("text").and_then(|v| v.as_str()).unwrap_or("");
882 ToolResult::success(call.id, format!("echo: {text}"))
883 }
884 }
885
886 struct Scripted(Mutex<Vec<Vec<String>>>);
889
890 #[async_trait]
891 impl ModelProvider for Scripted {
892 async fn complete(
893 &self,
894 req: ModelRequest,
895 _sink: &mut dyn EventSink,
896 ) -> Result<ModelOutput, ProviderError> {
897 let round = self.0.lock().unwrap().len();
898 self.0
899 .lock()
900 .unwrap()
901 .push(req.tools.iter().map(|t| t.name.clone()).collect());
902 if round == 0 {
903 Ok(ModelOutput {
904 message: Message {
905 role: Role::Assistant,
906 content: vec![Content::ToolCall(ToolCall {
907 id: "c1".into(),
908 name: "echo".into(),
909 args: json!({"text": "through the overlay"}),
910 ..ToolCall::default()
911 })],
912 },
913 stop_reason: StopReason::ToolUse,
914 usage: Usage::default(),
915 })
916 } else {
917 Ok(ModelOutput {
918 message: Message::text(Role::Assistant, "final"),
919 stop_reason: StopReason::EndTurn,
920 usage: Usage::default(),
921 })
922 }
923 }
924 }
925
926 #[tokio::test]
927 async fn the_extra_resolver_reaches_the_request_and_resolves() {
928 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
929 let config = CodingAgentConfig {
930 extra_resolver: Some(Arc::new(Echo)),
931 ..CodingAgentConfig::new("m", std::env::temp_dir())
932 };
933 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
934 let output = agent
935 .turn("use the tool", &mut NoopSink, Arc::new(AllowAll))
936 .await
937 .unwrap();
938
939 let text: String = output
942 .message
943 .content
944 .iter()
945 .filter_map(|c| match c {
946 Content::Text(text) => Some(text.as_str()),
947 _ => None,
948 })
949 .collect();
950 assert_eq!(text, "final");
951 let rounds = scripted.0.lock().unwrap();
952 assert_eq!(rounds.len(), 2);
953 assert!(
954 rounds[0].contains(&"echo".to_owned()) && rounds[0].contains(&"bash".to_owned()),
955 "the tools array carries both: {:?}",
956 rounds[0]
957 );
958 let prompt = agent
961 .system()
962 .expect("the agent was built with a system prompt");
963 assert!(prompt.contains("echo"), "{prompt}");
964 assert!(
965 prompt.contains("canary"),
966 "the extension's description is in the prompt: {prompt}"
967 );
968 }
969
970 #[tokio::test]
971 async fn the_default_advertises_zoom_and_none_skips_it() {
972 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
974 let mut agent = CodingAgent::new(
975 scripted.clone(),
976 CodingAgentConfig::new("m", std::env::temp_dir()),
977 )
978 .unwrap();
979 agent
980 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
981 .await
982 .unwrap();
983 {
984 let rounds = scripted.0.lock().unwrap();
985 assert_eq!(rounds[0], vec!["bash".to_owned(), "zoom".to_owned()]);
986 }
987 let prompt = agent
988 .system()
989 .expect("the agent was built with a system prompt");
990 assert!(prompt.contains("zoom"), "{prompt}");
991 assert!(
992 !agent
993 .system()
994 .is_some_and(|prompt| prompt.contains("canary"))
995 );
996
997 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
999 let config = CodingAgentConfig {
1000 transcript: None,
1001 ..CodingAgentConfig::new("m", std::env::temp_dir())
1002 };
1003 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
1004 agent
1005 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1006 .await
1007 .unwrap();
1008 let rounds = scripted.0.lock().unwrap();
1009 assert_eq!(rounds[0], vec!["bash".to_owned()]);
1010 let prompt = agent
1011 .system()
1012 .expect("the agent was built with a system prompt");
1013 assert!(!prompt.contains("zoom"), "{prompt}");
1014 assert!(
1015 !agent
1016 .system()
1017 .is_some_and(|prompt| prompt.contains("canary"))
1018 );
1019 }
1020
1021 #[test]
1022 fn pre_collected_skills_reach_the_prompt_without_a_filesystem() {
1023 let workspace = std::env::temp_dir();
1024 let skill = Skill {
1025 name: "deploy".into(),
1026 description: "Triggers on deploy tasks.".into(),
1027 license: None,
1028 path: workspace.join(".rhema/skills/deploy/SKILL.md"),
1029 body: "the deploy rules".into(),
1030 };
1031 let prompt = stamped_system_with_skills(None, &workspace, &[], &[skill]).unwrap();
1032 assert!(prompt.contains("deploy"), "{prompt}");
1033 assert!(prompt.contains("Triggers on deploy tasks."), "{prompt}");
1034 let without = stamped_system(None, &workspace, &[]).unwrap();
1036 assert!(!without.contains("Triggers on deploy tasks."), "{without}");
1037 }
1038
1039 #[test]
1040 fn stamped_system_carries_the_extra_specs() {
1041 let workspace = std::env::temp_dir();
1042 let spec = rhema_core::ToolSpec::new(
1043 "echo",
1044 "Repeats text back. The extension's canary.",
1045 json!({"type": "object"}),
1046 );
1047 let with = stamped_system(None, &workspace, std::slice::from_ref(&spec)).unwrap();
1048 assert!(with.contains("echo") && with.contains("canary"), "{with}");
1049 let without = stamped_system(None, &workspace, &[]).unwrap();
1050 assert!(!without.contains("canary"), "{without}");
1051 }
1052
1053 #[test]
1054 fn the_optional_sections_reach_the_prompt_where_the_template_places_them() {
1055 let workspace = std::env::temp_dir();
1056 let prompt = stamped_system_with_sections(
1057 None,
1058 &workspace,
1059 &[],
1060 &[],
1061 Some("## Environment\n\n- OS: test os"),
1062 Some("prefer concise answers"),
1063 &[],
1064 Some("pipe bytes into `see`"),
1065 )
1066 .unwrap();
1067 assert!(prompt.contains("- OS: test os"), "{prompt}");
1068 assert!(prompt.contains("prefer concise answers"), "{prompt}");
1069 assert!(prompt.contains("pipe bytes into `see`"), "{prompt}");
1070 let bare = stamped_system_with_skills(None, &workspace, &[], &[]).unwrap();
1073 assert!(!bare.contains("- OS: test os"), "{bare}");
1074 assert!(!bare.contains("prefer concise answers"), "{bare}");
1075 assert!(!bare.contains("pipe bytes into `see`"), "{bare}");
1076 }
1077}
1078
1079#[cfg(test)]
1080mod environment_stamp_tests {
1081 use super::*;
1082 use rhema_agent_tools::AllowAll;
1083 use rhema_core::{
1084 Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
1085 };
1086
1087 struct Ender;
1089 #[async_trait]
1090 impl ModelProvider for Ender {
1091 async fn complete(
1092 &self,
1093 _req: ModelRequest,
1094 _sink: &mut dyn EventSink,
1095 ) -> Result<ModelOutput, ProviderError> {
1096 Ok(ModelOutput {
1097 message: Message::text(Role::Assistant, "done"),
1098 stop_reason: StopReason::EndTurn,
1099 usage: Usage::default(),
1100 })
1101 }
1102 }
1103
1104 struct NoopSink;
1105 impl EventSink for NoopSink {
1106 fn emit(&mut self, _event: Event<'_>) {}
1107 }
1108
1109 fn config(report: &str) -> CodingAgentConfig {
1110 CodingAgentConfig {
1111 environment: Some(report.to_owned()),
1112 ..CodingAgentConfig::new("m", std::env::temp_dir())
1113 }
1114 }
1115
1116 #[test]
1119 fn a_fresh_conversation_stamps_at_construction() {
1120 let agent = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1121 assert_eq!(
1122 agent
1123 .session()
1124 .driver_state
1125 .get(ENVIRONMENT_STATE_KEY)
1126 .and_then(|v| v.as_str()),
1127 Some("REPORT-A")
1128 );
1129 }
1130
1131 #[tokio::test]
1136 async fn a_resumed_conversation_keeps_its_stamped_report() {
1137 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1138 first
1139 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1140 .await
1141 .unwrap();
1142 let session = first.session().clone();
1143
1144 let second =
1147 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1148 let prompt = second.system().expect("the agent has a prompt");
1149 assert!(prompt.contains("REPORT-A"), "{prompt}");
1150 assert!(!prompt.contains("REPORT-B"), "{prompt}");
1151 assert_eq!(
1152 second
1153 .session()
1154 .driver_state
1155 .get(ENVIRONMENT_STATE_KEY)
1156 .and_then(|v| v.as_str()),
1157 Some("REPORT-A"),
1158 "the stamp is idempotent across rebuilds"
1159 );
1160 }
1161
1162 #[tokio::test]
1165 async fn a_pre_report_session_stamps_the_fresh_render_once() {
1166 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1167 first
1168 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1169 .await
1170 .unwrap();
1171 let session = Session {
1173 driver_state: serde_json::Value::Null,
1174 ..first.session().clone()
1175 };
1176
1177 let second =
1178 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1179 let prompt = second.system().expect("the agent has a prompt");
1180 assert!(prompt.contains("REPORT-B"), "{prompt}");
1181 assert_eq!(
1182 second
1183 .session()
1184 .driver_state
1185 .get(ENVIRONMENT_STATE_KEY)
1186 .and_then(|v| v.as_str()),
1187 Some("REPORT-B"),
1188 "stamped once on arrival, carried forward"
1189 );
1190 }
1191
1192 #[test]
1196 fn driver_state_shapes_coexist_or_stay_untouched() {
1197 let session = Session {
1198 driver_state: serde_json::json!({ "someone-elses": "shape" }),
1199 ..Session::default()
1200 };
1201 let agent =
1202 CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), session).unwrap();
1203 let state = &agent.session().driver_state;
1204 assert_eq!(
1205 state.get("someone-elses"),
1206 Some(&serde_json::json!("shape")),
1207 "the other driver's key stands"
1208 );
1209 assert_eq!(
1210 state.get(ENVIRONMENT_STATE_KEY).and_then(|v| v.as_str()),
1211 Some("REPORT-A"),
1212 "the stamp coexists beside it"
1213 );
1214
1215 let opaque = Session {
1216 driver_state: serde_json::json!("a driver's opaque shape"),
1217 ..Session::default()
1218 };
1219 let agent = CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), opaque).unwrap();
1220 assert_eq!(
1221 agent.session().driver_state,
1222 serde_json::json!("a driver's opaque shape"),
1223 "a non-object shape is never clobbered"
1224 );
1225 assert!(
1226 agent
1227 .system()
1228 .is_some_and(|prompt| prompt.contains("REPORT-A")),
1229 "the section still renders for this conversation"
1230 );
1231 }
1232}
1233
1234#[cfg(test)]
1235mod transcript_tests {
1236 use super::*;
1237 use rhema_agent_tools::AllowAll;
1238 use rhema_core::{
1239 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
1240 ToolResult, Usage, async_trait,
1241 };
1242 use std::sync::Mutex;
1243
1244 struct Recording(Mutex<Vec<Vec<Message>>>);
1246
1247 #[async_trait]
1248 impl ModelProvider for Recording {
1249 async fn complete(
1250 &self,
1251 req: ModelRequest,
1252 _sink: &mut dyn EventSink,
1253 ) -> Result<ModelOutput, ProviderError> {
1254 self.0.lock().unwrap().push(req.messages);
1255 Ok(ModelOutput {
1256 message: Message::text(Role::Assistant, "done"),
1257 stop_reason: StopReason::EndTurn,
1258 usage: Usage::default(),
1259 })
1260 }
1261 }
1262
1263 struct NoopSink;
1264
1265 impl EventSink for NoopSink {
1266 fn emit(&mut self, _event: Event<'_>) {}
1267 }
1268
1269 fn big_history(rounds: usize) -> Session {
1271 let mut history = vec![Message::text(Role::User, "fix the build")];
1272 for i in 0..rounds {
1273 history.push(Message {
1274 role: Role::Assistant,
1275 content: vec![
1276 Content::Text(format!("round {i}")),
1277 Content::ToolCall(ToolCall {
1278 id: format!("c{i}"),
1279 name: "bash".into(),
1280 args: serde_json::json!({ "command": "make" }),
1281 provider_data: serde_json::Value::Null,
1282 }),
1283 ],
1284 });
1285 history.push(Message {
1286 role: Role::Tool,
1287 content: vec![Content::ToolResult(ToolResult::success(
1288 format!("c{i}"),
1289 "x".repeat(2048),
1290 ))],
1291 });
1292 }
1293 Session {
1294 history,
1295 ..Session::default()
1296 }
1297 }
1298
1299 #[tokio::test]
1300 async fn off_the_history_rides_verbatim() {
1301 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1302 let config = CodingAgentConfig {
1303 transcript: None,
1304 ..CodingAgentConfig::new("m", std::env::temp_dir())
1305 };
1306 let mut agent =
1307 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1308 agent
1309 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1310 .await
1311 .unwrap();
1312 let sent = recording.0.lock().unwrap();
1313 assert_eq!(sent.len(), 1);
1314 assert_eq!(sent[0].len(), 5 * 2 + 2);
1316 assert_eq!(sent[0][0], Message::text(Role::User, "fix the build"));
1317 }
1318
1319 #[tokio::test]
1320 async fn on_by_default_a_big_history_is_viewed_and_the_session_stays_whole() {
1321 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1322 let config = CodingAgentConfig {
1323 transcript: Some(rhema_transcript::TranscriptSettings {
1324 max_view_bytes: 4_000,
1325 ..rhema_transcript::TranscriptSettings::default()
1326 }),
1327 ..CodingAgentConfig::new("m", std::env::temp_dir())
1328 };
1329 let mut agent =
1330 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1331 let prompt = agent
1333 .system()
1334 .expect("the agent was built with a system prompt");
1335 assert!(prompt.contains("zoom"), "{prompt}");
1336 agent
1337 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1338 .await
1339 .unwrap();
1340
1341 let sent = recording.0.lock().unwrap();
1342 let messages = &sent[0];
1343 assert!(
1344 messages.len() < 5 * 2 + 2,
1345 "the view is shorter: {} messages",
1346 messages.len()
1347 );
1348 let rhema_core::Content::Text(text) = &messages[0].content[0] else {
1349 panic!("a text view message");
1350 };
1351 assert!(text.starts_with("<transcript>"), "{text}");
1352 assert_eq!(messages.last().unwrap(), &Message::text(Role::User, "next"));
1353
1354 assert_eq!(agent.session().history.len(), 5 * 2 + 3);
1357 }
1358}