mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
Merge origin/main into etraut/alarm
This commit is contained in:
@@ -86,6 +86,7 @@ fn transport_name(transport: AppServerTransport) -> &'static str {
|
||||
match transport {
|
||||
AppServerTransport::Stdio => "stdio",
|
||||
AppServerTransport::WebSocket { .. } => "websocket",
|
||||
AppServerTransport::Off => "off",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -20,6 +20,7 @@ use crate::thread_status::resolve_thread_status;
|
||||
use chrono::DateTime;
|
||||
use chrono::SecondsFormat;
|
||||
use chrono::Utc;
|
||||
use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_app_server_protocol::Account;
|
||||
use codex_app_server_protocol::AccountLoginCompletedNotification;
|
||||
use codex_app_server_protocol::AccountUpdatedNotification;
|
||||
@@ -29,6 +30,7 @@ use codex_app_server_protocol::AppsListParams;
|
||||
use codex_app_server_protocol::AppsListResponse;
|
||||
use codex_app_server_protocol::AskForApproval;
|
||||
use codex_app_server_protocol::AuthMode;
|
||||
use codex_app_server_protocol::AuthMode as CoreAuthMode;
|
||||
use codex_app_server_protocol::CancelLoginAccountParams;
|
||||
use codex_app_server_protocol::CancelLoginAccountResponse;
|
||||
use codex_app_server_protocol::CancelLoginAccountStatus;
|
||||
@@ -188,9 +190,7 @@ use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_backend_client::Client as BackendClient;
|
||||
use codex_chatgpt::connectors;
|
||||
use codex_cloud_requirements::cloud_requirements_loader;
|
||||
use codex_core::AnalyticsEventsClient;
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::CodexAuth;
|
||||
use codex_config::types::McpServerTransportConfig;
|
||||
use codex_core::CodexThread;
|
||||
use codex_core::Cursor as RolloutCursor;
|
||||
use codex_core::ForkSnapshot;
|
||||
@@ -201,23 +201,16 @@ use codex_core::SteerInputError;
|
||||
use codex_core::ThreadConfigSnapshot;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::ThreadSortKey as CoreThreadSortKey;
|
||||
use codex_core::auth::AuthMode as CoreAuthMode;
|
||||
use codex_core::auth::CLIENT_ID;
|
||||
use codex_core::auth::login_with_api_key;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigOverrides;
|
||||
use codex_core::config::NetworkProxyAuditMetadata;
|
||||
use codex_core::config::edit::ConfigEdit;
|
||||
use codex_core::config::edit::ConfigEditsBuilder;
|
||||
use codex_core::config::types::McpServerTransportConfig;
|
||||
use codex_core::config_loader::CloudRequirementsLoadError;
|
||||
use codex_core::config_loader::CloudRequirementsLoadErrorCode;
|
||||
use codex_core::config_loader::CloudRequirementsLoader;
|
||||
use codex_core::config_loader::LoaderOverrides;
|
||||
use codex_core::config_loader::load_config_layers_state;
|
||||
use codex_core::default_client::set_default_client_residency_requirement;
|
||||
use codex_core::error::CodexErr;
|
||||
use codex_core::error::Result as CodexResult;
|
||||
use codex_core::exec::ExecCapturePolicy;
|
||||
use codex_core::exec::ExecExpiration;
|
||||
use codex_core::exec::ExecParams;
|
||||
@@ -226,11 +219,6 @@ use codex_core::find_archived_thread_path_by_id_str;
|
||||
use codex_core::find_thread_name_by_id;
|
||||
use codex_core::find_thread_names_by_ids;
|
||||
use codex_core::find_thread_path_by_id_str;
|
||||
use codex_core::mcp::auth::discover_supported_scopes;
|
||||
use codex_core::mcp::auth::resolve_oauth_scopes;
|
||||
use codex_core::mcp::collect_mcp_snapshot;
|
||||
use codex_core::mcp::group_tools_by_server;
|
||||
use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig;
|
||||
use codex_core::parse_cursor;
|
||||
use codex_core::plugins::MarketplaceError;
|
||||
use codex_core::plugins::MarketplacePluginSource;
|
||||
@@ -245,9 +233,6 @@ use codex_core::read_head_for_summary;
|
||||
use codex_core::read_session_meta_line;
|
||||
use codex_core::rollout_date_parts;
|
||||
use codex_core::sandboxing::SandboxPermissions;
|
||||
use codex_core::state_db::StateDbHandle;
|
||||
use codex_core::state_db::get_state_db;
|
||||
use codex_core::state_db::reconcile_rollout;
|
||||
use codex_core::windows_sandbox::WindowsSandboxLevelExt;
|
||||
use codex_core::windows_sandbox::WindowsSandboxSetupMode as CoreWindowsSandboxSetupMode;
|
||||
use codex_core::windows_sandbox::WindowsSandboxSetupRequest;
|
||||
@@ -256,18 +241,33 @@ use codex_features::Feature;
|
||||
use codex_features::Stage;
|
||||
use codex_feedback::CodexFeedback;
|
||||
use codex_git_utils::git_diff_to_remote;
|
||||
use codex_git_utils::resolve_root_git_project_for_trust;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CLIENT_ID;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_login::ServerOptions as LoginServerOptions;
|
||||
use codex_login::ShutdownHandle;
|
||||
use codex_login::auth::login_with_chatgpt_auth_tokens;
|
||||
use codex_login::complete_device_code_login;
|
||||
use codex_login::default_client::set_default_client_residency_requirement;
|
||||
use codex_login::login_with_api_key;
|
||||
use codex_login::request_device_code;
|
||||
use codex_login::run_login_server;
|
||||
use codex_mcp::mcp::auth::discover_supported_scopes;
|
||||
use codex_mcp::mcp::auth::resolve_oauth_scopes;
|
||||
use codex_mcp::mcp::collect_mcp_snapshot;
|
||||
use codex_mcp::mcp::effective_mcp_servers;
|
||||
use codex_mcp::mcp::qualified_mcp_tool_name_prefix;
|
||||
use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::config_types::CollaborationMode;
|
||||
use codex_protocol::config_types::ForcedLoginMethod;
|
||||
use codex_protocol::config_types::Personality;
|
||||
use codex_protocol::config_types::TrustLevel;
|
||||
use codex_protocol::config_types::WindowsSandboxLevel;
|
||||
use codex_protocol::dynamic_tools::DynamicToolSpec as CoreDynamicToolSpec;
|
||||
use codex_protocol::error::CodexErr;
|
||||
use codex_protocol::error::Result as CodexResult;
|
||||
use codex_protocol::items::TurnItem;
|
||||
use codex_protocol::models::ResponseItem;
|
||||
use codex_protocol::protocol::AgentStatus;
|
||||
@@ -292,6 +292,9 @@ use codex_protocol::protocol::W3cTraceContext;
|
||||
use codex_protocol::user_input::MAX_USER_INPUT_TEXT_CHARS;
|
||||
use codex_protocol::user_input::UserInput as CoreInputItem;
|
||||
use codex_rmcp_client::perform_oauth_login_return_url;
|
||||
use codex_rollout::state_db::StateDbHandle;
|
||||
use codex_rollout::state_db::get_state_db;
|
||||
use codex_rollout::state_db::reconcile_rollout;
|
||||
use codex_state::StateRuntime;
|
||||
use codex_state::ThreadMetadata;
|
||||
use codex_state::ThreadMetadataBuilder;
|
||||
@@ -696,6 +699,7 @@ impl CodexMessageProcessor {
|
||||
connection_id: ConnectionId,
|
||||
request: ClientRequest,
|
||||
app_server_client_name: Option<String>,
|
||||
app_server_client_version: Option<String>,
|
||||
request_context: RequestContext,
|
||||
) {
|
||||
let to_connection_request_id = |request_id| ConnectionRequestId {
|
||||
@@ -712,6 +716,8 @@ impl CodexMessageProcessor {
|
||||
self.thread_start(
|
||||
to_connection_request_id(request_id),
|
||||
params,
|
||||
app_server_client_name.clone(),
|
||||
app_server_client_version.clone(),
|
||||
request_context,
|
||||
)
|
||||
.await;
|
||||
@@ -828,6 +834,7 @@ impl CodexMessageProcessor {
|
||||
to_connection_request_id(request_id),
|
||||
params,
|
||||
app_server_client_name.clone(),
|
||||
app_server_client_version.clone(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -2074,6 +2081,8 @@ impl CodexMessageProcessor {
|
||||
&self,
|
||||
request_id: ConnectionRequestId,
|
||||
params: ThreadStartParams,
|
||||
app_server_client_name: Option<String>,
|
||||
app_server_client_version: Option<String>,
|
||||
request_context: RequestContext,
|
||||
) {
|
||||
let ThreadStartParams {
|
||||
@@ -2129,6 +2138,8 @@ impl CodexMessageProcessor {
|
||||
runtime_feature_enablement,
|
||||
cloud_requirements,
|
||||
request_id,
|
||||
app_server_client_name,
|
||||
app_server_client_version,
|
||||
config,
|
||||
typesafe_overrides,
|
||||
dynamic_tools,
|
||||
@@ -2202,6 +2213,8 @@ impl CodexMessageProcessor {
|
||||
runtime_feature_enablement: BTreeMap<String, bool>,
|
||||
cloud_requirements: CloudRequirementsLoader,
|
||||
request_id: ConnectionRequestId,
|
||||
app_server_client_name: Option<String>,
|
||||
app_server_client_version: Option<String>,
|
||||
config_overrides: Option<HashMap<String, serde_json::Value>>,
|
||||
typesafe_overrides: ConfigOverrides,
|
||||
dynamic_tools: Option<Vec<ApiDynamicToolSpec>>,
|
||||
@@ -2210,10 +2223,11 @@ impl CodexMessageProcessor {
|
||||
experimental_raw_events: bool,
|
||||
request_trace: Option<W3cTraceContext>,
|
||||
) {
|
||||
let config = match derive_config_from_params(
|
||||
let requested_cwd = typesafe_overrides.cwd.clone();
|
||||
let mut config = match derive_config_from_params(
|
||||
&cli_overrides,
|
||||
config_overrides,
|
||||
typesafe_overrides,
|
||||
config_overrides.clone(),
|
||||
typesafe_overrides.clone(),
|
||||
&cloud_requirements,
|
||||
&listener_task_context.codex_home,
|
||||
&runtime_feature_enablement,
|
||||
@@ -2231,6 +2245,70 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
};
|
||||
|
||||
// The user may have requested WorkspaceWrite or DangerFullAccess via
|
||||
// the command line, though in the process of deriving the Config, it
|
||||
// could be downgraded to ReadOnly (perhaps there is no sandbox
|
||||
// available on Windows or the enterprise config disallows it). The cwd
|
||||
// should still be considered "trusted" in this case.
|
||||
let requested_sandbox_trusts_project = matches!(
|
||||
typesafe_overrides.sandbox_mode,
|
||||
Some(
|
||||
codex_protocol::config_types::SandboxMode::WorkspaceWrite
|
||||
| codex_protocol::config_types::SandboxMode::DangerFullAccess
|
||||
)
|
||||
);
|
||||
|
||||
if requested_cwd.is_some()
|
||||
&& !config.active_project.is_trusted()
|
||||
&& (requested_sandbox_trusts_project
|
||||
|| matches!(
|
||||
config.permissions.sandbox_policy.get(),
|
||||
codex_protocol::protocol::SandboxPolicy::WorkspaceWrite { .. }
|
||||
| codex_protocol::protocol::SandboxPolicy::DangerFullAccess
|
||||
| codex_protocol::protocol::SandboxPolicy::ExternalSandbox { .. }
|
||||
))
|
||||
{
|
||||
let trust_target = resolve_root_git_project_for_trust(config.cwd.as_path())
|
||||
.unwrap_or_else(|| config.cwd.to_path_buf());
|
||||
if let Err(err) = codex_core::config::set_project_trust_level(
|
||||
&listener_task_context.codex_home,
|
||||
trust_target.as_path(),
|
||||
TrustLevel::Trusted,
|
||||
) {
|
||||
let error = JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to persist trusted project state: {err}"),
|
||||
data: None,
|
||||
};
|
||||
listener_task_context
|
||||
.outgoing
|
||||
.send_error(request_id, error)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
|
||||
config = match derive_config_from_params(
|
||||
&cli_overrides,
|
||||
config_overrides,
|
||||
typesafe_overrides,
|
||||
&cloud_requirements,
|
||||
&listener_task_context.codex_home,
|
||||
&runtime_feature_enablement,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(config) => config,
|
||||
Err(err) => {
|
||||
let error = config_load_error(&err);
|
||||
listener_task_context
|
||||
.outgoing
|
||||
.send_error(request_id, error)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
let dynamic_tools = dynamic_tools.unwrap_or_default();
|
||||
let core_dynamic_tools = if dynamic_tools.is_empty() {
|
||||
Vec::new()
|
||||
@@ -2283,6 +2361,19 @@ impl CodexMessageProcessor {
|
||||
session_configured,
|
||||
..
|
||||
} = new_conv;
|
||||
if let Err(error) = Self::set_app_server_client_info(
|
||||
thread.as_ref(),
|
||||
app_server_client_name,
|
||||
app_server_client_version,
|
||||
)
|
||||
.await
|
||||
{
|
||||
listener_task_context
|
||||
.outgoing
|
||||
.send_error(request_id, error)
|
||||
.await;
|
||||
return;
|
||||
}
|
||||
let config_snapshot = thread
|
||||
.config_snapshot()
|
||||
.instrument(tracing::info_span!(
|
||||
@@ -3316,6 +3407,13 @@ impl CodexMessageProcessor {
|
||||
cwd,
|
||||
search_term,
|
||||
} = params;
|
||||
let cwd = match normalize_thread_list_cwd_filter(cwd) {
|
||||
Ok(cwd) => cwd,
|
||||
Err(error) => {
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let requested_page_size = limit
|
||||
.map(|value| value as usize)
|
||||
@@ -3334,7 +3432,7 @@ impl CodexMessageProcessor {
|
||||
model_providers,
|
||||
source_kinds,
|
||||
archived: archived.unwrap_or(false),
|
||||
cwd: cwd.map(PathBuf::from),
|
||||
cwd,
|
||||
search_term,
|
||||
},
|
||||
)
|
||||
@@ -3543,6 +3641,11 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
build_thread_from_snapshot(thread_uuid, &config_snapshot, loaded_rollout_path)
|
||||
};
|
||||
if thread.forked_from_id.is_none()
|
||||
&& let Some(rollout_path) = rollout_path.as_ref()
|
||||
{
|
||||
thread.forked_from_id = forked_from_id_from_rollout(rollout_path).await;
|
||||
}
|
||||
self.attach_thread_name(thread_uuid, &mut thread).await;
|
||||
|
||||
if include_turns && let Some(rollout_path) = rollout_path.as_ref() {
|
||||
@@ -4474,7 +4577,12 @@ impl CodexMessageProcessor {
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(summary) => summary_to_thread(summary),
|
||||
Ok(summary) => {
|
||||
let mut thread = summary_to_thread(summary);
|
||||
thread.forked_from_id =
|
||||
forked_from_id_from_rollout(fork_rollout_path.as_path()).await;
|
||||
thread
|
||||
}
|
||||
Err(err) => {
|
||||
self.send_internal_error(
|
||||
request_id,
|
||||
@@ -4508,6 +4616,14 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
};
|
||||
thread.preview = preview_from_rollout_items(&history_items);
|
||||
thread.forked_from_id = source_thread_id
|
||||
.or_else(|| {
|
||||
history_items.iter().find_map(|item| match item {
|
||||
RolloutItem::SessionMeta(meta_line) => Some(meta_line.meta.id),
|
||||
_ => None,
|
||||
})
|
||||
})
|
||||
.map(|id| id.to_string());
|
||||
if let Err(message) = populate_thread_turns(
|
||||
&mut thread,
|
||||
ThreadTurnSource::HistoryItems(&history_items),
|
||||
@@ -5161,9 +5277,12 @@ impl CodexMessageProcessor {
|
||||
return;
|
||||
}
|
||||
};
|
||||
let mcp_config = config.to_mcp_config(self.thread_manager.plugins_manager().as_ref());
|
||||
let auth = self.auth_manager.auth().await;
|
||||
|
||||
tokio::spawn(async move {
|
||||
Self::list_mcp_server_status_task(outgoing, request, params, config).await;
|
||||
Self::list_mcp_server_status_task(outgoing, request, params, config, mcp_config, auth)
|
||||
.await;
|
||||
});
|
||||
}
|
||||
|
||||
@@ -5172,15 +5291,66 @@ impl CodexMessageProcessor {
|
||||
request_id: ConnectionRequestId,
|
||||
params: ListMcpServerStatusParams,
|
||||
config: Config,
|
||||
mcp_config: codex_mcp::mcp::McpConfig,
|
||||
auth: Option<CodexAuth>,
|
||||
) {
|
||||
let snapshot = collect_mcp_snapshot(&config).await;
|
||||
let snapshot = collect_mcp_snapshot(
|
||||
&mcp_config,
|
||||
auth.as_ref(),
|
||||
request_id.request_id.to_string(),
|
||||
)
|
||||
.await;
|
||||
|
||||
let tools_by_server = group_tools_by_server(&snapshot.tools);
|
||||
// Rebuild the tool list per original server name instead of using
|
||||
// `group_tools_by_server()`: qualified tool names are sanitized for the
|
||||
// Responses API, so a config key like `some-server` is encoded as the
|
||||
// `mcp__some_server__` prefix. Matching with the original server name's
|
||||
// sanitized prefix preserves `/mcp` output for hyphenated names.
|
||||
let effective_servers = effective_mcp_servers(&mcp_config, auth.as_ref());
|
||||
let mut sanitized_prefix_counts = HashMap::<String, usize>::new();
|
||||
for name in effective_servers.keys() {
|
||||
let prefix = qualified_mcp_tool_name_prefix(name);
|
||||
*sanitized_prefix_counts.entry(prefix).or_default() += 1;
|
||||
}
|
||||
let tools_by_server = effective_servers
|
||||
.keys()
|
||||
.map(|name| {
|
||||
let prefix = qualified_mcp_tool_name_prefix(name);
|
||||
// If multiple server names normalize to the same prefix, the
|
||||
// qualified tool namespace is ambiguous (for example
|
||||
// `some-server` and `some_server` both become
|
||||
// `mcp__some_server__`). In that case, avoid attributing the
|
||||
// same tools to multiple servers.
|
||||
let tools = if sanitized_prefix_counts
|
||||
.get(&prefix)
|
||||
.copied()
|
||||
.unwrap_or_default()
|
||||
== 1
|
||||
{
|
||||
snapshot
|
||||
.tools
|
||||
.iter()
|
||||
.filter_map(|(qualified_name, tool)| {
|
||||
qualified_name
|
||||
.strip_prefix(&prefix)
|
||||
.map(|tool_name| (tool_name.to_string(), tool.clone()))
|
||||
})
|
||||
.collect::<HashMap<_, _>>()
|
||||
} else {
|
||||
HashMap::new()
|
||||
};
|
||||
(name.clone(), tools)
|
||||
})
|
||||
.collect::<HashMap<_, _>>();
|
||||
|
||||
let mut server_names: Vec<String> = config
|
||||
.mcp_servers
|
||||
.keys()
|
||||
.cloned()
|
||||
// Include built-in/plugin MCP servers that are present in the
|
||||
// effective runtime config even when they are not user-declared in
|
||||
// `config.mcp_servers`.
|
||||
.chain(effective_servers.keys().cloned())
|
||||
.chain(snapshot.auth_statuses.keys().cloned())
|
||||
.chain(snapshot.resources.keys().cloned())
|
||||
.chain(snapshot.resource_templates.keys().cloned())
|
||||
@@ -6449,6 +6619,7 @@ impl CodexMessageProcessor {
|
||||
request_id: ConnectionRequestId,
|
||||
params: TurnStartParams,
|
||||
app_server_client_name: Option<String>,
|
||||
app_server_client_version: Option<String>,
|
||||
) {
|
||||
if let Err(error) = Self::validate_v2_input_limit(¶ms.input) {
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
@@ -6461,8 +6632,12 @@ impl CodexMessageProcessor {
|
||||
return;
|
||||
}
|
||||
};
|
||||
if let Err(error) =
|
||||
Self::set_app_server_client_name(thread.as_ref(), app_server_client_name).await
|
||||
if let Err(error) = Self::set_app_server_client_info(
|
||||
thread.as_ref(),
|
||||
app_server_client_name,
|
||||
app_server_client_version,
|
||||
)
|
||||
.await
|
||||
{
|
||||
self.outgoing.send_error(request_id, error).await;
|
||||
return;
|
||||
@@ -6556,16 +6731,17 @@ impl CodexMessageProcessor {
|
||||
}
|
||||
}
|
||||
|
||||
async fn set_app_server_client_name(
|
||||
async fn set_app_server_client_info(
|
||||
thread: &CodexThread,
|
||||
app_server_client_name: Option<String>,
|
||||
app_server_client_version: Option<String>,
|
||||
) -> Result<(), JSONRPCErrorError> {
|
||||
thread
|
||||
.set_app_server_client_name(app_server_client_name)
|
||||
.set_app_server_client_info(app_server_client_name, app_server_client_version)
|
||||
.await
|
||||
.map_err(|err| JSONRPCErrorError {
|
||||
code: INTERNAL_ERROR_CODE,
|
||||
message: format!("failed to set app server client name: {err}"),
|
||||
message: format!("failed to set app server client info: {err}"),
|
||||
data: None,
|
||||
})
|
||||
}
|
||||
@@ -7718,6 +7894,57 @@ async fn rename_rollout_with_alarm_sidecar(from: &Path, to: &Path) -> std::io::R
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_thread_list_cwd_filter(
|
||||
cwd: Option<String>,
|
||||
) -> Result<Option<PathBuf>, JSONRPCErrorError> {
|
||||
let Some(cwd) = cwd else {
|
||||
return Ok(None);
|
||||
};
|
||||
AbsolutePathBuf::relative_to_current_dir(cwd.as_str())
|
||||
.map(AbsolutePathBuf::into_path_buf)
|
||||
.map(Some)
|
||||
.map_err(|err| JSONRPCErrorError {
|
||||
code: INVALID_PARAMS_ERROR_CODE,
|
||||
message: format!("invalid thread/list cwd filter `{cwd}`: {err}"),
|
||||
data: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod thread_list_cwd_filter_tests {
|
||||
use super::normalize_thread_list_cwd_filter;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use pretty_assertions::assert_eq;
|
||||
use std::path::PathBuf;
|
||||
|
||||
#[test]
|
||||
fn normalize_thread_list_cwd_filter_preserves_absolute_paths() {
|
||||
let cwd = if cfg!(windows) {
|
||||
String::from(r"C:\srv\repo-b")
|
||||
} else {
|
||||
String::from("/srv/repo-b")
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
normalize_thread_list_cwd_filter(Some(cwd.clone())).expect("cwd filter should parse"),
|
||||
Some(PathBuf::from(cwd))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_thread_list_cwd_filter_resolves_relative_paths_against_server_cwd()
|
||||
-> std::io::Result<()> {
|
||||
let expected = AbsolutePathBuf::relative_to_current_dir("repo-b")?.to_path_buf();
|
||||
|
||||
assert_eq!(
|
||||
normalize_thread_list_cwd_filter(Some(String::from("repo-b")))
|
||||
.expect("cwd filter should parse"),
|
||||
Some(expected)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn handle_thread_listener_command(
|
||||
conversation_id: ThreadId,
|
||||
@@ -8693,6 +8920,7 @@ async fn load_thread_summary_for_rollout(
|
||||
rollout_path.display()
|
||||
)
|
||||
})?;
|
||||
thread.forked_from_id = forked_from_id_from_rollout(rollout_path).await;
|
||||
if let Some(persisted_metadata) = persisted_metadata {
|
||||
merge_mutable_thread_metadata(
|
||||
&mut thread,
|
||||
@@ -8704,6 +8932,14 @@ async fn load_thread_summary_for_rollout(
|
||||
Ok(thread)
|
||||
}
|
||||
|
||||
async fn forked_from_id_from_rollout(path: &Path) -> Option<String> {
|
||||
read_session_meta_line(path)
|
||||
.await
|
||||
.ok()
|
||||
.and_then(|meta_line| meta_line.meta.forked_from_id)
|
||||
.map(|thread_id| thread_id.to_string())
|
||||
}
|
||||
|
||||
fn merge_mutable_thread_metadata(thread: &mut Thread, persisted_thread: Thread) {
|
||||
thread.git_info = persisted_thread.git_info;
|
||||
}
|
||||
@@ -8784,6 +9020,7 @@ fn build_thread_from_snapshot(
|
||||
let now = time::OffsetDateTime::now_utc().unix_timestamp();
|
||||
Thread {
|
||||
id: thread_id.to_string(),
|
||||
forked_from_id: None,
|
||||
preview: String::new(),
|
||||
ephemeral: config_snapshot.ephemeral,
|
||||
model_provider: config_snapshot.model_provider_id.clone(),
|
||||
@@ -8853,6 +9090,7 @@ pub(crate) fn summary_to_thread(summary: ConversationSummary) -> Thread {
|
||||
|
||||
Thread {
|
||||
id: conversation_id.to_string(),
|
||||
forked_from_id: None,
|
||||
preview,
|
||||
ephemeral: false,
|
||||
model_provider,
|
||||
@@ -9349,6 +9587,44 @@ mod tests {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_summary_from_rollout_preserves_forked_from_id() -> Result<()> {
|
||||
use codex_protocol::protocol::RolloutItem;
|
||||
use codex_protocol::protocol::RolloutLine;
|
||||
use codex_protocol::protocol::SessionMetaLine;
|
||||
use std::fs;
|
||||
|
||||
let temp_dir = TempDir::new()?;
|
||||
let path = temp_dir.path().join("rollout.jsonl");
|
||||
|
||||
let conversation_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
|
||||
let forked_from_id = ThreadId::from_string("ad7f0408-99b8-4f6e-a46f-bd0eec433370")?;
|
||||
let timestamp = "2025-09-05T16:53:11.850Z".to_string();
|
||||
|
||||
let session_meta = SessionMeta {
|
||||
id: conversation_id,
|
||||
forked_from_id: Some(forked_from_id),
|
||||
timestamp: timestamp.clone(),
|
||||
model_provider: Some("test-provider".to_string()),
|
||||
..SessionMeta::default()
|
||||
};
|
||||
|
||||
let line = RolloutLine {
|
||||
timestamp,
|
||||
item: RolloutItem::SessionMeta(SessionMetaLine {
|
||||
meta: session_meta,
|
||||
git: None,
|
||||
}),
|
||||
};
|
||||
fs::write(&path, format!("{}\n", serde_json::to_string(&line)?))?;
|
||||
|
||||
assert_eq!(
|
||||
forked_from_id_from_rollout(path.as_path()).await,
|
||||
Some(forked_from_id.to_string())
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn aborting_pending_request_clears_pending_state() -> Result<()> {
|
||||
let thread_id = ThreadId::from_string("bfd12a78-5900-467b-9bc5-d3d35df08191")?;
|
||||
|
||||
@@ -3,12 +3,12 @@ use std::sync::Arc;
|
||||
|
||||
use codex_app_server_protocol::McpServerOauthLoginCompletedNotification;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_config::types::McpServerConfig;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::types::McpServerConfig;
|
||||
use codex_core::mcp::auth::McpOAuthLoginSupport;
|
||||
use codex_core::mcp::auth::oauth_login_support;
|
||||
use codex_core::mcp::auth::resolve_oauth_scopes;
|
||||
use codex_core::mcp::auth::should_retry_without_scopes;
|
||||
use codex_mcp::mcp::auth::McpOAuthLoginSupport;
|
||||
use codex_mcp::mcp::auth::oauth_login_support;
|
||||
use codex_mcp::mcp::auth::resolve_oauth_scopes;
|
||||
use codex_mcp::mcp::auth::should_retry_without_scopes;
|
||||
use codex_rmcp_client::perform_oauth_login_silent;
|
||||
use tracing::warn;
|
||||
|
||||
|
||||
@@ -18,12 +18,12 @@ use codex_app_server_protocol::CommandExecWriteParams;
|
||||
use codex_app_server_protocol::CommandExecWriteResponse;
|
||||
use codex_app_server_protocol::JSONRPCErrorError;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_core::bytes_to_string_smart;
|
||||
use codex_core::config::StartedNetworkProxy;
|
||||
use codex_core::exec::DEFAULT_EXEC_COMMAND_TIMEOUT_MS;
|
||||
use codex_core::exec::ExecExpiration;
|
||||
use codex_core::exec::IO_DRAIN_TIMEOUT_MS;
|
||||
use codex_core::sandboxing::ExecRequest;
|
||||
use codex_protocol::exec_output::bytes_to_string_smart;
|
||||
use codex_sandboxing::SandboxType;
|
||||
use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP;
|
||||
use codex_utils_pty::ProcessHandle;
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use crate::error_code::INTERNAL_ERROR_CODE;
|
||||
use crate::error_code::INVALID_REQUEST_ERROR_CODE;
|
||||
use async_trait::async_trait;
|
||||
use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_app_server_protocol::ConfigBatchWriteParams;
|
||||
use codex_app_server_protocol::ConfigReadParams;
|
||||
use codex_app_server_protocol::ConfigReadResponse;
|
||||
@@ -16,7 +17,6 @@ use codex_app_server_protocol::NetworkDomainPermission;
|
||||
use codex_app_server_protocol::NetworkRequirements;
|
||||
use codex_app_server_protocol::NetworkUnixSocketPermission;
|
||||
use codex_app_server_protocol::SandboxMode;
|
||||
use codex_core::AnalyticsEventsClient;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigService;
|
||||
@@ -366,6 +366,12 @@ fn map_requirements_toml_to_api(requirements: ConfigRequirementsToml) -> ConfigR
|
||||
.map(codex_app_server_protocol::AskForApproval::from)
|
||||
.collect()
|
||||
}),
|
||||
allowed_approvals_reviewers: requirements.allowed_approvals_reviewers.map(|reviewers| {
|
||||
reviewers
|
||||
.into_iter()
|
||||
.map(codex_app_server_protocol::ApprovalsReviewer::from)
|
||||
.collect()
|
||||
}),
|
||||
allowed_sandbox_modes: requirements.allowed_sandbox_modes.map(|modes| {
|
||||
modes
|
||||
.into_iter()
|
||||
@@ -510,15 +516,16 @@ fn config_write_error(code: ConfigWriteErrorCode, message: impl Into<String>) ->
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use codex_core::AnalyticsEventsClient;
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::CodexAuth;
|
||||
use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_core::config_loader::NetworkDomainPermissionToml as CoreNetworkDomainPermissionToml;
|
||||
use codex_core::config_loader::NetworkDomainPermissionsToml as CoreNetworkDomainPermissionsToml;
|
||||
use codex_core::config_loader::NetworkRequirementsToml as CoreNetworkRequirementsToml;
|
||||
use codex_core::config_loader::NetworkUnixSocketPermissionToml as CoreNetworkUnixSocketPermissionToml;
|
||||
use codex_core::config_loader::NetworkUnixSocketPermissionsToml as CoreNetworkUnixSocketPermissionsToml;
|
||||
use codex_features::Feature;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::CodexAuth;
|
||||
use codex_protocol::config_types::ApprovalsReviewer as CoreApprovalsReviewer;
|
||||
use codex_protocol::protocol::AskForApproval as CoreAskForApproval;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
@@ -545,6 +552,10 @@ mod tests {
|
||||
CoreAskForApproval::Never,
|
||||
CoreAskForApproval::OnRequest,
|
||||
]),
|
||||
allowed_approvals_reviewers: Some(vec![
|
||||
CoreApprovalsReviewer::User,
|
||||
CoreApprovalsReviewer::GuardianSubagent,
|
||||
]),
|
||||
allowed_sandbox_modes: Some(vec![
|
||||
CoreSandboxModeRequirement::ReadOnly,
|
||||
CoreSandboxModeRequirement::ExternalSandbox,
|
||||
@@ -602,6 +613,13 @@ mod tests {
|
||||
codex_app_server_protocol::AskForApproval::OnRequest,
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
mapped.allowed_approvals_reviewers,
|
||||
Some(vec![
|
||||
codex_app_server_protocol::ApprovalsReviewer::User,
|
||||
codex_app_server_protocol::ApprovalsReviewer::GuardianSubagent,
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
mapped.allowed_sandbox_modes,
|
||||
Some(vec![SandboxMode::ReadOnly]),
|
||||
@@ -651,6 +669,7 @@ mod tests {
|
||||
fn map_requirements_toml_to_api_omits_unix_socket_none_entries_from_legacy_network_fields() {
|
||||
let requirements = ConfigRequirementsToml {
|
||||
allowed_approval_policies: None,
|
||||
allowed_approvals_reviewers: None,
|
||||
allowed_sandbox_modes: None,
|
||||
allowed_web_search_modes: None,
|
||||
guardian_developer_instructions: None,
|
||||
@@ -707,6 +726,7 @@ mod tests {
|
||||
fn map_requirements_toml_to_api_normalizes_allowed_web_search_modes() {
|
||||
let requirements = ConfigRequirementsToml {
|
||||
allowed_approval_policies: None,
|
||||
allowed_approvals_reviewers: None,
|
||||
allowed_sandbox_modes: None,
|
||||
allowed_web_search_modes: Some(Vec::new()),
|
||||
guardian_developer_instructions: None,
|
||||
|
||||
@@ -64,6 +64,7 @@ use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use crate::transport::CHANNEL_CAPACITY;
|
||||
use crate::transport::OutboundConnectionState;
|
||||
use crate::transport::route_outgoing_envelope;
|
||||
use codex_analytics::AppServerRpcTransport;
|
||||
use codex_app_server_protocol::ClientNotification;
|
||||
use codex_app_server_protocol::ClientRequest;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
@@ -74,12 +75,12 @@ use codex_app_server_protocol::Result;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_app_server_protocol::ServerRequest;
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_core::AppServerRpcTransport;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config_loader::CloudRequirementsLoader;
|
||||
use codex_core::config_loader::LoaderOverrides;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_feedback::CodexFeedback;
|
||||
use codex_login::AuthManager;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
@@ -379,6 +380,8 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
});
|
||||
|
||||
let processor_outgoing = Arc::clone(&outgoing_message_sender);
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(args.config.as_ref(), args.enable_codex_api_key_env);
|
||||
let (processor_tx, mut processor_rx) = mpsc::channel::<ProcessorCommand>(channel_capacity);
|
||||
let mut processor_handle = tokio::spawn(async move {
|
||||
let mut processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
@@ -393,7 +396,7 @@ fn start_uninitialized(args: InProcessStartArgs) -> InProcessClientHandle {
|
||||
log_db: None,
|
||||
config_warnings: args.config_warnings,
|
||||
session_source: args.session_source,
|
||||
enable_codex_api_key_env: args.enable_codex_api_key_env,
|
||||
auth_manager,
|
||||
rpc_transport: AppServerRpcTransport::InProcess,
|
||||
});
|
||||
let mut thread_created_rx = processor.thread_created_receiver();
|
||||
|
||||
@@ -2,12 +2,13 @@
|
||||
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_cloud_requirements::cloud_requirements_loader;
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigBuilder;
|
||||
use codex_core::config_loader::CloudRequirementsLoader;
|
||||
use codex_core::config_loader::ConfigLayerStackOrdering;
|
||||
use codex_core::config_loader::LoaderOverrides;
|
||||
use codex_features::Feature;
|
||||
use codex_login::AuthManager;
|
||||
use codex_utils_cli::CliConfigOverrides;
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
@@ -29,14 +30,15 @@ use crate::transport::OutboundConnectionState;
|
||||
use crate::transport::TransportEvent;
|
||||
use crate::transport::auth::policy_from_settings;
|
||||
use crate::transport::route_outgoing_envelope;
|
||||
use crate::transport::start_remote_control;
|
||||
use crate::transport::start_stdio_connection;
|
||||
use crate::transport::start_websocket_acceptor;
|
||||
use codex_analytics::AppServerRpcTransport;
|
||||
use codex_app_server_protocol::ConfigLayerSource;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::TextPosition as AppTextPosition;
|
||||
use codex_app_server_protocol::TextRange as AppTextRange;
|
||||
use codex_core::AppServerRpcTransport;
|
||||
use codex_core::ExecPolicyError;
|
||||
use codex_core::check_execpolicy_for_warnings;
|
||||
use codex_core::config_loader::ConfigLoadError;
|
||||
@@ -46,6 +48,7 @@ use codex_feedback::CodexFeedback;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_state::log_db;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
use toml::Value as TomlValue;
|
||||
@@ -396,11 +399,8 @@ pub async fn run_main_with_transport(
|
||||
}
|
||||
}
|
||||
|
||||
let auth_manager = AuthManager::shared(
|
||||
config.codex_home.clone(),
|
||||
/*enable_codex_api_key_env*/ false,
|
||||
config.cli_auth_credentials_store_mode,
|
||||
);
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false);
|
||||
cloud_requirements_loader(
|
||||
auth_manager,
|
||||
config.chatgpt_base_url,
|
||||
@@ -457,7 +457,9 @@ pub async fn run_main_with_transport(
|
||||
range: None,
|
||||
});
|
||||
}
|
||||
if let Some(warning) = codex_core::config::system_bwrap_warning() {
|
||||
if let Some(warning) =
|
||||
codex_core::config::system_bwrap_warning(config.permissions.sandbox_policy.get())
|
||||
{
|
||||
config_warnings.push(ConfigWarningNotification {
|
||||
summary: warning,
|
||||
details: None,
|
||||
@@ -500,13 +502,13 @@ pub async fn run_main_with_transport(
|
||||
|
||||
let feedback_layer = feedback.logger_layer();
|
||||
let feedback_metadata_layer = feedback.metadata_layer();
|
||||
let log_db = codex_state::StateRuntime::init(
|
||||
let state_db = codex_state::StateRuntime::init(
|
||||
config.sqlite_home.clone(),
|
||||
config.model_provider_id.clone(),
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
.map(log_db::start);
|
||||
.ok();
|
||||
let log_db = state_db.clone().map(log_db::start);
|
||||
let log_db_layer = log_db
|
||||
.clone()
|
||||
.map(|layer| layer.with_filter(Targets::new().with_default(Level::TRACE)));
|
||||
@@ -533,11 +535,18 @@ pub async fn run_main_with_transport(
|
||||
let single_client_mode = matches!(&transport, AppServerTransport::Stdio);
|
||||
let shutdown_when_no_connections = single_client_mode;
|
||||
let graceful_signal_restart_enabled = !single_client_mode;
|
||||
let mut app_server_client_name_rx = None;
|
||||
|
||||
match transport {
|
||||
AppServerTransport::Stdio => {
|
||||
start_stdio_connection(transport_event_tx.clone(), &mut transport_accept_handles)
|
||||
.await?;
|
||||
let (stdio_client_name_tx, stdio_client_name_rx) = oneshot::channel::<String>();
|
||||
app_server_client_name_rx = Some(stdio_client_name_rx);
|
||||
start_stdio_connection(
|
||||
transport_event_tx.clone(),
|
||||
&mut transport_accept_handles,
|
||||
stdio_client_name_tx,
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
AppServerTransport::WebSocket { bind_address } => {
|
||||
let accept_handle = start_websocket_acceptor(
|
||||
@@ -549,6 +558,29 @@ pub async fn run_main_with_transport(
|
||||
.await?;
|
||||
transport_accept_handles.push(accept_handle);
|
||||
}
|
||||
AppServerTransport::Off => {}
|
||||
}
|
||||
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false);
|
||||
|
||||
if config.features.enabled(Feature::RemoteControl) {
|
||||
let accept_handle = start_remote_control(
|
||||
config.chatgpt_base_url.clone(),
|
||||
state_db.clone(),
|
||||
auth_manager.clone(),
|
||||
transport_event_tx.clone(),
|
||||
transport_shutdown_token.clone(),
|
||||
app_server_client_name_rx,
|
||||
)
|
||||
.await?;
|
||||
transport_accept_handles.push(accept_handle);
|
||||
}
|
||||
if transport_accept_handles.is_empty() {
|
||||
return Err(std::io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
"no transport configured; use --listen or enable remote control",
|
||||
));
|
||||
}
|
||||
|
||||
let outbound_handle = tokio::spawn(async move {
|
||||
@@ -609,6 +641,8 @@ pub async fn run_main_with_transport(
|
||||
let processor_handle = tokio::spawn({
|
||||
let outgoing_message_sender = Arc::new(OutgoingMessageSender::new(outgoing_tx));
|
||||
let outbound_control_tx = outbound_control_tx;
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(&config, /*enable_codex_api_key_env*/ false);
|
||||
let cli_overrides: Vec<(String, TomlValue)> = cli_kv_overrides.clone();
|
||||
let loader_overrides = loader_overrides_for_config_api;
|
||||
let mut processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
@@ -623,7 +657,7 @@ pub async fn run_main_with_transport(
|
||||
log_db,
|
||||
config_warnings,
|
||||
session_source,
|
||||
enable_codex_api_key_env: false,
|
||||
auth_manager,
|
||||
rpc_transport: analytics_rpc_transport(transport),
|
||||
});
|
||||
let mut thread_created_rx = processor.thread_created_receiver();
|
||||
@@ -851,7 +885,9 @@ pub async fn run_main_with_transport(
|
||||
fn analytics_rpc_transport(transport: AppServerTransport) -> AppServerRpcTransport {
|
||||
match transport {
|
||||
AppServerTransport::Stdio => AppServerRpcTransport::Stdio,
|
||||
AppServerTransport::WebSocket { .. } => AppServerRpcTransport::Websocket,
|
||||
AppServerTransport::WebSocket { .. } | AppServerTransport::Off => {
|
||||
AppServerRpcTransport::Websocket
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ const MANAGED_CONFIG_PATH_ENV_VAR: &str = "CODEX_APP_SERVER_MANAGED_CONFIG_PATH"
|
||||
#[derive(Debug, Parser)]
|
||||
struct AppServerArgs {
|
||||
/// Transport endpoint URL. Supported values: `stdio://` (default),
|
||||
/// `ws://IP:PORT`.
|
||||
/// `ws://IP:PORT`, `off`.
|
||||
#[arg(
|
||||
long = "listen",
|
||||
value_name = "URL",
|
||||
|
||||
@@ -20,7 +20,10 @@ use crate::outgoing_message::OutgoingMessageSender;
|
||||
use crate::outgoing_message::RequestContext;
|
||||
use crate::transport::AppServerTransport;
|
||||
use async_trait::async_trait;
|
||||
use codex_analytics::AnalyticsEventsClient;
|
||||
use codex_analytics::AppServerRpcTransport;
|
||||
use codex_app_server_protocol::AppListUpdatedNotification;
|
||||
use codex_app_server_protocol::AuthMode as LoginAuthMode;
|
||||
use codex_app_server_protocol::ChatgptAuthTokensRefreshParams;
|
||||
use codex_app_server_protocol::ChatgptAuthTokensRefreshReason;
|
||||
use codex_app_server_protocol::ChatgptAuthTokensRefreshResponse;
|
||||
@@ -55,27 +58,24 @@ use codex_app_server_protocol::ServerRequestPayload;
|
||||
use codex_app_server_protocol::experimental_required_message;
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_chatgpt::connectors;
|
||||
use codex_core::AnalyticsEventsClient;
|
||||
use codex_core::AppServerRpcTransport;
|
||||
use codex_core::AuthManager;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config_loader::CloudRequirementsLoader;
|
||||
use codex_core::config_loader::LoaderOverrides;
|
||||
use codex_core::default_client::SetOriginatorError;
|
||||
use codex_core::default_client::USER_AGENT_SUFFIX;
|
||||
use codex_core::default_client::get_codex_user_agent;
|
||||
use codex_core::default_client::set_default_client_residency_requirement;
|
||||
use codex_core::default_client::set_default_originator;
|
||||
use codex_core::models_manager::collaboration_mode_presets::CollaborationModesConfig;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_features::Feature;
|
||||
use codex_feedback::CodexFeedback;
|
||||
use codex_login::AuthMode as LoginAuthMode;
|
||||
use codex_login::AuthManager;
|
||||
use codex_login::auth::ExternalAuth;
|
||||
use codex_login::auth::ExternalAuthRefreshContext;
|
||||
use codex_login::auth::ExternalAuthRefreshReason;
|
||||
use codex_login::auth::ExternalAuthTokens;
|
||||
use codex_login::default_client::SetOriginatorError;
|
||||
use codex_login::default_client::USER_AGENT_SUFFIX;
|
||||
use codex_login::default_client::get_codex_user_agent;
|
||||
use codex_login::default_client::set_default_client_residency_requirement;
|
||||
use codex_login::default_client::set_default_originator;
|
||||
use codex_models_manager::collaboration_mode_presets::CollaborationModesConfig;
|
||||
use codex_protocol::ThreadId;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
@@ -193,7 +193,7 @@ pub(crate) struct MessageProcessorArgs {
|
||||
pub(crate) log_db: Option<LogDbLayer>,
|
||||
pub(crate) config_warnings: Vec<ConfigWarningNotification>,
|
||||
pub(crate) session_source: SessionSource,
|
||||
pub(crate) enable_codex_api_key_env: bool,
|
||||
pub(crate) auth_manager: Arc<AuthManager>,
|
||||
pub(crate) rpc_transport: AppServerRpcTransport,
|
||||
}
|
||||
|
||||
@@ -213,17 +213,12 @@ impl MessageProcessor {
|
||||
log_db,
|
||||
config_warnings,
|
||||
session_source,
|
||||
enable_codex_api_key_env,
|
||||
auth_manager,
|
||||
rpc_transport,
|
||||
} = args;
|
||||
let auth_manager = AuthManager::shared_with_external_auth(
|
||||
config.codex_home.clone(),
|
||||
enable_codex_api_key_env,
|
||||
config.cli_auth_credentials_store_mode,
|
||||
Arc::new(ExternalAuthRefreshBridge {
|
||||
outgoing: outgoing.clone(),
|
||||
}),
|
||||
);
|
||||
auth_manager.set_external_auth(Arc::new(ExternalAuthRefreshBridge {
|
||||
outgoing: outgoing.clone(),
|
||||
}));
|
||||
let thread_manager = Arc::new(ThreadManager::new(
|
||||
config.as_ref(),
|
||||
auth_manager.clone(),
|
||||
@@ -235,7 +230,6 @@ impl MessageProcessor {
|
||||
},
|
||||
environment_manager,
|
||||
));
|
||||
auth_manager.set_forced_chatgpt_workspace_id(config.forced_chatgpt_workspace_id.clone());
|
||||
let analytics_events_client = AnalyticsEventsClient::new(
|
||||
Arc::clone(&auth_manager),
|
||||
config.chatgpt_base_url.trim_end_matches('/').to_string(),
|
||||
@@ -850,6 +844,7 @@ impl MessageProcessor {
|
||||
connection_id,
|
||||
other,
|
||||
session.app_server_client_name.clone(),
|
||||
session.client_version.clone(),
|
||||
request_context,
|
||||
)
|
||||
.boxed()
|
||||
|
||||
@@ -7,6 +7,7 @@ use crate::transport::AppServerTransport;
|
||||
use anyhow::Result;
|
||||
use app_test_support::create_mock_responses_server_repeating_assistant;
|
||||
use app_test_support::write_mock_responses_config_toml;
|
||||
use codex_analytics::AppServerRpcTransport;
|
||||
use codex_app_server_protocol::ClientInfo;
|
||||
use codex_app_server_protocol::ClientRequest;
|
||||
use codex_app_server_protocol::InitializeCapabilities;
|
||||
@@ -20,13 +21,13 @@ use codex_app_server_protocol::TurnStartParams;
|
||||
use codex_app_server_protocol::TurnStartResponse;
|
||||
use codex_app_server_protocol::UserInput;
|
||||
use codex_arg0::Arg0DispatchPaths;
|
||||
use codex_core::AppServerRpcTransport;
|
||||
use codex_core::config::Config;
|
||||
use codex_core::config::ConfigBuilder;
|
||||
use codex_core::config_loader::CloudRequirementsLoader;
|
||||
use codex_core::config_loader::LoaderOverrides;
|
||||
use codex_exec_server::EnvironmentManager;
|
||||
use codex_feedback::CodexFeedback;
|
||||
use codex_login::AuthManager;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use opentelemetry::global;
|
||||
@@ -234,6 +235,8 @@ fn build_test_processor(
|
||||
) {
|
||||
let (outgoing_tx, outgoing_rx) = mpsc::channel(16);
|
||||
let outgoing = Arc::new(OutgoingMessageSender::new(outgoing_tx));
|
||||
let auth_manager =
|
||||
AuthManager::shared_from_config(config.as_ref(), /*enable_codex_api_key_env*/ false);
|
||||
let processor = MessageProcessor::new(MessageProcessorArgs {
|
||||
outgoing,
|
||||
arg0_paths: Arg0DispatchPaths::default(),
|
||||
@@ -246,7 +249,7 @@ fn build_test_processor(
|
||||
log_db: None,
|
||||
config_warnings: Vec::new(),
|
||||
session_source: SessionSource::VSCode,
|
||||
enable_codex_api_key_env: false,
|
||||
auth_manager,
|
||||
rpc_transport: AppServerRpcTransport::Stdio,
|
||||
});
|
||||
(processor, outgoing_rx)
|
||||
|
||||
@@ -4,7 +4,7 @@ use codex_app_server_protocol::Model;
|
||||
use codex_app_server_protocol::ModelUpgradeInfo;
|
||||
use codex_app_server_protocol::ReasoningEffortOption;
|
||||
use codex_core::ThreadManager;
|
||||
use codex_core::models_manager::manager::RefreshStrategy;
|
||||
use codex_models_manager::manager::RefreshStrategy;
|
||||
use codex_protocol::openai_models::ModelPreset;
|
||||
use codex_protocol::openai_models::ReasoningEffortPreset;
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ use std::sync::Weak;
|
||||
use tokio::sync::Mutex;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tracing::error;
|
||||
|
||||
type PendingInterruptQueue = Vec<(
|
||||
ConnectionRequestId,
|
||||
@@ -116,6 +117,38 @@ impl ThreadState {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn resolve_server_request_on_thread_listener(
|
||||
thread_state: &Arc<Mutex<ThreadState>>,
|
||||
request_id: RequestId,
|
||||
) {
|
||||
let (completion_tx, completion_rx) = oneshot::channel();
|
||||
let listener_command_tx = {
|
||||
let state = thread_state.lock().await;
|
||||
state.listener_command_tx()
|
||||
};
|
||||
let Some(listener_command_tx) = listener_command_tx else {
|
||||
error!("failed to remove pending client request: thread listener is not running");
|
||||
return;
|
||||
};
|
||||
|
||||
if listener_command_tx
|
||||
.send(ThreadListenerCommand::ResolveServerRequest {
|
||||
request_id,
|
||||
completion_tx,
|
||||
})
|
||||
.is_err()
|
||||
{
|
||||
error!(
|
||||
"failed to remove pending client request: thread listener command channel is closed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if let Err(err) = completion_rx.await {
|
||||
error!("failed to remove pending client request: {err}");
|
||||
}
|
||||
}
|
||||
|
||||
struct ThreadEntry {
|
||||
state: Arc<Mutex<ThreadState>>,
|
||||
connection_ids: HashSet<ConnectionId>,
|
||||
|
||||
@@ -792,6 +792,7 @@ mod tests {
|
||||
fn test_thread(thread_id: &str, source: codex_app_server_protocol::SessionSource) -> Thread {
|
||||
Thread {
|
||||
id: thread_id.to_string(),
|
||||
forked_from_id: None,
|
||||
preview: String::new(),
|
||||
ephemeral: false,
|
||||
model_provider: "mock-provider".to_string(),
|
||||
|
||||
@@ -17,6 +17,7 @@ use std::str::FromStr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::RwLock;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
@@ -28,9 +29,11 @@ use tracing::warn;
|
||||
/// plenty for an interactive CLI.
|
||||
pub(crate) const CHANNEL_CAPACITY: usize = 128;
|
||||
|
||||
mod remote_control;
|
||||
mod stdio;
|
||||
mod websocket;
|
||||
|
||||
pub(crate) use remote_control::start_remote_control;
|
||||
pub(crate) use stdio::start_stdio_connection;
|
||||
pub(crate) use websocket::start_websocket_acceptor;
|
||||
|
||||
@@ -38,6 +41,7 @@ pub(crate) use websocket::start_websocket_acceptor;
|
||||
pub enum AppServerTransport {
|
||||
Stdio,
|
||||
WebSocket { bind_address: SocketAddr },
|
||||
Off,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Eq, PartialEq)]
|
||||
@@ -51,7 +55,7 @@ impl std::fmt::Display for AppServerTransportParseError {
|
||||
match self {
|
||||
AppServerTransportParseError::UnsupportedListenUrl(listen_url) => write!(
|
||||
f,
|
||||
"unsupported --listen URL `{listen_url}`; expected `stdio://` or `ws://IP:PORT`"
|
||||
"unsupported --listen URL `{listen_url}`; expected `stdio://`, `ws://IP:PORT`, or `off`"
|
||||
),
|
||||
AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!(
|
||||
f,
|
||||
@@ -71,6 +75,10 @@ impl AppServerTransport {
|
||||
return Ok(Self::Stdio);
|
||||
}
|
||||
|
||||
if listen_url == "off" {
|
||||
return Ok(Self::Off);
|
||||
}
|
||||
|
||||
if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
|
||||
let bind_address = socket_addr.parse::<SocketAddr>().map_err(|_| {
|
||||
AppServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string())
|
||||
@@ -166,6 +174,12 @@ impl OutboundConnectionState {
|
||||
}
|
||||
}
|
||||
|
||||
static CONNECTION_ID_COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
fn next_connection_id() -> ConnectionId {
|
||||
ConnectionId(CONNECTION_ID_COUNTER.fetch_add(1, Ordering::Relaxed))
|
||||
}
|
||||
|
||||
async fn forward_incoming_message(
|
||||
transport_event_tx: &mpsc::Sender<TransportEvent>,
|
||||
writer: &mpsc::Sender<QueuedOutgoingMessage>,
|
||||
@@ -378,8 +392,11 @@ pub(crate) async fn route_outgoing_envelope(
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::error_code::OVERLOADED_ERROR_CODE;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::JSONRPCNotification;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use codex_app_server_protocol::JSONRPCResponse;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use codex_utils_absolute_path::AbsolutePathBuf;
|
||||
use pretty_assertions::assert_eq;
|
||||
@@ -393,41 +410,10 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn app_server_transport_parses_stdio_listen_url() {
|
||||
let transport = AppServerTransport::from_listen_url(AppServerTransport::DEFAULT_LISTEN_URL)
|
||||
.expect("stdio listen URL should parse");
|
||||
assert_eq!(transport, AppServerTransport::Stdio);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn app_server_transport_parses_websocket_listen_url() {
|
||||
let transport = AppServerTransport::from_listen_url("ws://127.0.0.1:1234")
|
||||
.expect("websocket listen URL should parse");
|
||||
fn listen_off_parses_as_off_transport() {
|
||||
assert_eq!(
|
||||
transport,
|
||||
AppServerTransport::WebSocket {
|
||||
bind_address: "127.0.0.1:1234".parse().expect("valid socket address"),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn app_server_transport_rejects_invalid_websocket_listen_url() {
|
||||
let err = AppServerTransport::from_listen_url("ws://localhost:1234")
|
||||
.expect_err("hostname bind address should be rejected");
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
"invalid websocket --listen URL `ws://localhost:1234`; expected `ws://IP:PORT`"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn app_server_transport_rejects_unsupported_listen_url() {
|
||||
let err = AppServerTransport::from_listen_url("http://127.0.0.1:1234")
|
||||
.expect_err("unsupported scheme should fail");
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
"unsupported --listen URL `http://127.0.0.1:1234`; expected `stdio://` or `ws://IP:PORT`"
|
||||
AppServerTransport::from_listen_url("off"),
|
||||
Ok(AppServerTransport::Off)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -437,11 +423,10 @@ mod tests {
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(1);
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel(1);
|
||||
|
||||
let first_message =
|
||||
JSONRPCMessage::Notification(codex_app_server_protocol::JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let first_message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
transport_event_tx
|
||||
.send(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
@@ -450,8 +435,8 @@ mod tests {
|
||||
.await
|
||||
.expect("queue should accept first message");
|
||||
|
||||
let request = JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest {
|
||||
id: codex_app_server_protocol::RequestId::Integer(7),
|
||||
let request = JSONRPCMessage::Request(JSONRPCRequest {
|
||||
id: RequestId::Integer(7),
|
||||
method: "config/read".to_string(),
|
||||
params: Some(json!({ "includeLayers": false })),
|
||||
trace: None,
|
||||
@@ -499,11 +484,10 @@ mod tests {
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(1);
|
||||
let (writer_tx, _writer_rx) = mpsc::channel(1);
|
||||
|
||||
let first_message =
|
||||
JSONRPCMessage::Notification(codex_app_server_protocol::JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
let first_message = JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
});
|
||||
transport_event_tx
|
||||
.send(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
@@ -512,8 +496,8 @@ mod tests {
|
||||
.await
|
||||
.expect("queue should accept first message");
|
||||
|
||||
let response = JSONRPCMessage::Response(codex_app_server_protocol::JSONRPCResponse {
|
||||
id: codex_app_server_protocol::RequestId::Integer(7),
|
||||
let response = JSONRPCMessage::Response(JSONRPCResponse {
|
||||
id: RequestId::Integer(7),
|
||||
result: json!({"ok": true}),
|
||||
});
|
||||
let transport_event_tx_for_enqueue = transport_event_tx.clone();
|
||||
@@ -553,11 +537,10 @@ mod tests {
|
||||
match forwarded_event {
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: queued_connection_id,
|
||||
message:
|
||||
JSONRPCMessage::Response(codex_app_server_protocol::JSONRPCResponse { id, result }),
|
||||
message: JSONRPCMessage::Response(JSONRPCResponse { id, result }),
|
||||
} => {
|
||||
assert_eq!(queued_connection_id, connection_id);
|
||||
assert_eq!(id, codex_app_server_protocol::RequestId::Integer(7));
|
||||
assert_eq!(id, RequestId::Integer(7));
|
||||
assert_eq!(result, json!({"ok": true}));
|
||||
}
|
||||
_ => panic!("expected forwarded response message"),
|
||||
@@ -573,12 +556,10 @@ mod tests {
|
||||
transport_event_tx
|
||||
.send(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
message: JSONRPCMessage::Notification(
|
||||
codex_app_server_protocol::JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
},
|
||||
),
|
||||
message: JSONRPCMessage::Notification(JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
}),
|
||||
})
|
||||
.await
|
||||
.expect("transport queue should accept first message");
|
||||
@@ -597,15 +578,15 @@ mod tests {
|
||||
.await
|
||||
.expect("writer queue should accept first message");
|
||||
|
||||
let request = JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest {
|
||||
id: codex_app_server_protocol::RequestId::Integer(7),
|
||||
let request = JSONRPCMessage::Request(JSONRPCRequest {
|
||||
id: RequestId::Integer(7),
|
||||
method: "config/read".to_string(),
|
||||
params: Some(json!({ "includeLayers": false })),
|
||||
trace: None,
|
||||
});
|
||||
|
||||
let enqueue_result = tokio::time::timeout(
|
||||
std::time::Duration::from_millis(100),
|
||||
let enqueue_result = timeout(
|
||||
Duration::from_millis(100),
|
||||
enqueue_incoming_message(&transport_event_tx, &writer_tx, connection_id, request),
|
||||
)
|
||||
.await
|
||||
@@ -781,7 +762,7 @@ mod tests {
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id: codex_app_server_protocol::RequestId::Integer(1),
|
||||
request_id: RequestId::Integer(1),
|
||||
params: codex_app_server_protocol::CommandExecutionRequestApprovalParams {
|
||||
thread_id: "thr_123".to_string(),
|
||||
turn_id: "turn_123".to_string(),
|
||||
@@ -843,7 +824,7 @@ mod tests {
|
||||
OutgoingEnvelope::ToConnection {
|
||||
connection_id,
|
||||
message: OutgoingMessage::Request(ServerRequest::CommandExecutionRequestApproval {
|
||||
request_id: codex_app_server_protocol::RequestId::Integer(1),
|
||||
request_id: RequestId::Integer(1),
|
||||
params: codex_app_server_protocol::CommandExecutionRequestApprovalParams {
|
||||
thread_id: "thr_123".to_string(),
|
||||
turn_id: "turn_123".to_string(),
|
||||
|
||||
@@ -0,0 +1,568 @@
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::next_connection_id;
|
||||
use super::protocol::ClientEnvelope;
|
||||
pub use super::protocol::ClientEvent;
|
||||
pub use super::protocol::ClientId;
|
||||
use super::protocol::PongStatus;
|
||||
use super::protocol::ServerEvent;
|
||||
use super::protocol::StreamId;
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use crate::transport::remote_control::QueuedServerEnvelope;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use std::collections::HashMap;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::watch;
|
||||
use tokio::task::JoinSet;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::Instant;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
const REMOTE_CONTROL_CLIENT_IDLE_TIMEOUT: Duration = Duration::from_secs(10 * 60);
|
||||
pub(crate) const REMOTE_CONTROL_IDLE_SWEEP_INTERVAL: Duration = Duration::from_secs(30);
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct Stopped;
|
||||
|
||||
struct ClientState {
|
||||
connection_id: ConnectionId,
|
||||
disconnect_token: CancellationToken,
|
||||
last_activity_at: Instant,
|
||||
last_inbound_seq_id: Option<u64>,
|
||||
status_tx: watch::Sender<PongStatus>,
|
||||
}
|
||||
|
||||
pub(crate) struct ClientTracker {
|
||||
clients: HashMap<(ClientId, StreamId), ClientState>,
|
||||
legacy_stream_ids: HashMap<ClientId, StreamId>,
|
||||
join_set: JoinSet<(ClientId, StreamId)>,
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
}
|
||||
|
||||
impl ClientTracker {
|
||||
pub(crate) fn new(
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: &CancellationToken,
|
||||
) -> Self {
|
||||
Self {
|
||||
clients: HashMap::new(),
|
||||
legacy_stream_ids: HashMap::new(),
|
||||
join_set: JoinSet::new(),
|
||||
server_event_tx,
|
||||
transport_event_tx,
|
||||
shutdown_token: shutdown_token.child_token(),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn bookkeep_join_set(&mut self) -> Option<(ClientId, StreamId)> {
|
||||
while let Some(join_result) = self.join_set.join_next().await {
|
||||
let Ok(client_key) = join_result else {
|
||||
continue;
|
||||
};
|
||||
return Some(client_key);
|
||||
}
|
||||
futures::future::pending().await
|
||||
}
|
||||
|
||||
pub(crate) async fn shutdown(&mut self) {
|
||||
self.shutdown_token.cancel();
|
||||
|
||||
while let Some(client_key) = self.clients.keys().next().cloned() {
|
||||
let _ = self.close_client(&client_key).await;
|
||||
}
|
||||
|
||||
self.drain_join_set().await;
|
||||
}
|
||||
|
||||
async fn drain_join_set(&mut self) {
|
||||
while self.join_set.join_next().await.is_some() {}
|
||||
}
|
||||
|
||||
pub(crate) async fn handle_message(
|
||||
&mut self,
|
||||
client_envelope: ClientEnvelope,
|
||||
) -> Result<(), Stopped> {
|
||||
let ClientEnvelope {
|
||||
client_id,
|
||||
event,
|
||||
stream_id,
|
||||
seq_id,
|
||||
cursor: _,
|
||||
} = client_envelope;
|
||||
let is_legacy_stream_id = stream_id.is_none();
|
||||
let is_initialize = matches!(&event, ClientEvent::ClientMessage { message } if remote_control_message_starts_connection(message));
|
||||
let stream_id = match stream_id {
|
||||
Some(stream_id) => stream_id,
|
||||
None if is_initialize => {
|
||||
// TODO(ruslan): delete this fallback once all clients are updated to send stream_id.
|
||||
self.legacy_stream_ids
|
||||
.remove(&client_id)
|
||||
.unwrap_or_else(StreamId::new_random)
|
||||
}
|
||||
None => self
|
||||
.legacy_stream_ids
|
||||
.get(&client_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| {
|
||||
if matches!(&event, ClientEvent::Ping) {
|
||||
StreamId::new_random()
|
||||
} else {
|
||||
StreamId(String::new())
|
||||
}
|
||||
}),
|
||||
};
|
||||
if stream_id.0.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let client_key = (client_id.clone(), stream_id.clone());
|
||||
match event {
|
||||
ClientEvent::ClientMessage { message } => {
|
||||
if let Some(seq_id) = seq_id
|
||||
&& let Some(client) = self.clients.get(&client_key)
|
||||
&& client
|
||||
.last_inbound_seq_id
|
||||
.is_some_and(|last_seq_id| last_seq_id >= seq_id)
|
||||
&& !is_initialize
|
||||
{
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if is_initialize && self.clients.contains_key(&client_key) {
|
||||
self.close_client(&client_key).await?;
|
||||
}
|
||||
|
||||
if let Some(connection_id) = self.clients.get_mut(&client_key).map(|client| {
|
||||
client.last_activity_at = Instant::now();
|
||||
if let Some(seq_id) = seq_id {
|
||||
client.last_inbound_seq_id = Some(seq_id);
|
||||
}
|
||||
client.connection_id
|
||||
}) {
|
||||
self.send_transport_event(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
message,
|
||||
})
|
||||
.await?;
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if !is_initialize {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let connection_id = next_connection_id();
|
||||
let (writer_tx, writer_rx) =
|
||||
mpsc::channel::<QueuedOutgoingMessage>(CHANNEL_CAPACITY);
|
||||
let disconnect_token = self.shutdown_token.child_token();
|
||||
self.send_transport_event(TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
writer: writer_tx,
|
||||
disconnect_sender: Some(disconnect_token.clone()),
|
||||
})
|
||||
.await?;
|
||||
|
||||
let (status_tx, status_rx) = watch::channel(PongStatus::Active);
|
||||
self.join_set.spawn(Self::run_client_outbound(
|
||||
client_id.clone(),
|
||||
stream_id.clone(),
|
||||
self.server_event_tx.clone(),
|
||||
writer_rx,
|
||||
status_rx,
|
||||
disconnect_token.clone(),
|
||||
));
|
||||
self.clients.insert(
|
||||
client_key,
|
||||
ClientState {
|
||||
connection_id,
|
||||
disconnect_token,
|
||||
last_activity_at: Instant::now(),
|
||||
last_inbound_seq_id: if is_legacy_stream_id { None } else { seq_id },
|
||||
status_tx,
|
||||
},
|
||||
);
|
||||
if is_legacy_stream_id {
|
||||
self.legacy_stream_ids.insert(client_id.clone(), stream_id);
|
||||
}
|
||||
self.send_transport_event(TransportEvent::IncomingMessage {
|
||||
connection_id,
|
||||
message,
|
||||
})
|
||||
.await
|
||||
}
|
||||
ClientEvent::Ack => Ok(()),
|
||||
ClientEvent::Ping => {
|
||||
if let Some(client) = self.clients.get_mut(&client_key) {
|
||||
client.last_activity_at = Instant::now();
|
||||
let _ = client.status_tx.send(PongStatus::Active);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let server_event_tx = self.server_event_tx.clone();
|
||||
tokio::spawn(async move {
|
||||
let server_envelope = QueuedServerEnvelope {
|
||||
event: ServerEvent::Pong {
|
||||
status: PongStatus::Unknown,
|
||||
},
|
||||
client_id,
|
||||
stream_id,
|
||||
write_complete_tx: None,
|
||||
};
|
||||
let _ = server_event_tx.send(server_envelope).await;
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
ClientEvent::ClientClosed => self.close_client(&client_key).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_client_outbound(
|
||||
client_id: ClientId,
|
||||
stream_id: StreamId,
|
||||
server_event_tx: mpsc::Sender<QueuedServerEnvelope>,
|
||||
mut writer_rx: mpsc::Receiver<QueuedOutgoingMessage>,
|
||||
mut status_rx: watch::Receiver<PongStatus>,
|
||||
disconnect_token: CancellationToken,
|
||||
) -> (ClientId, StreamId) {
|
||||
loop {
|
||||
let (event, write_complete_tx) = tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
queued_message = writer_rx.recv() => {
|
||||
let Some(queued_message) = queued_message else {
|
||||
break;
|
||||
};
|
||||
let event = ServerEvent::ServerMessage {
|
||||
message: Box::new(queued_message.message),
|
||||
};
|
||||
(event, queued_message.write_complete_tx)
|
||||
}
|
||||
changed = status_rx.changed() => {
|
||||
if changed.is_err() {
|
||||
break;
|
||||
}
|
||||
let event = ServerEvent::Pong { status: status_rx.borrow().clone() };
|
||||
(event, None)
|
||||
}
|
||||
};
|
||||
let send_result = tokio::select! {
|
||||
_ = disconnect_token.cancelled() => {
|
||||
break;
|
||||
}
|
||||
send_result = server_event_tx.send(QueuedServerEnvelope {
|
||||
event,
|
||||
client_id: client_id.clone(),
|
||||
stream_id: stream_id.clone(),
|
||||
write_complete_tx,
|
||||
}) => send_result,
|
||||
};
|
||||
if send_result.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
(client_id, stream_id)
|
||||
}
|
||||
|
||||
pub(crate) async fn close_expired_clients(
|
||||
&mut self,
|
||||
) -> Result<Vec<(ClientId, StreamId)>, Stopped> {
|
||||
let now = Instant::now();
|
||||
let expired_client_ids: Vec<(ClientId, StreamId)> = self
|
||||
.clients
|
||||
.iter()
|
||||
.filter_map(|(client_key, client)| {
|
||||
(!remote_control_client_is_alive(client, now)).then_some(client_key.clone())
|
||||
})
|
||||
.collect();
|
||||
for client_key in &expired_client_ids {
|
||||
self.close_client(client_key).await?;
|
||||
}
|
||||
Ok(expired_client_ids)
|
||||
}
|
||||
|
||||
pub(super) async fn close_client(
|
||||
&mut self,
|
||||
client_key: &(ClientId, StreamId),
|
||||
) -> Result<(), Stopped> {
|
||||
let Some(client) = self.clients.remove(client_key) else {
|
||||
return Ok(());
|
||||
};
|
||||
if self
|
||||
.legacy_stream_ids
|
||||
.get(&client_key.0)
|
||||
.is_some_and(|stream_id| stream_id == &client_key.1)
|
||||
{
|
||||
self.legacy_stream_ids.remove(&client_key.0);
|
||||
}
|
||||
client.disconnect_token.cancel();
|
||||
self.send_transport_event(TransportEvent::ConnectionClosed {
|
||||
connection_id: client.connection_id,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
async fn send_transport_event(&self, event: TransportEvent) -> Result<(), Stopped> {
|
||||
self.transport_event_tx
|
||||
.send(event)
|
||||
.await
|
||||
.map_err(|_| Stopped)
|
||||
}
|
||||
}
|
||||
|
||||
fn remote_control_message_starts_connection(message: &JSONRPCMessage) -> bool {
|
||||
matches!(
|
||||
message,
|
||||
JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest { method, .. })
|
||||
if method == "initialize"
|
||||
)
|
||||
}
|
||||
|
||||
fn remote_control_client_is_alive(client: &ClientState, now: Instant) -> bool {
|
||||
now.duration_since(client.last_activity_at) < REMOTE_CONTROL_CLIENT_IDLE_TIMEOUT
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::outgoing_message::OutgoingMessage;
|
||||
use crate::transport::remote_control::protocol::ClientEnvelope;
|
||||
use crate::transport::remote_control::protocol::ClientEvent;
|
||||
use codex_app_server_protocol::ConfigWarningNotification;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use codex_app_server_protocol::RequestId;
|
||||
use codex_app_server_protocol::ServerNotification;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use tokio::time::timeout;
|
||||
|
||||
fn initialize_envelope(client_id: &str) -> ClientEnvelope {
|
||||
initialize_envelope_with_stream_id(client_id, /*stream_id*/ None)
|
||||
}
|
||||
|
||||
fn initialize_envelope_with_stream_id(
|
||||
client_id: &str,
|
||||
stream_id: Option<&str>,
|
||||
) -> ClientEnvelope {
|
||||
ClientEnvelope {
|
||||
event: ClientEvent::ClientMessage {
|
||||
message: JSONRPCMessage::Request(JSONRPCRequest {
|
||||
id: RequestId::Integer(1),
|
||||
method: "initialize".to_string(),
|
||||
params: Some(json!({
|
||||
"clientInfo": {
|
||||
"name": "remote-test-client",
|
||||
"version": "0.1.0"
|
||||
}
|
||||
})),
|
||||
trace: None,
|
||||
}),
|
||||
},
|
||||
client_id: ClientId(client_id.to_string()),
|
||||
stream_id: stream_id.map(|stream_id| StreamId(stream_id.to_string())),
|
||||
seq_id: Some(0),
|
||||
cursor: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancelled_outbound_task_emits_connection_closed() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
|
||||
let (connection_id, disconnect_sender) = match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection opened should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionOpened {
|
||||
connection_id,
|
||||
disconnect_sender: Some(disconnect_sender),
|
||||
..
|
||||
} => (connection_id, disconnect_sender),
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("initialize should be forwarded")
|
||||
{
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: incoming_connection_id,
|
||||
..
|
||||
} => assert_eq!(incoming_connection_id, connection_id),
|
||||
other => panic!("expected incoming initialize, got {other:?}"),
|
||||
}
|
||||
|
||||
disconnect_sender.cancel();
|
||||
let closed_client_id = timeout(Duration::from_secs(1), client_tracker.bookkeep_join_set())
|
||||
.await
|
||||
.expect("bookkeeping should process the closed task")
|
||||
.expect("closed task should return client id");
|
||||
assert_eq!(closed_client_id.0, ClientId("client-1".to_string()));
|
||||
client_tracker
|
||||
.close_client(&closed_client_id)
|
||||
.await
|
||||
.expect("closed client should emit connection closed");
|
||||
|
||||
match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection closed should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionClosed {
|
||||
connection_id: closed_connection_id,
|
||||
} => assert_eq!(closed_connection_id, connection_id),
|
||||
other => panic!("expected connection closed, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutdown_cancels_blocked_outbound_forwarding() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(1);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx.clone(), transport_event_tx, &shutdown_token);
|
||||
|
||||
server_event_tx
|
||||
.send(QueuedServerEnvelope {
|
||||
event: ServerEvent::Pong {
|
||||
status: PongStatus::Unknown,
|
||||
},
|
||||
client_id: ClientId("queued-client".to_string()),
|
||||
stream_id: StreamId("queued-stream".to_string()),
|
||||
write_complete_tx: None,
|
||||
})
|
||||
.await
|
||||
.expect("server event queue should accept prefill");
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
|
||||
let writer = match transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("connection opened should be sent")
|
||||
{
|
||||
TransportEvent::ConnectionOpened { writer, .. } => writer,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx
|
||||
.recv()
|
||||
.await
|
||||
.expect("initialize should be forwarded");
|
||||
|
||||
writer
|
||||
.send(QueuedOutgoingMessage::new(
|
||||
OutgoingMessage::AppServerNotification(ServerNotification::ConfigWarning(
|
||||
ConfigWarningNotification {
|
||||
summary: "test".to_string(),
|
||||
details: None,
|
||||
path: None,
|
||||
range: None,
|
||||
},
|
||||
)),
|
||||
))
|
||||
.await
|
||||
.expect("writer should accept queued message");
|
||||
|
||||
timeout(Duration::from_secs(1), client_tracker.shutdown())
|
||||
.await
|
||||
.expect("shutdown should not hang on blocked server forwarding");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn initialize_with_new_stream_id_opens_new_connection_for_same_client() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope_with_stream_id(
|
||||
"client-1",
|
||||
Some("stream-1"),
|
||||
))
|
||||
.await
|
||||
.expect("first initialize should open client");
|
||||
let first_connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx.recv().await.expect("initialize event");
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope_with_stream_id(
|
||||
"client-1",
|
||||
Some("stream-2"),
|
||||
))
|
||||
.await
|
||||
.expect("second initialize should open client");
|
||||
let second_connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
|
||||
assert_ne!(first_connection_id, second_connection_id);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn legacy_initialize_without_stream_id_resets_inbound_seq_id() {
|
||||
let (server_event_tx, _server_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let (transport_event_tx, mut transport_event_rx) = mpsc::channel(CHANNEL_CAPACITY);
|
||||
let shutdown_token = CancellationToken::new();
|
||||
let mut client_tracker =
|
||||
ClientTracker::new(server_event_tx, transport_event_tx, &shutdown_token);
|
||||
|
||||
client_tracker
|
||||
.handle_message(initialize_envelope("client-1"))
|
||||
.await
|
||||
.expect("initialize should open client");
|
||||
let connection_id = match transport_event_rx.recv().await.expect("open event") {
|
||||
TransportEvent::ConnectionOpened { connection_id, .. } => connection_id,
|
||||
other => panic!("expected connection opened, got {other:?}"),
|
||||
};
|
||||
let _ = transport_event_rx.recv().await.expect("initialize event");
|
||||
|
||||
client_tracker
|
||||
.handle_message(ClientEnvelope {
|
||||
event: ClientEvent::ClientMessage {
|
||||
message: JSONRPCMessage::Notification(
|
||||
codex_app_server_protocol::JSONRPCNotification {
|
||||
method: "initialized".to_string(),
|
||||
params: None,
|
||||
},
|
||||
),
|
||||
},
|
||||
client_id: ClientId("client-1".to_string()),
|
||||
stream_id: None,
|
||||
seq_id: Some(0),
|
||||
cursor: None,
|
||||
})
|
||||
.await
|
||||
.expect("legacy followup should be forwarded");
|
||||
|
||||
match transport_event_rx.recv().await.expect("followup event") {
|
||||
TransportEvent::IncomingMessage {
|
||||
connection_id: incoming_connection_id,
|
||||
..
|
||||
} => assert_eq!(incoming_connection_id, connection_id),
|
||||
other => panic!("expected incoming message, got {other:?}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,503 @@
|
||||
use super::protocol::EnrollRemoteServerRequest;
|
||||
use super::protocol::EnrollRemoteServerResponse;
|
||||
use super::protocol::RemoteControlTarget;
|
||||
use axum::http::HeaderMap;
|
||||
use codex_login::default_client::build_reqwest_client;
|
||||
use codex_state::RemoteControlEnrollmentRecord;
|
||||
use codex_state::StateRuntime;
|
||||
use gethostname::gethostname;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use tracing::info;
|
||||
use tracing::warn;
|
||||
|
||||
const REMOTE_CONTROL_ENROLL_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
|
||||
const REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES: usize = 4096;
|
||||
|
||||
const REQUEST_ID_HEADER: &str = "x-request-id";
|
||||
const OAI_REQUEST_ID_HEADER: &str = "x-oai-request-id";
|
||||
const CF_RAY_HEADER: &str = "cf-ray";
|
||||
pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct RemoteControlEnrollment {
|
||||
pub(super) account_id: String,
|
||||
pub(super) environment_id: String,
|
||||
pub(super) server_id: String,
|
||||
pub(super) server_name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct RemoteControlConnectionAuth {
|
||||
pub(super) bearer_token: String,
|
||||
pub(super) account_id: String,
|
||||
}
|
||||
|
||||
pub(super) async fn load_persisted_remote_control_enrollment(
|
||||
state_db: Option<&StateRuntime>,
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
account_id: &str,
|
||||
app_server_client_name: Option<&str>,
|
||||
) -> Option<RemoteControlEnrollment> {
|
||||
let Some(state_db) = state_db else {
|
||||
info!(
|
||||
"remote control enrollment cache unavailable because sqlite state db is disabled: websocket_url={}, account_id={}, app_server_client_name={:?}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
return None;
|
||||
};
|
||||
let enrollment = match state_db
|
||||
.get_remote_control_enrollment(
|
||||
&remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(enrollment) => enrollment,
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"failed to load persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, err={err}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
|
||||
match enrollment {
|
||||
Some(enrollment) => {
|
||||
info!(
|
||||
"reusing persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, server_id={}, environment_id={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.server_id,
|
||||
enrollment.environment_id
|
||||
);
|
||||
Some(RemoteControlEnrollment {
|
||||
account_id: enrollment.account_id,
|
||||
environment_id: enrollment.environment_id,
|
||||
server_id: enrollment.server_id,
|
||||
server_name: enrollment.server_name,
|
||||
})
|
||||
}
|
||||
None => {
|
||||
info!(
|
||||
"no persisted remote control enrollment found: websocket_url={}, account_id={}, app_server_client_name={:?}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn update_persisted_remote_control_enrollment(
|
||||
state_db: Option<&StateRuntime>,
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
account_id: &str,
|
||||
app_server_client_name: Option<&str>,
|
||||
enrollment: Option<&RemoteControlEnrollment>,
|
||||
) -> io::Result<()> {
|
||||
let Some(state_db) = state_db else {
|
||||
info!(
|
||||
"skipping remote control enrollment persistence because sqlite state db is disabled: websocket_url={}, account_id={}, app_server_client_name={:?}, has_enrollment={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.is_some()
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
if let &Some(enrollment) = &enrollment
|
||||
&& enrollment.account_id != account_id
|
||||
{
|
||||
return Err(io::Error::other(format!(
|
||||
"enrollment account_id does not match expected account_id `{account_id}`"
|
||||
)));
|
||||
}
|
||||
|
||||
if let Some(enrollment) = enrollment {
|
||||
state_db
|
||||
.upsert_remote_control_enrollment(&RemoteControlEnrollmentRecord {
|
||||
websocket_url: remote_control_target.websocket_url.clone(),
|
||||
account_id: account_id.to_string(),
|
||||
app_server_client_name: app_server_client_name.map(str::to_string),
|
||||
server_id: enrollment.server_id.clone(),
|
||||
environment_id: enrollment.environment_id.clone(),
|
||||
server_name: enrollment.server_name.clone(),
|
||||
})
|
||||
.await
|
||||
.map_err(io::Error::other)?;
|
||||
info!(
|
||||
"persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, server_id={}, environment_id={}",
|
||||
remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
enrollment.server_id,
|
||||
enrollment.environment_id
|
||||
);
|
||||
Ok(())
|
||||
} else {
|
||||
let rows_affected = state_db
|
||||
.delete_remote_control_enrollment(
|
||||
&remote_control_target.websocket_url,
|
||||
account_id,
|
||||
app_server_client_name,
|
||||
)
|
||||
.await
|
||||
.map_err(io::Error::other)?;
|
||||
info!(
|
||||
"cleared persisted remote control enrollment: websocket_url={}, account_id={}, app_server_client_name={:?}, rows_affected={rows_affected}",
|
||||
remote_control_target.websocket_url, account_id, app_server_client_name
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn preview_remote_control_response_body(body: &[u8]) -> String {
|
||||
let body = String::from_utf8_lossy(body);
|
||||
let trimmed = body.trim();
|
||||
if trimmed.is_empty() {
|
||||
return "<empty>".to_string();
|
||||
}
|
||||
if trimmed.len() <= REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES {
|
||||
return trimmed.to_string();
|
||||
}
|
||||
|
||||
let mut cut = REMOTE_CONTROL_RESPONSE_BODY_MAX_BYTES;
|
||||
while !trimmed.is_char_boundary(cut) {
|
||||
cut = cut.saturating_sub(1);
|
||||
}
|
||||
let mut truncated = trimmed[..cut].to_string();
|
||||
truncated.push_str("...");
|
||||
truncated
|
||||
}
|
||||
|
||||
pub(crate) fn format_headers(headers: &HeaderMap) -> String {
|
||||
let request_id_str = headers
|
||||
.get(REQUEST_ID_HEADER)
|
||||
.or_else(|| headers.get(OAI_REQUEST_ID_HEADER))
|
||||
.map(|value| value.to_str().unwrap_or("<invalid utf-8>").to_owned())
|
||||
.unwrap_or_else(|| "<none>".to_owned());
|
||||
let cf_ray_str = headers
|
||||
.get(CF_RAY_HEADER)
|
||||
.map(|value| value.to_str().unwrap_or("<invalid utf-8>").to_owned())
|
||||
.unwrap_or_else(|| "<none>".to_owned());
|
||||
format!("request-id: {request_id_str}, cf-ray: {cf_ray_str}")
|
||||
}
|
||||
|
||||
pub(super) async fn enroll_remote_control_server(
|
||||
remote_control_target: &RemoteControlTarget,
|
||||
auth: &RemoteControlConnectionAuth,
|
||||
) -> io::Result<RemoteControlEnrollment> {
|
||||
let enroll_url = &remote_control_target.enroll_url;
|
||||
let server_name = gethostname().to_string_lossy().trim().to_string();
|
||||
let request = EnrollRemoteServerRequest {
|
||||
name: server_name.clone(),
|
||||
os: std::env::consts::OS,
|
||||
arch: std::env::consts::ARCH,
|
||||
app_server_version: env!("CARGO_PKG_VERSION"),
|
||||
};
|
||||
let client = build_reqwest_client();
|
||||
let http_request = client
|
||||
.post(enroll_url)
|
||||
.timeout(REMOTE_CONTROL_ENROLL_TIMEOUT)
|
||||
.bearer_auth(&auth.bearer_token)
|
||||
.header(REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)
|
||||
.json(&request);
|
||||
|
||||
let response = http_request.send().await.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to enroll remote control server at `{enroll_url}`: {err}"
|
||||
))
|
||||
})?;
|
||||
let headers = response.headers().clone();
|
||||
let status = response.status();
|
||||
let body = response.bytes().await.map_err(|err| {
|
||||
io::Error::other(format!(
|
||||
"failed to read remote control enrollment response from `{enroll_url}`: {err}"
|
||||
))
|
||||
})?;
|
||||
let body_preview = preview_remote_control_response_body(&body);
|
||||
if !status.is_success() {
|
||||
let headers_str = format_headers(&headers);
|
||||
let error_kind = if matches!(status.as_u16(), 401 | 403) {
|
||||
ErrorKind::PermissionDenied
|
||||
} else {
|
||||
ErrorKind::Other
|
||||
};
|
||||
return Err(io::Error::new(
|
||||
error_kind,
|
||||
format!(
|
||||
"remote control server enrollment failed at `{enroll_url}`: HTTP {status}, {headers_str}, body: {body_preview}"
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
let enrollment = serde_json::from_slice::<EnrollRemoteServerResponse>(&body).map_err(|err| {
|
||||
let headers_str = format_headers(&headers);
|
||||
io::Error::other(format!(
|
||||
"failed to parse remote control enrollment response from `{enroll_url}`: HTTP {status}, {headers_str}, body: {body_preview}, decode error: {err}"
|
||||
))
|
||||
})?;
|
||||
|
||||
Ok(RemoteControlEnrollment {
|
||||
account_id: auth.account_id.clone(),
|
||||
environment_id: enrollment.environment_id,
|
||||
server_id: enrollment.server_id,
|
||||
server_name,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::transport::remote_control::protocol::normalize_remote_control_url;
|
||||
use codex_state::StateRuntime;
|
||||
use pretty_assertions::assert_eq;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tempfile::TempDir;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::time::Duration;
|
||||
use tokio::time::timeout;
|
||||
|
||||
async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> {
|
||||
StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string())
|
||||
.await
|
||||
.expect("state runtime should initialize")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persisted_remote_control_enrollment_round_trips_by_target_and_account() {
|
||||
let codex_home = TempDir::new().expect("temp dir should create");
|
||||
let state_db = remote_control_state_runtime(&codex_home).await;
|
||||
let first_target = normalize_remote_control_url("https://chatgpt.com/remote/control")
|
||||
.expect("first target should parse");
|
||||
let second_target =
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/other/control")
|
||||
.expect("second target should parse");
|
||||
let first_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_first".to_string(),
|
||||
server_id: "srv_e_first".to_string(),
|
||||
server_name: "first-server".to_string(),
|
||||
};
|
||||
let second_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_second".to_string(),
|
||||
server_id: "srv_e_second".to_string(),
|
||||
server_name: "second-server".to_string(),
|
||||
};
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
Some(&first_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("first enrollment should persist");
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
Some(&second_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("second enrollment should persist");
|
||||
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await,
|
||||
Some(first_enrollment.clone())
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-b",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await,
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
Some("desktop-client"),
|
||||
)
|
||||
.await,
|
||||
Some(second_enrollment)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn clearing_persisted_remote_control_enrollment_removes_only_matching_entry() {
|
||||
let codex_home = TempDir::new().expect("temp dir should create");
|
||||
let state_db = remote_control_state_runtime(&codex_home).await;
|
||||
let first_target = normalize_remote_control_url("https://chatgpt.com/remote/control")
|
||||
.expect("first target should parse");
|
||||
let second_target =
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/other/control")
|
||||
.expect("second target should parse");
|
||||
let first_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_first".to_string(),
|
||||
server_id: "srv_e_first".to_string(),
|
||||
server_name: "first-server".to_string(),
|
||||
};
|
||||
let second_enrollment = RemoteControlEnrollment {
|
||||
account_id: "account-a".to_string(),
|
||||
environment_id: "env_second".to_string(),
|
||||
server_id: "srv_e_second".to_string(),
|
||||
server_name: "second-server".to_string(),
|
||||
};
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
Some(&first_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("first enrollment should persist");
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
Some(&second_enrollment),
|
||||
)
|
||||
.await
|
||||
.expect("second enrollment should persist");
|
||||
|
||||
update_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
/*enrollment*/ None,
|
||||
)
|
||||
.await
|
||||
.expect("matching enrollment should clear");
|
||||
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&first_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
)
|
||||
.await,
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
load_persisted_remote_control_enrollment(
|
||||
Some(state_db.as_ref()),
|
||||
&second_target,
|
||||
"account-a",
|
||||
/*app_server_client_name*/ None,
|
||||
)
|
||||
.await,
|
||||
Some(second_enrollment)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn enroll_remote_control_server_parse_failure_includes_response_body() {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")
|
||||
.await
|
||||
.expect("listener should bind");
|
||||
let remote_control_url = format!(
|
||||
"http://127.0.0.1:{}/backend-api/",
|
||||
listener
|
||||
.local_addr()
|
||||
.expect("listener should have a local addr")
|
||||
.port()
|
||||
);
|
||||
let remote_control_target =
|
||||
normalize_remote_control_url(&remote_control_url).expect("target should parse");
|
||||
let enroll_url = remote_control_target.enroll_url.clone();
|
||||
let response_body = json!({
|
||||
"error": "not enrolled",
|
||||
});
|
||||
let expected_body = response_body.to_string();
|
||||
let server_task = tokio::spawn(async move {
|
||||
let stream = accept_http_request(&listener).await;
|
||||
respond_with_json(stream, response_body).await;
|
||||
});
|
||||
|
||||
let err = enroll_remote_control_server(
|
||||
&remote_control_target,
|
||||
&RemoteControlConnectionAuth {
|
||||
bearer_token: "Access Token".to_string(),
|
||||
account_id: "account_id".to_string(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("invalid response should fail to parse");
|
||||
|
||||
server_task.await.expect("server task should succeed");
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
format!(
|
||||
"failed to parse remote control enrollment response from `{enroll_url}`: HTTP 200 OK, request-id: <none>, cf-ray: <none>, body: {expected_body}, decode error: missing field `server_id` at line 1 column {}",
|
||||
expected_body.len()
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
async fn accept_http_request(listener: &TcpListener) -> TcpStream {
|
||||
let (stream, _) = timeout(Duration::from_secs(5), listener.accept())
|
||||
.await
|
||||
.expect("HTTP request should arrive in time")
|
||||
.expect("listener accept should succeed");
|
||||
let mut reader = BufReader::new(stream);
|
||||
|
||||
let mut request_line = String::new();
|
||||
reader
|
||||
.read_line(&mut request_line)
|
||||
.await
|
||||
.expect("request line should read");
|
||||
loop {
|
||||
let mut line = String::new();
|
||||
reader
|
||||
.read_line(&mut line)
|
||||
.await
|
||||
.expect("header line should read");
|
||||
if line == "\r\n" {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
reader.into_inner()
|
||||
}
|
||||
|
||||
async fn respond_with_json(mut stream: TcpStream, body: serde_json::Value) {
|
||||
let body = body.to_string();
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
|
||||
body.len()
|
||||
);
|
||||
stream
|
||||
.write_all(response.as_bytes())
|
||||
.await
|
||||
.expect("response should write");
|
||||
stream.flush().await.expect("response should flush");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
mod client_tracker;
|
||||
mod enroll;
|
||||
mod protocol;
|
||||
mod websocket;
|
||||
|
||||
use crate::transport::remote_control::websocket::RemoteControlWebsocket;
|
||||
use crate::transport::remote_control::websocket::load_remote_control_auth;
|
||||
|
||||
pub use self::protocol::ClientId;
|
||||
use self::protocol::ServerEvent;
|
||||
use self::protocol::StreamId;
|
||||
use self::protocol::normalize_remote_control_url;
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::next_connection_id;
|
||||
use codex_login::AuthManager;
|
||||
use codex_state::StateRuntime;
|
||||
use std::io;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::task::JoinHandle;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
pub(super) struct QueuedServerEnvelope {
|
||||
pub(super) event: ServerEvent,
|
||||
pub(super) client_id: ClientId,
|
||||
pub(super) stream_id: StreamId,
|
||||
pub(super) write_complete_tx: Option<oneshot::Sender<()>>,
|
||||
}
|
||||
|
||||
pub(crate) async fn start_remote_control(
|
||||
remote_control_url: String,
|
||||
state_db: Option<Arc<StateRuntime>>,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
shutdown_token: CancellationToken,
|
||||
app_server_client_name_rx: Option<oneshot::Receiver<String>>,
|
||||
) -> io::Result<JoinHandle<()>> {
|
||||
let remote_control_target = normalize_remote_control_url(&remote_control_url)?;
|
||||
validate_remote_control_auth(&auth_manager).await?;
|
||||
|
||||
Ok(tokio::spawn(async move {
|
||||
RemoteControlWebsocket::new(
|
||||
remote_control_target,
|
||||
state_db,
|
||||
auth_manager,
|
||||
transport_event_tx,
|
||||
shutdown_token,
|
||||
)
|
||||
.run(app_server_client_name_rx)
|
||||
.await;
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) async fn validate_remote_control_auth(
|
||||
auth_manager: &Arc<AuthManager>,
|
||||
) -> io::Result<()> {
|
||||
match load_remote_control_auth(auth_manager).await {
|
||||
Ok(_) => Ok(()),
|
||||
Err(err) if err.kind() == io::ErrorKind::WouldBlock => Ok(()),
|
||||
Err(err) => Err(err),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,252 @@
|
||||
use crate::outgoing_message::OutgoingMessage;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use serde::Deserialize;
|
||||
use serde::Serialize;
|
||||
use std::io;
|
||||
use std::io::ErrorKind;
|
||||
use url::Host;
|
||||
use url::Url;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(super) struct RemoteControlTarget {
|
||||
pub(super) websocket_url: String,
|
||||
pub(super) enroll_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub(super) struct EnrollRemoteServerRequest {
|
||||
pub(super) name: String,
|
||||
pub(super) os: &'static str,
|
||||
pub(super) arch: &'static str,
|
||||
pub(super) app_server_version: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(super) struct EnrollRemoteServerResponse {
|
||||
pub(super) server_id: String,
|
||||
pub(super) environment_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct ClientId(pub String);
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
|
||||
#[serde(transparent)]
|
||||
pub struct StreamId(pub String);
|
||||
|
||||
impl StreamId {
|
||||
pub fn new_random() -> Self {
|
||||
Self(uuid::Uuid::now_v7().to_string())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ClientEvent {
|
||||
ClientMessage {
|
||||
message: JSONRPCMessage,
|
||||
},
|
||||
/// Backend-generated acknowledgement for all server envelopes addressed to
|
||||
/// `client_id` whose envelope `seq_id` is less than or equal to this ack's
|
||||
/// `seq_id`. This cursor is client-scoped, not stream-scoped, so receivers
|
||||
/// must not use `stream_id` to partition acks.
|
||||
Ack,
|
||||
Ping,
|
||||
ClientClosed,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) struct ClientEnvelope {
|
||||
#[serde(flatten)]
|
||||
pub(crate) event: ClientEvent,
|
||||
#[serde(rename = "client_id")]
|
||||
pub(crate) client_id: ClientId,
|
||||
#[serde(rename = "stream_id", skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) stream_id: Option<StreamId>,
|
||||
/// For `Ack`, this is the backend-generated per-client cursor over
|
||||
/// `ServerEnvelope.seq_id`.
|
||||
#[serde(rename = "seq_id", skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) seq_id: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) cursor: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum PongStatus {
|
||||
Active,
|
||||
Unknown,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ServerEvent {
|
||||
ServerMessage {
|
||||
message: Box<OutgoingMessage>,
|
||||
},
|
||||
#[allow(dead_code)]
|
||||
Ack,
|
||||
Pong {
|
||||
status: PongStatus,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) struct ServerEnvelope {
|
||||
#[serde(flatten)]
|
||||
pub(crate) event: ServerEvent,
|
||||
#[serde(rename = "client_id")]
|
||||
pub(crate) client_id: ClientId,
|
||||
#[serde(rename = "stream_id")]
|
||||
pub(crate) stream_id: StreamId,
|
||||
#[serde(rename = "seq_id")]
|
||||
pub(crate) seq_id: u64,
|
||||
}
|
||||
|
||||
fn is_allowed_chatgpt_host(host: &Option<Host<&str>>) -> bool {
|
||||
let Some(Host::Domain(host)) = *host else {
|
||||
return false;
|
||||
};
|
||||
host == "chatgpt.com"
|
||||
|| host == "chatgpt-staging.com"
|
||||
|| host.ends_with(".chatgpt.com")
|
||||
|| host.ends_with(".chatgpt-staging.com")
|
||||
}
|
||||
|
||||
fn is_localhost(host: &Option<Host<&str>>) -> bool {
|
||||
match host {
|
||||
Some(Host::Domain("localhost")) => true,
|
||||
Some(Host::Ipv4(ip)) => ip.is_loopback(),
|
||||
Some(Host::Ipv6(ip)) => ip.is_loopback(),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn normalize_remote_control_url(
|
||||
remote_control_url: &str,
|
||||
) -> io::Result<RemoteControlTarget> {
|
||||
let map_url_parse_error = |err: url::ParseError| -> io::Error {
|
||||
io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!("invalid remote control URL `{remote_control_url}`: {err}"),
|
||||
)
|
||||
};
|
||||
let map_scheme_error = |_: ()| -> io::Error {
|
||||
io::Error::new(
|
||||
ErrorKind::InvalidInput,
|
||||
format!(
|
||||
"invalid remote control URL `{remote_control_url}`; expected HTTPS URL for chatgpt.com or chatgpt-staging.com, or HTTP/HTTPS URL for localhost"
|
||||
),
|
||||
)
|
||||
};
|
||||
|
||||
let mut remote_control_url = Url::parse(remote_control_url).map_err(map_url_parse_error)?;
|
||||
if !remote_control_url.path().ends_with('/') {
|
||||
let normalized_path = format!("{}/", remote_control_url.path());
|
||||
remote_control_url.set_path(&normalized_path);
|
||||
}
|
||||
|
||||
let enroll_url = remote_control_url
|
||||
.join("wham/remote/control/server/enroll")
|
||||
.map_err(map_url_parse_error)?;
|
||||
let mut websocket_url = remote_control_url
|
||||
.join("wham/remote/control/server")
|
||||
.map_err(map_url_parse_error)?;
|
||||
let host = enroll_url.host();
|
||||
match enroll_url.scheme() {
|
||||
"https" if is_localhost(&host) || is_allowed_chatgpt_host(&host) => {
|
||||
websocket_url.set_scheme("wss").map_err(map_scheme_error)?;
|
||||
}
|
||||
"http" if is_localhost(&host) => {
|
||||
websocket_url.set_scheme("ws").map_err(map_scheme_error)?;
|
||||
}
|
||||
_ => return Err(map_scheme_error(())),
|
||||
}
|
||||
|
||||
Ok(RemoteControlTarget {
|
||||
websocket_url: websocket_url.to_string(),
|
||||
enroll_url: enroll_url.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use pretty_assertions::assert_eq;
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_accepts_chatgpt_https_urls() {
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://chatgpt.com/backend-api")
|
||||
.expect("chatgpt.com URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "wss://chatgpt.com/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "https://chatgpt.com/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://api.chatgpt-staging.com/backend-api")
|
||||
.expect("chatgpt-staging.com subdomain URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url:
|
||||
"wss://api.chatgpt-staging.com/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url:
|
||||
"https://api.chatgpt-staging.com/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_accepts_localhost_urls() {
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("http://localhost:8080/backend-api")
|
||||
.expect("localhost http URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "ws://localhost:8080/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "http://localhost:8080/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
assert_eq!(
|
||||
normalize_remote_control_url("https://localhost:8443/backend-api")
|
||||
.expect("localhost https URL should normalize"),
|
||||
RemoteControlTarget {
|
||||
websocket_url: "wss://localhost:8443/backend-api/wham/remote/control/server"
|
||||
.to_string(),
|
||||
enroll_url: "https://localhost:8443/backend-api/wham/remote/control/server/enroll"
|
||||
.to_string(),
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn normalize_remote_control_url_rejects_unsupported_urls() {
|
||||
for remote_control_url in [
|
||||
"http://chatgpt.com/backend-api",
|
||||
"http://example.com/backend-api",
|
||||
"https://example.com/backend-api",
|
||||
"https://chatgpt.com.evil.com/backend-api",
|
||||
"https://evilchatgpt.com/backend-api",
|
||||
"https://foo.localhost/backend-api",
|
||||
] {
|
||||
let err = normalize_remote_control_url(remote_control_url)
|
||||
.expect_err("unsupported URL should be rejected");
|
||||
|
||||
assert_eq!(err.kind(), ErrorKind::InvalidInput);
|
||||
assert_eq!(
|
||||
err.to_string(),
|
||||
format!(
|
||||
"invalid remote control URL `{remote_control_url}`; expected HTTPS URL for chatgpt.com or chatgpt-staging.com, or HTTP/HTTPS URL for localhost"
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -1,9 +1,12 @@
|
||||
use super::CHANNEL_CAPACITY;
|
||||
use super::TransportEvent;
|
||||
use super::forward_incoming_message;
|
||||
use super::next_connection_id;
|
||||
use super::serialize_outgoing_message;
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
use codex_app_server_protocol::InitializeParams;
|
||||
use codex_app_server_protocol::JSONRPCMessage;
|
||||
use codex_app_server_protocol::JSONRPCRequest;
|
||||
use std::io::ErrorKind;
|
||||
use std::io::Result as IoResult;
|
||||
use tokio::io;
|
||||
@@ -11,6 +14,7 @@ use tokio::io::AsyncBufReadExt;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::io::BufReader;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::sync::oneshot;
|
||||
use tokio::task::JoinHandle;
|
||||
use tracing::debug;
|
||||
use tracing::error;
|
||||
@@ -19,8 +23,9 @@ use tracing::info;
|
||||
pub(crate) async fn start_stdio_connection(
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
stdio_handles: &mut Vec<JoinHandle<()>>,
|
||||
initialize_client_name_tx: oneshot::Sender<String>,
|
||||
) -> IoResult<()> {
|
||||
let connection_id = ConnectionId(0);
|
||||
let connection_id = next_connection_id();
|
||||
let (writer_tx, mut writer_rx) = mpsc::channel::<QueuedOutgoingMessage>(CHANNEL_CAPACITY);
|
||||
let writer_tx_for_reader = writer_tx.clone();
|
||||
transport_event_tx
|
||||
@@ -37,10 +42,16 @@ pub(crate) async fn start_stdio_connection(
|
||||
let stdin = io::stdin();
|
||||
let reader = BufReader::new(stdin);
|
||||
let mut lines = reader.lines();
|
||||
let mut initialize_client_name_tx = Some(initialize_client_name_tx);
|
||||
|
||||
loop {
|
||||
match lines.next_line().await {
|
||||
Ok(Some(line)) => {
|
||||
if let Some(client_name) = stdio_initialize_client_name(&line)
|
||||
&& let Some(initialize_client_name_tx) = initialize_client_name_tx.take()
|
||||
{
|
||||
let _ = initialize_client_name_tx.send(client_name);
|
||||
}
|
||||
if !forward_incoming_message(
|
||||
&transport_event_tx_for_reader,
|
||||
&writer_tx_for_reader,
|
||||
@@ -86,3 +97,15 @@ pub(crate) async fn start_stdio_connection(
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn stdio_initialize_client_name(line: &str) -> Option<String> {
|
||||
let message = serde_json::from_str::<JSONRPCMessage>(line).ok()?;
|
||||
let JSONRPCMessage::Request(JSONRPCRequest { method, params, .. }) = message else {
|
||||
return None;
|
||||
};
|
||||
if method != "initialize" {
|
||||
return None;
|
||||
}
|
||||
let params = serde_json::from_value::<InitializeParams>(params?).ok()?;
|
||||
Some(params.client_info.name)
|
||||
}
|
||||
|
||||
@@ -4,6 +4,7 @@ use super::auth::WebsocketAuthPolicy;
|
||||
use super::auth::authorize_upgrade;
|
||||
use super::auth::should_warn_about_unauthenticated_non_loopback_listener;
|
||||
use super::forward_incoming_message;
|
||||
use super::next_connection_id;
|
||||
use super::serialize_outgoing_message;
|
||||
use crate::outgoing_message::ConnectionId;
|
||||
use crate::outgoing_message::QueuedOutgoingMessage;
|
||||
@@ -32,8 +33,6 @@ use owo_colors::Style;
|
||||
use std::io::Result as IoResult;
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio::task::JoinHandle;
|
||||
@@ -75,7 +74,6 @@ fn print_websocket_startup_banner(addr: SocketAddr) {
|
||||
#[derive(Clone)]
|
||||
struct WebSocketListenerState {
|
||||
transport_event_tx: mpsc::Sender<TransportEvent>,
|
||||
connection_counter: Arc<AtomicU64>,
|
||||
auth_policy: Arc<WebsocketAuthPolicy>,
|
||||
}
|
||||
|
||||
@@ -113,7 +111,7 @@ async fn websocket_upgrade_handler(
|
||||
);
|
||||
return (err.status_code(), err.message()).into_response();
|
||||
}
|
||||
let connection_id = ConnectionId(state.connection_counter.fetch_add(1, Ordering::Relaxed));
|
||||
let connection_id = next_connection_id();
|
||||
info!(%peer_addr, "websocket client connected");
|
||||
websocket
|
||||
.on_upgrade(move |stream| async move {
|
||||
@@ -146,7 +144,6 @@ pub(crate) async fn start_websocket_acceptor(
|
||||
.layer(middleware::from_fn(reject_requests_with_origin_header))
|
||||
.with_state(WebSocketListenerState {
|
||||
transport_event_tx,
|
||||
connection_counter: Arc::new(AtomicU64::new(1)),
|
||||
auth_policy: Arc::new(auth_policy),
|
||||
});
|
||||
let server = axum::serve(
|
||||
|
||||
Reference in New Issue
Block a user