fix(assistant): avoid stream response cutoff

Author: Microck

CHANGELOG.md

@@ -11,6 +11,10 @@ Before `1.0.0`, breaking changes may still ship in minor releases.
 
 - `kagi search` and `kagi batch` accept `--limit <N>` to cap the number of results returned (truncated locally; Kagi's search endpoints have no native count parameter)
 
+### Fixed
+
+- `kagi assistant` no longer cuts off long streamed prompt responses at the generic 30 second API timeout
+
 ## [0.5.0]
 
 ### Added

Cargo.lock

@@ -28,7 +28,7 @@ clap_complete = "4.6.2"
 cliclack = "0.5.4"
 console = "0.16.3"
 ctrlc = "3.5.2"
-reqwest = { version = "0.12.15", default-features = false, features = ["json", "multipart", "rustls-tls"] }
+reqwest = { version = "0.12.15", default-features = false, features = ["brotli", "deflate", "gzip", "json", "multipart", "rustls-tls"] }
 scraper = "0.26.0"
 serde = { version = "1.0.228", features = ["derive"] }
 serde_json = "1.0.149"

Cargo.toml

@@ -3173,7 +3173,7 @@ async fn execute_assistant_stream(
         ));
     }
 
-    let client = build_client()?;
+    let client = http::client_assistant_stream()?;
     let response = client
         .post(url)
         .header(header::COOKIE, format!("kagi_session={token}"))
@@ -3200,7 +3200,7 @@ async fn execute_assistant_multipart_stream(
         ));
     }
 
-    let client = build_client()?;
+    let client = http::client_assistant_stream()?;
     let state_json = serde_json::to_vec(state).map_err(|error| {
         KagiError::Config(format!(
             "failed to serialize Assistant prompt upload state: {error}"
@@ -4479,7 +4479,7 @@ mod tests {
         Arc,
         atomic::{AtomicBool, Ordering},
     };
-    use std::time::{SystemTime, UNIX_EPOCH};
+    use std::time::{Duration, SystemTime, UNIX_EPOCH};
     use tempfile::TempDir;
 
     struct ScopedEnvVar {
@@ -4920,6 +4920,19 @@ mod tests {
         assert_eq!(parsed.message.trace_id.as_deref(), Some("trace-message-1"));
     }
 
+    #[test]
+    fn parses_assistant_prompt_stream_without_expires_at() {
+        let raw = concat!(
+            "hi:{\"v\":\"202603091651.stage.c128588\",\"trace\":\"trace-123\"}\0\n",
+            "thread.json:{\"id\":\"thread-1\",\"title\":\"Greeting\",\"ack\":\"2026-03-16T06:19:07Z\",\"created_at\":\"2026-03-16T06:19:07Z\",\"saved\":false,\"shared\":false,\"branch_id\":\"00000000-0000-4000-0000-000000000000\",\"tag_ids\":[]}\0\n",
+            "new_message.json:{\"id\":\"msg-1\",\"thread_id\":\"thread-1\",\"created_at\":\"2026-03-16T06:19:07Z\",\"branch_list\":[\"00000000-0000-4000-0000-000000000000\"],\"state\":\"done\",\"prompt\":\"Hello\",\"reply\":\"<p>Hi</p>\",\"md\":\"Hi\",\"references_html\":\"<ol><li>Doc</li></ol>\",\"references_md\":\"1. [Doc](https://example.com)\",\"metadata\":\"<li>meta</li>\",\"documents\":[],\"trace_id\":\"trace-message-1\"}\0\n"
+        );
+
+        let parsed = parse_assistant_prompt_stream(raw).expect("assistant stream parses");
+        assert_eq!(parsed.thread.id, "thread-1");
+        assert!(parsed.thread.expires_at.is_empty());
+    }
+
     #[test]
     fn parses_assistant_thread_cursor_from_cursor_payload() {
         assert_eq!(
@@ -5172,6 +5185,50 @@ mod tests {
         assert_eq!(response.message.markdown.as_deref(), Some("attached-note"));
     }
 
+    #[tokio::test]
+    #[allow(clippy::await_holding_lock)]
+    async fn assistant_prompt_accepts_delayed_stream_response() {
+        use httpmock::Method::POST;
+        use httpmock::MockServer;
+
+        let server = MockServer::start();
+        let _prompt = server.mock(|when, then| {
+            when.method(POST)
+                .path("/assistant/prompt")
+                .header("cookie", "kagi_session=test-session")
+                .header("accept", "application/vnd.kagi.stream");
+            then.status(200)
+                .header("content-type", "application/vnd.kagi.stream")
+                .delay(Duration::from_millis(200))
+                .body(concat!(
+                    "hi:{\"v\":\"test\",\"trace\":\"trace-delayed\"}\0\n",
+                    "thread.json:{\"id\":\"thread-delayed\",\"title\":\"Delayed test\",\"ack\":\"2026-05-01T00:00:00Z\",\"created_at\":\"2026-05-01T00:00:00Z\",\"saved\":false,\"shared\":false,\"branch_id\":\"00000000-0000-4000-0000-000000000000\",\"tag_ids\":[]}\0\n",
+                    "new_message.json:{\"id\":\"msg-delayed\",\"thread_id\":\"thread-delayed\",\"created_at\":\"2026-05-01T00:00:00Z\",\"state\":\"done\",\"prompt\":\"Hello\",\"md\":\"delayed-ok\",\"documents\":[]}\0\n"
+                ));
+        });
+
+        let _env_guard = lock_env();
+        let _base_url_env = set_env_var("KAGI_BASE_URL", &server.base_url());
+        let response = execute_assistant_prompt(
+            &AssistantPromptRequest {
+                query: "Hello".to_string(),
+                thread_id: None,
+                attachments: Vec::new(),
+                profile_id: None,
+                model: None,
+                lens_id: None,
+                internet_access: None,
+                personalizations: None,
+            },
+            "test-session",
+        )
+        .await
+        .expect("delayed assistant prompt should succeed");
+
+        assert_eq!(response.meta.trace.as_deref(), Some("trace-delayed"));
+        assert_eq!(response.message.markdown.as_deref(), Some("delayed-ok"));
+    }
+
     #[test]
     fn normalizes_custom_bang_trigger_and_redirect_rule() {
         assert_eq!(

src/api.rs

@@ -21,6 +21,7 @@ pub const KAGI_TRANSLATE_BASE_URL_ENV: &str = "KAGI_TRANSLATE_BASE_URL";
 
 static CLIENT_20S: OnceLock<Client> = OnceLock::new();
 static CLIENT_30S: OnceLock<Client> = OnceLock::new();
+static CLIENT_ASSISTANT_STREAM: OnceLock<Client> = OnceLock::new();
 
 /// Returns a shared HTTP client with a 20-second timeout.
 ///
@@ -38,6 +39,30 @@ pub fn client_30s() -> Result<Client, KagiError> {
     cached_client(&CLIENT_30S, Duration::from_secs(30))
 }
 
+/// Returns a shared HTTP client for Kagi Assistant streams.
+///
+/// Assistant responses can legitimately take longer than the short API deadline while the
+/// server continues streaming useful frames. Use connect and per-read timeouts instead of a
+/// total request timeout so long completions are not cut off after 30 seconds.
+///
+/// # Errors
+/// Returns `KagiError::Network` if the client cannot be constructed.
+pub fn client_assistant_stream() -> Result<Client, KagiError> {
+    if let Some(client) = CLIENT_ASSISTANT_STREAM.get() {
+        return Ok(client.clone());
+    }
+
+    let client = Client::builder()
+        .user_agent(USER_AGENT)
+        .connect_timeout(Duration::from_secs(20))
+        .read_timeout(Duration::from_secs(120))
+        .build()
+        .map_err(|error| KagiError::Network(format!("failed to build HTTP client: {error}")))?;
+
+    let _ = CLIENT_ASSISTANT_STREAM.set(client.clone());
+    Ok(client)
+}
+
 /// Maps a `reqwest::Error` to a domain-specific `KagiError`.
 ///
 /// # Arguments