Merge pull request #98 from ovh/feature/agent-config

Marlinski · web-flow · commit 08bbf61756bf · 2025-11-05T15:06:14.000+01:00
easier agent configuration
diff --git a/.ovh.config b/.ovh.config
@@ -1,14 +1,6 @@
 {
   "name": "ovh",
   "description": "OVH agent with OVH MCP server for cloud management and API calls",
-  "llm_provider": {
-    "provider": "ovhcloud",
-    "env_vars": {
-      "OVH_BASE_URL": "https://gpt-oss-120b.endpoints.kepler.ai.cloud.ovh.net/api/openai_compat/v1"
-    },
-    "model": "gpt-oss-120b",
-    "tool_method": "FunctionCall"
-  },
   "tools": {
     "builtin": ["*"],
     "mcp": {
@@ -17,7 +9,10 @@
           "type": "http",
           "url": "https://mcp.eu.ovhcloud.com/mcp"
         },
-        "enabled_tools": ["*"]
+        "enabled_tools": ["*"],
+        "excluded_tools": [
+          "get-cloud-project-flavor-list"
+        ]
       }
     }
   },
diff --git a/README.md b/README.md
@@ -137,16 +137,17 @@ Instead of a single global configuration, you can create custom agent in a separ
 
 [`.ovh.config`](./.ovh.config) contains an example of a custom configuration with an remote MCP server configured.
 
-Place this file in `~/.config/shai/agents/example.config`, you can then list the agents available with:
+Place this file in `~/.config/shai/agents/ovh.config`, you can then list the agents available with:
 
 ```bash
+curl https://raw.githubusercontent.com/ovh/shai/refs/heads/main/.ovh.config -o ˜/.config/shai/agents/ovh.config
 shai agent list
 ```
 
 You can run shai with this specific agent with the `agent` subcommand:
 
 ```bash
-shai agent example
+shai agent ovh
 ```
 
 ### OVHCloud Endpoints
diff --git a/shai-core/examples/oauth_test.rs b/shai-core/examples/oauth_test.rs
@@ -3,12 +3,22 @@ use shai_core::tools::mcp::mcp_oauth::signin_oauth;
 #[tokio::main]
 async fn main() {
     println!("🚀 Starting OAuth flow test...");
-    
+
     match signin_oauth("https://mcp.eu.ovhcloud.com/").await {
-        Ok(access_token) => {
+        Ok(token) => {
             println!("✅ OAuth flow completed successfully!");
-            println!("🎫 Access Token: {}", access_token);
-            println!("🔑 Token length: {} characters", access_token.len());
+            println!("🎫 Access Token: {}", token.access_token);
+            println!("🔑 Token length: {} characters", token.access_token.len());
+            if let Some(expires_at) = token.expires_at {
+                let now = std::time::SystemTime::now()
+                    .duration_since(std::time::UNIX_EPOCH)
+                    .unwrap()
+                    .as_secs() as i64;
+                let seconds_until_expiry = expires_at - now;
+                println!("⏰ Token expires in {} seconds", seconds_until_expiry);
+            } else {
+                println!("⏰ Token has no expiration");
+            }
         }
         Err(e) => {
             println!("❌ OAuth flow failed: {}", e);
diff --git a/shai-core/src/agent/builder.rs b/shai-core/src/agent/builder.rs
@@ -146,7 +146,7 @@ impl AgentBuilder {
             LlmClient::create_provider(&config.llm_provider.provider, &config.llm_provider.env_vars)
                 .map_err(|e| AgentError::LlmError(e.to_string()))?
         );
-        
+
         // Create brain with custom system prompt and temperature
         let brain = Box::new(CoderBrain::with_custom_prompt(
             llm_client.clone(),
@@ -275,43 +275,57 @@ impl AgentBuilder {
     /// Handle OAuth flow for MCP connections if needed
     async fn mcp_check_oauth(mcp_name: &str, mcp_config: &mut McpConfig) -> Result<bool, AgentError> {
         use crate::tools::mcp::McpConfig;
-        
         let mut config_changed = false;
-        
+
         // Only handle HTTP configs that might need OAuth
-        if let McpConfig::Http { url, bearer_token } = mcp_config {
-            // Test connection with current config
-            let test_config = McpConfig::Http { 
-                url: url.clone(), 
-                bearer_token: bearer_token.clone() 
-            };
-            let mut test_client = create_mcp_client(test_config);
-            match test_client.connect().await {
-                Ok(_) => {
-                    if bearer_token.is_some() {
+        if let McpConfig::Http { url, auth } = mcp_config {
+            let needs_new_token = match auth {
+                Some(token) if token.is_expired() => {
+                    eprintln!("\x1b[2m░ MCP '{}' token expired, refreshing...\x1b[0m", mcp_name);
+                    true
+                }
+                Some(_) => {
+                    // Test connection with existing token
+                    let test_config = McpConfig::Http { url: url.clone(), auth: auth.clone() };
+                    let mut test_client = create_mcp_client(test_config);
+                    if test_client.connect().await.is_ok() {
                         eprintln!("\x1b[2m░ MCP '{}' connected (authenticated)\x1b[0m", mcp_name);
+                        false
                     } else {
-                        eprintln!("\x1b[2m░ MCP '{}' connected (no auth)\x1b[0m", mcp_name);
+                        eprintln!("\x1b[2m░ MCP '{}' authentication failed, refreshing token...\x1b[0m", mcp_name);
+                        true
                     }
                 }
-                Err(_) => {
-                    eprintln!("\x1b[2m░ MCP '{}' connection failed, starting OAuth flow...\x1b[0m", mcp_name);
-                    let url_clone = url.clone();
-                    match signin_oauth(&url_clone).await {
-                        Ok(token) => {
-                            eprintln!("\x1b[2m░ MCP '{}' connected (OAuth successful)\x1b[0m", mcp_name);
-                            *bearer_token = Some(token);
-                            config_changed = true;
-                        }
-                        Err(e) => {
-                            return Err(AgentError::ConfigurationError(format!("OAuth failed for MCP '{}': {}", mcp_name, e)));
-                        }
+                None => {
+                    // Test connection without auth
+                    let test_config = McpConfig::Http { url: url.clone(), auth: None };
+                    let mut test_client = create_mcp_client(test_config);
+                    if test_client.connect().await.is_ok() {
+                        eprintln!("\x1b[2m░ MCP '{}' connected (no auth required)\x1b[0m", mcp_name);
+                        false
+                    } else {
+                        eprintln!("\x1b[2m░ MCP '{}' requires authentication, starting OAuth flow...\x1b[0m", mcp_name);
+                        true
+                    }
+                }
+            };
+
+            if needs_new_token {
+                let url_clone = url.clone();
+                match signin_oauth(&url_clone).await {
+                    Ok(token) => {
+                        eprintln!("\x1b[2m░ MCP '{}' OAuth successful\x1b[0m", mcp_name);
+                        *auth = Some(token);
+                        config_changed = true;
+                    }
+                    Err(e) => {
+                        return Err(AgentError::ConfigurationError(format!("OAuth failed for MCP '{}': {}", mcp_name, e)));
                     }
                 }
             }
         }
         // SSE and Stdio don't need OAuth handling for now
-        
+
         Ok(config_changed)
     }
 }
diff --git a/shai-core/src/config/agent.rs b/shai-core/src/config/agent.rs
@@ -3,6 +3,7 @@ use std::path::PathBuf;
 use serde::{Serialize, Deserialize};
 use shai_llm::ToolCallMethod;
 use crate::tools::mcp::McpConfig;
+use super::config::ShaiConfig;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct AgentProviderConfig {
@@ -35,6 +36,7 @@ pub struct AgentTools {
 pub struct AgentConfig {
     pub name: String,
     pub description: String,
+    #[serde(default = "default_llm_provider")]
     pub llm_provider: AgentProviderConfig,
     #[serde(default)]
     pub tools: AgentTools,
@@ -46,6 +48,23 @@ pub struct AgentConfig {
     pub temperature: f32,
 }
 
+fn default_llm_provider() -> AgentProviderConfig {
+    // Load the default provider from ShaiConfig
+    let shai_config = ShaiConfig::load()
+        .unwrap_or_else(|_| ShaiConfig::default());
+
+    let provider_config = shai_config
+        .get_selected_provider()
+        .expect("No provider configured in default config");
+
+    AgentProviderConfig {
+        provider: provider_config.provider.clone(),
+        env_vars: provider_config.env_vars.clone(),
+        model: provider_config.model.clone(),
+        tool_method: provider_config.tool_method.clone(),
+    }
+}
+
 fn default_system_prompt() -> String {
     "{{CODER_BASE_PROMPT}}".to_string()
 }
diff --git a/shai-core/src/config/config.rs b/shai-core/src/config/config.rs
@@ -199,8 +199,8 @@ impl ShaiConfig {
             .map(|(name, config)| {
                 let description = match config {
                     McpConfig::Stdio { command, .. } => format!("stdio: {}", command),
-                    McpConfig::Http { url, bearer_token } => {
-                        if bearer_token.is_some() {
+                    McpConfig::Http { url, auth } => {
+                        if auth.is_some() {
                             format!("http: {} (authenticated)", url)
                         } else {
                             format!("http: {}", url)
diff --git a/shai-core/src/tools/mcp/mcp_config.rs b/shai-core/src/tools/mcp/mcp_config.rs
@@ -3,24 +3,54 @@ use serde::{Serialize, Deserialize};
 
 use super::{StdioClient, HttpClient, SseClient};
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OAuthToken {
+    pub access_token: String,
+    /// Unix timestamp (seconds since epoch) when the token expires
+    pub expires_at: Option<i64>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "type")]
 pub enum McpConfig {
     #[serde(rename = "stdio")]
     Stdio { command: String, args: Vec<String> },
     #[serde(rename = "http")]
-    Http { url: String, bearer_token: Option<String> },
+    Http {
+        url: String,
+        #[serde(flatten)]
+        auth: Option<OAuthToken>
+    },
     #[serde(rename = "sse")]
     Sse { url: String },
 }
 
+impl OAuthToken {
+    /// Check if the token is expired or will expire within the next 60 seconds
+    pub fn is_expired(&self) -> bool {
+        if let Some(expires_at) = self.expires_at {
+            let now = std::time::SystemTime::now()
+                .duration_since(std::time::UNIX_EPOCH)
+                .unwrap()
+                .as_secs() as i64;
+
+            // Consider expired if it expires within the next 60 seconds (safety margin)
+            expires_at <= now + 60
+        } else {
+            // If no expiration time, assume it's still valid
+            false
+        }
+    }
+}
+
 /// Factory function to create an MCP client from configuration
 pub fn create_mcp_client(config: McpConfig) -> Box<dyn McpClient> {
     match config {
         McpConfig::Stdio { command, args } => {
             Box::new(StdioClient::new(command, args))
         }
-        McpConfig::Http { url, bearer_token } => {
+        McpConfig::Http { url, auth } => {
+            let bearer_token = auth.map(|t| t.access_token);
             Box::new(HttpClient::new_with_auth(url, bearer_token))
         }
         McpConfig::Sse { url } => {
diff --git a/shai-core/src/tools/mcp/mcp_oauth.rs b/shai-core/src/tools/mcp/mcp_oauth.rs
@@ -1,6 +1,6 @@
 use oauth2::{
     AuthUrl, TokenUrl, ClientId, ClientSecret, RedirectUrl, CsrfToken,
-    AuthorizationCode, PkceCodeChallenge, Scope, 
+    AuthorizationCode, PkceCodeChallenge, Scope,
     basic::BasicClient, reqwest::async_http_client, TokenResponse,
     AuthType, url::Url,
 };
@@ -9,6 +9,7 @@ use std::sync::{Arc, Mutex};
 use reqwest;
 use serde::{Deserialize, Serialize};
 use tokio::net::TcpListener;
+use super::mcp_config::OAuthToken;
 
 #[derive(Serialize)]
 struct ClientRegistrationRequest {
@@ -24,7 +25,7 @@ struct ClientRegistrationResponse {
     client_secret: Option<String>,
 }
 
-pub async fn signin_oauth(base_url: &str) -> anyhow::Result<String> {
+pub async fn signin_oauth(base_url: &str) -> anyhow::Result<OAuthToken> {
     // Extract the root domain for .well-known endpoint (OAuth standard)
     let url = Url::parse(base_url)?;
     let root_url = format!("{}://{}", url.scheme(), url.host_str().unwrap_or(""));
@@ -221,5 +222,17 @@ pub async fn signin_oauth(base_url: &str) -> anyhow::Result<String> {
         .request_async(async_http_client)
         .await?;
 
-    Ok(token_response.access_token().secret().to_string())
+    // Calculate expiration time
+    let expires_at = token_response.expires_in().map(|duration| {
+        let now = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap()
+            .as_secs() as i64;
+        now + duration.as_secs() as i64
+    });
+
+    Ok(OAuthToken {
+        access_token: token_response.access_token().secret().to_string(),
+        expires_at,
+    })
 }
diff --git a/shai-core/src/tools/mcp/mod.rs b/shai-core/src/tools/mcp/mod.rs
@@ -9,7 +9,7 @@ pub mod mcp_oauth;
 mod tests;
 
 pub use mcp::{McpClient, McpToolDescription, get_mcp_tools};
-pub use mcp_config::{McpConfig, create_mcp_client};
+pub use mcp_config::{McpConfig, OAuthToken, create_mcp_client};
 pub use mcp_stdio::StdioClient;
 pub use mcp_http::HttpClient;
 pub use mcp_sse::SseClient;
diff --git a/shai-core/src/tools/mcp/tests.rs b/shai-core/src/tools/mcp/tests.rs
@@ -148,7 +148,7 @@ mod tests {
         // Test HTTP config
         let http_config = McpConfig::Http {
             url: "http://localhost:8080".to_string(),
-            bearer_token: None
+            auth: None
         };
         let _http_client = create_mcp_client(http_config);
         println!("✅ Successfully created HttpClient via factory");
