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, CapturedText, HandoffCaptureSink, McpServerInfo, McpServerStatus,
19 TuiMessage, WorkflowState,
20};
21
22pub struct Session {
27 pub input: String,
29 pub output_lines: Vec<String>,
31 pub workflow_tx: mpsc::UnboundedSender<TuiMessage>,
33 pub workflow_state: WorkflowState,
35 pub config_toml: Option<String>,
37 pub secrets: Option<HashMap<String, String>>,
39 pub build_info: Option<BuildInfo>,
41 pub resume_info: Option<ResumeInfo>,
43 pub backup_resume_info: Option<ResumeInfo>,
46 pub todo_lines: Vec<String>,
48 pub cancel_token: CancellationToken,
50 pub pending_command: Option<String>,
52 pub handoff_pending: bool,
54 pub pending_tool_lines: Vec<(String, usize, Option<String>)>,
56 pub current_context_tokens: usize,
58 pub auto_handoff: AutoHandoffConfig,
60 pub mcp_servers: Vec<McpServerInfo>,
62 pub should_quit: bool,
64 pub auto_scroll: bool,
66
67 pub agent_name: String,
71 pub token_store: Option<Arc<dyn pep::token_store::TokenStore>>,
74
75 pub project_id: Option<String>,
78 pub project_name: Option<String>,
80 pub session_id: Option<String>,
82 pub session_name: Option<String>,
84 pub home_dir: Option<std::path::PathBuf>,
87
88 pub workflow_permit: Option<tokio::sync::OwnedSemaphorePermit>,
93}
94
95impl Session {
96 pub fn new() -> (Self, mpsc::UnboundedReceiver<TuiMessage>) {
101 let (workflow_tx, workflow_rx) = mpsc::unbounded_channel();
102 let session = Self {
103 input: String::new(),
104 output_lines: Vec::new(),
105 workflow_tx,
106 workflow_state: WorkflowState::Idle,
107 config_toml: None,
108 secrets: None,
109 build_info: None,
110 resume_info: None,
111 backup_resume_info: None,
112 todo_lines: Vec::new(),
113 cancel_token: CancellationToken::new(),
114 pending_command: None,
115 handoff_pending: false,
116 pending_tool_lines: Vec::new(),
117 current_context_tokens: 0,
118 auto_handoff: AutoHandoffConfig::default(),
119 mcp_servers: Vec::new(),
120 should_quit: false,
121 auto_scroll: true,
122 agent_name: "trustee".to_string(),
123 token_store: None,
124 project_id: None,
125 project_name: None,
126 session_id: None,
127 session_name: None,
128 home_dir: None,
129 workflow_permit: None,
130 };
131 (session, workflow_rx)
132 }
133
134 pub fn parse_auto_handoff_config(&mut self) {
136 if let Some(ref config_toml) = self.config_toml {
137 self.auto_handoff = crate::config::parse_auto_handoff_config(config_toml);
138 }
139 }
140
141 pub fn handle_workflow_message(&mut self, msg: TuiMessage) {
146 match msg {
147 TuiMessage::WorkflowCancelled => {
148 self.output_lines.push("⏹ Workflow cancelled".to_string());
149 self.output_lines.push("".to_string());
150 self.workflow_state = WorkflowState::Cancelling;
151 }
152 TuiMessage::OutputLine(line) => {
153 self.output_lines.push(line);
154 }
155 TuiMessage::StreamDelta(delta) => {
156 if let Some(last) = self.output_lines.last_mut() {
157 last.push_str(&delta);
158 } else {
159 self.output_lines.push(delta);
160 }
161 }
162 TuiMessage::ReasoningDelta(delta) => {
163 if let Some(last) = self.output_lines.last_mut() {
164 if !last.starts_with('\x01') {
165 last.insert(0, '\x01');
166 }
167 last.push_str(&delta);
168 } else {
169 self.output_lines.push(format!("\x01{}", delta));
170 }
171 }
172 TuiMessage::WorkflowCompleted => {
173 self.output_lines.push("✓ Workflow completed".to_string());
174 self.output_lines.push("".to_string());
175 if self.workflow_state == WorkflowState::Running {
176 self.workflow_state = WorkflowState::Cancelling;
177 }
178 }
179 TuiMessage::WorkflowError(err) => {
180 self.output_lines.push(format!("✗ Error: {}", err));
181 self.output_lines.push("".to_string());
182 if self.workflow_state == WorkflowState::Running {
183 self.workflow_state = WorkflowState::Cancelling;
184 }
185 }
186 TuiMessage::TodoUpdate(content) => {
187 self.todo_lines = content.lines().map(|l| l.to_string()).collect();
188 }
189 TuiMessage::ToolPending { tool_name, hint } => {
190 let label = match &hint {
191 Some(h) => format!("⠋ {} {}", tool_name, h),
192 None => format!("⠋ {}", tool_name),
193 };
194 let idx = self.output_lines.len();
195 self.output_lines.push(label);
196 self.pending_tool_lines.push((tool_name, idx, hint));
197 }
198 TuiMessage::ToolDone { tool_name, success, hint } => {
199 let status = if success { "✓" } else { "✗" };
200 if let Some(pos) = self.pending_tool_lines.iter().position(|(n, _, _)| *n == tool_name) {
201 let (_, idx, pending_hint) = self.pending_tool_lines.remove(pos);
202 let h = hint.or(pending_hint);
203 let label = match &h {
204 Some(h) => format!("{} {} {}", status, tool_name, h),
205 None => format!("{} {}", status, tool_name),
206 };
207 if idx < self.output_lines.len() {
208 self.output_lines[idx] = label;
209 return;
210 }
211 self.output_lines.push(label);
212 } else {
213 let label = match &hint {
214 Some(h) => format!("{} {} {}", status, tool_name, h),
215 None => format!("{} {}", status, tool_name),
216 };
217 self.output_lines.push(label);
218 }
219 }
220 TuiMessage::ResumeInfo(info) => {
221 if self.workflow_state == WorkflowState::Cancelling && info.is_none() {
222 self.resume_info = self.backup_resume_info.take();
223 } else if info.is_some() {
224 self.resume_info = info;
230 self.backup_resume_info = None;
231 }
232 if let Some(ref ri) = self.resume_info {
238 if self.session_id.is_none() {
239 self.session_id = Some(ri.session_id.clone());
240 }
241 }
242
243 if self.workflow_state == WorkflowState::Cancelling {
244 self.workflow_state = WorkflowState::Idle;
245 self.workflow_permit = None;
247 }
248 if self.resume_info.is_some() {
249 if std::env::var("RUST_LOG")
250 .map(|v| v.to_lowercase().contains("debug"))
251 .unwrap_or(false)
252 {
253 self.output_lines.push("🔄 Session preserved — next command will continue this session".to_string());
254 }
255 }
256 if self.workflow_state == WorkflowState::Idle && self.handoff_pending {
257 self.handoff_pending = false;
258 self.trigger_handoff(String::new());
259 } else if let Some(cmd) = self.pending_command.take() {
260 self.input = cmd;
261 self.execute_command();
262 }
263 }
264 TuiMessage::ContextTokensUpdated(count) => {
265 self.current_context_tokens = count;
266 if self.auto_handoff.enabled
267 && count >= self.auto_handoff.context_threshold
268 && self.workflow_state == WorkflowState::Running
269 && !self.handoff_pending
270 && self.resume_info.is_some()
271 {
272 self.handoff_pending = true;
273 self.cancel_token.cancel();
274 self.workflow_state = WorkflowState::Cancelling;
275 self.output_lines.push(format!(
276 "🔄 Auto-handoff: cancelling workflow, context tokens ({}) ≥ threshold ({})",
277 count, self.auto_handoff.context_threshold
278 ));
279 }
280 }
281 TuiMessage::McpServerStatus { name, connected, tool_count, error } => {
282 let status = if connected { McpServerStatus::Connected } else { McpServerStatus::Failed };
283 if let Some(existing) = self.mcp_servers.iter_mut().find(|s| s.name == name) {
284 existing.status = status;
285 existing.tool_count = tool_count;
286 existing.error = error;
287 } else {
288 self.mcp_servers.push(McpServerInfo { name, status, tool_count, error });
289 }
290 }
291 TuiMessage::HandoffReady(briefing) => {
292 self.workflow_state = WorkflowState::Idle;
293 self.resume_info = None;
294 self.input = briefing;
295 self.execute_command();
296 }
297 }
298 if self.auto_scroll {
299 }
302 }
303
304 pub fn execute_command(&mut self) {
309 let command = self.input.trim().to_string();
310
311 if self.workflow_state != WorkflowState::Idle {
312 self.pending_command = Some(command);
313 self.output_lines.push("⏳ Previous workflow finishing — command queued".to_string());
314 self.input.clear();
315 return;
316 }
317
318 let is_continuation = self.resume_info.is_some();
319
320 if !is_continuation {
321 self.output_lines.clear();
322 if self.session_id.is_none() {
328 let timestamp = chrono::Utc::now().format("%Y_%m_%d_%H_%M");
329 let uuid_suffix = uuid::Uuid::new_v4().simple().to_string();
330 let uuid8 = &uuid_suffix[..8];
331 self.session_id = Some(format!("session_{}_{}", timestamp, uuid8));
332 }
333 if self.session_name.is_none() {
334 let derived = if command.len() > 80 {
335 format!("{}...", &command[..77])
336 } else {
337 command.clone()
338 };
339 self.session_name = Some(derived);
340 }
341 }
342
343 self.output_lines.push(format!("> {}", command));
344
345 let config_toml = match &self.config_toml {
346 Some(c) => c.clone(),
347 None => {
348 self.output_lines.push("✗ Error: Configuration not loaded".to_string());
349 self.output_lines.push("".to_string());
350 return;
351 }
352 };
353
354 let secrets = self.secrets.clone().unwrap_or_default();
355 let build_info = self.build_info.clone();
356 let tx = self.workflow_tx.clone();
357
358 let agent_name = self.agent_name.clone();
359 let token_store = self.token_store.clone();
360 let project_id = self.project_id.clone();
361 let project_name = self.project_name.clone();
362 let session_id = self.session_id.clone();
363 let session_name = self.session_name.clone();
364 let home_dir = self.home_dir.clone();
365
366 self.backup_resume_info = self.resume_info.clone();
367 let resume_info = self.resume_info.take();
368
369 self.workflow_state = WorkflowState::Running;
370 self.auto_scroll = true;
371
372 self.cancel_token = CancellationToken::new();
373 let child_token = self.cancel_token.clone();
374
375 let (resume_tx, mut resume_rx) = mpsc::unbounded_channel();
376
377 let resume_forward_tx = tx.clone();
378 tokio::spawn(async move {
379 while let Some(info) = resume_rx.recv().await {
380 resume_forward_tx.send(TuiMessage::ResumeInfo(info)).ok();
381 }
382 });
383
384 tokio::spawn(async move {
385 let tui_sink: abk::orchestration::output::SharedSink =
386 Arc::new(crate::session::TuiForwardSink::new(tx.clone()));
387
388 let mut run_ctx = RunContext::new()
390 .with_agent_name(agent_name.clone());
391
392 if let Some(ref dir) = home_dir {
394 run_ctx = run_ctx.with_home_dir(dir.clone());
395 }
396
397 if project_id.is_some() || project_name.is_some() {
399 run_ctx = run_ctx.with_project(abk::context::ProjectIdentity {
400 id: project_id.unwrap_or_else(|| "default".to_string()),
401 name: project_name,
402 });
403 }
404
405 if session_id.is_some() || session_name.is_some() {
407 run_ctx = run_ctx.with_session(abk::context::SessionIdentity {
408 id: session_id.unwrap_or_else(|| "default".to_string()),
409 name: session_name,
410 });
411 }
412
413 #[cfg(feature = "registry-mcp-token")]
414 {
415 if let Some(ref ts) = token_store {
416 run_ctx = run_ctx.with_token_store(ts.clone());
417 }
418 }
419
420 let result = abk::observability::with_tui_mode(true, async {
424 abk::cli::run_task_from_raw_config(
425 &config_toml,
426 secrets,
427 build_info,
428 &command,
429 Some(tui_sink),
430 resume_info,
431 Some(resume_tx),
432 Some(child_token),
433 Some(&run_ctx),
434 )
435 .await
436 })
437 .await;
438
439 let task_result = result.unwrap_or_else(|e| abk::cli::TaskResult {
440 success: false,
441 error: Some(e.to_string()),
442 resume_info: None,
447 });
448
449 let msg = if task_result.success {
450 TuiMessage::WorkflowCompleted
451 } else {
452 TuiMessage::WorkflowError(task_result.error.unwrap_or_default())
453 };
454 tx.send(msg).ok();
455 tx.send(TuiMessage::ResumeInfo(task_result.resume_info)).ok();
456 });
457
458 self.input.clear();
459 }
460
461 pub fn trigger_handoff(&mut self, hint: String) {
466 if self.resume_info.is_none() {
467 self.output_lines.push("ℹ Nothing to hand off — run a task first".to_string());
468 return;
469 }
470
471 let config_toml = match &self.config_toml {
472 Some(c) => c.clone(),
473 None => {
474 self.output_lines.push("✗ Error: Configuration not loaded".to_string());
475 return;
476 }
477 };
478
479 let secrets = self.secrets.clone().unwrap_or_default();
480 let build_info = self.build_info.clone();
481 let tx = self.workflow_tx.clone();
482
483 let agent_name = self.agent_name.clone();
484 let token_store = self.token_store.clone();
485
486 let resume_info = self.resume_info.take();
487
488 self.workflow_state = WorkflowState::Running;
489 self.auto_scroll = true;
490 self.cancel_token = CancellationToken::new();
491 let child_token = self.cancel_token.clone();
492
493 self.output_lines.push("🔀 Generating session handoff briefing...".to_string());
494
495 tokio::spawn(async move {
496 let (cap_tx, mut cap_rx) = mpsc::unbounded_channel::<CapturedText>();
497 let cap_sink: abk::orchestration::output::SharedSink =
498 Arc::new(HandoffCaptureSink::new(cap_tx, child_token.clone()));
499
500 let base = "Output a session handoff briefing in at most 300 lines. \
501 Do NOT use any tools. Include: the FULL ABSOLUTE PATH of every \
502 project/repository being worked on (e.g. /Projects/Foo/bar — never \
503 omit the leading path), all project/task/workstream UUIDs referenced, \
504 every file created or modified with its full absolute path, all \
505 commands run and their outcomes, the current state of the work, any \
506 blockers, and the exact next action to take. \
507 Output ONLY the briefing text — no preamble, headers, or closing remarks.";
508 let prompt = if hint.is_empty() {
509 base.to_string()
510 } else {
511 format!("{base}\n\nIn the briefing also consider: {hint}")
512 };
513
514 let (dummy_tx, _dummy_rx) = mpsc::unbounded_channel();
515
516 let run_ctx = RunContext::new()
518 .with_agent_name(agent_name.clone());
519 #[cfg(feature = "registry-mcp-token")]
520 {
521 }
524
525 let _res = abk::observability::with_tui_mode(true, async {
527 abk::cli::run_task_from_raw_config(
528 &config_toml,
529 secrets,
530 build_info,
531 &prompt,
532 Some(cap_sink),
533 resume_info,
534 Some(dummy_tx),
535 Some(child_token),
536 Some(&run_ctx),
537 )
538 .await
539 })
540 .await;
541
542 let mut text_parts = String::new();
543 let mut reasoning_parts = String::new();
544 while let Ok(captured) = cap_rx.try_recv() {
545 match captured {
546 CapturedText::Text(s) => text_parts.push_str(&s),
547 CapturedText::Reasoning(s) => reasoning_parts.push_str(&s),
548 }
549 }
550
551 let briefing = if !text_parts.trim().is_empty() {
552 text_parts.trim().to_string()
553 } else if !reasoning_parts.trim().is_empty() {
554 reasoning_parts.trim().to_string()
555 } else {
556 "Session handoff: briefing unavailable — continue from previous context.".to_string()
557 };
558
559 tx.send(TuiMessage::HandoffReady(briefing)).ok();
560 });
561 }
562}
563
564impl Default for Session {
565 fn default() -> Self {
566 Self::new().0
567 }
568}
569
570pub struct TuiForwardSink {
576 tx: mpsc::UnboundedSender<TuiMessage>,
577 stream_state: AtomicU8,
578}
579
580const STREAM_IDLE: u8 = 0;
582const STREAM_REASONING: u8 = 1;
583const STREAM_CONTENT: u8 = 2;
584
585impl TuiForwardSink {
586 pub fn new(tx: mpsc::UnboundedSender<TuiMessage>) -> Self {
587 Self {
588 tx,
589 stream_state: AtomicU8::new(STREAM_IDLE),
590 }
591 }
592}
593
594impl abk::orchestration::output::OutputSink for TuiForwardSink {
595 fn emit(&self, event: abk::orchestration::output::OutputEvent) {
596 use abk::orchestration::output::OutputEvent;
597
598 let msg = match event {
599 OutputEvent::StreamingChunk { delta } => {
600 if delta.is_empty() {
601 return;
602 }
603 let prev = self.stream_state.swap(STREAM_CONTENT, Ordering::Relaxed);
604 if prev != STREAM_CONTENT {
605 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
606 }
607 let _ = self.tx.send(TuiMessage::StreamDelta(delta));
608 return;
609 }
610
611 OutputEvent::LlmResponse { text, model } => {
612 TuiMessage::OutputLine(format!("[{}] {}", model, text))
613 }
614
615 OutputEvent::Info { message } => {
616 if message.contains("API call completed successfully") {
618 return;
619 }
620 TuiMessage::OutputLine(message)
621 }
622
623 OutputEvent::WorkflowStarted { task_description } => {
624 TuiMessage::OutputLine(format!("🚀 Workflow started: {}", task_description))
625 }
626
627 OutputEvent::WorkflowCompleted { reason, iterations } => {
628 TuiMessage::OutputLine(format!(
629 "✅ Workflow completed after {} iterations: {}",
630 iterations, reason
631 ))
632 }
633
634 OutputEvent::IterationStarted { iteration, context_tokens } => {
635 let _ = self.tx.send(TuiMessage::ContextTokensUpdated(context_tokens));
636 TuiMessage::OutputLine(format!(
637 "📡 Iteration {} | Context = {} tokens",
638 iteration, context_tokens
639 ))
640 }
641
642 OutputEvent::ApiCallStarted {
643 call_number,
644 model,
645 tool_count,
646 streaming,
647 context_tokens,
648 tool_tokens,
649 } => {
650 let mode = if streaming { "Streaming" } else { "Non-streaming" };
651 let total = context_tokens + tool_tokens;
652 let _ = self.tx.send(TuiMessage::ContextTokensUpdated(total));
653 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
655 TuiMessage::OutputLine(format!(
656 "🔥 API Call {} | Ctx={}({}+{}) | {} | Model: {} | Tools: {}",
657 call_number, total, context_tokens, tool_tokens, mode, model, tool_count
658 ))
659 }
660
661 OutputEvent::ToolsExecuting { tool_names, hints } => {
662 for (name, hint) in tool_names.into_iter().zip(hints.into_iter()) {
663 let _ = self.tx.send(TuiMessage::ToolPending { tool_name: name, hint });
664 }
665 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
666 return;
667 }
668
669 OutputEvent::ToolCompleted {
670 tool_name,
671 success,
672 content,
673 description,
674 } => {
675 if tool_name == "todowrite" && success {
676 let _ = self.tx.send(TuiMessage::TodoUpdate(content.clone()));
677 }
678 let hint = description;
679 let _ = self.tx.send(TuiMessage::ToolDone { tool_name, success, hint });
680 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
681 return;
682 }
683
684 OutputEvent::Error { message, context } => {
685 if let Some(ctx) = context {
686 TuiMessage::OutputLine(format!("❌ Error: {} — {}", message, ctx))
687 } else {
688 TuiMessage::OutputLine(format!("❌ Error: {}", message))
689 }
690 }
691
692 OutputEvent::ReasoningChunk { delta } => {
693 if delta.is_empty() {
694 return;
695 }
696 let prev = self.stream_state.swap(STREAM_REASONING, Ordering::Relaxed);
697 if prev != STREAM_REASONING {
698 let _ = self.tx.send(TuiMessage::OutputLine(String::new()));
699 }
700 let _ = self.tx.send(TuiMessage::ReasoningDelta(delta));
701 return;
702 }
703
704 OutputEvent::McpServerStatus { name, connected, tool_count, error } => {
705 let _ = self.tx.send(TuiMessage::McpServerStatus {
706 name,
707 connected,
708 tool_count,
709 error,
710 });
711 return;
712 }
713 };
714
715 self.stream_state.store(STREAM_IDLE, Ordering::Relaxed);
716 let _ = self.tx.send(msg);
717 }
718}