Integrate kontext-dev SDK

README.md

@@ -9,6 +9,40 @@ If you want Codex in your code editor (VS Code, Cursor, Windsurf), <a href="http
 
 ---
 
+## Kontext-Dev fork
+
+This fork wires Codex CLI to the Kontext-Dev MCP server using the `kontext-dev`
+Rust SDK. It uses a single MCP server and refreshes SEARCH/EXECUTE tools before
+each user turn. The diff stays minimal so upstream updates are easy to rebase.
+
+### Configuration (local)
+
+Create agents at https://app.kontext.dev/ to get client credentials and
+visibility into your agents (like this Codex client). Then copy the credentials
+into `~/.codex/config.toml`:
+
+```toml
+[kontext-dev]
+mcp_url = "http://localhost:4000/mcp"
+token_url = "http://127.0.0.1:4444/oauth2/token"
+client_id = "<client_id>"
+client_secret = "<client_secret>"
+scope = "mcp:invoke"
+server_name = "kontext-dev"
+```
+
+Do not commit secrets.
+
+### Running locally
+
+```bash
+cd codex-rs
+cargo run --bin codex
+```
+
+Default branch: `kontext-dev`. `main` mirrors `openai/codex`, and `kontext-dev`
+is rebased onto `main` to stay current.
+
 ## Quickstart
 
 ### Installing and running Codex CLI

codex-rs/Cargo.toml

@@ -74,6 +74,7 @@ codex-chatgpt = { path = "chatgpt" }
 codex-client = { path = "codex-client" }
 codex-common = { path = "common" }
 codex-core = { path = "core" }
+kontext-dev = "0.1.1"
 codex-exec = { path = "exec" }
 codex-execpolicy = { path = "execpolicy" }
 codex-feedback = { path = "feedback" }

codex-rs/core/Cargo.toml

@@ -28,6 +28,7 @@ codex-client = { workspace = true }
 codex-execpolicy = { workspace = true }
 codex-file-search = { workspace = true }
 codex-git = { workspace = true }
+kontext-dev = { workspace = true }
 codex-keyring-store = { workspace = true }
 codex-otel = { workspace = true }
 codex-protocol = { workspace = true }

codex-rs/core/src/codex.rs

@@ -1,10 +1,12 @@
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::fmt::Debug;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU64;
 use std::sync::atomic::Ordering;
+use std::time::Duration;
 
 use crate::AuthManager;
 use crate::CodexAuth;
@@ -20,6 +22,7 @@ use crate::compact_remote::run_inline_remote_auto_compact_task;
 use crate::exec_policy::ExecPolicyManager;
 use crate::features::Feature;
 use crate::features::Features;
+use crate::kontext_dev;
 use crate::models_manager::manager::ModelsManager;
 use crate::parse_command::parse_command;
 use crate::parse_turn_item;
@@ -73,6 +76,7 @@ use tracing::info;
 use tracing::instrument;
 use tracing::trace_span;
 use tracing::warn;
+use uuid::Uuid;
 
 use crate::ModelProviderInfo;
 use crate::WireApi;
@@ -156,6 +160,7 @@ use codex_async_utils::OrCancelExt;
 use codex_otel::OtelManager;
 use codex_protocol::config_types::ReasoningSummary as ReasoningSummaryConfig;
 use codex_protocol::models::ContentItem;
+use codex_protocol::models::FunctionCallOutputPayload;
 use codex_protocol::models::ResponseInputItem;
 use codex_protocol::models::ResponseItem;
 use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
@@ -216,14 +221,19 @@ fn maybe_push_chat_wire_api_deprecation(
 impl Codex {
     /// Spawn a new [`Codex`] and initialize the session.
     pub(crate) async fn spawn(
-        config: Config,
+        mut config: Config,
         auth_manager: Arc<AuthManager>,
         models_manager: Arc<ModelsManager>,
         skills_manager: Arc<SkillsManager>,
         conversation_history: InitialHistory,
         session_source: SessionSource,
         agent_control: AgentControl,
     ) -> CodexResult<CodexSpawnOk> {
+        kontext_dev::attach_kontext_dev_mcp_server(&mut config)
+            .await
+            .map_err(|err| {
+                CodexErr::Fatal(format!("failed to attach Kontext-Dev MCP server: {err:#}"))
+            })?;
         let (tx_sub, rx_sub) = async_channel::bounded(SUBMISSION_CHANNEL_CAPACITY);
         let (tx_event, rx_event) = async_channel::unbounded();
 
@@ -1527,6 +1537,88 @@ impl Session {
             .await;
     }
 
+    async fn prefetch_mcp_tool_discovery(
+        self: &Arc<Self>,
+        turn_context: &TurnContext,
+        cancellation_token: CancellationToken,
+    ) {
+        let mcp_connection_manager = Arc::clone(&self.services.mcp_connection_manager);
+        let tools = match tokio::time::timeout(Duration::from_secs_f64(5.0), async {
+            mcp_connection_manager
+                .read()
+                .await
+                .list_all_tools()
+                .or_cancel(&cancellation_token)
+                .await
+        })
+        .await
+        {
+            Ok(Ok(tools)) => tools,
+            Ok(Err(codex_async_utils::CancelErr::Cancelled)) => return,
+            Err(_) => {
+                warn!("auto SEARCH_TOOLS prefetch: list_all_tools timed out");
+                return;
+            }
+        };
+
+        let search_tool_servers: HashSet<String> = tools
+            .values()
+            .filter(|tool| tool.tool_name == "SEARCH_TOOLS")
+            .map(|tool| tool.server_name.clone())
+            .collect();
+
+        for server in search_tool_servers {
+            let call_id = format!("{server}__SEARCH_TOOLS__{}", Uuid::new_v4().as_simple());
+            let call_arguments = serde_json::json!({
+                "limit": 200,
+            });
+            let call_arguments_str =
+                serde_json::to_string(&call_arguments).unwrap_or_else(|_| "{}".to_string());
+            let call_tool_result =
+                match tokio::time::timeout(Duration::from_secs_f64(15.0), async {
+                    mcp_connection_manager
+                        .read()
+                        .await
+                        .call_tool(&server, "SEARCH_TOOLS", Some(call_arguments.clone()))
+                        .or_cancel(&cancellation_token)
+                        .await
+                })
+                .await
+                {
+                    Ok(Ok(result)) => result,
+                    Ok(Err(codex_async_utils::CancelErr::Cancelled)) => return,
+                    Err(_) => {
+                        warn!("auto SEARCH_TOOLS prefetch for {server} timed out");
+                        continue;
+                    }
+                };
+
+            let call_tool_result = match call_tool_result {
+                Ok(result) => result,
+                Err(error) => {
+                    warn!("auto SEARCH_TOOLS prefetch for {server} failed: {error:#}");
+                    continue;
+                }
+            };
+
+            let call_item = ResponseItem::FunctionCall {
+                id: None,
+                name: format!("mcp__{server}__SEARCH_TOOLS"),
+                arguments: call_arguments_str,
+                call_id: call_id.clone(),
+            };
+            let output_item = ResponseItem::FunctionCallOutput {
+                call_id,
+                output: FunctionCallOutputPayload::from(&call_tool_result),
+            };
+
+            self.record_response_item_and_emit_turn_item(turn_context, call_item)
+                .await;
+            self.record_response_item_and_emit_turn_item(turn_context, output_item)
+                .await;
+        }
+    }
+
     /// Returns the input if there was no task running to inject into
     pub async fn inject_input(&self, input: Vec<UserInput>) -> Result<(), Vec<UserInput>> {
         let mut active = self.active_turn.lock().await;
@@ -2394,6 +2486,8 @@ pub(crate) async fn run_turn(
 
     sess.maybe_start_ghost_snapshot(Arc::clone(&turn_context), cancellation_token.child_token())
         .await;
+    sess.prefetch_mcp_tool_discovery(turn_context.as_ref(), cancellation_token.child_token())
+        .await;
     let mut last_agent_message: Option<String> = None;
     // Although from the perspective of codex.rs, TurnDiffTracker has the lifecycle of a Task which contains
     // many turns, from the perspective of the user, it is a single turn.

codex-rs/core/src/config/mod.rs

@@ -43,6 +43,7 @@ use codex_rmcp_client::OAuthCredentialsStoreMode;
 use codex_utils_absolute_path::AbsolutePathBuf;
 use codex_utils_absolute_path::AbsolutePathBufGuard;
 use dirs::home_dir;
+use kontext_dev::KontextDevConfig;
 use serde::Deserialize;
 use serde::Serialize;
 use similar::DiffableStr;
@@ -259,6 +260,9 @@ pub struct Config {
     /// Definition for MCP servers that Codex can reach out to for tool calls.
     pub mcp_servers: HashMap<String, McpServerConfig>,
 
+    /// Optional Kontext-Dev configuration that can attach a single MCP server.
+    pub kontext_dev: Option<KontextDevConfig>,
+
     /// Preferred store for MCP OAuth credentials.
     /// keyring: Use an OS-specific keyring service.
     ///          Credentials stored in the keyring will only be readable by Codex unless the user explicitly grants access via OS-level keyring access.
@@ -743,6 +747,10 @@ pub struct ConfigToml {
     #[serde(default)]
     pub mcp_servers: HashMap<String, McpServerConfig>,
 
+    /// Kontext-Dev configuration.
+    #[serde(default, rename = "kontext-dev")]
+    pub kontext_dev: Option<KontextDevConfig>,
+
     /// Preferred backend for storing MCP OAuth credentials.
     /// keyring: Use an OS-specific keyring service.
     ///          https://github.com/openai/codex/blob/main/codex-rs/rmcp-client/src/oauth.rs#L2
@@ -1358,6 +1366,7 @@ impl Config {
             // is important in code to differentiate the mode from the store implementation.
             cli_auth_credentials_store_mode: cfg.cli_auth_credentials_store.unwrap_or_default(),
             mcp_servers: cfg.mcp_servers,
+            kontext_dev: cfg.kontext_dev,
             // The config.toml omits "_mode" because it's a config file. However, "_mode"
             // is important in code to differentiate the mode from the store implementation.
             mcp_oauth_credentials_store_mode: cfg.mcp_oauth_credentials_store.unwrap_or_default(),
@@ -3244,6 +3253,7 @@ model_verbosity = "high"
                 cwd: fixture.cwd(),
                 cli_auth_credentials_store_mode: Default::default(),
                 mcp_servers: HashMap::new(),
+                kontext_dev: None,
                 mcp_oauth_credentials_store_mode: Default::default(),
                 model_providers: fixture.model_provider_map.clone(),
                 project_doc_max_bytes: PROJECT_DOC_MAX_BYTES,
@@ -3330,6 +3340,7 @@ model_verbosity = "high"
             cwd: fixture.cwd(),
             cli_auth_credentials_store_mode: Default::default(),
             mcp_servers: HashMap::new(),
+            kontext_dev: None,
             mcp_oauth_credentials_store_mode: Default::default(),
             model_providers: fixture.model_provider_map.clone(),
             project_doc_max_bytes: PROJECT_DOC_MAX_BYTES,
@@ -3431,6 +3442,7 @@ model_verbosity = "high"
             cwd: fixture.cwd(),
             cli_auth_credentials_store_mode: Default::default(),
             mcp_servers: HashMap::new(),
+            kontext_dev: None,
             mcp_oauth_credentials_store_mode: Default::default(),
             model_providers: fixture.model_provider_map.clone(),
             project_doc_max_bytes: PROJECT_DOC_MAX_BYTES,
@@ -3518,6 +3530,7 @@ model_verbosity = "high"
             cwd: fixture.cwd(),
             cli_auth_credentials_store_mode: Default::default(),
             mcp_servers: HashMap::new(),
+            kontext_dev: None,
             mcp_oauth_credentials_store_mode: Default::default(),
             model_providers: fixture.model_provider_map.clone(),
             project_doc_max_bytes: PROJECT_DOC_MAX_BYTES,

codex-rs/core/src/kontext_dev.rs

@@ -0,0 +1,69 @@
+use std::sync::OnceLock;
+use std::time::Duration;
+use std::time::Instant;
+
+use anyhow::Result;
+use tracing::debug;
+use tracing::info;
+
+use crate::config::Config;
+use crate::config::types::McpServerConfig;
+use crate::config::types::McpServerTransportConfig;
+use kontext_dev::build_mcp_url;
+use kontext_dev::request_access_token;
+
+const DEFAULT_TOKEN_TTL_SECONDS: i64 = 3600;
+
+static KONTEXT_DEV_TOKEN_EXPIRES_AT: OnceLock<Instant> = OnceLock::new();
+static KONTEXT_DEV_SERVER_NAME: OnceLock<String> = OnceLock::new();
+
+pub(crate) async fn attach_kontext_dev_mcp_server(config: &mut Config) -> Result<()> {
+    let Some(settings) = config.kontext_dev.clone() else {
+        debug!("Kontext-Dev not configured; skipping attachment.");
+        return Ok(());
+    };
+
+    let token = request_access_token(&settings).await?;
+    let url = build_mcp_url(&settings, token.access_token.as_str())?;
+    let server_name = settings.server_name.clone();
+
+    let transport = McpServerTransportConfig::StreamableHttp {
+        url,
+        bearer_token_env_var: None,
+        http_headers: None,
+        env_http_headers: None,
+    };
+
+    let server_config = McpServerConfig {
+        transport,
+        startup_timeout_sec: Some(Duration::from_secs_f64(30.0)),
+        tool_timeout_sec: None,
+        enabled: true,
+        enabled_tools: None,
+        disabled_tools: None,
+    };
+
+    config
+        .mcp_servers
+        .insert(server_name.clone(), server_config);
+
+    let expires_in = token.expires_in.unwrap_or(DEFAULT_TOKEN_TTL_SECONDS);
+    let expires_in = expires_in.max(0);
+    let expires_at = Instant::now() + Duration::from_secs_f64(expires_in as f64);
+    let _ = KONTEXT_DEV_TOKEN_EXPIRES_AT.set(expires_at);
+    let _ = KONTEXT_DEV_SERVER_NAME.set(server_name.clone());
+
+    info!("Attached Kontext-Dev MCP server '{server_name}'.");
+    Ok(())
+}
+
+pub(crate) fn kontext_dev_server_name() -> Option<&'static str> {
+    KONTEXT_DEV_SERVER_NAME.get().map(String::as_str)
+}
+
+pub(crate) fn kontext_dev_token_expired() -> bool {
+    KONTEXT_DEV_TOKEN_EXPIRES_AT
+        .get()
+        .map(|expires_at| Instant::now() >= *expires_at)
+        .unwrap_or(false)
+}

codex-rs/core/src/lib.rs

@@ -31,6 +31,7 @@ mod exec_policy;
 pub mod features;
 mod flags;
 pub mod git_info;
+mod kontext_dev;
 pub mod landlock;
 pub mod mcp;
 mod mcp_connection_manager;