fix: add qualifier=DEFAULT back

Author: chaodu-agent

Cargo.lock

@@ -27,6 +27,9 @@ image = { version = "0.25", default-features = false, features = ["jpeg", "png",
 unicode-width = "0.2"
 pulldown-cmark = { version = "0.13", default-features = false }
 tokio-tungstenite = { version = "0.21", features = ["rustls-tls-webpki-roots"] }
+rustls = { version = "0.22", optional = true }
+tokio-rustls = { version = "0.25", optional = true }
+webpki-roots = { version = "0.26", optional = true }
 cron = "0.16.0"
 chrono = { version = "0.4.44", features = ["serde"] }
 chrono-tz = "0.10.4"
@@ -43,7 +46,7 @@ http = { version = "1", optional = true }
 [features]
 default = ["secrets-aws", "agentcore"]
 secrets-aws = ["dep:aws-sdk-secretsmanager", "dep:aws-config"]
-agentcore = ["dep:aws-config", "dep:aws-sigv4", "dep:aws-credential-types", "dep:urlencoding", "dep:hex", "dep:http"]
+agentcore = ["dep:aws-config", "dep:aws-sigv4", "dep:aws-credential-types", "dep:urlencoding", "dep:hex", "dep:http", "dep:rustls", "dep:tokio-rustls", "dep:webpki-roots"]
 
 [target.'cfg(unix)'.dependencies]
 libc = "0.2"

Cargo.toml

@@ -108,7 +108,7 @@ struct ShellHandle {
 }
 
 type WsSink = futures_util::stream::SplitSink<
-    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
+    tokio_tungstenite::WebSocketStream<tokio_rustls::client::TlsStream<tokio::net::TcpStream>>,
     Message,
 >;
 
@@ -198,7 +198,7 @@ where
                         .unwrap_or_default()
                         .as_millis();
                     let acp_sid = format!("agentcore-{ts}");
-                    let runtime_sid = format!("oab-session-{ts}");
+                    let runtime_sid = format!("oab-session-{ts:020}-{ts:013x}");
 
                     // Eagerly open shell + initialize the agent
                     match self.open_shell(&runtime_sid).await {
@@ -376,6 +376,7 @@ where
         }
     }
 
+    #[allow(dead_code)]
     fn derive_runtime_session_id(&self, params: &Value) -> String {
         // Try to extract from sender_context in prompt blocks
         if let Some(blocks) = params.get("prompt").and_then(|p| p.as_array()) {
@@ -414,11 +415,38 @@ where
     }
 
     async fn open_shell(&self, session_id: &str) -> Result<ShellHandle> {
-        let request = build_signed_request(&self.runtime_arn, session_id, &self.region).await?;
+        let (request, host) = build_signed_request(&self.runtime_arn, session_id, &self.region).await?;
+
+        // Manual TLS connection — gives us full control, avoids connect_async host override
+        let tcp = tokio::net::TcpStream::connect(format!("{host}:443"))
+            .await
+            .map_err(|e| anyhow!("TCP connect to {host}:443 failed: {e}"))?;
+
+        let connector = tokio_tungstenite::Connector::Rustls(std::sync::Arc::new(
+            rustls::ClientConfig::builder()
+                .with_root_certificates(rustls::RootCertStore {
+                    roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(),
+                })
+                .with_no_client_auth(),
+        ));
+
+        let tls_stream = match connector {
+            tokio_tungstenite::Connector::Rustls(cfg) => {
+                let domain = rustls::pki_types::ServerName::try_from(host.as_str())
+                    .map_err(|e| anyhow!("bad DNS: {e}"))?
+                    .to_owned();
+                tokio_rustls::TlsConnector::from(cfg)
+                    .connect(domain, tcp)
+                    .await
+                    .map_err(|e| anyhow!("TLS failed: {e}"))?
+            }
+            _ => unreachable!(),
+        };
 
-        let (ws_stream, _) = tokio_tungstenite::connect_async(request)
+        // client_async performs the WebSocket upgrade using our exact request
+        let (ws_stream, _) = tokio_tungstenite::client_async(request, tls_stream)
             .await
-            .map_err(|e| anyhow!("WebSocket connect failed: {e}"))?;
+            .map_err(|e| anyhow!("WebSocket upgrade failed: {e}"))?;
 
         info!(session_id, "AgentCore shell connected");
 
@@ -571,7 +599,7 @@ async fn build_signed_request(
     arn: &str,
     session_id: &str,
     region: &str,
-) -> Result<http::Request<()>> {
+) -> Result<(http::Request<()>, String)> {
     let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
         .region(aws_config::Region::new(region.to_string()))
         .load()
@@ -594,15 +622,18 @@ async fn build_signed_request(
     let hash = Sha256::digest(session_id.as_bytes());
     let shell_id = format!("oab-{}", hex::encode(&hash[..8]));
 
+    // Include session_id in query so it's covered by the presigned signature
     let query = format!(
-        "shellId={shell_id}&X-Amzn-Bedrock-AgentCore-Runtime-Session-Id={}",
+        "qualifier=DEFAULT&shellId={shell_id}&X-Amzn-Bedrock-AgentCore-Runtime-Session-Id={}",
         urlencoding::encode(session_id),
     );
-
     let uri = format!("https://{host}{path}?{query}");
 
+    // Use presigned URL (auth in query string) — avoids tungstenite header issues
     let mut settings = SigningSettings::default();
     settings.expires_in = Some(std::time::Duration::from_secs(300));
+    settings.uri_path_normalization_mode =
+        aws_sigv4::http_request::UriPathNormalizationMode::Enabled;
 
     let signing_params = v4::SigningParams::builder()
         .identity(&identity)
@@ -613,30 +644,28 @@ async fn build_signed_request(
         .build()?;
 
     let headers = [("host", host.as_str())];
-
     let signable = SignableRequest::new("GET", &uri, headers.into_iter(), SignableBody::empty())?;
-
     let (instructions, _sig) = sign(signable, &signing_params.into())?.into_parts();
 
-    // Build presigned URL with auth in query params
-    let mut signed_query = query.clone();
-    let (_headers, params) = instructions.into_parts();
+    // Append SigV4 query params to build the presigned URL
+    let mut presigned_query = query.clone();
+    let (_hdrs, params) = instructions.into_parts();
     for (name, value) in params {
-        signed_query.push_str(&format!("&{name}={value}"));
+        presigned_query.push_str(&format!("&{name}={value}"));
     }
-    let wss_uri = format!("wss://{host}{path}?{signed_query}");
 
-    let builder = http::Request::builder()
+    let wss_uri = format!("wss://{host}{path}?{presigned_query}");
+
+    // Build request with WebSocket upgrade headers (not signed, just protocol-required)
+    let request = http::Request::builder()
         .method("GET")
         .uri(&wss_uri)
         .header("host", &host)
         .header("connection", "Upgrade")
         .header("upgrade", "websocket")
         .header("sec-websocket-version", "13")
-        .header(
-            "sec-websocket-key",
-            tokio_tungstenite::tungstenite::handshake::client::generate_key(),
-        );
+        .header("sec-websocket-key", tokio_tungstenite::tungstenite::handshake::client::generate_key())
+        .body(())?;
 
-    Ok(builder.body(())?)
+    Ok((request, host))
 }