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 {
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 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}