From eb4d1c02f4a6ed367fdbdf5fcf26713731aa101a Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Wed, 11 Mar 2026 13:07:27 -0600 Subject: [PATCH] feat: Support authenticating or refreshing OAuth for supported clients from within the REPL --- README.md | 1 + config.example.yaml | 3 ++- docs/REPL.md | 6 ++++++ docs/clients/CLIENTS.md | 10 ++++++---- src/client/claude.rs | 6 +++++- src/client/common.rs | 4 ++++ src/repl/mod.rs | 22 ++++++++++++++++++++-- 7 files changed, 44 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index e1c2316..026a718 100644 --- a/README.md +++ b/README.md @@ -166,6 +166,7 @@ clients: ```sh loki --authenticate my-claude-oauth +# Or via the REPL: .authenticate ``` For full details, see the [authentication documentation](./docs/clients/CLIENTS.md#authentication). diff --git a/config.example.yaml b/config.example.yaml index aadda92..4b843d2 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -210,7 +210,8 @@ clients: - type: claude api_base: https://api.anthropic.com/v1 # Optional api_key: '{{ANTHROPIC_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault - auth: null # When set to 'oauth', Loki will use OAuth instead of an API key; Authenticate with `loki --authenticate` + auth: null # When set to 'oauth', Loki will use OAuth instead of an API key + # Authenticate with `loki --authenticate` or `.authenticate` in the REPL # See https://docs.mistral.ai/ - type: openai-compatible diff --git a/docs/REPL.md b/docs/REPL.md index 47ffd14..809a02d 100644 --- a/docs/REPL.md +++ b/docs/REPL.md @@ -23,6 +23,7 @@ You can enter the REPL by simply typing `loki` without any follow-up flags or ar - [`.edit` - Modify configuration files](#edit---modify-configuration-files) - [`.delete` - Delete configurations from Loki](#delete---delete-configurations-from-loki) - [`.info` - Display information about the current mode](#info---display-information-about-the-current-mode) + - [`.authenticate` - Authenticate the current model client via OAuth](#authenticate---authenticate-the-current-model-client-via-oauth) - [`.exit` - Exit an agent/role/session/rag or the Loki REPL itself](#exit---exit-an-agentrolesessionrag-or-the-loki-repl-itself) - [`.help` - Show the help guide](#help---show-the-help-guide) @@ -237,6 +238,11 @@ The following entities are supported: | `.info agent` | Display information about the active agent | | `.info rag` | Display information about the active RAG | +### `.authenticate` - Authenticate the current model client via OAuth +The `.authenticate` command will start the OAuth flow for the current model client if +* The client supports OAuth (See the [clients documentation](./clients/CLIENTS.md#providers-that-support-oauth) for supported clients) +* The client is configured in your Loki configuration to use OAuth via the `auth: oauth` property + ### `.exit` - Exit an agent/role/session/rag or the Loki REPL itself The `.exit` command is used to move between modes in the Loki REPL. diff --git a/docs/clients/CLIENTS.md b/docs/clients/CLIENTS.md index 2baac80..576a6df 100644 --- a/docs/clients/CLIENTS.md +++ b/docs/clients/CLIENTS.md @@ -129,15 +129,17 @@ Run the `--authenticate` flag with the client name: loki --authenticate my-claude-oauth ``` -This opens your browser for the OAuth authorization flow. After authorizing, paste the authorization code back into -the terminal. Loki stores the tokens in `~/.cache/loki/oauth` and automatically refreshes them when they expire. - -If you have only one OAuth-configured client, you can omit the name: +Or if you have only one OAuth-configured client, you can omit the name: ```sh loki --authenticate ``` +Alternatively, you can use the REPL command `.authenticate`. + +This opens your browser for the OAuth authorization flow. After authorizing, paste the authorization code back into +the terminal. Loki stores the tokens in `~/.cache/loki/oauth` and automatically refreshes them when they expire. + **Step 3: Use normally** Once authenticated, the client works like any other. Loki uses the stored OAuth tokens automatically: diff --git a/src/client/claude.rs b/src/client/claude.rs index 283be26..cad5ade 100644 --- a/src/client/claude.rs +++ b/src/client/claude.rs @@ -27,7 +27,7 @@ pub struct ClaudeConfig { impl ClaudeClient { config_get_fn!(api_key, get_api_key); config_get_fn!(api_base, get_api_base); - + create_oauth_supported_client_config!(); } @@ -35,6 +35,10 @@ impl ClaudeClient { impl Client for ClaudeClient { client_common_fns!(); + fn supports_oauth(&self) -> bool { + self.config.auth.as_deref() == Some("oauth") + } + async fn chat_completions_inner( &self, client: &ReqwestClient, diff --git a/src/client/common.rs b/src/client/common.rs index c332dd9..202d732 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -47,6 +47,10 @@ pub trait Client: Sync + Send { fn model(&self) -> &Model; + fn supports_oauth(&self) -> bool { + false + } + fn build_client(&self) -> Result { let mut builder = ReqwestClient::builder(); let extra = self.extra_config(); diff --git a/src/repl/mod.rs b/src/repl/mod.rs index a49de40..452a0fd 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -6,7 +6,7 @@ use self::completer::ReplCompleter; use self::highlighter::ReplHighlighter; use self::prompt::ReplPrompt; -use crate::client::{call_chat_completions, call_chat_completions_streaming}; +use crate::client::{call_chat_completions, call_chat_completions_streaming, init_client, oauth}; use crate::config::{ AgentVariables, AssertState, Config, GlobalConfig, Input, LastMessage, StateFlags, macro_execute, @@ -17,6 +17,7 @@ use crate::utils::{ }; use crate::mcp::McpRegistry; +use crate::resolve_oauth_client; use anyhow::{Context, Result, bail}; use crossterm::cursor::SetCursorStyle; use fancy_regex::Regex; @@ -32,10 +33,15 @@ use std::{env, mem, process}; const MENU_NAME: &str = "completion_menu"; -static REPL_COMMANDS: LazyLock<[ReplCommand; 37]> = LazyLock::new(|| { +static REPL_COMMANDS: LazyLock<[ReplCommand; 38]> = LazyLock::new(|| { [ ReplCommand::new(".help", "Show this help guide", AssertState::pass()), ReplCommand::new(".info", "Show system info", AssertState::pass()), + ReplCommand::new( + ".authenticate", + "Authenticate the current model client via OAuth (if configured)", + AssertState::pass(), + ), ReplCommand::new( ".edit config", "Modify configuration file", @@ -421,6 +427,18 @@ pub async fn run_repl_command( } None => println!("Usage: .model "), }, + ".authenticate" => { + let client = init_client(config, None)?; + if !client.supports_oauth() { + bail!( + "Client '{}' doesn't either support OAuth or isn't configured to use it (i.e. uses an API key instead)", + client.name() + ); + } + let clients = config.read().clients.clone(); + let (client_name, provider) = resolve_oauth_client(Some(client.name()), &clients)?; + oauth::run_oauth_flow(&provider, &client_name).await?; + } ".prompt" => match args { Some(text) => { config.write().use_prompt(text)?;