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 agents = rhema_config::Config::agents_skills_dir().ok();
99 let global = rhema_config::Config::skills_dir().ok();
100 let skills =
101 rhema_agent_tools::skills::discover(agents.as_deref(), global.as_deref(), workspace);
102 stamped_system_with_sections(base, workspace, extra, &skills, None, None, &[], None)
103}
104
105pub fn stamped_system_with_skills(
112 base: Option<&str>,
113 workspace: &Path,
114 extra: &[ToolSpec],
115 skills: &[Skill],
116) -> Result<String, CodingAgentError> {
117 stamped_system_with_sections(base, workspace, extra, skills, None, None, &[], None)
118}
119
120#[expect(
127 clippy::too_many_arguments,
128 reason = "one parameter per prompt section keeps the render's inputs flat"
129)]
130pub fn stamped_system_with_sections(
131 base: Option<&str>,
132 workspace: &Path,
133 extra: &[ToolSpec],
134 skills: &[Skill],
135 environment: Option<&str>,
136 user: Option<&str>,
137 mcp: &[rhema_agent_tools::prompt::McpServerLine],
138 vision: Option<&str>,
139) -> Result<String, CodingAgentError> {
140 let mut tools = ShellTools::builtin_specs();
141 tools.extend_from_slice(extra);
142 rhema_agent_tools::render_system_with_sections(
143 base.unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
144 workspace,
145 skills,
146 &tools,
147 environment,
148 user,
149 mcp,
150 vision,
151 )
152 .map_err(CodingAgentError::from)
153}
154
155pub struct CodingAgent {
158 agent: Agent,
159 config: CodingAgentConfig,
160 session: Session,
161 tools: Arc<ShellTools>,
162 resolver: Arc<dyn ToolResolver>,
167 transcript: Option<Arc<Mutex<rhema_transcript::Transcript>>>,
170 turns: u64,
171}
172
173impl CodingAgent {
174 pub fn new(
180 provider: Arc<dyn rhema_core::ModelProvider>,
181 config: CodingAgentConfig,
182 ) -> Result<Self, CodingAgentError> {
183 let provider: Arc<dyn rhema_core::ModelProvider> = Arc::new(ExpandPastes::new(provider));
187 let tools = config.tools.clone().unwrap_or_else(|| {
188 Arc::new(ShellTools::with_env(
189 config.workspace.clone(),
190 &config.shell_env,
191 ))
192 });
193 let resolver: Arc<dyn ToolResolver> = match config.extra_resolver.clone() {
194 Some(extra) => Arc::new(Overlay::new(tools.clone(), extra)),
195 None => tools.clone(),
196 };
197 let mut extra: Vec<ToolSpec> = config
201 .extra_resolver
202 .as_ref()
203 .map(|resolver| resolver.specs())
204 .unwrap_or_default();
205 let (provider, resolver, transcript) = match config.transcript {
206 Some(settings) => {
207 extra.push(rhema_transcript::Zoom::spec());
208 let (provider, zoom, mirror) = rhema_transcript::attach_with_summarizer(
209 provider,
210 settings,
211 config.summary_worker.clone(),
212 );
213 let resolver = Arc::new(Overlay::new(resolver, Arc::new(zoom)));
214 (
215 provider,
216 Arc::clone(&resolver) as Arc<dyn ToolResolver>,
217 Some(mirror),
218 )
219 }
220 None => (provider, resolver, None),
221 };
222 let mut skills = match &config.skills {
227 Some(skills) => skills.clone(),
228 None => {
229 let agents = rhema_config::Config::agents_skills_dir().ok();
230 let global = rhema_config::Config::skills_dir().ok();
231 rhema_agent_tools::skills::discover(
232 agents.as_deref(),
233 global.as_deref(),
234 tools.workspace(),
235 )
236 }
237 };
238 skills.retain(|skill| !config.disabled_skills.contains(&skill.name));
239 let mcp: Vec<rhema_agent_tools::prompt::McpServerLine> = config
244 .mcp_servers
245 .iter()
246 .filter(|line| !config.disabled_mcp_servers.contains(&line.server))
247 .cloned()
248 .collect();
249 let system = match config.system.clone() {
250 Some(system) => system,
253 None => stamped_system_with_sections(
254 config.base_system.as_deref(),
255 tools.workspace(),
256 &extra,
257 &skills,
258 config.environment.as_deref(),
259 config.user_preferences.as_deref(),
260 &mcp,
261 config.vision.as_deref(),
262 )?,
263 };
264 let skills_bytes = match config.system {
269 Some(_) => 0,
270 None => {
271 let mut bare = ShellTools::builtin_specs();
278 bare.extend_from_slice(&extra);
279 rhema_agent_tools::render_system_with_sections(
280 config
281 .base_system
282 .as_deref()
283 .unwrap_or(crate::prompt::DEFAULT_SYSTEM_MD),
284 tools.workspace(),
285 &[],
286 &bare,
287 config.environment.as_deref(),
288 config.user_preferences.as_deref(),
289 &mcp,
290 config.vision.as_deref(),
291 )
292 .map(|without| system.len().saturating_sub(without.len()))
293 .unwrap_or(0)
294 }
295 };
296 if let Some(mirror) = &transcript {
297 let mut mirror = mirror.lock().expect("transcript lock poisoned");
298 mirror.set_skills_bytes(skills_bytes);
299 mirror.set_skills(
302 skills
303 .iter()
304 .map(|skill| rhema_transcript::SkillText {
305 name: skill.name.clone(),
306 body: skill.body.clone(),
307 })
308 .collect(),
309 );
310 }
311 let agent = Agent::builder()
312 .provider(provider)
313 .resolver(resolver.clone())
316 .system(system)
317 .params(ChatParams {
318 thinking: config.thinking,
319 ..ChatParams::new(config.model.clone())
320 })
321 .build()?;
322 let mut built = Self {
323 agent,
324 config,
325 session: Session::default(),
326 tools,
327 resolver,
328 transcript,
329 turns: 0,
330 };
331 built.stamp_environment();
332 Ok(built)
333 }
334
335 pub fn from_session(
348 provider: Arc<dyn rhema_core::ModelProvider>,
349 mut config: CodingAgentConfig,
350 session: Session,
351 ) -> Result<Self, CodingAgentError> {
352 if let Some(stamped) = session
353 .driver_state
354 .get(ENVIRONMENT_STATE_KEY)
355 .and_then(|value| value.as_str())
356 && config.environment.is_some()
357 {
358 config.environment = Some(stamped.to_owned());
359 }
360 let mut agent = Self::new(provider, config)?;
361 agent.session = session;
362 agent.stamp_environment();
363 Ok(agent)
364 }
365
366 fn stamp_environment(&mut self) {
372 let Some(report) = self.config.environment.clone() else {
373 return;
374 };
375 let state = &mut self.session.driver_state;
376 if !state.is_object() {
377 if !state.is_null() {
378 return;
381 }
382 *state = serde_json::json!({});
383 }
384 if state.get(ENVIRONMENT_STATE_KEY).is_none() {
385 state[ENVIRONMENT_STATE_KEY] = serde_json::json!(report);
386 }
387 }
388
389 pub fn set_commit_hook(&mut self, hook: Option<CommitHook>) {
396 self.agent.set_commit_hook(hook);
397 }
398
399 pub fn session(&self) -> &Session {
401 &self.session
402 }
403
404 pub fn system(&self) -> Option<&str> {
407 self.agent.system()
408 }
409
410 pub fn transcript_stats(&self) -> Option<rhema_transcript::TranscriptStats> {
414 self.transcript
415 .as_ref()
416 .map(|mirror| mirror.lock().expect("transcript lock poisoned").stats())
417 }
418
419 pub fn carry_transcript_stats(&mut self, saved: &rhema_transcript::TranscriptStats) {
424 if let Some(mirror) = &self.transcript {
425 mirror
426 .lock()
427 .expect("transcript lock poisoned")
428 .carry(saved);
429 }
430 }
431
432 pub fn transcript_sections(&self) -> Vec<rhema_transcript::Section> {
437 match &self.transcript {
438 Some(mirror) => mirror
439 .lock()
440 .expect("transcript lock poisoned")
441 .saved_sections(),
442 None => Vec::new(),
443 }
444 }
445
446 pub fn adopt_transcript_sections(&mut self, sections: Vec<rhema_transcript::Section>) {
451 if let Some(mirror) = &self.transcript {
452 mirror
453 .lock()
454 .expect("transcript lock poisoned")
455 .adopt_sections(sections);
456 }
457 }
458
459 pub fn reset_transcript_stats(&mut self) {
464 if let Some(mirror) = &self.transcript {
465 mirror
466 .lock()
467 .expect("transcript lock poisoned")
468 .reset_stats();
469 }
470 }
471
472 pub fn session_mut(&mut self) -> &mut Session {
475 &mut self.session
476 }
477
478 pub async fn turn(
485 &mut self,
486 input: impl Into<String>,
487 sink: &mut dyn EventSink,
488 approval: Arc<dyn Approval>,
489 ) -> Result<ModelOutput, CodingAgentError> {
490 self.turn_with_cancel(input, sink, approval, Arc::new(AtomicBool::new(false)))
491 .await
492 }
493
494 pub async fn turn_with_cancel(
501 &mut self,
502 input: impl Into<String>,
503 sink: &mut dyn EventSink,
504 approval: Arc<dyn Approval>,
505 cancelled: Arc<AtomicBool>,
506 ) -> Result<ModelOutput, CodingAgentError> {
507 self.turn_message_with_cancel(Message::text(Role::User, input), sink, approval, cancelled)
508 .await
509 }
510
511 pub async fn turn_message(
515 &mut self,
516 message: Message,
517 sink: &mut dyn EventSink,
518 approval: Arc<dyn Approval>,
519 ) -> Result<ModelOutput, CodingAgentError> {
520 self.turn_message_with_cancel(message, sink, approval, Arc::new(AtomicBool::new(false)))
521 .await
522 }
523
524 pub async fn turn_message_with_cancel(
527 &mut self,
528 message: Message,
529 sink: &mut dyn EventSink,
530 approval: Arc<dyn Approval>,
531 cancelled: Arc<AtomicBool>,
532 ) -> Result<ModelOutput, CodingAgentError> {
533 self.turns += 1;
534 let ctx = ToolContext {
535 cwd: self.tools.workspace().to_path_buf(),
536 run_id: format!("turn-{}", self.turns),
537 cancelled,
538 };
539 self.tools.install_approval(approval);
543 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
544 let result = match self.config.steering.clone() {
548 Some(steering) => {
549 self.steps(Some(message), sink, &ctx, Some(steering.as_ref()))
550 .await
551 }
552 None => {
553 run_turn(
554 &mut self.agent,
555 &mut self.session,
556 message,
557 sink,
558 &ctx,
559 self.config.max_turn_steps,
560 )
561 .await
562 }
563 };
564 self.tools.clear_approval();
565 Ok(result?)
566 }
567
568 pub async fn retry_turn_with_cancel(
579 &mut self,
580 sink: &mut dyn EventSink,
581 approval: Arc<dyn Approval>,
582 cancelled: Arc<AtomicBool>,
583 ) -> Result<ModelOutput, CodingAgentError> {
584 self.turns += 1;
585 let ctx = ToolContext {
586 cwd: self.tools.workspace().to_path_buf(),
587 run_id: format!("turn-{}", self.turns),
588 cancelled,
589 };
590 self.tools.install_approval(approval);
594 revive_dangling_calls(&mut self.session, self.resolver.as_ref(), &ctx).await;
595 let steering = self.config.steering.clone();
596 let result = self.steps(None, sink, &ctx, steering.as_deref()).await;
597 self.tools.clear_approval();
598 Ok(result?)
599 }
600
601 async fn steps(
613 &mut self,
614 input: Option<Message>,
615 sink: &mut dyn EventSink,
616 ctx: &ToolContext,
617 steering: Option<&dyn Steering>,
618 ) -> Result<ModelOutput, AgentError> {
619 let mut announced: Vec<String> = Vec::new();
622 if let Some(steering) = steering {
623 self.append_notifications(steering, &mut announced);
624 }
625 if let Some(input) = input {
626 self.session.history.push(input);
627 }
628 for _ in 0..self.config.max_turn_steps {
629 let step = self.agent.step(&mut self.session, sink, ctx).await?;
630 if step.output.stop_reason != StopReason::ToolUse {
631 return Ok(step.output);
632 }
633 if let Some(steering) = steering {
634 self.append_notifications(steering, &mut announced);
635 }
636 }
637 Err(AgentError::MaxIterations(self.config.max_turn_steps))
638 }
639
640 fn append_notifications(&mut self, steering: &dyn Steering, announced: &mut Vec<String>) {
646 let lines: Vec<String> = steering
647 .drain()
648 .into_iter()
649 .filter(|line| !announced.contains(line))
650 .collect();
651 if lines.is_empty() {
652 return;
653 }
654 announced.extend(lines.iter().cloned());
655 self.session
656 .history
657 .push(Message::text(Role::User, lines.join("\n")));
658 }
659
660 pub fn save_session_to(&self, path: &impl AsRef<Path>) -> io::Result<()> {
667 let json = serde_json::to_string_pretty(&self.session)
668 .map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))?;
669 rhema_config::privacy::private_write_atomic(path.as_ref(), json.as_bytes())
670 }
671
672 pub fn load_session_from(path: &impl AsRef<Path>) -> io::Result<Session> {
674 let raw = std::fs::read_to_string(path)?;
675 serde_json::from_str(&raw).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
676 }
677}
678
679#[cfg(test)]
680mod paste_tests {
681 use super::*;
682 use rhema_agent_tools::AllowAll;
683 use rhema_core::{
684 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
685 };
686 use std::sync::Mutex;
687
688 struct Recording(Mutex<Vec<String>>);
690
691 #[async_trait]
692 impl ModelProvider for Recording {
693 async fn complete(
694 &self,
695 req: ModelRequest,
696 _sink: &mut dyn EventSink,
697 ) -> Result<ModelOutput, ProviderError> {
698 let mut user = String::new();
699 for message in &req.messages {
700 if message.role != Role::User {
701 continue;
702 }
703 for content in &message.content {
704 if let Content::Text(text) = content {
705 user.push_str(text);
706 }
707 }
708 }
709 self.0.lock().unwrap().push(user);
710 Ok(ModelOutput {
711 message: Message::text(Role::Assistant, "done"),
712 stop_reason: StopReason::EndTurn,
713 usage: Usage::default(),
714 })
715 }
716 }
717
718 struct NoopSink;
719
720 impl EventSink for NoopSink {
721 fn emit(&mut self, _event: Event<'_>) {}
722 }
723
724 #[tokio::test]
725 async fn the_provider_sees_pastes_expanded_and_the_session_keeps_the_chips() {
726 let payload = "l1\nl2\nl3\nl4";
727 let embedded = rhema_core::paste::embed(
728 &format!("look at {} this", rhema_core::paste::MARK),
729 &[payload.to_owned()],
730 );
731 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
732 let mut agent = CodingAgent::new(
733 recording.clone(),
734 CodingAgentConfig::new("m", std::env::temp_dir()),
735 )
736 .unwrap();
737 agent
738 .turn(&embedded, &mut NoopSink, Arc::new(AllowAll))
739 .await
740 .unwrap();
741
742 assert_eq!(
744 agent.session().history[0],
745 Message::text(Role::User, embedded.clone())
746 );
747 let seen = recording.0.lock().unwrap();
749 assert_eq!(seen.len(), 1);
750 assert_eq!(seen[0], format!("look at {payload} this"));
751 }
752
753 struct RecordingBlocks(Mutex<Vec<Vec<Content>>>);
756
757 #[async_trait]
758 impl ModelProvider for RecordingBlocks {
759 async fn complete(
760 &self,
761 req: ModelRequest,
762 _sink: &mut dyn EventSink,
763 ) -> Result<ModelOutput, ProviderError> {
764 let mut blocks = Vec::new();
765 for message in &req.messages {
766 if message.role == Role::User {
767 blocks.extend(message.content.iter().cloned());
768 }
769 }
770 self.0.lock().unwrap().push(blocks);
771 Ok(ModelOutput {
772 message: Message::text(Role::Assistant, "done"),
773 stop_reason: StopReason::EndTurn,
774 usage: Usage::default(),
775 })
776 }
777 }
778
779 fn blocks_agent(recording: &Arc<RecordingBlocks>) -> CodingAgent {
780 CodingAgent::new(
781 recording.clone(),
782 CodingAgentConfig::new("m", std::env::temp_dir()),
783 )
784 .unwrap()
785 }
786
787 #[tokio::test]
788 async fn turn_message_delivers_text_then_image_blocks_to_the_provider() {
789 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
790 let mut agent = blocks_agent(&recording);
791 let message = Message::text_with_images(
792 Role::User,
793 "what is in this picture?",
794 vec![
795 rhema_core::Image::new(rhema_core::ImageMediaType::Png, "aGVsbG8="),
796 rhema_core::Image::new(rhema_core::ImageMediaType::Jpeg, "d29ybGQ="),
797 ],
798 );
799 agent
800 .turn_message(message, &mut NoopSink, Arc::new(AllowAll))
801 .await
802 .unwrap();
803
804 let seen = recording.0.lock().unwrap();
805 assert_eq!(seen.len(), 1);
806 assert_eq!(
807 seen[0],
808 vec![
809 Content::Text("what is in this picture?".into()),
810 Content::Image(rhema_core::Image::new(
811 rhema_core::ImageMediaType::Png,
812 "aGVsbG8="
813 )),
814 Content::Image(rhema_core::Image::new(
815 rhema_core::ImageMediaType::Jpeg,
816 "d29ybGQ="
817 )),
818 ]
819 );
820 assert_eq!(agent.session().history[0].content, seen[0]);
823 }
824
825 #[tokio::test]
826 async fn the_text_entry_still_sends_exactly_one_text_block() {
827 let recording = Arc::new(RecordingBlocks(Mutex::new(Vec::new())));
828 let mut agent = blocks_agent(&recording);
829 agent
830 .turn("plain words", &mut NoopSink, Arc::new(AllowAll))
831 .await
832 .unwrap();
833
834 let seen = recording.0.lock().unwrap();
835 assert_eq!(seen.len(), 1);
836 assert_eq!(seen[0], vec![Content::Text("plain words".into())]);
837 }
838}
839
840#[cfg(test)]
841mod extra_tools_tests {
842 use super::*;
843 use rhema_agent_tools::AllowAll;
844 use rhema_core::{
845 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
846 ToolResult, Usage, async_trait,
847 };
848 use std::sync::Mutex;
849
850 use serde_json::json;
851
852 struct NoopSink;
853
854 impl EventSink for NoopSink {
855 fn emit(&mut self, _event: Event<'_>) {}
856 }
857
858 struct Echo;
860
861 #[async_trait]
862 impl rhema_core::ToolResolver for Echo {
863 fn specs(&self) -> Vec<rhema_core::ToolSpec> {
864 vec![rhema_core::ToolSpec::new(
865 "echo",
866 "Repeats the text argument back. The extension's canary.",
867 json!({
868 "type": "object",
869 "properties": { "text": { "type": "string" } },
870 "required": ["text"]
871 }),
872 )]
873 }
874
875 async fn resolve(&self, call: ToolCall, _ctx: &ToolContext) -> ToolResult {
876 let text = call.args.get("text").and_then(|v| v.as_str()).unwrap_or("");
877 ToolResult::success(call.id, format!("echo: {text}"))
878 }
879 }
880
881 struct Scripted(Mutex<Vec<Vec<String>>>);
884
885 #[async_trait]
886 impl ModelProvider for Scripted {
887 async fn complete(
888 &self,
889 req: ModelRequest,
890 _sink: &mut dyn EventSink,
891 ) -> Result<ModelOutput, ProviderError> {
892 let round = self.0.lock().unwrap().len();
893 self.0
894 .lock()
895 .unwrap()
896 .push(req.tools.iter().map(|t| t.name.clone()).collect());
897 if round == 0 {
898 Ok(ModelOutput {
899 message: Message {
900 role: Role::Assistant,
901 content: vec![Content::ToolCall(ToolCall {
902 id: "c1".into(),
903 name: "echo".into(),
904 args: json!({"text": "through the overlay"}),
905 ..ToolCall::default()
906 })],
907 },
908 stop_reason: StopReason::ToolUse,
909 usage: Usage::default(),
910 })
911 } else {
912 Ok(ModelOutput {
913 message: Message::text(Role::Assistant, "final"),
914 stop_reason: StopReason::EndTurn,
915 usage: Usage::default(),
916 })
917 }
918 }
919 }
920
921 #[tokio::test]
922 async fn the_extra_resolver_reaches_the_request_and_resolves() {
923 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
924 let config = CodingAgentConfig {
925 extra_resolver: Some(Arc::new(Echo)),
926 ..CodingAgentConfig::new("m", std::env::temp_dir())
927 };
928 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
929 let output = agent
930 .turn("use the tool", &mut NoopSink, Arc::new(AllowAll))
931 .await
932 .unwrap();
933
934 let text: String = output
937 .message
938 .content
939 .iter()
940 .filter_map(|c| match c {
941 Content::Text(text) => Some(text.as_str()),
942 _ => None,
943 })
944 .collect();
945 assert_eq!(text, "final");
946 let rounds = scripted.0.lock().unwrap();
947 assert_eq!(rounds.len(), 2);
948 assert!(
949 rounds[0].contains(&"echo".to_owned()) && rounds[0].contains(&"bash".to_owned()),
950 "the tools array carries both: {:?}",
951 rounds[0]
952 );
953 let prompt = agent
956 .system()
957 .expect("the agent was built with a system prompt");
958 assert!(prompt.contains("echo"), "{prompt}");
959 assert!(
960 prompt.contains("canary"),
961 "the extension's description is in the prompt: {prompt}"
962 );
963 }
964
965 #[tokio::test]
966 async fn the_default_advertises_zoom_and_none_skips_it() {
967 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
969 let mut agent = CodingAgent::new(
970 scripted.clone(),
971 CodingAgentConfig::new("m", std::env::temp_dir()),
972 )
973 .unwrap();
974 agent
975 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
976 .await
977 .unwrap();
978 {
979 let rounds = scripted.0.lock().unwrap();
980 assert_eq!(rounds[0], vec!["bash".to_owned(), "zoom".to_owned()]);
981 }
982 let prompt = agent
983 .system()
984 .expect("the agent was built with a system prompt");
985 assert!(prompt.contains("zoom"), "{prompt}");
986 assert!(
987 !agent
988 .system()
989 .is_some_and(|prompt| prompt.contains("canary"))
990 );
991
992 let scripted = Arc::new(Scripted(Mutex::new(Vec::new())));
994 let config = CodingAgentConfig {
995 transcript: None,
996 ..CodingAgentConfig::new("m", std::env::temp_dir())
997 };
998 let mut agent = CodingAgent::new(scripted.clone(), config).unwrap();
999 agent
1000 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1001 .await
1002 .unwrap();
1003 let rounds = scripted.0.lock().unwrap();
1004 assert_eq!(rounds[0], vec!["bash".to_owned()]);
1005 let prompt = agent
1006 .system()
1007 .expect("the agent was built with a system prompt");
1008 assert!(!prompt.contains("zoom"), "{prompt}");
1009 assert!(
1010 !agent
1011 .system()
1012 .is_some_and(|prompt| prompt.contains("canary"))
1013 );
1014 }
1015
1016 #[test]
1017 fn pre_collected_skills_reach_the_prompt_without_a_filesystem() {
1018 let workspace = std::env::temp_dir();
1019 let skill = Skill {
1020 name: "deploy".into(),
1021 description: "Triggers on deploy tasks.".into(),
1022 license: None,
1023 path: workspace.join(".rhema/skills/deploy/SKILL.md"),
1024 body: "the deploy rules".into(),
1025 };
1026 let prompt = stamped_system_with_skills(None, &workspace, &[], &[skill]).unwrap();
1027 assert!(prompt.contains("deploy"), "{prompt}");
1028 assert!(prompt.contains("Triggers on deploy tasks."), "{prompt}");
1029 let without = stamped_system(None, &workspace, &[]).unwrap();
1031 assert!(!without.contains("Triggers on deploy tasks."), "{without}");
1032 }
1033
1034 #[test]
1035 fn stamped_system_carries_the_extra_specs() {
1036 let workspace = std::env::temp_dir();
1037 let spec = rhema_core::ToolSpec::new(
1038 "echo",
1039 "Repeats text back. The extension's canary.",
1040 json!({"type": "object"}),
1041 );
1042 let with = stamped_system(None, &workspace, std::slice::from_ref(&spec)).unwrap();
1043 assert!(with.contains("echo") && with.contains("canary"), "{with}");
1044 let without = stamped_system(None, &workspace, &[]).unwrap();
1045 assert!(!without.contains("canary"), "{without}");
1046 }
1047
1048 #[test]
1049 fn the_optional_sections_reach_the_prompt_where_the_template_places_them() {
1050 let workspace = std::env::temp_dir();
1051 let prompt = stamped_system_with_sections(
1052 None,
1053 &workspace,
1054 &[],
1055 &[],
1056 Some("## Environment\n\n- OS: test os"),
1057 Some("prefer concise answers"),
1058 &[],
1059 Some("pipe bytes into `see`"),
1060 )
1061 .unwrap();
1062 assert!(prompt.contains("- OS: test os"), "{prompt}");
1063 assert!(prompt.contains("prefer concise answers"), "{prompt}");
1064 assert!(prompt.contains("pipe bytes into `see`"), "{prompt}");
1065 let bare = stamped_system_with_skills(None, &workspace, &[], &[]).unwrap();
1068 assert!(!bare.contains("- OS: test os"), "{bare}");
1069 assert!(!bare.contains("prefer concise answers"), "{bare}");
1070 assert!(!bare.contains("pipe bytes into `see`"), "{bare}");
1071 }
1072}
1073
1074#[cfg(test)]
1075mod environment_stamp_tests {
1076 use super::*;
1077 use rhema_agent_tools::AllowAll;
1078 use rhema_core::{
1079 Event, ModelProvider, ModelRequest, ProviderError, StopReason, Usage, async_trait,
1080 };
1081
1082 struct Ender;
1084 #[async_trait]
1085 impl ModelProvider for Ender {
1086 async fn complete(
1087 &self,
1088 _req: ModelRequest,
1089 _sink: &mut dyn EventSink,
1090 ) -> Result<ModelOutput, ProviderError> {
1091 Ok(ModelOutput {
1092 message: Message::text(Role::Assistant, "done"),
1093 stop_reason: StopReason::EndTurn,
1094 usage: Usage::default(),
1095 })
1096 }
1097 }
1098
1099 struct NoopSink;
1100 impl EventSink for NoopSink {
1101 fn emit(&mut self, _event: Event<'_>) {}
1102 }
1103
1104 fn config(report: &str) -> CodingAgentConfig {
1105 CodingAgentConfig {
1106 environment: Some(report.to_owned()),
1107 ..CodingAgentConfig::new("m", std::env::temp_dir())
1108 }
1109 }
1110
1111 #[test]
1114 fn a_fresh_conversation_stamps_at_construction() {
1115 let agent = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1116 assert_eq!(
1117 agent
1118 .session()
1119 .driver_state
1120 .get(ENVIRONMENT_STATE_KEY)
1121 .and_then(|v| v.as_str()),
1122 Some("REPORT-A")
1123 );
1124 }
1125
1126 #[tokio::test]
1131 async fn a_resumed_conversation_keeps_its_stamped_report() {
1132 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1133 first
1134 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1135 .await
1136 .unwrap();
1137 let session = first.session().clone();
1138
1139 let second =
1142 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1143 let prompt = second.system().expect("the agent has a prompt");
1144 assert!(prompt.contains("REPORT-A"), "{prompt}");
1145 assert!(!prompt.contains("REPORT-B"), "{prompt}");
1146 assert_eq!(
1147 second
1148 .session()
1149 .driver_state
1150 .get(ENVIRONMENT_STATE_KEY)
1151 .and_then(|v| v.as_str()),
1152 Some("REPORT-A"),
1153 "the stamp is idempotent across rebuilds"
1154 );
1155 }
1156
1157 #[tokio::test]
1160 async fn a_pre_report_session_stamps_the_fresh_render_once() {
1161 let mut first = CodingAgent::new(Arc::new(Ender), config("REPORT-A")).unwrap();
1162 first
1163 .turn("hi", &mut NoopSink, Arc::new(AllowAll))
1164 .await
1165 .unwrap();
1166 let session = Session {
1168 driver_state: serde_json::Value::Null,
1169 ..first.session().clone()
1170 };
1171
1172 let second =
1173 CodingAgent::from_session(Arc::new(Ender), config("REPORT-B"), session).unwrap();
1174 let prompt = second.system().expect("the agent has a prompt");
1175 assert!(prompt.contains("REPORT-B"), "{prompt}");
1176 assert_eq!(
1177 second
1178 .session()
1179 .driver_state
1180 .get(ENVIRONMENT_STATE_KEY)
1181 .and_then(|v| v.as_str()),
1182 Some("REPORT-B"),
1183 "stamped once on arrival, carried forward"
1184 );
1185 }
1186
1187 #[test]
1191 fn driver_state_shapes_coexist_or_stay_untouched() {
1192 let session = Session {
1193 driver_state: serde_json::json!({ "someone-elses": "shape" }),
1194 ..Session::default()
1195 };
1196 let agent =
1197 CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), session).unwrap();
1198 let state = &agent.session().driver_state;
1199 assert_eq!(
1200 state.get("someone-elses"),
1201 Some(&serde_json::json!("shape")),
1202 "the other driver's key stands"
1203 );
1204 assert_eq!(
1205 state.get(ENVIRONMENT_STATE_KEY).and_then(|v| v.as_str()),
1206 Some("REPORT-A"),
1207 "the stamp coexists beside it"
1208 );
1209
1210 let opaque = Session {
1211 driver_state: serde_json::json!("a driver's opaque shape"),
1212 ..Session::default()
1213 };
1214 let agent = CodingAgent::from_session(Arc::new(Ender), config("REPORT-A"), opaque).unwrap();
1215 assert_eq!(
1216 agent.session().driver_state,
1217 serde_json::json!("a driver's opaque shape"),
1218 "a non-object shape is never clobbered"
1219 );
1220 assert!(
1221 agent
1222 .system()
1223 .is_some_and(|prompt| prompt.contains("REPORT-A")),
1224 "the section still renders for this conversation"
1225 );
1226 }
1227}
1228
1229#[cfg(test)]
1230mod transcript_tests {
1231 use super::*;
1232 use rhema_agent_tools::AllowAll;
1233 use rhema_core::{
1234 Content, Event, ModelProvider, ModelRequest, ProviderError, StopReason, ToolCall,
1235 ToolResult, Usage, async_trait,
1236 };
1237 use std::sync::Mutex;
1238
1239 struct Recording(Mutex<Vec<Vec<Message>>>);
1241
1242 #[async_trait]
1243 impl ModelProvider for Recording {
1244 async fn complete(
1245 &self,
1246 req: ModelRequest,
1247 _sink: &mut dyn EventSink,
1248 ) -> Result<ModelOutput, ProviderError> {
1249 self.0.lock().unwrap().push(req.messages);
1250 Ok(ModelOutput {
1251 message: Message::text(Role::Assistant, "done"),
1252 stop_reason: StopReason::EndTurn,
1253 usage: Usage::default(),
1254 })
1255 }
1256 }
1257
1258 struct NoopSink;
1259
1260 impl EventSink for NoopSink {
1261 fn emit(&mut self, _event: Event<'_>) {}
1262 }
1263
1264 fn big_history(rounds: usize) -> Session {
1266 let mut history = vec![Message::text(Role::User, "fix the build")];
1267 for i in 0..rounds {
1268 history.push(Message {
1269 role: Role::Assistant,
1270 content: vec![
1271 Content::Text(format!("round {i}")),
1272 Content::ToolCall(ToolCall {
1273 id: format!("c{i}"),
1274 name: "bash".into(),
1275 args: serde_json::json!({ "command": "make" }),
1276 provider_data: serde_json::Value::Null,
1277 }),
1278 ],
1279 });
1280 history.push(Message {
1281 role: Role::Tool,
1282 content: vec![Content::ToolResult(ToolResult::success(
1283 format!("c{i}"),
1284 "x".repeat(2048),
1285 ))],
1286 });
1287 }
1288 Session {
1289 history,
1290 ..Session::default()
1291 }
1292 }
1293
1294 #[tokio::test]
1295 async fn off_the_history_rides_verbatim() {
1296 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1297 let config = CodingAgentConfig {
1298 transcript: None,
1299 ..CodingAgentConfig::new("m", std::env::temp_dir())
1300 };
1301 let mut agent =
1302 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1303 agent
1304 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1305 .await
1306 .unwrap();
1307 let sent = recording.0.lock().unwrap();
1308 assert_eq!(sent.len(), 1);
1309 assert_eq!(sent[0].len(), 5 * 2 + 2);
1311 assert_eq!(sent[0][0], Message::text(Role::User, "fix the build"));
1312 }
1313
1314 #[tokio::test]
1315 async fn on_by_default_a_big_history_is_viewed_and_the_session_stays_whole() {
1316 let recording = Arc::new(Recording(Mutex::new(Vec::new())));
1317 let config = CodingAgentConfig {
1318 transcript: Some(rhema_transcript::TranscriptSettings {
1319 max_view_bytes: 4_000,
1320 ..rhema_transcript::TranscriptSettings::default()
1321 }),
1322 ..CodingAgentConfig::new("m", std::env::temp_dir())
1323 };
1324 let mut agent =
1325 CodingAgent::from_session(recording.clone(), config, big_history(5)).unwrap();
1326 let prompt = agent
1328 .system()
1329 .expect("the agent was built with a system prompt");
1330 assert!(prompt.contains("zoom"), "{prompt}");
1331 agent
1332 .turn("next", &mut NoopSink, Arc::new(AllowAll))
1333 .await
1334 .unwrap();
1335
1336 let sent = recording.0.lock().unwrap();
1337 let messages = &sent[0];
1338 assert!(
1339 messages.len() < 5 * 2 + 2,
1340 "the view is shorter: {} messages",
1341 messages.len()
1342 );
1343 let rhema_core::Content::Text(text) = &messages[0].content[0] else {
1344 panic!("a text view message");
1345 };
1346 assert!(text.starts_with("<transcript>"), "{text}");
1347 assert_eq!(messages.last().unwrap(), &Message::text(Role::User, "next"));
1348
1349 assert_eq!(agent.session().history.len(), 5 * 2 + 3);
1352 }
1353}