fix(ai): Codex CLI provider reliability fixes for DBX Desktop

* fix(ai): resolve Codex CLI from shell PATH on macOS

* fix(ai): hide Codex CLI subprocess windows on Windows

* feat(ai): allow to set environment variable for Codex CLI

* fix(ai): parse Codex MCP tool events

* fix(mcp): honor Postgres TLS settings in direct queries
This commit is contained in:
Oleksandr Liakhevych authored and GitHub committed 2026-06-25 22:31:14 +08:00
1 parent 2759480331
commit 034cce47e0
17 files changed
+613 -30

No files matched your search

+14
View File
@@ -120,6 +120,8 @@ pub struct AiConfig {
pub context_window: Option<u32>,
#[serde(default)]
pub codex_cli_path: Option<String>,
#[serde(default)]
pub codex_cli_env: HashMap<String, String>,
}
fn default_enable_thinking() -> bool {
@@ -2118,6 +2120,7 @@ mod tests {
assert_eq!(config.proxy_url, "");
assert!(config.enable_thinking);
assert_eq!(config.auth_method, AiAuthMethod::ApiKey);
assert!(config.codex_cli_env.is_empty());
}
#[test]
@@ -2135,6 +2138,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
let err = build_ai_http_client(&config, 1).unwrap_err();
@@ -2157,6 +2161,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
build_ai_http_client(&config, 1).unwrap();
@@ -2177,6 +2182,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
build_ai_http_client(&config, 1).unwrap();
@@ -2197,6 +2203,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert_eq!(
@@ -2217,6 +2224,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert_eq!(resolve_endpoint(&ollama), "http://localhost:11434/v1/chat/completions");
@@ -2238,6 +2246,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert_eq!(resolve_model_list_endpoint(&openai).unwrap(), "https://api.openai.com/v1/models");
@@ -2254,6 +2263,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert_eq!(resolve_model_list_endpoint(&claude).unwrap(), "https://api.anthropic.com/v1/models");
}
@@ -2274,6 +2284,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert_eq!(resolve_endpoint(&config), "https://api.example.com/v1/chat/completions");
assert_eq!(resolve_model_list_endpoint(&config).unwrap(), "https://api.example.com/v1/models");
@@ -2329,6 +2340,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
let api_key_headers = claude_headers(&config).unwrap();
@@ -2400,6 +2412,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
assert!(!supports_temperature(&config));
@@ -2466,6 +2479,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
};
let mut body = serde_json::json!({
"model": &config.model,
+55 -3
View File
@@ -2,11 +2,15 @@ use crate::agent_events::AgentEvent;
use crate::ai::{AiMessage, AiModelInfo};
use crate::token_usage::TokenUsage;
use serde_json::Value;
use std::ffi::OsStr;
use std::process::Stdio;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, BufReader};
use tokio::process::Command;
use tokio::sync::Notify;
#[cfg(windows)]
const CREATE_NO_WINDOW: u32 = 0x08000000;
#[derive(Debug, Clone)]
pub struct CliAgentRunOptions {
pub connection_id: String,
@@ -96,12 +100,28 @@ pub fn model_infos(ids: &[&str]) -> Vec<AiModelInfo> {
ids.iter().map(|id| AiModelInfo { id: (*id).to_string(), display_name: None }).collect()
}
pub fn cli_command(program: impl AsRef<OsStr>) -> Command {
let command = Command::new(program);
configure_cli_command(command)
}
#[cfg(windows)]
fn configure_cli_command(mut command: Command) -> Command {
command.creation_flags(CREATE_NO_WINDOW);
command
}
#[cfg(not(windows))]
fn configure_cli_command(command: Command) -> Command {
command
}
pub async fn list_json_models_or_default(
program: String,
args: impl IntoIterator<Item = String>,
default_models: &[&str],
) -> Result<Vec<AiModelInfo>, String> {
let output = Command::new(program).args(args).output().await;
let output = cli_command(program).args(args).output().await;
let Ok(output) = output else {
return Ok(model_infos(default_models));
@@ -285,6 +305,7 @@ fn cli_item_id(item: &Value) -> String {
fn cli_tool_name(item: &Value) -> String {
item.get("tool_name")
.and_then(Value::as_str)
.or_else(|| item.get("tool").and_then(Value::as_str))
.or_else(|| item.get("name").and_then(Value::as_str))
.or_else(|| item.get("server_tool_name").and_then(Value::as_str))
.unwrap_or("cli_tool")
@@ -300,7 +321,16 @@ fn cli_tool_args(item: &Value) -> Value {
}
fn cli_tool_result(item: &Value) -> Value {
item.get("result").or_else(|| item.get("output")).or_else(|| item.get("content")).cloned().unwrap_or(Value::Null)
non_null_field(item, "result")
.or_else(|| non_null_field(item, "output"))
.or_else(|| non_null_field(item, "content"))
.or_else(|| non_null_field(item, "error"))
.cloned()
.unwrap_or(Value::Null)
}
fn non_null_field<'a>(value: &'a Value, key: &str) -> Option<&'a Value> {
value.get(key).filter(|field| !field.is_null())
}
fn cli_text(item: &Value) -> Option<String> {
@@ -326,7 +356,8 @@ pub async fn run_cli_jsonl_agent(
cancelled: &Notify,
on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
) -> Result<String, String> {
let mut child = Command::new(&spec.command.program)
let mut command = cli_command(&spec.command.program);
let mut child = command
.args(&spec.command.args)
.envs(spec.env.iter().map(|(key, value)| (key.as_str(), value.as_str())))
.stdout(Stdio::piped())
@@ -427,6 +458,27 @@ mod tests {
.unwrap_or(false)
}
#[tokio::test]
async fn jsonl_agent_applies_env_to_child() {
let spec = CliAgentProcessSpec {
command: CliAgentCommandSpec {
program: "sh".to_string(),
args: vec![
"-c".to_string(),
"printf '%s\n' \"{\\\"type\\\":\\\"item.completed\\\",\\\"item\\\":{\\\"type\\\":\\\"agent_message\\\",\\\"text\\\":\\\"$DBX_TEST_ENV\\\"}}\" \"{\\\"type\\\":\\\"turn.completed\\\"}\"".to_string(),
],
},
env: vec![("DBX_TEST_ENV".to_string(), "from-env".to_string())],
dialect: CliAgentJsonlDialect::CodexExec,
classify_spawn_error,
classify_run_error,
};
let result = run_cli_jsonl_agent(spec, &Notify::new(), |_| {}).await.unwrap();
assert_eq!(result, "from-env");
}
#[tokio::test]
async fn jsonl_error_kills_and_waits_for_child() {
let pid_file = std::env::temp_dir().join(format!(
+171 -16
View File
@@ -1,12 +1,12 @@
use crate::agent_events::AgentEvent;
use crate::ai::{AiConfig, AiModelInfo, AiTestConnectionResult};
use crate::ai_cli_agent::{
append_config_overrides, build_cli_agent_prompt, dbx_mcp_enabled_tools, dbx_mcp_scope_env, model_infos,
parse_cli_jsonl_event, run_cli_jsonl_agent, toml_string, toml_string_array, CliAgentCommandSpec,
append_config_overrides, build_cli_agent_prompt, cli_command, dbx_mcp_enabled_tools, dbx_mcp_scope_env,
model_infos, parse_cli_jsonl_event, run_cli_jsonl_agent, toml_string, toml_string_array, CliAgentCommandSpec,
CliAgentJsonlDialect, CliAgentProcessSpec, CliAgentRunOptions,
};
use serde_json::Value;
use std::collections::BTreeSet;
use std::collections::{BTreeMap, BTreeSet};
use std::env;
use std::path::{Path, PathBuf};
use std::process::Stdio;
@@ -36,8 +36,13 @@ async fn resolve_codex_command(config: &AiConfig) -> CodexCommandSpec {
}
}
fn command_env(command: &CodexCommandSpec) -> Vec<(String, String)> {
command.parent_dir().map(|dir| ("PATH".to_string(), merged_path_with_dir(&dir))).into_iter().collect()
fn codex_process_env(config: &AiConfig, command: &CodexCommandSpec) -> Result<Vec<(String, String)>, String> {
let mut env = BTreeMap::from_iter(codex_cli_env(config)?);
if let Some(dir) = command.parent_dir() {
let user_path = env.get("PATH").map(String::as_str);
env.insert("PATH".to_string(), merged_path_with_dir(&dir, user_path));
}
Ok(env.into_iter().collect())
}
trait CommandParentDir {
@@ -167,9 +172,12 @@ fn common_executable_dirs() -> Vec<PathBuf> {
dirs
}
fn merged_path_with_dir(dir: &str) -> String {
fn merged_path_with_dir(dir: &str, user_path: Option<&str>) -> String {
let mut seen = BTreeSet::new();
let mut dirs = vec![PathBuf::from(dir)];
if let Some(path) = user_path {
dirs.extend(env::split_paths(path));
}
dirs.extend(common_executable_dirs());
let paths = dirs.into_iter().filter(|path| seen.insert(path.clone())).collect::<Vec<_>>();
env::join_paths(paths).unwrap_or_default().to_string_lossy().to_string()
@@ -227,6 +235,58 @@ fn shell_quote(value: &str) -> String {
format!("'{}'", value.replace('\'', "'\"'\"'"))
}
fn validate_codex_program(config: &AiConfig) -> Result<String, String> {
let program = codex_program(config);
if starts_with_env_assignment(&program) {
return Err("[codexCliPathInvalid] Codex CLI path should contain only the executable path. Add environment variables in the Codex CLI environment variables section.".to_string());
}
Ok(program)
}
pub fn codex_cli_env(config: &AiConfig) -> Result<Vec<(String, String)>, String> {
let mut env = BTreeMap::new();
for (key, value) in &config.codex_cli_env {
let key = key.trim();
if key.is_empty() {
continue;
}
if !is_env_var_name(key) {
return Err(format!(
"[codexEnvInvalid] Invalid Codex CLI environment variable name `{key}`. Use names like HTTPS_PROXY."
));
}
if is_reserved_dbx_mcp_env_name(key) {
return Err(format!(
"[codexEnvReserved] `{key}` is managed by DBX for the scoped MCP server and cannot be set here."
));
}
env.insert(key.to_string(), value.clone());
}
Ok(env.into_iter().collect())
}
fn starts_with_env_assignment(program: &str) -> bool {
let Some(first_token) = program.split_whitespace().next() else {
return false;
};
let Some((key, _)) = first_token.split_once('=') else {
return false;
};
is_env_var_name(key)
}
fn is_env_var_name(name: &str) -> bool {
let mut chars = name.chars();
let Some(first) = chars.next() else {
return false;
};
(first == '_' || first.is_ascii_alphabetic()) && chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}
fn is_reserved_dbx_mcp_env_name(name: &str) -> bool {
name.to_ascii_uppercase().starts_with("DBX_MCP_")
}
pub fn codex_enabled_tools(agent_mode: bool) -> Vec<&'static str> {
dbx_mcp_enabled_tools(agent_mode)
}
@@ -239,7 +299,7 @@ fn codex_mcp_config_overrides(options: &CodexRunOptions) -> Vec<String> {
"mcp_servers.dbx.required=true".to_string(),
"mcp_servers.dbx.startup_timeout_sec=20".to_string(),
"mcp_servers.dbx.tool_timeout_sec=120".to_string(),
"mcp_servers.dbx.default_tools_approval_mode=\"auto\"".to_string(),
"mcp_servers.dbx.default_tools_approval_mode=\"approve\"".to_string(),
format!("mcp_servers.dbx.enabled_tools={}", toml_string_array(&dbx_mcp_enabled_tools(options.agent_mode))),
];
if let Some(command) = options.mcp_server_command.as_ref().filter(|command| !command.args.is_empty()) {
@@ -284,10 +344,11 @@ pub fn build_codex_prompt(system_prompt: &str, messages: &[crate::ai::AiMessage]
}
pub async fn list_codex_models(config: &AiConfig) -> Result<Vec<AiModelInfo>, String> {
validate_codex_program(config)?;
let command = resolve_codex_command(config).await;
let output = Command::new(&command.program)
let output = cli_command(&command.program)
.args(["debug", "models"])
.envs(command_env(&command).iter().map(|(key, value)| (key.as_str(), value.as_str())))
.envs(codex_process_env(config, &command)?.iter().map(|(key, value)| (key.as_str(), value.as_str())))
.output()
.await;
@@ -335,10 +396,11 @@ fn parse_codex_models(stdout: &str) -> Option<Vec<AiModelInfo>> {
pub async fn test_codex_connection(config: &AiConfig) -> Result<AiTestConnectionResult, String> {
let start = Instant::now();
validate_codex_program(config)?;
let codex_command = resolve_codex_command(config).await;
let mut command = Command::new(&codex_command.program);
let mut command = cli_command(&codex_command.program);
command.args(["exec", "--json", "--skip-git-repo-check", "--sandbox", "read-only"]);
command.envs(command_env(&codex_command).iter().map(|(key, value)| (key.as_str(), value.as_str())));
command.envs(codex_process_env(config, &codex_command)?.iter().map(|(key, value)| (key.as_str(), value.as_str())));
let model = config.model.trim();
if !model.is_empty() && !model.eq_ignore_ascii_case("default") {
@@ -405,10 +467,11 @@ pub async fn run_codex_agent(
cancelled: &Notify,
on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
) -> Result<String, String> {
validate_codex_program(config)?;
let mut command = build_codex_exec_command(config, prompt, &options);
let resolved_command = resolve_codex_command(config).await;
command.program = resolved_command.program;
let env = command_env(&command);
let env = codex_process_env(config, &command)?;
run_cli_jsonl_agent(
CliAgentProcessSpec {
command,
@@ -428,8 +491,9 @@ mod tests {
#[cfg(not(windows))]
use super::shell_quote;
use super::{
build_codex_exec_command, codex_enabled_tools, command_env, common_executable_dirs, is_path_like_program,
merged_path_with_dir, parse_codex_jsonl_event, parse_codex_models, CodexRunOptions, DEFAULT_CODEX_MODELS,
build_codex_exec_command, codex_cli_env, codex_enabled_tools, codex_process_env, common_executable_dirs,
is_path_like_program, merged_path_with_dir, parse_codex_jsonl_event, parse_codex_models,
validate_codex_program, CodexRunOptions, DEFAULT_CODEX_MODELS,
};
use crate::agent_events::AgentEvent;
use crate::ai::{AiApiStyle, AiAuthMethod, AiConfig, AiProvider, AiReasoningLevel};
@@ -449,6 +513,7 @@ mod tests {
reasoning_level: AiReasoningLevel::Default,
context_window: None,
codex_cli_path: None,
codex_cli_env: Default::default(),
}
}
@@ -471,6 +536,7 @@ mod tests {
assert!(!spec.args.contains(&"--model".to_string()));
assert!(!spec.args.contains(&"--ask-for-approval".to_string()));
assert!(spec.args.contains(&"mcp_servers.dbx.command=\"dbx-mcp-server\"".to_string()));
assert!(spec.args.contains(&"mcp_servers.dbx.default_tools_approval_mode=\"approve\"".to_string()));
assert!(spec.args.contains(&"mcp_servers.dbx.env.DBX_MCP_ALLOW_WRITES=\"0\"".to_string()));
assert!(spec.args.contains(&"mcp_servers.dbx.env.DBX_MCP_SCOPE_CONNECTION_ID=\"conn-1\"".to_string()));
assert!(spec.args.iter().any(|arg| arg.contains("dbx_execute_query")));
@@ -533,8 +599,9 @@ mod tests {
#[test]
#[cfg(not(windows))]
fn codex_command_env_prepends_resolved_program_dir_and_keeps_node_dirs() {
let config = codex_config("default");
let command = CliAgentCommandSpec { program: "/opt/homebrew/bin/codex".to_string(), args: Vec::new() };
let env = command_env(&command);
let env = codex_process_env(&config, &command).unwrap();
let path = env.iter().find(|(key, _)| key == "PATH").map(|(_, value)| value).unwrap();
let dirs = std::env::split_paths(path).collect::<Vec<_>>();
@@ -545,13 +612,28 @@ mod tests {
#[test]
#[cfg(not(windows))]
fn merged_path_deduplicates_codex_dir() {
let path = merged_path_with_dir("/opt/homebrew/bin");
let path = merged_path_with_dir("/opt/homebrew/bin", None);
let dirs = std::env::split_paths(&path).collect::<Vec<_>>();
let count = dirs.iter().filter(|dir| *dir == std::path::Path::new("/opt/homebrew/bin")).count();
assert_eq!(count, 1);
}
#[test]
#[cfg(not(windows))]
fn codex_process_env_keeps_resolved_program_dir_before_user_path() {
let mut config = codex_config("default");
config.codex_cli_env.insert("PATH".to_string(), "/custom/bin:/usr/bin".to_string());
let command = CliAgentCommandSpec { program: "/opt/homebrew/bin/codex".to_string(), args: Vec::new() };
let env = codex_process_env(&config, &command).unwrap();
let path = env.iter().find(|(key, _)| key == "PATH").map(|(_, value)| value).unwrap();
let dirs = std::env::split_paths(path).collect::<Vec<_>>();
assert_eq!(dirs.first().unwrap(), std::path::Path::new("/opt/homebrew/bin"));
assert!(dirs.iter().any(|dir| dir == std::path::Path::new("/custom/bin")));
}
#[test]
#[cfg(not(windows))]
fn shell_quote_handles_codex_program_names() {
@@ -564,6 +646,55 @@ mod tests {
assert_eq!(model_infos(DEFAULT_CODEX_MODELS), model_infos(&["default", "gpt-5.5", "gpt-5.4-mini"]));
}
#[test]
fn normalizes_codex_cli_env() {
let mut config = codex_config("default");
config.codex_cli_env.insert(" HTTPS_PROXY ".to_string(), "http://proxy:9800".to_string());
config.codex_cli_env.insert("NO_PROXY".to_string(), "localhost,127.0.0.1".to_string());
config.codex_cli_env.insert("".to_string(), "ignored".to_string());
let env = codex_cli_env(&config).unwrap();
assert_eq!(
env,
vec![
("HTTPS_PROXY".to_string(), "http://proxy:9800".to_string()),
("NO_PROXY".to_string(), "localhost,127.0.0.1".to_string())
]
);
}
#[test]
fn rejects_invalid_codex_cli_env_name() {
let mut config = codex_config("default");
config.codex_cli_env.insert("BAD-NAME".to_string(), "value".to_string());
let err = codex_cli_env(&config).unwrap_err();
assert!(err.contains("[codexEnvInvalid]"));
}
#[test]
fn rejects_reserved_dbx_mcp_env_name() {
let mut config = codex_config("default");
config.codex_cli_env.insert("DBX_MCP_ALLOW_DANGEROUS_SQL".to_string(), "1".to_string());
let err = codex_cli_env(&config).unwrap_err();
assert!(err.contains("[codexEnvReserved]"));
}
#[test]
fn rejects_shell_style_env_prefix_in_codex_cli_path() {
let mut config = codex_config("default");
config.codex_cli_path = Some("HTTPS_PROXY=http://proxy:9800 /opt/homebrew/bin/codex".to_string());
let err = validate_codex_program(&config).unwrap_err();
assert!(err.contains("[codexCliPathInvalid]"));
assert!(err.contains("environment variables section"));
}
#[test]
fn parses_codex_model_catalog_without_service_tiers() {
let models = parse_codex_models(
@@ -603,6 +734,30 @@ mod tests {
matches!(&tool_done[0], AgentEvent::ToolCallEnd { tool_name, is_error, .. } if tool_name == "dbx_list_tables" && !is_error)
);
let current_started = parse_codex_jsonl_event(
r#"{"type":"item.started","item":{"id":"item_0","type":"mcp_tool_call","server":"dbx","tool":"dbx_execute_query","arguments":{"sql":"SELECT 1"},"result":null,"error":null,"status":"in_progress"}}"#,
)
.unwrap();
assert!(
matches!(&current_started[0], AgentEvent::ToolCallStart { tool_name, args, .. } if tool_name == "dbx_execute_query" && args["sql"] == "SELECT 1")
);
let current_tool_done = parse_codex_jsonl_event(
r#"{"type":"item.completed","item":{"id":"item_0","type":"mcp_tool_call","server":"dbx","tool":"dbx_execute_query","arguments":{"sql":"SELECT 1"},"result":{"content":[{"type":"text","text":"ok"}],"structured_content":null},"error":null,"status":"completed"}}"#,
)
.unwrap();
assert!(
matches!(&current_tool_done[0], AgentEvent::ToolCallEnd { tool_name, result, is_error, .. } if tool_name == "dbx_execute_query" && result["content"][0]["text"] == "ok" && !is_error)
);
let current_tool_error = parse_codex_jsonl_event(
r#"{"type":"item.completed","item":{"id":"item_0","type":"mcp_tool_call","server":"dbx","tool":"dbx_execute_query","arguments":{"sql":"SELECT 1"},"result":null,"error":{"message":"user cancelled MCP tool call"},"status":"failed"}}"#,
)
.unwrap();
assert!(
matches!(&current_tool_error[0], AgentEvent::ToolCallEnd { tool_name, result, is_error, .. } if tool_name == "dbx_execute_query" && result["message"] == "user cancelled MCP tool call" && *is_error)
);
let turn_done =
parse_codex_jsonl_event(r#"{"type":"turn.completed","usage":{"input_tokens":12,"output_tokens":3}}"#)
.unwrap();