diff --git a/src/main.rs b/src/main.rs index b17f1af..f585a7e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -398,12 +398,7 @@ async fn run() -> Result<(), KagiError> { }, } } else { - let query = args.query.ok_or_else(|| { - KagiError::Config( - "assistant prompt mode requires a QUERY unless a thread subcommand is used" - .to_string(), - ) - })?; + let query = read_assistant_prompt_query(args.query)?; let request = AssistantPromptRequest { query, thread_id: args.thread_id, @@ -1165,6 +1160,21 @@ fn build_translate_request(args: TranslateArgs) -> Result) -> Result { + let query = match query { + Some(query) => query, + None => read_stdin_to_string()?.trim().to_string(), + }; + if query.trim().is_empty() { + return Err(KagiError::Config( + "assistant prompt mode requires a QUERY or non-empty stdin unless a thread subcommand is used" + .to_string(), + )); + } + + Ok(query.trim().to_string()) +} + fn normalize_optional_string(value: Option) -> Option { value .map(|value| value.trim().to_string()) diff --git a/tests/integration-cli.rs b/tests/integration-cli.rs index bd2e8ca..4eb17df 100644 --- a/tests/integration-cli.rs +++ b/tests/integration-cli.rs @@ -252,6 +252,47 @@ fn session_env(server: &MockServer) -> Vec<(&'static str, String)> { ] } +#[test] +fn assistant_prompt_stream_reads_query_from_stdin() { + 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") + .json_body(json!({ + "focus": { + "thread_id": null, + "branch_id": "00000000-0000-4000-0000-000000000000", + "prompt": "do a little dance", + "message_id": null, + }, + "profile": {}, + })); + then.status(200) + .header("content-type", "application/vnd.kagi.stream") + .body(concat!( + "hi:{\"v\":\"test\",\"trace\":\"trace-stdin\"}\0\n", + "thread.json:{\"id\":\"thread-stdin\",\"title\":\"Stdin test\",\"ack\":\"2026-06-07T00:00:00Z\",\"created_at\":\"2026-06-07T00:00:00Z\",\"saved\":false,\"shared\":false,\"branch_id\":\"00000000-0000-4000-0000-000000000000\",\"tag_ids\":[]}\0\n", + "new_message.json:{\"id\":\"msg-stdin\",\"thread_id\":\"thread-stdin\",\"created_at\":\"2026-06-07T00:00:00Z\",\"state\":\"streaming\",\"prompt\":\"do a little dance\",\"md\":\"dance\",\"documents\":[]}\0\n", + "new_message.json:{\"id\":\"msg-stdin\",\"thread_id\":\"thread-stdin\",\"created_at\":\"2026-06-07T00:00:00Z\",\"state\":\"done\",\"prompt\":\"do a little dance\",\"md\":\"dance-ok\",\"documents\":[]}\0\n", + )); + }); + + let tempdir = TempDir::new().expect("tempdir"); + let env = session_env(&server); + let output = run_kagi_with_stdin( + &["assistant", "--stream"], + "do a little dance\n", + &env_refs(&env), + tempdir.path(), + ); + + assert_success(&output); + prompt.assert_calls(1); + assert_eq!(String::from_utf8_lossy(&output.stdout), "dance-ok\n"); +} + fn api_meta() -> Value { json!({ "id": "req-1",