1mod adapter;
2mod adapters;
3mod backend;
4mod terminal;
5
6use std::path::Path;
7use std::sync::Arc;
8use std::sync::LazyLock;
9use std::sync::Mutex as StdMutex;
10use std::time::Duration;
11use std::time::{SystemTime, UNIX_EPOCH};
12
13use ab_glyph::{Font, FontVec, PxScale, ScaleFont, point};
14use anyhow::{Context, Result};
15use beam_core::{
16 BackendType, BeamPaths, CliUsageLimitKind, CliUsageLimitState, DaemonToWorker, DisplayMode,
17 InitConfig, ScreenAnalyzerConfig, ScreenStatus, TermActionKey, TuiPromptOption, WorkerToDaemon,
18};
19use image::{ColorType, ImageBuffer, ImageEncoder, Rgba, codecs::png::PngEncoder};
20use reqwest::multipart::{Form, Part};
21use reqwest::{Client, header::HeaderMap};
22use serde::Deserialize;
23use sha2::{Digest, Sha256};
24use terminal::{TerminalState, serve};
25use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
26use tokio::sync::{Mutex, RwLock, broadcast};
27use tracing::{info, warn};
28use unicode_width::UnicodeWidthChar;
29use uuid::Uuid;
30
31use crate::adapter::CliAdapter;
32use crate::backend::{
33 PtyBackend, SessionBackend, SpawnOpts, TmuxPipeBackend, ZellijBackend, ZellijObserveBackend,
34};
35
36fn render_screen_for_display_mode(screen: &str, mode: DisplayMode) -> String {
37 match mode {
38 DisplayMode::Hidden => "[screen hidden]".to_string(),
39 DisplayMode::Screenshot => screen.to_string(),
40 }
41}
42
43const SCREEN_ANALYZER_SYSTEM_PROMPT: &str = "You are a terminal screen analyzer. Determine whether the CLI is showing a blocking interactive prompt. Return only JSON with fields needsInteraction, description, options, multiSelect, toggleKey, confirmKey, checkAgainWhen. checkAgainWhen must be one of content_changed, after_5s, after_10s, not_needed.";
44
45#[derive(Debug, Clone, Default)]
46struct AnalyzerRuntime {
47 last_snapshot: String,
48 stable_count: u32,
49 last_analyzed_snapshot: String,
50 waiting_for_content_change: bool,
51 cooldown_until_ms: u64,
52 is_analyzing: bool,
53 prompt_active: bool,
54}
55
56#[derive(Debug, Deserialize)]
57struct AnalyzerChatResponse {
58 choices: Vec<AnalyzerChoice>,
59}
60
61#[derive(Debug, Deserialize)]
62struct AnalyzerChoice {
63 message: AnalyzerMessage,
64}
65
66#[derive(Debug, Deserialize)]
67struct AnalyzerMessage {
68 content: String,
69}
70
71#[derive(Debug, Deserialize)]
72struct AnalyzerParsedResponse {
73 #[serde(rename = "needsInteraction")]
74 needs_interaction: Option<bool>,
75 description: Option<String>,
76 options: Option<Vec<AnalyzerParsedOption>>,
77 #[serde(rename = "multiSelect")]
78 multi_select: Option<bool>,
79 #[serde(rename = "toggleKey")]
80 toggle_key: Option<String>,
81 #[serde(rename = "confirmKey")]
82 confirm_key: Option<String>,
83 #[serde(rename = "checkAgainWhen")]
84 check_again_when: Option<String>,
85}
86
87#[derive(Debug, Deserialize)]
88struct AnalyzerParsedOption {
89 label: Option<String>,
90 text: Option<String>,
91 selected: Option<bool>,
92 #[serde(rename = "type")]
93 option_type: Option<String>,
94 index: Option<usize>,
95}
96
97#[derive(Debug, Clone, PartialEq, Eq)]
98struct AnalyzerResult {
99 needs_interaction: bool,
100 description: Option<String>,
101 options: Vec<TuiPromptOption>,
102 multi_select: bool,
103 check_again_when: String,
104}
105
106fn screen_analyzer_enabled(cfg: &ScreenAnalyzerConfig) -> bool {
107 cfg.enabled
108 && !cfg.base_url.trim().is_empty()
109 && !cfg.api_key.trim().is_empty()
110 && !cfg.model.trim().is_empty()
111}
112
113fn now_ms() -> u64 {
114 SystemTime::now()
115 .duration_since(UNIX_EPOCH)
116 .unwrap_or_default()
117 .as_millis() as u64
118}
119
120fn parse_screen_analyzer_response(content: &str) -> AnalyzerResult {
121 let json_str = content
122 .trim()
123 .trim_start_matches("```json")
124 .trim_start_matches("```")
125 .trim_end_matches("```")
126 .trim();
127 let parsed = serde_json::from_str::<AnalyzerParsedResponse>(json_str).ok();
128 let check_again_when = parsed
129 .as_ref()
130 .and_then(|parsed| parsed.check_again_when.clone())
131 .filter(|value| {
132 matches!(
133 value.as_str(),
134 "content_changed" | "after_5s" | "after_10s" | "not_needed"
135 )
136 })
137 .unwrap_or_else(|| "content_changed".to_string());
138 let toggle_key = parsed
139 .as_ref()
140 .and_then(|parsed| parsed.toggle_key.clone())
141 .unwrap_or_else(|| "Space".to_string());
142 let confirm_key = parsed
143 .as_ref()
144 .and_then(|parsed| parsed.confirm_key.clone())
145 .unwrap_or_else(|| "Enter".to_string());
146 let options = parsed
147 .as_ref()
148 .and_then(|parsed| parsed.options.as_ref())
149 .map(|options| build_tui_prompt_options(options, &toggle_key, &confirm_key))
150 .unwrap_or_default();
151 AnalyzerResult {
152 needs_interaction: parsed
153 .as_ref()
154 .and_then(|parsed| parsed.needs_interaction)
155 .unwrap_or(false),
156 description: parsed
157 .as_ref()
158 .and_then(|parsed| parsed.description.clone()),
159 options,
160 multi_select: parsed
161 .as_ref()
162 .and_then(|parsed| parsed.multi_select)
163 .unwrap_or(false),
164 check_again_when,
165 }
166}
167
168fn apply_screen_analyzer_result(
169 runtime: &mut AnalyzerRuntime,
170 check_again_when: &str,
171 now_ms: u64,
172) {
173 match check_again_when {
174 "after_5s" => {
175 runtime.cooldown_until_ms = now_ms + 5_000;
176 runtime.waiting_for_content_change = false;
177 }
178 "after_10s" => {
179 runtime.cooldown_until_ms = now_ms + 10_000;
180 runtime.waiting_for_content_change = false;
181 }
182 "not_needed" | "content_changed" => {
183 runtime.waiting_for_content_change = true;
184 runtime.cooldown_until_ms = 0;
185 }
186 _ => {
187 runtime.waiting_for_content_change = true;
188 runtime.cooldown_until_ms = 0;
189 }
190 }
191}
192
193fn build_tui_prompt_options(
194 options: &[AnalyzerParsedOption],
195 toggle_key: &str,
196 confirm_key: &str,
197) -> Vec<TuiPromptOption> {
198 options
199 .iter()
200 .enumerate()
201 .map(|(i, option)| {
202 let index = option.index.unwrap_or(i);
203 let option_type = option
204 .option_type
205 .clone()
206 .filter(|kind| matches!(kind.as_str(), "select" | "toggle" | "confirm" | "input"))
207 .unwrap_or_else(|| "select".to_string());
208 let mut keys = Vec::new();
209 for _ in 0..index {
210 keys.push("Down".to_string());
211 }
212 match option_type.as_str() {
213 "toggle" => {
214 keys.push(toggle_key.to_string());
215 for _ in 0..index {
216 keys.push("Up".to_string());
217 }
218 }
219 "select" | "confirm" | "input" => {
220 keys.push(confirm_key.to_string());
221 }
222 _ => {}
223 }
224 TuiPromptOption {
225 label: option.label.clone().or_else(|| Some((i + 1).to_string())),
226 text: option
227 .text
228 .clone()
229 .unwrap_or_default()
230 .replace('\n', " ")
231 .trim()
232 .to_string(),
233 selected: option.selected.unwrap_or(false),
234 option_type: Some(option_type),
235 keys,
236 }
237 })
238 .collect()
239}
240
241async fn call_screen_analyzer(
242 client: &Client,
243 cfg: &ScreenAnalyzerConfig,
244 snapshot: &str,
245) -> Result<AnalyzerResult> {
246 let url = format!("{}/chat/completions", cfg.base_url.trim_end_matches('/'));
247 let mut body = serde_json::json!({
248 "model": cfg.model,
249 "messages": [
250 { "role": "system", "content": SCREEN_ANALYZER_SYSTEM_PROMPT },
251 { "role": "user", "content": snapshot },
252 ],
253 "temperature": 0,
254 "max_tokens": 2048,
255 });
256 if let Some(map) = body.as_object_mut() {
257 for (key, value) in &cfg.extra_body {
258 map.insert(key.clone(), value.clone());
259 }
260 }
261
262 let mut headers = HeaderMap::new();
263 headers.insert("content-type", "application/json".parse()?);
264 headers.insert("authorization", format!("Bearer {}", cfg.api_key).parse()?);
265 for (key, value) in &cfg.extra_headers {
266 headers.insert(key.parse::<reqwest::header::HeaderName>()?, value.parse()?);
267 }
268
269 let response = client
270 .post(url)
271 .headers(headers)
272 .json(&body)
273 .timeout(Duration::from_secs(15))
274 .send()
275 .await?;
276 let status = response.status();
277 if !status.is_success() {
278 let text = response.text().await.unwrap_or_default();
279 anyhow::bail!(
280 "screen analyzer API {}: {}",
281 status,
282 text.chars().take(200).collect::<String>()
283 );
284 }
285 let payload = response.json::<AnalyzerChatResponse>().await?;
286 let content = payload
287 .choices
288 .first()
289 .map(|choice| choice.message.content.clone())
290 .unwrap_or_default();
291 Ok(parse_screen_analyzer_response(&content))
292}
293
294fn term_action_keys(key: TermActionKey) -> Vec<String> {
295 match key {
296 TermActionKey::Esc => vec!["Escape".to_string()],
297 TermActionKey::CtrlC => vec!["C-c".to_string()],
298 TermActionKey::Tab => vec!["Tab".to_string()],
299 TermActionKey::Enter => vec!["Enter".to_string()],
300 TermActionKey::Space => vec!["Space".to_string()],
301 TermActionKey::Up => vec!["Up".to_string()],
302 TermActionKey::Down => vec!["Down".to_string()],
303 TermActionKey::Left => vec!["Left".to_string()],
304 TermActionKey::Right => vec!["Right".to_string()],
305 TermActionKey::HalfPageUp => vec!["PageUp".to_string()],
306 TermActionKey::HalfPageDown => vec!["PageDown".to_string()],
307 }
308}
309
310fn shell_quote(input: &str) -> String {
311 if input.is_empty() {
312 return "''".to_string();
313 }
314 if !input.bytes().any(|b| {
315 matches!(
316 b,
317 b' ' | b'\t'
318 | b'\n'
319 | b'\''
320 | b'"'
321 | b'\\'
322 | b'$'
323 | b'`'
324 | b'!'
325 | b'&'
326 | b'|'
327 | b';'
328 | b'<'
329 | b'>'
330 | b'('
331 | b')'
332 | b'['
333 | b']'
334 | b'{'
335 | b'}'
336 | b'*'
337 | b'?'
338 | b'#'
339 )
340 }) {
341 return input.to_string();
342 }
343 format!("'{}'", input.replace('\'', "'\"'\"'"))
344}
345
346fn has_pattern(text: &str, patterns: &[&str]) -> bool {
347 let lower = text.to_ascii_lowercase();
348 patterns.iter().any(|pattern| lower.contains(pattern))
349}
350
351#[derive(Debug, serde::Deserialize)]
352struct LarkTokenResponse {
353 code: i32,
354 msg: Option<String>,
355 tenant_access_token: Option<String>,
356}
357
358#[derive(Debug, serde::Deserialize)]
359struct LarkImageUploadResponse {
360 code: i32,
361 msg: Option<String>,
362 image_key: Option<String>,
363 data: Option<LarkImageUploadData>,
364}
365
366#[derive(Debug, serde::Deserialize)]
367struct LarkImageUploadData {
368 image_key: Option<String>,
369}
370
371fn parse_retry_time(text: &str, now_ms: u64) -> Option<(u64, String)> {
372 let lower = text.to_ascii_lowercase();
373 let marker = ["try again at ", "resets at ", "reset at ", "resets "]
374 .into_iter()
375 .find_map(|needle| lower.find(needle).map(|idx| (idx, needle)))?;
376 let start = marker.0 + marker.1.len();
377 let tail = text.get(start..)?.trim_start();
378 let mut chars = tail.chars().peekable();
379 let mut hour = String::new();
380 while let Some(ch) = chars.peek().copied() {
381 if ch.is_ascii_digit() {
382 hour.push(ch);
383 chars.next();
384 } else {
385 break;
386 }
387 }
388 if hour.is_empty() {
389 return None;
390 }
391 let mut minute = String::new();
392 if chars.peek() == Some(&':') {
393 chars.next();
394 while let Some(ch) = chars.peek().copied() {
395 if ch.is_ascii_digit() {
396 minute.push(ch);
397 chars.next();
398 } else {
399 break;
400 }
401 }
402 }
403 while let Some(ch) = chars.peek().copied() {
404 if ch.is_ascii_whitespace() {
405 chars.next();
406 } else {
407 break;
408 }
409 }
410 let mut meridiem = String::new();
411 while let Some(ch) = chars.peek().copied() {
412 if matches!(ch.to_ascii_lowercase(), 'a' | 'p' | 'm' | '.') {
413 meridiem.push(ch);
414 chars.next();
415 } else {
416 break;
417 }
418 }
419 let meridiem = meridiem.to_ascii_lowercase().replace('.', "");
420 if meridiem != "am" && meridiem != "pm" {
421 return None;
422 }
423 let raw_hour = hour.parse::<u32>().ok()?;
424 let minute = if minute.is_empty() {
425 0
426 } else {
427 minute.parse::<u32>().ok()?
428 };
429 if !(1..=12).contains(&raw_hour) || minute > 59 {
430 return None;
431 }
432 let now = chrono::DateTime::<chrono::Utc>::from(
433 SystemTime::UNIX_EPOCH + Duration::from_millis(now_ms),
434 );
435 let mut hour24 = raw_hour % 12;
436 if meridiem == "pm" {
437 hour24 += 12;
438 }
439 let mut retry_at = now
440 .date_naive()
441 .and_hms_opt(hour24, minute, 0)?
442 .and_utc()
443 .timestamp_millis() as u64;
444 if retry_at < now_ms && hour24 < 12 {
445 retry_at += 24 * 60 * 60 * 1000;
446 }
447 let label = tail
448 .split_whitespace()
449 .take(2)
450 .collect::<Vec<_>>()
451 .join(" ")
452 .trim_end_matches(|ch: char| ch == '.' || ch == ',' || ch == ';')
453 .to_string();
454 Some((retry_at, label))
455}
456
457fn detect_cli_usage_limit(text: &str, now_ms: u64) -> Option<CliUsageLimitState> {
458 if !text.to_ascii_lowercase().contains("again") && !text.to_ascii_lowercase().contains("reset")
459 {
460 return None;
461 }
462 let (retry_at_ms, retry_label) = parse_retry_time(text, now_ms)?;
463 let kind = if has_pattern(
464 text,
465 &["rate limit reached", "rate limit exceeded", "rate limited"],
466 ) {
467 CliUsageLimitKind::Rate
468 } else if has_pattern(
469 text,
470 &[
471 "hit your usage limit",
472 "hit usage limit",
473 "usage limit reached",
474 "usage limit exceeded",
475 "quota reached",
476 "quota exceeded",
477 "limit reached",
478 "limit exceeded",
479 "reached your usage limit",
480 "exceeded your usage limit",
481 ],
482 ) {
483 CliUsageLimitKind::Usage
484 } else {
485 return None;
486 };
487 Some(CliUsageLimitState {
488 limited: true,
489 kind,
490 retry_at_ms,
491 retry_label,
492 retry_ready: now_ms >= retry_at_ms,
493 })
494}
495
496fn usage_limit_state_key(state: &CliUsageLimitState) -> String {
497 format!(
498 "{:?}:{}:{}",
499 state.kind, state.retry_at_ms, state.retry_label
500 )
501}
502
503static PRIMARY_FONT: LazyLock<StdMutex<Option<FontVec>>> = LazyLock::new(|| StdMutex::new(None));
504static CJK_FONT: LazyLock<StdMutex<Option<FontVec>>> = LazyLock::new(|| StdMutex::new(None));
505const FONT_SIZE: f32 = 14.0;
506const CELL_W: f32 = 8.4;
507const CELL_H: f32 = 18.0;
508const PADDING: u32 = 12;
509const BG_COLOR: Rgba<u8> = Rgba([26, 27, 38, 255]);
510const FG_COLOR: Rgba<u8> = Rgba([169, 177, 214, 255]);
511
512fn home_font_dir() -> Option<std::path::PathBuf> {
513 std::env::var("HOME")
514 .ok()
515 .map(|h| std::path::PathBuf::from(h).join(".beam").join("fonts"))
516}
517
518fn load_font_files() {
519 let mut primary = PRIMARY_FONT.lock().unwrap();
520 if primary.is_some() {
521 return;
522 }
523
524 let search_paths: Vec<std::path::PathBuf> = {
525 let mut paths = Vec::new();
526 if let Some(d) = home_font_dir() {
527 paths.push(d.join("JetBrainsMono-Regular.ttf"));
528 paths.push(d.join("DejaVuSansMono.ttf"));
529 paths.push(d.join("NotoSansMonoCJKsc-Regular.otf"));
530 }
531 paths.push("/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf".into());
532 paths.push("/usr/share/fonts/dejavu/DejaVuSansMono.ttf".into());
533 paths.push("/usr/share/fonts/truetype/liberation/LiberationMono-Regular.ttf".into());
534 paths.push("/usr/share/fonts/liberation/LiberationMono-Regular.ttf".into());
535 paths.push("/usr/share/fonts/truetype/jetbrains-mono/JetBrainsMono-Regular.ttf".into());
536 paths
537 };
538
539 for path in &search_paths {
540 if let Ok(data) = std::fs::read(path) {
541 if let Ok(font) = FontVec::try_from_vec(data) {
542 *primary = Some(font);
543 break;
544 }
545 }
546 }
547
548 let cjk_search: Vec<std::path::PathBuf> = {
549 let mut paths = Vec::new();
550 if let Some(d) = home_font_dir() {
551 paths.push(d.join("NotoSansMonoCJKsc-Regular.otf"));
552 }
553 paths.push("/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc".into());
554 paths.push("/usr/share/fonts/noto-cjk/NotoSansCJK-Regular.ttc".into());
555 paths.push("/usr/share/fonts/google-noto-cjk/NotoSansCJK-Regular.ttc".into());
556 paths
557 };
558
559 let mut cjk = CJK_FONT.lock().unwrap();
560 for path in &cjk_search {
561 if let Ok(data) = std::fs::read(path) {
562 if let Ok(font) = FontVec::try_from_vec(data) {
563 *cjk = Some(font);
564 break;
565 }
566 }
567 }
568}
569
570fn is_fullwidth(ch: char) -> bool {
571 matches!(UnicodeWidthChar::width(ch), Some(2))
572}
573
574fn strip_ansi(s: &str) -> String {
575 let mut out = String::with_capacity(s.len());
576 let mut chars = s.chars().peekable();
577 while let Some(ch) = chars.next() {
578 if ch == '\x1b' {
579 if chars.peek() == Some(&'[') {
580 chars.next();
581 while let Some(&c) = chars.peek() {
582 if c.is_ascii_alphanumeric() || c == ';' {
583 chars.next();
584 if c.is_ascii_alphabetic() {
585 break;
586 }
587 } else {
588 break;
589 }
590 }
591 } else if chars.peek() == Some(&']') {
592 chars.next();
593 while let Some(&c) = chars.peek() {
594 chars.next();
595 if c == '\x07' || (c == '\x1b' && chars.peek() == Some(&'\\')) {
596 if c == '\x1b' {
597 chars.next();
598 }
599 break;
600 }
601 }
602 }
603 } else {
604 out.push(ch);
605 }
606 }
607 out
608}
609
610fn find_glyph_font<'a>(
611 ch: char,
612 primary: &'a FontVec,
613 cjk: Option<&'a FontVec>,
614) -> (&'a FontVec, f32) {
615 let primary_id = primary.glyph_id(ch);
616 if primary_id.0 != 0 {
617 return (primary, 1.0);
618 }
619 if let Some(cjk_font) = cjk {
620 let cjk_id = cjk_font.glyph_id(ch);
621 if cjk_id.0 != 0 {
622 return (cjk_font, if is_fullwidth(ch) { 2.0 } else { 1.0 });
623 }
624 }
625 (primary, 1.0)
626}
627
628fn render_text_screenshot_png(screen_raw: &str) -> Result<Vec<u8>> {
629 load_font_files();
630
631 let screen = strip_ansi(screen_raw);
632 let lines: Vec<&str> = screen.lines().collect();
633 let rows = lines.len().max(1);
634
635 let primary_guard = PRIMARY_FONT.lock().unwrap();
636 let cjk_guard = CJK_FONT.lock().unwrap();
637 let primary = primary_guard.as_ref();
638 let cjk = cjk_guard.as_ref();
639
640 let primary = match primary {
641 Some(f) => f,
642 None => return fallback_bitmap_png(&screen),
643 };
644
645 let scale = PxScale::from(FONT_SIZE);
646 let scaled = primary.as_scaled(scale);
647 let baseline_offset =
648 ((CELL_H - (scaled.ascent() + scaled.descent())).max(0.0) / 2.0) + scaled.ascent();
649
650 let cols = lines
651 .iter()
652 .map(|line| {
653 line.chars()
654 .map(|ch| if is_fullwidth(ch) { 2u32 } else { 1u32 })
655 .sum::<u32>()
656 })
657 .max()
658 .unwrap_or(1)
659 .max(1);
660
661 let width = ((cols as f32 * CELL_W).ceil() as u32 + PADDING * 2).clamp(64, 2200);
662 let height = ((rows as f32 * CELL_H).ceil() as u32 + PADDING * 2).clamp(32, 2200);
663
664 let mut image = ImageBuffer::from_pixel(width, height, BG_COLOR);
665
666 for (row, line) in lines.iter().enumerate() {
667 let mut col_cells: u32 = 0;
668 for ch in line.chars() {
669 let (font, char_width) = find_glyph_font(ch, primary, cjk);
670 let scaled = font.as_scaled(scale);
671 let x = PADDING as f32 + col_cells as f32 * CELL_W;
672 let y = PADDING as f32 + row as f32 * CELL_H;
673
674 if ch != ' ' {
675 let cell_px = char_width * CELL_W;
676 let advance = scaled.h_advance(scaled.glyph_id(ch));
677 let glyph_x = x + ((cell_px - advance).max(0.0) / 2.0);
678 let baseline = y + baseline_offset;
679 let mut glyph = scaled.scaled_glyph(ch);
680 glyph.position = point(glyph_x, baseline);
681 if let Some(outline) = font.outline_glyph(glyph) {
682 let bounds = outline.px_bounds();
683 outline.draw(|gx, gy, cv| {
684 let px = bounds.min.x as i32 + gx as i32;
685 let py = bounds.min.y as i32 + gy as i32;
686 if px >= 0
687 && py >= 0
688 && (px as u32) < width
689 && (py as u32) < height
690 && cv > 0.0
691 {
692 let alpha = (cv * 255.0).min(255.0) as u8;
693 if alpha == 255 {
694 image.put_pixel(px as u32, py as u32, FG_COLOR);
695 } else {
696 let existing = image.get_pixel(px as u32, py as u32);
697 let blended = blend_alpha(*existing, FG_COLOR, alpha);
698 image.put_pixel(px as u32, py as u32, blended);
699 }
700 }
701 });
702 }
703 }
704
705 col_cells += char_width.ceil() as u32;
706 }
707 }
708
709 let mut out = Vec::new();
710 let encoder = PngEncoder::new(&mut out);
711 encoder.write_image(image.as_raw(), width, height, ColorType::Rgba8.into())?;
712 Ok(out)
713}
714
715fn blend_alpha(bg: Rgba<u8>, fg: Rgba<u8>, alpha: u8) -> Rgba<u8> {
716 let a = alpha as f32 / 255.0;
717 let r = (fg.0[0] as f32 * a + bg.0[0] as f32 * (1.0 - a)) as u8;
718 let g = (fg.0[1] as f32 * a + bg.0[1] as f32 * (1.0 - a)) as u8;
719 let b = (fg.0[2] as f32 * a + bg.0[2] as f32 * (1.0 - a)) as u8;
720 Rgba([r, g, b, 255])
721}
722
723fn fallback_bitmap_png(screen: &str) -> Result<Vec<u8>> {
724 use font8x8::UnicodeFonts;
725
726 let lines: Vec<&str> = screen.lines().collect();
727 let rows = lines.len().max(1);
728 let cols = lines
729 .iter()
730 .map(|line| line.chars().count())
731 .max()
732 .unwrap_or(1)
733 .max(1);
734 let scale = 2u32;
735 let glyph_w = 8u32 * scale;
736 let glyph_h = 8u32 * scale;
737 let width = (cols as u32 * glyph_w + PADDING * 2).clamp(64, 2200);
738 let height = (rows as u32 * glyph_h + PADDING * 2).clamp(32, 2200);
739 let bg = Rgba([15, 23, 42, 255]);
740 let fg = Rgba([226, 232, 240, 255]);
741 let mut image = ImageBuffer::from_pixel(width, height, bg);
742
743 for (row, line) in lines.iter().enumerate() {
744 for (col, ch) in line.chars().take(cols as usize).enumerate() {
745 let glyph = font8x8::BASIC_FONTS
746 .get(ch)
747 .or_else(|| font8x8::BASIC_FONTS.get('?'))
748 .unwrap_or([0; 8]);
749 for (gy, bits) in glyph.iter().enumerate() {
750 for gx in 0..8 {
751 if (bits >> gx) & 1 == 0 {
752 continue;
753 }
754 for sy in 0..scale {
755 for sx in 0..scale {
756 let x = PADDING + col as u32 * glyph_w + (7 - gx) as u32 * scale + sx;
757 let y = PADDING + row as u32 * glyph_h + gy as u32 * scale + sy;
758 if x < width && y < height {
759 image.put_pixel(x, y, fg);
760 }
761 }
762 }
763 }
764 }
765 }
766 }
767
768 let mut out = Vec::new();
769 let encoder = PngEncoder::new(&mut out);
770 encoder.write_image(image.as_raw(), width, height, ColorType::Rgba8.into())?;
771 Ok(out)
772}
773
774fn lark_base_url() -> &'static str {
775 "https://open.feishu.cn/open-apis"
776}
777
778async fn lark_tenant_token(app_id: &str, secret: &str) -> Result<String> {
779 let body = reqwest::Client::new()
780 .post(format!(
781 "{}/auth/v3/tenant_access_token/internal",
782 lark_base_url()
783 ))
784 .json(&serde_json::json!({
785 "app_id": app_id,
786 "app_secret": secret,
787 }))
788 .send()
789 .await?
790 .json::<LarkTokenResponse>()
791 .await?;
792 if body.code != 0 {
793 anyhow::bail!(
794 "lark tenant_access_token failed: {}",
795 body.msg.unwrap_or_else(|| "unknown error".to_string())
796 );
797 }
798 body.tenant_access_token
799 .context("lark tenant_access_token missing")
800}
801
802async fn upload_image_buffer(app_id: &str, secret: &str, image: Vec<u8>) -> Result<String> {
803 let token = lark_tenant_token(app_id, secret).await?;
804 let form = Form::new().text("image_type", "message").part(
805 "image",
806 Part::bytes(image)
807 .file_name("screen.png")
808 .mime_str("image/png")?,
809 );
810 let body = reqwest::Client::new()
811 .post(format!("{}/im/v1/images", lark_base_url()))
812 .bearer_auth(token)
813 .multipart(form)
814 .send()
815 .await?
816 .json::<LarkImageUploadResponse>()
817 .await?;
818 if body.code != 0 {
819 anyhow::bail!(
820 "lark image upload failed: {}",
821 body.msg.unwrap_or_else(|| "unknown error".to_string())
822 );
823 }
824 body.image_key
825 .or_else(|| body.data.and_then(|data| data.image_key))
826 .context("lark image upload missing image_key")
827}
828
829async fn maybe_send_screenshot_upload(
830 stdout: &Arc<Mutex<tokio::io::Stdout>>,
831 app_id: &str,
832 app_secret: &str,
833 screen: &str,
834 status: ScreenStatus,
835 usage_limit: Option<CliUsageLimitState>,
836 last_uploaded_hash: &Arc<Mutex<Option<String>>>,
837) {
838 if app_id == "local" || app_secret.is_empty() {
839 return;
840 }
841 let hash = format!("{:x}", Sha256::digest(screen.as_bytes()));
842 {
843 let guard = last_uploaded_hash.lock().await;
844 if guard.as_deref() == Some(hash.as_str()) {
845 return;
846 }
847 }
848 let png = match render_text_screenshot_png(screen) {
849 Ok(png) => png,
850 Err(err) => {
851 warn!("failed to render terminal screenshot: {err:#}");
852 return;
853 }
854 };
855 let image_key = match upload_image_buffer(app_id, app_secret, png).await {
856 Ok(image_key) => image_key,
857 Err(err) => {
858 warn!("failed to upload terminal screenshot: {err:#}");
859 return;
860 }
861 };
862 let _ = send_message(
863 stdout,
864 &WorkerToDaemon::ScreenshotUploaded {
865 image_key,
866 status,
867 usage_limit,
868 },
869 )
870 .await;
871 *last_uploaded_hash.lock().await = Some(hash);
872}
873
874#[derive(Debug, Default)]
875struct UsageLimitTracker {
876 turn_seq: u64,
877 detected_turn: Option<u64>,
878 suppressed_retry_ready_key: Option<String>,
879}
880
881impl UsageLimitTracker {
882 fn begin_turn(&mut self, snapshot: &str, now_ms: u64) -> u64 {
883 self.turn_seq += 1;
884 self.detected_turn = None;
885 self.suppressed_retry_ready_key = detect_cli_usage_limit(snapshot, now_ms)
886 .filter(|state| state.retry_ready)
887 .map(|state| usage_limit_state_key(&state));
888 self.turn_seq
889 }
890
891 fn classify(
892 &mut self,
893 content: &str,
894 status: ScreenStatus,
895 now_ms: u64,
896 ) -> (ScreenStatus, Option<CliUsageLimitState>) {
897 let Some(detected) = detect_cli_usage_limit(content, now_ms) else {
898 return (status, None);
899 };
900 let key = usage_limit_state_key(&detected);
901 if detected.retry_ready && self.suppressed_retry_ready_key.as_deref() == Some(key.as_str())
902 {
903 return (status, None);
904 }
905 self.suppressed_retry_ready_key = None;
906 self.detected_turn = Some(self.turn_seq);
907 (ScreenStatus::Limited, Some(detected))
908 }
909}
910
911async fn prepare_wrapper(init: &InitConfig, paths: &BeamPaths) -> Result<std::path::PathBuf> {
912 tokio::fs::create_dir_all(paths.run_dir()).await?;
913 let wrapper = paths.worker_wrapper_sh(&init.session_id);
914 let exe_path = std::env::current_exe()
915 .ok()
916 .map(|p| p.display().to_string())
917 .unwrap_or_else(|| "beam".to_string());
918 let content = format!(
919 "#!/bin/sh\ncd {cwd}\nexport BEAM_SESSION_ID={sid}\nexport BEAM_HOME={home}\nexport BEAM_BIN={exe}\nif [ -n \"$PATH\" ]; then\n export PATH={bindir}:$PATH\nelse\n export PATH={bindir}\nfi\nexec \"$@\"\n",
920 cwd = shell_quote(&init.working_dir),
921 sid = shell_quote(&init.session_id),
922 home = shell_quote(&paths.root().display().to_string()),
923 exe = shell_quote(&exe_path),
924 bindir = shell_quote(
925 &std::env::current_exe()
926 .ok()
927 .and_then(|p| p.parent().map(|v| v.display().to_string()))
928 .unwrap_or_default()
929 ),
930 );
931 tokio::fs::write(&wrapper, content).await?;
932 #[cfg(unix)]
933 {
934 use std::os::unix::fs::PermissionsExt;
935 let perms = std::fs::Permissions::from_mode(0o755);
936 tokio::fs::set_permissions(&wrapper, perms).await?;
937 }
938 Ok(wrapper)
939}
940
941async fn send_message(stdout: &Arc<Mutex<tokio::io::Stdout>>, msg: &WorkerToDaemon) -> Result<()> {
942 let mut out = stdout.lock().await;
943 out.write_all(serde_json::to_string(msg)?.as_bytes())
944 .await?;
945 out.write_all(b"\n").await?;
946 out.flush().await?;
947 Ok(())
948}
949
950async fn handle_tui_keys(
951 backend: &Arc<Mutex<Box<dyn SessionBackend>>>,
952 analyzer_runtime: &Arc<RwLock<AnalyzerRuntime>>,
953 keys: &[String],
954 is_final: bool,
955) -> Result<()> {
956 if keys.is_empty() {
957 return Ok(());
958 }
959 let guard = backend.lock().await;
960 for key in keys {
961 guard.send_special_keys(std::slice::from_ref(key)).await?;
962 tokio::time::sleep(Duration::from_millis(100)).await;
963 }
964 drop(guard);
965 if is_final {
966 analyzer_runtime.write().await.prompt_active = false;
967 }
968 Ok(())
969}
970
971async fn handle_tui_text_input(
972 backend: &Arc<Mutex<Box<dyn SessionBackend>>>,
973 adapter: &Arc<Mutex<CliAdapter>>,
974 analyzer_runtime: &Arc<RwLock<AnalyzerRuntime>>,
975 keys: &[String],
976 text: &str,
977) -> Result<()> {
978 let nav_keys = if keys.last().map(String::as_str) == Some("Enter") {
979 &keys[..keys.len().saturating_sub(1)]
980 } else {
981 keys
982 };
983 if !nav_keys.is_empty() {
984 let guard = backend.lock().await;
985 for key in nav_keys {
986 guard.send_special_keys(std::slice::from_ref(key)).await?;
987 tokio::time::sleep(Duration::from_millis(100)).await;
988 }
989 }
990 analyzer_runtime.write().await.prompt_active = false;
991 tokio::time::sleep(Duration::from_millis(200)).await;
992 let guard = backend.lock().await;
993 let _ = adapter
994 .lock()
995 .await
996 .write_input(guard.as_ref(), text)
997 .await?;
998 Ok(())
999}
1000
1001async fn handle_tui_prompt_override(
1002 stdout: &Arc<Mutex<tokio::io::Stdout>>,
1003 analyzer_runtime: &Arc<RwLock<AnalyzerRuntime>>,
1004) {
1005 let was_active = {
1006 let mut runtime = analyzer_runtime.write().await;
1007 if runtime.prompt_active {
1008 runtime.prompt_active = false;
1009 true
1010 } else {
1011 false
1012 }
1013 };
1014 if was_active {
1015 let _ = send_message(
1016 stdout,
1017 &WorkerToDaemon::TuiPromptResolved {
1018 selected_text: Some("user-override".to_string()),
1019 },
1020 )
1021 .await;
1022 }
1023}
1024
1025pub async fn run(init: InitConfig) -> Result<()> {
1026 let stdout = Arc::new(Mutex::new(tokio::io::stdout()));
1027 let session_name = format!("bmx-{}", &init.session_id[..8.min(init.session_id.len())]);
1028 let paths = BeamPaths::discover()?;
1029 let adapter = Arc::new(Mutex::new(CliAdapter::from_init(&init)?));
1030 let wrapper = if init.resume || init.adopted_from.is_some() {
1031 None
1032 } else {
1033 Some(prepare_wrapper(&init, &paths).await?)
1034 };
1035 let (mut backend_impl, attach_context): (Box<dyn SessionBackend>, &'static str) =
1036 match init.backend_type {
1037 BackendType::Tmux => {
1038 let tmux = if let Some(target) = init
1039 .adopted_from
1040 .as_ref()
1041 .and_then(|adopted| adopted.tmux_target.clone())
1042 {
1043 TmuxPipeBackend::attach_existing(target)
1044 } else if init.resume {
1045 TmuxPipeBackend::attach_existing(session_name)
1046 } else {
1047 TmuxPipeBackend::new(session_name)
1048 };
1049 (
1050 Box::new(tmux),
1051 if init.resume || init.adopted_from.is_some() {
1052 "attach"
1053 } else {
1054 "spawn"
1055 },
1056 )
1057 }
1058 BackendType::Pty => (Box::new(PtyBackend::new()), "spawn"),
1059 BackendType::Zellij => {
1060 if let Some(adopted) = init.adopted_from.as_ref() {
1061 if let Some(pane_id) = adopted.zellij_pane_id.clone() {
1062 let session = adopted.zellij_session.clone().unwrap_or_else(|| {
1063 format!("bmx-{}", &init.session_id[..8.min(init.session_id.len())])
1064 });
1065 let observe = ZellijObserveBackend::new(
1066 session,
1067 pane_id,
1068 u32::try_from(adopted.original_cli_pid).ok(),
1069 );
1070 (Box::new(observe), "observe")
1071 } else {
1072 let zellij = ZellijBackend::new(format!(
1073 "bmx-{}",
1074 &init.session_id[..8.min(init.session_id.len())]
1075 ));
1076 (Box::new(zellij), "spawn")
1077 }
1078 } else {
1079 let zellij = ZellijBackend::new(format!(
1080 "bmx-{}",
1081 &init.session_id[..8.min(init.session_id.len())]
1082 ));
1083 (Box::new(zellij), "spawn")
1084 }
1085 }
1086 };
1087 let spawn_spec = adapter.lock().await.build_spawn_spec(&init);
1088 let args = if let Some(wrapper) = wrapper {
1089 let mut args = Vec::with_capacity(2 + init.cli_args.len());
1090 args.push(wrapper.display().to_string());
1091 args.push(spawn_spec.bin.clone());
1092 args.extend(spawn_spec.args.clone());
1093 ("/bin/sh".to_string(), args)
1094 } else {
1095 (spawn_spec.bin, spawn_spec.args)
1096 };
1097 backend_impl
1098 .spawn(
1099 &args.0,
1100 &args.1,
1101 SpawnOpts {
1102 cwd: init.working_dir.clone(),
1103 cols: 160,
1104 rows: 40,
1105 env: Vec::new(),
1106 },
1107 )
1108 .await
1109 .with_context(|| format!("failed to {} session {}", attach_context, init.session_id))?;
1110
1111 let backend: Arc<Mutex<Box<dyn SessionBackend>>> = Arc::new(Mutex::new(backend_impl));
1112 let mut cli_pid_marker = None;
1113 let child_pid = backend.lock().await.child_pid().await?;
1114 adapter.lock().await.on_spawned(child_pid);
1115 if let Some(pid) = child_pid {
1116 tokio::fs::create_dir_all(paths.cli_pid_markers_dir()).await?;
1117 let marker = paths.cli_pid_markers_dir().join(pid.to_string());
1118 tokio::fs::write(&marker, init.session_id.as_bytes()).await?;
1119 cli_pid_marker = Some(marker);
1120 }
1121 let latest_screen = Arc::new(RwLock::new(String::new()));
1122 let latest_raw_screen = Arc::new(RwLock::new(String::new()));
1123 let display_mode = Arc::new(RwLock::new(DisplayMode::Hidden));
1124 let analyzer_runtime = Arc::new(RwLock::new(AnalyzerRuntime::default()));
1125 let usage_limit_tracker = Arc::new(Mutex::new(UsageLimitTracker::default()));
1126 let current_turn_id = Arc::new(RwLock::new(String::new()));
1127 let (updates, _) = broadcast::channel::<String>(256);
1128 let token = Uuid::new_v4().to_string();
1129 let addr = serve(TerminalState {
1130 backend: backend.clone(),
1131 latest_screen: latest_screen.clone(),
1132 token: token.clone(),
1133 updates: updates.clone(),
1134 })
1135 .await?;
1136
1137 send_message(
1138 &stdout,
1139 &WorkerToDaemon::Ready {
1140 port: addr.port(),
1141 token,
1142 },
1143 )
1144 .await?;
1145
1146 let sample_backend = backend.clone();
1147 let sample_screen = latest_screen.clone();
1148 let sample_raw_screen = latest_raw_screen.clone();
1149 let sample_updates = updates.clone();
1150 let sample_stdout = stdout.clone();
1151 let sample_adapter = adapter.clone();
1152 let sample_display_mode = display_mode.clone();
1153 let sample_usage_limit_tracker = usage_limit_tracker.clone();
1154 let sample_current_turn_id = current_turn_id.clone();
1155 let screenshot_app_id = init.lark_app_id.clone();
1156 let screenshot_app_secret = init.lark_app_secret.clone();
1157 let last_uploaded_hash = Arc::new(Mutex::new(None::<String>));
1158 let sample_last_uploaded_hash = last_uploaded_hash.clone();
1159 let last_broadcast_hash: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
1160 let sample_last_broadcast_hash = last_broadcast_hash.clone();
1161 let sample_analyzer_runtime = analyzer_runtime.clone();
1162 let screen_capture_task = tokio::spawn(async move {
1163 let mut last_emitted_status = ScreenStatus::Starting;
1164 let mut last_emitted_usage_limit: Option<CliUsageLimitState> = None;
1165 loop {
1166 let (screen, alive) = {
1167 let guard = sample_backend.lock().await;
1168 let screen = guard.capture_viewport().await.unwrap_or_default();
1169 let alive = guard.is_alive().await.unwrap_or(false);
1170 (screen, alive)
1171 };
1172
1173 let hash_changed;
1174
1175 {
1176 *sample_raw_screen.write().await = screen.clone();
1177 let mode = *sample_display_mode.read().await;
1178 let rendered = render_screen_for_display_mode(&screen, mode);
1179 let now_ms = now_ms();
1180 let analyzing = sample_analyzer_runtime.read().await.is_analyzing;
1181 let base_status = if analyzing {
1182 ScreenStatus::Analyzing
1183 } else {
1184 ScreenStatus::Working
1185 };
1186 let (status, usage_limit) =
1187 sample_usage_limit_tracker
1188 .lock()
1189 .await
1190 .classify(&screen, base_status, now_ms);
1191 let rendered_hash = format!("{:x}", Sha256::digest(rendered.as_bytes()));
1192 {
1193 let guard = sample_last_broadcast_hash.lock().await;
1194 hash_changed = guard.as_deref() != Some(&rendered_hash);
1195 }
1196
1197 if hash_changed
1198 || last_emitted_status != status
1199 || last_emitted_usage_limit != usage_limit
1200 {
1201 *sample_last_broadcast_hash.lock().await = Some(rendered_hash.clone());
1202 let mut current = sample_screen.write().await;
1203 *current = rendered.clone();
1204 let _ = sample_updates.send(rendered.clone());
1205 let _ = send_message(
1206 &sample_stdout,
1207 &WorkerToDaemon::ScreenUpdate {
1208 content: rendered.clone(),
1209 status,
1210 usage_limit: usage_limit.clone(),
1211 },
1212 )
1213 .await;
1214 last_emitted_status = status;
1215 last_emitted_usage_limit = usage_limit.clone();
1216 }
1217
1218 if mode == DisplayMode::Screenshot
1219 && screenshot_app_id != "local"
1220 && !screenshot_app_secret.is_empty()
1221 {
1222 maybe_send_screenshot_upload(
1223 &sample_stdout,
1224 &screenshot_app_id,
1225 &screenshot_app_secret,
1226 &screen,
1227 status,
1228 usage_limit.clone(),
1229 &sample_last_uploaded_hash,
1230 )
1231 .await;
1232 }
1233 }
1234
1235 if let Ok(poll) = sample_adapter.lock().await.poll() {
1236 if let Some(cli_session_id) = poll.cli_session_id {
1237 let _ = send_message(
1238 &sample_stdout,
1239 &WorkerToDaemon::CliSessionId { cli_session_id },
1240 )
1241 .await;
1242 }
1243 if let Some((user_text, assistant_text)) = poll.adopt_preamble {
1244 let _ = send_message(
1245 &sample_stdout,
1246 &WorkerToDaemon::AdoptPreamble {
1247 user_text,
1248 assistant_text,
1249 },
1250 )
1251 .await;
1252 }
1253 if let Some(content) = poll.final_output {
1254 let turn_id = sample_current_turn_id.read().await.clone();
1255 let _ = send_message(
1256 &sample_stdout,
1257 &WorkerToDaemon::FinalOutput {
1258 content,
1259 turn_id,
1260 kind: poll.final_output_kind,
1261 user_text: poll.final_output_user_text,
1262 },
1263 )
1264 .await;
1265 }
1266 if poll.prompt_ready {
1267 let _ = send_message(&sample_stdout, &WorkerToDaemon::PromptReady).await;
1268 let rendered = sample_screen.read().await.clone();
1269 let raw = sample_raw_screen.read().await.clone();
1270 let now_ms = now_ms();
1271 let analyzing = sample_analyzer_runtime.read().await.is_analyzing;
1272 let base_status = if analyzing {
1273 ScreenStatus::Analyzing
1274 } else {
1275 ScreenStatus::Idle
1276 };
1277 let (status, usage_limit) =
1278 sample_usage_limit_tracker
1279 .lock()
1280 .await
1281 .classify(&raw, base_status, now_ms);
1282 let _ = send_message(
1283 &sample_stdout,
1284 &WorkerToDaemon::ScreenUpdate {
1285 content: rendered,
1286 status,
1287 usage_limit,
1288 },
1289 )
1290 .await;
1291 }
1292 }
1293
1294 if !alive {
1295 let _ = send_message(
1296 &sample_stdout,
1297 &WorkerToDaemon::CliExit {
1298 code: Some(0),
1299 signal: None,
1300 },
1301 )
1302 .await;
1303 break;
1304 }
1305
1306 tokio::time::sleep(Duration::from_millis(5000)).await;
1307 }
1308 });
1309 let mut worker_joins = tokio::task::JoinSet::new();
1310 worker_joins.spawn(async move {
1311 let _ = screen_capture_task.await;
1312 });
1313
1314 if screen_analyzer_enabled(&init.screen_analyzer) {
1315 let analyzer_cfg = init.screen_analyzer.clone();
1316 let analyzer_raw_screen = latest_raw_screen.clone();
1317 let analyzer_runtime_state = analyzer_runtime.clone();
1318 let analyzer_stdout = stdout.clone();
1319 let analyzer_task = tokio::spawn(async move {
1320 let client = Client::new();
1321 loop {
1322 tokio::time::sleep(Duration::from_millis(analyzer_cfg.interval_ms)).await;
1323 let snapshot = analyzer_raw_screen.read().await.clone();
1324 if snapshot.is_empty() {
1325 continue;
1326 }
1327 let truncated = if snapshot.len() > analyzer_cfg.snapshot_max_chars {
1328 snapshot[snapshot.len() - analyzer_cfg.snapshot_max_chars..].to_string()
1329 } else {
1330 snapshot
1331 };
1332 let now = now_ms();
1333 {
1334 let mut runtime = analyzer_runtime_state.write().await;
1335 if truncated == runtime.last_snapshot {
1336 runtime.stable_count = runtime.stable_count.saturating_add(1);
1337 } else {
1338 runtime.stable_count = 1;
1339 runtime.last_snapshot = truncated.clone();
1340 if runtime.waiting_for_content_change {
1341 runtime.waiting_for_content_change = false;
1342 }
1343 }
1344 if runtime.stable_count < analyzer_cfg.stable_count {
1345 continue;
1346 }
1347 if runtime.waiting_for_content_change
1348 && truncated == runtime.last_analyzed_snapshot
1349 {
1350 continue;
1351 }
1352 if runtime.cooldown_until_ms > now {
1353 continue;
1354 }
1355 runtime.is_analyzing = true;
1356 runtime.last_analyzed_snapshot = truncated.clone();
1357 }
1358
1359 let result = call_screen_analyzer(&client, &analyzer_cfg, &truncated).await;
1360
1361 let mut runtime = analyzer_runtime_state.write().await;
1362 runtime.is_analyzing = false;
1363 match result {
1364 Ok(analysis) => {
1365 apply_screen_analyzer_result(
1366 &mut runtime,
1367 &analysis.check_again_when,
1368 now_ms(),
1369 );
1370 if analysis.needs_interaction && !analysis.options.is_empty() {
1371 if !runtime.prompt_active {
1372 runtime.prompt_active = true;
1373 let _ = send_message(
1374 &analyzer_stdout,
1375 &WorkerToDaemon::TuiPrompt {
1376 description: analysis.description.clone().unwrap_or_else(
1377 || "CLI needs your selection".to_string(),
1378 ),
1379 options: analysis.options.clone(),
1380 multi_select: analysis.multi_select,
1381 },
1382 )
1383 .await;
1384 }
1385 } else if runtime.prompt_active {
1386 runtime.prompt_active = false;
1387 let _ = send_message(
1388 &analyzer_stdout,
1389 &WorkerToDaemon::TuiPromptResolved {
1390 selected_text: None,
1391 },
1392 )
1393 .await;
1394 }
1395 }
1396 Err(_) => {
1397 runtime.waiting_for_content_change = true;
1398 runtime.cooldown_until_ms = 0;
1399 }
1400 }
1401 }
1402 });
1403 worker_joins.spawn(async move {
1404 let _ = analyzer_task.await;
1405 });
1406 }
1407
1408 if !init.prompt.is_empty() && !crate::adapters::passes_initial_prompt_via_args(&init.cli_id) {
1409 usage_limit_tracker.lock().await.begin_turn(
1410 "",
1411 SystemTime::now()
1412 .duration_since(UNIX_EPOCH)
1413 .unwrap_or_default()
1414 .as_millis() as u64,
1415 );
1416 *current_turn_id.write().await = init
1417 .prompt_turn_id
1418 .clone()
1419 .unwrap_or_else(|| Uuid::new_v4().to_string());
1420 *last_uploaded_hash.lock().await = None;
1421 let guard = backend.lock().await;
1422 let submit = adapter
1423 .lock()
1424 .await
1425 .write_input(guard.as_ref(), &init.prompt)
1426 .await?;
1427 if let Some(cli_session_id) = submit.cli_session_id {
1428 send_message(&stdout, &WorkerToDaemon::CliSessionId { cli_session_id }).await?;
1429 }
1430 }
1431
1432 let stdin = BufReader::new(tokio::io::stdin());
1433 let mut lines = stdin.lines();
1434 loop {
1435 let line = match lines.next_line().await {
1436 Ok(Some(line)) => line,
1437 Ok(None) => break,
1438 Err(_) => break,
1439 };
1440 if line.trim().is_empty() {
1441 continue;
1442 }
1443 let msg: DaemonToWorker = serde_json::from_str(&line)?;
1444 match msg {
1445 DaemonToWorker::Message { content, turn_id } => {
1446 handle_tui_prompt_override(&stdout, &analyzer_runtime).await;
1447 let snapshot = latest_raw_screen.read().await.clone();
1448 usage_limit_tracker.lock().await.begin_turn(
1449 &snapshot,
1450 SystemTime::now()
1451 .duration_since(UNIX_EPOCH)
1452 .unwrap_or_default()
1453 .as_millis() as u64,
1454 );
1455 *current_turn_id.write().await = turn_id;
1456 *last_uploaded_hash.lock().await = None;
1457 let guard = backend.lock().await;
1458 let submit = adapter
1459 .lock()
1460 .await
1461 .write_input(guard.as_ref(), &content)
1462 .await?;
1463 if let Some(cli_session_id) = submit.cli_session_id {
1464 send_message(&stdout, &WorkerToDaemon::CliSessionId { cli_session_id }).await?;
1465 }
1466 if !submit.submitted {
1467 let message = submit
1468 .failure_reason
1469 .unwrap_or_else(|| "CLI submit could not be confirmed".to_string());
1470 send_message(&stdout, &WorkerToDaemon::UserNotify { message }).await?;
1471 }
1472 }
1473 DaemonToWorker::RawInput { content, turn_id } => {
1474 handle_tui_prompt_override(&stdout, &analyzer_runtime).await;
1475 let snapshot = latest_raw_screen.read().await.clone();
1476 usage_limit_tracker.lock().await.begin_turn(
1477 &snapshot,
1478 SystemTime::now()
1479 .duration_since(UNIX_EPOCH)
1480 .unwrap_or_default()
1481 .as_millis() as u64,
1482 );
1483 *current_turn_id.write().await = turn_id;
1484 *last_uploaded_hash.lock().await = None;
1485 let guard = backend.lock().await;
1486 guard.raw_input(&content).await?;
1487 }
1488 DaemonToWorker::Close => {
1489 let mut guard = backend.lock().await;
1490 guard.destroy_session().await?;
1491 break;
1492 }
1493 DaemonToWorker::Restart => {
1494 let mut guard = backend.lock().await;
1495 guard.destroy_session().await?;
1496 break;
1497 }
1498 DaemonToWorker::RefreshScreen => {
1499 let guard = backend.lock().await;
1500 let screen = guard.capture_viewport().await?;
1501 *latest_raw_screen.write().await = screen.clone();
1502 let mode = *display_mode.read().await;
1503 let rendered = render_screen_for_display_mode(&screen, mode);
1504 let now_ms = now_ms();
1505 let analyzing = analyzer_runtime.read().await.is_analyzing;
1506 let base_status = if analyzing {
1507 ScreenStatus::Analyzing
1508 } else {
1509 ScreenStatus::Working
1510 };
1511 let (status, usage_limit) =
1512 usage_limit_tracker
1513 .lock()
1514 .await
1515 .classify(&screen, base_status, now_ms);
1516 *latest_screen.write().await = rendered.clone();
1517 let _ = updates.send(rendered.clone());
1518 send_message(
1519 &stdout,
1520 &WorkerToDaemon::ScreenUpdate {
1521 content: rendered.clone(),
1522 status,
1523 usage_limit: usage_limit.clone(),
1524 },
1525 )
1526 .await?;
1527 let rendered_hash = format!("{:x}", Sha256::digest(rendered.as_bytes()));
1528 *last_broadcast_hash.lock().await = Some(rendered_hash);
1529 if mode == DisplayMode::Screenshot {
1530 maybe_send_screenshot_upload(
1531 &stdout,
1532 &init.lark_app_id,
1533 &init.lark_app_secret,
1534 &screen,
1535 status,
1536 usage_limit,
1537 &last_uploaded_hash,
1538 )
1539 .await;
1540 }
1541 }
1542 DaemonToWorker::SetDisplayMode { mode } => {
1543 *display_mode.write().await = mode;
1544 let raw = latest_raw_screen.read().await.clone();
1545 let rendered = render_screen_for_display_mode(&raw, mode);
1546 let now_ms = now_ms();
1547 let analyzing = analyzer_runtime.read().await.is_analyzing;
1548 let base_status = if analyzing {
1549 ScreenStatus::Analyzing
1550 } else {
1551 ScreenStatus::Working
1552 };
1553 let (status, usage_limit) =
1554 usage_limit_tracker
1555 .lock()
1556 .await
1557 .classify(&raw, base_status, now_ms);
1558 *latest_screen.write().await = rendered.clone();
1559 let _ = updates.send(rendered.clone());
1560 send_message(
1561 &stdout,
1562 &WorkerToDaemon::ScreenUpdate {
1563 content: rendered,
1564 status,
1565 usage_limit: usage_limit.clone(),
1566 },
1567 )
1568 .await?;
1569 if mode == DisplayMode::Screenshot {
1570 maybe_send_screenshot_upload(
1571 &stdout,
1572 &init.lark_app_id,
1573 &init.lark_app_secret,
1574 &raw,
1575 status,
1576 usage_limit,
1577 &last_uploaded_hash,
1578 )
1579 .await;
1580 }
1581 }
1582 DaemonToWorker::TermAction { key } => {
1583 let keys = term_action_keys(key);
1584 let guard = backend.lock().await;
1585 guard.send_special_keys(&keys).await?;
1586 }
1587 DaemonToWorker::SpecialKeys { keys } => {
1588 let guard = backend.lock().await;
1589 guard.send_special_keys(&keys).await?;
1590 }
1591 DaemonToWorker::TuiKeys { keys, is_final } => {
1592 handle_tui_keys(&backend, &analyzer_runtime, &keys, is_final).await?;
1593 }
1594 DaemonToWorker::TuiTextInput { keys, text } => {
1595 handle_tui_text_input(&backend, &adapter, &analyzer_runtime, &keys, &text).await?;
1596 }
1597 DaemonToWorker::Init(_) => {}
1598 }
1599 }
1600
1601 worker_joins.abort_all();
1602 while worker_joins.join_next().await.is_some() {}
1603 {
1604 let mut guard = backend.lock().await;
1605 let _ = guard.kill().await;
1606 }
1607 if let Some(marker) = cli_pid_marker {
1608 let _ = tokio::fs::remove_file(marker).await;
1609 }
1610
1611 info!("worker exiting");
1612 Ok(())
1613}
1614
1615#[cfg(test)]
1616mod tests {
1617 use super::*;
1618
1619 #[test]
1620 fn render_screen_for_display_mode_hides_or_shows_content() {
1621 assert_eq!(
1622 render_screen_for_display_mode("hello", DisplayMode::Hidden),
1623 "[screen hidden]"
1624 );
1625 assert_eq!(
1626 render_screen_for_display_mode("hello", DisplayMode::Screenshot),
1627 "hello"
1628 );
1629 }
1630
1631 #[test]
1632 fn term_action_keys_maps_supported_actions() {
1633 assert_eq!(term_action_keys(TermActionKey::Esc), vec!["Escape"]);
1634 assert_eq!(term_action_keys(TermActionKey::CtrlC), vec!["C-c"]);
1635 assert_eq!(term_action_keys(TermActionKey::Enter), vec!["Enter"]);
1636 assert_eq!(
1637 term_action_keys(TermActionKey::HalfPageDown),
1638 vec!["PageDown"]
1639 );
1640 }
1641
1642 #[test]
1643 fn detect_cli_usage_limit_finds_usage_and_rate_limits() {
1644 let usage = detect_cli_usage_limit(
1645 "You have hit your usage limit. Try again at 3:15 PM.",
1646 1_700_000_000_000,
1647 )
1648 .expect("usage limit detected");
1649 assert_eq!(usage.kind, CliUsageLimitKind::Usage);
1650 assert_eq!(usage.retry_label, "3:15 PM");
1651 assert!(usage.limited);
1652
1653 let rate = detect_cli_usage_limit("Rate limited. Resets at 11:00 AM.", 1_700_000_000_000)
1654 .expect("rate limit detected");
1655 assert_eq!(rate.kind, CliUsageLimitKind::Rate);
1656 assert_eq!(rate.retry_label, "11:00 AM");
1657 }
1658
1659 #[test]
1660 fn usage_limit_tracker_suppresses_stale_retry_ready_banner_on_new_turn() {
1661 let now_ms = 1_700_000_000_000;
1662 let text = "Usage limit reached. Try again at 3:15 PM.";
1663 let mut tracker = UsageLimitTracker::default();
1664 let initial = detect_cli_usage_limit(text, now_ms).expect("limit");
1665 tracker.begin_turn(text, initial.retry_at_ms + 1);
1666 let (status, usage_limit) =
1667 tracker.classify(text, ScreenStatus::Working, initial.retry_at_ms + 1);
1668 assert_eq!(status, ScreenStatus::Working);
1669 assert_eq!(usage_limit, None);
1670 }
1671
1672 #[test]
1673 fn render_text_screenshot_png_produces_png_bytes() {
1674 let png = render_text_screenshot_png("hello\nworld").expect("png rendered");
1675 assert!(png.starts_with(&[0x89, b'P', b'N', b'G']));
1676 assert!(png.len() > 64);
1677 }
1678
1679 #[test]
1680 fn screen_analyzer_enablement_requires_complete_config() {
1681 let mut cfg = ScreenAnalyzerConfig::default();
1682 assert!(!screen_analyzer_enabled(&cfg));
1683 cfg.enabled = true;
1684 cfg.base_url = "https://example.com".to_string();
1685 cfg.api_key = "k".to_string();
1686 cfg.model = "m".to_string();
1687 assert!(screen_analyzer_enabled(&cfg));
1688 }
1689
1690 #[test]
1691 fn parse_screen_analyzer_response_accepts_markdown_wrapped_json() {
1692 let content = "```json\n{\"needsInteraction\":false,\"checkAgainWhen\":\"after_5s\"}\n```";
1693 assert_eq!(
1694 parse_screen_analyzer_response(content).check_again_when,
1695 "after_5s"
1696 );
1697 assert_eq!(
1698 parse_screen_analyzer_response("{\"needsInteraction\":false}").check_again_when,
1699 "content_changed"
1700 );
1701 }
1702
1703 #[test]
1704 fn parse_screen_analyzer_response_builds_tui_prompt_keys() {
1705 let content = r#"{
1706 "needsInteraction": true,
1707 "description": "pick one",
1708 "multiSelect": false,
1709 "confirmKey": "Enter",
1710 "options": [
1711 { "label": "1", "text": "alpha", "type": "select", "index": 0 },
1712 { "label": "2", "text": "beta", "type": "confirm", "index": 1 }
1713 ],
1714 "checkAgainWhen": "content_changed"
1715 }"#;
1716 let parsed = parse_screen_analyzer_response(content);
1717 assert!(parsed.needs_interaction);
1718 assert_eq!(parsed.description.as_deref(), Some("pick one"));
1719 assert_eq!(parsed.options.len(), 2);
1720 assert_eq!(parsed.options[0].keys, vec!["Enter"]);
1721 assert_eq!(parsed.options[1].keys, vec!["Down", "Enter"]);
1722 }
1723}
1724
1725pub async fn run_from_init_path(path: &Path) -> Result<()> {
1726 let payload = tokio::fs::read_to_string(path).await?;
1727 let init = serde_json::from_str::<InitConfig>(&payload)?;
1728 run(init).await
1729}