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