1use std::collections::HashMap;
7use std::sync::Arc;
8
9use tokio::sync::mpsc;
10use tokio_util::sync::CancellationToken;
11
12use abk::cli::ResumeInfo;
13use abk::context::RunContext;
14
15use std::sync::atomic::{AtomicU8, Ordering};
16
17use crate::types::{
18 AutoHandoffConfig, BuildInfo, McpServerInfo, McpServerStatus,
19 TuiMessage, WorkflowState,
20};
21
22pub fn truncate_session_name(command: &str) -> String {
33 if command.chars().count() <= 80 {
34 command.to_string()
35 } else {
36 let end = command
37 .char_indices()
38 .map(|(i, _)| i)
39 .take_while(|&i| i <= 77)
40 .last()
41 .unwrap_or(0);
42 format!("{}...", &command[..end])
43 }
44}
45
46#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
56pub struct TextFileAttachment {
57 pub mime: String,
59 pub filename: String,
61 pub content: String,
63}
64
65fn longest_backtick_run(content: &str) -> usize {
69 let mut max = 0;
70 let mut run = 0;
71 for c in content.chars() {
72 if c == '`' {
73 run += 1;
74 max = max.max(run);
75 } else {
76 run = 0;
77 }
78 }
79 max
80}
81
82fn fence_for(content: &str) -> String {
85 "`".repeat(3.max(longest_backtick_run(content) + 1))
86}
87
88pub fn compose_task_with_text_files(
93 command: &str,
94 files: &[TextFileAttachment],
95) -> String {
96 if files.is_empty() {
97 return command.to_string();
98 }
99 let mut composed = String::with_capacity(command.len() + 64);
100 composed.push_str(command);
101 for f in files {
102 let fence = fence_for(&f.content);
103 let info = match f.filename.rsplit_once('.') {
104 Some((stem, ext)) if !stem.is_empty() && !ext.is_empty() && ext.len() <= 8 => ext,
105 _ => "text",
106 };
107 composed.push_str(&format!(
108 "\n\n[Attached file: {} ({}, {} bytes)]\n{}{}\n{}\n{}",
109 f.filename,
110 f.mime,
111 f.content.len(),
112 fence,
113 info,
114 f.content,
115 fence,
116 ));
117 }
118 composed
119}
120
121pub struct Session {
124 pub input: String,
126 pub input_images: Vec<umf::chatml::ImageAttachment>,
130 pub input_text_files: Vec<TextFileAttachment>,
135 pub output_lines: Vec<String>,
137 pub workflow_tx: mpsc::UnboundedSender<TuiMessage>,
139 pub workflow_state: WorkflowState,
141 pub config_toml: Option<String>,
143 pub secrets: Option<HashMap<String, String>>,
145 pub build_info: Option<BuildInfo>,
147 pub resume_info: Option<ResumeInfo>,
149 pub backup_resume_info: Option<ResumeInfo>,
152 pub todo_lines: Vec<String>,
154 pub cancel_token: CancellationToken,
156 pub pending_command: Option<String>,
158 pub handoff_pending: bool,
160 pub pending_tool_lines: Vec<(String, usize, Option<String>)>,
162 pub current_context_tokens: usize,
164 pub auto_handoff: AutoHandoffConfig,
166 pub mcp_servers: Vec<McpServerInfo>,
168 pub should_quit: bool,
170 pub auto_scroll: bool,
172
173 pub agent_name: String,
177 pub token_store: Option<Arc<dyn pep::token_store::TokenStore>>,
180
181 #[cfg(feature = "registry-mcp-token")]
186 pub mcp_loader: Option<std::sync::Arc<abk::agent::McpToolLoader>>,
187 #[cfg(not(feature = "registry-mcp-token"))]
188 pub mcp_loader: Option<()>,
189
190 pub project_id: Option<String>,
193 pub project_name: Option<String>,
195 pub session_id: Option<String>,
197 pub session_name: Option<String>,
199 pub home_dir: Option<std::path::PathBuf>,
202
203 pub workflow_permit: Option<tokio::sync::OwnedSemaphorePermit>,
208
209 pub identity: Option<String>,
214
215 pub model: Option<String>,
221
222 pub rotating_from_handoff: bool,
228 pub briefing_born: bool,
233 pub handoff_count: u32,
237}
238
239impl Session {
240 pub fn new() -> (Self, mpsc::UnboundedReceiver<TuiMessage>) {
245 let (workflow_tx, workflow_rx) = mpsc::unbounded_channel();
246 let session = Self {
247 input: String::new(),
248 output_lines: Vec::new(),
249 workflow_tx,
250 workflow_state: WorkflowState::Idle,
251 config_toml: None,
252 secrets: None,
253 build_info: None,
254 resume_info: None,
255 backup_resume_info: None,
256 todo_lines: Vec::new(),
257 cancel_token: CancellationToken::new(),
258 input_images: Vec::new(),
259 input_text_files: Vec::new(),
260 pending_command: None,
261 handoff_pending: false,
262 pending_tool_lines: Vec::new(),
263 current_context_tokens: 0,
264 auto_handoff: AutoHandoffConfig::default(),
265 mcp_servers: Vec::new(),
266 should_quit: false,
267 auto_scroll: true,
268 agent_name: "trustee".to_string(),
269 token_store: None,
270 mcp_loader: None,
271 project_id: None,
272 project_name: None,
273 session_id: None,
274 session_name: None,
275 home_dir: None,
276 workflow_permit: None,
277 identity: None,
278 model: None,
279 rotating_from_handoff: false,
280 briefing_born: false,
281 handoff_count: 0,
282 };
283 (session, workflow_rx)
284 }
285
286 pub fn parse_auto_handoff_config(&mut self) {
288 if let Some(ref config_toml) = self.config_toml {
289 self.auto_handoff = crate::config::parse_auto_handoff_config(config_toml);
290 }
291 }
292
293 pub fn handle_workflow_message(&mut self, msg: TuiMessage) {
298 match msg {
299 TuiMessage::WorkflowCancelled => {
300 self.output_lines.push("⏹ Workflow cancelled".to_string());
301 self.output_lines.push("".to_string());
302 self.workflow_state = WorkflowState::Cancelling;
303 }
304 TuiMessage::OutputLine(line) => {
305 self.output_lines.push(line);
306 }
307 TuiMessage::StreamDelta(delta) => {
308 if let Some(last) = self.output_lines.last_mut() {
309 last.push_str(&delta);
310 } else {
311 self.output_lines.push(delta);
312 }
313 }
314 TuiMessage::ReasoningDelta(delta) => {
315 if let Some(last) = self.output_lines.last_mut() {
316 if !last.starts_with('\x01') {
317 last.insert(0, '\x01');
318 }
319 last.push_str(&delta);
320 } else {
321 self.output_lines.push(format!("\x01{}", delta));
322 }
323 }
324 TuiMessage::WorkflowCompleted => {
325 self.output_lines.push("✓ Workflow completed".to_string());
326 self.output_lines.push("".to_string());
327 if self.workflow_state == WorkflowState::Running {
328 self.workflow_state = WorkflowState::Cancelling;
329 }
330 }
331 TuiMessage::WorkflowError(err) => {
332 self.output_lines.push(format!("✗ Error: {}", err));
333 self.output_lines.push("".to_string());
334 if self.workflow_state == WorkflowState::Running {
335 self.workflow_state = WorkflowState::Cancelling;
336 }
337 }
338 TuiMessage::TodoUpdate(content) => {
339 self.todo_lines = content.lines().map(|l| l.to_string()).collect();
340 }
341 TuiMessage::ToolPending { tool_name, hint } => {
342 let label = match &hint {
343 Some(h) => format!("⠋ {} {}", tool_name, h),
344 None => format!("⠋ {}", tool_name),
345 };
346 let idx = self.output_lines.len();
347 self.output_lines.push(label);
348 self.pending_tool_lines.push((tool_name, idx, hint));
349 }
350 TuiMessage::ToolDone { tool_name, success, hint } => {
351 let status = if success { "✓" } else { "✗" };
352 if let Some(pos) = self.pending_tool_lines.iter().position(|(n, _, _)| *n == tool_name) {
353 let (_, idx, pending_hint) = self.pending_tool_lines.remove(pos);
354 let h = hint.or(pending_hint);
355 let label = match &h {
356 Some(h) => format!("{} {} {}", status, tool_name, h),
357 None => format!("{} {}", status, tool_name),
358 };
359 if idx < self.output_lines.len() {
360 self.output_lines[idx] = label;
361 return;
362 }
363 self.output_lines.push(label);
364 } else {
365 let label = match &hint {
366 Some(h) => format!("{} {} {}", status, tool_name, h),
367 None => format!("{} {}", status, tool_name),
368 };
369 self.output_lines.push(label);
370 }
371 }
372 TuiMessage::ResumeInfo(info) => {
373 if self.workflow_state == WorkflowState::Cancelling && info.is_none() {
374 self.resume_info = self.backup_resume_info.take();
375 } else if info.is_some() {
376 self.resume_info = info;
382 self.backup_resume_info = None;
383 }
384 if let Some(ref ri) = self.resume_info {
390 if self.session_id.is_none() {
391 self.session_id = Some(ri.session_id.clone());
392 }
393 }
394
395 if self.workflow_state == WorkflowState::Cancelling {
396 self.workflow_state = WorkflowState::Idle;
397 self.workflow_permit = None;
399 }
400 if self.resume_info.is_some() {
401 if std::env::var("RUST_LOG")
402 .map(|v| v.to_lowercase().contains("debug"))
403 .unwrap_or(false)
404 {
405 self.output_lines.push("🔄 Session preserved — next command will continue this session".to_string());
406 }
407 }
408 if self.workflow_state == WorkflowState::Idle && self.handoff_pending {
409 self.handoff_pending = false;
410 self.trigger_handoff(String::new());
411 } else if let Some(cmd) = self.pending_command.take() {
412 self.input = cmd;
413 self.execute_command();
414 }
415 }
416 TuiMessage::ContextTokensUpdated(count) => {
417 self.current_context_tokens = count;
418 if self.auto_handoff.enabled
419 && count >= self.auto_handoff.context_threshold
420 && self.workflow_state == WorkflowState::Running
421 && !self.handoff_pending
422 && self.resume_info.is_some()
423 && !self.briefing_born
428 {
429 self.handoff_pending = true;
430 self.cancel_token.cancel();
431 self.workflow_state = WorkflowState::Cancelling;
432 self.output_lines.push(format!(
433 "🔄 Auto-handoff: cancelling workflow, context tokens ({}) ≥ threshold ({})",
434 count, self.auto_handoff.context_threshold
435 ));
436 }
437 }
438 TuiMessage::McpServerStatus { name, connected, tool_count, error } => {
439 let status = if connected { McpServerStatus::Connected } else { McpServerStatus::Failed };
440 if let Some(existing) = self.mcp_servers.iter_mut().find(|s| s.name == name) {
441 existing.status = status;
442 existing.tool_count = tool_count;
443 existing.error = error;
444 } else {
445 self.mcp_servers.push(McpServerInfo { name, status, tool_count, error });
446 }
447 }
448 TuiMessage::HandoffReady(briefing) => {
449 self.workflow_state = WorkflowState::Idle;
450 self.resume_info = None;
451 let old_session_id = self.session_id.clone();
454 self.session_id = None;
457 self.session_name = None;
458 self.rotating_from_handoff = true;
463 self.briefing_born = true;
464 self.handoff_count += 1;
465 self.input = briefing;
466 self.execute_command();
467 let _ = self.workflow_tx.send(TuiMessage::SessionRotated {
471 old: old_session_id,
472 new: self.session_id.clone(),
473 });
474 }
475 TuiMessage::SessionRotated { .. } => {
476 }
480 TuiMessage::HandoffFailed => {
481 self.workflow_state = WorkflowState::Idle;
485 self.output_lines.push(
486 "✗ Handoff briefing unavailable — session preserved, try again.".to_string(),
487 );
488 self.output_lines.push("".to_string());
489 }
490 TuiMessage::SessionTitleUpdated(title) => {
491 self.session_name = Some(title);
493 }
494 }
495 if self.auto_scroll {
496 }
499 }
500
501 pub fn execute_command(&mut self) {
506 let command = self.input.trim().to_string();
507
508 if self.workflow_state != WorkflowState::Idle {
509 self.pending_command = Some(command);
510 self.output_lines.push("⏳ Previous workflow finishing — command queued".to_string());
511 self.input.clear();
512 return;
513 }
514
515 let is_continuation = self.resume_info.is_some();
516
517 let rotating_from_handoff = self.rotating_from_handoff;
522 self.rotating_from_handoff = false;
523
524 if !rotating_from_handoff {
528 self.briefing_born = false;
529 }
530
531if !is_continuation && !rotating_from_handoff && self.handoff_count == 0 {
538 self.output_lines.clear();
539 }
540 if !is_continuation {
544 if self.session_id.is_none() {
550 let timestamp = chrono::Utc::now().format("%Y_%m_%d_%H_%M");
551 let uuid_suffix = uuid::Uuid::new_v4().simple().to_string();
552 let uuid8 = &uuid_suffix[..8];
553 self.session_id = Some(format!("session_{}_{}", timestamp, uuid8));
554 }
555 if self.session_name.is_none() {
556 self.session_name = Some(truncate_session_name(&command));
557 }
558 }
559
560 self.output_lines.push(format!("> {}", command));
561
562 let input_images = std::mem::take(&mut self.input_images);
565 if !input_images.is_empty() {
566 self.output_lines
567 .push(format!("📎 {} image attachment(s)", input_images.len()));
568 }
569
570 let input_text_files = std::mem::take(&mut self.input_text_files);
574 if !input_text_files.is_empty() {
575 self.output_lines
576 .push(format!("📎 {} file attachment(s)", input_text_files.len()));
577 }
578 let task_text = compose_task_with_text_files(&command, &input_text_files);
579
580 let config_toml = match &self.config_toml {
581 Some(c) => c.clone(),
582 None => {
583 self.output_lines.push("✗ Error: Configuration not loaded".to_string());
584 self.output_lines.push("".to_string());
585 return;
586 }
587 };
588
589 let config_toml = inject_identity(config_toml, &self.identity);
591
592 let model_override = self.model.take();
594 let config_toml = inject_model(config_toml, model_override);
595
596 let secrets = self.secrets.clone().unwrap_or_default();
597 let title_secrets = secrets.clone();
599 let build_info = self.build_info.clone();
600 let tx = self.workflow_tx.clone();
601
602 let agent_name = self.agent_name.clone();
603 let token_store = self.token_store.clone();
604 let project_id = self.project_id.clone();
605 let project_name = self.project_name.clone();
606 let session_id = self.session_id.clone();
607 let session_name = self.session_name.clone();
608 let home_dir = self.home_dir.clone();
609 let mcp_loader = self.mcp_loader.clone();
610
611 self.backup_resume_info = self.resume_info.clone();
612 let resume_info = self.resume_info.take();
613
614 self.workflow_state = WorkflowState::Running;
615 self.auto_scroll = true;
616
617 self.cancel_token = CancellationToken::new();
618 let child_token = self.cancel_token.clone();
619
620 let (resume_tx, mut resume_rx) = mpsc::unbounded_channel();
621
622 let resume_forward_tx = tx.clone();
623 tokio::spawn(async move {
624 while let Some(info) = resume_rx.recv().await {
625 resume_forward_tx.send(TuiMessage::ResumeInfo(info)).ok();
626 }
627 });
628
629 tokio::spawn(async move {
630 let tui_sink: abk::orchestration::output::SharedSink =
631 Arc::new(crate::session::TuiForwardSink::new(tx.clone()));
632
633 let mut run_ctx = RunContext::new()
635 .with_agent_name(agent_name.clone());
636
637 if let Some(ref dir) = home_dir {
639 run_ctx = run_ctx.with_home_dir(dir.clone());
640 }
641
642 if project_id.is_some() || project_name.is_some() {
644 run_ctx = run_ctx.with_project(abk::context::ProjectIdentity {
645 id: project_id.unwrap_or_else(|| "default".to_string()),
646 name: project_name,
647 });
648 }
649
650 if session_id.is_some() || session_name.is_some() {
652 run_ctx = run_ctx.with_session(abk::context::SessionIdentity {
653 id: session_id.unwrap_or_else(|| "default".to_string()),
654 name: session_name,
655 });
656 }
657
658 #[cfg(feature = "registry-mcp-token")]
659 {
660 if let Some(ref ts) = token_store {
661 run_ctx = run_ctx.with_token_store(ts.clone());
662 }
663 }
664
665 #[cfg(feature = "registry-mcp-token")]
668 if let Some(loader) = &mcp_loader {
669 run_ctx = run_ctx.with_mcp_loader(loader.clone());
670 }
671
672 let scope_logger = {
677 abk::observability::Logger::with_agent_name(
678 None::<&std::path::Path>,
679 Some("INFO"),
680 Some(&agent_name),
681 ).unwrap_or_else(|_| abk::observability::Logger::new(None, Some("INFO")).unwrap())
682 };
683
684 let result = abk::observability::with_logger(scope_logger, async {
685 abk::observability::with_tui_mode(true, async {
686 abk::cli::run_task_from_raw_config(
687 &config_toml,
688 secrets,
689 build_info,
690 &task_text,
691 Vec::new(),
695 input_images,
696 Some(tui_sink),
697 resume_info,
698 Some(resume_tx),
699 Some(child_token),
700 Some(&run_ctx),
701 )
702 .await
703 })
704 .await
705 })
706 .await;
707
708 let task_result = result.unwrap_or_else(|e| abk::cli::TaskResult {
709 success: false,
710 error: Some(e.to_string()),
711 resume_info: None,
716 });
717
718 let msg = if task_result.success {
719 TuiMessage::WorkflowCompleted
720 } else {
721 TuiMessage::WorkflowError(task_result.error.unwrap_or_default())
722 };
723
724 let title_session_id = task_result.resume_info
726 .as_ref()
727 .map(|ri| ri.session_id.clone());
728
729 tx.send(msg).ok();
730 tx.send(TuiMessage::ResumeInfo(task_result.resume_info)).ok();
731
732 if task_result.success {
739 let title_tx = tx.clone();
740 let title_config = config_toml.clone();
741 let title_command = command.clone();
742 let title_ctx = run_ctx.clone();
743
744 tokio::spawn(async move {
745 if let Some(ref sid) = title_session_id {
749 if !abk::cli::should_generate_title(&title_ctx, sid, &title_command).await {
750 return;
751 }
752 } else {
753 return; }
755
756 tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
758
759 if let Some(ref sid) = title_session_id {
761 if !abk::cli::should_generate_title(&title_ctx, sid, &title_command).await {
762 return;
763 }
764 }
765
766 match abk::cli::generate_session_title(
767 &title_config,
768 title_secrets,
769 &title_command,
770 )
771 .await
772 {
773 Ok(Some(title)) => {
774 if let Some(ref sid) = title_session_id {
776 if let Err(e) = abk::cli::persist_session_title(
777 &title_ctx,
778 &title_config,
779 sid,
780 &title,
781 ).await {
782 let _ = e; }
784 }
785 title_tx.send(TuiMessage::SessionTitleUpdated(title)).ok();
787 }
788 Ok(None) => {}
789 Err(e) => { let _ = e; }
790 }
791 });
792 }
793 });
794
795 self.input.clear();
796 }
797
798 pub fn request_handoff(&mut self, hint: String) {
809 if self.briefing_born {
815 self.workflow_tx
816 .send(TuiMessage::OutputLine(
817 "ℹ Handoff unavailable — this session started from a handoff briefing. Run a command first.".to_string(),
818 ))
819 .ok();
820 return;
821 }
822 match self.workflow_state {
823 WorkflowState::Idle => self.trigger_handoff(hint),
824 WorkflowState::Running => {
825 self.cancel_token.cancel();
826 self.workflow_state = WorkflowState::Cancelling;
827 self.handoff_pending = true;
828 self.workflow_tx
833 .send(TuiMessage::OutputLine("⏹ Cancelling before handoff...".to_string()))
834 .ok();
835 }
836 WorkflowState::Cancelling => {
837 self.handoff_pending = true;
838 }
839 }
840 }
841
842 pub fn trigger_handoff(&mut self, hint: String) {
850 if self.workflow_state != WorkflowState::Idle {
854 self.output_lines
855 .push("⏳ Workflow still running — handoff will fire after it stops".to_string());
856 return;
857 }
858 if self.briefing_born {
862 self.output_lines
863 .push("ℹ Handoff unavailable — this session started from a handoff briefing. Run a command first.".to_string());
864 return;
865 }
866 if self.resume_info.is_none() {
867 self.output_lines.push("ℹ Nothing to hand off — run a task first".to_string());
868 return;
869 }
870
871 let config_toml = match &self.config_toml {
872 Some(c) => c.clone(),
873 None => {
874 self.output_lines.push("✗ Error: Configuration not loaded".to_string());
875 return;
876 }
877 };
878
879 let config_toml = inject_identity(config_toml, &self.identity);
881
882 let model_override = self.model.clone();
884 let config_toml = inject_model(config_toml, model_override);
885
886 let tx = self.workflow_tx.clone();
887 let agent_name = self.agent_name.clone();
888 let project_id = self.project_id.clone();
889 let project_name = self.project_name.clone();
890 let session_id = self.session_id.clone();
891 let session_name = self.session_name.clone();
892 let home_dir = self.home_dir.clone();
893
894 let backup_resume_info = self.resume_info.clone();
897 let resume_info = match self.resume_info.take() {
898 Some(ri) => ri,
899 None => {
900 self.output_lines.push("ℹ Nothing to hand off — run a task first".to_string());
902 return;
903 }
904 };
905
906 self.workflow_state = WorkflowState::Running;
907 self.auto_scroll = true;
908 self.cancel_token = CancellationToken::new();
909
910 self.workflow_tx
917 .send(TuiMessage::OutputLine("🔀 Generating session handoff briefing...".to_string()))
918 .ok();
919
920 tokio::spawn(async move {
921 let base = "Output a session handoff briefing in at most 300 lines. \
922 Do NOT use any tools. The FIRST line must be exactly: \
923 \"Session Title: <a concise descriptive title of this work, maximum 50 characters>\" \
924 — nothing else on that line. The rest of the briefing must include: \
925 the FULL ABSOLUTE PATH of every \
926 project/repository being worked on (e.g. /Projects/Foo/bar — never \
927 omit the leading path), all project/task/workstream UUIDs referenced, \
928 every file created or modified with its full absolute path, all \
929 commands run and their outcomes, the current state of the work, any \
930 blockers, and the exact next action to take. \
931 Output ONLY the briefing text — the title line first, then the body, \
932 with no other headers, preamble, or closing remarks.";
933 let prompt = if hint.is_empty() {
934 base.to_string()
935 } else {
936 format!("{base}\n\nIn the briefing also consider: {hint}")
937 };
938
939 let mut run_ctx = RunContext::new()
941 .with_agent_name(agent_name.clone());
942
943 if let Some(ref dir) = home_dir {
945 run_ctx = run_ctx.with_home_dir(dir.clone());
946 }
947
948 if project_id.is_some() || project_name.is_some() {
950 run_ctx = run_ctx.with_project(abk::context::ProjectIdentity {
951 id: project_id.unwrap_or_else(|| "default".to_string()),
952 name: project_name,
953 });
954 }
955
956 if session_id.is_some() || session_name.is_some() {
958 run_ctx = run_ctx.with_session(abk::context::SessionIdentity {
959 id: session_id.unwrap_or_else(|| "default".to_string()),
960 name: session_name,
961 });
962 }
963
964 let scope_logger = abk::observability::Logger::with_agent_name(
966 None,
967 Some("INFO"),
968 Some(&agent_name),
969 ).unwrap_or_else(|_| abk::observability::Logger::new(None, Some("INFO")).unwrap());
970
971 let result = abk::observability::with_logger(scope_logger, async {
975 abk::observability::with_tui_mode(true, async {
976 abk::cli::generate_handoff_briefing(
977 &config_toml,
978 &run_ctx,
979 &resume_info,
980 &prompt,
981 )
982 .await
983 })
984 .await
985 })
986 .await;
987
988 match result {
989 Ok(briefing) if !briefing.trim().is_empty() => {
990 tx.send(TuiMessage::HandoffReady(briefing)).ok();
991 }
992 _ => {
993 tx.send(TuiMessage::ResumeInfo(backup_resume_info)).ok();
997 tx.send(TuiMessage::HandoffFailed).ok();
998 }
999 }
1000 });
1001 }
1002}
1003
1004impl Default for Session {
1005 fn default() -> Self {
1006 Self::new().0
1007 }
1008}
1009
1010fn inject_identity(config_toml: String, identity: &Option<String>) -> String {
1015 let Some(identity) = identity.as_ref().filter(|s| !s.is_empty()) else {
1016 return config_toml;
1017 };
1018
1019 let Ok(mut table) = config_toml.parse::<toml::Value>() else {
1020 return config_toml;
1021 };
1022
1023 let lifecycle = table
1024 .get_mut("lifecycle")
1025 .and_then(|v| v.as_table_mut());
1026
1027 if let Some(lifecycle) = lifecycle {
1028 let existing = lifecycle
1029 .get("system_template")
1030 .and_then(|v| v.as_str())
1031 .unwrap_or("");
1032 let combined = format!("{}\n\n{}", identity, existing);
1033 lifecycle.insert(
1034 "system_template".to_string(),
1035 toml::Value::String(combined),
1036 );
1037 } else {
1038 let mut ltable = toml::value::Table::new();
1040 ltable.insert(
1041 "system_template".to_string(),
1042 toml::Value::String(identity.clone()),
1043 );
1044 if let Some(table) = table.as_table_mut() {
1045 table.insert("lifecycle".to_string(), toml::Value::Table(ltable));
1046 }
1047 }
1048
1049 toml::to_string(&table).unwrap_or(config_toml)
1050}
1051
1052fn inject_model(config_toml: String, model: Option<String>) -> String {
1066 let Some(model_name) = model.as_ref().filter(|s| !s.is_empty()) else {
1067 return config_toml;
1068 };
1069
1070 let Ok(mut table) = config_toml.parse::<toml::Value>() else {
1071 return config_toml;
1072 };
1073
1074 let llm = match table.get_mut("llm").and_then(|v| v.as_table_mut()) {
1075 Some(l) => l,
1076 None => return config_toml,
1077 };
1078
1079 let looks_like_provider = |v: &toml::Value| -> bool {
1081 v.as_table()
1082 .map(|t| {
1083 t.contains_key("base_url")
1084 || t.contains_key("api_key")
1085 || t.contains_key("models")
1086 })
1087 .unwrap_or(false)
1088 };
1089 let offers_model = |v: &toml::Value| -> bool {
1091 v.get("models")
1092 .and_then(|m| m.as_array())
1093 .map(|arr| arr.iter().any(|m| m.as_str() == Some(model_name)))
1094 .unwrap_or(false)
1095 || v.get("model").and_then(|m| m.as_str()) == Some(model_name.as_str())
1096 };
1097
1098 let providers = llm.get("providers").cloned();
1099
1100 if let Some(ref ps) = providers {
1102 if let Some(selected) = ps.get(model_name.as_str()) {
1103 if looks_like_provider(selected) {
1104 let mut selected = selected.clone();
1105 if let Some(st) = selected.as_table_mut() {
1106 st.entry("provider_type".to_string())
1107 .or_insert_with(|| toml::Value::String("openai".to_string()));
1108 }
1109 llm.insert("provider".to_string(), selected);
1110 return toml::to_string(&table).unwrap_or(config_toml);
1111 }
1112 }
1113 }
1114
1115 if let Some(ref ps) = providers {
1117 if let Some(ptable) = ps.as_table() {
1118 for (_key, entry) in ptable {
1119 if looks_like_provider(entry) && offers_model(entry) {
1120 let mut selected = entry.clone();
1121 if let Some(st) = selected.as_table_mut() {
1122 st.insert(
1123 "model".to_string(),
1124 toml::Value::String(model_name.clone()),
1125 );
1126 st.entry("provider_type".to_string())
1127 .or_insert_with(|| toml::Value::String("openai".to_string()));
1128 }
1129 llm.insert("provider".to_string(), selected);
1130 return toml::to_string(&table).unwrap_or(config_toml);
1131 }
1132 }
1133 }
1134 }
1135
1136 if let Some(provider) = llm.get_mut("provider").and_then(|p| p.as_table_mut()) {
1138 provider.insert(
1139 "model".to_string(),
1140 toml::Value::String(model_name.clone()),
1141 );
1142 }
1143
1144 toml::to_string(&table).unwrap_or(config_toml)
1145}
1146pub struct TuiForwardSink {
1151 tx: mpsc::UnboundedSender<TuiMessage>,
1152 stream_state: AtomicU8,
1153}
1154
1155const STREAM_IDLE: u8 = 0;
1157const STREAM_REASONING: u8 = 1;
1158const STREAM_CONTENT: u8 = 2;
1159
1160impl TuiForwardSink {
1161 pub fn new(tx: mpsc::UnboundedSender<TuiMessage>) -> Self {
1162 Self {
1163 tx,
1164 stream_state: AtomicU8::new(STREAM_IDLE),
1165 }
1166 }
1167}
1168
1169impl abk::orchestration::output::OutputSink for TuiForwardSink {
1170 fn emit(&self, event: abk::orchestration::output::OutputEvent) {
1171 use abk::orchestration::output::OutputEvent;
1172
1173 let msg = match event {
1174 OutputEvent::StreamingChunk { delta } => {
1175 if delta.is_empty() {
1176 return;
1177 }
1178 let prev = self.stream_state.swap(STREAM_CONTENT, Ordering::Relaxed);
1179 if prev != STREAM_CONTENT {
1180 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
1181 }
1182 let _ = self.tx.send(TuiMessage::StreamDelta(delta));
1183 return;
1184 }
1185
1186 OutputEvent::LlmResponse { text, model } => {
1187 TuiMessage::OutputLine(format!("[{}] {}", model, text))
1188 }
1189
1190 OutputEvent::Info { message } => {
1191 if message.contains("API call completed successfully") {
1193 return;
1194 }
1195 TuiMessage::OutputLine(message)
1196 }
1197
1198 OutputEvent::WorkflowStarted { task_description } => {
1199 TuiMessage::OutputLine(format!("🚀 Workflow started: {}", task_description))
1200 }
1201
1202 OutputEvent::WorkflowCompleted { reason, iterations } => {
1203 TuiMessage::OutputLine(format!(
1204 "✅ Workflow completed after {} iterations: {}",
1205 iterations, reason
1206 ))
1207 }
1208
1209 OutputEvent::IterationStarted { iteration, context_tokens } => {
1210 let _ = self.tx.send(TuiMessage::ContextTokensUpdated(context_tokens));
1211 TuiMessage::OutputLine(format!(
1212 "📡 Iteration {} | Context = {} tokens",
1213 iteration, context_tokens
1214 ))
1215 }
1216
1217 OutputEvent::ApiCallStarted {
1218 call_number,
1219 model,
1220 tool_count,
1221 streaming,
1222 context_tokens,
1223 tool_tokens,
1224 } => {
1225 let mode = if streaming { "Streaming" } else { "Non-streaming" };
1226 let total = context_tokens + tool_tokens;
1227 let _ = self.tx.send(TuiMessage::ContextTokensUpdated(total));
1228 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
1230 TuiMessage::OutputLine(format!(
1231 "🔥 API Call {} | Ctx={}({}+{}) | {} | Model: {} | Tools: {}",
1232 call_number, total, context_tokens, tool_tokens, mode, model, tool_count
1233 ))
1234 }
1235
1236 OutputEvent::ToolsExecuting { tool_names, hints } => {
1237 for (name, hint) in tool_names.into_iter().zip(hints.into_iter()) {
1238 let _ = self.tx.send(TuiMessage::ToolPending { tool_name: name, hint });
1239 }
1240 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
1241 return;
1242 }
1243
1244 OutputEvent::ToolCompleted {
1245 tool_name,
1246 success,
1247 content,
1248 description,
1249 } => {
1250 if tool_name == "todowrite" && success {
1251 let _ = self.tx.send(TuiMessage::TodoUpdate(content.clone()));
1252 }
1253 let hint = description;
1254 let _ = self.tx.send(TuiMessage::ToolDone { tool_name, success, hint });
1255 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
1256 return;
1257 }
1258
1259 OutputEvent::Error { message, context } => {
1260 if let Some(ctx) = context {
1261 TuiMessage::OutputLine(format!("❌ Error: {} — {}", message, ctx))
1262 } else {
1263 TuiMessage::OutputLine(format!("❌ Error: {}", message))
1264 }
1265 }
1266
1267 OutputEvent::ReasoningChunk { delta } => {
1268 if delta.is_empty() {
1269 return;
1270 }
1271 let prev = self.stream_state.swap(STREAM_REASONING, Ordering::Relaxed);
1272 if prev != STREAM_REASONING {
1273 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
1274 }
1275 let _ = self.tx.send(TuiMessage::ReasoningDelta(delta));
1276 return;
1277 }
1278
1279 OutputEvent::McpServerStatus { name, connected, tool_count, error } => {
1280 let _ = self.tx.send(TuiMessage::McpServerStatus {
1281 name,
1282 connected,
1283 tool_count,
1284 error,
1285 });
1286 return;
1287 }
1288 };
1289
1290 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
1291 let _ = self.tx.send(msg);
1292 }
1293}
1294
1295#[cfg(test)]
1296mod inject_model_tests {
1297 use super::*;
1298
1299 #[test]
1305 fn provider_swap_preserves_danger_accept_invalid_certs() {
1306 let config_toml = r#"
1307[llm.provider]
1308name = "cloud"
1309base_url = "https://api.openai.com/v1"
1310api_key = "${OPENAI_API_KEY}"
1311model = "gpt-4o"
1312
1313[llm.providers.ninfer]
1314name = "ninfer"
1315provider_type = "openai"
1316base_url = "https://10.253.1.9/v1"
1317api_key = "${NINFER_API_KEY}"
1318model = "qwen3.8-27b"
1319models = ["qwen3.8-27b"]
1320danger_accept_invalid_certs = true
1321"#;
1322
1323 let swapped = inject_model(config_toml.to_string(), Some("ninfer".to_string()));
1326 let table = swapped.parse::<toml::Value>().expect("swapped config parses");
1327 let provider = table
1328 .get("llm")
1329 .and_then(|l| l.get("provider"))
1330 .and_then(|p| p.as_table())
1331 .expect("llm.provider present after key swap");
1332 assert_eq!(
1333 provider.get("base_url").and_then(|v| v.as_str()),
1334 Some("https://10.253.1.9/v1"),
1335 "swap selected the ninfer provider"
1336 );
1337 assert_eq!(
1338 provider.get("danger_accept_invalid_certs").and_then(|v| v.as_bool()),
1339 Some(true),
1340 "TLS flag survives the provider swap"
1341 );
1342
1343 let literal = inject_model(config_toml.to_string(), Some("qwen3.8-27b".to_string()));
1346 let table = literal.parse::<toml::Value>().expect("literal config parses");
1347 let providers = table
1348 .get("llm")
1349 .and_then(|l| l.get("providers"))
1350 .and_then(|p| p.get("ninfer"))
1351 .and_then(|p| p.as_table())
1352 .expect("providers table untouched by literal override");
1353 assert_eq!(
1354 providers.get("danger_accept_invalid_certs").and_then(|v| v.as_bool()),
1355 Some(true),
1356 "TLS flag untouched in the providers table on literal override"
1357 );
1358 }
1359}
1360
1361#[cfg(test)]
1362mod compose_task_tests {
1363 use super::*;
1364
1365 fn file(mime: &str, filename: &str, content: &str) -> TextFileAttachment {
1366 TextFileAttachment {
1367 mime: mime.to_string(),
1368 filename: filename.to_string(),
1369 content: content.to_string(),
1370 }
1371 }
1372
1373 #[test]
1374 fn no_files_returns_command_unchanged() {
1375 assert_eq!(compose_task_with_text_files("do the thing", &[]), "do the thing");
1376 }
1377
1378 #[test]
1379 fn single_file_uses_standard_fence() {
1380 let out = compose_task_with_text_files(
1381 "summarize this",
1382 &[file("text/markdown", "notes.md", "# Hello\n\nworld")],
1383 );
1384 assert!(out.starts_with("summarize this"), "got: {}", out);
1385 assert!(out.contains("[Attached file: notes.md (text/markdown, 14 bytes)]"), "got: {}", out);
1386 assert!(out.contains("```md\n# Hello\n\nworld\n```"), "got: {}", out);
1387 }
1388
1389 #[test]
1390 fn content_with_fences_grows_the_fence() {
1391 let content = "```rust\nfn main() {}\n```\ndone";
1392 let out = compose_task_with_text_files("check", &[file("text/plain", "code.txt", content)]);
1393 assert!(out.contains("````txt\n"), "got: {}", out);
1396 assert!(out.ends_with("\n````"), "got: {}", out);
1397 }
1398
1399 #[test]
1400 fn extensionless_filename_falls_back_to_text_info_string() {
1401 let out = compose_task_with_text_files("read", &[file("text/plain", "Makefile", "all:")]);
1402 assert!(out.contains("```text\nall:\n```"), "got: {}", out);
1403 }
1404
1405 #[test]
1406 fn multiple_files_preserve_order() {
1407 let out = compose_task_with_text_files(
1408 "compare",
1409 &[
1410 file("text/plain", "a.txt", "AAA"),
1411 file("text/markdown", "b.md", "BBB"),
1412 ],
1413 );
1414 let a = out.find("[Attached file: a.txt").expect("a header");
1415 let b = out.find("[Attached file: b.md").expect("b header");
1416 assert!(a < b, "files must appear in attachment order");
1417 }
1418
1419 #[test]
1420 fn bytes_in_header_reflect_content_length() {
1421 let content = "héllo"; let out = compose_task_with_text_files("t", &[file("text/plain", "u.txt", content)]);
1423 assert!(out.contains("(text/plain, 6 bytes)"), "got: {}", out);
1424 }
1425}