Skip to main content

beam_worker/
lib.rs

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}