From 25e7f1cb5091ee8681bcf2411984828ef80b4541 Mon Sep 17 00:00:00 2001 From: Microck Date: Thu, 28 May 2026 10:19:09 +0000 Subject: [PATCH] feat(cli): add auth fallback, completion install, and assistant streaming --- README.md | 23 ++- docs/SKILL.md | 21 ++- docs/commands/assistant.mdx | 24 +++ docs/commands/completion.mdx | 69 ++++++++ docs/commands/search.mdx | 2 +- docs/guides/installation.mdx | 26 ++- docs/llms.txt | 1 + docs/reference/coverage.mdx | 1 + docs/reference/output-contract.mdx | 2 + src/cli.rs | 53 ++++++- src/main.rs | 246 +++++++++++++++++++++++++++-- tests/integration-cli.rs | 123 ++++++++++++++- 12 files changed, 556 insertions(+), 35 deletions(-) create mode 100644 docs/commands/completion.mdx diff --git a/README.md b/README.md index a7a264b..85e228b 100644 --- a/README.md +++ b/README.md @@ -157,6 +157,7 @@ notes: - environment variables override `.kagi.toml` - base `kagi search` defaults to the session-token path when both credentials are present - set `[auth] preferred_auth = "api"` if you want base search to prefer the API path instead +- API-first base search falls back to the session-token path when the API key is rejected, including quota and rate-limit failures - `search --lens`, `--time`, `--order`, `--verbatim`, and personalization flags require `KAGI_SESSION_TOKEN` - `search --region`, `--from-date`, and `--to-date` use V1 API filters when routed through `KAGI_API_KEY` - `auth check` validates the selected primary credential without using search fallback logic @@ -170,6 +171,7 @@ for the full command-to-token matrix, use the [`auth-matrix`](https://kagi.micr. | `kagi search` | search Kagi with `json` by default, or render as `toon`, `pretty`, `compact`, `markdown`, or `csv` | | `kagi batch` | run multiple searches in parallel with JSON, TOON, compact, pretty, markdown, or csv output and shared filters | | `kagi auth` | launch the auth wizard, or inspect, validate, and save credentials | +| `kagi completion` | generate or install shell completions for bash, zsh, fish, or PowerShell | | `kagi summarize` | use the paid public summarizer API or the subscriber summarizer with `--subscriber` | | `kagi extract` | extract a page's full content as markdown through the current paid API, using `KAGI_API_KEY` directly | | `kagi watch` | rerun a search on an interval and emit added/removed result URLs | @@ -193,22 +195,31 @@ for automation, stdout stays JSON by default. Use `--format toon` for token-effi ## shell completion -generate a completion script and install it with your shell of choice: +install completions automatically for your detected shell: + +```bash +kagi completion install +kagi completion install --shell fish +``` + +or generate a completion script and install it yourself: ```bash # bash -kagi --generate-completion bash > ~/.local/share/bash-completion/completions/kagi +kagi completion generate bash > ~/.local/share/bash-completion/completions/kagi # zsh -kagi --generate-completion zsh > ~/.zsh/completion/_kagi +kagi completion generate zsh > ~/.zsh/completions/_kagi # fish -kagi --generate-completion fish > ~/.config/fish/completions/kagi.fish +kagi completion generate fish > ~/.config/fish/completions/kagi.fish # powershell -kagi --generate-completion powershell >> $PROFILE +kagi completion generate powershell >> $PROFILE ``` +`kagi --generate-completion ` remains available as a shortcut for script generation. + see the [installation guide](https://kagi.micr.dev/guides/installation) for platform-specific setup details. ## examples @@ -296,6 +307,8 @@ continue research with assistant: ```bash kagi assistant "plan a focused research session in the terminal" +kagi assistant --stream "write a short implementation plan" +kagi assistant --stream --stream-output json "write a short implementation plan" ``` run assistant with a saved assistant profile and markdown output: diff --git a/docs/SKILL.md b/docs/SKILL.md index 505ffe6..5e0c173 100644 --- a/docs/SKILL.md +++ b/docs/SKILL.md @@ -180,6 +180,12 @@ kagi assistant --thread-id "" "Give me an example" # Use a saved assistant profile with prompt overrides kagi assistant --assistant research --model gpt-5-mini --web-access --no-personalized "Summarize the latest Rust release" +# Stream markdown deltas as they arrive +kagi assistant --stream "Explain quantum computing" + +# Stream structured events for scripts +kagi assistant --stream --stream-output json "Explain quantum computing" + # List threads kagi assistant thread list @@ -344,20 +350,27 @@ kagi search "query" --format csv > results.csv ## Shell Completions +Install completions for the detected shell: + +```bash +kagi completion install +kagi completion install --shell fish +``` + Generate completion scripts for Bash, Zsh, Fish, and PowerShell: ```bash # Bash -kagi --generate-completion bash > ~/.local/share/bash-completion/completions/kagi +kagi completion generate bash > ~/.local/share/bash-completion/completions/kagi # Zsh -kagi --generate-completion zsh > ~/.zsh/completion/_kagi +kagi completion generate zsh > ~/.zsh/completions/_kagi # Fish -kagi --generate-completion fish > ~/.config/fish/completions/kagi.fish +kagi completion generate fish > ~/.config/fish/completions/kagi.fish # PowerShell -kagi --generate-completion powershell >> $PROFILE +kagi completion generate powershell >> $PROFILE ``` ## Common Workflows diff --git a/docs/commands/assistant.mdx b/docs/commands/assistant.mdx index de7e7e0..34646b3 100644 --- a/docs/commands/assistant.mdx +++ b/docs/commands/assistant.mdx @@ -13,6 +13,7 @@ Prompt Kagi Assistant, continue an existing thread, use a saved assistant profil ```bash kagi assistant [OPTIONS] +kagi assistant --stream [--stream-output text|json] kagi assistant thread list kagi assistant thread get kagi assistant thread delete @@ -79,6 +80,29 @@ kagi assistant --format markdown "Write a release summary" Disable ANSI colors in `--format pretty`. +#### `--stream` + +Stream Assistant output as it arrives instead of waiting for the final response. By default this writes incremental markdown deltas to stdout and flushes after each update: + +```bash +kagi assistant --stream "Write a 3-line release note" +``` + +Use JSON mode when a script needs structured stream events: + +```bash +kagi assistant --stream --stream-output json "Write a 3-line release note" +``` + +#### `--stream-output ` + +Select the streamed output mode. This option requires `--stream`. + +Possible values: + +- `text` - default, prints only incremental markdown deltas +- `json` - newline-delimited compact JSON events with `md_delta`, `thread`, `message`, and `meta` + #### `--model ` Override the model slug for a single prompt. diff --git a/docs/commands/completion.mdx b/docs/commands/completion.mdx new file mode 100644 index 0000000..21fd401 --- /dev/null +++ b/docs/commands/completion.mdx @@ -0,0 +1,69 @@ +--- +title: "completion" +description: "Complete reference for *kagi* completion command - generate and install shell completion scripts." +--- + +# `kagi completion` + +Generate or install shell completion scripts for `kagi`. + +## Synopsis + +```bash +kagi completion generate +kagi completion install [--shell ] [--dir ] +``` + +Supported shells: + +- `bash` +- `zsh` +- `fish` +- `powershell` + +## `kagi completion install` + +Install a completion script for your active shell. When `--shell` is omitted, `kagi` detects the shell from environment variables such as `SHELL`. + +```bash +kagi completion install +``` + +Choose the shell explicitly when detection is unavailable or when installing for a different shell: + +```bash +kagi completion install --shell bash +kagi completion install --shell zsh +kagi completion install --shell fish +kagi completion install --shell powershell +``` + +Default install targets: + +| Shell | Target | +| --- | --- | +| bash | `$XDG_DATA_HOME/bash-completion/completions/kagi`, or `~/.local/share/bash-completion/completions/kagi` | +| zsh | `~/.zsh/completions/_kagi` | +| fish | `$XDG_CONFIG_HOME/fish/completions/kagi.fish`, or `~/.config/fish/completions/kagi.fish` | +| powershell | `$XDG_CONFIG_HOME/powershell/kagi-completions.ps1`, or `~/.config/powershell/kagi-completions.ps1` | + +For zsh, the installer also adds the completion directory to `~/.zshrc` and ensures `compinit` is loaded. For PowerShell, it adds a dot-source line to the standard profile path under the same config directory. + +Use `--dir` to write the completion file somewhere else: + +```bash +kagi completion install --shell fish --dir ./completions +``` + +## `kagi completion generate` + +Print a completion script to stdout without changing any files: + +```bash +kagi completion generate bash > ~/.local/share/bash-completion/completions/kagi +kagi completion generate zsh > ~/.zsh/completions/_kagi +kagi completion generate fish > ~/.config/fish/completions/kagi.fish +kagi completion generate powershell >> $PROFILE +``` + +`kagi --generate-completion ` is still available as a shortcut for script generation. diff --git a/docs/commands/search.mdx b/docs/commands/search.mdx index 82ee877..6d3c207 100644 --- a/docs/commands/search.mdx +++ b/docs/commands/search.mdx @@ -39,7 +39,7 @@ Plain base search keeps the existing repo behavior: - default session-first search when a session token is available - API-first only when you configure `[auth.preferred_auth] = "api"` and provide `KAGI_API_KEY` -- session fallback only when the API path was selected first and rejected +- session fallback only when the API path was selected first and rejected, including `429 Too Many Requests` and quota-style API failures ### Session-Only Search Options diff --git a/docs/guides/installation.mdx b/docs/guides/installation.mdx index 18ce048..178b58c 100644 --- a/docs/guides/installation.mdx +++ b/docs/guides/installation.mdx @@ -369,19 +369,31 @@ The binary is statically linked and has no runtime dependencies beyond the Linux **Shell Completion:** -To enable tab completion, add to your shell configuration: +To enable tab completion, install a generated completion script for your detected shell: ```bash -# Bash (~/.bashrc) -eval "$(kagi --generate-completion bash)" +kagi completion install +``` -# Zsh (~/.zshrc) -eval "$(kagi --generate-completion zsh)" +You can also choose the shell explicitly: -# Fish (~/.config/fish/config.fish) -kagi --generate-completion fish | source +```bash +kagi completion install --shell bash +kagi completion install --shell zsh +kagi completion install --shell fish +kagi completion install --shell powershell ``` +For manual installs, write the generated script yourself: + +```bash +kagi completion generate bash > ~/.local/share/bash-completion/completions/kagi +kagi completion generate zsh > ~/.zsh/completions/_kagi +kagi completion generate fish > ~/.config/fish/completions/kagi.fish +``` + +The older `kagi --generate-completion ` shortcut still prints the generated script to stdout. + ### Windows **PowerShell Execution Policy:** diff --git a/docs/llms.txt b/docs/llms.txt index 7fcfec8..3e66f8a 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -16,6 +16,7 @@ - [search](https://kagi.micr.dev/commands/search): Complete reference for *kagi* search command - syntax, flags, authentication, output formats, and real-world examples. - [auth](https://kagi.micr.dev/commands/auth): Complete reference for *kagi* auth command - credential management, validation, and configuration. +- [completion](https://kagi.micr.dev/commands/completion): Complete reference for *kagi* completion command - generate and install shell completion scripts. - [summarize](https://kagi.micr.dev/commands/summarize): Complete reference for *kagi* summarize command - summarize URLs and text using subscriber or public API modes. - [news](https://kagi.micr.dev/commands/news): Complete reference for *kagi* news command - fetch Kagi News from public JSON endpoints with category filtering. - [smallweb](https://kagi.micr.dev/commands/smallweb): Complete reference for *kagi* smallweb command - fetch the Kagi Small Web feed of independent websites. diff --git a/docs/reference/coverage.mdx b/docs/reference/coverage.mdx index 0581bab..652433d 100644 --- a/docs/reference/coverage.mdx +++ b/docs/reference/coverage.mdx @@ -78,6 +78,7 @@ These require no authentication: | `batch` | Parallel search with shared search options | API or Session | ✅ | | `batch` via stdin | Parallel search from stdin lines | API or Session | ✅ | | `auth` | Credential management | None | ✅ | +| `completion` | Generate and install shell completions | None | Implemented | | `--profile` | Named auth profiles from `.kagi.toml` | None | ✅ | | `summarize` | Public API summarizer | API | ✅ | | `summarize --subscriber` | Web summarizer | Session | ✅ | diff --git a/docs/reference/output-contract.mdx b/docs/reference/output-contract.mdx index 3c29bdb..ebd634b 100644 --- a/docs/reference/output-contract.mdx +++ b/docs/reference/output-contract.mdx @@ -164,6 +164,8 @@ Subscriber mode: } ``` +`kagi assistant --stream` writes incremental markdown deltas to stdout and flushes after each update. `kagi assistant --stream --stream-output json` writes the same stream as newline-delimited compact JSON events with an `md_delta` field plus the current `meta`, `thread`, and `message` snapshot. + ### `kagi assistant custom` `list` returns an array of assistant summaries: diff --git a/src/cli.rs b/src/cli.rs index 7593ec8..b5972e7 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -17,6 +17,15 @@ pub enum CompletionShell { PowerShell, } +#[derive(Debug, Clone, ValueEnum)] +/// Output mode for streamed assistant prompt updates. +pub enum AssistantStreamOutput { + /// Write only incremental markdown deltas to stdout. + Text, + /// Write each stream update as one compact JSON object per line. + Json, +} + #[derive(Debug, Clone, ValueEnum)] /// Output format for assistant thread exports. pub enum AssistantThreadExportFormat { @@ -251,6 +260,8 @@ pub enum Commands { History(HistoryCommand), /// Manage local domain preferences used by CLI search output SitePref(SitePrefCommand), + /// Generate or install shell completion scripts + Completion(CompletionCommand), /// Manage Kagi search lenses Lens(LensCommand), /// Manage custom bangs @@ -551,6 +562,42 @@ pub struct AuthSetArgs { pub session_token: Option, } +#[derive(Debug, Args)] +/// Arguments for the `completion` command group. +pub struct CompletionCommand { + #[command(subcommand)] + pub command: CompletionSubcommand, +} + +#[derive(Debug, Subcommand)] +/// Subcommands for shell completion management. +pub enum CompletionSubcommand { + /// Print a completion script to stdout + Generate(CompletionGenerateArgs), + /// Install a completion script for the active shell + Install(CompletionInstallArgs), +} + +#[derive(Debug, Args)] +/// Arguments for `completion generate`. +pub struct CompletionGenerateArgs { + /// Shell to generate completions for + #[arg(value_name = "SHELL", value_enum)] + pub shell: CompletionShell, +} + +#[derive(Debug, Args)] +/// Arguments for `completion install`. +pub struct CompletionInstallArgs { + /// Shell to install completions for; detected from SHELL when omitted + #[arg(long, value_name = "SHELL", value_enum)] + pub shell: Option, + + /// Override the completion directory + #[arg(long, value_name = "DIR")] + pub dir: Option, +} + #[derive(Debug, Args)] /// Arguments for the `summarize` subcommand. pub struct SummarizeArgs { @@ -755,10 +802,14 @@ pub struct AssistantArgs { #[arg(long, value_name = "FORMAT", default_value_t = AssistantOutputFormat::Json)] pub format: AssistantOutputFormat, - /// Emit prompt updates as newline-delimited JSON + /// Stream prompt updates as text deltas or newline-delimited JSON #[arg(long, conflicts_with = "export")] pub stream: bool, + /// Stream output mode used with --stream + #[arg(long, value_name = "MODE", value_enum, default_value_t = AssistantStreamOutput::Text, requires = "stream")] + pub stream_output: AssistantStreamOutput, + /// Disable colored terminal output (only affects pretty format) #[arg(long)] pub no_color: bool, diff --git a/src/main.rs b/src/main.rs index 424b0f6..7ee5cca 100644 --- a/src/main.rs +++ b/src/main.rs @@ -37,9 +37,10 @@ use crate::auth::{ }; use crate::auth_wizard::{run_auth_wizard, supports_interactive_auth, validate_credential}; use crate::cli::{ - AssistantCustomSubcommand, AssistantOutputFormat, AssistantReplArgs, AssistantSubcommand, - AssistantThreadExportFormat, AssistantThreadSubcommand, AuthSetArgs, AuthSubcommand, - BangSubcommand, Cli, Commands, CompletionShell, CustomBangSubcommand, EnrichSubcommand, + AssistantCustomSubcommand, AssistantOutputFormat, AssistantReplArgs, AssistantStreamOutput, + AssistantSubcommand, AssistantThreadExportFormat, AssistantThreadSubcommand, AuthSetArgs, + AuthSubcommand, BangSubcommand, Cli, Commands, CompletionCommand, CompletionInstallArgs, + CompletionShell, CompletionSubcommand, CustomBangSubcommand, EnrichSubcommand, HistorySubcommand, McpArgs, NotifyArgs, OutputFormat, SearchArgs, SearchOrder, SearchTime, SitePrefMode, SitePrefSubcommand, TranslateArgs, WatchArgs, }; @@ -60,6 +61,7 @@ use std::env; use std::fs; use std::future::Future; use std::io::{self, BufRead, Read, Write}; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::Semaphore; @@ -192,6 +194,7 @@ async fn run() -> Result<(), KagiError> { AuthSubcommand::Check => run_auth_check(profile.as_deref()).await, AuthSubcommand::Set(args) => run_auth_set(args, profile.as_deref()), }, + Commands::Completion(args) => run_completion(args), Commands::Summarize(args) => { args.validate().map_err(KagiError::Config)?; @@ -419,18 +422,20 @@ async fn run() -> Result<(), KagiError> { }, }; if args.once { - let response = - execute_once_assistant_prompt(&request, args.stream, &token).await?; + let response = execute_once_assistant_prompt( + &request, + args.stream.then_some(args.stream_output), + &token, + ) + .await?; if args.stream { Ok(()) } else { print_assistant_response(&response, args.format, !args.no_color) } } else if args.stream { - execute_assistant_prompt_stream(&request, &token, |event| { - print_compact_json(event) - }) - .await?; + execute_streaming_assistant_prompt(&request, &token, args.stream_output) + .await?; Ok(()) } else { let response = execute_assistant_prompt(&request, &token).await?; @@ -795,16 +800,193 @@ fn is_bare_auth_invocation_from(args: &[&str]) -> bool { } fn print_completion(shell: CompletionShell) { + let output = completion_script(shell); + print!("{output}"); +} + +fn run_completion(args: CompletionCommand) -> Result<(), KagiError> { + match args.command { + CompletionSubcommand::Generate(args) => { + print_completion(args.shell); + Ok(()) + } + CompletionSubcommand::Install(args) => install_completion(args), + } +} + +fn completion_script(shell: CompletionShell) -> String { let mut cmd = Cli::command(); + let mut buffer = Vec::new(); match shell { - CompletionShell::Bash => generate(shells::Bash, &mut cmd, "kagi", &mut std::io::stdout()), - CompletionShell::Zsh => generate(shells::Zsh, &mut cmd, "kagi", &mut std::io::stdout()), - CompletionShell::Fish => generate(shells::Fish, &mut cmd, "kagi", &mut std::io::stdout()), + CompletionShell::Bash => generate(shells::Bash, &mut cmd, "kagi", &mut buffer), + CompletionShell::Zsh => generate(shells::Zsh, &mut cmd, "kagi", &mut buffer), + CompletionShell::Fish => generate(shells::Fish, &mut cmd, "kagi", &mut buffer), CompletionShell::PowerShell => { - generate(shells::PowerShell, &mut cmd, "kagi", &mut std::io::stdout()); + generate(shells::PowerShell, &mut cmd, "kagi", &mut buffer); } } + + String::from_utf8(buffer).expect("clap completion scripts are valid UTF-8") +} + +fn install_completion(args: CompletionInstallArgs) -> Result<(), KagiError> { + let shell = args.shell.or_else(detect_completion_shell).ok_or_else(|| { + KagiError::Config( + "could not detect shell; rerun with `kagi completion install --shell `" + .to_string(), + ) + })?; + let target_dir = args.dir.unwrap_or_else(|| default_completion_dir(&shell)); + let target_path = target_dir.join(completion_filename(&shell)); + + fs::create_dir_all(&target_dir).map_err(|error| { + KagiError::Config(format!( + "failed to create completion directory {}: {error}", + target_dir.display() + )) + })?; + fs::write(&target_path, completion_script(shell.clone())).map_err(|error| { + KagiError::Config(format!( + "failed to write completion file {}: {error}", + target_path.display() + )) + })?; + + maybe_update_zshrc(&shell, &target_dir)?; + maybe_update_powershell_profile(&shell, &target_path)?; + + println!( + "installed {shell_name} completions to {}", + target_path.display(), + shell_name = completion_shell_name(&shell) + ); + Ok(()) +} + +fn detect_completion_shell() -> Option { + let shell = env::var("SHELL") + .ok() + .or_else(|| env::var("ComSpec").ok()) + .or_else(|| env::var("PSModulePath").ok())?; + let name = Path::new(&shell) + .file_stem() + .and_then(|value| value.to_str()) + .unwrap_or(shell.as_str()) + .to_ascii_lowercase(); + + match name.as_str() { + "bash" => Some(CompletionShell::Bash), + "zsh" => Some(CompletionShell::Zsh), + "fish" => Some(CompletionShell::Fish), + "pwsh" | "powershell" | "powershell.exe" | "pwsh.exe" => Some(CompletionShell::PowerShell), + _ => None, + } +} + +fn default_completion_dir(shell: &CompletionShell) -> PathBuf { + match shell { + CompletionShell::Bash => xdg_data_home().join("bash-completion").join("completions"), + CompletionShell::Zsh => home_dir().join(".zsh").join("completions"), + CompletionShell::Fish => xdg_config_home().join("fish").join("completions"), + CompletionShell::PowerShell => xdg_config_home().join("powershell"), + } +} + +fn completion_filename(shell: &CompletionShell) -> &'static str { + match shell { + CompletionShell::Bash => "kagi", + CompletionShell::Zsh => "_kagi", + CompletionShell::Fish => "kagi.fish", + CompletionShell::PowerShell => "kagi-completions.ps1", + } +} + +fn completion_shell_name(shell: &CompletionShell) -> &'static str { + match shell { + CompletionShell::Bash => "bash", + CompletionShell::Zsh => "zsh", + CompletionShell::Fish => "fish", + CompletionShell::PowerShell => "powershell", + } +} + +fn maybe_update_zshrc(shell: &CompletionShell, target_dir: &Path) -> Result<(), KagiError> { + if !matches!(shell, CompletionShell::Zsh) { + return Ok(()); + } + + let zshrc = home_dir().join(".zshrc"); + let line = format!("fpath=({} $fpath)", target_dir.display()); + append_line_if_missing(&zshrc, &line)?; + append_line_if_missing(&zshrc, "autoload -Uz compinit && compinit")?; + Ok(()) +} + +fn maybe_update_powershell_profile( + shell: &CompletionShell, + target_path: &Path, +) -> Result<(), KagiError> { + if !matches!(shell, CompletionShell::PowerShell) { + return Ok(()); + } + + let profile = xdg_config_home() + .join("powershell") + .join("Microsoft.PowerShell_profile.ps1"); + let line = format!(". '{}'", target_path.display()); + append_line_if_missing(&profile, &line) +} + +fn append_line_if_missing(path: &Path, line: &str) -> Result<(), KagiError> { + let existing = match fs::read_to_string(path) { + Ok(raw) => raw, + Err(error) if error.kind() == io::ErrorKind::NotFound => String::new(), + Err(error) => { + return Err(KagiError::Config(format!( + "failed to read {}: {error}", + path.display() + ))); + } + }; + + if existing.lines().any(|existing_line| existing_line == line) { + return Ok(()); + } + + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).map_err(|error| { + KagiError::Config(format!("failed to create {}: {error}", parent.display())) + })?; + } + + let mut updated = existing; + if !updated.is_empty() && !updated.ends_with('\n') { + updated.push('\n'); + } + updated.push_str(line); + updated.push('\n'); + fs::write(path, updated) + .map_err(|error| KagiError::Config(format!("failed to update {}: {error}", path.display()))) +} + +fn xdg_data_home() -> PathBuf { + env::var_os("XDG_DATA_HOME") + .map(PathBuf::from) + .unwrap_or_else(|| home_dir().join(".local").join("share")) +} + +fn xdg_config_home() -> PathBuf { + env::var_os("XDG_CONFIG_HOME") + .map(PathBuf::from) + .unwrap_or_else(|| home_dir().join(".config")) +} + +fn home_dir() -> PathBuf { + env::var_os("HOME") + .map(PathBuf::from) + .or_else(|| env::var_os("USERPROFILE").map(PathBuf::from)) + .unwrap_or_else(|| PathBuf::from(".")) } fn run_auth_status(profile: Option<&str>) -> Result<(), KagiError> { @@ -1070,7 +1252,7 @@ fn print_json(value: &T) -> Result<(), KagiError> { async fn execute_once_assistant_prompt( request: &AssistantPromptRequest, - stream: bool, + stream_output: Option, token: &str, ) -> Result { let model = request @@ -1108,8 +1290,8 @@ async fn execute_once_assistant_prompt( prompt_request.profile_id = Some(delete_target.clone()); prompt_request.model = None; - let prompt_result = if stream { - execute_assistant_prompt_stream(&prompt_request, token, print_compact_json).await + let prompt_result = if let Some(stream_output) = stream_output { + execute_streaming_assistant_prompt(&prompt_request, token, stream_output).await } else { execute_assistant_prompt(&prompt_request, token).await }; @@ -1130,6 +1312,38 @@ fn temporary_assistant_name() -> String { format!("kagi-cli-once-{millis}-{}", std::process::id()) } +async fn execute_streaming_assistant_prompt( + request: &AssistantPromptRequest, + token: &str, + stream_output: AssistantStreamOutput, +) -> Result { + let response = match stream_output { + AssistantStreamOutput::Text => { + let mut saw_text = false; + let response = execute_assistant_prompt_stream(request, token, |event| { + if !event.md_delta.is_empty() { + saw_text = true; + print!("{}", event.md_delta); + io::stdout().flush().map_err(|error| { + KagiError::Network(format!("failed to flush assistant stream: {error}")) + })?; + } + Ok(()) + }) + .await?; + if saw_text { + println!(); + } + response + } + AssistantStreamOutput::Json => { + execute_assistant_prompt_stream(request, token, print_compact_json).await? + } + }; + + Ok(response) +} + fn print_compact_json(value: &T) -> Result<(), KagiError> { let output = serde_json::to_string(value) .map_err(|error| KagiError::Parse(format!("failed to serialize JSON output: {error}")))?; diff --git a/tests/integration-cli.rs b/tests/integration-cli.rs index 1fd67bb..7b4f54b 100644 --- a/tests/integration-cli.rs +++ b/tests/integration-cli.rs @@ -23,6 +23,9 @@ fn run_kagi(args: &[&str], envs: &[(&str, &str)], cwd: &Path) -> Output { "KAGI_NEWS_BASE_URL", "KAGI_TRANSLATE_BASE_URL", "KAGI_CACHE_DIR", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + "SHELL", ] { command.env_remove(key); } @@ -51,6 +54,9 @@ fn run_kagi_with_stdin(args: &[&str], stdin: &str, envs: &[(&str, &str)], cwd: & "KAGI_NEWS_BASE_URL", "KAGI_TRANSLATE_BASE_URL", "KAGI_CACHE_DIR", + "XDG_CONFIG_HOME", + "XDG_DATA_HOME", + "SHELL", ] { command.env_remove(key); } @@ -168,6 +174,17 @@ fn search_payload(title: &str, url: &str, snippet: &str) -> Value { }) } +fn search_html_fixture() -> &'static str { + r#" + +
+ Session Result +
Served by session fallback.
+
+ + "# +} + fn news_latest_batch() -> Value { json!({ "createdAt": "2026-04-06T00:00:00Z", @@ -419,6 +436,54 @@ fn search_command_limit_truncates_results() { assert_eq!(data[1]["title"], "B"); } +#[test] +fn search_command_falls_back_to_session_when_api_is_rate_limited() { + let server = MockServer::start(); + let api_search = server.mock(|when, then| { + when.method(POST) + .path("/api/v1/search") + .json_body(json!({ "query": "rust programming" })) + .header("authorization", "Bearer test-api-key"); + then.status(429) + .header("content-type", "application/json") + .json_body(json!({ + "error": [{ "msg": "rate limit exceeded" }] + })); + }); + let session_search = server.mock(|when, then| { + when.method(GET) + .path("/html/search") + .query_param("q", "rust programming") + .header("cookie", "kagi_session=test-session"); + then.status(200) + .header("content-type", "text/html") + .body(search_html_fixture()); + }); + + let tempdir = TempDir::new().expect("tempdir"); + fs::write( + tempdir.path().join(".kagi.toml"), + "[auth]\npreferred_auth = \"api\"\n", + ) + .expect("config should write"); + let env = vec![ + ("KAGI_API_KEY", API_KEY.to_string()), + ("KAGI_SESSION_TOKEN", "test-session".to_string()), + ("KAGI_BASE_URL", server.base_url()), + ]; + let output = run_kagi( + &["search", "rust programming", "--format", "json"], + &env_refs(&env), + tempdir.path(), + ); + + assert_success(&output); + api_search.assert_calls(1); + session_search.assert_calls(1); + let body: Value = serde_json::from_slice(&output.stdout).expect("json output should parse"); + assert_eq!(body["data"][0]["title"], "Session Result"); +} + #[test] fn batch_command_returns_queries_and_results() { let server = MockServer::start(); @@ -1044,7 +1109,7 @@ fn assistant_models_prints_json_catalog() { } #[test] -fn assistant_stream_prints_ndjson_updates() { +fn assistant_stream_prints_text_deltas_by_default() { let server = MockServer::start(); let _prompt = server.mock(|when, then| { when.method(POST) @@ -1070,6 +1135,37 @@ fn assistant_stream_prints_ndjson_updates() { tempdir.path(), ); + assert_success(&output); + assert_eq!(String::from_utf8_lossy(&output.stdout), "Hello\n"); +} + +#[test] +fn assistant_stream_can_print_ndjson_updates() { + 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") + .header("content-type", "application/json"); + then.status(200) + .header("content-type", "application/vnd.kagi.stream") + .body(concat!( + "hi:{\"v\":\"test\",\"trace\":\"trace-stream\"}\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\",\"state\":\"streaming\",\"prompt\":\"Hello\",\"md\":\"Hel\",\"documents\":[]}\0\n", + "new_message.json:{\"id\":\"msg-1\",\"thread_id\":\"thread-1\",\"created_at\":\"2026-03-16T06:19:07Z\",\"state\":\"done\",\"prompt\":\"Hello\",\"md\":\"Hello\",\"documents\":[]}\0\n" + )); + }); + + let tempdir = TempDir::new().expect("tempdir"); + let env = session_env(&server); + let output = run_kagi( + &["assistant", "--stream", "--stream-output", "json", "Hello"], + &env_refs(&env), + tempdir.path(), + ); + assert_success(&output); let lines = String::from_utf8_lossy(&output.stdout) .lines() @@ -1081,6 +1177,31 @@ fn assistant_stream_prints_ndjson_updates() { assert_eq!(lines[1]["message"]["state"], "done"); } +#[test] +fn completion_install_detects_fish_and_writes_completion_file() { + let tempdir = TempDir::new().expect("tempdir"); + let config_home = tempdir.path().join("config"); + let config_home_value = config_home.to_string_lossy().to_string(); + let env = vec![ + ("SHELL", "/usr/bin/fish".to_string()), + ("XDG_CONFIG_HOME", config_home_value), + ]; + + let output = run_kagi(&["completion", "install"], &env_refs(&env), tempdir.path()); + + assert_success(&output); + let target = config_home + .join("fish") + .join("completions") + .join("kagi.fish"); + let completion = fs::read_to_string(&target).expect("completion file should exist"); + assert!( + completion.contains("complete"), + "expected fish completion script, got:\n{completion}" + ); + assert!(String::from_utf8_lossy(&output.stdout).contains(target.to_string_lossy().as_ref())); +} + #[test] fn assistant_once_creates_prompts_and_deletes_temporary_profile() { let server = MockServer::start();