Skip to main content

beam_worker/
lib.rs

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}