spawn tasks for mcp connection (stdio bottleneck)

.github/workflows/integration-tests.yaml

@@ -7,6 +7,14 @@ on:
       - "Cargo.toml"
       - "crates/**"
       - "pctx-py/**"
+      - "scripts/test-mcp-cli.sh"
+      - ".github/workflows/integration-tests.yaml"
+  pull_request:
+    paths:
+      - "Cargo.toml"
+      - "crates/**"
+      - "pctx-py/**"
+      - "scripts/test-mcp-cli.sh"
       - ".github/workflows/integration-tests.yaml"
   workflow_dispatch:
 
@@ -19,7 +27,51 @@ env:
   CARGO_INCREMENTAL: 0
 
 jobs:
-  integration-tests:
+  # CLI integration tests - tests pctx mcp start command
+  cli-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      # Install build dependencies
+      - name: Install build dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential libclang-dev libc6-dev
+
+      # Set up Rust with caching
+      - uses: dtolnay/rust-toolchain@stable
+      - uses: Swatinem/rust-cache@v2
+        with:
+          cache-on-failure: false
+          shared-key: "integration-tests"
+
+      # Build pctx binary in release mode
+      - name: Build pctx
+        run: cargo build --release --bin pctx
+
+      # Run the CLI integration test script with built binary
+      - name: Run CLI integration tests
+        env:
+          PCTX_CMD: ./target/release/pctx
+        run: ./scripts/test-mcp-cli.sh
+
+      # Show server logs on failure
+      - name: Show server logs
+        if: failure()
+        run: |
+          # Find the temp directory log file
+          LOG_FILE=$(find /tmp -name "pctx-test.log" 2>/dev/null | head -1)
+          if [ -n "$LOG_FILE" ] && [ -f "$LOG_FILE" ]; then
+            echo "=== PCTX MCP Server Logs ==="
+            cat "$LOG_FILE"
+          else
+            echo "No log file found"
+          fi
+
+  # Python client integration tests - tests HTTP client
+  python-client-tests:
+    needs: cli-tests
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
@@ -30,14 +82,15 @@ jobs:
           sudo apt-get update
           sudo apt-get install -y build-essential libclang-dev libc6-dev
 
-      # Set up Rust
+      # Set up Rust with caching
       - uses: dtolnay/rust-toolchain@stable
       - uses: Swatinem/rust-cache@v2
         with:
           cache-on-failure: false
+          shared-key: "integration-tests"
 
-      # Build and install pctx server
-      - name: Build pctx server
+      # Build pctx binary in release mode (should hit cache from cli-tests)
+      - name: Build pctx
         run: cargo build --release --bin pctx
 
       # Set up Python
@@ -79,7 +132,7 @@ jobs:
           fi
 
       # Run integration tests
-      - name: Run integration tests
+      - name: Run Python integration tests
         working-directory: pctx-py
         run: uv run pytest --integration -v
 

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help release docs test-python
+.PHONY: help release docs test-python test-cli
 
 # Default target - show help when running just 'make'
 .DEFAULT_GOAL := help
@@ -10,6 +10,7 @@ help:
 	@echo "  make docs                    - Generate CLI and Python documentation"
 	@echo "  make test-python             - Run Python client tests"
 	@echo "  make test-python-integration - Run Python client tests with integration testing"
+	@echo "  make test-cli                - Run CLI integration tests (pctx mcp start)"
 	@echo "  make release                 - Interactive release script (bump version, update changelog)"
 	@echo ""
 
@@ -30,6 +31,10 @@ test-python:
 test-python-integration:
 	@cd pctx-py && uv run pytest tests/ --integration -v
 
+# Run CLI integration tests
+test-cli:
+	@./scripts/test-mcp-cli.sh
+
 # Interactive release workflow
 release:
 	@./release.sh

crates/pctx/src/commands/mcp/start.rs

@@ -2,7 +2,7 @@ use anyhow::Result;
 use clap::Parser;
 use pctx_code_mode::CodeMode;
 use pctx_config::Config;
-use tracing::{debug, info, warn};
+use tracing::{info, warn};
 
 use pctx_mcp_server::PctxMcpServer;
 
@@ -27,26 +27,36 @@ pub struct StartCmd {
 
 impl StartCmd {
     pub(crate) async fn load_code_mode(cfg: &Config) -> Result<CodeMode> {
-        // Connect to each MCP server and fetch their tool definitions
+        // Connect to each MCP server and fetch their tool definitions in parallel
         info!(
-            "Creating code mode interface for {} upstream MCP servers",
+            "Creating code mode interface for {} upstream MCP servers (parallel)",
             cfg.servers.len()
         );
         let mut code_mode = CodeMode::default();
 
-        for server in &cfg.servers {
-            debug!("Creating code mode interface for {}", &server.name);
-            if let Err(e) = code_mode.add_server(server).await {
-                warn!(
-                    err =? e,
-                    server.name =? &server.name,
-                    server.target =? server.display_target(),
-                    "Failed creating creating code mode for `{}` MCP server",
-                    &server.name
-                );
-            }
+        // Use parallel registration with 30 second timeout per server
+        let mut results =
+            pctx_code_mode::parallel_registration::register_servers_parallel(&cfg.servers, 30)
+                .await;
+
+        // Add successful registrations to code_mode
+        let registered = results.add_to_code_mode(&mut code_mode);
+
+        // Log failures
+        for failure in &results.failed {
+            warn!(
+                server.name = failure.server_name,
+                error = failure.error_message,
+                "Failed creating code mode for MCP server"
+            );
         }
 
+        info!(
+            "Code mode initialized with {}/{} MCP servers",
+            registered,
+            cfg.servers.len()
+        );
+
         Ok(code_mode)
     }
 

crates/pctx_code_mode/Cargo.toml

@@ -19,6 +19,7 @@ serde_json = { workspace = true }
 serde = { workspace = true }
 utoipa = { workspace = true }
 tokio = { workspace = true }
+futures = "0.3"
 schemars = "1.1.0"
 
 [dev-dependencies]

crates/pctx_code_mode/src/code_mode.rs

@@ -112,7 +112,7 @@ impl CodeMode {
     ) -> Result<ExecuteOutput> {
         let registry = callback_registry.unwrap_or_default();
         // Format for logging only
-        let formatted_code = codegen::format::format_ts(&code);
+        let formatted_code = codegen::format::format_ts(code);
 
         debug!(
             code_from_llm = %code,
@@ -183,77 +183,6 @@ impl CodeMode {
         })
     }
 
-    // Generates a ToolSet from the given MCP server config
-    pub async fn add_server(&mut self, server: &ServerConfig) -> Result<()> {
-        if self.tool_sets.iter().any(|t| t.name == server.name) {
-            return Err(Error::Message(format!(
-                "ToolSet with name `{}` already exists, MCP servers must have unique names",
-                &server.name
-            )));
-        }
-
-        // initialize and list tools
-        debug!(
-            "Fetching tools from MCP '{}'({})...",
-            &server.name,
-            server.display_target()
-        );
-        let mcp_client = server.connect().await?;
-        debug!(
-            "Successfully connected to '{}', inspecting tools...",
-            server.name
-        );
-        let listed_tools = mcp_client.list_all_tools().await?;
-        debug!("Found {} tools", listed_tools.len());
-
-        // convert tools into codegen tools
-        let mut codegen_tools = vec![];
-        for mcp_tool in listed_tools {
-            let input_schema: codegen::RootSchema =
-                serde_json::from_value(json!(mcp_tool.input_schema)).map_err(|e| {
-                    Error::Message(format!(
-                        "Failed parsing inputSchema as json schema for tool `{}`: {e}",
-                        &mcp_tool.name
-                    ))
-                })?;
-
-            let output_schema = if let Some(o) = mcp_tool.output_schema {
-                Some(
-                    serde_json::from_value::<codegen::RootSchema>(json!(o)).map_err(|e| {
-                        Error::Message(format!(
-                            "Failed parsing outputSchema as json schema for tool `{}`: {e}",
-                            &mcp_tool.name
-                        ))
-                    })?,
-                )
-            } else {
-                None
-            };
-
-            codegen_tools.push(codegen::Tool::new_mcp(
-                &mcp_tool.name,
-                mcp_tool.description.map(String::from),
-                input_schema,
-                output_schema,
-            )?);
-        }
-
-        let description = mcp_client
-            .peer_info()
-            .and_then(|p| p.server_info.title.clone())
-            .unwrap_or(format!("MCP server at {}", server.display_target()));
-
-        // add toolset & it's server configuration
-        self.tool_sets.push(codegen::ToolSet::new(
-            &server.name,
-            &description,
-            codegen_tools,
-        ));
-        self.servers.push(server.clone());
-
-        Ok(())
-    }
-
     // Generates a Tool and add it to the correct Toolset from the given callback config
     pub fn add_callback(&mut self, cfg: &CallbackConfig) -> Result<()> {
         // find the correct toolset & check for clashes

crates/pctx_code_mode/src/lib.rs

@@ -1,5 +1,6 @@
 mod code_mode;
 pub mod model;
+pub mod parallel_registration;
 pub use code_mode::CodeMode;
 use codegen::CodegenError;