cookie/
main.rs

1use std::{
2    collections::BTreeMap,
3    env, fmt,
4    io::{self, IsTerminal, Write},
5    path::{Path, PathBuf},
6    sync::Arc,
7    time::Duration,
8};
9
10use anyhow::Context as _;
11use clap::{Parser, Subcommand};
12use cookie_agent::run;
13use cookie_agent_engine::{Engine, EngineOptions};
14use cookie_agent_models::{
15    ModelManager,
16    catalog::{CatalogManager, CatalogTransport, HttpCatalogTransport},
17    provider_store::ProviderStore,
18};
19use cookie_agent_protocol::{
20    AuthMethodDescriptor, ClientConnectId, ClientRequestId, EffectiveAuthState, McpAuthBeginParams,
21    McpServerState, ProviderConfigurationState, ProviderConnectResult, ProviderDescriptor,
22    ProviderDisconnectParams, ProviderDisconnectResult, ProviderId, ProviderSupportState,
23    RuntimeSnapshotResult, SafeCode, SafeSetupValue, parse_setup_value, paths, setup_value_text,
24};
25use cookie_agent_server::{
26    Client, ClientProtocol, Server, generate_token, in_process_pair, ready_line,
27    validate_websocket_url,
28};
29use cookie_agent_tools::{
30    BuiltinTools, delegate::DelegateToolProvider, message::MessageToolProvider,
31};
32use serde::Serialize;
33use tokio_util::sync::CancellationToken;
34use uuid::Uuid;
35use zeroize::{Zeroize, Zeroizing};
36
37#[cfg(all(test, unix))]
38use cookie_agent_protocol::{SetupFieldDescriptor, SetupFieldType};
39#[cfg(all(test, unix))]
40use std::sync::atomic::{AtomicUsize as TestAtomicUsize, Ordering as TestOrdering};
41
42const DEFAULT_WEBSOCKET_URL: &str = "ws://127.0.0.1:7419/ws";
43const CATALOG_REFRESH_INTERVAL: Duration = Duration::from_secs(60 * 60);
44
45enum RunCatalogTransport {
46    Http(HttpCatalogTransport),
47    #[cfg(debug_assertions)]
48    Bundled,
49}
50
51impl CatalogTransport for RunCatalogTransport {
52    fn fetch(
53        &self,
54        request: cookie_agent_models::catalog::CatalogRequest,
55    ) -> cookie_agent_models::catalog::CatalogTransportFuture<'_> {
56        match self {
57            Self::Http(transport) => transport.fetch(request),
58            #[cfg(debug_assertions)]
59            Self::Bundled => Box::pin(async {
60                Ok(
61                    cookie_agent_models::catalog::CatalogTransportResponse::from_bytes(
62                        200,
63                        cookie_agent_models::catalog::MODELS_DEV_BOOTSTRAP.to_vec(),
64                    ),
65                )
66            }),
67        }
68    }
69}
70
71fn run_catalog_transport() -> anyhow::Result<RunCatalogTransport> {
72    #[cfg(debug_assertions)]
73    if matches!(
74        env::var("COOKIE_AGENT_TEST_BUNDLED_CATALOG").as_deref(),
75        Ok("1")
76    ) {
77        return Ok(RunCatalogTransport::Bundled);
78    }
79    HttpCatalogTransport::new()
80        .context("construct fixed catalog transport")
81        .map(RunCatalogTransport::Http)
82}
83
84#[cfg(all(test, unix))]
85static SECRET_VALUES_WIPED: TestAtomicUsize = TestAtomicUsize::new(0);
86
87#[derive(Debug, Parser)]
88#[command(name = "cookie", version = env!("COOKIE_VERSION"))]
89struct Cli {
90    #[command(subcommand)]
91    command: Option<Command>,
92}
93
94#[derive(Debug, PartialEq, Eq, Subcommand)]
95enum Command {
96    /// Run one prompt without the interactive TUI.
97    Run {
98        #[command(flatten)]
99        args: Box<run::RunArgs>,
100    },
101    /// Serve the exact cookie-agent protocol 20 JSON-RPC WebSocket daemon on localhost.
102    Daemon {
103        /// Localhost port to bind. `0` selects an ephemeral port; the ready line
104        /// reports the real one.
105        #[arg(long)]
106        port: Option<u16>,
107    },
108    /// Attach the TUI to an existing daemon.
109    Attach {
110        #[arg(long, default_value = DEFAULT_WEBSOCKET_URL)]
111        url: String,
112        #[arg(long, env = "COOKIE_DAEMON_TOKEN", hide_env_values = true)]
113        token: String,
114    },
115    /// Securely create or update a durable global managed-provider connection.
116    Connect {
117        /// Exact runtime provider ID. Omit to choose interactively.
118        provider_id: Option<String>,
119        #[arg(long, default_value = DEFAULT_WEBSOCKET_URL)]
120        url: String,
121        #[arg(long, env = "COOKIE_DAEMON_TOKEN", hide_env_values = true)]
122        token: String,
123    },
124    /// Remove a durable global managed-provider connection.
125    Disconnect {
126        /// Exact runtime provider ID. Omit to choose interactively.
127        provider_id: Option<String>,
128        #[arg(long, default_value = DEFAULT_WEBSOCKET_URL)]
129        url: String,
130        #[arg(long, env = "COOKIE_DAEMON_TOKEN", hide_env_values = true)]
131        token: String,
132    },
133    /// List MCP servers or start remote-server OAuth.
134    Mcp {
135        #[command(subcommand)]
136        command: McpCommand,
137        #[arg(long, default_value = DEFAULT_WEBSOCKET_URL)]
138        url: String,
139        #[arg(long, env = "COOKIE_DAEMON_TOKEN", hide_env_values = true)]
140        token: String,
141    },
142}
143
144#[derive(Debug, PartialEq, Eq, Subcommand)]
145enum McpCommand {
146    /// List configured MCP servers.
147    List,
148    /// Start OAuth authorization for a remote MCP server.
149    Auth { server: String },
150}
151
152struct Runtime {
153    engine: Engine,
154    server: Arc<Server>,
155    port: u16,
156    catalog_refresh_shutdown: CancellationToken,
157    catalog_refresh_task: Option<tokio::task::JoinHandle<()>>,
158}
159
160impl Runtime {
161    async fn stop_catalog_refresh(&mut self) {
162        self.catalog_refresh_shutdown.cancel();
163        if let Some(task) = self.catalog_refresh_task.take() {
164            let _ = task.await;
165        }
166    }
167}
168
169impl Drop for Runtime {
170    fn drop(&mut self) {
171        self.catalog_refresh_shutdown.cancel();
172    }
173}
174
175#[derive(Default, Serialize)]
176#[serde(transparent)]
177struct SecretValues(BTreeMap<String, String>);
178
179impl fmt::Debug for SecretValues {
180    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
181        formatter.write_str("SecretValues(<redacted>)")
182    }
183}
184
185impl Drop for SecretValues {
186    fn drop(&mut self) {
187        for value in self.0.values_mut() {
188            value.zeroize();
189        }
190        #[cfg(all(test, unix))]
191        SECRET_VALUES_WIPED.fetch_add(1, TestOrdering::SeqCst);
192    }
193}
194
195#[derive(Serialize)]
196struct SensitiveProviderConnectParams {
197    provider_id: ProviderId,
198    expected_catalog_revision: cookie_agent_protocol::CatalogRevision,
199    setup_values: BTreeMap<cookie_agent_protocol::SetupFieldId, SafeSetupValue>,
200    auth_method: cookie_agent_protocol::AuthMethodId,
201    auth_values: SecretValues,
202    client_connect_id: ClientConnectId,
203}
204
205trait ConnectIo {
206    fn write_line(&mut self, line: &str) -> anyhow::Result<()>;
207    fn read_public(&mut self, prompt: &str) -> anyhow::Result<String>;
208    fn read_secret(&mut self, prompt: &str) -> anyhow::Result<Zeroizing<String>>;
209}
210
211struct StdioConnectIo;
212
213impl ConnectIo for StdioConnectIo {
214    fn write_line(&mut self, line: &str) -> anyhow::Result<()> {
215        println!("{line}");
216        Ok(())
217    }
218
219    fn read_public(&mut self, prompt: &str) -> anyhow::Result<String> {
220        read_public_line(prompt)
221    }
222
223    fn read_secret(&mut self, prompt: &str) -> anyhow::Result<Zeroizing<String>> {
224        read_secret_line(prompt)
225    }
226}
227
228#[tokio::main]
229async fn main() {
230    if let Err(error) = main_result().await {
231        eprintln!(
232            "cookie: {}",
233            cookie_agent_protocol::diagnostics::detail(&format!("{error:#}"))
234        );
235        std::process::exit(1);
236    }
237}
238
239async fn main_result() -> anyhow::Result<()> {
240    tracing_subscriber::fmt::init();
241    let Cli { command } = Cli::parse();
242    if let Some(Command::Run { args }) = &command
243        && let Err(message) = args.validate_cli()
244    {
245        clap::Error::raw(clap::error::ErrorKind::ArgumentConflict, message).exit();
246    }
247    let workspace = if matches!(&command, Some(Command::Run { .. })) {
248        None
249    } else {
250        local_workspace(&command, env::current_dir)?
251    };
252    match command {
253        Some(Command::Run { args }) => {
254            let workspace = match env::current_dir().context("determine current run workspace") {
255                Ok(path) => path,
256                Err(error) => {
257                    eprintln!(
258                        "cookie run: {}",
259                        cookie_agent_protocol::diagnostics::detail(&format!("{error:#}"))
260                    );
261                    std::process::exit(run::EXIT_ENVIRONMENT);
262                }
263            };
264            let data_dir = match args.data_dir.clone().map_or_else(data_dir, Ok) {
265                Ok(path) => path,
266                Err(error) => {
267                    eprintln!(
268                        "cookie run: {}",
269                        cookie_agent_protocol::diagnostics::detail(&format!("{error:#}"))
270                    );
271                    std::process::exit(run::EXIT_ENVIRONMENT);
272                }
273            };
274            let mut runtime = match compose_with(
275                &workspace,
276                run_catalog_transport,
277                CatalogManager::standard,
278                ProviderStore::standard,
279                || Ok(data_dir),
280            )
281            .await
282            {
283                Ok(runtime) => runtime,
284                Err(error) => {
285                    eprintln!(
286                        "cookie run: {}",
287                        cookie_agent_protocol::diagnostics::detail(&format!("{error:#}"))
288                    );
289                    std::process::exit(run::EXIT_ENVIRONMENT);
290                }
291            };
292            let exit_code = run::execute(&runtime.engine, *args).await;
293            runtime.stop_catalog_refresh().await;
294            runtime.engine.shutdown().await;
295            if exit_code != 0 {
296                std::process::exit(exit_code);
297            }
298            Ok(())
299        }
300        Some(Command::Daemon { port }) => {
301            run_daemon(
302                compose(workspace.as_deref().expect("daemon workspace")).await?,
303                port,
304            )
305            .await
306        }
307        Some(Command::Connect {
308            provider_id,
309            url,
310            token,
311        }) => run_connect(&url, &token, provider_id).await,
312        Some(Command::Disconnect {
313            provider_id,
314            url,
315            token,
316        }) => run_disconnect(&url, &token, provider_id).await,
317        Some(Command::Mcp {
318            command,
319            url,
320            token,
321        }) => run_mcp(&url, &token, command).await,
322        Some(Command::Attach { url, token }) => run_attached_tui(&url, &token).await,
323        None => {
324            run_local_frontend(compose(workspace.as_deref().expect("local workspace")).await?).await
325        }
326    }
327}
328
329fn local_workspace(
330    command: &Option<Command>,
331    current_dir: impl FnOnce() -> io::Result<PathBuf>,
332) -> anyhow::Result<Option<PathBuf>> {
333    if matches!(
334        command,
335        Some(
336            Command::Attach { .. }
337                | Command::Connect { .. }
338                | Command::Disconnect { .. }
339                | Command::Mcp { .. },
340        )
341    ) {
342        Ok(None)
343    } else {
344        current_dir()
345            .context("determine current workspace")
346            .map(Some)
347    }
348}
349
350async fn compose(workspace: &Path) -> anyhow::Result<Runtime> {
351    compose_with(
352        workspace,
353        || HttpCatalogTransport::new().context("construct fixed catalog transport"),
354        CatalogManager::standard,
355        ProviderStore::standard,
356        data_dir,
357    )
358    .await
359}
360
361async fn compose_with<T: CatalogTransport + 'static>(
362    workspace: &Path,
363    open_transport: impl FnOnce() -> anyhow::Result<T>,
364    open_catalog: impl FnOnce(T) -> CatalogManager<T>,
365    open_provider_store: impl FnOnce() -> Result<
366        ProviderStore,
367        cookie_agent_models::provider_store::ProviderStoreError,
368    >,
369    open_data_dir: impl FnOnce() -> anyhow::Result<PathBuf>,
370) -> anyhow::Result<Runtime> {
371    let configuration =
372        cookie_agent_config::load(workspace).context("load workspace configuration and agents")?;
373    compose_with_configuration(
374        workspace,
375        configuration,
376        open_transport,
377        open_catalog,
378        open_provider_store,
379        open_data_dir,
380    )
381    .await
382}
383
384async fn compose_with_configuration<T: CatalogTransport + 'static>(
385    workspace: &Path,
386    configuration: cookie_agent_config::LoadedConfiguration,
387    open_transport: impl FnOnce() -> anyhow::Result<T>,
388    open_catalog: impl FnOnce(T) -> CatalogManager<T>,
389    open_provider_store: impl FnOnce() -> Result<
390        ProviderStore,
391        cookie_agent_models::provider_store::ProviderStoreError,
392    >,
393    open_data_dir: impl FnOnce() -> anyhow::Result<PathBuf>,
394) -> anyhow::Result<Runtime> {
395    if configuration.runtime.server.host != "127.0.0.1" {
396        anyhow::bail!("server.host must be exactly 127.0.0.1");
397    }
398    let port = configuration.runtime.server.port;
399
400    // Startup never waits on the network: it serves the validated cache (or
401    // the bundled catalog) and the refresh loop checks models.dev right away.
402    let catalog_manager = open_catalog(open_transport()?);
403    let catalog = Arc::new(
404        catalog_manager
405            .load_cached()
406            .context("load cached models.dev catalog")?,
407    );
408    let provider_store = open_provider_store().context("open provider store 3")?;
409    let model_manager = Arc::new(
410        ModelManager::new_with_headers(
411            configuration.runtime.providers.clone(),
412            configuration.runtime.headers.clone(),
413            catalog,
414            provider_store,
415        )
416        .context("compose effective providers and current model runtime")?,
417    );
418    let engine = Engine::open(EngineOptions {
419        data_dir: open_data_dir()?,
420        cwd: workspace.to_owned(),
421        config: configuration,
422        model_manager,
423        tools: vec![Arc::new(BuiltinTools::new(workspace))],
424    })
425    .context("open manifests, rehydrate project state, and reconcile engine")?;
426    if let Err(error) = engine
427        .try_register_tool_provider(Arc::new(DelegateToolProvider::new(engine.clone())))
428        .context("register delegate tools")
429    {
430        engine.shutdown().await;
431        return Err(error);
432    }
433    if let Err(error) = engine
434        .try_register_tool_provider(Arc::new(MessageToolProvider::new(engine.clone())))
435        .context("register message tools")
436    {
437        engine.shutdown().await;
438        return Err(error);
439    }
440    if let Err(error) = engine
441        .try_register_tool_provider(Arc::new(cookie_agent_tools::skill::SkillTool::new(
442            engine.clone(),
443        )))
444        .context("register skill tool")
445    {
446        engine.shutdown().await;
447        return Err(error);
448    }
449    if let Err(error) = engine
450        .try_register_tool_provider(Arc::new(cookie_agent_tools::goal::GoalTools::new(
451            engine.clone(),
452        )))
453        .context("register goal tools")
454    {
455        engine.shutdown().await;
456        return Err(error);
457    }
458    let server = Arc::new(Server::new(engine.clone()));
459    let catalog_refresh_shutdown = CancellationToken::new();
460    let catalog_refresh_task = tokio::spawn(run_catalog_refresh_loop(
461        catalog_manager,
462        engine.clone(),
463        catalog_refresh_shutdown.clone(),
464        CATALOG_REFRESH_INTERVAL,
465    ));
466    Ok(Runtime {
467        engine,
468        server,
469        port,
470        catalog_refresh_shutdown,
471        catalog_refresh_task: Some(catalog_refresh_task),
472    })
473}
474
475async fn run_catalog_refresh_loop<T: CatalogTransport + 'static>(
476    catalog_manager: CatalogManager<T>,
477    engine: Engine,
478    shutdown: CancellationToken,
479    cadence: Duration,
480) {
481    // The first check runs immediately, since startup served the cache.
482    let mut first = true;
483    loop {
484        if !std::mem::take(&mut first) {
485            tokio::select! {
486                () = shutdown.cancelled() => return,
487                () = tokio::time::sleep(cadence) => {}
488            }
489        }
490        let refreshed = tokio::select! {
491            () = shutdown.cancelled() => return,
492            refreshed = catalog_manager.refresh() => refreshed,
493        };
494        let Ok(catalog) = refreshed else {
495            continue;
496        };
497        if !catalog_publication_changed(&engine, &catalog) {
498            continue;
499        }
500        if engine.refresh_catalog(Arc::new(catalog)).is_err() {
501            continue;
502        }
503    }
504}
505
506/// Republishing recompiles every model, so it happens only when the models
507/// could differ: a new catalog revision or a change in availability. A server
508/// confirming the cached body (source `cache` becoming `network`) is not one.
509fn catalog_publication_changed(
510    engine: &Engine,
511    catalog: &cookie_agent_models::catalog::CatalogSnapshot,
512) -> bool {
513    let current = engine.current_runtime();
514    let current = current.models.catalog();
515    current.revision != catalog.revision || current.state.availability != catalog.state.availability
516}
517
518fn data_dir() -> anyhow::Result<PathBuf> {
519    paths::user_data_root().context("determine home directory for cookie agent data")
520}
521
522#[cfg(feature = "tui")]
523async fn run_local_frontend(mut runtime: Runtime) -> anyhow::Result<()> {
524    let (client_stream, server_stream) = in_process_pair(128);
525    let server_task = tokio::spawn(runtime.server.clone().serve_stream(server_stream));
526    let result = async {
527        let client = cookie_agent_tui::Client::connect_stream(client_stream);
528        client.handshake().await.context("handshake with daemon")?;
529        cookie_agent_tui::run_with_new_session(client).await
530    }
531    .await;
532    runtime.server.shutdown();
533    runtime.stop_catalog_refresh().await;
534    let server_result = server_task
535        .await
536        .context("join in-process server task")?
537        .context("run in-process server task");
538    runtime.engine.shutdown().await;
539    result.and(server_result)
540}
541
542#[cfg(not(feature = "tui"))]
543async fn run_local_frontend(mut runtime: Runtime) -> anyhow::Result<()> {
544    runtime.server.shutdown();
545    runtime.stop_catalog_refresh().await;
546    runtime.engine.shutdown().await;
547    anyhow::bail!(
548        "cookie was built without TUI support; use `cookie daemon`, `cookie connect`, or `cookie disconnect`"
549    )
550}
551
552#[cfg(feature = "tui")]
553async fn run_attached_tui(url: &str, token: &str) -> anyhow::Result<()> {
554    validate_websocket_url(url)?;
555    let client = cookie_agent_tui::Client::connect_websocket_with_token(url, token)
556        .await
557        .context("connect to daemon WebSocket")?;
558    client.handshake().await.context("handshake with daemon")?;
559    cookie_agent_tui::run_with_client(client).await
560}
561
562#[cfg(not(feature = "tui"))]
563async fn run_attached_tui(url: &str, _token: &str) -> anyhow::Result<()> {
564    validate_websocket_url(url)?;
565    anyhow::bail!("cookie was built without TUI support")
566}
567
568async fn runtime_snapshot(
569    client: &(impl ClientProtocol + ?Sized),
570) -> anyhow::Result<RuntimeSnapshotResult> {
571    client
572        .runtime_snapshot()
573        .await
574        .context("runtime.snapshot.get failed")
575}
576
577async fn run_connect(url: &str, token: &str, provider_id: Option<String>) -> anyhow::Result<()> {
578    require_interactive_tty(
579        io::stdin().is_terminal(),
580        io::stdout().is_terminal(),
581        io::stderr().is_terminal(),
582        "connect",
583    )?;
584    let client = Client::connect_websocket_with_token(url, token).await?;
585    let mut io = StdioConnectIo;
586    run_connect_with(&client, provider_id, &mut io).await
587}
588
589async fn run_connect_with<I: ConnectIo>(
590    client: &Client,
591    provider_id: Option<String>,
592    io: &mut I,
593) -> anyhow::Result<()> {
594    client.handshake().await.context("handshake with daemon")?;
595    let runtime = runtime_snapshot(client).await?;
596    let provider = choose_provider(&runtime.snapshot.providers, provider_id, io)?;
597    print_provider_details(provider, runtime.snapshot.catalog_revision.as_str(), io)?;
598    ensure_supported(provider)?;
599    io.write_line("Scope: durable, global (available to all workspaces).")?;
600    if !matches!(
601        io.read_public(connect_confirmation_prompt(provider))?
602            .to_ascii_lowercase()
603            .as_str(),
604        "y" | "yes"
605    ) {
606        anyhow::bail!("provider connection cancelled");
607    }
608
609    let setup_values = collect_setup_values(provider, |prompt| io.read_public(prompt))?;
610    let auth_method = choose_auth_method(provider, io)?;
611    let auth_values = collect_auth_values(auth_method, |prompt| io.read_secret(prompt))?;
612    let result: ProviderConnectResult = client
613        .call_sensitive(
614            "provider.connect",
615            SensitiveProviderConnectParams {
616                provider_id: provider.id.clone(),
617                expected_catalog_revision: runtime.snapshot.catalog_revision,
618                setup_values,
619                auth_method: auth_method.id.clone(),
620                auth_values,
621                client_connect_id: ClientConnectId::new(Uuid::now_v7().to_string())
622                    .expect("UUID is a valid connect ID"),
623            },
624        )
625        .await
626        .context("provider.connect failed")?;
627    io.write_line(&format!(
628        "provider.connect succeeded for {} at runtime revision {}{}.",
629        result.durable_connection.provider_id,
630        result.runtime.runtime_revision,
631        if result.replayed { " (replayed)" } else { "" }
632    ))?;
633    Ok(())
634}
635
636async fn run_disconnect(url: &str, token: &str, provider_id: Option<String>) -> anyhow::Result<()> {
637    require_interactive_tty(
638        io::stdin().is_terminal(),
639        io::stdout().is_terminal(),
640        io::stderr().is_terminal(),
641        "disconnect",
642    )?;
643    let client = Client::connect_websocket_with_token(url, token).await?;
644    let mut io = StdioConnectIo;
645    client.handshake().await.context("handshake with daemon")?;
646    let runtime = runtime_snapshot(&client).await?;
647    let provider = choose_provider(&runtime.snapshot.providers, provider_id, &mut io)?;
648    print_provider_details(
649        provider,
650        runtime.snapshot.catalog_revision.as_str(),
651        &mut io,
652    )?;
653    println!("Scope: durable, global (available to all workspaces).");
654    if !prompt_confirmation("Disconnect this managed provider? [y/N] ")? {
655        anyhow::bail!("provider disconnection cancelled");
656    }
657    let result: ProviderDisconnectResult = client
658        .disconnect_provider(ProviderDisconnectParams {
659            provider_id: provider.id.clone(),
660            expected_runtime_revision: runtime.snapshot.runtime_revision,
661            expected_provider_state_revision: runtime.snapshot.provider_state_revision,
662            expected_connection_generation: provider
663                .durable_connection
664                .as_ref()
665                .map(|connection| connection.connection_generation),
666            client_request_id: ClientRequestId::new(Uuid::now_v7().to_string())
667                .expect("UUID is a valid request ID"),
668        })
669        .await
670        .context("provider.disconnect failed")?;
671    println!(
672        "provider.disconnect succeeded for {} at runtime revision {}{}.",
673        result.provider_id,
674        result.runtime.snapshot.runtime_revision,
675        if result.replayed { " (replayed)" } else { "" }
676    );
677    Ok(())
678}
679
680async fn run_mcp(url: &str, token: &str, command: McpCommand) -> anyhow::Result<()> {
681    let client = Client::connect_websocket_with_token(url, token).await?;
682    client.handshake().await.context("handshake with daemon")?;
683    match command {
684        McpCommand::List => {
685            let servers = client
686                .list_mcp_servers()
687                .await
688                .context("mcp.server.list failed")?
689                .servers;
690            if servers.is_empty() {
691                println!("No MCP servers are configured.");
692            }
693            for server in servers {
694                println!(
695                    "{}\t{} tools\t{:?}",
696                    server.name, server.tool_count, server.state
697                );
698            }
699            Ok(())
700        }
701        McpCommand::Auth { server } => {
702            let status = client
703                .list_mcp_servers()
704                .await
705                .context("mcp.server.list failed")?
706                .servers
707                .into_iter()
708                .find(|candidate| candidate.name == server)
709                .with_context(|| format!("unknown MCP server `{server}`"))?;
710            if status.state != McpServerState::NeedsAuth {
711                anyhow::bail!(
712                    "MCP server `{server}` does not currently require OAuth authorization"
713                );
714            }
715            let result = client
716                .begin_mcp_auth(McpAuthBeginParams { server })
717                .await
718                .context("mcp.auth.begin failed")?;
719            println!(
720                "Open this URL to authenticate MCP server `{}`:",
721                result.server
722            );
723            println!("{}", result.authorization_url);
724            Ok(())
725        }
726    }
727}
728
729fn ensure_supported(provider: &ProviderDescriptor) -> anyhow::Result<()> {
730    if provider.support.state != ProviderSupportState::Supported {
731        let reason = provider
732            .support
733            .reason
734            .as_ref()
735            .map_or("unspecified", SafeCode::as_str);
736        anyhow::bail!("provider is not connectable: {reason}");
737    }
738    Ok(())
739}
740
741fn connect_confirmation_prompt(provider: &ProviderDescriptor) -> &'static str {
742    if matches!(
743        provider.configuration,
744        ProviderConfigurationState::Authored | ProviderConfigurationState::AuthoredAndStored
745    ) && provider.effective_auth_state == EffectiveAuthState::Unavailable
746    {
747        "Complete setup and authentication for this authored provider? [y/N] "
748    } else if provider.durable_connection.is_some() {
749        "Reconnect or update this provider? [y/N] "
750    } else {
751        "Connect this provider? [y/N] "
752    }
753}
754
755fn collect_setup_values(
756    provider: &ProviderDescriptor,
757    mut read: impl FnMut(&str) -> anyhow::Result<String>,
758) -> anyhow::Result<BTreeMap<cookie_agent_protocol::SetupFieldId, SafeSetupValue>> {
759    let stored = provider
760        .durable_connection
761        .as_ref()
762        .map(|connection| &connection.setup_values);
763    provider
764        .setup_fields
765        .iter()
766        .filter_map(|field| {
767            let existing = stored.and_then(|values| values.get(&field.id));
768            let fallback = existing.or(field.default.as_ref());
769            let prompt = if let Some(value) = fallback {
770                format!("{} [{}]: ", field.display_name, setup_value_text(value))
771            } else {
772                format!("{}: ", field.display_name)
773            };
774            let answer = match read(&prompt) {
775                Ok(answer) => answer,
776                Err(error) => return Some(Err(error)),
777            };
778            if answer.is_empty() {
779                return fallback
780                    .cloned()
781                    .map(|value| Ok((field.id.clone(), value)))
782                    .or_else(|| {
783                        field.required.then(|| {
784                            Err(anyhow::anyhow!(
785                                "required setup field `{}` was blank",
786                                field.id
787                            ))
788                        })
789                    });
790            }
791            Some(
792                parse_setup_value(field, &answer)
793                    .map_err(anyhow::Error::from)
794                    .map(|value| (field.id.clone(), value)),
795            )
796        })
797        .collect()
798}
799
800fn choose_auth_method<'a>(
801    provider: &'a ProviderDescriptor,
802    io: &mut impl ConnectIo,
803) -> anyhow::Result<&'a AuthMethodDescriptor> {
804    if provider.auth_methods.is_empty() {
805        anyhow::bail!("provider has no supported authentication method");
806    }
807    if provider.auth_methods.len() == 1 {
808        return Ok(&provider.auth_methods[0]);
809    }
810    io.write_line("Authentication methods:")?;
811    for (index, method) in provider.auth_methods.iter().enumerate() {
812        io.write_line(&format!(
813            "  {}. {} ({})",
814            index + 1,
815            method.display_name,
816            method.id
817        ))?;
818    }
819    let answer = io.read_public("Authentication method number or ID: ")?;
820    if let Ok(index) = answer.parse::<usize>()
821        && let Some(method) = index
822            .checked_sub(1)
823            .and_then(|index| provider.auth_methods.get(index))
824    {
825        return Ok(method);
826    }
827    provider
828        .auth_methods
829        .iter()
830        .find(|method| method.id.as_str() == answer)
831        .with_context(|| format!("authentication method `{answer}` was not found"))
832}
833
834fn collect_auth_values(
835    method: &AuthMethodDescriptor,
836    mut read: impl FnMut(&str) -> anyhow::Result<Zeroizing<String>>,
837) -> anyhow::Result<SecretValues> {
838    let mut values = SecretValues::default();
839    for field in &method.credentials {
840        let mut value = read(&format!(
841            "{} (secret, blank does not retain): ",
842            field.display_name
843        ))?;
844        if value.is_empty() {
845            if field.required {
846                anyhow::bail!("required credential field `{}` was blank", field.id);
847            }
848        } else {
849            values
850                .0
851                .insert(field.id.as_str().to_owned(), std::mem::take(&mut *value));
852        }
853    }
854    Ok(values)
855}
856
857fn require_interactive_tty(
858    stdin: bool,
859    stdout: bool,
860    stderr: bool,
861    command: &str,
862) -> anyhow::Result<()> {
863    if !(stdin && stdout && stderr) {
864        anyhow::bail!(
865            "cookie {command} requires an interactive TTY; credentials are never accepted as command-line arguments"
866        );
867    }
868    Ok(())
869}
870
871fn choose_provider<'a>(
872    providers: &'a [ProviderDescriptor],
873    requested: Option<String>,
874    io: &mut impl ConnectIo,
875) -> anyhow::Result<&'a ProviderDescriptor> {
876    if let Some(requested) = requested {
877        return providers
878            .iter()
879            .find(|provider| provider.id.as_str() == requested)
880            .with_context(|| format!("runtime provider `{requested}` was not found"));
881    }
882    if providers.is_empty() {
883        anyhow::bail!("the runtime has no providers");
884    }
885    io.write_line("Runtime providers:")?;
886    for (index, provider) in providers.iter().enumerate() {
887        let support = provider.support.reason.as_ref().map_or_else(
888            || format!("{:?}", provider.support.state).to_ascii_lowercase(),
889            |reason| format!("{:?}: {}", provider.support.state, reason).to_ascii_lowercase(),
890        );
891        io.write_line(&format!(
892            "  {}. {} ({}) — {}",
893            index + 1,
894            provider.display_name,
895            provider.id,
896            support
897        ))?;
898    }
899    let answer = io.read_public("Provider number or ID: ")?;
900    if let Ok(index) = answer.parse::<usize>()
901        && let Some(provider) = index.checked_sub(1).and_then(|index| providers.get(index))
902    {
903        return Ok(provider);
904    }
905    providers
906        .iter()
907        .find(|provider| provider.id.as_str() == answer)
908        .with_context(|| format!("runtime provider `{answer}` was not found"))
909}
910
911fn print_provider_details(
912    provider: &ProviderDescriptor,
913    revision: &str,
914    io: &mut impl ConnectIo,
915) -> anyhow::Result<()> {
916    io.write_line(&format!("Provider ID: {}", provider.id))?;
917    io.write_line(&format!("Name: {}", provider.display_name))?;
918    io.write_line(&format!("Catalog revision: {revision}"))?;
919    io.write_line(&format!("Presence: {:?}", provider.presence))?;
920    io.write_line(&format!("Support: {:?}", provider.support.state))?;
921    if let Some(reason) = &provider.support.reason {
922        io.write_line(&format!("Support reason: {reason}"))?;
923    }
924    io.write_line(&format!("Configuration: {:?}", provider.configuration))?;
925    Ok(())
926}
927
928fn read_public_line(prompt: &str) -> anyhow::Result<String> {
929    print!("{prompt}");
930    io::stdout().flush().context("flush prompt")?;
931    let mut answer = String::new();
932    io::stdin().read_line(&mut answer).context("read input")?;
933    Ok(answer.trim().to_owned())
934}
935
936fn prompt_confirmation(prompt: &str) -> anyhow::Result<bool> {
937    Ok(matches!(
938        read_public_line(prompt)?.to_ascii_lowercase().as_str(),
939        "y" | "yes"
940    ))
941}
942
943#[cfg(unix)]
944fn read_secret_line(prompt: &str) -> anyhow::Result<Zeroizing<String>> {
945    struct EchoGuard(libc::termios);
946    impl Drop for EchoGuard {
947        fn drop(&mut self) {
948            // SAFETY: stdin remains a valid process file descriptor.
949            unsafe {
950                libc::tcsetattr(libc::STDIN_FILENO, libc::TCSAFLUSH, &self.0);
951            }
952        }
953    }
954
955    print!("{prompt}");
956    io::stdout().flush().context("flush credential prompt")?;
957    // SAFETY: initialized by tcgetattr before use.
958    let mut original = unsafe { std::mem::zeroed::<libc::termios>() };
959    // SAFETY: connect verifies stdin is a TTY.
960    if unsafe { libc::tcgetattr(libc::STDIN_FILENO, &mut original) } != 0 {
961        return Err(io::Error::last_os_error()).context("disable credential echo");
962    }
963    let guard = EchoGuard(original);
964    let mut hidden = original;
965    hidden.c_lflag &= !libc::ECHO;
966    // SAFETY: both termios values are initialized.
967    if unsafe { libc::tcsetattr(libc::STDIN_FILENO, libc::TCSAFLUSH, &hidden) } != 0 {
968        return Err(io::Error::last_os_error()).context("disable credential echo");
969    }
970    let mut value = String::new();
971    let read = io::stdin().read_line(&mut value).context("read credential");
972    drop(guard);
973    println!();
974    read?;
975    while matches!(value.as_bytes().last(), Some(b'\n' | b'\r')) {
976        value.pop();
977    }
978    Ok(Zeroizing::new(value))
979}
980
981#[cfg(windows)]
982fn read_secret_line(prompt: &str) -> anyhow::Result<Zeroizing<String>> {
983    use windows_sys::Win32::{
984        Foundation::{HANDLE, INVALID_HANDLE_VALUE},
985        System::Console::{
986            ENABLE_ECHO_INPUT, GetConsoleMode, GetStdHandle, STD_INPUT_HANDLE, SetConsoleMode,
987        },
988    };
989
990    struct ConsoleModeGuard {
991        handle: HANDLE,
992        mode: u32,
993    }
994    impl Drop for ConsoleModeGuard {
995        fn drop(&mut self) {
996            // SAFETY: the console handle and original mode remain valid for this process.
997            unsafe {
998                SetConsoleMode(self.handle, self.mode);
999            }
1000        }
1001    }
1002
1003    print!("{prompt}");
1004    io::stdout().flush().context("flush credential prompt")?;
1005    // SAFETY: GetStdHandle requires no caller-owned pointers.
1006    let handle = unsafe { GetStdHandle(STD_INPUT_HANDLE) };
1007    if handle.is_null() || handle == INVALID_HANDLE_VALUE {
1008        return Err(io::Error::last_os_error()).context("open credential console");
1009    }
1010    let mut original = 0;
1011    // SAFETY: original points to writable mode storage.
1012    if unsafe { GetConsoleMode(handle, &mut original) } == 0 {
1013        return Err(io::Error::last_os_error()).context("read credential console mode");
1014    }
1015    // SAFETY: handle is a console input handle and the mode only disables echo.
1016    if unsafe { SetConsoleMode(handle, original & !ENABLE_ECHO_INPUT) } == 0 {
1017        return Err(io::Error::last_os_error()).context("disable credential echo");
1018    }
1019    let guard = ConsoleModeGuard {
1020        handle,
1021        mode: original,
1022    };
1023    let mut value = String::new();
1024    let read = io::stdin().read_line(&mut value).context("read credential");
1025    drop(guard);
1026    println!();
1027    read?;
1028    while matches!(value.as_bytes().last(), Some(b'\n' | b'\r')) {
1029        value.pop();
1030    }
1031    Ok(Zeroizing::new(value))
1032}
1033
1034async fn run_daemon(mut runtime: Runtime, port_override: Option<u16>) -> anyhow::Result<()> {
1035    let token = generate_token().context("generate daemon authentication token")?;
1036    let port = port_override.unwrap_or(runtime.port);
1037    let listener = match runtime.server.clone().serve(port, token.clone()).await {
1038        Ok(listener) => listener,
1039        Err(error) => {
1040            runtime.server.shutdown();
1041            runtime.stop_catalog_refresh().await;
1042            runtime.engine.shutdown().await;
1043            return Err(anyhow::Error::new(error).context("start WebSocket daemon"));
1044        }
1045    };
1046    let url = format!("ws://{}/ws", listener.address());
1047    println!("{}", ready_line(&url, token.as_str()));
1048    drop(token);
1049    println!(
1050        "cookie daemon listening on {url} (protocol {})",
1051        cookie_agent_protocol::PROTOCOL_VERSION
1052    );
1053    let signal = tokio::signal::ctrl_c().await;
1054    runtime.server.shutdown();
1055    runtime.stop_catalog_refresh().await;
1056    listener.wait().await;
1057    runtime.engine.shutdown().await;
1058    signal.context("wait for daemon shutdown signal")
1059}
1060
1061#[cfg(test)]
1062#[path = "../../../test-support/config_harness.rs"]
1063mod config_harness;
1064
1065#[cfg(all(test, unix))]
1066mod tests;
cookie/main.rs

cookie/
main.rs
