use std::collections::VecDeque;
use std::fs::{File, read_dir};
use std::io::Read;
use std::io::Seek;
use std::io::SeekFrom;
use std::path::{Path, PathBuf};
use std::time::Duration;
use anyhow::Result;
use async_trait::async_trait;
use beam_core::{FinalOutputKind, InitConfig};
use serde_json::Value;
use crate::adapter::{
Adapter, PendingTurnKind, PollResult, ResolveOutcome, SpawnSpec, SubmitResult,
TranscriptCursor, file_size, is_uuid_like, normalize_history_text,
};
use crate::backend::SessionBackend;
use crate::composer::{CODEX_COMPOSER, confirm_typed_submit, sample_draft_fgs, screen_looks_busy};
const COMPOSER_WAIT_ATTEMPTS: usize = 20;
const COMPOSER_WAIT_INTERVAL: Duration = Duration::from_millis(500);
async fn wait_for_codex_composer(backend: &dyn SessionBackend) -> bool {
for attempt in 0..COMPOSER_WAIT_ATTEMPTS {
let screen = backend.capture_viewport().await.unwrap_or_default();
if screen.contains('›') {
tracing::debug!(attempt, "codex composer visible");
return true;
}
tokio::time::sleep(COMPOSER_WAIT_INTERVAL).await;
}
tracing::debug!("codex composer not visible after wait");
false
}
#[derive(Debug, Clone)]
pub(crate) struct CodexState {
home_dir: PathBuf,
history_path: PathBuf,
rollout_path: Option<PathBuf>,
cli_pid: Option<u32>,
cli_session_id: Option<String>,
cursor: TranscriptCursor,
adopt_mode: bool,
adopt_restored_from_metadata: bool,
adopt_preamble_emitted: bool,
pending_remote_user_inputs: VecDeque<String>,
active_turn: Option<PendingTurnKind>,
}
fn state_from_init(init: &InitConfig) -> CodexState {
let codex_home = PathBuf::from(
std::env::var("CODEX_HOME")
.unwrap_or_else(|_| format!("{}/.codex", std::env::var("HOME").unwrap_or_default())),
);
create_state_with_paths(init, codex_home.join("history.jsonl"), codex_home)
}
fn traex_state_from_init(init: &InitConfig) -> CodexState {
let home = std::env::var("HOME").unwrap_or_default();
let (history_path, home_dir) = traex_paths(Path::new(&home));
create_state_with_paths(init, history_path, home_dir)
}
pub fn create(init: &InitConfig) -> Box<dyn Adapter> {
Box::new(state_from_init(init))
}
pub fn create_traex(init: &InitConfig) -> Box<dyn Adapter> {
Box::new(traex_state_from_init(init))
}
fn traex_paths(home: &Path) -> (PathBuf, PathBuf) {
(home.join(".trae/cli/history.json"), home.join(".traex/cli"))
}
fn create_state_with_paths(
init: &InitConfig,
history_path: PathBuf,
home_dir: PathBuf,
) -> CodexState {
CodexState {
history_path,
home_dir,
rollout_path: None,
cli_pid: None,
cli_session_id: init.cli_session_id.clone(),
cursor: TranscriptCursor::new(),
adopt_mode: init.adopted_from.is_some(),
adopt_restored_from_metadata: init.adopt_restored_from_metadata,
adopt_preamble_emitted: false,
pending_remote_user_inputs: VecDeque::new(),
active_turn: None,
}
}
#[async_trait]
impl Adapter for CodexState {
fn build_spawn_spec(&self, init: &InitConfig) -> SpawnSpec {
let mut args = Vec::new();
if init.resume
&& let Some(cli_session_id) = init.cli_session_id.clone().or_else(|| {
latest_codex_session_for_beam_session(&self.history_path, &init.session_id)
})
{
args.push("resume".to_string());
args.push(cli_session_id);
}
args.push("-C".to_string());
args.push(init.working_dir.clone());
args.extend(init.cli_args.clone());
SpawnSpec {
bin: init.cli_bin.clone(),
args,
}
}
async fn write_input(
&mut self,
backend: &dyn SessionBackend,
content: &str,
) -> Result<SubmitResult> {
if self.adopt_mode {
self.pending_remote_user_inputs
.push_back(normalize_history_text(content));
}
if !wait_for_codex_composer(backend).await {
tracing::debug!("codex composer not visible, refusing to type");
return Ok(SubmitResult {
submitted: false,
cli_session_id: self.cli_session_id.clone(),
failure_reason: Some("Codex TUI is not ready".to_string()),
});
}
let history_boundary = capture_history_boundary(&self.history_path)?;
backend.paste_text(content).await?;
tokio::time::sleep(Duration::from_millis(200)).await;
let typed_screen = backend.capture_viewport().await.unwrap_or_default();
let draft_fgs = sample_draft_fgs(&typed_screen, CODEX_COMPOSER);
let queue = screen_looks_busy(&typed_screen);
let submit_via = if queue { "tab" } else { "enter" };
tracing::debug!(busy = queue, submit_via, "codex submit key");
if queue {
backend.send_special_keys(&["Tab".to_string()]).await?;
} else {
backend.send_enter().await?;
}
let mut confirmed_session_id: Option<String> = None;
confirm_typed_submit(
backend,
CODEX_COMPOSER,
&draft_fgs,
submit_via,
|| {
let found = codex_history_match(&self.history_path, &history_boundary, content)?;
if let Some(cli_session_id) = found {
confirmed_session_id = Some(cli_session_id);
return Ok(true);
}
Ok(false)
},
|| async {
if queue {
backend.send_special_keys(&["Tab".to_string()]).await
} else {
backend.send_enter().await
}
},
)
.await?;
if let Some(cli_session_id) = confirmed_session_id {
self.cli_session_id = Some(cli_session_id.clone());
return Ok(SubmitResult {
submitted: true,
cli_session_id: Some(cli_session_id),
..Default::default()
});
}
Ok(SubmitResult {
submitted: false,
cli_session_id: self.cli_session_id.clone(),
failure_reason: Some("Codex did not accept the input".to_string()),
})
}
fn poll(&mut self) -> Result<PollResult> {
if self.rollout_path.is_none() {
if let Some(cli_session_id) = self.cli_session_id.clone() {
self.rollout_path =
find_codex_rollout_by_session_id(&self.home_dir, &cli_session_id);
}
if self.rollout_path.is_none()
&& let Some(pid) = self.cli_pid
&& let ResolveOutcome::Found((path, cli_session_id)) =
find_codex_rollout_by_pid(pid, &self.home_dir)
{
self.rollout_path = Some(path);
self.cli_session_id = Some(cli_session_id);
}
}
let mut result = PollResult {
cli_session_id: self.cli_session_id.clone(),
final_output: None,
final_output_kind: None,
final_output_user_text: None,
adopt_preamble: None,
prompt_ready: false,
};
let Some(path) = self.rollout_path.clone() else {
return Ok(result);
};
if self.adopt_mode && !self.adopt_preamble_emitted {
if !self.adopt_restored_from_metadata {
result.adopt_preamble = baseline_codex_adopt_preamble(&path)?;
}
self.cursor.skip_to(file_size(&path));
self.adopt_preamble_emitted = true;
return Ok(result);
}
let lines = self.cursor.drain(&path)?;
for line in &lines {
let Ok(value) = serde_json::from_str::<Value>(line) else {
continue;
};
if value.get("type").and_then(Value::as_str) != Some("response_item") {
continue;
}
let payload = value.get("payload").unwrap_or(&Value::Null);
if payload.get("type").and_then(Value::as_str) != Some("message") {
continue;
}
if let Some(role) = payload.get("role").and_then(Value::as_str) {
match role {
"user" if self.adopt_mode => {
let text = extract_codex_message_text(payload.get("content"));
if !text.trim().is_empty() {
let normalized = normalize_history_text(&text);
let kind = if self
.pending_remote_user_inputs
.front()
.map(|expected| *expected == normalized)
.unwrap_or(false)
{
let _ = self.pending_remote_user_inputs.pop_front();
PendingTurnKind::Remote
} else {
PendingTurnKind::Local { user_text: text }
};
self.active_turn = Some(kind);
}
}
"assistant"
if payload.get("phase").and_then(Value::as_str) == Some("final_answer") =>
{
let text = extract_codex_text(payload.get("content"), "output_text");
if let Some(text) = self.cursor.emit_if_new(&text) {
let kind = self.active_turn.take().or({
if self.adopt_mode {
Some(PendingTurnKind::LocalHeadless)
} else {
None
}
});
result.final_output = Some(text);
match kind {
Some(PendingTurnKind::Local { user_text }) => {
result.final_output_kind = Some(FinalOutputKind::LocalTurn);
result.final_output_user_text = Some(user_text);
}
Some(PendingTurnKind::LocalHeadless) => {
result.final_output_kind =
Some(FinalOutputKind::LocalTurnHeadless);
}
_ => {}
}
result.prompt_ready = true;
}
}
_ => {}
}
}
}
Ok(result)
}
fn on_spawned(&mut self, child_pid: Option<u32>) {
self.cli_pid = child_pid;
}
}
fn latest_codex_session_for_beam_session(
history_path: &Path,
beam_session_id: &str,
) -> Option<String> {
let entries = read_history_entries(history_path).ok()?;
for value in entries.iter().rev() {
let Some(text) = history_text(value) else {
continue;
};
if !text.contains(beam_session_id) {
continue;
}
if let Some(session_id) = history_session_id(value) {
return Some(session_id.to_string());
}
}
None
}
#[derive(Debug, Clone)]
enum HistoryBoundary {
Byte(u64),
DocumentEntries(Vec<Value>),
}
fn capture_history_boundary(history_path: &Path) -> Result<HistoryBoundary> {
let raw = match std::fs::read_to_string(history_path) {
Ok(raw) => raw,
Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
return Ok(
if history_path
.extension()
.is_some_and(|extension| extension == "json")
{
HistoryBoundary::DocumentEntries(Vec::new())
} else {
HistoryBoundary::Byte(0)
},
);
}
Err(error) => return Err(error.into()),
};
Ok(parse_history_document(&raw)
.map(HistoryBoundary::DocumentEntries)
.unwrap_or_else(|| HistoryBoundary::Byte(raw.len() as u64)))
}
fn codex_history_match(
history_path: &Path,
boundary: &HistoryBoundary,
expected_text: &str,
) -> Result<Option<String>> {
if !history_path.exists() {
return Ok(None);
}
let expected = normalize_history_text(expected_text);
for value in read_recent_history_entries(history_path, boundary)?
.into_iter()
.rev()
{
let Some(actual) = history_text(&value) else {
continue;
};
if normalize_history_text(actual) == expected {
return Ok(history_session_id(&value).map(ToOwned::to_owned));
}
}
Ok(None)
}
fn read_recent_history_entries(
history_path: &Path,
boundary: &HistoryBoundary,
) -> Result<Vec<Value>> {
let raw = std::fs::read_to_string(history_path)?;
if let Some(entries) = parse_history_document(&raw) {
return match boundary {
HistoryBoundary::DocumentEntries(previous) if entries.starts_with(previous) => {
Ok(entries[previous.len()..].to_vec())
}
HistoryBoundary::DocumentEntries(_) => Ok(Vec::new()),
HistoryBoundary::Byte(_) => Ok(Vec::new()),
};
}
let HistoryBoundary::Byte(from_byte) = boundary else {
return Ok(Vec::new());
};
let size = file_size(history_path);
if size <= *from_byte {
return Ok(Vec::new());
}
let mut file = File::open(history_path)?;
file.seek(SeekFrom::Start(*from_byte))?;
let mut text = String::new();
file.read_to_string(&mut text)?;
Ok(parse_history_jsonl(&text))
}
fn read_history_entries(history_path: &Path) -> Result<Vec<Value>> {
let raw = std::fs::read_to_string(history_path)?;
if let Some(entries) = parse_history_document(&raw) {
return Ok(entries);
}
Ok(parse_history_jsonl(&raw))
}
fn parse_history_document(raw: &str) -> Option<Vec<Value>> {
let value = serde_json::from_str::<Value>(raw).ok()?;
match value {
Value::Array(items) => Some(items),
Value::Object(map) => {
for key in ["history", "entries", "items", "data"] {
if let Some(Value::Array(items)) = map.get(key) {
return Some(items.clone());
}
}
if map.contains_key("text")
|| map.contains_key("session_id")
|| map.contains_key("sessionId")
|| map.contains_key("message")
{
return Some(vec![Value::Object(map)]);
}
None
}
_ => None,
}
}
fn parse_history_jsonl(raw: &str) -> Vec<Value> {
raw.lines()
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
.collect()
}
fn history_text(value: &Value) -> Option<&str> {
value
.get("text")
.and_then(Value::as_str)
.or_else(|| value.get("message").and_then(Value::as_str))
.or_else(|| value.get("content").and_then(Value::as_str))
}
fn history_session_id(value: &Value) -> Option<&str> {
value
.get("session_id")
.and_then(Value::as_str)
.or_else(|| value.get("sessionId").and_then(Value::as_str))
}
fn extract_codex_text(content: Option<&Value>, block_type: &str) -> String {
let Some(items) = content.and_then(Value::as_array) else {
return String::new();
};
items
.iter()
.filter_map(|block| {
if block.get("type").and_then(Value::as_str) == Some(block_type) {
block
.get("text")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
} else {
None
}
})
.collect::<Vec<_>>()
.join("")
}
fn extract_codex_message_text(content: Option<&Value>) -> String {
let Some(items) = content.and_then(Value::as_array) else {
return String::new();
};
items
.iter()
.filter_map(|block| {
block
.get("text")
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
.collect::<Vec<_>>()
.join("")
}
fn baseline_codex_adopt_preamble(path: &Path) -> Result<Option<(String, String)>> {
if !path.exists() {
return Ok(None);
}
let raw = std::fs::read_to_string(path)?;
let mut pending_user: Option<String> = None;
let mut latest_pair: Option<(String, String)> = None;
for line in raw.lines() {
let Ok(value) = serde_json::from_str::<Value>(line) else {
continue;
};
if value.get("type").and_then(Value::as_str) != Some("response_item") {
continue;
}
let payload = value.get("payload").unwrap_or(&Value::Null);
if payload.get("type").and_then(Value::as_str) != Some("message") {
continue;
}
let Some(role) = payload.get("role").and_then(Value::as_str) else {
continue;
};
match role {
"user" => {
let text = extract_codex_message_text(payload.get("content"));
if !text.trim().is_empty() {
pending_user = Some(text);
}
}
"assistant" if payload.get("phase").and_then(Value::as_str) == Some("final_answer") => {
let text = extract_codex_text(payload.get("content"), "output_text");
if !text.trim().is_empty()
&& let Some(user_text) = pending_user.take()
{
latest_pair = Some((user_text, text));
}
}
_ => {}
}
}
Ok(latest_pair)
}
fn find_codex_rollout_by_session_id(home_dir: &Path, cli_session_id: &str) -> Option<PathBuf> {
let root = home_dir.join("sessions");
if !root.exists() {
return None;
}
let suffix = format!("-{}.jsonl", cli_session_id);
let mut stack = vec![root];
while let Some(dir) = stack.pop() {
let Ok(entries) = read_dir(dir) else {
continue;
};
for entry in entries.flatten() {
let path = entry.path();
let Ok(meta) = entry.metadata() else {
continue;
};
if meta.is_dir() {
stack.push(path);
} else if meta.is_file()
&& path
.file_name()
.and_then(|s| s.to_str())
.map(|name| name.ends_with(&suffix))
.unwrap_or(false)
{
return Some(path);
}
}
}
None
}
fn find_codex_rollout_by_pid(pid: u32, home_dir: &Path) -> ResolveOutcome<(PathBuf, String)> {
find_codex_rollout_by_fd_dir(
&PathBuf::from(format!("/proc/{}/fd", pid)),
&home_dir.join("sessions"),
)
}
fn find_codex_rollout_by_fd_dir(
fd_dir: &Path,
sessions_root: &Path,
) -> ResolveOutcome<(PathBuf, String)> {
let entries = match read_dir(fd_dir) {
Ok(e) => e,
Err(e) => {
return ResolveOutcome::NotFound {
reason: format!("cannot read /proc fd dir: {}", e),
};
}
};
for entry in entries.flatten() {
let target = match std::fs::read_link(entry.path()) {
Ok(t) => t,
Err(_) => continue,
};
if !target
.extension()
.and_then(|ext| ext.to_str())
.map(|ext| ext == "jsonl")
.unwrap_or(false)
|| !target.starts_with(sessions_root)
{
continue;
}
let target_str = target.to_string_lossy();
if let Some(session_id) = codex_session_id_from_rollout_path(&target_str) {
return ResolveOutcome::Found((target, session_id));
}
}
ResolveOutcome::NotFound {
reason: "no codex rollout found in pid fd dir".to_string(),
}
}
fn codex_session_id_from_rollout_path(path: &str) -> Option<String> {
let base = Path::new(path).file_name()?.to_str()?;
if !base.starts_with("rollout-") || !base.ends_with(".jsonl") {
return None;
}
let trimmed = base.strip_suffix(".jsonl")?;
let tail = trimmed.rsplit_once('-')?.1;
if is_uuid_like(tail) {
return Some(tail.to_string());
}
if trimmed.len() >= 36 {
let candidate = &trimmed[trimmed.len() - 36..];
if is_uuid_like(candidate) {
return Some(candidate.to_string());
}
}
None
}
#[cfg(test)]
#[path = "codex_tests.rs"]
mod tests;