use super::background::BackgroundProcessManager;
use super::error::{ErrorCategory, format_error_for_llm};
use crate::agent::ui::prometheus_display::{ConnectionMode, PrometheusConnectionDisplay};
use crate::analyzer::k8s_optimize::{PrometheusAuth, PrometheusClient};
use rig::completion::ToolDefinition;
use rig::tool::Tool;
use serde::Deserialize;
use serde_json::json;
use std::sync::Arc;
#[derive(Debug, Deserialize)]
pub struct PrometheusConnectArgs {
#[serde(default)]
pub service: Option<String>,
#[serde(default)]
pub namespace: Option<String>,
#[serde(default)]
pub url: Option<String>,
#[serde(default)]
pub port: Option<u16>,
#[serde(default)]
pub auth_type: Option<String>,
#[serde(default)]
pub username: Option<String>,
#[serde(default)]
pub password: Option<String>,
#[serde(default)]
pub token: Option<String>,
}
#[derive(Debug, thiserror::Error)]
#[error("Prometheus connect error: {0}")]
pub struct PrometheusConnectError(String);
#[derive(Clone)]
pub struct PrometheusConnectTool {
bg_manager: Arc<BackgroundProcessManager>,
}
impl PrometheusConnectTool {
pub fn new(bg_manager: Arc<BackgroundProcessManager>) -> Self {
Self { bg_manager }
}
fn validate_port(port: u16) -> Result<(), String> {
if port == 0 {
return Err("Port must be between 1 and 65535 (got 0)".to_string());
}
Ok(())
}
fn validate_url(url: &str) -> Result<(), String> {
let url_lower = url.to_lowercase();
if !url_lower.starts_with("http://") && !url_lower.starts_with("https://") {
return Err(format!(
"URL must start with http:// or https:// (got '{}')",
url
));
}
Ok(())
}
fn build_auth(args: &PrometheusConnectArgs) -> PrometheusAuth {
match args.auth_type.as_deref() {
Some("basic") => {
if let (Some(username), Some(password)) = (&args.username, &args.password) {
PrometheusAuth::Basic {
username: username.clone(),
password: password.clone(),
}
} else {
PrometheusAuth::None
}
}
Some("bearer") => {
if let Some(token) = &args.token {
PrometheusAuth::Bearer(token.clone())
} else {
PrometheusAuth::None
}
}
_ => PrometheusAuth::None,
}
}
async fn test_connection(url: &str, auth: PrometheusAuth) -> bool {
match PrometheusClient::with_auth(url, auth) {
Ok(client) => client.is_available().await,
Err(_) => false,
}
}
}
impl Tool for PrometheusConnectTool {
const NAME: &'static str = "prometheus_connect";
type Args = PrometheusConnectArgs;
type Output = String;
type Error = PrometheusConnectError;
async fn definition(&self, _prompt: String) -> ToolDefinition {
ToolDefinition {
name: Self::NAME.to_string(),
description: r#"Connect to Prometheus for K8s optimization analysis.
**Use after prometheus_discover or when user provides a URL.**
**Connection Methods (in order of preference):**
1. **Port-forward** (recommended) - NO authentication needed
- Provide: service, namespace, port
- Starts kubectl port-forward in background
- Direct pod connection bypasses auth
2. **External URL** - May require authentication
- Provide: url
- Optional: auth_type, username/password or token
**Examples:**
Port-forward (no auth):
```json
{"service": "prometheus-server", "namespace": "monitoring", "port": 9090}
```
External URL without auth:
```json
{"url": "http://prometheus.example.com"}
```
External URL with basic auth:
```json
{"url": "https://prometheus.example.com", "auth_type": "basic", "username": "admin", "password": "secret"}
```
**Returns:**
- Connection URL for use with k8s_optimize
- Connection mode (port-forward or direct)
- Local port (if port-forward)"#
.to_string(),
parameters: json!({
"type": "object",
"properties": {
"service": {
"type": "string",
"description": "Kubernetes service name (for port-forward)"
},
"namespace": {
"type": "string",
"description": "Kubernetes namespace (for port-forward)"
},
"url": {
"type": "string",
"description": "External Prometheus URL (alternative to port-forward)"
},
"port": {
"type": "integer",
"description": "Target port (default: 9090)"
},
"auth_type": {
"type": "string",
"description": "Authentication type for external URL: 'none', 'basic', 'bearer'",
"enum": ["none", "basic", "bearer"]
},
"username": {
"type": "string",
"description": "Username for basic auth (only for external URL)"
},
"password": {
"type": "string",
"description": "Password for basic auth (only for external URL)"
},
"token": {
"type": "string",
"description": "Bearer token (only for external URL)"
}
}
}),
}
}
async fn call(&self, args: Self::Args) -> Result<Self::Output, Self::Error> {
if let Some(port) = args.port {
if let Err(e) = Self::validate_port(port) {
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::ValidationFailed,
&e,
Some(vec![
"Port must be a valid TCP port between 1 and 65535",
"Common Prometheus port is 9090 (default if not specified)",
]),
));
}
}
if let Some(ref url) = args.url {
if let Err(e) = Self::validate_url(url) {
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::ValidationFailed,
&e,
Some(vec![
"URL must start with http:// or https://",
"Example: http://prometheus.example.com or https://prometheus.example.com",
]),
));
}
}
let target_port = args.port.unwrap_or(9090);
if let (Some(service), Some(namespace)) = (&args.service, &args.namespace) {
let resource = format!("svc/{}", service);
let display = PrometheusConnectionDisplay::new(ConnectionMode::PortForward);
let target = format!("{}/{}", namespace, service);
display.start(&target);
match self
.bg_manager
.start_port_forward("prometheus-port-forward", &resource, namespace, target_port)
.await
{
Ok(local_port) => {
let url = format!("http://localhost:{}", local_port);
display.port_forward_established(local_port, service, namespace);
tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
display.testing_connection();
let mut connected = false;
for attempt in 0..6 {
if Self::test_connection(&url, PrometheusAuth::None).await {
connected = true;
break;
}
let delay = match attempt {
0 | 1 => 1000,
2 | 3 => 2000,
_ => 3000,
};
tokio::time::sleep(tokio::time::Duration::from_millis(delay)).await;
}
if connected {
display.connected(&url, false);
display.background_process_info("prometheus-port-forward");
display.ready_for_use(&url);
let response = json!({
"connected": true,
"url": url,
"mode": "port-forward",
"local_port": local_port,
"service": service,
"namespace": namespace,
"process_id": "prometheus-port-forward",
"note": "Port-forward established. No authentication needed.",
"usage": {
"k8s_optimize": {
"prometheus": url
}
}
});
return Ok(serde_json::to_string_pretty(&response)
.unwrap_or_else(|_| "{}".to_string()));
} else {
let _ = self.bg_manager.stop("prometheus-port-forward").await;
display.connection_failed(
"Port-forward started but Prometheus not responding",
&[
"Verify the service is correct",
"Check if Prometheus pod is running",
"The service might need more time to start",
],
);
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::NetworkError,
"Port-forward started but Prometheus not responding",
Some(vec![
&format!(
"Verify the service is correct: kubectl get svc -n {}",
namespace
),
&format!(
"Check if Prometheus pod is running: kubectl get pods -n {} | grep prometheus",
namespace
),
"The service might need more time to start - try again in a few seconds",
]),
));
}
}
Err(e) => {
display.connection_failed(
&format!("Port-forward failed: {}", e),
&[
"Check if kubectl is configured correctly",
"Verify the service exists",
"Try providing an external URL instead",
],
);
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::ExternalCommandFailed,
&format!("Port-forward failed: {}", e),
Some(vec![
"Check if kubectl is configured correctly: kubectl config current-context",
&format!(
"Verify the service exists: kubectl get svc -n {}",
namespace
),
"Try providing an external URL instead",
]),
));
}
}
}
if let Some(url) = &args.url {
let display = PrometheusConnectionDisplay::new(ConnectionMode::DirectUrl);
display.start(url);
display.testing_connection();
if Self::test_connection(url, PrometheusAuth::None).await {
display.connected(url, false);
display.ready_for_use(url);
let response = json!({
"connected": true,
"url": url,
"mode": "direct",
"authenticated": false,
"note": "Connected without authentication",
"usage": {
"k8s_optimize": {
"prometheus": url
}
}
});
return Ok(
serde_json::to_string_pretty(&response).unwrap_or_else(|_| "{}".to_string())
);
}
let auth = Self::build_auth(&args);
if !matches!(auth, PrometheusAuth::None) && Self::test_connection(url, auth).await {
display.connected(url, true);
display.ready_for_use(url);
let response = json!({
"connected": true,
"url": url,
"mode": "direct",
"authenticated": true,
"auth_type": args.auth_type,
"note": "Connected with authentication",
"usage": {
"k8s_optimize": {
"prometheus": url,
"auth_type": args.auth_type,
"username": args.username,
}
}
});
return Ok(
serde_json::to_string_pretty(&response).unwrap_or_else(|_| "{}".to_string())
);
}
if args.auth_type.is_none() {
display.auth_required();
display.connection_failed(
"Connection failed - URL may require authentication",
&[
"Try with auth_type='basic' and username/password",
"Or try auth_type='bearer' with a token",
"Verify the URL is correct and accessible",
],
);
let test_url_suggestion =
format!("Test URL manually: curl -s {}/api/v1/status/config", url);
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::NetworkError,
"Connection failed - URL may require authentication",
Some(vec![
"Try with auth_type='basic' and username/password",
"Or try auth_type='bearer' with a token",
"Verify the URL is correct and accessible",
&test_url_suggestion,
]),
));
} else {
display.connection_failed(
"Connection failed - authentication credentials may be incorrect",
&[
"Verify the username/password or token",
"Check if the auth_type matches what the server expects",
"Ensure the user has permission to access Prometheus API",
],
);
return Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::NetworkError,
"Connection failed - authentication credentials may be incorrect",
Some(vec![
"Verify the username/password or token",
"Check if the auth_type matches what the server expects",
"Ensure the user has permission to access Prometheus API",
]),
));
}
}
Ok(format_error_for_llm(
"prometheus_connect",
ErrorCategory::ValidationFailed,
"No service or URL provided",
Some(vec![
"Provide service + namespace for port-forward: {\"service\": \"prometheus-server\", \"namespace\": \"monitoring\"}",
"Or provide url for external Prometheus: {\"url\": \"http://prometheus.example.com\"}",
"Use prometheus_discover to find available Prometheus instances",
]),
))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tool_name() {
assert_eq!(PrometheusConnectTool::NAME, "prometheus_connect");
}
#[test]
fn test_validate_port_valid() {
assert!(PrometheusConnectTool::validate_port(9090).is_ok());
assert!(PrometheusConnectTool::validate_port(1).is_ok());
assert!(PrometheusConnectTool::validate_port(65535).is_ok());
}
#[test]
fn test_validate_port_invalid() {
let result = PrometheusConnectTool::validate_port(0);
assert!(result.is_err());
assert!(
result
.unwrap_err()
.contains("Port must be between 1 and 65535")
);
}
#[test]
fn test_validate_url_valid() {
assert!(PrometheusConnectTool::validate_url("http://prometheus.example.com").is_ok());
assert!(PrometheusConnectTool::validate_url("https://prometheus.example.com").is_ok());
assert!(PrometheusConnectTool::validate_url("HTTP://PROMETHEUS.EXAMPLE.COM").is_ok());
assert!(PrometheusConnectTool::validate_url("HTTPS://prometheus.example.com").is_ok());
}
#[test]
fn test_validate_url_invalid() {
let result = PrometheusConnectTool::validate_url("prometheus.example.com");
assert!(result.is_err());
assert!(
result
.unwrap_err()
.contains("must start with http:// or https://")
);
let result = PrometheusConnectTool::validate_url("ftp://prometheus.example.com");
assert!(result.is_err());
let result = PrometheusConnectTool::validate_url("/api/v1/query");
assert!(result.is_err());
}
#[tokio::test]
async fn test_missing_service_and_url_error() {
let bg_manager = Arc::new(BackgroundProcessManager::new());
let tool = PrometheusConnectTool::new(bg_manager);
let args = PrometheusConnectArgs {
service: None,
namespace: None,
url: None,
port: None,
auth_type: None,
username: None,
password: None,
token: None,
};
let result = tool.call(args).await.unwrap();
assert!(result.contains("\"error\": true"));
assert!(result.contains("VALIDATION_FAILED"));
assert!(result.contains("No service or URL provided"));
assert!(result.contains("suggestions"));
}
#[tokio::test]
async fn test_invalid_port_validation() {
let bg_manager = Arc::new(BackgroundProcessManager::new());
let tool = PrometheusConnectTool::new(bg_manager);
let args = PrometheusConnectArgs {
service: Some("prometheus".to_string()),
namespace: Some("monitoring".to_string()),
url: None,
port: Some(0), auth_type: None,
username: None,
password: None,
token: None,
};
let result = tool.call(args).await.unwrap();
assert!(result.contains("\"error\": true"));
assert!(result.contains("VALIDATION_FAILED"));
assert!(result.contains("Port must be between 1 and 65535"));
}
#[tokio::test]
async fn test_malformed_url_validation() {
let bg_manager = Arc::new(BackgroundProcessManager::new());
let tool = PrometheusConnectTool::new(bg_manager);
let args = PrometheusConnectArgs {
service: None,
namespace: None,
url: Some("prometheus.example.com".to_string()), port: None,
auth_type: None,
username: None,
password: None,
token: None,
};
let result = tool.call(args).await.unwrap();
assert!(result.contains("\"error\": true"));
assert!(result.contains("VALIDATION_FAILED"));
assert!(result.contains("must start with http:// or https://"));
assert!(result.contains("suggestions"));
}
#[test]
fn test_build_auth_none() {
let args = PrometheusConnectArgs {
service: None,
namespace: None,
url: Some("http://localhost".to_string()),
port: None,
auth_type: None,
username: None,
password: None,
token: None,
};
let auth = PrometheusConnectTool::build_auth(&args);
assert!(matches!(auth, PrometheusAuth::None));
}
#[test]
fn test_build_auth_basic() {
let args = PrometheusConnectArgs {
service: None,
namespace: None,
url: Some("http://localhost".to_string()),
port: None,
auth_type: Some("basic".to_string()),
username: Some("admin".to_string()),
password: Some("secret".to_string()),
token: None,
};
let auth = PrometheusConnectTool::build_auth(&args);
match auth {
PrometheusAuth::Basic { username, password } => {
assert_eq!(username, "admin");
assert_eq!(password, "secret");
}
_ => panic!("Expected Basic auth"),
}
}
#[test]
fn test_build_auth_bearer() {
let args = PrometheusConnectArgs {
service: None,
namespace: None,
url: Some("http://localhost".to_string()),
port: None,
auth_type: Some("bearer".to_string()),
username: None,
password: None,
token: Some("mytoken".to_string()),
};
let auth = PrometheusConnectTool::build_auth(&args);
match auth {
PrometheusAuth::Bearer(token) => {
assert_eq!(token, "mytoken");
}
_ => panic!("Expected Bearer auth"),
}
}
}