Skip to main content

adk_cli/
launcher.rs

1//! Simple launcher for ADK agents with CLI support.
2//!
3//! Provides a one-liner to run agents with console or web server modes,
4//! similar to adk-go's launcher pattern.
5//!
6//! # Example
7//!
8//! ```ignore
9//! use adk_cli::Launcher;
10//! use std::sync::Arc;
11//!
12//! #[tokio::main]
13//! async fn main() -> adk_core::Result<()> {
14//!     let agent = /* create your agent */;
15//!
16//!     // Run with CLI support (console by default, or `serve` for web)
17//!     Launcher::new(Arc::new(agent)).run().await
18//! }
19//! ```
20//!
21//! # CLI Usage
22//!
23//! ```bash
24//! # Interactive console (default)
25//! cargo run
26//!
27//! # Web server with UI
28//! cargo run -- serve
29//! cargo run -- serve --port 3000
30//! ```
31
32use adk_artifact::ArtifactService;
33use adk_core::{
34    Agent, CacheCapable, Content, ContextCacheConfig, EventsCompactionConfig, Memory, Part, Result,
35    RunConfig, StreamingMode,
36};
37use adk_runner::{Runner, RunnerConfig};
38use adk_server::{
39    RequestContextExtractor, SecurityConfig, ServerConfig, create_app, create_app_with_a2a,
40    shutdown_signal,
41};
42use adk_session::{CreateRequest, InMemorySessionService, SessionService};
43use axum::Router;
44use clap::{Parser, Subcommand};
45use futures::StreamExt;
46use rustyline::DefaultEditor;
47use serde_json::Value;
48use std::collections::HashMap;
49use std::io::{self, Write};
50use std::sync::Arc;
51use std::time::Duration;
52use tokio_util::sync::CancellationToken;
53use tracing::{error, warn};
54
55/// CLI arguments for the launcher.
56#[derive(Parser)]
57#[command(name = "agent")]
58#[command(about = "ADK Agent", long_about = None)]
59struct LauncherCli {
60    #[command(subcommand)]
61    command: Option<LauncherCommand>,
62}
63
64#[derive(Subcommand)]
65enum LauncherCommand {
66    /// Run interactive console (default if no command specified)
67    Chat,
68    /// Start web server with UI
69    Serve {
70        /// Server port
71        #[arg(long, default_value_t = 8080)]
72        port: u16,
73    },
74}
75
76/// Controls how the console renders thinking/reasoning content.
77#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
78pub enum ThinkingDisplayMode {
79    /// Show thinking when the model emits it.
80    #[default]
81    Auto,
82    /// Always surface emitted thinking.
83    Show,
84    /// Hide emitted thinking from the console output.
85    Hide,
86}
87
88/// Controls how serve mode initializes telemetry.
89#[derive(Debug, Clone, PartialEq, Eq)]
90pub enum TelemetryConfig {
91    /// Use the in-memory ADK span exporter for debug endpoints.
92    AdkExporter { service_name: String },
93    /// Export telemetry to an OTLP collector.
94    Otlp { service_name: String, endpoint: String },
95    /// Disable telemetry initialization in the launcher.
96    None,
97}
98
99impl Default for TelemetryConfig {
100    fn default() -> Self {
101        Self::AdkExporter { service_name: "adk-server".to_string() }
102    }
103}
104
105/// Launcher for running ADK agents with CLI support.
106///
107/// Provides console and web server modes out of the box.
108///
109/// # Console mode
110///
111/// Uses a `rustyline` REPL with history, Ctrl+C handling, and streaming
112/// output that renders `<think>` blocks, tool calls, and inline data.
113///
114/// # Serve mode
115///
116/// Starts an Axum HTTP server with the ADK web UI.
117///
118/// # Note on `with_streaming_mode`
119///
120/// `with_streaming_mode` currently only affects console mode. `ServerConfig`
121/// does not accept a `RunConfig`, so the setting is not forwarded to the
122/// server. This will be addressed when `ServerConfig` gains that field.
123pub struct Launcher {
124    agent: Arc<dyn Agent>,
125    app_name: Option<String>,
126    session_service: Option<Arc<dyn SessionService>>,
127    artifact_service: Option<Arc<dyn ArtifactService>>,
128    memory_service: Option<Arc<dyn Memory>>,
129    compaction_config: Option<EventsCompactionConfig>,
130    context_cache_config: Option<ContextCacheConfig>,
131    cache_capable: Option<Arc<dyn CacheCapable>>,
132    security_config: Option<SecurityConfig>,
133    request_context_extractor: Option<Arc<dyn RequestContextExtractor>>,
134    a2a_base_url: Option<String>,
135    telemetry_config: TelemetryConfig,
136    shutdown_grace_period: Duration,
137    run_config: Option<RunConfig>,
138    thinking_mode: ThinkingDisplayMode,
139}
140
141impl Launcher {
142    /// Create a new launcher with the given agent.
143    pub fn new(agent: Arc<dyn Agent>) -> Self {
144        Self {
145            agent,
146            app_name: None,
147            session_service: None,
148            artifact_service: None,
149            memory_service: None,
150            compaction_config: None,
151            context_cache_config: None,
152            cache_capable: None,
153            security_config: None,
154            request_context_extractor: None,
155            a2a_base_url: None,
156            telemetry_config: TelemetryConfig::default(),
157            shutdown_grace_period: Duration::from_secs(30),
158            run_config: None,
159            thinking_mode: ThinkingDisplayMode::Auto,
160        }
161    }
162
163    /// Set a custom application name (defaults to agent name).
164    pub fn app_name(mut self, name: impl Into<String>) -> Self {
165        self.app_name = Some(name.into());
166        self
167    }
168
169    /// Set a custom artifact service.
170    pub fn with_artifact_service(mut self, service: Arc<dyn ArtifactService>) -> Self {
171        self.artifact_service = Some(service);
172        self
173    }
174
175    /// Set a custom session service.
176    pub fn with_session_service(mut self, service: Arc<dyn SessionService>) -> Self {
177        self.session_service = Some(service);
178        self
179    }
180
181    /// Set a custom memory service.
182    pub fn with_memory_service(mut self, service: Arc<dyn Memory>) -> Self {
183        self.memory_service = Some(service);
184        self
185    }
186
187    /// Enable runner-level context compaction in serve mode.
188    pub fn with_compaction(mut self, config: EventsCompactionConfig) -> Self {
189        self.compaction_config = Some(config);
190        self
191    }
192
193    /// Enable automatic prompt cache lifecycle management in serve mode.
194    pub fn with_context_cache(
195        mut self,
196        config: ContextCacheConfig,
197        cache_capable: Arc<dyn CacheCapable>,
198    ) -> Self {
199        self.context_cache_config = Some(config);
200        self.cache_capable = Some(cache_capable);
201        self
202    }
203
204    /// Set custom server security settings.
205    pub fn with_security_config(mut self, config: SecurityConfig) -> Self {
206        self.security_config = Some(config);
207        self
208    }
209
210    /// Set a request context extractor for authenticated deployments.
211    pub fn with_request_context_extractor(
212        mut self,
213        extractor: Arc<dyn RequestContextExtractor>,
214    ) -> Self {
215        self.request_context_extractor = Some(extractor);
216        self
217    }
218
219    /// Enable A2A routes when building or serving the HTTP app.
220    pub fn with_a2a_base_url(mut self, base_url: impl Into<String>) -> Self {
221        self.a2a_base_url = Some(base_url.into());
222        self
223    }
224
225    /// Configure how serve mode initializes telemetry.
226    pub fn with_telemetry(mut self, config: TelemetryConfig) -> Self {
227        self.telemetry_config = config;
228        self
229    }
230
231    /// Set the maximum graceful shutdown window for the web server.
232    pub fn with_shutdown_grace_period(mut self, grace_period: Duration) -> Self {
233        self.shutdown_grace_period = grace_period;
234        self
235    }
236
237    /// Set streaming mode for console mode.
238    ///
239    /// Note: this currently only affects console mode. The server does not
240    /// yet accept a `RunConfig`.
241    pub fn with_streaming_mode(mut self, mode: StreamingMode) -> Self {
242        self.run_config = Some(RunConfig { streaming_mode: mode, ..RunConfig::default() });
243        self
244    }
245
246    /// Control how emitted thinking content is rendered in console mode.
247    pub fn with_thinking_mode(mut self, mode: ThinkingDisplayMode) -> Self {
248        self.thinking_mode = mode;
249        self
250    }
251
252    /// Run the launcher, parsing CLI arguments.
253    ///
254    /// - No arguments or `chat`: Interactive console
255    /// - `serve [--port PORT]`: Web server with UI
256    pub async fn run(self) -> Result<()> {
257        let cli = LauncherCli::parse();
258
259        match cli.command.unwrap_or(LauncherCommand::Chat) {
260            LauncherCommand::Chat => self.run_console_directly().await,
261            LauncherCommand::Serve { port } => self.run_serve_directly(port).await,
262        }
263    }
264
265    /// Run in interactive console mode without parsing CLI arguments.
266    ///
267    /// Use this when you already know you want console mode (e.g. from
268    /// your own CLI parser). [`run`](Self::run) calls this internally.
269    pub async fn run_console_directly(self) -> Result<()> {
270        let app_name = self.app_name.unwrap_or_else(|| self.agent.name().to_string());
271        let user_id = "user".to_string();
272        let thinking_mode = self.thinking_mode;
273        let agent = self.agent;
274        let artifact_service = self.artifact_service;
275        let memory_service = self.memory_service;
276        let run_config = self.run_config;
277
278        let session_service =
279            self.session_service.unwrap_or_else(|| Arc::new(InMemorySessionService::new()));
280
281        let session = session_service
282            .create(CreateRequest {
283                app_name: app_name.clone(),
284                user_id: user_id.clone(),
285                session_id: None,
286                state: HashMap::new(),
287            })
288            .await?;
289
290        let session_id = session.id().to_string();
291
292        let mut rl = DefaultEditor::new()
293            .map_err(|e| adk_core::AdkError::Config(format!("failed to init readline: {e}")))?;
294
295        print_banner(agent.name());
296
297        loop {
298            let readline = rl.readline("\x1b[36mYou >\x1b[0m ");
299            match readline {
300                Ok(line) => {
301                    let trimmed = line.trim().to_string();
302                    if trimmed.is_empty() {
303                        continue;
304                    }
305                    if is_exit_command(&trimmed) {
306                        println!("\nGoodbye.\n");
307                        break;
308                    }
309                    if trimmed == "/help" {
310                        print_help();
311                        continue;
312                    }
313                    if trimmed == "/clear" {
314                        println!("(conversation cleared — session state is unchanged)");
315                        continue;
316                    }
317
318                    let _ = rl.add_history_entry(&line);
319
320                    let user_content = Content::new("user").with_text(trimmed);
321                    println!();
322
323                    let cancellation_token = CancellationToken::new();
324                    let runner = Runner::new(RunnerConfig {
325                        app_name: app_name.clone(),
326                        agent: agent.clone(),
327                        session_service: session_service.clone(),
328                        artifact_service: artifact_service.clone(),
329                        memory_service: memory_service.clone(),
330                        plugin_manager: None,
331                        run_config: run_config.clone(),
332                        compaction_config: None,
333                        context_cache_config: None,
334                        cache_capable: None,
335                        request_context: None,
336                        cancellation_token: Some(cancellation_token.clone()),
337                    })?;
338                    let mut events =
339                        runner.run(user_id.clone(), session_id.clone(), user_content).await?;
340                    let mut printer = StreamPrinter::new(thinking_mode);
341                    let mut current_agent = String::new();
342                    let mut printed_header = false;
343                    let mut interrupted = false;
344
345                    loop {
346                        tokio::select! {
347                            event = events.next() => {
348                                let Some(event) = event else {
349                                    break;
350                                };
351
352                                match event {
353                                    Ok(evt) => {
354                                        // Track agent switches in multi-agent workflows
355                                        if !evt.author.is_empty()
356                                            && evt.author != "user"
357                                            && evt.author != current_agent
358                                        {
359                                            if !current_agent.is_empty() {
360                                                println!();
361                                            }
362                                            current_agent = evt.author.clone();
363                                            // Only show agent label in multi-agent scenarios
364                                            if printed_header {
365                                                print!("\x1b[33m[{current_agent}]\x1b[0m ");
366                                                let _ = io::stdout().flush();
367                                            }
368                                            printed_header = true;
369                                        }
370
371                                        // Show agent transfer requests
372                                        if let Some(target) = &evt.actions.transfer_to_agent {
373                                            print!("\x1b[90m[transfer -> {target}]\x1b[0m ");
374                                            let _ = io::stdout().flush();
375                                        }
376
377                                        if let Some(content) = &evt.llm_response.content {
378                                            for part in &content.parts {
379                                                printer.handle_part(part);
380                                            }
381                                        }
382                                    }
383                                    Err(e) => {
384                                        error!("stream error: {e}");
385                                    }
386                                }
387                            }
388                            signal = tokio::signal::ctrl_c() => {
389                                match signal {
390                                    Ok(()) => {
391                                        cancellation_token.cancel();
392                                        interrupted = true;
393                                        break;
394                                    }
395                                    Err(err) => {
396                                        error!("failed to listen for Ctrl+C: {err}");
397                                    }
398                                }
399                            }
400                        }
401                    }
402
403                    printer.finish();
404                    if interrupted {
405                        println!("\nInterrupted.\n");
406                        continue;
407                    }
408
409                    println!("\n");
410                }
411                Err(rustyline::error::ReadlineError::Interrupted) => {
412                    println!("\nInterrupted. Type exit to quit.\n");
413                    continue;
414                }
415                Err(rustyline::error::ReadlineError::Eof) => {
416                    println!("\nGoodbye.\n");
417                    break;
418                }
419                Err(err) => {
420                    error!("readline error: {err}");
421                    break;
422                }
423            }
424        }
425
426        Ok(())
427    }
428
429    fn init_telemetry(&self) -> Option<Arc<adk_telemetry::AdkSpanExporter>> {
430        match &self.telemetry_config {
431            TelemetryConfig::AdkExporter { service_name } => {
432                match adk_telemetry::init_with_adk_exporter(service_name) {
433                    Ok(exporter) => Some(exporter),
434                    Err(e) => {
435                        warn!("failed to initialize telemetry: {e}");
436                        None
437                    }
438                }
439            }
440            TelemetryConfig::Otlp { service_name, endpoint } => {
441                if let Err(e) = adk_telemetry::init_with_otlp(service_name, endpoint) {
442                    warn!("failed to initialize otlp telemetry: {e}");
443                }
444                None
445            }
446            TelemetryConfig::None => None,
447        }
448    }
449
450    fn into_server_config(
451        self,
452        span_exporter: Option<Arc<adk_telemetry::AdkSpanExporter>>,
453    ) -> ServerConfig {
454        let session_service =
455            self.session_service.unwrap_or_else(|| Arc::new(InMemorySessionService::new()));
456        let agent_loader = Arc::new(adk_core::SingleAgentLoader::new(self.agent));
457
458        let mut config = ServerConfig::new(agent_loader, session_service)
459            .with_artifact_service_opt(self.artifact_service);
460
461        if let Some(memory_service) = self.memory_service {
462            config = config.with_memory_service(memory_service);
463        }
464
465        if let Some(compaction_config) = self.compaction_config {
466            config = config.with_compaction(compaction_config);
467        }
468
469        if let (Some(context_cache_config), Some(cache_capable)) =
470            (self.context_cache_config, self.cache_capable)
471        {
472            config = config.with_context_cache(context_cache_config, cache_capable);
473        }
474
475        if let Some(security) = self.security_config {
476            config = config.with_security(security);
477        }
478
479        if let Some(extractor) = self.request_context_extractor {
480            config = config.with_request_context(extractor);
481        }
482
483        if let Some(exporter) = span_exporter {
484            config = config.with_span_exporter(exporter);
485        }
486
487        config
488    }
489
490    /// Build the Axum application without serving it.
491    ///
492    /// This is the production escape hatch for adding custom routes,
493    /// middleware, metrics, or owning the serve loop yourself.
494    pub fn build_app(self) -> Result<Router> {
495        let span_exporter = self.init_telemetry();
496        let a2a_base_url = self.a2a_base_url.clone();
497        let config = self.into_server_config(span_exporter);
498
499        Ok(match a2a_base_url {
500            Some(base_url) => create_app_with_a2a(config, Some(&base_url)),
501            None => create_app(config),
502        })
503    }
504
505    /// Build the Axum application with A2A routes enabled.
506    pub fn build_app_with_a2a(mut self, base_url: impl Into<String>) -> Result<Router> {
507        self.a2a_base_url = Some(base_url.into());
508        self.build_app()
509    }
510
511    /// Run web server without parsing CLI arguments.
512    ///
513    /// Use this when you already know you want serve mode (e.g. from
514    /// your own CLI parser). [`run`](Self::run) calls this internally.
515    pub async fn run_serve_directly(self, port: u16) -> Result<()> {
516        let app = self.build_app()?;
517
518        let addr = format!("0.0.0.0:{port}");
519        let listener = tokio::net::TcpListener::bind(&addr).await?;
520
521        println!("ADK Server starting on http://localhost:{port}");
522        println!("Open http://localhost:{port} in your browser");
523        println!("Press Ctrl+C to stop\n");
524
525        axum::serve(listener, app).with_graceful_shutdown(shutdown_signal()).await?;
526
527        Ok(())
528    }
529}
530
531/// Print the ADK-Rust welcome banner.
532fn print_banner(agent_name: &str) {
533    let version = env!("CARGO_PKG_VERSION");
534    let title = format!("ADK-Rust  v{version}");
535    let subtitle = "Rust Agent Development Kit";
536    // Box inner width = 49 (between the ║ chars)
537    let inner: usize = 49;
538    let pad_title = (inner.saturating_sub(title.len())) / 2;
539    let pad_subtitle = (inner.saturating_sub(subtitle.len())) / 2;
540
541    println!();
542    println!("    ╔{:═<inner$}╗", "");
543    println!(
544        "    ║{:>w$}{title}{:<r$}║",
545        "",
546        "",
547        w = pad_title,
548        r = inner - pad_title - title.len()
549    );
550    println!(
551        "    ║{:>w$}{subtitle}{:<r$}║",
552        "",
553        "",
554        w = pad_subtitle,
555        r = inner - pad_subtitle - subtitle.len()
556    );
557    println!("    ╚{:═<inner$}╝", "");
558    println!();
559    println!("  Agent    : {agent_name}");
560    println!("  Runtime  : Tokio async, streaming responses");
561    println!("  Features : tool calling, multi-provider, multi-agent, think blocks");
562    println!();
563    println!("  Type a message to chat. /help for commands.");
564    println!();
565}
566
567/// Print available REPL commands.
568fn print_help() {
569    println!();
570    println!("  Commands:");
571    println!("    /help   Show this help");
572    println!("    /clear  Clear display (session state is kept)");
573    println!("    quit    Exit the REPL");
574    println!("    exit    Exit the REPL");
575    println!("    /quit   Exit the REPL");
576    println!("    /exit   Exit the REPL");
577    println!();
578    println!("  Tips:");
579    println!("    - Up/Down arrows browse history");
580    println!("    - Ctrl+C interrupts the current operation");
581    println!("    - Multi-agent workflows show [agent_name] on handoff");
582    println!();
583}
584
585fn is_exit_command(input: &str) -> bool {
586    matches!(input, "quit" | "exit" | "/quit" | "/exit")
587}
588
589/// Handles streaming output with special rendering for think blocks,
590/// tool calls, function responses, and inline/file data.
591pub struct StreamPrinter {
592    thinking_mode: ThinkingDisplayMode,
593    in_think_block: bool,
594    in_thinking_part_stream: bool,
595    think_buffer: String,
596}
597
598impl StreamPrinter {
599    /// Create a printer with the selected thinking display mode.
600    pub fn new(thinking_mode: ThinkingDisplayMode) -> Self {
601        Self {
602            thinking_mode,
603            in_think_block: false,
604            in_thinking_part_stream: false,
605            think_buffer: String::new(),
606        }
607    }
608
609    /// Process a single response part, rendering it to stdout.
610    pub fn handle_part(&mut self, part: &Part) {
611        match part {
612            Part::Text { text } => {
613                self.flush_part_thinking_if_needed();
614                self.handle_text_chunk(text);
615            }
616            Part::Thinking { thinking, .. } => {
617                if matches!(self.thinking_mode, ThinkingDisplayMode::Hide) {
618                    return;
619                }
620                if !self.in_thinking_part_stream {
621                    print!("\n[thinking] ");
622                    let _ = io::stdout().flush();
623                    self.in_thinking_part_stream = true;
624                }
625                self.think_buffer.push_str(thinking);
626                print!("{thinking}");
627                let _ = io::stdout().flush();
628            }
629            Part::FunctionCall { name, args, .. } => {
630                self.flush_pending_thinking();
631                print!("\n[tool-call] {name} {args}\n");
632                let _ = io::stdout().flush();
633            }
634            Part::FunctionResponse { function_response, .. } => {
635                self.flush_pending_thinking();
636                self.print_tool_response(&function_response.name, &function_response.response);
637            }
638            Part::InlineData { mime_type, data } => {
639                self.flush_pending_thinking();
640                print!("\n[inline-data] mime={mime_type} bytes={}\n", data.len());
641                let _ = io::stdout().flush();
642            }
643            Part::FileData { mime_type, file_uri } => {
644                self.flush_pending_thinking();
645                print!("\n[file-data] mime={mime_type} uri={file_uri}\n");
646                let _ = io::stdout().flush();
647            }
648        }
649    }
650
651    fn handle_text_chunk(&mut self, chunk: &str) {
652        if matches!(self.thinking_mode, ThinkingDisplayMode::Hide) {
653            let mut visible = String::with_capacity(chunk.len());
654            let mut remaining = chunk;
655
656            while let Some(start_idx) = remaining.find("<think>") {
657                visible.push_str(&remaining[..start_idx]);
658                let after_start = &remaining[start_idx + "<think>".len()..];
659                if let Some(end_idx) = after_start.find("</think>") {
660                    remaining = &after_start[end_idx + "</think>".len()..];
661                } else {
662                    remaining = "";
663                    break;
664                }
665            }
666
667            visible.push_str(remaining);
668            if !visible.is_empty() {
669                print!("{visible}");
670                let _ = io::stdout().flush();
671            }
672            return;
673        }
674
675        const THINK_START: &str = "<think>";
676        const THINK_END: &str = "</think>";
677
678        let mut remaining = chunk;
679
680        while !remaining.is_empty() {
681            if self.in_think_block {
682                if let Some(end_idx) = remaining.find(THINK_END) {
683                    self.think_buffer.push_str(&remaining[..end_idx]);
684                    self.flush_think();
685                    self.in_think_block = false;
686                    remaining = &remaining[end_idx + THINK_END.len()..];
687                } else {
688                    self.think_buffer.push_str(remaining);
689                    break;
690                }
691            } else if let Some(start_idx) = remaining.find(THINK_START) {
692                let visible = &remaining[..start_idx];
693                if !visible.is_empty() {
694                    print!("{visible}");
695                    let _ = io::stdout().flush();
696                }
697                self.in_think_block = true;
698                self.think_buffer.clear();
699                remaining = &remaining[start_idx + THINK_START.len()..];
700            } else {
701                print!("{remaining}");
702                let _ = io::stdout().flush();
703                break;
704            }
705        }
706    }
707
708    fn flush_think(&mut self) {
709        let content = self.think_buffer.trim();
710        if !content.is_empty() {
711            print!("\n[think] {content}\n");
712            let _ = io::stdout().flush();
713        }
714        self.think_buffer.clear();
715    }
716
717    /// Flush any pending think block content.
718    pub fn finish(&mut self) {
719        self.flush_pending_thinking();
720    }
721
722    fn print_tool_response(&self, name: &str, response: &Value) {
723        print!("\n[tool-response] {name} {response}\n");
724        let _ = io::stdout().flush();
725    }
726
727    fn flush_part_thinking_if_needed(&mut self) {
728        if self.in_thinking_part_stream {
729            println!();
730            let _ = io::stdout().flush();
731            self.think_buffer.clear();
732            self.in_thinking_part_stream = false;
733        }
734    }
735
736    fn flush_pending_thinking(&mut self) {
737        self.flush_part_thinking_if_needed();
738        if self.in_think_block {
739            self.flush_think_with_label("think");
740            self.in_think_block = false;
741        }
742    }
743
744    fn flush_think_with_label(&mut self, label: &str) {
745        let content = self.think_buffer.trim();
746        if !content.is_empty() {
747            print!("\n[{label}] {content}\n");
748            let _ = io::stdout().flush();
749        }
750        self.think_buffer.clear();
751    }
752}
753
754impl Default for StreamPrinter {
755    fn default() -> Self {
756        Self::new(ThinkingDisplayMode::Auto)
757    }
758}
759
760#[cfg(test)]
761mod tests {
762    use super::*;
763    use adk_core::{Agent, EventStream, InvocationContext, Result as AdkResult};
764    use async_trait::async_trait;
765    use axum::{
766        body::{Body, to_bytes},
767        http::{Request, StatusCode},
768    };
769    use futures::stream;
770    use std::sync::Arc;
771    use tower::ServiceExt;
772
773    struct TestAgent;
774
775    #[async_trait]
776    impl Agent for TestAgent {
777        fn name(&self) -> &str {
778            "launcher_test_agent"
779        }
780
781        fn description(&self) -> &str {
782            "launcher test agent"
783        }
784
785        fn sub_agents(&self) -> &[Arc<dyn Agent>] {
786            &[]
787        }
788
789        async fn run(&self, _ctx: Arc<dyn InvocationContext>) -> AdkResult<EventStream> {
790            Ok(Box::pin(stream::empty()))
791        }
792    }
793
794    fn test_launcher() -> Launcher {
795        Launcher::new(Arc::new(TestAgent)).with_telemetry(TelemetryConfig::None)
796    }
797
798    #[test]
799    fn stream_printer_tracks_think_block_state() {
800        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
801        assert!(!printer.in_think_block);
802
803        // Opening a think block sets the flag
804        printer.handle_text_chunk("<think>reasoning");
805        assert!(printer.in_think_block);
806        assert_eq!(printer.think_buffer, "reasoning");
807
808        // Closing the think block clears the flag and buffer
809        printer.handle_text_chunk(" more</think>visible");
810        assert!(!printer.in_think_block);
811        assert!(printer.think_buffer.is_empty());
812    }
813
814    #[test]
815    fn stream_printer_handles_think_block_across_chunks() {
816        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
817
818        printer.handle_text_chunk("before<think>start");
819        assert!(printer.in_think_block);
820        assert_eq!(printer.think_buffer, "start");
821
822        printer.handle_text_chunk(" middle");
823        assert!(printer.in_think_block);
824        assert_eq!(printer.think_buffer, "start middle");
825
826        printer.handle_text_chunk(" end</think>after");
827        assert!(!printer.in_think_block);
828        assert!(printer.think_buffer.is_empty());
829    }
830
831    #[test]
832    fn stream_printer_finish_flushes_open_think_block() {
833        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
834
835        printer.handle_text_chunk("<think>unclosed reasoning");
836        assert!(printer.in_think_block);
837
838        printer.finish();
839        assert!(!printer.in_think_block);
840        assert!(printer.think_buffer.is_empty());
841    }
842
843    #[test]
844    fn stream_printer_finish_is_noop_when_no_think_block() {
845        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
846        printer.finish();
847        assert!(!printer.in_think_block);
848        assert!(printer.think_buffer.is_empty());
849    }
850
851    #[test]
852    fn stream_printer_handles_multiple_think_blocks_in_one_chunk() {
853        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
854
855        printer.handle_text_chunk("a<think>first</think>b<think>second</think>c");
856        assert!(!printer.in_think_block);
857        assert!(printer.think_buffer.is_empty());
858    }
859
860    #[test]
861    fn stream_printer_handles_empty_think_block() {
862        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
863
864        printer.handle_text_chunk("<think></think>after");
865        assert!(!printer.in_think_block);
866        assert!(printer.think_buffer.is_empty());
867    }
868
869    #[test]
870    fn stream_printer_handles_all_part_types() {
871        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
872
873        // Text
874        printer.handle_part(&Part::Text { text: "hello".into() });
875        assert!(!printer.in_think_block);
876
877        // Thinking
878        printer.handle_part(&Part::Thinking { thinking: "reasoning".into(), signature: None });
879        assert!(printer.in_thinking_part_stream);
880
881        // FunctionCall
882        printer.handle_part(&Part::FunctionCall {
883            name: "get_weather".into(),
884            args: serde_json::json!({"city": "Seattle"}),
885            id: None,
886            thought_signature: None,
887        });
888
889        // FunctionResponse
890        printer.handle_part(&Part::FunctionResponse {
891            function_response: adk_core::FunctionResponseData {
892                name: "get_weather".into(),
893                response: serde_json::json!({"temp": 72}),
894            },
895            id: None,
896        });
897
898        // InlineData
899        printer
900            .handle_part(&Part::InlineData { mime_type: "image/png".into(), data: vec![0u8; 100] });
901
902        // FileData
903        printer.handle_part(&Part::FileData {
904            mime_type: "audio/wav".into(),
905            file_uri: "gs://bucket/file.wav".into(),
906        });
907
908        // No panics, no state corruption
909        assert!(!printer.in_think_block);
910        assert!(!printer.in_thinking_part_stream);
911    }
912
913    #[test]
914    fn stream_printer_text_without_think_tags_leaves_state_clean() {
915        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
916        printer.handle_text_chunk("just plain text with no tags");
917        assert!(!printer.in_think_block);
918        assert!(printer.think_buffer.is_empty());
919    }
920
921    #[test]
922    fn stream_printer_coalesces_streamed_thinking_parts() {
923        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
924
925        printer.handle_part(&Part::Thinking { thinking: "Okay".into(), signature: None });
926        printer.handle_part(&Part::Thinking { thinking: ", the".into(), signature: None });
927        printer.handle_part(&Part::Thinking { thinking: " user".into(), signature: None });
928
929        assert!(printer.in_thinking_part_stream);
930        assert_eq!(printer.think_buffer, "Okay, the user");
931
932        printer.handle_part(&Part::Text { text: "hello".into() });
933
934        assert!(!printer.in_thinking_part_stream);
935        assert!(printer.think_buffer.is_empty());
936    }
937
938    #[test]
939    fn stream_printer_finish_closes_streamed_thinking_state() {
940        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Auto);
941
942        printer.handle_part(&Part::Thinking { thinking: "reasoning".into(), signature: None });
943        assert!(printer.in_thinking_part_stream);
944
945        printer.finish();
946
947        assert!(!printer.in_thinking_part_stream);
948        assert!(printer.think_buffer.is_empty());
949    }
950
951    #[test]
952    fn stream_printer_hide_mode_ignores_emitted_thinking_state() {
953        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Hide);
954
955        printer.handle_part(&Part::Thinking { thinking: "secret".into(), signature: None });
956
957        assert!(!printer.in_thinking_part_stream);
958        assert!(printer.think_buffer.is_empty());
959    }
960
961    #[test]
962    fn stream_printer_hide_mode_drops_think_tags_from_text() {
963        let mut printer = StreamPrinter::new(ThinkingDisplayMode::Hide);
964
965        printer.handle_text_chunk("visible<think>hidden</think>after");
966
967        assert!(!printer.in_think_block);
968        assert!(printer.think_buffer.is_empty());
969    }
970
971    #[test]
972    fn exit_command_helper_accepts_plain_and_slash_variants() {
973        for command in ["quit", "exit", "/quit", "/exit"] {
974            assert!(is_exit_command(command));
975        }
976
977        assert!(!is_exit_command("hello"));
978    }
979
980    #[tokio::test]
981    async fn build_app_includes_health_route() {
982        let app = test_launcher().build_app().expect("launcher app should build");
983
984        let response = app
985            .oneshot(Request::builder().uri("/api/health").body(Body::empty()).unwrap())
986            .await
987            .expect("health request should succeed");
988
989        assert_eq!(response.status(), StatusCode::OK);
990
991        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
992        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
993        assert_eq!(json["status"], "healthy");
994    }
995
996    #[tokio::test]
997    async fn build_app_does_not_enable_a2a_routes_by_default() {
998        let app = test_launcher().build_app().expect("launcher app should build");
999
1000        let response = app
1001            .oneshot(Request::builder().uri("/.well-known/agent.json").body(Body::empty()).unwrap())
1002            .await
1003            .expect("agent card request should complete");
1004
1005        assert_eq!(response.status(), StatusCode::NOT_FOUND);
1006    }
1007
1008    #[tokio::test]
1009    async fn build_app_with_a2a_enables_agent_card_route() {
1010        let app = test_launcher()
1011            .build_app_with_a2a("http://localhost:8080")
1012            .expect("launcher app with a2a should build");
1013
1014        let response = app
1015            .oneshot(Request::builder().uri("/.well-known/agent.json").body(Body::empty()).unwrap())
1016            .await
1017            .expect("agent card request should complete");
1018
1019        assert_eq!(response.status(), StatusCode::OK);
1020
1021        let body = to_bytes(response.into_body(), usize::MAX).await.unwrap();
1022        let json: serde_json::Value = serde_json::from_slice(&body).unwrap();
1023        assert_eq!(json["name"], "launcher_test_agent");
1024        assert_eq!(json["description"], "launcher test agent");
1025    }
1026}