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