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, 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, 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 jobs: Option<&str>,
141) -> Result<String, CodingAgentError> {
142 let mut tools = ShellTools::builtin_specs();
143 tools.extend_from_slice(extra);
144 rhema_agent_tools::render_system_with_sections(
145 base.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
146 workspace,
147 skills,
148 &tools,
149 environment,
150 user,
151 mcp,
152 vision,
153 jobs,
154 )
155 .map_err(CodingAgentError::from)
156}
157
158pub struct CodingAgent {
161 agent: Agent,
162 config: CodingAgentConfig,
163 session: Session,
164 tools: Arc<ShellTools>,
165 resolver: Arc<dyn ToolResolver>,
170 transcript: Option<Arc<Mutex<rhema_transcript::Transcript>>>,
173 turns: u64,
174}
175
176impl CodingAgent {
177 pub fn new(
183 provider: Arc<dyn rhema_core::ModelProvider>,
184 config: CodingAgentConfig,
185 ) -> Result<Self, CodingAgentError> {
186 let provider: Arc<dyn rhema_core::ModelProvider> = Arc::new(ExpandPastes::new(provider));
190 let tools = config.tools.clone().unwrap_or_else(|| {
191 Arc::new(ShellTools::with_env(
192 config.workspace.clone(),
193 &config.shell_env,
194 ))
195 });
196 let resolver: Arc<dyn ToolResolver> = match config.extra_resolver.clone() {
197 Some(extra) => Arc::new(Overlay::new(tools.clone(), extra)),
198 None => tools.clone(),
199 };
200 let mut extra: Vec<ToolSpec> = config
204 .extra_resolver
205 .as_ref()
206 .map(|resolver| resolver.specs())
207 .unwrap_or_default();
208 let (provider, resolver, transcript) = match config.transcript {
209 Some(settings) => {
210 extra.push(rhema_transcript::Zoom::spec());
211 let (provider, zoom, mirror) = rhema_transcript::attach_with_summarizer(
212 provider,
213 settings,
214 config.summary_worker.clone(),
215 );
216 let resolver = Arc::new(Overlay::new(resolver, Arc::new(zoom)));
217 (
218 provider,
219 Arc::clone(&resolver) as Arc<dyn ToolResolver>,
220 Some(mirror),
221 )
222 }
223 None => (provider, resolver, None),
224 };
225 let mut skills = match &config.skills {
230 Some(skills) => skills.clone(),
231 None => {
232 let home = config
233 .home
234 .clone()
235 .or_else(|| rhema_config::Home::discover().ok());
236 let agents = home.as_ref().map(|h| h.agents_skills());
237 let global = home.as_ref().map(|h| h.skills());
238 rhema_agent_tools::skills::discover(
239 agents.as_deref(),
240 global.as_deref(),
241 tools.workspace(),
242 )
243 }
244 };
245 skills.retain(|skill| !config.disabled_skills.contains(&skill.name));
246 let mcp: Vec<rhema_agent_tools::prompt::McpServerLine> = config
251 .mcp_servers
252 .iter()
253 .filter(|line| !config.disabled_mcp_servers.contains(&line.server))
254 .cloned()
255 .collect();
256 let system = match config.system.clone() {
257 Some(system) => system,
260 None => stamped_system_with_sections(
261 config.base_system.as_deref(),
262 tools.workspace(),
263 &extra,
264 &skills,
265 config.environment.as_deref(),
266 config.user_preferences.as_deref(),
267 &mcp,
268 config.vision.as_deref(),
269 config.jobs.as_deref(),
270 )?,
271 };
272 let skills_bytes = match config.system {
277 Some(_) => 0,
278 None => {
279 let mut bare = ShellTools::builtin_specs();
286 bare.extend_from_slice(&extra);
287 rhema_agent_tools::render_system_with_sections(
288 config
289 .base_system
290 .as_deref()
291 .unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
292 tools.workspace(),
293 &[],
294 &bare,
295 config.environment.as_deref(),
296 config.user_preferences.as_deref(),
297 &mcp,
298 config.vision.as_deref(),
299 config.jobs.as_deref(),
300 )
301 .map(|without| system.len().saturating_sub(without.len()))
302 .unwrap_or(0)
303 }
304 };
305 if let Some(mirror) = &transcript {
306 let mut mirror = mirror.lock().expect("transcript lock poisoned");
307 mirror.set_skills_bytes(skills_bytes);
308 mirror.set_skills(
311 skills
312 .iter()
313 .map(|skill| rhema_transcript::SkillText {
314 name: skill.name.clone(),
315 body: skill.body.clone(),
316 })
317 .collect(),
318 );
319 }
320 let agent = Agent::builder()
321 .provider(provider)
322 .resolver(resolver.clone())
325 .system(system)
326 .params(ChatParams {
327 thinking: config.thinking,
328 ..ChatParams::new(config.model.clone())
329 })
330 .build()?;
331 let mut built = Self {
332 agent,
333 config,
334 session: Session::default(),
335 tools,
336 resolver,
337 transcript,
338 turns: 0,
339 };
340 built.stamp_environment();
341 Ok(built)
342 }
343
344 pub fn from_session(
357 provider: Arc<dyn rhema_core::ModelProvider>,
358 mut config: CodingAgentConfig,
359 session: Session,
360 ) -> Result<Self, CodingAgentError> {
361 if let Some(stamped) = session
362 .driver_state
363 .get(ENVIRONMENT_STATE_KEY)
364 .and_then(|value| value.as_str())
365 && config.environment.is_some()
366 {
367 config.environment = Some(stamped.to_owned());
368 }
369 let mut agent = Self::new(provider, config)?;
370 agent.session = session;
371 agent.stamp_environment();
372 Ok(agent)
373 }
374
375 fn stamp_environment(&mut self) {
381 let Some(report) = self.config.environment.clone() else {
382 return;
383 };
384 let state = &mut self.session.driver_state;
385 if !state.is_object() {
386 if !state.is_null() {
387 return;
390 }
391 *state = serde_json::json!({});
392 }
393 if state.get(ENVIRONMENT_STATE_KEY).is_none() {
394 state[ENVIRONMENT_STATE_KEY] = serde_json::json!(report);
395 }
396 }
397
398 pub fn set_commit_hook(&mut self, hook: Option<CommitHook>) {
405 self.agent.set_commit_hook(hook);
406 }
407
408 pub fn session(&self) -> &Session {
410 &self.session
411 }
412
413 pub fn system(&self) -> Option<&str> {
416 self.agent.system()
417 }
418
419 pub fn transcript_stats(&self) -> Option<rhema_transcript::TranscriptStats> {
423 self.transcript
424 .as_ref()
425 .map(|mirror| mirror.lock().expect("transcript lock poisoned").stats())
426 }
427
428 pub fn carry_transcript_stats(&mut self, saved: &rhema_transcript::TranscriptStats) {
433 if let Some(mirror) = &self.transcript {
434 mirror
435 .lock()
436 .expect("transcript lock poisoned")
437 .carry(saved);
438 }
439 }
440
441 pub fn transcript_sections(&self) -> Vec<rhema_transcript::Section> {
446 match &self.transcript {
447 Some(mirror) => mirror
448 .lock()
449 .expect("transcript lock poisoned")
450 .saved_sections(),
451 None => Vec::new(),
452 }
453 }
454
455 pub fn adopt_transcript_sections(&mut self, sections: Vec<rhema_transcript::Section>) {
460 if let Some(mirror) = &self.transcript {
461 mirror
462 .lock()
463 .expect("transcript lock poisoned")
464 .adopt_sections(sections);
465 }
466 }
467
468 pub fn reset_transcript_stats(&mut self) {
473 if let Some(mirror) = &self.transcript {
474 mirror
475 .lock()
476 .expect("transcript lock poisoned")
477 .reset_stats();
478 }
479 }
480
481 pub fn session_mut(&mut self) -> &mut Session {
484 &mut self.session
485 }
486
487 pub async fn turn(
494 &mut self,
495 input: impl Into<String>,
496 sink: &mut dyn EventSink,
497 approval: Arc<dyn Approval>,
498 ) -> Result<ModelOutput, CodingAgentError> {
499 self.turn_with_cancel(input, sink, approval, Arc::new(AtomicBool::new(false)))
500 .await
501 }
502
503 pub async fn turn_with_cancel(
510 &mut self,
511 input: impl Into<String>,
512 sink: &mut dyn EventSink,
513 approval: Arc<dyn Approval>,
514 cancelled: Arc<AtomicBool>,
515 ) -> Result<ModelOutput, CodingAgentError> {
516 self.turn_message_with_cancel(Message::text(Role::User, input), sink, approval, cancelled)
517 .await
518 }
519
520 pub async fn turn_message(
524 &mut self,
525 message: Message,
526 sink: &mut dyn EventSink,
527 approval: Arc<dyn Approval>,
528 ) -> Result<ModelOutput, CodingAgentError> {
529 self.turn_message_with_cancel(message, sink, approval, Arc::new(AtomicBool::new(false)))
530 .await
531 }
532
533 pub async fn turn_message_with_cancel(
536 &mut self,
537 message: Message,
538 sink: &mut dyn EventSink,
539 approval: Arc<dyn Approval>,
540 cancelled: Arc<AtomicBool>,
541 ) -> Result<ModelOutput, CodingAgentError> {
542 self.turns += 1;
543 let ctx = ToolContext {
544 cwd: self.tools.workspace().to_path_buf(),
545 run_id: format!("turn-{}", self.turns),
546 cancelled,
547 };
548 self.tools.install_approval(approval);
552 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
553 let result = match self.config.steering.clone() {
557 Some(steering) => {
558 self.steps(Some(message), sink, &ctx, Some(steering.as_ref()))
559 .await
560 }
561 None => {
562 run_turn(
563 &mut self.agent,
564 &mut self.session,
565 message,
566 sink,
567 &ctx,
568 self.config.max_turn_steps,
569 )
570 .await
571 }
572 };
573 self.tools.clear_approval();
574 Ok(result?)
575 }
576
577 pub async fn retry_turn_with_cancel(
588 &mut self,
589 sink: &mut dyn EventSink,
590 approval: Arc<dyn Approval>,
591 cancelled: Arc<AtomicBool>,
592 ) -> Result<ModelOutput, CodingAgentError> {
593 self.turns += 1;
594 let ctx = ToolContext {
595 cwd: self.tools.workspace().to_path_buf(),
596 run_id: format!("turn-{}", self.turns),
597 cancelled,
598 };
599 self.tools.install_approval(approval);
603 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
604 let steering = self.config.steering.clone();
605 let result = self.steps(None, sink, &ctx, steering.as_deref()).await;
606 self.tools.clear_approval();
607 Ok(result?)
608 }
609
610 async fn steps(
622 &mut self,
623 input: Option<Message>,
624 sink: &mut dyn EventSink,
625 ctx: &ToolContext,
626 steering: Option<&dyn Steering>,
627 ) -> Result<ModelOutput, AgentError> {
628 let mut announced: Vec<String> = Vec::new();
631 if let Some(steering) = steering {
632 self.append_notifications(steering, &mut announced);
633 }
634 if let Some(input) = input {
635 self.session.history.push(input);
636 }
637 for _ in 0..self.config.max_turn_steps {
638 let step = self.agent.step(&mut self.session, sink, ctx).await?;
639 if step.output.stop_reason != StopReason::ToolUse {
640 return Ok(step.output);
641 }
642 if let Some(steering) = steering {
643 self.append_notifications(steering, &mut announced);
644 }
645 }
646 Err(AgentError::MaxIterations(self.config.max_turn_steps))
647 }
648
649 fn append_notifications(&mut self, steering: &dyn Steering, announced: &mut Vec<String>) {
655 let lines: Vec<String> = steering
656 .drain()
657 .into_iter()
658 .filter(|line| !announced.contains(line))
659 .collect();
660 if lines.is_empty() {
661 return;
662 }
663 announced.extend(lines.iter().cloned());
664 self.session
665 .history
666 .push(Message::text(Role::User, lines.join("\n")));
667 }
668
669 pub fn save_session_to(&self, path: &impl AsRef<Path>) -> io::Result<()> {
676 let json = serde_json::to_string_pretty(&self.session)
677 .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
678 rhema_config::privacy::private_write_atomic(path.as_ref(), json.as_bytes())
679 }
680
681 pub fn load_session_from(path: &impl AsRef<Path>) -> io::Result<Session> {
683 let raw = std::fs::read_to_string(path)?;
684 serde_json::from_str(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
685 }
686}
687
688#[cfg(test)]
689mod paste_tests {
690 use super::*;
691 use rhema_agent_tools::AllowAll;
692 use rhema_core::{
693 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
694 };
695 use std::sync::Mutex;
696
697 struct Recording(Mutex<Vec<String>>);
699
700 #[async_trait]
701 impl ModelProvider for Recording {
702 async fn complete(
703 &self,
704 req: ModelRequest,
705 _sink: &mut dyn EventSink,
706 ) -> Result<ModelOutput, ProviderError> {
707 let mut user = String::new();
708 for message in &req.messages {
709 if message.role != Role::User {
710 continue;
711 }
712 for content in &message.content {
713 if let Content::Text(text) = content {
714 user.push_str(text);
715 }
716 }
717 }
718 self.0.lock().unwrap().push(user);
719 Ok(ModelOutput {
720 message: Message::text(Role::Assistant, "done"),
721 stop_reason: StopReason::EndTurn,
722 usage: Usage::default(),
723 })
724 }
725 }
726
727 struct NoopSink;
728
729 impl EventSink for NoopSink {
730 fn emit(&mut self, _event: Event<'_>) {}
731 }
732
733 #[tokio::test]
734 async fn the_provider_sees_pastes_expanded_and_the_session_keeps_the_chips() {
735 let payload = "l1\nl2\nl3\nl4";
736 let embedded = rhema_core::paste::embed(
737 &format!("look at {} this", rhema_core::paste::MARK),
738 &[payload.to_owned()],
739 );
740 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
741 let mut agent = CodingAgent::new(
742 recording.clone(),
743 CodingAgentConfig::new("m", std::env::temp_dir()),
744 )
745 .unwrap();
746 agent
747 .turn(&embedded, &mut NoopSink, Arc::new(AllowAll))
748 .await
749 .unwrap();
750
751 assert_eq!(
753 agent.session().history[0],
754 Message::text(Role::User, embedded.clone())
755 );
756 let seen = recording.0.lock().unwrap();
758 assert_eq!(seen.len(), 1);
759 assert_eq!(seen[0], format!("look at {payload} this"));
760 }
761
762 struct RecordingBlocks(Mutex<Vec<Vec<Content>>>);
765
766 #[async_trait]
767 impl ModelProvider for RecordingBlocks {
768 async fn complete(
769 &self,
770 req: ModelRequest,
771 _sink: &mut dyn EventSink,
772 ) -> Result<ModelOutput, ProviderError> {
773 let mut blocks = Vec::new();
774 for message in &req.messages {
775 if message.role == Role::User {
776 blocks.extend(message.content.iter().cloned());
777 }
778 }
779 self.0.lock().unwrap().push(blocks);
780 Ok(ModelOutput {
781 message: Message::text(Role::Assistant, "done"),
782 stop_reason: StopReason::EndTurn,
783 usage: Usage::default(),
784 })
785 }
786 }
787
788 fn blocks_agent(recording: &Arc<RecordingBlocks>) -> CodingAgent {
789 CodingAgent::new(
790 recording.clone(),
791 CodingAgentConfig::new("m", std::env::temp_dir()),
792 )
793 .unwrap()
794 }
795
796 #[tokio::test]
797 async fn turn_message_delivers_text_then_image_blocks_to_the_provider() {
798 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
799 let mut agent = blocks_agent(&recording);
800 let message = Message::text_with_images(
801 Role::User,
802 "what is in this picture?",
803 vec![
804 rhema_core::Image::new(rhema_core::ImageMediaType::Png, "aGVsbG8="),
805 rhema_core::Image::new(rhema_core::ImageMediaType::Jpeg, "d29ybGQ="),
806 ],
807 );
808 agent
809 .turn_message(message, &mut NoopSink, Arc::new(AllowAll))
810 .await
811 .unwrap();
812
813 let seen = recording.0.lock().unwrap();
814 assert_eq!(seen.len(), 1);
815 assert_eq!(
816 seen[0],
817 vec![
818 Content::Text("what is in this picture?".into()),
819 Content::Image(rhema_core::Image::new(
820 rhema_core::ImageMediaType::Png,
821 "aGVsbG8="
822 )),
823 Content::Image(rhema_core::Image::new(
824 rhema_core::ImageMediaType::Jpeg,
825 "d29ybGQ="
826 )),
827 ]
828 );
829 assert_eq!(agent.session().history[0].content, seen[0]);
832 }
833
834 #[tokio::test]
835 async fn the_text_entry_still_sends_exactly_one_text_block() {
836 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
837 let mut agent = blocks_agent(&recording);
838 agent
839 .turn("plain words", &mut NoopSink, Arc::new(AllowAll))
840 .await
841 .unwrap();
842
843 let seen = recording.0.lock().unwrap();
844 assert_eq!(seen.len(), 1);
845 assert_eq!(seen[0], vec![Content::Text("plain words".into())]);
846 }
847}
848
849#[cfg(test)]
850mod extra_tools_tests {
851 use super::*;
852 use rhema_agent_tools::AllowAll;
853 use rhema_core::{
854 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
855 ToolResult, Usage, async_trait,
856 };
857 use std::sync::Mutex;
858
859 use serde_json::json;
860
861 struct NoopSink;
862
863 impl EventSink for NoopSink {
864 fn emit(&mut self, _event: Event<'_>) {}
865 }
866
867 struct Echo;
869
870 #[async_trait]
871 impl rhema_core::ToolResolver for Echo {
872 fn specs(&self) -> Vec<rhema_core::ToolSpec> {
873 vec![rhema_core::ToolSpec::new(
874 "echo",
875 "Repeats the text argument back. The extension's canary.",
876 json!({
877 "type": "object",
878 "properties": { "text": { "type": "string" } },
879 "required": ["text"]
880 }),
881 )]
882 }
883
884 async fn resolve(&self, call: ToolCall, _ctx: &ToolContext) -> ToolResult {
885 let text = call.args.get("text").and_then(|v| v.as_str()).unwrap_or("");
886 ToolResult::success(call.id, format!("echo: {text}"))
887 }
888 }
889
890 struct Scripted(Mutex<Vec<Vec<String>>>);
893
894 #[async_trait]
895 impl ModelProvider for Scripted {
896 async fn complete(
897 &self,
898 req: ModelRequest,
899 _sink: &mut dyn EventSink,
900 ) -> Result<ModelOutput, ProviderError> {
901 let round = self.0.lock().unwrap().len();
902 self.0
903 .lock()
904 .unwrap()
905 .push(req.tools.iter().map(|t| t.name.clone()).collect());
906 if round == 0 {
907 Ok(ModelOutput {
908 message: Message {
909 role: Role::Assistant,
910 content: vec![Content::ToolCall(ToolCall {
911 id: "c1".into(),
912 name: "echo".into(),
913 args: json!({"text": "through the overlay"}),
914 ..ToolCall::default()
915 })],
916 },
917 stop_reason: StopReason::ToolUse,
918 usage: Usage::default(),
919 })
920 } else {
921 Ok(ModelOutput {
922 message: Message::text(Role::Assistant, "final"),
923 stop_reason: StopReason::EndTurn,
924 usage: Usage::default(),
925 })
926 }
927 }
928 }
929
930 #[tokio::test]
931 async fn the_extra_resolver_reaches_the_request_and_resolves() {
932 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
933 let config = CodingAgentConfig {
934 extra_resolver: Some(Arc::new(Echo)),
935 ..CodingAgentConfig::new("m", std::env::temp_dir())
936 };
937 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
938 let output = agent
939 .turn("use the tool", &mut NoopSink, Arc::new(AllowAll))
940 .await
941 .unwrap();
942
943 let text: String = output
946 .message
947 .content
948 .iter()
949 .filter_map(|c| match c {
950 Content::Text(text) => Some(text.as_str()),
951 _ => None,
952 })
953 .collect();
954 assert_eq!(text, "final");
955 let rounds = scripted.0.lock().unwrap();
956 assert_eq!(rounds.len(), 2);
957 assert!(
958 rounds[0].contains(&"echo".to_owned()) && rounds[0].contains(&"bash".to_owned()),
959 "the tools array carries both: {:?}",
960 rounds[0]
961 );
962 let prompt = agent
965 .system()
966 .expect("the agent was built with a system prompt");
967 assert!(prompt.contains("echo"), "{prompt}");
968 assert!(
969 prompt.contains("canary"),
970 "the extension's description is in the prompt: {prompt}"
971 );
972 }
973
974 #[tokio::test]
975 async fn the_default_advertises_zoom_and_none_skips_it() {
976 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
978 let mut agent = CodingAgent::new(
979 scripted.clone(),
980 CodingAgentConfig::new("m", std::env::temp_dir()),
981 )
982 .unwrap();
983 agent
984 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
985 .await
986 .unwrap();
987 {
988 let rounds = scripted.0.lock().unwrap();
989 assert_eq!(rounds[0], vec!["bash".to_owned(), "zoom".to_owned()]);
990 }
991 let prompt = agent
992 .system()
993 .expect("the agent was built with a system prompt");
994 assert!(prompt.contains("zoom"), "{prompt}");
995 assert!(
996 !agent
997 .system()
998 .is_some_and(|prompt| prompt.contains("canary"))
999 );
1000
1001 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
1003 let config = CodingAgentConfig {
1004 transcript: None,
1005 ..CodingAgentConfig::new("m", std::env::temp_dir())
1006 };
1007 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
1008 agent
1009 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1010 .await
1011 .unwrap();
1012 let rounds = scripted.0.lock().unwrap();
1013 assert_eq!(rounds[0], vec!["bash".to_owned()]);
1014 let prompt = agent
1015 .system()
1016 .expect("the agent was built with a system prompt");
1017 assert!(!prompt.contains("zoom"), "{prompt}");
1018 assert!(
1019 !agent
1020 .system()
1021 .is_some_and(|prompt| prompt.contains("canary"))
1022 );
1023 }
1024
1025 #[test]
1026 fn pre_collected_skills_reach_the_prompt_without_a_filesystem() {
1027 let workspace = std::env::temp_dir();
1028 let skill = Skill {
1029 name: "deploy".into(),
1030 description: "Triggers on deploy tasks.".into(),
1031 license: None,
1032 path: workspace.join(".rhema/skills/deploy/SKILL.md"),
1033 body: "the deploy rules".into(),
1034 };
1035 let prompt = stamped_system_with_skills(None, &workspace, &[], &[skill]).unwrap();
1036 assert!(prompt.contains("deploy"), "{prompt}");
1037 assert!(prompt.contains("Triggers on deploy tasks."), "{prompt}");
1038 let without = stamped_system(None, &workspace, &[]).unwrap();
1040 assert!(!without.contains("Triggers on deploy tasks."), "{without}");
1041 }
1042
1043 #[test]
1044 fn stamped_system_carries_the_extra_specs() {
1045 let workspace = std::env::temp_dir();
1046 let spec = rhema_core::ToolSpec::new(
1047 "echo",
1048 "Repeats text back. The extension's canary.",
1049 json!({"type": "object"}),
1050 );
1051 let with = stamped_system(None, &workspace, std::slice::from_ref(&spec)).unwrap();
1052 assert!(with.contains("echo") && with.contains("canary"), "{with}");
1053 let without = stamped_system(None, &workspace, &[]).unwrap();
1054 assert!(!without.contains("canary"), "{without}");
1055 }
1056
1057 #[test]
1058 fn the_optional_sections_reach_the_prompt_where_the_template_places_them() {
1059 let workspace = std::env::temp_dir();
1060 let prompt = stamped_system_with_sections(
1061 None,
1062 &workspace,
1063 &[],
1064 &[],
1065 Some("## Environment\n\n- OS: test os"),
1066 Some("prefer concise answers"),
1067 &[],
1068 Some("pipe bytes into `see`"),
1069 Some("`job run` starts it"),
1070 )
1071 .unwrap();
1072 assert!(prompt.contains("- OS: test os"), "{prompt}");
1073 assert!(prompt.contains("prefer concise answers"), "{prompt}");
1074 assert!(prompt.contains("pipe bytes into `see`"), "{prompt}");
1075 assert!(prompt.contains("`job run` starts it"), "{prompt}");
1076 let bare = stamped_system_with_skills(None, &workspace, &[], &[]).unwrap();
1079 assert!(!bare.contains("- OS: test os"), "{bare}");
1080 assert!(!bare.contains("`job run` starts it"), "{bare}");
1081 assert!(!bare.contains("prefer concise answers"), "{bare}");
1082 assert!(!bare.contains("pipe bytes into `see`"), "{bare}");
1083 }
1084}
1085
1086#[cfg(test)]
1087mod environment_stamp_tests {
1088 use super::*;
1089 use rhema_agent_tools::AllowAll;
1090 use rhema_core::{
1091 Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
1092 };
1093
1094 struct Ender;
1096 #[async_trait]
1097 impl ModelProvider for Ender {
1098 async fn complete(
1099 &self,
1100 _req: ModelRequest,
1101 _sink: &mut dyn EventSink,
1102 ) -> Result<ModelOutput, ProviderError> {
1103 Ok(ModelOutput {
1104 message: Message::text(Role::Assistant, "done"),
1105 stop_reason: StopReason::EndTurn,
1106 usage: Usage::default(),
1107 })
1108 }
1109 }
1110
1111 struct NoopSink;
1112 impl EventSink for NoopSink {
1113 fn emit(&mut self, _event: Event<'_>) {}
1114 }
1115
1116 fn config(report: &str) -> CodingAgentConfig {
1117 CodingAgentConfig {
1118 environment: Some(report.to_owned()),
1119 ..CodingAgentConfig::new("m", std::env::temp_dir())
1120 }
1121 }
1122
1123 #[test]
1126 fn a_fresh_conversation_stamps_at_construction() {
1127 let agent = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1128 assert_eq!(
1129 agent
1130 .session()
1131 .driver_state
1132 .get(ENVIRONMENT_STATE_KEY)
1133 .and_then(|v| v.as_str()),
1134 Some("REPORT-A")
1135 );
1136 }
1137
1138 #[tokio::test]
1143 async fn a_resumed_conversation_keeps_its_stamped_report() {
1144 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1145 first
1146 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1147 .await
1148 .unwrap();
1149 let session = first.session().clone();
1150
1151 let second =
1154 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1155 let prompt = second.system().expect("the agent has a prompt");
1156 assert!(prompt.contains("REPORT-A"), "{prompt}");
1157 assert!(!prompt.contains("REPORT-B"), "{prompt}");
1158 assert_eq!(
1159 second
1160 .session()
1161 .driver_state
1162 .get(ENVIRONMENT_STATE_KEY)
1163 .and_then(|v| v.as_str()),
1164 Some("REPORT-A"),
1165 "the stamp is idempotent across rebuilds"
1166 );
1167 }
1168
1169 #[tokio::test]
1172 async fn a_pre_report_session_stamps_the_fresh_render_once() {
1173 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1174 first
1175 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1176 .await
1177 .unwrap();
1178 let session = Session {
1180 driver_state: serde_json::Value::Null,
1181 ..first.session().clone()
1182 };
1183
1184 let second =
1185 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1186 let prompt = second.system().expect("the agent has a prompt");
1187 assert!(prompt.contains("REPORT-B"), "{prompt}");
1188 assert_eq!(
1189 second
1190 .session()
1191 .driver_state
1192 .get(ENVIRONMENT_STATE_KEY)
1193 .and_then(|v| v.as_str()),
1194 Some("REPORT-B"),
1195 "stamped once on arrival, carried forward"
1196 );
1197 }
1198
1199 #[test]
1203 fn driver_state_shapes_coexist_or_stay_untouched() {
1204 let session = Session {
1205 driver_state: serde_json::json!({ "someone-elses": "shape" }),
1206 ..Session::default()
1207 };
1208 let agent =
1209 CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), session).unwrap();
1210 let state = &agent.session().driver_state;
1211 assert_eq!(
1212 state.get("someone-elses"),
1213 Some(&serde_json::json!("shape")),
1214 "the other driver's key stands"
1215 );
1216 assert_eq!(
1217 state.get(ENVIRONMENT_STATE_KEY).and_then(|v| v.as_str()),
1218 Some("REPORT-A"),
1219 "the stamp coexists beside it"
1220 );
1221
1222 let opaque = Session {
1223 driver_state: serde_json::json!("a driver's opaque shape"),
1224 ..Session::default()
1225 };
1226 let agent = CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), opaque).unwrap();
1227 assert_eq!(
1228 agent.session().driver_state,
1229 serde_json::json!("a driver's opaque shape"),
1230 "a non-object shape is never clobbered"
1231 );
1232 assert!(
1233 agent
1234 .system()
1235 .is_some_and(|prompt| prompt.contains("REPORT-A")),
1236 "the section still renders for this conversation"
1237 );
1238 }
1239}
1240
1241#[cfg(test)]
1242mod transcript_tests {
1243 use super::*;
1244 use rhema_agent_tools::AllowAll;
1245 use rhema_core::{
1246 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
1247 ToolResult, Usage, async_trait,
1248 };
1249 use std::sync::Mutex;
1250
1251 struct Recording(Mutex<Vec<Vec<Message>>>);
1253
1254 #[async_trait]
1255 impl ModelProvider for Recording {
1256 async fn complete(
1257 &self,
1258 req: ModelRequest,
1259 _sink: &mut dyn EventSink,
1260 ) -> Result<ModelOutput, ProviderError> {
1261 self.0.lock().unwrap().push(req.messages);
1262 Ok(ModelOutput {
1263 message: Message::text(Role::Assistant, "done"),
1264 stop_reason: StopReason::EndTurn,
1265 usage: Usage::default(),
1266 })
1267 }
1268 }
1269
1270 struct NoopSink;
1271
1272 impl EventSink for NoopSink {
1273 fn emit(&mut self, _event: Event<'_>) {}
1274 }
1275
1276 fn big_history(rounds: usize) -> Session {
1278 let mut history = vec![Message::text(Role::User, "fix the build")];
1279 for i in 0..rounds {
1280 history.push(Message {
1281 role: Role::Assistant,
1282 content: vec![
1283 Content::Text(format!("round {i}")),
1284 Content::ToolCall(ToolCall {
1285 id: format!("c{i}"),
1286 name: "bash".into(),
1287 args: serde_json::json!({ "command": "make" }),
1288 provider_data: serde_json::Value::Null,
1289 }),
1290 ],
1291 });
1292 history.push(Message {
1293 role: Role::Tool,
1294 content: vec![Content::ToolResult(ToolResult::success(
1295 format!("c{i}"),
1296 "x".repeat(2048),
1297 ))],
1298 });
1299 }
1300 Session {
1301 history,
1302 ..Session::default()
1303 }
1304 }
1305
1306 #[tokio::test]
1307 async fn off_the_history_rides_verbatim() {
1308 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1309 let config = CodingAgentConfig {
1310 transcript: None,
1311 ..CodingAgentConfig::new("m", std::env::temp_dir())
1312 };
1313 let mut agent =
1314 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1315 agent
1316 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1317 .await
1318 .unwrap();
1319 let sent = recording.0.lock().unwrap();
1320 assert_eq!(sent.len(), 1);
1321 assert_eq!(sent[0].len(), 5 * 2 + 2);
1323 assert_eq!(sent[0][0], Message::text(Role::User, "fix the build"));
1324 }
1325
1326 #[tokio::test]
1327 async fn on_by_default_a_big_history_is_viewed_and_the_session_stays_whole() {
1328 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1329 let config = CodingAgentConfig {
1330 transcript: Some(rhema_transcript::TranscriptSettings {
1331 max_view_bytes: 4_000,
1332 ..rhema_transcript::TranscriptSettings::default()
1333 }),
1334 ..CodingAgentConfig::new("m", std::env::temp_dir())
1335 };
1336 let mut agent =
1337 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1338 let prompt = agent
1340 .system()
1341 .expect("the agent was built with a system prompt");
1342 assert!(prompt.contains("zoom"), "{prompt}");
1343 agent
1344 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1345 .await
1346 .unwrap();
1347
1348 let sent = recording.0.lock().unwrap();
1349 let messages = &sent[0];
1350 assert!(
1351 messages.len() < 5 * 2 + 2,
1352 "the view is shorter: {} messages",
1353 messages.len()
1354 );
1355 let rhema_core::Content::Text(text) = &messages[0].content[0] else {
1356 panic!("a text view message");
1357 };
1358 assert!(text.starts_with("<transcript>"), "{text}");
1359 assert_eq!(messages.last().unwrap(), &Message::text(Role::User, "next"));
1360
1361 assert_eq!(agent.session().history.len(), 5 * 2 + 3);
1364 }
1365}